refactor(ocr): use ppocrv6 small model only

Author: A-nony-mous

src/one_dragon/base/config/basic_model_config.py

@@ -1,7 +1,13 @@
 from one_dragon.base.config.config_item import ConfigItem
 from one_dragon.base.config.yaml_config import YamlConfig
-from one_dragon.base.matcher.ocr.onnx_ocr_matcher import DEFAULT_OCR_MODEL_NAME, get_ocr_model_dir, \
-    get_ocr_download_url_github, get_ocr_download_url_gitee, get_final_file_list
+from one_dragon.base.matcher.ocr.onnx_ocr_matcher import (
+    DEFAULT_OCR_MODEL_NAME,
+    PPOCRV6_MODEL_NAME,
+    get_final_file_list,
+    get_ocr_download_url_gitee,
+    get_ocr_download_url_github,
+    get_ocr_model_dir,
+)
 from one_dragon.base.web.common_downloader import CommonDownloaderParam
 
 
@@ -34,7 +40,7 @@ def using_old_model(self) -> bool:
         pass
 
 def get_ocr_opts() -> list[ConfigItem]:
-    models_list = [DEFAULT_OCR_MODEL_NAME, 'ppocrv6_small', 'ppocrv6_tiny']
+    models_list = [DEFAULT_OCR_MODEL_NAME, PPOCRV6_MODEL_NAME]
     config_list: list[ConfigItem] = []
     for model in models_list:
         model_dir = get_ocr_model_dir(model)

src/one_dragon/base/matcher/ocr/onnx_ocr_matcher.py

@@ -18,11 +18,22 @@
 from one_dragon.utils.log_utils import log
 
 DEFAULT_OCR_MODEL_NAME: str = 'ppocrv5'
+PPOCRV6_MODEL_NAME: str = 'ppocrv6'
 GITHUB_DOWNLOAD_URL: str = 'https://github.com/OneDragon-Anything/OneDragon-Env/releases/download'
 GITEE_DOWNLOAD_URL: str = 'https://gitee.com/OneDragon-Anything/OneDragon-Env/releases/download'
 
 
+def normalize_ocr_model_name(ocr_model_name: str) -> str:
+    """
+    规范化 OCR 模型名。
+    """
+    if ocr_model_name.startswith(PPOCRV6_MODEL_NAME):
+        return PPOCRV6_MODEL_NAME
+    return ocr_model_name
+
+
 def get_ocr_model_dir(ocr_model_name: str) -> str:
+    ocr_model_name = normalize_ocr_model_name(ocr_model_name)
     return os_utils.get_path_under_work_dir('assets', 'models', 'onnx_ocr', ocr_model_name)
 
 
@@ -35,6 +46,7 @@ def get_ocr_download_url_gitee(ocr_model_name: str) -> str:
 
 
 def get_ocr_download_url(website: str, ocr_model_name: str) -> str:
+    ocr_model_name = normalize_ocr_model_name(ocr_model_name)
     return f'{website}/{ocr_model_name}/{ocr_model_name}.zip'
 
 
@@ -88,13 +100,13 @@ def __init__(
             det_limit_side_len: float = 960.0,
             ocr_model_size: str | None = None,
     ):
-        self.ocr_model_name: str = ocr_model_name
+        self.ocr_model_name: str = normalize_ocr_model_name(ocr_model_name)
         self.models_dir: str = get_ocr_model_dir(ocr_model_name)
         if dict_name is None:
             dict_name = get_ocr_model_dict_name(ocr_model_name)
             if dict_name is None:
                 # 首次运行未下载时，根据模型名推导一个默认的字典文件名，避免崩溃
-                dict_name = f"{ocr_model_name}_dict.txt"
+                dict_name = f"{self.ocr_model_name}_dict.txt"
         # ===================================================================
         # I. 设备与性能 (Device & Performance)
         # ===================================================================
