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


import cv2
import numpy as np


def detect_light_color(image_path):
    # 读取图像
    image = cv2.imread(image_path)

    # 将图像从BGR转换到HSV
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    # 为红色定义HSV范围
    lower_red1 = np.array([0, 50, 50])
    upper_red1 = np.array([10, 255, 255])
    lower_red2 = np.array([160, 50, 50])
    upper_red2 = np.array([180, 255, 255])

    # 为绿色定义HSV范围
    lower_green = np.array([40, 50, 50])
    upper_green = np.array([90, 255, 255])

    # 为黄色定义HSV范围
    lower_yellow = np.array([20, 50, 50])
    upper_yellow = np.array([40, 255, 255])

    # 创建掩模
    mask_red1 = cv2.inRange(hsv, lower_red1, upper_red1)
    mask_red2 = cv2.inRange(hsv, lower_red2, upper_red2)
    mask_red = mask_red1 + mask_red2
    mask_green = cv2.inRange(hsv, lower_green, upper_green)
    mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)

    # 计算掩模中的白色像素数量
    red_count = cv2.countNonZero(mask_red)
    green_count = cv2.countNonZero(mask_green)
    yellow_count = cv2.countNonZero(mask_yellow)

    # 根据像素数量判断颜色
    if red_count > green_count and red_count > yellow_count:
        return "Red"
    elif green_count > red_count and green_count > yellow_count:
        return "Green"
    elif yellow_count > red_count and yellow_count > green_count:
        return "Yellow"
    else:
        return "Unknown"


# 测试
for image_path in ["cropped_image_red.jpg", "img.png", "cropped_image_green.jpg"]:
    print(f"{image_path}: {detect_light_color(image_path)}")