Files
agent_test/video_to_frames.py

167 lines
4.9 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""
视频切割脚本 - 将视频逐帧切割为图片
支持扫描目录下所有视频文件,多进程并行处理,已切割的自动跳过
"""
import os
import sys
import cv2
import argparse
from multiprocessing import Pool, cpu_count
from pathlib import Path
# 支持的视频扩展名
VIDEO_EXTENSIONS = {'.mp4', '.avi', '.mkv', '.mov', '.flv', '.wmv', '.webm', '.ts', '.m4v'}
def scan_videos(root_path: str) -> list[str]:
"""扫描目录下所有视频文件"""
videos = []
root = Path(root_path)
if root.is_file() and root.suffix.lower() in VIDEO_EXTENSIONS:
return [str(root)]
for p in sorted(root.rglob('*')):
if p.is_file() and p.suffix.lower() in VIDEO_EXTENSIONS:
videos.append(str(p))
return videos
def extract_frames(video_path: str, output_dir: str = None) -> dict:
"""
将单个视频按每秒1帧切割为图片
Args:
video_path: 视频文件路径
output_dir: 输出目录(默认在视频同目录下创建同名文件夹)
Returns:
处理结果字典
"""
video_path = os.path.abspath(video_path)
video_name = Path(video_path).stem
if output_dir is None:
output_dir = os.path.join(os.path.dirname(video_path), video_name + '_frames')
# 已经切割过则跳过
if os.path.isdir(output_dir) and any(Path(output_dir).glob('*.jpg')):
existing = len(list(Path(output_dir).glob('*.jpg')))
return {
'video': video_path,
'status': 'skipped',
'reason': f'已存在,共 {existing} 张图片',
'output': output_dir,
}
os.makedirs(output_dir, exist_ok=True)
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
return {
'video': video_path,
'status': 'error',
'reason': '无法打开视频文件',
'output': output_dir,
}
fps = cap.get(cv2.CAP_PROP_FPS) or 25
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
duration = total_frames / fps if fps else 0
# 每秒取一帧根据fps算出每隔多少帧取一帧
frame_interval = int(fps)
frame_idx = 0
saved = 0
while True:
ret, frame = cap.read()
if not ret:
break
if frame_idx % frame_interval == 0:
# 用视频名+时间戳命名第01集_01m23s_456.jpg
sec = frame_idx / fps
mins = int(sec // 60)
secs = int(sec % 60)
ms = int((sec % 1) * 1000)
filename = os.path.join(output_dir, f'{video_name}_{mins:02d}m{secs:02d}s_{ms:03d}.jpg')
cv2.imwrite(filename, frame)
saved += 1
frame_idx += 1
cap.release()
return {
'video': video_path,
'status': 'done',
'frames_total': total_frames,
'frames_saved': saved,
'duration': f'{duration:.1f}s',
'fps': fps,
'output': output_dir,
}
def _worker(args):
"""多进程 worker 包装"""
video_path, output_dir = args
try:
return extract_frames(video_path, output_dir)
except Exception as e:
return {
'video': video_path,
'status': 'error',
'reason': str(e),
}
def main(video_path, output=None, interval=1, jobs=None):
videos = scan_videos(video_path)
if not videos:
print(f'未在 {video_path} 中找到视频文件')
return
jobs = jobs or cpu_count()
print(f'找到 {len(videos)} 个视频,使用 {jobs} 个进程并行处理')
tasks = []
for v in videos:
if output:
video_name = Path(v).stem
out = os.path.join(output, video_name + '_frames')
else:
out = None
tasks.append((v, out))
done = 0
skipped = 0
errors = 0
with Pool(jobs) as pool:
for result in pool.imap_unordered(_worker, tasks):
status = result['status']
video = result['video']
if status == 'done':
done += 1
print(f'[完成] {video} -> {result["frames_saved"]}'
f'(总{result["frames_total"]}, {result["duration"]})')
elif status == 'skipped':
skipped += 1
print(f'[跳过] {video} ({result["reason"]})')
else:
errors += 1
print(f'[错误] {video} ({result["reason"]})')
print(f'\n处理完毕: 完成 {done}, 跳过 {skipped}, 错误 {errors}, 共 {len(videos)}')
# ========== 运行设置(改这里)==========
VIDEO_DIR = "./shipin" # 视频文件或目录路径
OUTPUT_DIR = None # 输出目录None 则在视频同目录
INTERVAL = 1 # 每隔 N 帧取一帧
JOBS = None # 并行进程数None 则用 CPU 核心数
# ======================================
if __name__ == '__main__':
main(VIDEO_DIR, OUTPUT_DIR, INTERVAL, JOBS)