config.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. from pathlib import Path
  2. from typing import Dict, List
  3. from enum import Enum
  4. # 定义一个模型的配置结构
  5. class CardModelConfig:
  6. pth_path: str
  7. class_dict: dict
  8. img_size: dict
  9. confidence: float
  10. input_channels: int
  11. class Settings:
  12. API_Inference_prefix: str = "/api/card_inference"
  13. API_Score_prefix: str = "/api/card_score"
  14. BASE_PATH = Path(__file__).parent.parent.parent.absolute()
  15. TEMP_WORK_DIR = BASE_PATH / "_temp_work"
  16. SCORE_CONFIG_PATH = BASE_PATH / "app/core/scoring_config.json"
  17. # 图片像素与真实图片缩放比例
  18. PIXEL_RESOLUTION = 24.54
  19. CORNER_SIZE_MM = 3.0
  20. # 使用一个字典来管理所有卡片检测模型
  21. # key (如 'outer_box') 将成为 API 路径中的 {inference_type}
  22. CARD_MODELS_CONFIG: Dict[str, CardModelConfig] = {
  23. "pokemon_outer_box": {
  24. "pth_path": "Model/pokemon_outer_box.pth",
  25. "class_dict": {1: 'outer_box'},
  26. "img_size": {'width': 1280, 'height': 1280},
  27. "confidence": 0.5,
  28. "input_channels": 3,
  29. },
  30. "pokemon_inner_box": {
  31. "pth_path": "Model/pokemon_inner_box.pth",
  32. "class_dict": {1: 'inner_box'},
  33. "img_size": {'width': 1280, 'height': 1280},
  34. "confidence": 0.5,
  35. "input_channels": 3,
  36. },
  37. "pokemon_back_corner_defect": {
  38. "pth_path": "Model/pokemon_back_corner_defect.pth",
  39. "class_dict": {
  40. 1: 'wear', 2: 'wear_and_impact', 3: 'impact',
  41. 4: 'damaged', 5: 'wear_and_stain',
  42. },
  43. "img_size": {'width': 512, 'height': 512},
  44. "confidence": 0.5,
  45. "input_channels": 3,
  46. },
  47. "pokemon_front_corner_reflect_defect": {
  48. "pth_path": "Model/pokemon_front_corner_reflect_defect.pth",
  49. "class_dict": {"1": "impact", "2": "wear_and_impact", "3": "wear"},
  50. "img_size": {'width': 512, 'height': 512},
  51. "confidence": 0.5,
  52. "input_channels": 3,
  53. },
  54. "pokemon_front_corner_no_reflect_defect": {
  55. "pth_path": "Model/pokemon_front_corner_no_reflect_defect.pth",
  56. "class_dict": {"1": "wear", "2": "wear_and_impact", "3": "impact", "4": "damaged"},
  57. "img_size": {'width': 512, 'height': 512},
  58. "confidence": 0.5,
  59. "input_channels": 3,
  60. },
  61. "pokemon_front_face_no_reflect_defect": {
  62. "pth_path": "Model/pokemon_front_face_no_reflect_defect.pth",
  63. "class_dict": {"1": "scratch", "2": "pit", "3": "stain"},
  64. "img_size": {'width': 512, 'height': 512},
  65. "confidence": 0.5,
  66. "input_channels": 3,
  67. },
  68. }
  69. # 包含, 环形光居中计算, 环形光正反边角缺陷, 同轴光正反表面缺陷
  70. # 里面存储需要用到的模型类型
  71. DEFECT_TYPE: Dict[str, dict] = {
  72. "pokemon_card_center": {
  73. "inner_box": "pokemon_inner_box",
  74. "outer_box": "pokemon_outer_box",
  75. },
  76. "pokemon_back_corner_defect": {
  77. 'model_type': "pokemon_back_corner_defect"
  78. },
  79. "pokemon_front_corner_reflect_defect": {
  80. "model_type": "pokemon_front_corner_reflect_defect"
  81. },
  82. "pokemon_front_corner_no_reflect_defect": {
  83. "model_type": "pokemon_front_corner_no_reflect_defect",
  84. },
  85. "pokemon_front_face_no_reflect_defect": {
  86. "model_type": "pokemon_front_face_no_reflect_defect",
  87. }
  88. }
  89. SCORE_TYPE: List[str] = ["front_corner_edge", "front_face",
  90. "back_corner_edge", "back_face"]
  91. settings = Settings()
  92. print(settings.BASE_PATH)
  93. # DefectType = Enum("InferenceType", {name: name for name in settings.DEFECT_TYPE})
  94. # print()