57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
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'") |