diff --git a/skills/kb-ocr-processor/ocr_processor.py b/skills/kb-ocr-processor/ocr_processor.py index b793ae6..d905a05 100644 --- a/skills/kb-ocr-processor/ocr_processor.py +++ b/skills/kb-ocr-processor/ocr_processor.py @@ -93,46 +93,81 @@ class OCRProcessor: "confidence": 0.0 } - def process_frames(self, input_dir: str, fps: float = 0.1) -> List[Dict]: + 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 [] - - print(f"Found {len(image_files)} images in {input_dir}") + + 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): - print(f"Processing [{i+1}/{len(image_files)}]: {image_path}") + 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"] = os.path.basename(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: @@ -205,11 +240,12 @@ def main(): 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"], + 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="+", + parser.add_argument("--languages", "-l", nargs="+", default=["ch_sim", "en"], help="语言列表") + parser.add_argument("--progress", "-p", help="可选:进度文件路径(会实时更新)") args = parser.parse_args() @@ -217,7 +253,7 @@ def main(): processor = OCRProcessor(engine=args.engine, languages=args.languages) # 处理帧 - results = processor.process_frames(args.input, fps=args.fps) + results = processor.process_frames(args.input, fps=args.fps, progress_file=args.progress) if results: # 保存结果