@@ -113,15 +125,8 @@ def __init__(
         # III. 核心功能开关 (Core Feature Switches)
         # ===================================================================
         self.use_angle_cls = use_angle_cls  # 是否加载并使用方向分类模型
-        if ocr_model_size is None:
-            if 'medium' in ocr_model_name:
-                ocr_model_size = 'medium'
-            elif 'tiny' in ocr_model_name:
-                ocr_model_size = 'tiny'
-            elif 'small' in ocr_model_name:
-                ocr_model_size = 'small'
-            else:
-                ocr_model_size = 'small'
+        if self.ocr_model_name == PPOCRV6_MODEL_NAME or ocr_model_size is None:
+            ocr_model_size = 'small'
         self.ocr_model_size: str | None = ocr_model_size
 
         # ===================================================================
@@ -400,7 +405,7 @@ def match_words(
         """
         all_match_result: dict = self.run_ocr(image, threshold, merge_line_distance=merge_line_distance)
         match_key = set()
-        for k in all_match_result.keys():
+        for k in all_match_result:
             for w in words:
                 ocr_result: str = k
                 ocr_target = gt(w, 'ocr')
@@ -533,7 +538,7 @@ def _emit_overlay_vision_from_ocr_results(
             return
 
         offset_x, offset_y = bus.crop_offset
-        for i, result in enumerate(ocr_results[:60]):
+        for result in ocr_results[:60]:
             label = str(result.data or "").strip()
             if len(label) > 32:
                 label = label[:29] + "..."

src/onnxocr/onnx_paddleocr.py

@@ -12,60 +12,49 @@
 
 log = get_logger("onnx_paddleocr")
 
-PPOCRV6_MODEL_CONFIGS = {
-    "medium": {
-        "det_db_box_thresh": 0.45,
-        "rec_char_dict_path": "ppocrv6_dict.txt",
-    },
-    "small": {
-        "det_db_box_thresh": 0.45,
-        "rec_char_dict_path": "ppocrv6_dict.txt",
-    },
-    "tiny": {
-        "det_db_box_thresh": 0.4,
-        "rec_char_dict_path": "ppocrv6_tiny_dict.txt",
-    },
+PPOCRV6_MODEL_CONFIG = {
+    "det_db_box_thresh": 0.45,
+    "rec_char_dict_path": "ppocrv6_dict.txt",
 }
 
 
-def _normalize_ppocrv6_size(model_name=None, model_size=None):
+def _normalize_ppocrv6_size(model_name: str | None = None, model_size: str | None = None) -> str | None:
     if not model_name:
         return None
-    if model_size:
-        size = str(model_size).lower()
-    else:
-        normalized = str(model_name).lower()
-        size = next((name for name in PPOCRV6_MODEL_CONFIGS if name in normalized), None)
-
-    if size not in PPOCRV6_MODEL_CONFIGS:
+    if "ppocrv6" not in str(model_name).lower():
         return None
-    return size
+    return "small"
+
+
+def _normalize_ppocrv6_model_name(model_name: str) -> str:
+    if "ppocrv6" in model_name.lower():
+        return "ppocrv6"
+    return model_name
 
 
-def _build_ppocrv6_defaults(kwargs):
+def _build_ppocrv6_defaults(kwargs: dict[str, Any]) -> dict[str, Any]:
     """为 PP-OCRv6 模型构建默认参数。如果不是 v6 模型则返回空字典，保证 v5 兼容。"""
     model_name = kwargs.pop("ocr_model_name", None)
     if not model_name:
         kwargs.pop("ocr_model_size", None)
         return {}
     model_size = kwargs.pop("ocr_model_size", None)
-    size = _normalize_ppocrv6_size(model_name=model_name, model_size=model_size)
-    if not size:
+    if _normalize_ppocrv6_size(model_name=model_name, model_size=model_size) is None:
         return {}
+    model_name = _normalize_ppocrv6_model_name(model_name)
 
     from one_dragon.utils import os_utils
     model_root = Path(os_utils.get_path_under_work_dir('assets', 'models', 'onnx_ocr', model_name))
-    config = PPOCRV6_MODEL_CONFIGS[size]
 
     defaults = {
         "det_model_dir": str(model_root / "det.onnx"),
         "rec_model_dir": str(model_root / "rec.onnx"),
-        "rec_char_dict_path": str(model_root / config["rec_char_dict_path"]),
+        "rec_char_dict_path": str(model_root / PPOCRV6_MODEL_CONFIG["rec_char_dict_path"]),
         "rec_image_shape": "3, 48, 320",
         "det_limit_side_len": 960,
         "det_limit_type": "max",
         "det_db_thresh": 0.3,
-        "det_db_box_thresh": config["det_db_box_thresh"],
+        "det_db_box_thresh": PPOCRV6_MODEL_CONFIG["det_db_box_thresh"],
         "det_db_unclip_ratio": 1.5,
         "det_db_max_candidates": 1000,
     }