361 lines
12 KiB
Python
361 lines
12 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 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(json.dumps({
|
|
"success": False,
|
|
"error": f"视频文件不存在: {args.video}"
|
|
}, ensure_ascii=False, indent=2))
|
|
return 1
|
|
|
|
# 创建临时帧目录
|
|
frames_dir = Path(args.video).stem + "_frames"
|
|
input_dir = extract_frames(args.video, frames_dir, fps=args.fps)
|
|
|
|
if not input_dir:
|
|
print(json.dumps({
|
|
"success": False,
|
|
"error": "截帧失败"
|
|
}, ensure_ascii=False, indent=2))
|
|
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
|
|
|
|
output = {
|
|
"success": True,
|
|
"filename": Path(args.output).name,
|
|
"filepath": str(Path(args.output).resolve()),
|
|
"message": f"OCR处理完成",
|
|
"stats": {
|
|
"total_frames": len(results),
|
|
"frames_with_text": total_texts,
|
|
"avg_confidence": round(avg_conf, 4)
|
|
}
|
|
}
|
|
print(json.dumps(output, ensure_ascii=False, indent=2))
|
|
return 0
|
|
else:
|
|
print(json.dumps({
|
|
"success": False,
|
|
"error": "未处理到任何结果"
|
|
}, ensure_ascii=False, indent=2))
|
|
return 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main())
|