|
|
@@ -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:
|