src/det_and_ocr_pt.py

import os
import re
import cv2
import numpy as np
import onnxruntime as ort
from ultralytics import YOLO
from PIL import Image

DET_MODEL_PATH = "models/yolo11-nano-page-best.pt"
# DET_MODEL_PATH = "models/yolo11-small-page-best.pt"
# DET_MODEL_PATH = "models/yolo11-medium-page-best.pt"

OCR_MODEL_PATH = "models/ch_PP-OCRv3_rec_infer.onnx"
OCR_KEYS_PATH = "models/ppocr_keys_v1.txt"


def load_det_model(model_path=DET_MODEL_PATH):
    return YOLO(model_path)


def load_ocr_model(model_path=OCR_MODEL_PATH, keys_path=OCR_KEYS_PATH):
    sess = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
    with open(keys_path, "r", encoding="utf-8") as f:
        keys = [line.strip() for line in f.readlines()]
    return sess, keys


class OCRProcessor:
    def __init__(self, model_path=OCR_MODEL_PATH, keys_path=OCR_KEYS_PATH):
        sess_options = ort.SessionOptions()
        sess_options.log_severity_level = 3
        self.sess = ort.InferenceSession(model_path, sess_options)
        self.input_name = self.sess.get_inputs()[0].name
        self.keys = self._load_keys(keys_path)

    def _load_keys(self, keys_path):
        with open(keys_path, "r", encoding="utf-8") as f:
            keys = [line.strip() for line in f.readlines()]
        out_dims = self.sess.get_outputs()[0].shape[-1]
        keys = ["[blank]"] + keys
        while len(keys) < out_dims:
            keys.append(f"[unused{len(keys) - 1}]")
        return keys[:out_dims]

    def _preprocess(self, img):
        h, w = img.shape[:2]
        if len(img.shape) == 3 and img.shape[2] == 3:
            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        elif len(img.shape) == 2:
            img_rgb = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
        else:
            img_rgb = img

        gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
        _, thresh = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)
        coords = cv2.findNonZero(thresh)
        if coords is not None:
            x, y, cw, ch = cv2.boundingRect(coords)
            margin = 2
            x = max(0, x - margin)
            y = max(0, y - margin)
            cw = min(w - x, cw + 2 * margin)
            ch = min(h - y, ch + 2 * margin)
            if cw > 0 and ch > 0:
                img_rgb = img_rgb[y:y+ch, x:x+cw]
                w, h = img_rgb.shape[1], img_rgb.shape[0]

        ratio = w / float(h)
        new_h = 48
        new_w = max(32, int(new_h * ratio))
        new_w = (new_w + 31) // 32 * 32

        if h > new_h or w > new_w:
            interp = cv2.INTER_AREA
        else:
            interp = cv2.INTER_LANCZOS4

        resized = cv2.resize(img_rgb, (new_w, new_h), interpolation=interp)
        img_mean = (0.5, 0.5, 0.5)
        img_std = (0.5, 0.5, 0.5)
        resized = (resized.astype(np.float32) / 255.0 - img_mean) / img_std
        resized = resized.transpose(2, 0, 1)
        resized = np.expand_dims(resized, axis=0).astype(np.float32)
        return resized

    def _postprocess(self, preds):
        preds = preds.reshape(-1, preds.shape[-1])
        pred_id = np.argmax(preds, axis=1)
        char_list = []
        conf_list = []
        for i in range(len(pred_id)):
            if pred_id[i] > 0 and (i == 0 or pred_id[i] != pred_id[i - 1]):
                char_list.append(self.keys[pred_id[i]])
                conf_list.append(float(preds[i][pred_id[i]]))
        text = "".join(char_list)
        text = re.sub(r'\[unused\d+\]', ' ', text)
        avg_conf = float(np.mean(conf_list)) if conf_list else 0.0
        return text, avg_conf

    def recognize(self, img):
        if isinstance(img, str):
            img = cv2.imread(img)
        elif isinstance(img, Image.Image):
            img = np.array(img)
            if len(img.shape) == 2:
                img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
        elif isinstance(img, np.ndarray) and len(img.shape) == 3 and img.shape[2] == 3:
            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        input_data = self._preprocess(img)
        preds = self.sess.run(None, {self.input_name: input_data})[0]
        return self._postprocess(preds)


def filter_text(text):
    # 判断文本中是否包含数字和字母
    has_digit = bool(re.search(r"\d", text))
    has_letter = bool(re.search(r"[a-zA-Z]", text))

    # 情况1：既没有数字也没有字母，返回空
    if not has_digit and not has_letter:
        return ""

    ## always has outlier, so don't use this
    # # 情况2：只有数字没有字母，保留数字和空格，其他字符替换为空
    # if has_digit and not has_letter:
    #     return "".join(ch if ch.isdigit() or ch == " " else "" for ch in text)

    # # 情况3：既有数字又有字母，去掉标点符号 ; , . *
    # if has_digit and has_letter:
    #     return re.sub(r"[;,.+-]", "", text)

    return text


def process_image(image_path, det_model=None, ocr_processor=None):
    if det_model is None:
        det_model = load_det_model()
    if ocr_processor is None:
        ocr_processor = OCRProcessor()

    img_pil = Image.open(image_path).convert("RGB")
    img_w, img_h = img_pil.size
    results = det_model(img_pil, verbose=False)

    output = []
    picture_name = os.path.splitext(os.path.basename(image_path))[0]
    os.makedirs("./tmp", exist_ok=True)
    idx = 0

    for result in results:
        boxes = result.boxes
        if boxes is None:
            continue
        for box in boxes:
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            conf = float(box.conf[0])

            x1_i, y1_i, x2_i, y2_i = int(x1), int(y1), int(x2), int(y2)
            crop = img_pil.crop((x1_i, y1_i, x2_i, y2_i))
            crop.save(f"./tmp/{picture_name}_c{idx}.jpg", "JPEG")

            crop_np = np.array(crop)
            text, _ = ocr_processor.recognize(crop_np)
            text = filter_text(text)

            x1_n, y1_n = x1 / img_w, y1 / img_h
            x2_n, y2_n = x2 / img_w, y2 / img_h
            output.append([x1_n, y1_n, x2_n, y2_n, conf, text])
            idx += 1

    return output


def process_directory(dir_path, det_model=None, ocr_processor=None):
    if det_model is None:
        det_model = load_det_model()
    if ocr_processor is None:
        ocr_processor = OCRProcessor()

    exts = (".png", ".jpg", ".jpeg", ".bmp")
    img_files = sorted([f for f in os.listdir(dir_path) if f.lower().endswith(exts)])

    results = {}
    for fname in img_files:
        img_path = os.path.join(dir_path, fname)
        res = process_image(img_path, det_model, ocr_processor)
        results[fname] = res

    return results


if __name__ == "__main__":
    import sys
    import json

    det_model = load_det_model()
    ocr_processor = OCRProcessor()

    if len(sys.argv) == 2:
        path = sys.argv[1]
        if os.path.isfile(path):
            res = process_image(path, det_model, ocr_processor)
            print(json.dumps(res, ensure_ascii=False, indent=2))
        elif os.path.isdir(path):
            results = process_directory(path, det_model, ocr_processor)
            for fname, res in results.items():
                print(f"[{fname}]")
                print(json.dumps(res, ensure_ascii=False, indent=2))
    else:
        print("Usage: python det_and_ocr.py <image_path_or_directory>")