纯代码干货-Python基于YOLOv5的车辆检测并计数

首先,你需要安装YOLOv5,可以通过官方文档了解安装步骤和相关依赖。接着,你需要准备训练好的车辆检测模型和测试图片或视频。

纯代码干货-Python基于YOLOv5的车辆检测并计数_第1张图片

1. 导入库

```
import cv2
import torch
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.plots import Annotator
```

2. 加载模型

```
weights = 'path/to/weights' # 模型权重文件
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = attempt_load(weights, map_location=device)
```

3. 定义检测函数

```
def detect(model, img, conf_thres=0.25, iou_thres=0.45):
    """使用YOLOv5模型检测车辆并计数"""
    # Resize input image
    img = cv2.resize(img, (640, 640))
    # Convert to RGB format and normalize
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
    # Convert to tensor and unsqueeze for batch dimension
    img = torch.from_numpy(img).to(device).float().unsqueeze(0)
    # Inference
    pred = model(img)[0]
    # Process detections
    pred = non_max_suppression(pred, conf_thres, iou_thres)
    # Count vehicles
    count = len(pred)
    # Draw bounding boxes
    annotator = Annotator(img, line_width=3, font_scale=0.5)
    for det in pred:
        annotator.box(det[:4], label=f'VEHICLE {det[5]:.2f}')
    # Convert annotated image back to BGR format
    img = cv2.cvtColor(annotator.image, cv2.COLOR_RGB2B
4. 加载图片并调用检测函数



```

img_path = 'path/to/image'  # 测试图片路径

img = cv2.imread(img_path)

count = detect(model, img)

print(f'车辆数量:{count}')

cv2.imshow('Detection result', img)

cv2.waitKey()

以上代码片段演示了如何使用YOLOv5模型进行车辆检测,并在检测结果上绘制边界框并输出车辆数量。你可以根据实际需求调整检测函数的参数,比如置信度阈值和重叠度阈值,以及绘制边界框的样式和标签。同时,你也可以扩展这个脚本,实现批量处理图片或视频,输出更详细的统计信息,或者将检测结果存储到数据库或文本文件中。

你可能感兴趣的:(yolov5,深度学习,Python,python,计算机视觉,人工智能,opencv,图像处理)