57 lines
1.6 KiB
Python
57 lines
1.6 KiB
Python
|
# -*- 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)}")
|