nanocode/urinecheck/backend/ml_config.py

"""
ML配置加载器 — 解析 /usr/local/Models/ml.json
直接翻译 C++ UrineCheck14::initconfig() + UrineCheck_Class::initconfig()

Author: 小月
Date: 2026-05-13
"""

import json
import logging
from collections import defaultdict
from typing import Dict, List, Optional, Callable, Tuple

import cv2
import numpy as np

logger = logging.getLogger("ml_config")

# 常量（与C++一致）
MAX_PAPER_TYPE = 21
MAX_PAPER_BLOCK = 15
MAX_ALG_TAG = 8

# 各检测项在 ml.json 中的 name → 结果索引函数类型映射
# algorithm=0 → Red, algorithm=1 → Green/Blue (原始)
ITEM_FUNC_MAP = {
    "uro": "uro",
    "glu": "glu",
    "ket": "ket",
    "bil": "bil",
    "pro": "pro",
    "nit": "nit",
    "ph":  "ph",
    "bld": "bld",
    "sg":  "sg",
    "wbc": "wbc",
    "vc":  "vc",
    "cre": "cre",
    "uca": "uca",
    "mau": "mau",
}


class MLConfigLoader:
    """ML模型配置加载器 — 全局单例"""

    def __init__(self, config_path: str = "/usr/local/Models/ml.json"):
        self.config_path = config_path

        # 三维查表: [paperType][position][rawIndex] → outputIndex
        self.itemsResultIdxMapTbl: List[List[List[int]]] = [
            [[-1] * MAX_ALG_TAG for _ in range(MAX_PAPER_BLOCK)]
            for _ in range(MAX_PAPER_TYPE)
        ]

        # 二维: [paperType][position] → 输出id
        self.itemsStatIds: List[List[int]] = [
            [-1] * MAX_PAPER_BLOCK for _ in range(MAX_PAPER_TYPE)
        ]

        # [paperType][position] → 已加载的 OpenCV RTrees 模型
        self.models: List[Dict[int, cv2.ml.RTrees]] = [
            {} for _ in range(MAX_PAPER_TYPE)
        ]

        # [paperType][position] → 结果索引函数类型名
        # "uro"|"glu"|...|"red_uro"|"red_glu"|...|"pet_glu"|"pet_sg"
        self.result_funcs: List[Dict[int, str]] = [
            {} for _ in range(MAX_PAPER_TYPE)
        ]

        # PRO/MAU/SG/pH 在各 paperType 上的 position
        self.positionsPro: List[int] = [-1] * MAX_PAPER_TYPE
        self.positionsMau: List[int] = [-1] * MAX_PAPER_TYPE
        self.positionsSg:  List[int] = [-1] * MAX_PAPER_TYPE
        self.positionsPh:  List[int] = [-1] * MAX_PAPER_TYPE

        # 文件名缓存（避免重复加载同一个 .data 文件）
        self._file_cache: Dict[str, cv2.ml.RTrees] = {}

        self._loaded = False

    def load(self) -> bool:
        """加载 ml.json 并初始化所有配置"""
        if self._loaded:
            return True

        try:
            with open(self.config_path, 'r') as f:
                config = json.load(f)
        except Exception as e:
            logger.error(f"无法加载 {self.config_path}: {e}")
            return False

        ml = config.get("ml", [])
        if not ml:
            logger.error("ml.json 中未找到 'ml' 数组")
            return False

        for item in ml:
            name = item.get("name", "")
            models_list = item.get("model", [])

            if name not in ITEM_FUNC_MAP:
                logger.warning(f"未知检测项目: {name}")
                continue

            for paper_type, block in enumerate(models_list):
                if paper_type >= MAX_PAPER_TYPE:
                    break

                position = block.get("positon", -1)
                if position < 0:
                    continue

                item_id = block.get("id", 0)
                filepath = block.get("filepath", "")
                algorithm = block.get("algorithm", 1)
                indexmap_str = block.get("indexmap", "0")

                # 记录 PRO/MAU/SG/pH 的 position
                if name == "pro":
                    self.positionsPro[paper_type] = position
                elif name == "mau":
                    self.positionsMau[paper_type] = position
                elif name == "sg":
                    self.positionsSg[paper_type] = position
                elif name == "ph":
                    self.positionsPh[paper_type] = position

                # 记录 output id
                self.itemsStatIds[paper_type][position] = item_id

                # 解析 indexmap
                index_map = [int(x.strip()) for x in indexmap_str.split(",")]
                for k, val in enumerate(index_map):
                    if k < MAX_ALG_TAG:
                        self.itemsResultIdxMapTbl[paper_type][position][k] = val

                # 加载模型文件（带缓存）
                if filepath not in self._file_cache:
                    try:
                        model = cv2.ml.RTrees_load(filepath)
                        if model is None or model.empty():
                            logger.error(f"模型加载失败: {filepath}")
                            continue
                        self._file_cache[filepath] = model
                    except Exception as e:
                        logger.error(f"模型加载异常 {filepath}: {e}")
                        continue

                self.models[paper_type][position] = self._file_cache[filepath]

                # 确定结果索引函数
                func_type = self._resolve_func_type(name, algorithm)
                self.result_funcs[paper_type][position] = func_type

        self._loaded = True
        count = sum(len(m) for m in self.models)
        logger.info(f"ML配置加载完成: {count} 个模型, {len(self._file_cache)} 个唯一文件")
        return True

    def _resolve_func_type(self, name: str, algorithm: int) -> str:
        """根据试纸类型确定结果索引函数名"""
        base = ITEM_FUNC_MAP[name]
        if algorithm == 0:
            # 红色试纸
            return f"red_{base}"
        elif algorithm == 2:
            # PET 试纸 (仅 glu/sg 有)
            return f"pet_{base}"
        else:
            # 绿色/蓝色试纸（原始）
            return base

    # ---- 便捷查询方法 ----

    def get_model(self, paper_type: int, position: int) -> Optional[cv2.ml.RTrees]:
        if paper_type < 0 or paper_type >= MAX_PAPER_TYPE:
            return None
        return self.models[paper_type].get(position)

    def get_func_type(self, paper_type: int, position: int) -> Optional[str]:
        if paper_type < 0 or paper_type >= MAX_PAPER_TYPE:
            return None
        return self.result_funcs[paper_type].get(position)

    def get_stat_id(self, paper_type: int, position: int) -> int:
        if paper_type < 0 or paper_type >= MAX_PAPER_TYPE:
            return -1
        return self.itemsStatIds[paper_type][position]

    def get_output_index(self, paper_type: int, position: int, raw_index: int) -> int:
        """通过 indexmap 将算法输出index映射为最终输出index"""
        if paper_type < 0 or paper_type >= MAX_PAPER_TYPE:
            return 0
        if position < 0 or position >= MAX_PAPER_BLOCK:
            return 0
        if raw_index < 0 or raw_index >= MAX_ALG_TAG:
            return 0
        mapped = self.itemsResultIdxMapTbl[paper_type][position][raw_index]
        return mapped if mapped >= 0 else 0

    def get_max_mau(self, paper_type: int) -> int:
        """MAU最大等级: red=3, others=2"""
        return 2 if paper_type == 0 else 3

    def get_max_sg(self, paper_type: int) -> int:
        """SG最大等级: red=6, others=5"""
        return 5 if paper_type == 0 else 6

    def has_model_at(self, paper_type: int, position: int) -> bool:
        if paper_type < 0 or paper_type >= MAX_PAPER_TYPE:
            return False
        return position in self.models[paper_type]