Yolov5 AI学习笔记

Yolov5 AI学习笔记

环境准备

需要Python的开发环境，安装Anaconda。
Anaconda的一些命令：

# 创建虚拟环境
conda create -n yolo_cpu python=3.9
# 查看虚拟环境
conda env list
# 激活虚拟环境
conda activate <env_name>

Yolov5上手

下载源码

Yolov5源码：https://github.com/ultralytics/yolov5。

git clone https://github.com/ultralytics/yolov5

下载模型

使用已经训练好的模型：https://github.com/ultralytics/yolov5/releases。

例如yolov5s.pt，放置在yolov5目录。

安装依赖环境

在requirements.txt里面记录了yolo的依赖环境，可以使用pip安装。

# 使用国内源，不一定可以用
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
# 或者
pip install -r requirements.txt 

测试

原理其实就是是在电脑上用python加载训练好的模型，然后将输入送到模型中获取识别结果。

图片

目录下有提供的测试图片。

python detect.py --source ./data/images/bus.jpg

识别结果路径见打印。这里用的是CPU， 因为requirements.txt 依赖包里默认安装的 pytorch-cpu 版。

视频

在电脑上随便找了个mp4文件。

python detect.py --source data/mp4/GiantKillerPlants_60sec_5p1_AAC.mp4

使用GPU

确定Cuda版本

nvidia-smi

官网查看安装命令 https://pytorch.org/get-started/locally/

# 下载会有点慢
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 

装好之后可以找到GPU了，但是报错了。

conda list查看当前环境的包，发现torchvision不是cuda版本。

卸载后重装。

pip uninstall torchvision 
pip3 install torchvision --index-url https://download.pytorch.org/whl/cu121 

测试

python detect.py   --source ./data/images/bus.jpg

识别结果一致，耗时少了40ms。

训练模型

官方说明
https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data

labelImg标注工具安装

git clone https://gitee.com/ochip/labelImg.git

# conda环境
conda install pyqt=5
conda install -c anaconda lxml
pyrcc5 -o libs/resources.py resources.qrc

# 运行软件
python labelImg.py

软件预定义了一些类别，可以进行添加。

后面需要设置图片的文件夹和保存的文件夹，格式选择为yolo，可以在view勾选自动保存。

准备数据集

这里做一个识别用过的卫生纸的功能，把predefined_classes.txt里内容删除掉，将waste paper towels加入，然后用手机拍一些照片。

文件夹的名字需要是images。

使用labelImg进行标注，保存到labels文件夹，创建框选择类别即可。

创建yaml文件。

文件内容：

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../my_model/images/
val: ../my_model/images/
# number of classes
nc: 1

# class names
names: ['waste paper towels']

模型训练

在命令行里输入。

# 报没有内存
python train.py --img 640 --batch 50 --epochs 100 --data ../my_model/my_model.yaml --weights yolov5s.pt --nosave --cache
# 改成320正常运行
python train.py --img 320 --batch 50 --epochs 100 --data ../my_model/my_model.yaml --weights yolov5s.pt --nosave --cache

测试

直接用训练的图片进行测试。

python detect.py --weights runs/train/exp21/weights/best.pt --img 640 --conf 0.25 --source ..\my_model\images\2024_08_25_19_34_IMG_0249.JPG

参考文章

1. 教程：超详细从零开始yolov5模型训练
2. YOLOv8问题分析和解决