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 ")