525 lines
17 KiB
Python
525 lines
17 KiB
Python
import cv2
|
||
import numpy as np
|
||
import os
|
||
import subprocess
|
||
import shutil
|
||
from datetime import timedelta
|
||
import argparse
|
||
from sklearn.metrics.pairwise import cosine_similarity
|
||
from skimage.metrics import structural_similarity as ssim
|
||
|
||
# 设置固定的输入输出路径
|
||
INPUT_VIDEO_PATH = "/root/autodl-tmp/hot_video_analyse/source/sample_demo_1.mp4" # 修改为视频文件夹路径
|
||
OUTPUT_DIR = "/root/autodl-tmp/hot_video_analyse/source/Splitter" # 请修改为您的实际输出目录路径
|
||
|
||
# 支持的视频格式
|
||
VIDEO_EXTENSIONS = ['.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv']
|
||
|
||
# 设置参数
|
||
SAMPLE_RATE = 1 # 帧采样率
|
||
METHOD = "ssim" # 比较方法，可选 "ssim" 或 "cosine"
|
||
## 有分镜0.2 无分镜0.7
|
||
THRESHOLD = 0.8 # 相似度阈值
|
||
|
||
# 启用详细日志输出
|
||
VERBOSE = True
|
||
|
||
# FFMPEG可能的路径
|
||
FFMPEG_PATHS = [
|
||
'ffmpeg', # 系统路径
|
||
'/usr/bin/ffmpeg', # 常见Linux路径
|
||
'/usr/local/bin/ffmpeg', # 常见macOS路径
|
||
'C:\\ffmpeg\\bin\\ffmpeg.exe', # 常见Windows路径
|
||
]
|
||
|
||
def find_ffmpeg():
|
||
"""查找系统中可用的ffmpeg路径"""
|
||
# 首先尝试使用which/where命令
|
||
try:
|
||
if os.name == 'nt': # Windows
|
||
result = subprocess.run(['where', 'ffmpeg'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||
if result.returncode == 0:
|
||
return result.stdout.strip().split('\n')[0]
|
||
else: # Linux/Mac
|
||
result = subprocess.run(['which', 'ffmpeg'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||
if result.returncode == 0:
|
||
return result.stdout.strip()
|
||
except Exception:
|
||
pass
|
||
|
||
# 然后检查预定义的路径
|
||
for path in FFMPEG_PATHS:
|
||
if shutil.which(path):
|
||
return path
|
||
|
||
return None
|
||
|
||
def extract_frames(video_path, output_dir, sample_rate=1):
|
||
"""
|
||
从视频中提取帧并保存到指定目录
|
||
|
||
Args:
|
||
video_path: 视频文件路径
|
||
output_dir: 输出帧的目录
|
||
sample_rate: 采样率（每N帧提取一帧）
|
||
|
||
Returns:
|
||
frames_info: 包含帧信息的列表 [(frame_number, timestamp, frame_path), ...]
|
||
"""
|
||
if not os.path.exists(output_dir):
|
||
os.makedirs(output_dir)
|
||
|
||
# 打开视频
|
||
cap = cv2.VideoCapture(video_path)
|
||
|
||
# 获取视频属性
|
||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
duration = frame_count / fps
|
||
|
||
print(f"视频信息：{frame_count}帧, {fps}fps, 时长:{timedelta(seconds=duration)}")
|
||
|
||
frames_info = []
|
||
frame_number = 0
|
||
saved_count = 0
|
||
|
||
while True:
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
break
|
||
|
||
if frame_number % sample_rate == 0:
|
||
# 计算时间戳（秒）
|
||
timestamp = frame_number / fps
|
||
# 保存帧
|
||
frame_path = os.path.join(output_dir, f"frame_{saved_count:05d}.jpg")
|
||
cv2.imwrite(frame_path, frame)
|
||
# 记录帧信息
|
||
frames_info.append((frame_number, timestamp, frame_path))
|
||
saved_count += 1
|
||
|
||
frame_number += 1
|
||
|
||
# 显示进度
|
||
if frame_number % 100 == 0:
|
||
print(f"处理进度: {frame_number}/{frame_count} ({frame_number/frame_count*100:.2f}%)")
|
||
|
||
cap.release()
|
||
print(f"共提取了 {saved_count} 帧")
|
||
return frames_info
|
||
|
||
def compare_frames(frame1_path, frame2_path, method='ssim', threshold=0.85):
|
||
"""
|
||
比较两帧的相似度
|
||
|
||
Args:
|
||
frame1_path: 第一帧路径
|
||
frame2_path: 第二帧路径
|
||
method: 比较方法，'ssim'或'cosine'
|
||
threshold: 相似度阈值
|
||
|
||
Returns:
|
||
is_similar: 是否相似
|
||
similarity: 相似度值
|
||
"""
|
||
# 读取帧
|
||
frame1 = cv2.imread(frame1_path)
|
||
frame2 = cv2.imread(frame2_path)
|
||
|
||
# 调整大小以加快处理
|
||
frame1_resized = cv2.resize(frame1, (320, 180))
|
||
frame2_resized = cv2.resize(frame2, (320, 180))
|
||
|
||
# 转换为灰度图
|
||
gray1 = cv2.cvtColor(frame1_resized, cv2.COLOR_BGR2GRAY)
|
||
gray2 = cv2.cvtColor(frame2_resized, cv2.COLOR_BGR2GRAY)
|
||
|
||
if method == 'ssim':
|
||
# 使用结构相似性指数
|
||
similarity, _ = ssim(gray1, gray2, full=True)
|
||
else: # cosine
|
||
# 使用余弦相似度
|
||
flat1 = gray1.flatten().reshape(1, -1)
|
||
flat2 = gray2.flatten().reshape(1, -1)
|
||
similarity = cosine_similarity(flat1, flat2)[0][0]
|
||
|
||
is_similar = similarity >= threshold
|
||
return is_similar, similarity
|
||
|
||
def detect_scene_changes(frames_info, method='ssim', threshold=0.85):
|
||
"""
|
||
检测场景变化
|
||
|
||
Args:
|
||
frames_info: 帧信息列表
|
||
method: 比较方法
|
||
threshold: 相似度阈值
|
||
|
||
Returns:
|
||
scenes: 场景信息列表 [(start_frame, end_frame, start_time, end_time), ...]
|
||
"""
|
||
global SCENE_START_FRAMES # 声明使用全局变量
|
||
|
||
if len(frames_info) < 2:
|
||
return []
|
||
|
||
scenes = []
|
||
clips = []
|
||
scene_start = frames_info[0]
|
||
|
||
for i in range(1, len(frames_info)):
|
||
prev_frame_path = frames_info[i-1][2]
|
||
curr_frame_path = frames_info[i][2]
|
||
|
||
is_similar, similarity = compare_frames(
|
||
prev_frame_path, curr_frame_path, method, threshold
|
||
)
|
||
|
||
if not is_similar:
|
||
# 场景变化，记录上一个场景
|
||
scene_end = frames_info[i-1]
|
||
scene_duration = scene_end[1] - scene_start[1]
|
||
clips.append(scene_start[0])
|
||
# 只记录持续时间超过1秒的场景
|
||
if scene_duration >= 0.2:
|
||
scenes.append((
|
||
scene_start[0], # 开始帧号
|
||
scene_end[0], # 结束帧号
|
||
scene_start[1], # 开始时间
|
||
scene_end[1], # 结束时间
|
||
scene_duration # 持续时间
|
||
))
|
||
|
||
# 开始新场景
|
||
scene_start = frames_info[i]
|
||
|
||
if VERBOSE:
|
||
print(f"检测到场景变化点: 帧 {scene_end[0]}, 时间 {timedelta(seconds=scene_end[1])}, 相似度: {similarity:.4f}")
|
||
|
||
# 添加最后一个场景
|
||
scene_end = frames_info[-1]
|
||
scene_duration = scene_end[1] - scene_start[1]
|
||
if scene_duration >= 0.2:
|
||
scenes.append((
|
||
scene_start[0],
|
||
scene_end[0],
|
||
scene_start[1],
|
||
scene_end[1],
|
||
scene_duration
|
||
))
|
||
|
||
print(f"\n场景变化统计:")
|
||
print(f"检测到 {len(scenes)} 个场景, 平均时长: {sum(s[4] for s in scenes)/max(1, len(scenes)):.2f}秒")
|
||
return scenes , clips
|
||
|
||
def extract_video_clips(video_path, scenes, output_dir, ffmpeg_path=None):
|
||
"""
|
||
根据场景信息提取视频片段
|
||
|
||
Args:
|
||
video_path: 视频路径
|
||
scenes: 场景信息列表
|
||
output_dir: 输出目录
|
||
ffmpeg_path: ffmpeg可执行文件路径
|
||
"""
|
||
if not os.path.exists(output_dir):
|
||
os.makedirs(output_dir)
|
||
|
||
# 检查ffmpeg是否可用
|
||
if ffmpeg_path is None:
|
||
ffmpeg_path = find_ffmpeg()
|
||
if ffmpeg_path is None:
|
||
print("错误: 找不到ffmpeg。请安装ffmpeg并确保它在系统路径中。")
|
||
print("您可以从 https://ffmpeg.org/download.html 下载ffmpeg")
|
||
print("或使用包管理器安装: apt-get install ffmpeg / brew install ffmpeg 等")
|
||
return []
|
||
|
||
print(f"\n开始切割视频: {video_path}")
|
||
print(f"输出目录: {output_dir}")
|
||
print(f"使用ffmpeg: {ffmpeg_path}")
|
||
print("-" * 60)
|
||
|
||
clips_info = []
|
||
|
||
for i, scene in enumerate(scenes):
|
||
start_time = scene[2]
|
||
end_time = scene[3]
|
||
duration = scene[4]
|
||
|
||
output_file = os.path.join(output_dir, f"clip_{i:03d}_{duration:.2f}s.mp4")
|
||
|
||
# 构建ffmpeg命令
|
||
cmd = [
|
||
ffmpeg_path,
|
||
'-i', video_path,
|
||
'-ss', f"{start_time:.2f}",
|
||
'-to', f"{end_time:.2f}",
|
||
'-c:v', 'libx264',
|
||
'-c:a', 'aac',
|
||
'-y', # 覆盖已存在的文件
|
||
output_file
|
||
]
|
||
|
||
try:
|
||
print(f"\n切割片段 {i+1}/{len(scenes)}:")
|
||
print(f" 开始时间: {timedelta(seconds=start_time)}")
|
||
print(f" 结束时间: {timedelta(seconds=end_time)}")
|
||
print(f" 时长: {duration:.2f}秒")
|
||
print(f" 输出文件: {os.path.basename(output_file)}")
|
||
|
||
if VERBOSE:
|
||
print(f" 执行命令: {' '.join(cmd)}")
|
||
|
||
# 使用subprocess执行命令，可以获取输出和错误
|
||
result = subprocess.run(
|
||
cmd,
|
||
stdout=subprocess.PIPE if not VERBOSE else None,
|
||
stderr=subprocess.PIPE,
|
||
text=True
|
||
)
|
||
|
||
if result.returncode == 0:
|
||
# 获取输出文件信息
|
||
file_size = os.path.getsize(output_file) / (1024 * 1024) # MB
|
||
print(f" ✓ 切割成功: {os.path.basename(output_file)} ({file_size:.2f} MB)")
|
||
|
||
clips_info.append({
|
||
'index': i,
|
||
'file': output_file,
|
||
'start': start_time,
|
||
'end': end_time,
|
||
'duration': duration,
|
||
'size_mb': file_size
|
||
})
|
||
else:
|
||
print(f" ✗ 切割失败: {result.stderr}")
|
||
except Exception as e:
|
||
print(f" ✗ 切割失败: {str(e)}")
|
||
|
||
# 汇总信息
|
||
if clips_info:
|
||
total_size = sum(clip['size_mb'] for clip in clips_info)
|
||
total_duration = sum(clip['duration'] for clip in clips_info)
|
||
|
||
print("\n切割完成汇总:")
|
||
print(f" 成功切割片段数: {len(clips_info)}/{len(scenes)}")
|
||
print(f" 总时长: {timedelta(seconds=total_duration)}")
|
||
print(f" 总文件大小: {total_size:.2f} MB")
|
||
|
||
# 列出所有文件
|
||
print("\n输出文件列表:")
|
||
for clip in clips_info:
|
||
print(f" {os.path.basename(clip['file'])} - {clip['duration']:.2f}秒 ({clip['size_mb']:.2f} MB)")
|
||
else:
|
||
print("\n没有成功切割任何片段")
|
||
|
||
return clips_info
|
||
|
||
def get_video_files(directory):
|
||
"""
|
||
获取目录中所有视频文件
|
||
|
||
Args:
|
||
directory: 目录路径
|
||
|
||
Returns:
|
||
视频文件路径列表
|
||
"""
|
||
video_files = []
|
||
|
||
# 检查是否是单个文件
|
||
if os.path.isfile(directory):
|
||
ext = os.path.splitext(directory)[1].lower()
|
||
if ext in VIDEO_EXTENSIONS:
|
||
return [directory]
|
||
|
||
# 遍历目录
|
||
for root, _, files in os.walk(directory):
|
||
for file in files:
|
||
# 检查文件扩展名
|
||
ext = os.path.splitext(file)[1].lower()
|
||
if ext in VIDEO_EXTENSIONS:
|
||
video_files.append(os.path.join(root, file))
|
||
|
||
return video_files
|
||
|
||
def process_video(video_path, output_base_dir, sample_rate, method, threshold, ffmpeg_path):
|
||
"""
|
||
处理单个视频文件
|
||
|
||
Args:
|
||
video_path: 视频文件路径
|
||
output_base_dir: 基础输出目录
|
||
sample_rate: 帧采样率
|
||
method: 比较方法
|
||
threshold: 相似度阈值
|
||
ffmpeg_path: ffmpeg路径
|
||
|
||
Returns:
|
||
处理是否成功
|
||
"""
|
||
# 获取视频文件名（不含扩展名）
|
||
video_filename = os.path.splitext(os.path.basename(video_path))[0]
|
||
|
||
# 为当前视频创建输出目录
|
||
video_output_dir = os.path.join(output_base_dir, video_filename)
|
||
if not os.path.exists(video_output_dir):
|
||
os.makedirs(video_output_dir)
|
||
|
||
# 创建输出子目录
|
||
frames_dir = os.path.join(video_output_dir, 'frames')
|
||
clips_dir = os.path.join(video_output_dir, 'clips')
|
||
|
||
if not os.path.exists(frames_dir):
|
||
os.makedirs(frames_dir)
|
||
if not os.path.exists(clips_dir):
|
||
os.makedirs(clips_dir)
|
||
|
||
print("\n处理参数:")
|
||
print(f"输入视频: {os.path.abspath(video_path)}")
|
||
print(f"输出目录: {os.path.abspath(video_output_dir)}")
|
||
print(f"帧采样率: 每{sample_rate}帧")
|
||
print(f"比较方法: {method}")
|
||
print(f"相似度阈值: {threshold}")
|
||
print("-" * 60)
|
||
|
||
try:
|
||
# 步骤1: 提取帧
|
||
print("\n步骤1: 正在提取视频帧...")
|
||
frames_info = extract_frames(video_path, frames_dir, sample_rate)
|
||
|
||
# 步骤2: 检测场景变化
|
||
print("\n步骤2: 正在检测场景变化...")
|
||
scenes = detect_scene_changes(frames_info, method, threshold)
|
||
|
||
# 输出场景信息
|
||
print(f"\n检测到 {len(scenes)} 个场景:")
|
||
for i, scene in enumerate(scenes):
|
||
start_time = timedelta(seconds=scene[2])
|
||
end_time = timedelta(seconds=scene[3])
|
||
duration = scene[4]
|
||
print(f"场景 {i+1}: {start_time} - {end_time} (时长: {duration:.2f}s)")
|
||
|
||
# 如果没有ffmpeg，提前报错
|
||
if not ffmpeg_path:
|
||
print("\n错误: 缺少ffmpeg，无法继续视频切割步骤。")
|
||
print("请安装ffmpeg后重试。")
|
||
return False
|
||
|
||
# 步骤3: 提取视频片段
|
||
print("\n步骤3: 正在提取视频片段...")
|
||
clips_info = extract_video_clips(video_path, scenes, clips_dir, ffmpeg_path)
|
||
|
||
# 创建结果摘要文件
|
||
summary_file = os.path.join(video_output_dir, 'summary.txt')
|
||
try:
|
||
with open(summary_file, 'w', encoding='utf-8') as f:
|
||
f.write("智能视频切割结果摘要\n")
|
||
f.write("=" * 40 + "\n\n")
|
||
f.write(f"输入视频: {os.path.abspath(video_path)}\n")
|
||
f.write(f"处理时间: {os.path.getmtime(summary_file)}\n\n")
|
||
f.write(f"检测到场景数: {len(scenes)}\n")
|
||
f.write(f"生成片段数: {len(clips_info)}\n\n")
|
||
|
||
f.write("片段详情:\n")
|
||
for i, clip in enumerate(clips_info):
|
||
f.write(f"{i+1}. {os.path.basename(clip['file'])}\n")
|
||
f.write(f" 开始: {timedelta(seconds=clip['start'])}\n")
|
||
f.write(f" 结束: {timedelta(seconds=clip['end'])}\n")
|
||
f.write(f" 时长: {clip['duration']:.2f}秒\n")
|
||
f.write(f" 大小: {clip['size_mb']:.2f} MB\n\n")
|
||
|
||
print(f"\n已保存处理摘要到: {summary_file}")
|
||
except Exception as e:
|
||
print(f"保存摘要文件失败: {str(e)}")
|
||
|
||
print("\n处理完成!")
|
||
print(f"帧提取目录: {os.path.abspath(frames_dir)}")
|
||
print(f"视频片段目录: {os.path.abspath(clips_dir)}")
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"\n处理视频 {video_path} 时发生错误: {str(e)}")
|
||
return False
|
||
|
||
def get_parent_folder_name(path):
|
||
"""
|
||
获取路径中 'video' 上一级文件夹的名字
|
||
"""
|
||
abs_path = os.path.abspath(path)
|
||
# 如果是文件夹，直接用
|
||
if os.path.isdir(abs_path):
|
||
parent = os.path.dirname(abs_path.rstrip('/'))
|
||
folder_name = os.path.basename(parent)
|
||
else:
|
||
# 如果是文件，取其父目录的父目录
|
||
parent = os.path.dirname(os.path.dirname(abs_path))
|
||
folder_name = os.path.basename(parent)
|
||
return folder_name
|
||
|
||
def main():
|
||
# 欢迎信息
|
||
print("=" * 60)
|
||
print("智能视频切割工具 - 批量处理版")
|
||
print("=" * 60)
|
||
|
||
# 查找ffmpeg
|
||
ffmpeg_path = find_ffmpeg()
|
||
if ffmpeg_path:
|
||
print(f"已找到ffmpeg: {ffmpeg_path}")
|
||
else:
|
||
print("警告: 未找到ffmpeg，视频切割功能将不可用")
|
||
print("请安装ffmpeg并确保它在系统路径中")
|
||
|
||
# 获取输入目录中的所有视频文件
|
||
video_files = get_video_files(INPUT_VIDEO_PATH)
|
||
|
||
if not video_files:
|
||
print(f"错误: 在 '{INPUT_VIDEO_PATH}' 中没有找到视频文件")
|
||
print(f"支持的视频格式: {', '.join(VIDEO_EXTENSIONS)}")
|
||
return
|
||
|
||
# 自动获取video上一级文件夹名
|
||
parent_folder_name = get_parent_folder_name(INPUT_VIDEO_PATH)
|
||
output_dir = os.path.join(OUTPUT_DIR, parent_folder_name)
|
||
if not os.path.exists(output_dir):
|
||
os.makedirs(output_dir)
|
||
|
||
print(f"\n输出目录: {output_dir}")
|
||
|
||
# 处理每个视频文件
|
||
successful = 0
|
||
failed = 0
|
||
|
||
for i, video_path in enumerate(video_files):
|
||
print("\n" + "=" * 60)
|
||
print(f"正在处理视频 [{i+1}/{len(video_files)}]: {os.path.basename(video_path)}")
|
||
print("=" * 60)
|
||
|
||
success = process_video(
|
||
video_path=video_path,
|
||
output_base_dir=output_dir,
|
||
sample_rate=SAMPLE_RATE,
|
||
method=METHOD,
|
||
threshold=THRESHOLD,
|
||
ffmpeg_path=ffmpeg_path
|
||
)
|
||
|
||
if success:
|
||
successful += 1
|
||
else:
|
||
failed += 1
|
||
|
||
# 打印批量处理总结
|
||
print("\n" + "=" * 60)
|
||
print("批量处理完成!")
|
||
print("=" * 60)
|
||
print(f"总共处理: {len(video_files)} 个视频文件")
|
||
print(f"成功: {successful} 个")
|
||
print(f"失败: {failed} 个")
|
||
print(f"输出目录: {os.path.abspath(output_dir)}")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|