【源码】基于Python框架的YOLOv5目标检测项目

项目简介

本项目是基于YOLOv5的目标检测项目，构建了一个灵活且易于扩展的模型架构。它支持多种骨干网络，用户还能方便地添加自定义模块。同时，该项目提供了模型的训练、测试、部署功能，并且支持使用TensorRT进行模型优化和推理。

项目的主要特性和功能

1. 灵活的模型架构：把YOLOv5模型拆分为backbone、neck和head，便于用户灵活修改和扩展网络。
2. 多种骨干网络支持：支持ResNet、MobileNet、ShuffleNet等多种骨干网络，还可轻松集成其他骨干网络。
3. 自定义模块：能添加SE、CBAM等自定义模块，增强模型表现力。
4. 模型训练与测试：具备与YOLOv5相同的训练和测试流程，确保兼容性和易用性。
5. 模型部署：支持将模型导出为ONNX格式，并用TensorRT进行优化和推理。
6. 量化支持：可直接对ONNX模型进行量化，提升推理速度。

安装使用步骤

1. 环境准备

确保已安装Python环境，并根据requirements.txt安装项目所需的依赖库。

2. 数据准备

使用od/data/transform_voc.py脚本将VOC数据转换为YOLOv5数据格式。

3. 模型训练

1. 检查并配置configs/data.yaml，替换为自己的数据集路径和类别数量。
2. 选择骨干网络并配置configs/model_*.yaml，修改类别数量。
3. 运行训练脚本： shell python scripts/train.py --batch 16 --epochs 5 --data configs/data.yaml --cfg configs/model_XXX.yaml

4. 模型测试与可视化

暂未详细说明，可根据实际情况进行相关操作。

5. 模型部署

1. 导出ONNX模型： shell python scripts/deploy/export.py --model_path runs/train/exp1/weights/best.pt --output_path best.onnx
2. 使用TensorRT进行量化： shell python scripts/trt_quant/convert_trt_quant.py --img_dir /XXXX/train/ --img_size 640 --batch_size 6 --batch 200 --onnx_model runs/train/exp1/weights/best.onnx --mode int8
3. 运行TensorRT推理脚本： shell python scripts/trt_quant/trt_infer.py --model_path best_int8.trt --img_path test.jpg

注意事项

• CBAM、SE、DCN、coord conv目前未添加到所有网络中，可能需用户自行修改代码。
• 其他骨干网络的PAN和FPN默认gw和gd与YOLOv5_L相同，若需要更小更快的模型，需修改FPN和PAN中的self.gw和self.gd。

参考

• Mobilenet v3
• resnet
• hrnet
• shufflenet

下载地址

点击下载 【提取码: 4003】