给我一个使用 yolov5 模型进行道路缺陷检测的代码
时间: 2024-02-13 15:07:08 浏览: 24
以下是使用 YOLOv5 模型进行道路缺陷检测的 Python 代码示例:
```python
import torch
import cv2
import numpy as np
from pathlib import Path
# 加载 YOLOv5 模型
model_path = 'path/to/yolov5.pt'
model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)
# 定义道路缺陷类别
class_names = ['crack', 'pothole']
# 加载测试图片
test_image_path = 'path/to/test/image.jpg'
test_image = cv2.imread(test_image_path)
# 进行推理
results = model(test_image)
# 处理检测结果
predictions = results.xyxy[0]
for pred in predictions:
label = int(pred[5])
class_name = class_names[label]
confidence = pred[4]
x1, y1, x2, y2 = map(int, pred[:4])
cv2.rectangle(test_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(test_image, f"{class_name} {confidence:.2f}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
# 显示结果
cv2.imshow('result', test_image)
cv2.waitKey(0)
cv2.destroyAllWindows()
```
在这个代码中,我们首先加载了预训练的 YOLOv5 模型,然后定义了道路缺陷的类别。接着,我们加载测试图片并进行推理,最后处理检测结果并显示在图片上。注意,这个示例代码中的模型路径和测试图片路径需要替换为你自己的路径。