【源码】基于PyTorch框架的YOLOv5视觉AI项目

项目简介

本项目是基于PyTorch框架开发的YOLOv5视觉AI系统，代表了Ultralytics对未来视觉人工智能方法的开源研究。它融合了数千小时研发积累的经验与最佳实践，可处理目标检测、实例分割、图像分类等多种视觉任务，具有简单易上手、注重实际效果的特点。

项目的主要特性和功能

1. 预训练模型丰富：提供多种预训练模型（如YOLOv5n、YOLOv5s等），可直接用于推理，满足不同场景需求。
2. 支持多任务处理：能进行目标检测、实例分割和图像分类任务。实例分割模型速度快且精度高，分类模型也能高效准确分类。
3. 训练便捷：支持自动下载部分数据集，可使用单GPU或多GPU进行训练，且提供详细的训练参数设置。
4. 推理灵活：可通过多种方式进行推理，如使用PyTorch Hub、detect.py脚本等，支持图像、视频、摄像头等多种输入源。
5. 模型导出方便：可将模型导出为ONNX、TensorRT等格式，便于在不同平台部署。
6. 集成多种工具：与Roboflow、ClearML、Comet、Neural Magic等工具集成，实现数据集标注与导出、训练过程跟踪与可视化、模型保存与调试等功能。

安装使用步骤

安装

假设用户已经下载了本项目的源码文件，首先需要复制项目仓库，并在Python>=3.8.0且PyTorch>=1.8的环境中安装所需依赖： bash cd yolov5 pip install -r requirements.txt

推理

使用PyTorch Hub推理

```python import torch

model = torch.hub.load("ultralytics/yolov5", "yolov5s")

img = "https://ultralytics.com/images/zidane.jpg"

results = model(img)

results.print() ```

使用detect.py推理

bash python detect.py --weights yolov5s.pt --source 0 # 摄像头 python detect.py --weights yolov5s.pt --source img.jpg # 图像 python detect.py --weights yolov5s.pt --source vid.mp4 # 视频

训练

以下命令可重现YOLOv5在COCO数据集上的训练结果： bash python train.py --data coco.yaml --epochs 300 --weights '' --cfg yolov5n.yaml --batch-size 128

实例分割使用步骤

训练

```bash python segment/train.py --data coco128-seg.yaml --weights yolov5s-seg.pt --img 640

python -m torch.distributed.run --nproc_per_node 4 --master_port 1 segment/train.py --data coco128-seg.yaml --weights yolov5s-seg.pt --img 640 --device 0,1,2,3 ```

验证

bash bash data/scripts/get_coco.sh --val --segments python segment/val.py --weights yolov5s-seg.pt --data coco.yaml --img 640

预测

bash python segment/predict.py --weights yolov5m-seg.pt --source data/images/bus.jpg

模型导出

bash python export.py --weights yolov5s-seg.pt --include onnx engine --img 640 --device 0

分类网络使用步骤

训练

```bash python classify/train.py --model yolov5s-cls.pt --data cifar100 --epochs 5 --img 224 --batch 128

python -m torch.distributed.run --nproc_per_node 4 --master_port 1 classify/train.py --model yolov5s-cls.pt --data imagenet --epochs 5 --img 224 --device 0,1,2,3 ```

验证

bash bash data/scripts/get_imagenet.sh --val python classify/val.py --weights yolov5m-cls.pt --data../datasets22/imagenet --img 224

预测

bash python classify/predict.py --weights yolov5s-cls.pt --source data/images/bus.jpg

模型导出

bash python export.py --weights yolov5s-cls.pt resnet50.pt efficientnet_b0.pt --include onnx engine --img 224

下载地址

点击下载 【提取码: 4003】【解压密码: www.makuang.net】