CheckCardBoxAndDefectServer/app/utils/defect_inference/ClassifyEdgeCorner.py

import json
import cv2
import numpy as np
import math
import logging

logger = logging.getLogger('ClassifyEdgeCorner')


class ClassifyEdgeCorner:
    def __init__(self, PIXEL_RESOLUTION_UM: float, PIXEL_CENTER_UM: float, EDGE_SIZE_MM: float):
        """
        Args:
            PIXEL_RESOLUTION_UM: 像素分辨率 (um/pixel)
            PIXEL_CENTER_UM: 角区定义的尺寸 (mm) - 对应原本的 CORNER_SIZE_MM
            EDGE_SIZE_MM: 边区向内缩进的尺寸 (mm)
        """
        self.PIXEL_RESOLUTION_UM = PIXEL_RESOLUTION_UM
        self.CORNER_SIZE_MM = PIXEL_CENTER_UM
        self.EDGE_SIZE_MM = EDGE_SIZE_MM  # 新增：边的宽度

    def get_rotated_rect_data(self, outer_box_data):
        """
        从外框数据中提取最小外接矩形的原始数据 ((x,y), (w,h), angle)。
        """
        if not outer_box_data or 'shapes' not in outer_box_data or not outer_box_data['shapes']:
            raise ValueError("无效的外框数据格式。")

        get_point = outer_box_data['shapes'][0].get('points', None)
        if get_point is None:
            logger.info("外框计算非第一次, 取 rect_box")
            get_point = outer_box_data['shapes'][0].get('rect_box', None)
            if get_point is None:
                raise ValueError("外框数据缺失: rect_box")

        points = get_point
        contour = np.array(points, dtype=np.int32)

        # 返回 ((cx, cy), (w, h), angle)
        return cv2.minAreaRect(contour)

    def create_corner_regions(self, outer_box_corners, corner_size_px):
        """
        根据外框的角点和指定的尺寸（像素），创建四个角区的多边形。
        (逻辑保持不变)
        """
        corner_regions = []
        num_corners = len(outer_box_corners)

        for i in range(num_corners):
            p_current = outer_box_corners[i]
            p_prev = outer_box_corners[(i - 1 + num_corners) % num_corners]
            p_next = outer_box_corners[(i + 1) % num_corners]

            vec1 = p_prev - p_current
            vec2 = p_next - p_current

            norm1 = np.linalg.norm(vec1)
            norm2 = np.linalg.norm(vec2)

            # 防止除以0的保护
            if norm1 == 0 or norm2 == 0:
                unit_vec1 = np.array([0, 0])
                unit_vec2 = np.array([0, 0])
            else:
                unit_vec1 = vec1 / norm1
                unit_vec2 = vec2 / norm2

            q0 = p_current
            q1 = p_current + corner_size_px * unit_vec1
            q2 = p_current + corner_size_px * unit_vec2
            q3 = q1 + corner_size_px * unit_vec2

            corner_poly = np.array([q0, q1, q3, q2], dtype=np.float32)
            corner_regions.append(corner_poly)

        return corner_regions

    def create_inner_face_polygon(self, rect_data, edge_size_px):
        """
        根据外框旋转矩形和边宽，计算内缩后的“面”区域多边形。
        """
        (cx, cy), (w, h), angle = rect_data

        # 计算内框的宽和高 (两边各减去 edge_size)
        inner_w = w - 2 * edge_size_px
        inner_h = h - 2 * edge_size_px

        # 如果边设置得太宽，导致内框消失，则返回 None 或极小的框
        if inner_w <= 0 or inner_h <= 0:
            logger.warning("警告: EDGE_SIZE_MM 设置过大，超过了卡牌尺寸，面(Face)区域将不存在。")
            return None

        # 生成内框的旋转矩形
        inner_rect = ((cx, cy), (inner_w, inner_h), angle)

        # 获取内框的四个角点
        inner_box = cv2.boxPoints(inner_rect)
        inner_box = np.int32(inner_box)  # 转换为整数点以便于 polygon test
        return inner_box

    def classify_defects_location(self, defect_data, outer_box_data):
        """
        主函数：分类为 corner, face, edge
        """
        if not defect_data or 'defects' not in defect_data or not defect_data['defects']:
            # logger.warn("无缺陷数据")
            return defect_data

        # 1. 单位转换
        pixel_to_mm = self.PIXEL_RESOLUTION_UM / 1000.0
        corner_size_px = self.CORNER_SIZE_MM / pixel_to_mm
        edge_size_px = self.EDGE_SIZE_MM / pixel_to_mm

        logger.info(
            f"Param Info: Corner={self.CORNER_SIZE_MM}mm ({corner_size_px:.1f}px), Edge={self.EDGE_SIZE_MM}mm ({edge_size_px:.1f}px)")

        # 2. 获取外框几何信息
        try:
            # 获取旋转矩形完整数据 ((x,y), (w,h), angle)
            rect_data = self.get_rotated_rect_data(outer_box_data)
            # 获取外框的四个角点 (用于角区计算)
            outer_box_corners = cv2.boxPoints(rect_data)
        except ValueError as e:
            logger.error(f"外框数据错误: {e}")
            return None

        # 3. 定义区域
        # A. 定义四个角区
        corner_regions = self.create_corner_regions(outer_box_corners, corner_size_px)

        # B. 定义中间的面区 (内缩矩形)
        face_region = self.create_inner_face_polygon(rect_data, edge_size_px)

        # 4. 遍历并分类所有缺陷
        processed_count = 0
        for defect in defect_data['defects']:
            if 'min_rect' not in defect or not defect['min_rect']:
                continue

            # 获取缺陷中心点
            center_point = tuple(defect['min_rect'][0])

            # --- 判定逻辑优先级 ---

            # 优先级 1: 角 (Corner)
            is_corner = False
            for region in corner_regions:
                # pointPolygonTest: >0 内部, =0 边界, <0 外部
                if cv2.pointPolygonTest(region, center_point, False) >= 0:
                    is_corner = True
                    break

            if is_corner:
                defect['defect_type'] = 'corner'
            else:
                # 优先级 2: 面 (Face)
                # 如果点在内缩矩形内部，且不是角，那就是面
                is_face = False
                if face_region is not None:
                    if cv2.pointPolygonTest(face_region, center_point, False) >= 0:
                        is_face = True

                if is_face:
                    defect['defect_type'] = 'face'
                else:
                    # 优先级 3: 边 (Edge)
                    # 既不在角区，也不在面区（内框）里，但还在处理范围内，那就是边
                    # (注：这里默认缺陷是在外框内的，如果是外框外的噪声，通常也会被算作 Edge 或需要额外过滤)
                    defect['defect_type'] = 'edge'

            processed_count += 1

        logger.info(f"分类完成: {processed_count} 个缺陷处理完毕。")
        return defect_data