from pathlib import Path from typing import Dict, List from enum import Enum # 定义一个模型的配置结构 class CardModelConfig: pth_path: str class_dict: dict img_size: dict confidence: float input_channels: int class Settings: API_prefix: str = "/api/card_inference" BASE_PATH = Path(__file__).parent.parent.parent.absolute() TEMP_WORK_DIR = BASE_PATH / "_temp_work" pixel_resolution = 24.54 # 使用一个字典来管理所有卡片检测模型 # key (如 'outer_box') 将成为 API 路径中的 {inference_type} CARD_MODELS_CONFIG: Dict[str, CardModelConfig] = { "pokemon_outer_box": { "pth_path": "Model/pokemon_outer_box.pth", "class_dict": {1: 'outer_box'}, "img_size": {'width': 1280, 'height': 1280}, "confidence": 0.5, "input_channels": 3, }, "pokemon_inner_box": { "pth_path": "Model/pokemon_inner_box.pth", "class_dict": {1: 'inner_box'}, "img_size": {'width': 1280, 'height': 1280}, "confidence": 0.5, "input_channels": 3, }, "pokemon_back_corner_defect": { "pth_path": "Model/pokemon_back_corner_defect.pth", "class_dict": { 1: 'wear', 2: 'wear_and_impact', 3: 'impact', 4: 'damaged', 5: 'wear_and_stain', }, "img_size": {'width': 512, 'height': 512}, "confidence": 0.5, "input_channels": 3, }, "pokemon_front_corner_reflect_defect": { "pth_path": "Model/pokemon_front_corner_reflect_defect.pth", "class_dict": {"1": "impact", "2": "wear_and_impact", "3": "wear"}, "img_size": {'width': 512, 'height': 512}, "confidence": 0.5, "input_channels": 3, }, "pokemon_front_corner_no_reflect_defect": { "pth_path": "Model/pokemon_front_corner_no_reflect_defect.pth", "class_dict": {"1": "wear", "2": "wear_and_impact", "3": "impact", "4": "damaged"}, "img_size": {'width': 512, 'height': 512}, "confidence": 0.5, "input_channels": 3, }, "pokemon_front_face_no_reflect_defect": { "pth_path": "Model/pokemon_front_face_no_reflect_defect.pth", "class_dict": {"1": "scratch", "2": "pit", "3": "stain"}, "img_size": {'width': 512, 'height': 512}, "confidence": 0.5, "input_channels": 3, }, } # 包含, 环形光居中计算, 环形光正反边角缺陷, 同轴光正反表面缺陷 # 里面存储需要用到的模型类型 DEFECT_TYPE: Dict[str, CardModelConfig] = { "pokemon_card_center": { "inner_box": "pokemon_inner_box", "outer_box": "pokemon_outer_box", }, "pokemon_back_corner_defect": { 'model_type': "pokemon_back_corner_defect" }, "pokemon_front_corner_reflect_defect": { "model_type": "pokemon_front_corner_reflect_defect" }, "pokemon_front_corner_no_reflect_defect": { "model_type": "pokemon_front_corner_no_reflect_defect", }, "pokemon_front_face_no_reflect_defect": { "model_type": "pokemon_front_face_no_reflect_defect", } } settings = Settings() print(settings.BASE_PATH) # DefectType = Enum("InferenceType", {name: name for name in settings.DEFECT_TYPE}) # print()