CheckCardBoxAndDefectServer/app/core/model_loader.py

from typing import Dict
from .config import settings
from ..utils.card_inference.fry_bisenetv2_predictor_V04_250819 import FryBisenetV2Predictor

# 全局的模型预测器字典
predictors: Dict[str, FryBisenetV2Predictor] = {}
'''
predictors['outer_box'] ...
'''


def load_models():
    print("--- 开始加载卡片识别模型 ---")
    for name, config in settings.CARD_MODELS_CONFIG.items():
        print(f"... 正在加载模型: {name} ...")
        try:
            predictor = FryBisenetV2Predictor(
                pth_path=config['pth_path'],
                real_seg_class_dict=config['class_dict'],
                imgSize_train_dict=config['img_size'],
                confidence=config['confidence'],
                input_channels=config['input_channels']
            )
            predictors[name] = predictor
            print(f"--- 模型 '{name}' 加载成功 ---")
        except Exception as e:
            print(f"!!! 模型 '{name}' 加载失败: {e} !!!")


def unload_models():
    """在应用关闭时清理资源"""
    print("... 卸载模型 ...")
    predictors.clear()


def get_predictor(name: str) -> FryBisenetV2Predictor:
    """获取一个已加载的预测器实例"""
    predictor = predictors.get(name)
    if not predictor:
        raise ValueError(f"模型 '{name}' 不存在或未成功加载。可用模型: {list(predictors.keys())}")
    return predictor