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

首先，你需要准备一个包含红绿灯图像的数据集。然后，你可以使用YOLOv5训练一个物体检测模型来识别红绿灯。

以下是一个简单的例子，假设你的数据集已准备好，包含两个类别：红灯和绿灯。你可以使用YOLOv5 PyTorch实现库进行训练和推理。

import torch
import cv2
import numpy as np
from models.experimental import attempt_load
from utils.general import non_max_suppression


def detect_traffic_lights(image_path, conf_threshold=0.5, iou_threshold=0.5):
    # 加载模型
    model = attempt_load('yolov5s.pt', map_location=torch.device('cpu'))

    # 加载图像
    img = cv2.imread(image_path)

    # 图像预处理
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.transpose(2, 0, 1)
    img = np.expand_dims(img, axis=0)
    img = img / 255.0

    # 推理
    with torch.no_grad():
        detections = model(torch.from_numpy(img).float())

    # 非极大值抑制
    detections = non_max_suppression(detections, conf_threshold, iou_threshold)[0]

    # 可视化结果
    for det in detections:
        x1, y1, x2, y2, conf, cls = det.tolist()
        class_name = '红灯' if cls == 0 else '绿灯'
        print(f'类别: {class_name}, 置信度: {conf:.2f}, 位置: ({x1:.2f}, {y1:.2f}) - ({x2:.2f}, {y2:.2f})')
        cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 2)

    # 显示图像
    img = img.squeeze()
    img = img.transpose(1, 2, 0)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    cv2.imshow('image', img)
    cv2.waitKey(0)


if __name__ == '__main__':
    detect_traffic_lights('traffic_light.jpg')

在上述代码中，我们首先加载YOLOv5模型，然后加载图像并进行预处理。接下来，我们使用模型进行推断，并对检测到的对象进行非极大值抑制。最后，我们可视化结果并显示图像。