OpenClaw Agent
279 lines
9.2 KiB
Python
279 lines
9.2 KiB
Python
#!/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())
|