写一个基于yolov5和opencv的训练红绿灯py代码

以下是基于 YOLOv5 和 OpenCV 的红绿灯检测和识别的 Python 代码示例：

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

# 加载模型
model = attempt_load(Path('path/to/yolov5s.pt'), map_location=torch.device('cpu'))
model.eval()

# 定义类别
classes = ['red', 'yellow', 'green']

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

# 加载输入图像
img_path = 'path/to/image.jpg'
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# 图像预处理
img = cv2.resize(img, (640, 640))
img = img.astype(np.float32) / 255.0
img = np.transpose(img, (2, 0, 1))
img = np.expand_dims(img, axis=0)

# 模型预测
with torch.no_grad():
    pred = model(torch.from_numpy(img))
    pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.5)

# 遍历预测结果
for i, det in enumerate(pred):
    if len(det):
        for *xyxy, conf, cls in reversed(det):
            # 提取坐标和类别
            x1, y1, x2, y2 = map(int, xyxy)
            cls_name = classes[int(cls)]

            # 绘制边框和类别
            cv2.rectangle(img, (x1, y1), (x2, y2), colors[int(cls)], 2)
            cv2.putText(img, cls_name, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, colors[int(cls)], 2)

# 显示结果
cv2.imshow('image', img)
cv2.waitKey(0)
cv2.destroyAllWindows()

在这个示例中，我们首先加载了 YOLOv5 模型并定义了类别和颜色。然后，我们加载输入图像并进行预处理。接下来，我们使用模型对图像进行预测，并使用非最大抑制（NMS）来过滤预测结果。最后，我们遍历预测结果并绘制边框和类别。最终，我们显示结果。

请注意，这个示例只是一个简单的演示，你需要根据实际情况来调整参数和代码。同时，由于 YOLOv5 和 OpenCV 的版本可能会不同，你需要根据实际情况来调整代码和模型。