Browse Source

增加环光面模型

AnlaAnla 4 tuần trước cách đây
mục cha
commit
f9649555b9

BIN
Model/pokemon_back_face_ring_light_defect.pth


BIN
Model/pokemon_front_face_reflect_ring_light_defect.pth


+ 0 - 1
app/api/card_inference.py

@@ -53,7 +53,6 @@ async def card_model_inference(
              description="环形光居中计算, 环形光正反边角缺陷, 同轴光正反表面缺陷")
 async def card_model_inference(
         defect_type: DefectType,
-        # service: DefectInferenceService = Depends(lambda: defect_service),
         file: UploadFile = File(...),
         is_draw_image: bool = False,
 ):

+ 26 - 2
app/core/config.py

@@ -108,6 +108,23 @@ class Settings:
             "confidence": 0.5,
             "input_channels": 3,
         },
+        # 2个环光的面模型
+        "pokemon_back_face_ring_light_defect": {
+            "pth_path": "Model/pokemon_back_face_ring_light_defect.pth",
+            "class_dict": {"1": "wear", "2": "scratch", "3": "stain", "4": "protrudent",
+                           "5": "damaged", "6": "wear_and_impact", "7": "impact"},
+            "img_size": {'width': 512, 'height': 512},
+            "confidence": 0.5,
+            "input_channels": 3,
+        },
+        "pokemon_front_face_reflect_ring_light_defect": {
+            "pth_path": "Model/pokemon_front_face_reflect_ring_light_defect.pth",
+            "class_dict": {"1": "impact", "2": "wear", "3": "scratch", "4": "wear_and_impact",
+                           "5": "stain", "6": "damaged", "7": "protrudent"},
+            "img_size": {'width': 512, 'height': 512},
+            "confidence": 0.5,
+            "input_channels": 3,
+        }
     }
 
     # 包含, 环形光居中计算, 环形光正反边角缺陷, 同轴光正反表面缺陷
@@ -138,11 +155,18 @@ class Settings:
         },
         "pokemon_front_face_no_reflect_defect": {
             "model_type": "pokemon_front_face_no_reflect_defect",
+        },
+        # 2个环光的
+        "pokemon_back_face_ring_light_defect": {
+            "model_type": "pokemon_back_face_ring_light_defect"
+        },
+        "pokemon_front_face_reflect_ring_light_defect": {
+            "model_type": "pokemon_front_face_reflect_ring_light_defect",
         }
     }
 
-    SCORE_TYPE: List[str] = ["front_corner_edge", "front_face",
-                             "back_corner_edge", "back_face"]
+    SCORE_TYPE: List[str] = ["front_corner_edge", "front_face", "front_face_ring_light",
+                             "back_corner_edge", "back_face", "back_face_ring_light"]
 
 
 settings = Settings()

+ 3 - 1
app/services/defect_service.py

@@ -55,7 +55,9 @@ class DefectInferenceService:
         # 面
         if (inference_type == "pokemon_front_face_no_reflect_defect"
                 or inference_type == "pokemon_front_face_reflect_defect"
-                or inference_type == "pokemon_back_face_defect"):
+                or inference_type == "pokemon_back_face_defect"
+                or inference_type == "pokemon_back_face_ring_light_defect"
+                or inference_type == "pokemon_front_face_reflect_ring_light_defect"):
             # 1. 获取对应的预测器实例
             predictor = get_predictor(inference_type)
 

+ 14 - 2
app/services/score_service.py

@@ -27,7 +27,7 @@ class ScoreService:
         imageHeight, imageWidth = img_bgr.shape[:2]
 
         logger.info("开始进行卡片分数推理")
-        if score_type == 'front_corner_edge' or score_type == 'front_face':
+        if score_type == 'front_corner_edge' or score_type == 'front_face' or score_type == "front_face_ring_light":
             center_data = self.defect_service.defect_inference("pokemon_front_card_center", img_bgr.copy())
         else:
             center_data = self.defect_service.defect_inference("pokemon_back_card_center", img_bgr.copy())
@@ -42,6 +42,12 @@ class ScoreService:
                 defect_data = self.defect_service.defect_inference('pokemon_back_corner_defect', img_bgr.copy())
             elif score_type == 'back_face':
                 defect_data = self.defect_service.defect_inference('pokemon_back_face_defect', img_bgr.copy())
+            elif score_type == 'front_face_ring_light':
+                defect_data = self.defect_service.defect_inference('pokemon_front_face_reflect_ring_light_defect',
+                                                                   img_bgr.copy())
+            elif score_type == 'back_face_ring_light':
+                defect_data = self.defect_service.defect_inference('pokemon_back_face_ring_light_defect',
+                                                                   img_bgr.copy())
             else:
                 return {}
         else:
@@ -55,11 +61,17 @@ class ScoreService:
                 defect_data = self.defect_service.defect_inference('pokemon_back_corner_defect', img_bgr.copy())
             elif score_type == 'back_face':
                 defect_data = self.defect_service.defect_inference('pokemon_back_face_defect', img_bgr.copy())
+            elif score_type == 'front_face_ring_light':
+                defect_data = self.defect_service.defect_inference('pokemon_front_face_reflect_ring_light_defect',
+                                                                   img_bgr.copy())
+            elif score_type == 'back_face_ring_light':
+                defect_data = self.defect_service.defect_inference('pokemon_back_face_ring_light_defect',
+                                                                   img_bgr.copy())
             else:
                 return {}
 
         logger.info("模型推理结束, 开始计算分数")
-        if score_type == 'front_corner_edge' or score_type == 'front_face':
+        if score_type == 'front_corner_edge' or score_type == 'front_face' or score_type=="front_face_ring_light":
             card_aspect = "front"
         else:
             card_aspect = "back"