Преглед изворни кода

增加外框外的缺陷过滤

AnlaAnla пре 1 недеља
родитељ
комит
89631f2117

BIN
Model/pokemon_front_face_no_reflect_defect.pth


BIN
Model/pokemon_front_face_reflect_defect.pth


+ 5 - 4
app/core/config.py

@@ -81,28 +81,29 @@ class Settings:
         },
         "pokemon_front_face_reflect_defect": {
             "pth_path": "Model/pokemon_front_face_reflect_defect.pth",
-            "class_dict": {"1": "stain","2": "scratch","3": "impact","4": "wear","5": "wear"},
+            "class_dict": {"1": "stain", "2": "scratch", "3": "impact", "4": "wear", "5": "wear"},
             "img_size": {'width': 512, 'height': 512},
             "confidence": 0.5,
             "input_channels": 3,
         },
         "pokemon_front_corner_reflect_defect": {
             "pth_path": "Model/pokemon_front_corner_reflect_defect.pth",
-            "class_dict": {"1": "impact","2": "wear_and_impact","3": "wear"},
+            "class_dict": {"1": "impact", "2": "wear_and_impact", "3": "wear"},
             "img_size": {'width': 512, 'height': 512},
             "confidence": 0.5,
             "input_channels": 3,
         },
         "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","5": "stain"},
+            "class_dict": {"1": "wear", "2": "wear_and_impact", "3": "impact", "4": "damaged", "5": "stain"},
             "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": "wear", "3": "stain", "4": "damaged", "5": "impact"},
+            "class_dict": {"1": "scratch", "2": "wear", "3": "stain", "4": "damaged", "5": "impact", "6": "protrudent",
+                           "7": "wear_and_impact"},
             "img_size": {'width': 512, 'height': 512},
             "confidence": 0.5,
             "input_channels": 3,

+ 0 - 15
app/core/缺陷英文与中文对应.json

@@ -1,15 +0,0 @@
-{
-  "face": "面",
-  "wear": "磨损",
-  "scratch": "划痕",
-  "stain": "污渍",
-  "scuff": "磨痕",
-  "impact": "撞击痕迹",
-  "damaged": "损坏",
-  "wear_and_impact": "磨损和撞击痕迹",
-  "pit": "凹坑",
-  "corner": "角",
-  "wear_and_stain": "磨损和污渍",
-  "inner_box": "内框",
-  "outer_box": "外框"
-}

+ 23 - 5
app/services/defect_service.py

@@ -4,10 +4,11 @@ from ..core.model_loader import get_predictor
 from app.utils.defect_inference.CardDefectAggregator import CardDefectAggregator
 from app.utils.defect_inference.arean_anylize_draw import DefectProcessor, DrawingParams
 from app.utils.defect_inference.AnalyzeCenter import (
-    analyze_centering_rotated,analyze_centering_rect ,formate_center_data)
+    analyze_centering_rotated, analyze_centering_rect, formate_center_data)
 from app.utils.defect_inference.DrawCenterInfo import draw_boxes_and_center_info
 from app.utils.defect_inference.ClassifyEdgeCorner import ClassifyEdgeCorner
 from app.utils.json_data_formate import formate_face_data, formate_add_edit_type
+from app.utils.defect_inference.FilterOutsideDefects import FilterOutsideDefects
 from app.utils.simplify_points import SimplifyPoints
 from app.core.config import settings
 from app.core.logger import get_logger
