90 lines
3.3 KiB
Python
90 lines
3.3 KiB
Python
#!/usr/bin/env python3
|
||
"""KOCR v5 — OCR 识别模块（简化版：原图直出，无预处理）"""
|
||
|
||
import cv2
|
||
from paddleocr import PaddleOCR
|
||
|
||
|
||
_ocr = None
|
||
|
||
def get_ocr_instance(config: dict):
|
||
"""懒加载 PaddleOCR 实例"""
|
||
global _ocr
|
||
if _ocr is None:
|
||
_ocr = PaddleOCR(
|
||
use_angle_cls=config.get('use_angle_cls', True),
|
||
lang=config.get('lang', 'ch'),
|
||
use_gpu=config.get('use_gpu', False),
|
||
show_log=config.get('show_log', False),
|
||
det_db_thresh=config.get('det_db_thresh', 0.3),
|
||
det_db_box_thresh=config.get('det_db_box_thresh', 0.5),
|
||
det_db_unclip_ratio=config.get('det_db_unclip_ratio', 2.2),
|
||
det_limit_side_len=config.get('max_side_len', 1600),
|
||
det_db_score_mode=config.get('det_db_score_mode', 'slow'),
|
||
min_size=config.get('min_size', 8),
|
||
rec_drop_score=config.get('rec_drop_score', 0.5),
|
||
)
|
||
return _ocr
|
||
|
||
|
||
def run_ocr(image, ocr_config: dict, photo_type: str = 'normal') -> list:
|
||
"""对图像执行 OCR，返回 [(text, confidence, cx, cy, box), ...]
|
||
|
||
输入:
|
||
image: numpy array (cv2.imread) 或 图像文件路径
|
||
ocr_config: dict（PaddleOCR 参数）
|
||
photo_type: str — 照片类型，用于动态调参
|
||
|
||
返回:
|
||
排序后的 lines: [(text, conf, cx, cy, box), ...]
|
||
按 (y, x) 排序
|
||
"""
|
||
# 根据照片类型动态调整 OCR 参数
|
||
config_override = dict(ocr_config)
|
||
if photo_type == 'flash':
|
||
config_override['det_db_thresh'] = config_override.get('det_db_thresh_flash', 0.18)
|
||
elif photo_type == 'dark':
|
||
config_override['det_db_thresh'] = config_override.get('det_db_thresh_dark', 0.25)
|
||
elif photo_type == 'blurry':
|
||
config_override['det_db_thresh'] = config_override.get('det_db_thresh_blurry', 0.2)
|
||
# normal 保持默认
|
||
|
||
ocr = get_ocr_instance(config_override)
|
||
|
||
# 支持 numpy array 和 文件路径
|
||
result = ocr(image)
|
||
|
||
# PaddleOCR v2.8 返回格式多变，兼容处理
|
||
lines_raw = []
|
||
if result is None:
|
||
return []
|
||
|
||
# tuple of 3: (boxes, texts, elapsed)
|
||
if isinstance(result, tuple) and len(result) >= 2:
|
||
boxes = result[0]
|
||
texts = result[1]
|
||
if boxes and texts:
|
||
for box, (text, conf) in zip(boxes, texts):
|
||
if conf < 0.4:
|
||
continue
|
||
cx = float((box[0][0] + box[2][0]) / 2)
|
||
cy = float((box[0][1] + box[2][1]) / 2)
|
||
lines_raw.append((text, conf, cx, cy, box))
|
||
|
||
# 旧版格式: list of [[box], (text, conf), ...]
|
||
elif isinstance(result, list):
|
||
for item in result:
|
||
if isinstance(item, (list, tuple)) and len(item) >= 2:
|
||
box = item[0]
|
||
if isinstance(box, (list, tuple)) and len(box) >= 1:
|
||
text, conf = item[1]
|
||
if conf < 0.4:
|
||
continue
|
||
cx = float((box[0][0] + box[2][0]) / 2)
|
||
cy = float((box[0][1] + box[2][1]) / 2)
|
||
lines_raw.append((text, conf, cx, cy, box))
|
||
|
||
# 按 y 排序（从上到下），同 y 按 x 排序（从左到右）
|
||
lines_raw.sort(key=lambda l: (l[3], l[2]))
|
||
return lines_raw
|