yolov5红绿灯识别

Question

请问怎然才能返回目标物体的类别，比如这个怎样才能返回red，或者green，用的yolov5-master源代码检测

import torch
from matplotlib import pyplot as plt
import numpy as np
import cv2
import time
import string
import pygame

#model = torch.hub.load("ultralytics/yolov5","yolov5s")
model = torch.hub.load('D:/ks/yolov5-master','custom','runs/train/exp11/weights/best.pt', source='local')
model.conf=0.4

video_path = "./data/images/hld.mp4"
cap = cv2.VideoCapture(video_path)
#cap = cv2.VideoCapture(0)


while cap.isOpened():
    # 参数ret 为True 或者False,代表有没有读取到图片帧
    # 第二个参数frame表示截取到一帧的图片
    ret,frame = cap.read()
    # 将图片帧输入模型
    frame=cv2.resize(frame,(500,500))
    results = model(frame)
    print(results)

    #LOGGER.info(f"{names[int(c)]}"))
    #print(results.names)
    
    cv2.imshow('YOLO',np.squeeze(results.render()))
    # 退出
    if cv2.waitKey(10) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

Answer 1

YOLOv5的预训练模型里，对交通灯可能没有区分红绿这些颜色，可能就是统一的object「交通灯」。
但是用YOLO你拿到了目标检测的box，你把交通灯区域取出来，用opencv也能识别颜色。要不然你就需要准备一些不同颜色交通灯图片，标注一下，重新train一下YOLOv5

import cv2

# 读取图像
image = cv2.imread("traffic_light.jpg")

# 转换到 HSV 颜色空间
hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

# 设定红色的阈值
lower_red = np.array([0, 50, 50])
upper_red = np.array([10, 255, 255])

# 根据阈值构建掩模
mask = cv2.inRange(hsv, lower_red, upper_red)

# 对原图像和掩模进行位运算
res = cv2.bitwise_and(image, image, mask=mask)

# 显示图像
cv2.imshow("Image", image)
cv2.imshow("Mask", mask)
cv2.imshow("Result", res)

# 等待用户按键
cv2.waitKey(0)