@@ -30,6 +31,7 @@ class DefectInferenceService:
             一个包含推理结果的字典。
         """
         simplifyPoints = SimplifyPoints()
+        outside_filter = FilterOutsideDefects(expansion_pixel=30)
 
         # 面
         if (inference_type == "pokemon_front_face_no_reflect_defect"
@@ -62,6 +64,17 @@ class DefectInferenceService:
                 new_num1 = len(simplify_points)
                 logger.info(f"num: {num1}, new_num1: {new_num1}")
 
+            logger.info("开始执行外框过滤...")
+            try:
+                # 1. 因为面原本没有外框推理,这里需要专门加载并推理一次
+                predictor_outer = get_predictor("outer_box")
+                outer_result = predictor_outer.predict_from_image(img_bgr)
+
+                # 2. 执行过滤
+                json_data = outside_filter.execute(json_data, outer_result)
+            except Exception as e:
+                logger.error(f"面流程外框过滤失败: {e}")
+
             merge_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-merge.json'
             with open(merge_json_path, 'w', encoding='utf-8') as f:
                 json.dump(json_data, f, ensure_ascii=False, indent=4)
@@ -104,6 +117,7 @@ class DefectInferenceService:
                 mode='edge'
             )
 
+            # 简化点数
             for shapes in json_data["shapes"]:
                 points = shapes["points"]
                 num1 = len(points)
@@ -117,6 +131,14 @@ class DefectInferenceService:
             #     json.dump(json_data, f, ensure_ascii=False, indent=4)
             # logger.info(f"合并结束")
 
+            logger.info("开始执行外框过滤...")
+            # 外框推理
+            predictor_outer = get_predictor("outer_box")
+            outer_result = predictor_outer.predict_from_image(img_bgr)
+
+            # 2. 执行过滤
+            json_data = outside_filter.execute(json_data, outer_result)
+
             processor = DefectProcessor(pixel_resolution=settings.PIXEL_RESOLUTION)
             area_json_path = settings.TEMP_WORK_DIR / f'{inference_type}-corner_result.json'
             if is_draw_image:
@@ -129,10 +151,6 @@ class DefectInferenceService:
                 area_json: dict = processor.analyze_from_json(json_data)
             logger.info("边角缺陷面积计算结束")
 
-            # 推理外框
-            predictor_outer = get_predictor("outer_box")
-            outer_result = predictor_outer.predict_from_image(img_bgr)
-
             classifier = ClassifyEdgeCorner(settings.PIXEL_RESOLUTION, settings.CORNER_SIZE_MM)
             edge_corner_data = classifier.classify_defects_location(area_json, outer_result)
 

+ 83 - 0
app/utils/defect_inference/FilterOutsideDefects.py

@@ -0,0 +1,83 @@
+import cv2
+import numpy as np
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class FilterOutsideDefects:
+    def __init__(self, expansion_pixel: int = 50):
+        """
+        Args:
+            expansion_pixel (int): 外框向外扩大的像素距离 (默认为 50 像素)
+        """
+        self.expansion_pixel = expansion_pixel
+
+    def _get_expanded_polygon(self, outer_box_data):
+        """
+        获取扩大后的外框多边形
+        """
+        if not outer_box_data or 'shapes' not in outer_box_data or not outer_box_data['shapes']:
+            raise ValueError("外框数据为空或格式不正确")
+
+        # 1. 获取原始外框轮廓
+        points = outer_box_data['shapes'][0]['points']
+        contour = np.array(points, dtype=np.int32)
+
+        # 2. 计算最小外接矩形 (Center, (Width, Height), Angle)
+        rect = cv2.minAreaRect(contour)
+        (center, (width, height), angle) = rect
+
+        # 3. 扩大尺寸 (长宽各增加 2 * expansion_pixel)
+        new_size = (width + 2 * self.expansion_pixel, height + 2 * self.expansion_pixel)
+        expanded_rect = (center, new_size, angle)
+
+        # 4. 获取新的矩形 4 个角点
+        box_points = cv2.boxPoints(expanded_rect)
+        return box_points
+
+    def execute(self, defect_json: dict, outer_box_data: dict) -> dict:
+        """
+        执行过滤逻辑
+
+        Args:
+            defect_json: 包含所有缺陷形状的 JSON 数据 (在 Aggregator 之后的数据)
+            outer_box_data: 外框模型的推理结果
+        Returns:
+            过滤后的 defect_json
+        """
+        if not defect_json or 'shapes' not in defect_json:
+            return defect_json
+
+        try:
+            # 获取扩大后的外框多边形
+            expanded_poly = self._get_expanded_polygon(outer_box_data)
+        except Exception as e:
+            logger.error(f"外框数据处理失败,跳过过滤: {e}")
+            return defect_json
+
+        valid_shapes = []
+        removed_count = 0
+
+        for shape in defect_json['shapes']:
+            points = shape.get("points", [])
+            if not points:
+                continue
+
+            # 计算当前缺陷形状的几何中心
+            pts_array = np.array(points, dtype=np.float32)
+            center_point = tuple(np.mean(pts_array, axis=0))
+
+            # 判断点是否在扩大后的多边形内
+            # pointPolygonTest: >= 0 表示在内部或边缘
+            if cv2.pointPolygonTest(expanded_poly, center_point, False) >= 0:
+                valid_shapes.append(shape)
+            else:
+                removed_count += 1
+
+        # 更新 shapes
+        defect_json['shapes'] = valid_shapes
+        logger.info(
+            f"[FilterOutsideDefects] 过滤完成: 保留 {len(valid_shapes)} 个, 移除 {removed_count} 个 (外扩 {self.expansion_pixel} px)")
+
+        return defect_json

+ 1 - 1
app/utils/score_inference/CardScorer.py

@@ -88,7 +88,7 @@ class CardScorer:
                     defect_type = "wear_area"
                 elif defect['label'] in ['scratch', 'scuff']:
                     defect_type = "scratch_length"
-                elif defect['label'] in ['pit', 'impact']:
+                elif defect['label'] in ['pit', 'impact', 'protrudent']:
                     defect_type = "pit_area"
                 elif defect['label'] in ['stain']:
                     defect_type = "stain_area"