#!/usr/bin/env python3 """ 视频预处理脚本 - 将MP4视频转换为224x224帧图像 支持多线程并发处理、进度条显示和中断恢复功能 """ import os import sys import json import argparse import subprocess import threading from pathlib import Path from concurrent.futures import ThreadPoolExecutor, as_completed from tqdm import tqdm import time from typing import List, Dict, Optional class VideoPreprocessor: """视频预处理器，支持多线程和中断恢复""" def __init__(self, input_dir: str, output_dir: str, frame_size: int = 224, fps: int = 30, num_workers: int = 4, quality: int = 2, resume: bool = True): """ 初始化预处理器 Args: input_dir: 输入视频目录 output_dir: 输出帧目录 frame_size: 帧大小（正方形） fps: 提取帧率 num_workers: 并发工作线程数 quality: JPEG质量 (1-31, 数值越小质量越高) resume: 是否启用中断恢复 """ self.input_dir = Path(input_dir) self.output_dir = Path(output_dir) self.frame_size = frame_size self.fps = fps self.num_workers = num_workers self.quality = quality self.resume = resume # 状态文件路径 self.state_file = self.output_dir / ".preprocessing_state.json" # 创建输出目录 self.output_dir.mkdir(parents=True, exist_ok=True) # 初始化状态 self.state = self._load_state() # 收集所有视频文件 self.video_files = self._collect_video_files() def _load_state(self) -> Dict: """加载处理状态""" if self.resume and self.state_file.exists(): try: with open(self.state_file, 'r') as f: return json.load(f) except (json.JSONDecodeError, IOError): print(f"警告: 无法读取状态文件，将重新开始处理") return { "completed": [], "failed": [], "total_processed": 0, "start_time": None, "last_update": None } def _save_state(self): """保存处理状态""" self.state["last_update"] = time.time() try: with open(self.state_file, 'w') as f: json.dump(self.state, f, indent=2) except IOError as e: print(f"警告: 无法保存状态文件: {e}") def _collect_video_files(self) -> List[Path]: """收集所有需要处理的视频文件""" video_files = [] for file_path in self.input_dir.glob("*.mp4"): if file_path.name not in self.state["completed"]: video_files.append(file_path) return sorted(video_files) def _parse_video_name(self, video_path: Path) -> Dict[str, str]: """解析视频文件名，使用完整文件名作为ID""" name_without_ext = video_path.stem # 直接使用完整文件名作为ID，确保每个mp4文件有独立的输出目录 return { "video_id": name_without_ext, "start_frame": "unknown", "end_frame": "unknown", "full_name": name_without_ext } def _extract_frames(self, video_path: Path) -> bool: """提取单个视频的帧""" try: # 解析视频名称 video_info = self._parse_video_name(video_path) output_subdir = self.output_dir / video_info["video_id"] output_subdir.mkdir(exist_ok=True) # 构建FFmpeg命令 output_pattern = output_subdir / "frame_%04d.jpg" cmd = [ "ffmpeg", "-i", str(video_path), "-vf", f"fps={self.fps},scale={self.frame_size}:{self.frame_size}", "-q:v", str(self.quality), "-y", # 覆盖输出文件 str(output_pattern) ] # 执行FFmpeg命令 result = subprocess.run( cmd, capture_output=True, text=True, timeout=300 # 5分钟超时 ) if result.returncode != 0: print(f"FFmpeg错误处理 {video_path.name}: {result.stderr}") return False # 验证输出帧数量 output_frames = list(output_subdir.glob("frame_*.jpg")) if len(output_frames) == 0: print(f"警告: {video_path.name} 没有生成任何帧") return False return True except subprocess.TimeoutExpired: print(f"超时处理 {video_path.name}") return False except Exception as e: print(f"处理 {video_path.name} 时发生错误: {e}") return False def _process_video(self, video_path: Path) -> tuple[bool, str]: """处理单个视频文件""" video_name = video_path.name try: success = self._extract_frames(video_path) if success: self.state["completed"].append(video_name) if video_name in self.state["failed"]: self.state["failed"].remove(video_name) return True, video_name else: self.state["failed"].append(video_name) return False, video_name except Exception as e: print(f"处理 {video_name} 时发生异常: {e}") self.state["failed"].append(video_name) return False, video_name def process_all_videos(self): """处理所有视频文件""" if not self.video_files: print("没有找到需要处理的视频文件") return print(f"找到 {len(self.video_files)} 个待处理视频文件") print(f"输出目录: {self.output_dir}") print(f"帧大小: {self.frame_size}x{self.frame_size}") print(f"帧率: {self.fps} fps") print(f"并发线程数: {self.num_workers}") if self.state["completed"]: print(f"跳过 {len(self.state['completed'])} 个已处理的视频") # 记录开始时间 if self.state["start_time"] is None: self.state["start_time"] = time.time() # 创建进度条 with tqdm(total=len(self.video_files), desc="处理视频", unit="个") as pbar: with ThreadPoolExecutor(max_workers=self.num_workers) as executor: # 提交所有任务 future_to_video = { executor.submit(self._process_video, video_path): video_path for video_path in self.video_files } # 处理完成的任务 for future in as_completed(future_to_video): video_path = future_to_video[future] try: success, video_name = future.result() if success: pbar.set_postfix({"状态": "成功", "文件": video_name[:20]}) else: pbar.set_postfix({"状态": "失败", "文件": video_name[:20]}) except Exception as e: print(f"处理 {video_path.name} 时发生异常: {e}") pbar.set_postfix({"状态": "异常", "文件": video_path.name[:20]}) pbar.update(1) self.state["total_processed"] += 1 # 定期保存状态 if self.state["total_processed"] % 5 == 0: self._save_state() # 最终保存状态 self._save_state() # 打印处理结果 self._print_summary() def _print_summary(self): """打印处理摘要""" print("\n" + "="*50) print("处理完成摘要:") print(f"总处理视频数: {len(self.state['completed'])}") print(f"失败视频数: {len(self.state['failed'])}") if self.state["failed"]: print("\n失败的视频:") for video_name in self.state["failed"]: print(f" - {video_name}") if self.state["start_time"]: elapsed_time = time.time() - self.state["start_time"] print(f"\n总耗时: {elapsed_time:.2f} 秒") if self.state["total_processed"] > 0: avg_time = elapsed_time / self.state["total_processed"] print(f"平均每个视频: {avg_time:.2f} 秒") print("="*50) def main(): """主函数""" parser = argparse.ArgumentParser(description="视频预处理脚本") parser.add_argument("--input_dir", type=str, default="/home/hexone/Workplace/ws_asmo/vhead/sekai-real-drone/sekai-real-drone", help="输入视频目录") parser.add_argument("--output_dir", type=str, default="/home/hexone/Workplace/ws_asmo/vhead/sekai-real-drone/processed", help="输出帧目录") parser.add_argument("--size", type=int, default=224, help="帧大小 (默认: 224)") parser.add_argument("--fps", type=int, default=10, help="提取帧率 (默认: 30)") parser.add_argument("--workers", type=int, default=32, help="并发线程数 (默认: 4)") parser.add_argument("--quality", type=int, default=2, help="JPEG质量 1-31 (默认: 2)") parser.add_argument("--no-resume", action="store_true", help="不启用中断恢复") args = parser.parse_args() # 检查输入目录 if not Path(args.input_dir).exists(): print(f"错误: 输入目录不存在: {args.input_dir}") sys.exit(1) # 检查FFmpeg是否可用 try: subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True) except (subprocess.CalledProcessError, FileNotFoundError): print("错误: FFmpeg未安装或不在PATH中") sys.exit(1) # 创建预处理器并开始处理 preprocessor = VideoPreprocessor( input_dir=args.input_dir, output_dir=args.output_dir, frame_size=args.size, fps=args.fps, num_workers=args.workers, quality=args.quality, resume=not args.no_resume ) try: preprocessor.process_all_videos() except KeyboardInterrupt: print("\n\n用户中断处理，状态已保存") preprocessor._save_state() print("可以使用相同命令恢复处理") except Exception as e: print(f"\n处理过程中发生错误: {e}") preprocessor._save_state() sys.exit(1) if __name__ == "__main__": main()