CheckCardBoxAndDefectServer/app/utils/card_inference/handle_result.py

import cv2
import numpy as np
from typing import List, Dict, Tuple


class ShapeDrawer:
    """形状绘制工具类"""

    @staticmethod
    def points_to_contour(points: List[List[int]]) -> np.ndarray:
        """
        将点集转换为轮廓

        Args:
            points: 点集列表，每个点是 [x, y] 格式

        Returns:
            np.ndarray: OpenCV格式的轮廓
        """
        # 确保点集是整数类型的numpy数组
        contour = np.array(points, dtype=np.int32)
        # 重塑为OpenCV轮廓格式 (N, 1, 2)
        return contour.reshape((-1, 1, 2))

    @staticmethod
    def contour_to_mask(contour: np.ndarray,
                        image_shape: Tuple[int, int]) -> np.ndarray:
        """
        将轮廓转换为掩码

        Args:
            contour: OpenCV格式的轮廓
            image_shape: 图像形状 (height, width)

        Returns:
            np.ndarray: 二值掩码
        """
        # 创建空白掩码
        mask = np.zeros(image_shape, dtype=np.uint8)
        # 填充轮廓
        cv2.fillPoly(mask, [contour], 255)
        return mask

    @staticmethod
    def apply_mask_to_image(image: np.ndarray,
                            mask: np.ndarray,
                            color: Tuple[int, int, int],
                            alpha: float = 0.5) -> np.ndarray:
        """
        将掩码应用到图像上

        Args:
            image: 原始图像
            mask: 二值掩码
            color: BGR颜色
            alpha: 透明度

        Returns:
            np.ndarray: 带有掩码的图像
        """
        # 创建图像副本
        result = image.copy()

        # 创建彩色掩码
        color_mask = np.zeros_like(image)
        # color_mask[mask > 0] = color

        result[mask > 0] = color

        # 混合图像
        # cv2.addWeighted(color_mask, alpha, result, 1 - alpha, 0, result)

        return result

    @staticmethod
    def draw_shape_from_points(image: np.ndarray,
                               points: List[List[int]],
                               color: Tuple[int, int, int] = (0, 255, 0),
                               alpha: float = 0.5,
                               draw_contour: bool = True) -> np.ndarray:
        """
        从点集绘制形状到图像上

        Args:
            image: 原始图像
            points: 点集列表
            color: BGR颜色
            alpha: 透明度
            draw_contour: 是否绘制轮廓线

        Returns:
            np.ndarray: 处理后的图像
        """
        # 1. 转换为轮廓
        contour = ShapeDrawer.points_to_contour(points)

        # 2. 创建掩码
        mask = ShapeDrawer.contour_to_mask(contour, image.shape[:2])

        # 3. 应用掩码
        result = ShapeDrawer.apply_mask_to_image(image, mask, color, alpha)

        # 4. 可选：绘制轮廓线
        if draw_contour:
            cv2.drawContours(result, [contour], -1, color, 2)

        return result


def process_detection_result(image: np.ndarray,
                             detection: Dict,
                             label_colors: Dict
                             ) -> np.ndarray:
    """
    处理检测结果并在图像上绘制

    Args:
        image: 原始图像
        detection: 检测结果字典

    Returns:
        np.ndarray: 处理后的图像
    """
    result = image.copy()

    # 处理每个检测到的形状
    for shape in detection['shapes']:
        points = shape['points']
        label = shape['label']
        conf = shape['probability']

        # 获取颜色（默认绿色）
        color = label_colors.get(label, (0, 255, 0))

        # 绘制形状
        result = ShapeDrawer.draw_shape_from_points(
            result,
            points,
            color=color,
            alpha=0.3,
            draw_contour=True
        )

        # 添加标签和置信度
        first_point = points[0]
        is_print_text = False
        if is_print_text:
            cv2.putText(
                result,
                f"{label}: {conf:.2f}",
                (first_point[0], first_point[1] - 10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                color,
                2
            )

    return result


# 示例使用
if __name__ == "__main__":
    # 加载示例图像
    image = cv2.imread("example.jpg")

    # 为不同类别设置不同颜色
    label_colors = {
        'baseboard': (0, 255, 0),
    }

    # 示例检测结果
    detection_result = {
        'num': 1,
        'cls': [1],
        'names': ['baseboard'],
        'conf': 0.9966249431977074,
        'shapes': [{
            'class_num': 1,
            'label': 'baseboard',
            'probability': 0.9966249431977074,
            'points': [[5, 171], [4, 172], [0, 172], [0, 487], [34, 487],
                       # ... 其他点 ...
                       [1019, 172], [1018, 171]]
        }]
    }

    # 处理图像
    result = process_detection_result(image, detection_result, label_colors)

    # 显示结果
    cv2.imshow("Result", result)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    # 保存结果
    # cv2.imwrite("result.jpg", result)