Explorar o código

多余外框过滤

AnlaAnla hai 1 mes
pai
achega
b09306c40c

+ 1 - 1
Test/model_test01.py

@@ -55,7 +55,7 @@ def predict_single_image(config_params: dict,
 
 
 if __name__ == '__main__':
-    big_img_path = r"C:\Code\ML\Image\Card\_250915_many_capture_img\_250915_1743_reflect_nature_defrct\1_front_coaxial_1_0.jpg"
+    big_img_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\_temp_work\rectify_center_img.jpg"
 
     config = settings.CARD_MODELS_CONFIG
     # predict_single_image(config['pokemon_front_inner_box'],

+ 32 - 0
app/services/card_rectify_and_center.py

@@ -15,6 +15,36 @@ class CardRectifyAndCenter:
         self.inference_type = "outer_box"
         self.outer_box_model = get_predictor(self.inference_type)
 
+    def _keep_max_prob_shape(self, json_data: dict) -> dict:
+        """
+        辅助函数:过滤掉低置信度的框,只保留 probability 最大的那一个。
+        """
+        if not json_data or 'shapes' not in json_data or not json_data['shapes']:
+            return json_data
+
+        shapes = json_data['shapes']
+        # 如果只有一个,直接返回
+        if len(shapes) == 1:
+            return json_data
+
+        logger.info(f"检测到 {len(shapes)} 个外框,开始筛选最大置信度目标...")
+
+        # 使用 max 函数找出 probability 最大的 shape
+        # x.get('probability', 0) 确保如果字段不存在默认为0
+        best_shape = max(shapes, key=lambda x: x.get('probability', 0))
+
+        logger.info(f"筛选结果: 保留置信度为 {best_shape.get('probability')} 的目标,"
+                    f"丢弃了其他 {len(shapes) - 1} 个目标。")
+
+        # 重新赋值给 shapes,保持列表结构
+        json_data['shapes'] = [best_shape]
+
+        # 同时更新 num 数量,保持数据一致性
+        if 'num' in json_data:
+            json_data['num'] = 1
+
+        return json_data
+
     def rectify_and_center(self, img_bgr: np.ndarray) -> np.ndarray:
         # 1. 设置处理参数
         params = FryCardProcessParams(
@@ -28,6 +58,8 @@ class CardRectifyAndCenter:
         processor = FryCardProcessor()
 
         seg_json = self.outer_box_model.predict_from_image(img_bgr)
+        # 过滤噪点,只留最大置信度的外框
+        seg_json = self._keep_max_prob_shape(seg_json)
 
         # 4. 执行处理
         final_image = processor.process_image_with_json(img_bgr, seg_json, params)

+ 27 - 1
app/services/defect_service.py

@@ -18,6 +18,25 @@ logger = get_logger(__name__)
 
 
 class DefectInferenceService:
+    def _filter_max_prob_shape(self, json_data: dict, tag: str = "unknown") -> dict:
+        """
+        通用过滤函数:只保留 shapes 中 probability 最大的那一个。
+        """
+        if not json_data or 'shapes' not in json_data or not json_data['shapes']:
+            return json_data
+
+        shapes = json_data['shapes']
+        if len(shapes) <= 1:
+            return json_data
+
+        # 按 probability 降序排列,取第一个
+        best_shape = max(shapes, key=lambda x: x.get('probability', 0))
+
+        logger.info(f"[{tag}] 过滤多余框: 原有 {len(shapes)} 个, 保留最大置信度 {best_shape.get('probability'):.4f}")
+
+        json_data['shapes'] = [best_shape]
+        return json_data
+
     def defect_inference(self, inference_type: str, img_bgr: np.ndarray,
                          is_draw_image=True) -> dict:
         """
@@ -136,7 +155,10 @@ class DefectInferenceService:
             predictor_outer = get_predictor("outer_box")
             outer_result = predictor_outer.predict_from_image(img_bgr)
 
-            # 2. 执行过滤
+            # 过滤外框,只留一个
+            outer_result = self._filter_max_prob_shape(outer_result, tag="Corner流程-外框")
+
+            # 2. 执行过滤, 去掉外框外的缺陷
             json_data = outside_filter.execute(json_data, outer_result)
 
             processor = DefectProcessor(pixel_resolution=settings.PIXEL_RESOLUTION)
@@ -170,6 +192,10 @@ class DefectInferenceService:
             inner_result = predictor_inner.predict_from_image(img_bgr)
             outer_result = predictor_outer.predict_from_image(img_bgr)
 
+            # 过滤内框和外框,只留最大置信度的框
+            inner_result = self._filter_max_prob_shape(inner_result, tag="Center流程-内框")
+            outer_result = self._filter_max_prob_shape(outer_result, tag="Center流程-外框")
+
             # temp_inner_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-inner_result.json'
             # temp_outer_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-outer_result.json'
             # with open(temp_inner_json_path, 'w', encoding='utf-8') as f: