import cv2
import numpy as np

class ImagePreprocessor:
    """图像预处理：分离目标颜色并生成掩码"""
    @staticmethod
    def process_red_light(frame):
        """红色灯条预处理"""
        red_only = frame.copy()
        red_only[:, :, 0] = 0  # B=0
        red_only[:, :, 1] = 0  # G=0

        hsv = cv2.cvtColor(red_only, cv2.COLOR_BGR2HSV)
        v_mask = cv2.inRange(hsv[:, :, 2], 70, 190)

        kernel_close = np.ones((7, 7), np.uint8)
        kernel_dilate = np.ones((2, 2), np.uint8)
        mask = cv2.morphologyEx(v_mask, cv2.MORPH_CLOSE, kernel_close)
        mask = cv2.dilate(mask, kernel_dilate, iterations=1)

        kernel = np.ones((4, 4), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=1)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
        mask = cv2.dilate(mask, kernel, iterations=1)

        return mask, red_only

    @staticmethod
    def process_blue_light(frame):
        """蓝色灯条预处理"""
        blue_only = frame.copy()
        blue_only[:, :, 1] = 0  # G=0
        blue_only[:, :, 2] = 0  # R=0

        hsv = cv2.cvtColor(blue_only, cv2.COLOR_BGR2HSV)
        v_mask = cv2.inRange(hsv[:, :, 2], 120, 230)

        kernel_close = np.ones((7, 7), np.uint8)
        kernel_dilate = np.ones((2, 2), np.uint8)
        mask = cv2.morphologyEx(v_mask, cv2.MORPH_CLOSE, kernel_close)
        mask = cv2.dilate(mask, kernel_dilate, iterations=1)

        kernel = np.ones((4, 4), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=1)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
        mask = cv2.dilate(mask, kernel, iterations=1)

        return mask, blue_only

    def get_mask(self, frame, target_color):
        """根据颜色选择预处理逻辑"""
        if target_color == "red":
            return self.process_red_light(frame)
        elif target_color == "blue":
            return self.process_blue_light(frame)
        else:
            raise ValueError("仅支持 'red' 或 'blue'")