模型更新

Test/model_test01.py

@@ -3,27 +3,23 @@ from pathlib import Path
 from app.utils.fry_bisenetv2_predictor_V04_250819 import FryBisenetV2Predictor
 
 BASE_PATH = Path(__file__).parent.parent.absolute()
-# model_inner_box_path = BASE_PATH / "Model/pokemon_inner_box.pth"
-# model_outer_box_path = BASE_PATH / "Model/pokemon_outer_box.pth"
 
 
-def _test_pokemon_inner_box(model_path=None,
-                            real_seg_class_dict = {1: 'inner_box'},
-imgSize_train_dict = {'width': 1280, 'height': 1280}
-                            ):
+def predict_single_image(config_params: dict,
+                         img_path: str,
+                         output_dir: str):
     # 配置参数
-    input_channels = 3
-
-    real_seg_class_dict = {1: 'inner_box'}
+    model_path = BASE_PATH / config_params["pth_path"]
+    real_seg_class_dict = config_params['class_dict']
+    imgSize_train_dict = config_params['img_size']
+    confidence = config_params['confidence']
+    input_channels = config_params['input_channels']
 
     # 为不同类别设置不同颜色（可选）
     label_colors_dict = {
         'outer_box': (255, 0, 0),
     }
 
-    imgSize_train_dict = {'width': 1280, 'height': 1280}
-    confidence = 0.5
-
     # 创建预测器
     predictor = FryBisenetV2Predictor(
         pth_path=str(model_path),
@@ -36,8 +32,8 @@ imgSize_train_dict = {'width': 1280, 'height': 1280}
 
     # 单张图片预测
     print("=== 单张图片预测 ===")
-    now_img_path = r"C:\Code\ML\Project\StitchImageServer\temp\output\0_half_importance_add_weight\final_stitched_image.jpg"
-    answer_json_dir_str = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\inner"
+    now_img_path = img_path
+    answer_json_dir_str = output_dir
 
     result = predictor.predict_single_image(
         img_path=now_img_path,
@@ -63,8 +59,8 @@ if __name__ == '__main__':
             "confidence": 0.5,
             "input_channels": 3,
         },
-        "pokemon_back_defect": {
-            "pth_path": "Model/pokemon_back_defect.pth",
+        "pokemon_back_corner_defect": {
+            "pth_path": "Model/pokemon_back_corner_defect.pth",
             "class_dict": {
                 1: 'wear', 2: 'wear_and_impact', 3: 'impact',
                 4: 'damaged', 5: 'wear_and_stain',
@@ -73,22 +69,43 @@ if __name__ == '__main__':
             "confidence": 0.5,
             "input_channels": 3,
         },
-        "pokemon_reflect_front_defect": {
-            "pth_path": "Model/pokemon_reflect_front_defect.pth",
-            "class_dict": {"1": "impact",
-                           "2": "wear_and_impact",
-                           "3": "wear"},
+        "pokemon_front_corner_reflect_defect": {
+            "pth_path": "Model/pokemon_front_corner_reflect_defect.pth",
+            "class_dict": {"1": "impact", "2": "wear_and_impact", "3": "wear"},
             "img_size": {'width': 512, 'height': 512},
             "confidence": 0.5,
             "input_channels": 3,
         },
-        "pokemon_no_reflect_front_defect": {
-            "pth_path": "Model/pokemon_reflect_front_defect.pth",
-            "class_dict": {1: 'scratch',
-                           2: 'pit',
-                           3: 'stain'},
+        "pokemon_front_corner_no_reflect_defect": {
+            "pth_path": "Model/pokemon_front_corner_no_reflect_defect.pth",
+            "class_dict": {"1": "wear", "2": "wear_and_impact", "3": "impact", "4": "damaged"},
             "img_size": {'width': 512, 'height': 512},
             "confidence": 0.5,
             "input_channels": 3,
-        }}
-    _test_pokemon_inner_box()
+        },
+        "pokemon_front_face_no_reflect_defect": {
+            "pth_path": "Model/pokemon_front_face_no_reflect_defect.pth",
+            "class_dict": {"1": "scratch", "2": "pit", "3": "stain"},
+            "img_size": {'width': 512, 'height': 512},
+            "confidence": 0.5,
+            "input_channels": 3,
+        },
+    }
+    # predict_single_image(config['pokemon_inner_box'],
+    #                      img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\img.png",
+    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\inner")
+
+    # predict_single_image(config['pokemon_outer_box'],
+    #                      img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\img.png",
+    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\outer")
+
+    # predict_single_image(config['pokemon_back_corner_defect'],
+    #                      img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\img_2.png",
+    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\back_corner")
+
+    # predict_single_image(config['pokemon_front_face_no_reflect_defect'],
+    #                      img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\250805_pokemon_0001_grid_r5_c2.jpg",
+    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\face_no_reflect")
+
+
+

Test/test01.py

@@ -3,4 +3,4 @@ import json
 with open("../temp/test_return.json", 'r') as f:
     data = json.load(f)
 
-print(data)
+print(data)

app/core/config.py

@@ -32,8 +32,8 @@ class Settings:
             "confidence": 0.5,
             "input_channels": 3,
         },
-        "pokemon_back_defect": {
-            "pth_path": "Model/pokemon_back_defect.pth",
+        "pokemon_back_corner_defect": {
+            "pth_path": "Model/pokemon_back_corner_defect.pth",
             "class_dict": {
                 1: 'wear', 2: 'wear_and_impact', 3: 'impact',
                 4: 'damaged', 5: 'wear_and_stain',
@@ -42,24 +42,27 @@ class Settings:
             "confidence": 0.5,
             "input_channels": 3,
         },
-        "pokemon_reflect_front_defect": {
-            "pth_path": "Model/pokemon_reflect_front_defect.pth",
-            "class_dict": {"1": "impact",
-                           "2": "wear_and_impact",
-                           "3": "wear"},
+        "pokemon_front_corner_reflect_defect": {
+            "pth_path": "Model/pokemon_front_corner_reflect_defect.pth",
+            "class_dict": {"1": "impact", "2": "wear_and_impact", "3": "wear"},
             "img_size": {'width': 512, 'height': 512},
             "confidence": 0.5,
             "input_channels": 3,
         },
-        "pokemon_no_reflect_front_defect": {
-            "pth_path": "Model/pokemon_reflect_front_defect.pth",
-            "class_dict": {1: 'scratch',
-                           2: 'pit',
-                           3: 'stain'},
+        "pokemon_front_corner_no_reflect_defect": {
+            "pth_path": "Model/pokemon_front_corner_no_reflect_defect.pth",
+            "class_dict": {"1": "wear", "2": "wear_and_impact", "3": "impact", "4": "damaged"},
             "img_size": {'width': 512, 'height': 512},
             "confidence": 0.5,
             "input_channels": 3,
-        }
+        },
+        "pokemon_front_face_no_reflect_defect": {
+            "pth_path": "Model/pokemon_front_face_no_reflect_defect.pth",
+            "class_dict": {"1": "scratch", "2": "pit", "3": "stain"},
+            "img_size": {'width': 512, 'height': 512},
+            "confidence": 0.5,
+            "input_channels": 3,
+        },
     }