337 lines
14 KiB
Python
337 lines
14 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
批量视频字幕提取器
|
|||
|
|
支持批量处理多个视频文件，提取字幕
|
|||
|
|
支持PaddleOCR、EasyOCR和CnOCR三种引擎
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import sys
|
|||
|
|
import time
|
|||
|
|
import json
|
|||
|
|
import argparse
|
|||
|
|
from pathlib import Path
|
|||
|
|
from datetime import datetime
|
|||
|
|
import logging
|
|||
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|||
|
|
|
|||
|
|
# 添加当前目录到路径以导入OCR模块
|
|||
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
|||
|
|
|
|||
|
|
from ocr_subtitle_extractor import VideoSubtitleExtractor
|
|||
|
|
|
|||
|
|
# 设置OCR模型路径环境变量
|
|||
|
|
os.environ['EASYOCR_MODULE_PATH'] = '/root/autodl-tmp/llm/easyocr'
|
|||
|
|
os.environ['CNOCR_HOME'] = '/root/autodl-tmp/llm/cnocr'
|
|||
|
|
|
|||
|
|
# 设置日志
|
|||
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
class BatchSubtitleExtractor:
|
|||
|
|
"""批量视频字幕提取器"""
|
|||
|
|
|
|||
|
|
def __init__(self, ocr_engine="paddleocr", language="ch", max_workers=2):
|
|||
|
|
"""
|
|||
|
|
初始化批量提取器
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
ocr_engine: OCR引擎 ("paddleocr", "easyocr", "cnocr", "all")
|
|||
|
|
language: 语言设置 ("ch", "en", "ch_en")
|
|||
|
|
max_workers: 最大并行工作数
|
|||
|
|
"""
|
|||
|
|
self.ocr_engine = ocr_engine
|
|||
|
|
self.language = language
|
|||
|
|
self.max_workers = max_workers
|
|||
|
|
self.extractor = VideoSubtitleExtractor(ocr_engine=ocr_engine, language=language)
|
|||
|
|
|
|||
|
|
def find_video_files(self, input_dir):
|
|||
|
|
"""查找目录中的所有视频文件"""
|
|||
|
|
video_extensions = ['.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv', '.webm']
|
|||
|
|
video_files = []
|
|||
|
|
|
|||
|
|
input_path = Path(input_dir)
|
|||
|
|
|
|||
|
|
if input_path.is_file():
|
|||
|
|
# 单个文件
|
|||
|
|
if input_path.suffix.lower() in video_extensions:
|
|||
|
|
video_files.append(input_path)
|
|||
|
|
elif input_path.is_dir():
|
|||
|
|
# 目录中的所有视频文件
|
|||
|
|
for ext in video_extensions:
|
|||
|
|
video_files.extend(input_path.glob(f"*{ext}"))
|
|||
|
|
video_files.extend(input_path.glob(f"*{ext.upper()}"))
|
|||
|
|
|
|||
|
|
return sorted(video_files)
|
|||
|
|
|
|||
|
|
def extract_single_video(self, video_path, output_dir, **kwargs):
|
|||
|
|
"""
|
|||
|
|
处理单个视频文件
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
video_path: 视频文件路径
|
|||
|
|
output_dir: 输出目录
|
|||
|
|
**kwargs: 其他参数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
dict: 处理结果
|
|||
|
|
"""
|
|||
|
|
video_path = Path(video_path)
|
|||
|
|
video_name = video_path.stem
|
|||
|
|
|
|||
|
|
logger.info(f"开始处理视频: {video_path}")
|
|||
|
|
start_time = time.time()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 提取字幕
|
|||
|
|
results = self.extractor.extract_subtitles_from_video(
|
|||
|
|
str(video_path),
|
|||
|
|
sample_interval=kwargs.get('interval', 30),
|
|||
|
|
confidence_threshold=kwargs.get('confidence', 0.5),
|
|||
|
|
subtitle_position=kwargs.get('position', 'bottom')
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 保存结果
|
|||
|
|
output_path = Path(output_dir)
|
|||
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
for format_type in kwargs.get('formats', ['json']):
|
|||
|
|
output_file = output_path / f"{video_name}_subtitles.{format_type}"
|
|||
|
|
self.extractor.save_results(results, output_file, format_type)
|
|||
|
|
|
|||
|
|
process_time = time.time() - start_time
|
|||
|
|
results['process_time'] = process_time
|
|||
|
|
results['video_path'] = str(video_path)
|
|||
|
|
results['success'] = True
|
|||
|
|
|
|||
|
|
# 统计位置信息
|
|||
|
|
subtitles_with_bbox = [s for s in results['subtitles'] if s.get('bbox')]
|
|||
|
|
bbox_coverage = len(subtitles_with_bbox) / len(results['subtitles']) * 100 if results['subtitles'] else 0
|
|||
|
|
|
|||
|
|
logger.info(f"完成处理视频: {video_path} (耗时: {process_time:.2f}秒)")
|
|||
|
|
logger.info(f" 字幕总数: {len(results['subtitles'])}")
|
|||
|
|
logger.info(f" 有位置信息: {len(subtitles_with_bbox)}")
|
|||
|
|
logger.info(f" 位置信息覆盖率: {bbox_coverage:.1f}%")
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'video_path': str(video_path),
|
|||
|
|
'success': True,
|
|||
|
|
'process_time': process_time,
|
|||
|
|
'subtitle_count': results['stats']['filtered_detections'],
|
|||
|
|
'text_length': results['stats']['text_length'],
|
|||
|
|
'total_subtitles': len(results['subtitles']),
|
|||
|
|
'subtitles_with_bbox': len(subtitles_with_bbox),
|
|||
|
|
'bbox_coverage': bbox_coverage,
|
|||
|
|
'output_files': [str(output_path / f"{video_name}_subtitles.{fmt}") for fmt in kwargs.get('formats', ['json'])]
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
error_msg = f"处理视频 {video_path} 时出错: {str(e)}"
|
|||
|
|
logger.error(error_msg)
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'video_path': str(video_path),
|
|||
|
|
'success': False,
|
|||
|
|
'error': error_msg,
|
|||
|
|
'process_time': time.time() - start_time
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def extract_batch(self, input_dir, output_dir, parallel=True, **kwargs):
|
|||
|
|
"""
|
|||
|
|
批量提取字幕
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
input_dir: 输入目录或文件
|
|||
|
|
output_dir: 输出目录
|
|||
|
|
parallel: 是否并行处理
|
|||
|
|
**kwargs: 其他参数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
dict: 批量处理结果
|
|||
|
|
"""
|
|||
|
|
logger.info(f"开始批量字幕提取")
|
|||
|
|
logger.info(f"输入: {input_dir}")
|
|||
|
|
logger.info(f"输出目录: {output_dir}")
|
|||
|
|
logger.info(f"OCR引擎: {self.ocr_engine}")
|
|||
|
|
logger.info(f"字幕位置: {kwargs.get('position', 'bottom')}")
|
|||
|
|
logger.info(f"并行处理: {parallel}")
|
|||
|
|
|
|||
|
|
start_time = time.time()
|
|||
|
|
|
|||
|
|
# 查找视频文件
|
|||
|
|
video_files = self.find_video_files(input_dir)
|
|||
|
|
|
|||
|
|
if not video_files:
|
|||
|
|
logger.warning(f"在 {input_dir} 中未找到视频文件")
|
|||
|
|
return {
|
|||
|
|
'success': False,
|
|||
|
|
'message': '未找到视频文件',
|
|||
|
|
'total_files': 0,
|
|||
|
|
'results': []
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
logger.info(f"找到 {len(video_files)} 个视频文件")
|
|||
|
|
|
|||
|
|
results = []
|
|||
|
|
|
|||
|
|
if parallel and len(video_files) > 1:
|
|||
|
|
# 并行处理
|
|||
|
|
logger.info(f"使用 {self.max_workers} 个并行工作进程")
|
|||
|
|
|
|||
|
|
with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
|
|||
|
|
# 提交任务
|
|||
|
|
future_to_video = {
|
|||
|
|
executor.submit(self.extract_single_video, video_file, output_dir, **kwargs): video_file
|
|||
|
|
for video_file in video_files
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 收集结果
|
|||
|
|
for future in as_completed(future_to_video):
|
|||
|
|
video_file = future_to_video[future]
|
|||
|
|
try:
|
|||
|
|
result = future.result()
|
|||
|
|
results.append(result)
|
|||
|
|
|
|||
|
|
# 显示进度
|
|||
|
|
progress = len(results) / len(video_files) * 100
|
|||
|
|
logger.info(f"批量处理进度: {progress:.1f}% ({len(results)}/{len(video_files)})")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"处理视频 {video_file} 时发生异常: {str(e)}")
|
|||
|
|
results.append({
|
|||
|
|
'video_path': str(video_file),
|
|||
|
|
'success': False,
|
|||
|
|
'error': str(e)
|
|||
|
|
})
|
|||
|
|
else:
|
|||
|
|
# 串行处理
|
|||
|
|
for i, video_file in enumerate(video_files, 1):
|
|||
|
|
logger.info(f"处理第 {i}/{len(video_files)} 个视频")
|
|||
|
|
result = self.extract_single_video(video_file, output_dir, **kwargs)
|
|||
|
|
results.append(result)
|
|||
|
|
|
|||
|
|
# 显示进度
|
|||
|
|
progress = i / len(video_files) * 100
|
|||
|
|
logger.info(f"批量处理进度: {progress:.1f}%")
|
|||
|
|
|
|||
|
|
total_time = time.time() - start_time
|
|||
|
|
|
|||
|
|
# 统计结果
|
|||
|
|
success_count = sum(1 for r in results if r['success'])
|
|||
|
|
failed_count = len(results) - success_count
|
|||
|
|
|
|||
|
|
total_subtitles = sum(r.get('subtitle_count', 0) for r in results if r['success'])
|
|||
|
|
total_text_length = sum(r.get('text_length', 0) for r in results if r['success'])
|
|||
|
|
|
|||
|
|
# 统计位置信息
|
|||
|
|
total_subtitles_raw = sum(r.get('total_subtitles', 0) for r in results if r['success'])
|
|||
|
|
total_subtitles_with_bbox = sum(r.get('subtitles_with_bbox', 0) for r in results if r['success'])
|
|||
|
|
overall_bbox_coverage = total_subtitles_with_bbox / total_subtitles_raw * 100 if total_subtitles_raw > 0 else 0
|
|||
|
|
|
|||
|
|
batch_result = {
|
|||
|
|
'success': True,
|
|||
|
|
'total_time': total_time,
|
|||
|
|
'total_files': len(video_files),
|
|||
|
|
'success_count': success_count,
|
|||
|
|
'failed_count': failed_count,
|
|||
|
|
'total_subtitles': total_subtitles,
|
|||
|
|
'total_text_length': total_text_length,
|
|||
|
|
'total_subtitles_raw': total_subtitles_raw,
|
|||
|
|
'total_subtitles_with_bbox': total_subtitles_with_bbox,
|
|||
|
|
'overall_bbox_coverage': overall_bbox_coverage,
|
|||
|
|
'output_directory': output_dir,
|
|||
|
|
'ocr_engine': self.ocr_engine,
|
|||
|
|
'timestamp': datetime.now().isoformat(),
|
|||
|
|
'results': results
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 保存批量处理报告
|
|||
|
|
report_file = Path(output_dir) / "batch_report.json"
|
|||
|
|
with open(report_file, 'w', encoding='utf-8') as f:
|
|||
|
|
json.dump(batch_result, f, ensure_ascii=False, indent=2)
|
|||
|
|
|
|||
|
|
logger.info(f"批量处理完成！")
|
|||
|
|
logger.info(f"总文件数: {len(video_files)}")
|
|||
|
|
logger.info(f"成功: {success_count}, 失败: {failed_count}")
|
|||
|
|
logger.info(f"总耗时: {total_time:.2f} 秒")
|
|||
|
|
logger.info(f"提取字幕: {total_subtitles} 个")
|
|||
|
|
logger.info(f"文本长度: {total_text_length} 字符")
|
|||
|
|
logger.info(f"位置信息统计:")
|
|||
|
|
logger.info(f" 总字幕数: {total_subtitles_raw}")
|
|||
|
|
logger.info(f" 有位置信息: {total_subtitles_with_bbox}")
|
|||
|
|
logger.info(f" 位置信息覆盖率: {overall_bbox_coverage:.1f}%")
|
|||
|
|
logger.info(f"处理报告: {report_file}")
|
|||
|
|
|
|||
|
|
return batch_result
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
"""主函数"""
|
|||
|
|
parser = argparse.ArgumentParser(description="批量视频字幕提取器")
|
|||
|
|
parser.add_argument("input", help="输入视频文件或目录")
|
|||
|
|
parser.add_argument("-e", "--engine", default="cnocr",
|
|||
|
|
choices=["paddleocr", "easyocr", "cnocr", "all"],
|
|||
|
|
help="OCR引擎 (默认: paddleocr)")
|
|||
|
|
parser.add_argument("-l", "--language", default="ch",
|
|||
|
|
choices=["ch", "en", "ch_en"],
|
|||
|
|
help="语言设置 (默认: ch)")
|
|||
|
|
parser.add_argument("-i", "--interval", type=int, default=30,
|
|||
|
|
help="帧采样间隔 (默认: 30)")
|
|||
|
|
parser.add_argument("-c", "--confidence", type=float, default=0.5,
|
|||
|
|
help="置信度阈值 (默认: 0.5)")
|
|||
|
|
parser.add_argument("-o", "--output", default="batch_subtitles",
|
|||
|
|
help="输出目录 (默认: batch_subtitles)")
|
|||
|
|
parser.add_argument("-f", "--formats", nargs='+', default=["json"],
|
|||
|
|
choices=["json", "txt", "srt"],
|
|||
|
|
help="输出格式 (默认: json)")
|
|||
|
|
parser.add_argument("--position", default="full",
|
|||
|
|
choices=["full", "center", "bottom"],
|
|||
|
|
help="字幕区域位置 (full=全屏, center=居中0.5-0.8, bottom=居下0.7-1.0)")
|
|||
|
|
parser.add_argument("--workers", type=int, default=2,
|
|||
|
|
help="并行工作进程数 (默认: 2)")
|
|||
|
|
parser.add_argument("--no-parallel", action="store_true",
|
|||
|
|
help="禁用并行处理")
|
|||
|
|
|
|||
|
|
args = parser.parse_args()
|
|||
|
|
|
|||
|
|
# 创建批量提取器
|
|||
|
|
batch_extractor = BatchSubtitleExtractor(
|
|||
|
|
ocr_engine=args.engine,
|
|||
|
|
language=args.language,
|
|||
|
|
max_workers=args.workers
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 执行批量提取
|
|||
|
|
result = batch_extractor.extract_batch(
|
|||
|
|
input_dir=args.input,
|
|||
|
|
output_dir=args.output,
|
|||
|
|
parallel=not args.no_parallel,
|
|||
|
|
interval=args.interval,
|
|||
|
|
confidence=args.confidence,
|
|||
|
|
formats=args.formats,
|
|||
|
|
position=args.position
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if result['success']:
|
|||
|
|
print(f"\n✅ 批量字幕提取完成!")
|
|||
|
|
print(f"📁 输出目录: {args.output}")
|
|||
|
|
print(f"📊 成功处理: {result['success_count']}/{result['total_files']} 个视频")
|
|||
|
|
if result['failed_count'] > 0:
|
|||
|
|
print(f"❌ 失败: {result['failed_count']} 个")
|
|||
|
|
print(f"⏱️ 总耗时: {result['total_time']:.2f} 秒")
|
|||
|
|
print(f"📝 字幕片段: {result['total_subtitles']} 个")
|
|||
|
|
print(f"📏 文本长度: {result['total_text_length']} 字符")
|
|||
|
|
print(f"📍 位置信息统计:")
|
|||
|
|
print(f" 总字幕数: {result['total_subtitles_raw']}")
|
|||
|
|
print(f" 有位置信息: {result['total_subtitles_with_bbox']}")
|
|||
|
|
print(f" 位置信息覆盖率: {result['overall_bbox_coverage']:.1f}%")
|
|||
|
|
else:
|
|||
|
|
print(f"\n❌ 批量处理失败: {result.get('message', '未知错误')}")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"批量处理出错: {str(e)}")
|
|||
|
|
print(f"\n❌ 批量处理出错: {str(e)}")
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|