config.py 1.7 KB

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