config.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  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. '''
  23. face: "1": "wear", "2": "scratch", "3": "stain",
  24. "4": "scuff", "5": "impact", "6": "damaged",
  25. "7": "wear_and_impact"
  26. "8": "stain", "9": "pit"
  27. corner: "1": "wear", "2": "wear_and_impact", "3": "damaged",
  28. "4": "impact", "5": "wear_and_stain"
  29. '''
  30. CARD_MODELS_CONFIG: Dict[str, CardModelConfig] = {
  31. "outer_box": {
  32. "pth_path": "Model/outer_box.pth",
  33. "class_dict": {1: 'outer_box'},
  34. "img_size": {'width': 1280, 'height': 1280},
  35. "confidence": 0.5,
  36. "input_channels": 3,
  37. },
  38. "pokemon_front_inner_box": {
  39. "pth_path": "Model/pokemon_front_inner_box.pth",
  40. "class_dict": {1: 'inner_box'},
  41. "img_size": {'width': 1280, 'height': 1280},
  42. "confidence": 0.5,
  43. "input_channels": 3,
  44. },
  45. "pokemon_back_inner_box": {
  46. "pth_path": "Model/pokemon_back_inner_box.pth",
  47. "class_dict": {1: 'inner_box'},
  48. "img_size": {'width': 1280, 'height': 1280},
  49. "confidence": 0.5,
  50. "input_channels": 3,
  51. },
  52. "pokemon_back_corner_defect": {
  53. "pth_path": "Model/pokemon_back_corner_defect.pth",
  54. "class_dict": {"1": "wear", "2": "wear_and_impact", "3": "damaged",
  55. "4": "impact", "5": "wear_and_stain"},
  56. "img_size": {'width': 512, 'height': 512},
  57. "confidence": 0.5,
  58. "input_channels": 3,
  59. },
  60. "pokemon_back_face_defect": {
  61. "pth_path": "Model/pokemon_back_face_defect.pth",
  62. "class_dict": {"1": "wear", "2": "scratch", "3": "stain",
  63. "4": "scuff", "5": "impact", "6": "damaged",
  64. "7": "wear_and_impact"},
  65. "img_size": {'width': 512, 'height': 512},
  66. "confidence": 0.5,
  67. "input_channels": 3,
  68. },
  69. "pokemon_front_face_reflect_defect": {
  70. "pth_path": "Model/pokemon_front_face_reflect_defect.pth",
  71. "class_dict": {"1": "scratch", "2": "stain", "3": "wear",
  72. "4": "impact", "5": "stain_and_scratch"},
  73. "img_size": {'width': 512, 'height': 512},
  74. "confidence": 0.5,
  75. "input_channels": 3,
  76. },
  77. "pokemon_front_corner_reflect_defect": {
  78. "pth_path": "Model/pokemon_front_corner_reflect_defect.pth",
  79. "class_dict": {"1": "impact", "2": "wear_and_impact", "3": "wear"},
  80. "img_size": {'width': 512, 'height': 512},
  81. "confidence": 0.5,
  82. "input_channels": 3,
  83. },
  84. "pokemon_front_corner_no_reflect_defect": {
  85. "pth_path": "Model/pokemon_front_corner_no_reflect_defect.pth",
  86. "class_dict": {"1": "wear", "2": "wear_and_impact", "3": "impact",
  87. "4": "damaged", "5": "stain"},
  88. "img_size": {'width': 512, 'height': 512},
  89. "confidence": 0.5,
  90. "input_channels": 3,
  91. },
  92. "pokemon_front_face_no_reflect_defect": {
  93. "pth_path": "Model/pokemon_front_face_no_reflect_defect.pth",
  94. "class_dict": {"1": "wear", "2": "scratch", "3": "damaged",
  95. "4": "stain", "5": "impact", "6": "pit"},
  96. "img_size": {'width': 512, 'height': 512},
  97. "confidence": 0.5,
  98. "input_channels": 3,
  99. },
  100. }
  101. # 包含, 环形光居中计算, 环形光正反边角缺陷, 同轴光正反表面缺陷
  102. # 里面存储需要用到的模型类型
  103. DEFECT_TYPE: Dict[str, dict] = {
  104. "pokemon_front_card_center": {
  105. "inner_box": "pokemon_front_inner_box",
  106. "outer_box": "outer_box",
  107. },
  108. "pokemon_back_card_center": {
  109. "inner_box": "pokemon_back_inner_box",
  110. "outer_box": "outer_box",
  111. },
  112. "pokemon_back_face_defect": {
  113. 'model_type': "pokemon_back_face_defect"
  114. },
  115. "pokemon_back_corner_defect": {
  116. 'model_type': "pokemon_back_corner_defect"
  117. },
  118. "pokemon_front_corner_reflect_defect": {
  119. "model_type": "pokemon_front_corner_reflect_defect"
  120. },
  121. "pokemon_front_face_reflect_defect": {
  122. "model_type": "pokemon_front_face_reflect_defect"
  123. },
  124. "pokemon_front_corner_no_reflect_defect": {
  125. "model_type": "pokemon_front_corner_no_reflect_defect",
  126. },
  127. "pokemon_front_face_no_reflect_defect": {
  128. "model_type": "pokemon_front_face_no_reflect_defect",
  129. }
  130. }
  131. SCORE_TYPE: List[str] = ["front_corner_edge", "front_face",
  132. "back_corner_edge", "back_face"]
  133. settings = Settings()
  134. print(f"项目根目录: {settings.BASE_PATH}")
  135. print(f"数据存储目录: {settings.DATA_DIR}")
  136. # DefectType = Enum("InferenceType", {name: name for name in settings.DEFECT_TYPE})
  137. # print()