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

Hack Hui

已于 2024-04-28 09:28:42 修改

阅读量2.9k

点赞数

CC 4.0 BY-SA版权

分类专栏： yolov5 深度学习 Python 文章标签： python 计算机视觉人工智能 opencv 图像处理

于 2023-03-22 21:48:16 首次发布

本文链接：https://blog.csdn.net/m0_58508552/article/details/129720196

Python 同时被 3 个专栏收录

11 篇文章

订阅专栏

yolov5

9 篇文章

订阅专栏

深度学习

7 篇文章

订阅专栏

该代码示例展示了如何利用YOLOv5模型进行车辆检测。首先安装模型，然后加载预训练权重，定义检测函数，对输入图片进行处理，通过非极大值抑制减少边界框，最后在图像上绘制检测结果并输出车辆数量。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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


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