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)

            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

            # 确保生成的点是 float32，用于相交计算
            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

        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)
        # 修改：为了支持 intersectConvexConvex，这里转换为 float32 而不是 int32
        inner_box = np.array(inner_box, dtype=np.float32)
        return inner_box

    def check_intersection(self, poly1, poly2):
        """
        检测两个凸多边形是否相交（重合面积 > 0）
        poly1, poly2: np.float32 类型的多边形点集
        """
        if poly1 is None or poly2 is None:
            return False

        # intersectConvexConvex 返回 (intersection_area, intersection_polygon)
        # 我们只需要面积即可。如果面积 > 0 (考虑到浮点误差，可以用 > 0.01) 即视为相交
        try:
            intersection_area, _ = cv2.intersectConvexConvex(poly1, poly2)
            return intersection_area > 0.1  # 设置一个微小的阈值防止浮点噪声
        except Exception as e:
            # 极少数情况下如果多边形退化可能报错，视为不相交
            logger.debug(f"Intersection check failed: {e}")
            return False

    def classify_defects_location(self, defect_data, outer_box_data):
        """
        主函数：分类为 corner, face, edge (基于重合逻辑)
        优先级：Face (只要碰到面就算面) > Corner (没碰到面但碰到角) > Edge (其余)
        """
        if not defect_data or 'defects' not in defect_data or not defect_data['defects']:
            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:
            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

            # --- 步骤变更：获取缺陷的多边形轮廓 ---
            # defect['min_rect'] 结构通常为 ((x,y), (w,h), angle)
            defect_rect = defect['min_rect']
            defect_box = cv2.boxPoints(defect_rect)
            defect_poly = np.array(defect_box, dtype=np.float32)

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

            # 优先级 1: 面 (Face) - 只要碰到 Face 区域就算 Face
            is_face = False
            if face_region is not None:
                if self.check_intersection(face_region, defect_poly):
                    is_face = True

            if is_face:
                defect['defect_type'] = 'face'
            else:
                # 优先级 2: 角 (Corner) - 没碰到 Face，但碰到了 Corner 区域
                is_corner = False
                for region in corner_regions:
                    if self.check_intersection(region, defect_poly):
                        is_corner = True
                        break

                if is_corner:
                    defect['defect_type'] = 'corner'
                else:
                    # 优先级 3: 边 (Edge)
                    # 既没碰到面，也没碰到角，说明完全处于边区域内（或者在外框外）
                    defect['defect_type'] = 'edge'

            processed_count += 1

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