测试文件夹架构3

新建文件夹/image_preprocessor.py

@@ -0,0 +1,57 @@
+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'")