#!/usr/bin/env python3 """ OCR Processor - 视频帧文字识别工具 支持 EasyOCR 和 Tesseract 引擎 """ import argparse import json import os import glob from pathlib import Path from typing import List, Dict, Optional import easyocr import pytesseract from PIL import Image class OCRProcessor: """OCR 文字识别处理器""" def __init__(self, engine: str = "easyocr", languages: List[str] = ["ch_sim", "en"]): """ 初始化 OCR 处理器 Args: engine: OCR 引擎 ("easyocr" 或 "tesseract") languages: 语言列表 """ self.engine = engine self.languages = languages self.reader = None if engine == "easyocr": self.reader = easyocr.Reader(languages, gpu=False) def process_image(self, image_path: str) -> Dict: """ 处理单张图片进行 OCR 识别 Args: image_path: 图片路径 Returns: 包含 text 和 confidence 的字典 """ if not os.path.exists(image_path): return {"text": "", "confidence": 0.0, "error": "File not found"} try: if self.engine == "easyocr": return self._process_with_easyocr(image_path) elif self.engine == "tesseract": return self._process_with_tesseract(image_path) else: return {"text": "", "confidence": 0.0, "error": f"Unknown engine: {self.engine}"} except Exception as e: return {"text": "", "confidence": 0.0, "error": str(e)} def _process_with_easyocr(self, image_path: str) -> Dict: """使用 EasyOCR 处理图片""" results = self.reader.readtext(image_path) if not results: return {"text": "", "confidence": 0.0} texts = [] confidences = [] for bbox, text, conf in results: texts.append(text) confidences.append(conf) merged_text = "\n".join(texts) avg_confidence = sum(confidences) / len(confidences) if confidences else 0.0 return { "text": merged_text, "confidence": round(avg_confidence, 4), "details": [ {"text": text, "confidence": round(conf, 4)} for text, conf in zip(texts, confidences) ] } def _process_with_tesseract(self, image_path: str) -> Dict: """使用 Tesseract 处理图片""" image = Image.open(image_path) text = pytesseract.image_to_string(image, lang="+".join(self.languages)) # Tesseract 不直接返回置信度，返回 0.0 return { "text": text.strip(), "confidence": 0.0 } def process_frames(self, input_dir: str, fps: float = 0.1, progress_file: Optional[str] = None) -> List[Dict]: """ 处理目录下所有帧图片 Args: input_dir: 帧图片目录 fps: 采样帧率 progress_file: 可选，进度文件路径（会实时更新） Returns: OCR 结果列表 """ if not os.path.isdir(input_dir): print(f"Error: Directory not found: {input_dir}") return [] # 支持的图片格式 image_extensions = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp"] image_files = [] for ext in image_extensions: image_files.extend(glob.glob(os.path.join(input_dir, ext))) image_files.extend(glob.glob(os.path.join(input_dir, "**", ext), recursive=True)) # 按文件名排序 image_files = sorted(image_files) if not image_files: print(f"No images found in {input_dir}") return [] total = len(image_files) print(f"Found {total} images in {input_dir}") print(f"Processing with {self.engine} engine...") def update_progress(current, status="running", last_result=None): """更新进度文件""" if progress_file: import json import time prog = { "current": current, "total": total, "percent": round(current / total * 100, 1) if total > 0 else 0, "status": status, "last_result": last_result, "updated_at": time.strftime("%Y-%m-%d %H:%M:%S") } with open(progress_file, 'w', encoding='utf-8') as f: json.dump(prog, f, ensure_ascii=False, indent=2) # 初始化进度 update_progress(0, "started") results = [] for i, image_path in enumerate(image_files, 1): basename = os.path.basename(image_path) print(f"Processing [{i}/{total}]: {basename}") ocr_result = self.process_image(image_path) ocr_result["source"] = basename # 记录结果摘要 text_count = len(ocr_result.get("text", "").split()) if ocr_result.get("text") else 0 last_result = f"✓ {text_count} 处文字" if ocr_result.get("text") else "✗ 无文字" results.append(ocr_result) # 实时更新进度 update_progress(i, "running", last_result) # 完成 update_progress(total, "completed", f"完成 {total} 张") print(f"\n✅ Processed {total} frames") return results def merge_texts(self, results: List[Dict]) -> str: """ 合并多个 OCR 结果的文本 Args: results: OCR 结果列表 Returns: 合并后的文本 """ texts = [] for result in results: if result.get("text"): texts.append(result["text"]) return "\n\n---\n\n".join(texts) def save_results(self, results: List[Dict], output_path: str, merge: bool = False): """ 保存 OCR 结果到 JSON 文件 Args: results: OCR 结果列表 output_path: 输出文件路径 merge: 是否同时保存合并后的文本文件 """ output_data = { "engine": self.engine, "total_frames": len(results), "results": results } # 保存 JSON with open(output_path, "w", encoding="utf-8") as f: json.dump(output_data, f, ensure_ascii=False, indent=2) print(f"Results saved to {output_path}") # 如果需要，同时保存合并的文本 if merge: text_path = output_path.replace(".json", ".txt") merged_text = self.merge_texts(results) with open(text_path, "w", encoding="utf-8") as f: f.write(merged_text) print(f"Merged text saved to {text_path}") def extract_frames(video_path: str, output_dir: str, fps: float = 1.0) -> str: """ 使用 ffmpeg 从视频中截取帧图片 Args: video_path: 视频文件路径 output_dir: 输出目录 fps: 每秒截取帧数 (默认: 1) Returns: 帧图片目录路径 """ import subprocess output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) # ffmpeg 命令: -i 输入 -vf fps=fps 输出格式 cmd = [ "ffmpeg", "-y", "-i", video_path, "-vf", f"fps={fps}", "-q:v", "2", # 质量等级 (2=高质量) str(output_dir / "%04d.jpg") ] print(f"🔄 截取视频帧: {video_path}") print(f" 输出目录: {output_dir}") try: result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode == 0: # 统计截取的帧数 frame_count = len(list(output_dir.glob("*.jpg"))) print(f"✅ 截取完成: {frame_count} 帧") return str(output_dir) else: print(f"❌ ffmpeg 错误: {result.stderr}") return "" except FileNotFoundError: print("❌ 错误: 未找到 ffmpeg，请先安装: apt install ffmpeg") return "" def main(): """命令行入口""" parser = argparse.ArgumentParser( description="OCR 视频帧文字识别工具", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" 示例: # 使用 EasyOCR 处理帧目录 python ocr_processor.py --input ./frames/ --output ./ocr.json # 直接从视频截帧并 OCR python ocr_processor.py --video ./video.flv --output ./ocr.json # 使用 Tesseract 处理 python ocr_processor.py --input ./frames/ --output ./ocr.json --engine tesseract # 指定截帧帧率 python ocr_processor.py --video ./video.flv --output ./ocr.json --fps 2 # 同时生成合并的文本文件 python ocr_processor.py --input ./frames/ --output ./ocr.json --merge """ ) # 输入源：帧目录 或 视频文件（二选一） input_group = parser.add_mutually_exclusive_group(required=True) input_group.add_argument("--input", "-i", help="输入帧图片目录") input_group.add_argument("--video", "-v", help="输入视频文件（会自动截帧）") parser.add_argument("--output", "-o", required=True, help="输出 JSON 文件路径") parser.add_argument("--fps", "-f", type=float, default=1.0, help="截帧帧率 (默认: 1帧/秒)") parser.add_argument("--engine", "-e", choices=["easyocr", "tesseract"], default="easyocr", help="OCR 引擎 (默认: easyocr)") parser.add_argument("--merge", action="store_true", help="同时保存合并的文本文件") parser.add_argument("--languages", "-l", nargs="+", default=["ch_sim", "en"], help="语言列表") parser.add_argument("--progress", "-p", help="可选：进度文件路径（会实时更新）") args = parser.parse_args() # 如果指定了视频文件，先截帧 input_dir = args.input if args.video: if not os.path.exists(args.video): print(f"❌ 视频文件不存在: {args.video}") return 1 # 创建临时帧目录 frames_dir = Path(args.video).stem + "_frames" input_dir = extract_frames(args.video, frames_dir, fps=args.fps) if not input_dir: return 1 # 初始化处理器 processor = OCRProcessor(engine=args.engine, languages=args.languages) # 处理帧 results = processor.process_frames(input_dir, fps=args.fps, progress_file=args.progress) if results: # 保存结果 processor.save_results(results, args.output, merge=args.merge) # 打印摘要 total_texts = sum(1 for r in results if r.get("text")) avg_conf = sum(r.get("confidence", 0) for r in results) / len(results) if results else 0 print(f"\n=== 摘要 ===") print(f"处理帧数: {len(results)}") print(f"识别到文字的帧: {total_texts}") print(f"平均置信度: {avg_conf:.2%}") else: print("未处理到任何结果") return 1 return 0 if __name__ == "__main__": exit(main())