#!/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 main(): """命令行入口""" parser = argparse.ArgumentParser( description="OCR 视频帧文字识别工具", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" 示例: # 使用 EasyOCR 处理帧目录 python ocr_processor.py --input ./frames/ --output ./ocr.json # 使用 Tesseract 处理 python ocr_processor.py --input ./frames/ --output ./ocr.json --engine tesseract # 指定采样帧率 python ocr_processor.py --input ./frames/ --output ./ocr.json --fps 0.5 # 同时生成合并的文本文件 python ocr_processor.py --input ./frames/ --output ./ocr.json --merge """ ) parser.add_argument("--input", "-i", required=True, help="输入帧图片目录") parser.add_argument("--output", "-o", required=True, help="输出 JSON 文件路径") parser.add_argument("--fps", "-f", type=float, default=0.1, help="采样帧率 (默认: 0.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() # 初始化处理器 processor = OCRProcessor(engine=args.engine, languages=args.languages) # 处理帧 results = processor.process_frames(args.input, 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())