| 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" |
| |
| |
|
|
| 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)) |
|
|
| |
| if not has_digit and not has_letter: |
| return "" |
|
|
| |
| |
| |
| |
|
|
| |
| |
| |
|
|
| 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>") |
|
|
