scripts/pdf_ocr_v6.py
14.8 KB| 1 | #!/usr/bin/env python3 |
| 2 | """ |
| 3 | pdf_ocr_v6.py — PP-OCRv6 三档模型 OCR（ONNX Runtime） |
| 4 | |
| 5 | 使用 PP-OCRv6 模型家族进行本地 OCR，支持 Tiny / Small / Medium 三档选择。 |
| 6 | |
| 7 | 用法： |
| 8 | # Tiny 模型（极速，1.5MB） |
| 9 | python scripts/pdf_ocr_v6.py input.jpg --tier tiny |
| 10 | |
| 11 | # Small 模型（均衡，7.7MB） |
| 12 | python scripts/pdf_ocr_v6.py input.jpg --tier small -o output.txt |
| 13 | |
| 14 | # Medium 模型（高精度，34.5MB） |
| 15 | python scripts/pdf_ocr_v6.py input.jpg --tier medium -o output.txt |
| 16 | |
| 17 | # PDF 批量处理 |
| 18 | python scripts/pdf_ocr_v6.py input.pdf --tier medium -o ./output_dir |
| 19 | """ |
| 20 | |
| 21 | import argparse |
| 22 | import json |
| 23 | import os |
| 24 | import sys |
| 25 | from pathlib import Path |
| 26 | |
| 27 | try: |
| 28 | import cv2 |
| 29 | import numpy as np |
| 30 | import onnxruntime |
| 31 | except ImportError: |
| 32 | print("请先安装依赖: pip install onnxruntime opencv-python-headless numpy Pillow pdf2image") |
| 33 | sys.exit(1) |
| 34 | |
| 35 | |
| 36 | # ── 模型路径配置 ────────────────────────────────────────────── |
| 37 | |
| 38 | MODELS_BASE = os.path.expanduser(r"~\.workbuddy\.venvs\ppocrv6_models") |
| 39 | |
| 40 | MODEL_TIERS = { |
| 41 | "tiny": { |
| 42 | "dir": "ppocrv6_onnx", |
| 43 | "size": "1.5 MB", |
| 44 | "desc": "极速，可浏览器端运行", |
| 45 | }, |
| 46 | "small": { |
| 47 | "dir": "ppocrv6_small_onnx", |
| 48 | "size": "7.7 MB", |
| 49 | "desc": "性能均衡", |
| 50 | }, |
| 51 | "medium": { |
| 52 | "dir": "ppocrv6_medium_onnx", |
| 53 | "size": "34.5 MB", |
| 54 | "desc": "精度最高", |
| 55 | }, |
| 56 | } |
| 57 | |
| 58 | |
| 59 | # ── OCR 引擎 ────────────────────────────────────────────────── |
| 60 | |
| 61 | class PaddleOCREngine: |
| 62 | """PP-OCRv6 ONNX Runtime 推理引擎""" |
| 63 | |
| 64 | def __init__(self, tier: str = "tiny", device: str = "cpu"): |
| 65 | if tier not in MODEL_TIERS: |
| 66 | raise ValueError(f"不支持的模型档位: {tier}，可选: {list(MODEL_TIERS.keys())}") |
| 67 | |
| 68 | model_dir = os.path.join(MODELS_BASE, MODEL_TIERS[tier]["dir"]) |
| 69 | det_path = os.path.join(model_dir, "det", "inference.onnx") |
| 70 | rec_path = os.path.join(model_dir, "rec", "inference.onnx") |
| 71 | |
| 72 | if not os.path.exists(det_path) or not os.path.exists(rec_path): |
| 73 | raise FileNotFoundError( |
| 74 | f"模型文件未找到: {model_dir}\n" |
| 75 | f"请先下载模型: cd {MODELS_BASE} && 从 https://github.com/andyhuo520/ppocrv6-studio/releases 下载 {tier}.tar.gz" |
| 76 | ) |
| 77 | |
| 78 | # 选择推理后端 |
| 79 | if device == "gpu": |
| 80 | try: |
| 81 | import onnxruntime as _ort |
| 82 | if "DmlExecutionProvider" not in _ort.get_available_providers(): |
| 83 | print("⚠️ DirectML 不可用，回退到 CPU。请安装: pip install onnxruntime-directml") |
| 84 | device = "cpu" |
| 85 | except ImportError: |
| 86 | print("⚠️ onnxruntime 未安装，无法检查 DirectML") |
| 87 | device = "cpu" |
| 88 | |
| 89 | if device == "gpu": |
| 90 | providers = ["DmlExecutionProvider", "CPUExecutionProvider"] |
| 91 | print(f" 🔧 后端: DirectML (Intel Arc GPU)") |
| 92 | else: |
| 93 | providers = ["CPUExecutionProvider"] |
| 94 | print(f" 🔧 后端: CPU") |
| 95 | |
| 96 | self.det_session = onnxruntime.InferenceSession(det_path, providers=providers) |
| 97 | self.rec_session = onnxruntime.InferenceSession(rec_path, providers=providers) |
| 98 | |
| 99 | # 获取输入输出名称 |
| 100 | self.det_input = self.det_session.get_inputs()[0].name |
| 101 | self.rec_input = self.rec_session.get_inputs()[0].name |
| 102 | |
| 103 | # 加载字符字典 (PP-OCRv6 的 char_dict.json 是 list, 不是 dict) |
| 104 | # Tiny 有独立的 char_dict.json; Small/Medium 需从 inference.yml 解析 |
| 105 | char_dict_path = os.path.join(model_dir, "rec", "char_dict.json") |
| 106 | yml_path = os.path.join(model_dir, "rec", "inference.yml") |
| 107 | |
| 108 | if os.path.exists(char_dict_path): |
| 109 | with open(char_dict_path, "r", encoding="utf-8") as f: |
| 110 | self.char_dict = json.load(f) |
| 111 | if isinstance(self.char_dict, dict): |
| 112 | self.char_dict = list(self.char_dict.values()) |
| 113 | elif os.path.exists(yml_path): |
| 114 | # 从 inference.yml 解析 character_dict 列表 |
| 115 | self.char_dict = self._parse_char_dict_from_yml(yml_path) |
| 116 | else: |
| 117 | self.char_dict = None |
| 118 | |
| 119 | if self.char_dict: |
| 120 | print(f" 📖 字符字典: {len(self.char_dict)} 个字符") |
| 121 | |
| 122 | self.tier = tier |
| 123 | self.device = device |
| 124 | |
| 125 | # ── det 预处理参数（来自 inference.yml）── |
| 126 | DET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) |
| 127 | DET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) |
| 128 | DET_THRESH = 0.2 # binarization threshold (from config) |
| 129 | DET_BOX_THRESH = 0.4 # box threshold (from config) |
| 130 | DET_MIN_AREA = 20 |
| 131 | DET_TARGET_SIZE = 640 |
| 132 | |
| 133 | # ── rec 参数（来自 inference.yml）── |
| 134 | REC_H, REC_W = 48, 320 |
| 135 | |
| 136 | |
| 137 | def _parse_char_dict_from_yml(yml_path: str) -> list: |
| 138 | """从 inference.yml 的 PostProcess.character_dict 列表中提取字符""" |
| 139 | import re |
| 140 | with open(yml_path, "r", encoding="utf-8") as f: |
| 141 | lines = f.readlines() |
| 142 | |
| 143 | # 找到 character_dict: 行，然后收集后续的 " - " 列表项 |
| 144 | chars = [] |
| 145 | in_dict = False |
| 146 | for line in lines: |
| 147 | if "character_dict:" in line: |
| 148 | in_dict = True |
| 149 | continue |
| 150 | if in_dict: |
| 151 | # 列表项格式: " - 'x'" 或 " - x" |
| 152 | m = re.match(r'\s*-\s+(.*)', line) |
| 153 | if m: |
| 154 | val = m.group(1).strip() |
| 155 | # 去除引号 |
| 156 | if (val.startswith("'") and val.endswith("'")) or \ |
| 157 | (val.startswith('"') and val.endswith('"')): |
| 158 | val = val[1:-1] |
| 159 | chars.append(val) |
| 160 | elif line.strip() and not line.strip().startswith("#"): |
| 161 | # 遇到非列表项且非注释，说明 character_dict 段结束 |
| 162 | if not line.startswith(" " * 4): # 缩进减少，段结束 |
| 163 | break |
| 164 | return chars |
| 165 | |
| 166 | def detect(self, img: np.ndarray) -> list: |
| 167 | """文本检测: 返回文本框列表 [[x1,y1,x2,y2,...], ...]""" |
| 168 | h, w = img.shape[:2] |
| 169 | # 保持长宽比 resize 到 640 |
| 170 | scale = self.DET_TARGET_SIZE / max(h, w) |
| 171 | new_h, new_w = int(h * scale), int(w * scale) |
| 172 | resized = cv2.resize(img, (new_w, new_h)) |
| 173 | # pad 到 640x640 |
| 174 | canvas = np.zeros((self.DET_TARGET_SIZE, self.DET_TARGET_SIZE, 3), dtype=np.uint8) |
| 175 | canvas[:new_h, :new_w] = resized |
| 176 | |
| 177 | # 归一化: (img/255 - mean) / std, 然后 HWC→CHW |
| 178 | inp = canvas.astype(np.float32) / 255.0 |
| 179 | inp = (inp - self.DET_MEAN) / self.DET_STD |
| 180 | inp = inp.transpose(2, 0, 1) # HWC → CHW |
| 181 | inp = np.expand_dims(inp, axis=0) |
| 182 | |
| 183 | out = self.det_session.run(None, {self.det_input: inp})[0] |
| 184 | # 输出 shape: (1, 1, 640, 640) → squeeze 到 2D |
| 185 | prob_map = out.squeeze() # (640, 640) |
| 186 | return self._db_postprocess(prob_map, (w, h), (new_h, new_w)) |
| 187 | |
| 188 | def _db_postprocess(self, prob_map: np.ndarray, orig_size: tuple, resized_size: tuple, |
| 189 | threshold: float = None) -> list: |
| 190 | """DB 后处理: prob_map → 文本框""" |
| 191 | if threshold is None: |
| 192 | threshold = self.DET_THRESH |
| 193 | |
| 194 | h_prob, w_prob = prob_map.shape |
| 195 | bin_map = (prob_map > threshold).astype(np.uint8) * 255 |
| 196 | contours, _ = cv2.findContours(bin_map, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
| 197 | |
| 198 | ow, oh = orig_size |
| 199 | new_h, new_w = resized_size |
| 200 | boxes = [] |
| 201 | for cnt in contours: |
| 202 | area = cv2.contourArea(cnt) |
| 203 | if area < self.DET_MIN_AREA: |
| 204 | continue |
| 205 | # 计算最小外接矩形 |
| 206 | rect = cv2.minAreaRect(cnt) |
| 207 | pts = cv2.boxPoints(rect) |
| 208 | # 只取 prob_map 有效区域内的框（排除 padding 区产生的假框） |
| 209 | if pts[:, 0].max() > new_w + 5 or pts[:, 1].max() > new_h + 5: |
| 210 | continue |
| 211 | # 缩放回原图尺寸 |
| 212 | pts[:, 0] = pts[:, 0] / new_w * ow |
| 213 | pts[:, 1] = pts[:, 1] / new_h * oh |
| 214 | boxes.append(pts.astype(np.int32).tolist()) |
| 215 | return boxes |
| 216 | |
| 217 | def recognize(self, img: np.ndarray, box: list) -> tuple[str, float]: |
| 218 | """文本框识别: 返回 (文本, 置信度)""" |
| 219 | pts = np.array(box, dtype=np.float32) |
| 220 | x, y, bw, bh = cv2.boundingRect(pts) |
| 221 | pad = max(3, int(bh * 0.1)) |
| 222 | x1 = max(0, x - pad) |
| 223 | y1 = max(0, y - pad) |
| 224 | x2 = min(img.shape[1], x + bw + pad) |
| 225 | y2 = min(img.shape[0], y + bh + pad) |
| 226 | crop = img[y1:y2, x1:x2] |
| 227 | |
| 228 | if crop.size == 0 or crop.shape[0] < 5 or crop.shape[1] < 5: |
| 229 | return "", 0.0 |
| 230 | |
| 231 | # 保持宽高比 resize 到 (48, 320) |
| 232 | h, w = crop.shape[:2] |
| 233 | ratio = w / h |
| 234 | target_w = min(int(self.REC_H * ratio), self.REC_W) |
| 235 | target_w = max(target_w, 1) |
| 236 | resized = cv2.resize(crop, (target_w, self.REC_H)) |
| 237 | # pad 宽度到 320 |
| 238 | canvas = np.zeros((self.REC_H, self.REC_W, 3), dtype=np.uint8) |
| 239 | canvas[:, :target_w] = resized |
| 240 | |
| 241 | # 归一化到 [0, 1], HWC → CHW |
| 242 | inp = canvas.astype(np.float32) / 255.0 |
| 243 | inp = inp.transpose(2, 0, 1) |
| 244 | inp = np.expand_dims(inp, axis=0) |
| 245 | |
| 246 | outs = self.rec_session.run(None, {self.rec_input: inp}) |
| 247 | probs = outs[0][0] # shape: (T, num_classes) |
| 248 | pred_idx = probs.argmax(axis=-1) |
| 249 | |
| 250 | # CTC 解码: char_dict 是 list, index 0 = blank |
| 251 | text = "" |
| 252 | conf_sum = 0.0 |
| 253 | count = 0 |
| 254 | last_idx = 0 # 0 = blank |
| 255 | |
| 256 | for t in range(len(pred_idx)): |
| 257 | idx = int(pred_idx[t]) |
| 258 | if idx != 0 and idx != last_idx: |
| 259 | # char_dict[idx-1] 对应模型输出 idx (idx 0 = blank) |
| 260 | if self.char_dict and idx - 1 < len(self.char_dict): |
| 261 | char = self.char_dict[idx - 1] |
| 262 | text += char |
| 263 | conf_sum += float(probs[t, idx]) |
| 264 | count += 1 |
| 265 | last_idx = idx |
| 266 | |
| 267 | conf = conf_sum / max(count, 1) |
| 268 | return text.strip(), conf |
| 269 | |
| 270 | def ocr_image(self, img: np.ndarray, min_conf: float = 0.5) -> list[dict]: |
| 271 | """完整 OCR 流程: 检测 → 识别""" |
| 272 | boxes = self.detect(img) |
| 273 | results = [] |
| 274 | for box in boxes: |
| 275 | text, conf = self.recognize(img, box) |
| 276 | if text and conf >= min_conf: |
| 277 | results.append({ |
| 278 | "text": text, |
| 279 | "confidence": round(float(conf), 4), |
| 280 | "box": box, |
| 281 | }) |
| 282 | return results |
| 283 | |
| 284 | |
| 285 | # ── CLI ─────────────────────────────────────────────────────── |
| 286 | |
| 287 | def main(): |
| 288 | parser = argparse.ArgumentParser( |
| 289 | description="PP-OCRv6 三档模型 OCR — ONNX Runtime 本地推理", |
| 290 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 291 | epilog=""" |
| 292 | 模型档位: |
| 293 | tiny (1.5 MB) 极速，可浏览器端运行 |
| 294 | small (7.7 MB) 性能均衡 |
| 295 | medium (34.5 MB) 精度最高 |
| 296 | |
| 297 | 示例: |
| 298 | %(prog)s input.jpg --tier medium |
| 299 | %(prog)s input.jpg --tier tiny -o output.txt |
| 300 | %(prog)s input.pdf --tier medium -o ./output |
| 301 | %(prog)s input.jpg --tier small --device gpu |
| 302 | """, |
| 303 | ) |
| 304 | parser.add_argument("input", help="输入文件 (图片或 PDF)") |
| 305 | parser.add_argument("--tier", "-t", choices=["tiny", "small", "medium"], |
| 306 | default="tiny", help="模型档位 (默认: tiny)") |
| 307 | parser.add_argument("--device", "-d", choices=["cpu", "gpu"], |
| 308 | default="gpu", help="推理后端 (默认: gpu (DirectML)，cpu 用 CPU)") |
| 309 | parser.add_argument("--output", "-o", help="输出文件或目录") |
| 310 | parser.add_argument("--min-conf", type=float, default=0.5, |
| 311 | help="最低置信度阈值 (默认: 0.5)") |
| 312 | |
| 313 | args = parser.parse_args() |
| 314 | |
| 315 | input_path = args.input |
| 316 | if not os.path.exists(input_path): |
| 317 | print(f"❌ 文件不存在: {input_path}") |
| 318 | sys.exit(1) |
| 319 | |
| 320 | # 加载引擎 |
| 321 | print(f"🔧 加载 PP-OCRv6 {args.tier} 模型...") |
| 322 | try: |
| 323 | engine = PaddleOCREngine(tier=args.tier, device=args.device) |
| 324 | except FileNotFoundError as e: |
| 325 | print(f"❌ {e}") |
| 326 | sys.exit(1) |
| 327 | print(f"✅ 模型已加载") |
| 328 | |
| 329 | # PDF 或图片 |
| 330 | ext = input_path.lower() |
| 331 | if ext.endswith(".pdf"): |
| 332 | try: |
| 333 | from pdf2image import convert_from_path |
| 334 | except ImportError: |
| 335 | print("❌ PDF 处理需要 pdf2image: pip install pdf2image") |
| 336 | sys.exit(1) |
| 337 | print(f"📄 转换 PDF → 图片...") |
| 338 | images = convert_from_path(input_path, dpi=300) |
| 339 | imgs = [cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) for img in images] |
| 340 | else: |
| 341 | img = cv2.imread(input_path) |
| 342 | if img is None: |
| 343 | print(f"❌ 无法读取图片: {input_path}") |
| 344 | sys.exit(1) |
| 345 | imgs = [img] |
| 346 | |
| 347 | # OCR |
| 348 | all_results = [] |
| 349 | for i, img in enumerate(imgs): |
| 350 | print(f" OCR 第 {i+1}/{len(imgs)} 页...") |
| 351 | results = engine.ocr_image(img, min_conf=args.min_conf) |
| 352 | all_results.append({"page": i + 1, "results": results}) |
| 353 | print(f" 识别 {len(results)} 个文本块") |
| 354 | |
| 355 | # 输出 |
| 356 | basename = os.path.splitext(os.path.basename(input_path))[0] |
| 357 | |
| 358 | if args.output: |
| 359 | out_path = args.output |
| 360 | if os.path.isdir(out_path) or len(imgs) > 1: |
| 361 | os.makedirs(out_path, exist_ok=True) |
| 362 | json_path = os.path.join(out_path, f"{basename}.json") |
| 363 | text_path = os.path.join(out_path, f"{basename}.txt") |
| 364 | else: |
| 365 | json_path = out_path.replace(".txt", ".json") if out_path.endswith(".txt") else out_path + ".json" |
| 366 | text_path = out_path |
| 367 | else: |
| 368 | json_path = f"{basename}_ocr.json" |
| 369 | text_path = f"{basename}_ocr.txt" |
| 370 | |
| 371 | # 保存 JSON |
| 372 | output_data = { |
| 373 | "engine": f"PP-OCRv6 {args.tier}", |
| 374 | "model_size": MODEL_TIERS[args.tier]["size"], |
| 375 | "pages": all_results, |
| 376 | } |
| 377 | with open(json_path, "w", encoding="utf-8") as f: |
| 378 | json.dump(output_data, f, ensure_ascii=False, indent=2) |
| 379 | print(f"\n ✅ JSON: {json_path}") |
| 380 | |
| 381 | # 保存文本 |
| 382 | lines = [] |
| 383 | for page in all_results: |
| 384 | lines.append(f"--- 第 {page['page']} 页 ---") |
| 385 | for r in page["results"]: |
| 386 | lines.append(r["text"]) |
| 387 | lines.append("") |
| 388 | Path(text_path).write_text("\n".join(lines), encoding="utf-8") |
| 389 | print(f" ✅ 文本: {text_path}") |
| 390 | |
| 391 | # 汇总统计 |
| 392 | total_blocks = sum(len(p["results"]) for p in all_results) |
| 393 | avg_conf = sum(r["confidence"] for p in all_results for r in p["results"]) / max(total_blocks, 1) |
| 394 | print(f"\n📊 统计: {total_blocks} 个文本块, 平均置信度 {avg_conf:.2f}") |
| 395 | |
| 396 | |
| 397 | if __name__ == "__main__": |
| 398 | main() |
文件内容是平台检查时读出的文本，只供查看。平台不会运行包里的任何脚本或程序。
