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):
        self.PIXEL_RESOLUTION_UM = PIXEL_RESOLUTION_UM  # 单位: μm/pixel
        # 角区的定义尺寸
        self.CORNER_SIZE_MM = PIXEL_CENTER_UM

    def get_outer_box_corners(self, outer_box_data):
        """从外框数据中计算并返回精确的四个角点坐标。"""
        if not outer_box_data or 'shapes' not in outer_box_data or not outer_box_data['shapes']:
            raise ValueError("无效的外框数据格式。")

        points = outer_box_data['shapes'][0]['points']
        contour = np.array(points, dtype=np.int32)

        # 计算最小面积旋转矩形
        rotated_rect = cv2.minAreaRect(contour)

        # 获取矩形的4个角点
        # boxPoints返回的角点顺序是可预测的，通常是顺时针或逆时针
        box_corners = cv2.boxPoints(rotated_rect)
        return box_corners

    def create_corner_regions(self, outer_box_corners, corner_size_px):
        """
        根据外框的角点和指定的尺寸（像素），创建四个角区的多边形。

        Args:
            outer_box_corners (np.array): 外框的4个角点坐标。
            corner_size_px (float): 角区尺寸（像素单位）。

        Returns:
            list: 包含四个角区（每个角区都是一个4x2的np.array）的列表。
        """
        corner_regions = []
        num_corners = len(outer_box_corners)

        for i in range(num_corners):
            # 获取当前角点 p_i 和它的两个相邻角点 p_{i-1} 和 p_{i+1}
            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

            # 将向量归一化为单位向量
            unit_vec1 = vec1 / np.linalg.norm(vec1)
            unit_vec2 = vec2 / np.linalg.norm(vec2)

            # 计算角区正方形的四个顶点
            # q0: 当前角点
            # q1: 沿一个边移动指定距离
            # q2: 沿另一个边移动指定距离
            # q3: 两个向量相加得到的内部点
            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 classify_defects_location(self, defect_data, outer_box_data):
        """
        主函数，对缺陷数据进行分类并添加 "defect_type" 标签。
        """
        if not defect_data or 'defects' not in defect_data:
            logger.warn("警告: 缺陷数据为空或格式不正确，跳过处理。")
            return defect_data

        if 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
        logger.info(f"3mm角区尺寸已转换为 {corner_size_px:.2f} 像素。")

        # 2. 获取外框几何信息
        try:
            outer_box_corners = self.get_outer_box_corners(outer_box_data)
        except ValueError as e:
            logger.info(f"错误: {e}")
            return None

        # 3. 定义四个角区
        corner_regions = self.create_corner_regions(outer_box_corners, corner_size_px)
        logger.info(f"已成功定义 {len(corner_regions)} 个角区。")

        # 4. 遍历并分类所有缺陷
        processed_count = 0
        for defect in defect_data['defects']:
            # 使用JSON中预先计算好的min_rect中心点，更高效
            if 'min_rect' not in defect or not defect['min_rect']:
                logger.warn(f"警告: 缺陷 '{defect.get('label')}' 缺少 'min_rect' 信息，跳过。")
                continue

            center_point = tuple(defect['min_rect'][0])
            is_in_corner = False

            # 5. 判断缺陷中心点是否在任何一个角区内
            for region in corner_regions:
                # pointPolygonTest 返回值: >0 (内部), 0 (边上), <0 (外部)
                if cv2.pointPolygonTest(region, center_point, False) >= 0:
                    is_in_corner = True
                    break

            # 6. 添加新标签
            if is_in_corner:
                defect['defect_type'] = 'corner'
            else:
                defect['defect_type'] = 'edge'

            processed_count += 1

        logger.info(f"处理完成！共为 {processed_count} 个缺陷添加了 'defect_type' 标签。")
        return defect_data