# -*- coding:utf-8 -*-
# @Author len
# @Create 2023/10/27 15:24

import cv2
import numpy as np

# 读取图像
image = cv2.imread('trafficLight_red.png')

# 将图像转换为HSV色彩空间
hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

# 定义红色在HSV中的范围
lower_red_1 = np.array([0,50,50])
upper_red_1 = np.array([10,255,255])

lower_red_2 = np.array([160,50,50])
upper_red_2 = np.array([180,255,255])

# 创建掩码
mask1 = cv2.inRange(hsv, lower_red_1, upper_red_1)
mask2 = cv2.inRange(hsv, lower_red_2, upper_red_2)

mask = mask1 + mask2

# 对掩码进行腐蚀和膨胀操作，以消除噪声
mask = cv2.erode(mask, None, iterations=2)
mask = cv2.dilate(mask, None, iterations=2)

# 找到掩码中的轮廓
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

# 遍历轮廓并在原始图像上绘制
for contour in contours:
    if cv2.contourArea(contour) > 100: # 你可以根据需要调整这个值
        x, y, w, h = cv2.boundingRect(contour)
        cv2.rectangle(image, (x, y), (x+w, y+h), (0, 255, 0), 2)

# 显示结果
cv2.imshow('Detected Red Light', image)
cv2.waitKey(0)
cv2.destroyAllWindows()