用yolo5写一个红绿灯图像识别的程序

好的，以下是基于yolo5的红绿灯图像识别程序的Python代码：

import torch
import cv2
import numpy as np

# 加载yolo5模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

# 定义类别
classes = ['red_light', 'yellow_light', 'green_light']

# 定义颜色
colors = [[0, 0, 255], [0, 255, 255], [0, 255, 0]]

# 加载图像
img = cv2.imread('test.jpg')

# 调整图像大小
img = cv2.resize(img, (640, 640))

# 预测结果
results = model(img)

# 获取预测框和类别
boxes = results.xyxy[0][:, :4].numpy()
pred_classes = results.xyxy[0][:, 5].numpy()

# 遍历预测框和类别
for i in range(len(pred_classes)):
    # 获取预测框坐标和类别
    x1, y1, x2, y2 = boxes[i]
    pred_class = int(pred_classes[i])
    # 如果类别是红灯、黄灯或绿灯
    if pred_class in [0, 1, 2]:
        # 绘制预测框和类别标签
        cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), colors[pred_class], 2)
        cv2.putText(img, classes[pred_class], (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors[pred_class], 2)

# 显示图像
cv2.imshow('image', img)
cv2.waitKey(0)
cv2.destroyAllWindows()

注意事项：

1. 需要安装yolo5和OpenCV库；
2. 图像路径需要根据实际情况修改；
3. 代码中的颜色和类别可能需要根据实际情况进行调整。

希望能够帮到你！