config.py 1.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. from pathlib import Path
  2. from typing import Dict
  3. # 定义一个模型的配置结构
  4. class CardModelConfig:
  5. pth_path: str
  6. class_dict: dict
  7. img_size: dict
  8. confidence: float
  9. input_channels: int
  10. class Settings:
  11. API_prefix: str = "/api/card_inference"
  12. BASE_PATH = Path(__file__).parent.parent.absolute()
  13. # 使用一个字典来管理所有卡片检测模型
  14. # key (如 'outer_box') 将成为 API 路径中的 {inference_type}
  15. CARD_MODELS_CONFIG: Dict[str, CardModelConfig] = {
  16. "outer_box": {
  17. "pth_path": "Model/outer_box.pth",
  18. "class_dict": {1: 'outer_box'},
  19. "img_size": {'width': 1280, 'height': 1280},
  20. "confidence": 0.5,
  21. "input_channels": 3,
  22. },
  23. "inner_box": {
  24. "pth_path": "Model/inner_box.pth",
  25. "class_dict": {1: 'inner_box'},
  26. "img_size": {'width': 1280, 'height': 1280},
  27. "confidence": 0.5,
  28. "input_channels": 3,
  29. },
  30. "back_defect": {
  31. "pth_path": "Model/back_defect.pth",
  32. "class_dict": {
  33. 1: 'wear', 2: 'wear_and_impact', 3: 'impact',
  34. 4: 'damaged', 5: 'wear_and_stain',
  35. },
  36. "img_size": {'width': 512, 'height': 512},
  37. "confidence": 0.5,
  38. "input_channels": 3,
  39. },
  40. "no_reflect_front_defect": {
  41. "pth_path": "Model/no_reflect_front_defect.pth",
  42. "class_dict": {1: 'scratch',
  43. 2: 'pit',
  44. 3: 'stain'},
  45. "img_size": {'width': 512, 'height': 512},
  46. "confidence": 0.5,
  47. "input_channels": 3,
  48. }
  49. }
  50. settings = Settings()
  51. print(settings.BASE_PATH)