CheckCardBoxAndDefectServer/app/core/config.py

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_Inference_prefix: str = "/api/card_inference"
    API_Score_prefix: str = "/api/card_score"
    API_Config_prefix: str = "/api/config"

    BASE_PATH = Path(__file__).parent.parent.parent.absolute()

    TEMP_WORK_DIR = BASE_PATH / "_temp_work"
    SCORE_CONFIG_PATH = BASE_PATH / "app/core/scoring_config.json"

    # 图片像素与真实图片缩放比例
    PIXEL_RESOLUTION = 24.54
    CORNER_SIZE_MM = 3.0

    # 使用一个字典来管理所有卡片检测模型
    # key (如 'outer_box') 将成为 API 路径中的 {inference_type}
    '''
    face: "1": "wear", "2": "scratch", "3": "stain",
        "4": "scuff", "5": "impact", "6": "damaged",
        "7": "wear_and_impact"
        "8": "stain", "9": "pit"
    
    corner: "1": "wear", "2": "wear_and_impact", "3": "damaged",
                           "4": "impact", "5": "wear_and_stain"
    '''
    CARD_MODELS_CONFIG: Dict[str, CardModelConfig] = {
        "outer_box": {
            "pth_path": "Model/outer_box.pth",
            "class_dict": {1: 'outer_box'},
            "img_size": {'width': 1280, 'height': 1280},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_front_inner_box": {
            "pth_path": "Model/pokemon_front_inner_box.pth",
            "class_dict": {1: 'inner_box'},
            "img_size": {'width': 1280, 'height': 1280},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_back_inner_box": {
            "pth_path": "Model/pokemon_back_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": "damaged",
                           "4": "impact", "5": "wear_and_stain"},
            "img_size": {'width': 512, 'height': 512},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_back_face_defect": {
            "pth_path": "Model/pokemon_back_face_defect.pth",
            "class_dict": {"1": "wear", "2": "scratch", "3": "stain",
                           "4": "scuff", "5": "impact", "6": "damaged",
                           "7": "wear_and_impact"},
            "img_size": {'width': 512, 'height': 512},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_front_face_reflect_defect": {
            "pth_path": "Model/pokemon_front_face_reflect_defect.pth",
            "class_dict": {"1": "stain", "2": "scratch", "3": "impact", "4": "wear"},
            "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", "5": "stain"},
            "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": "wear", "3": "stain", "4": "damaged",
                           "5": "impact", "6": "wear_and_impact",
                           "7": "chip", "8": "protrudent", "9": "wear_and_stain"},
            "img_size": {'width': 512, 'height': 512},
            "confidence": 0.5,
            "input_channels": 3,
        },
    }

    # 包含, 环形光居中计算, 环形光正反边角缺陷, 同轴光正反表面缺陷
    # 里面存储需要用到的模型类型
    DEFECT_TYPE: Dict[str, dict] = {
        "pokemon_front_card_center": {
            "inner_box": "pokemon_front_inner_box",
            "outer_box": "outer_box",
        },
        "pokemon_back_card_center": {
            "inner_box": "pokemon_back_inner_box",
            "outer_box": "outer_box",
        },
        "pokemon_back_face_defect": {
            'model_type': "pokemon_back_face_defect"
        },
        "pokemon_back_corner_defect": {
            'model_type': "pokemon_back_corner_defect"
        },
        "pokemon_front_corner_reflect_defect": {
            "model_type": "pokemon_front_corner_reflect_defect"
        },
        "pokemon_front_face_reflect_defect": {
            "model_type": "pokemon_front_face_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",
        }
    }

    SCORE_TYPE: List[str] = ["front_corner_edge", "front_face",
                             "back_corner_edge", "back_face"]


settings = Settings()
print(f"项目根目录: {settings.BASE_PATH}")

# DefectType = Enum("InferenceType", {name: name for name in settings.DEFECT_TYPE})
# print()