config.py 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161
  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. API_Config_prefix: str = "/api/config"
  15. BASE_PATH = Path(__file__).parent.parent.parent.absolute()
  16. TEMP_WORK_DIR = BASE_PATH / "_temp_work"
  17. SCORE_CONFIG_PATH = BASE_PATH / "app/core/scoring_config.json"
  18. SCORE_BACKUP_PATH = BASE_PATH / "score_config_backup"
  19. # 图片像素与真实图片缩放比例
  20. PIXEL_RESOLUTION = 24.54
  21. CORNER_SIZE_MM = 2.0
  22. EDGE_SIZE_MM = 1
  23. # 使用一个字典来管理所有卡片检测模型
  24. # key (如 'outer_box') 将成为 API 路径中的 {inference_type}
  25. '''
  26. face: "1": "wear", "2": "scratch", "3": "stain",
  27. "4": "scuff", "5": "impact", "6": "damaged",
  28. "7": "wear_and_impact", "8": "pit"
  29. corner: "1": "wear", "2": "wear_and_impact", "3": "damaged",
  30. "4": "impact", "5": "wear_and_stain"
  31. 简化版本: 边和角: wear, impact, damaged
  32. 面: wear, scratch, pit, impact, stain
  33. '''
  34. CARD_MODELS_CONFIG: Dict[str, CardModelConfig] = {
  35. "outer_box": {
  36. "pth_path": "Model/outer_box.pth",
  37. "class_dict": {1: 'outer_box'},
  38. "img_size": {'width': 1280, 'height': 1280},
  39. "confidence": 0.5,
  40. "input_channels": 3,
  41. },
  42. "pokemon_front_inner_box": {
  43. "pth_path": "Model/pokemon_front_inner_box.pth",
  44. "class_dict": {1: 'inner_box'},
  45. "img_size": {'width': 1280, 'height': 1280},
  46. "confidence": 0.5,
  47. "input_channels": 3,
  48. },
  49. "pokemon_back_inner_box": {
  50. "pth_path": "Model/pokemon_back_inner_box.pth",
  51. "class_dict": {1: 'inner_box'},
  52. "img_size": {'width': 1280, 'height': 1280},
  53. "confidence": 0.5,
  54. "input_channels": 3,
  55. },
  56. # 同轴光模型
  57. "pokemon_back_face_coaxial_light_defect": {
  58. "pth_path": "Model/pokemon_back_face_coaxial_light_defect.pth",
  59. "class_dict": {"1": "wear", "2": "scratch", "3": "stain",
  60. "4": "scuff", "5": "impact", "6": "damaged",
  61. "7": "wear_and_impact"},
  62. "img_size": {'width': 512, 'height': 512},
  63. "confidence": 0.5,
  64. "input_channels": 3,
  65. },
  66. "pokemon_front_face_reflect_coaxial_light_defect": {
  67. "pth_path": "Model/pokemon_front_face_reflect_coaxial_light_defect.pth",
  68. "class_dict": {"1": "stain", "2": "scratch", "3": "impact", "4": "wear", "5": "wear"},
  69. "img_size": {'width': 512, 'height': 512},
  70. "confidence": 0.5,
  71. "input_channels": 3,
  72. },
  73. "pokemon_front_face_no_reflect_coaxial_light_defect": {
  74. "pth_path": "Model/pokemon_front_face_no_reflect_coaxial_light_defect.pth",
  75. "class_dict": {"1": "scratch", "2": "wear", "3": "stain", "4": "damaged", "5": "impact", "6": "protrudent",
  76. "7": "wear_and_impact"},
  77. "img_size": {'width': 512, 'height': 512},
  78. "confidence": 0.5,
  79. "input_channels": 3,
  80. },
  81. # 环光模型
  82. "pokemon_back_face_ring_light_defect": {
  83. "pth_path": "Model/pokemon_back_face_ring_light_defect.pth",
  84. "class_dict": {"1": "wear", "2": "scratch", "3": "stain", "4": "impact",
  85. "5": "damaged", "6": "wear_and_impact", "7": "impact"},
  86. "img_size": {'width': 512, 'height': 512},
  87. "confidence": 0.5,
  88. "input_channels": 3,
  89. },
  90. "pokemon_front_face_reflect_ring_light_defect": {
  91. "pth_path": "Model/pokemon_front_face_reflect_ring_light_defect.pth",
  92. "class_dict": {"1": "impact", "2": "wear", "3": "scratch", "4": "wear_and_impact",
  93. "5": "stain", "6": "damaged", "7": "impact"},
  94. "img_size": {'width': 512, 'height': 512},
  95. "confidence": 0.5,
  96. "input_channels": 3,
  97. },
  98. "pokemon_front_face_no_reflect_ring_light_defect": {
  99. "pth_path": "Model/pokemon_front_face_no_reflect_ring_light_defect.pth",
  100. "class_dict": {"1": "damaged", "2": "scratch", "3": "wear", "4": "stain", "5": "impact",
  101. "6": "wear_and_stain", "7": "wear_and_impact"}
  102. ,
  103. "img_size": {'width': 512, 'height': 512},
  104. "confidence": 0.5,
  105. "input_channels": 3,
  106. }
  107. }
  108. # 包含, 环形光居中计算, 环形光正反边角缺陷, 同轴光正反表面缺陷
  109. # 里面存储需要用到的模型类型
  110. DEFECT_TYPE: Dict[str, dict] = {
  111. "pokemon_front_card_center": {
  112. "inner_box": "pokemon_front_inner_box",
  113. "outer_box": "outer_box",
  114. },
  115. "pokemon_back_card_center": {
  116. "inner_box": "pokemon_back_inner_box",
  117. "outer_box": "outer_box",
  118. },
  119. # 同轴
  120. "pokemon_back_face_coaxial_light_defect": {
  121. 'model_type': "pokemon_back_face_coaxial_light_defect"
  122. },
  123. "pokemon_front_face_reflect_coaxial_light_defect": {
  124. "model_type": "pokemon_front_face_reflect_coaxial_light_defect"
  125. },
  126. "pokemon_front_face_no_reflect_coaxial_light_defect": {
  127. "model_type": "pokemon_front_face_no_reflect_coaxial_light_defect",
  128. },
  129. # 环光
  130. "pokemon_back_face_ring_light_defect": {
  131. "model_type": "pokemon_back_face_ring_light_defect"
  132. },
  133. "pokemon_front_face_reflect_ring_light_defect": {
  134. "model_type": "pokemon_front_face_reflect_ring_light_defect"
  135. },
  136. "pokemon_front_face_no_reflect_ring_light_defect": {
  137. "model_type": "pokemon_front_face_no_reflect_ring_light_defect"
  138. }
  139. }
  140. SCORE_TYPE: List[str] = ["front_coaxial", "front_ring",
  141. "back_coaxial", "back_ring"]
  142. settings = Settings()
  143. print(f"项目根目录: {settings.BASE_PATH}")
  144. # DefectType = Enum("InferenceType", {name: name for name in settings.DEFECT_TYPE})
  145. # print()