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
参考文章
- 教程：超详细从零开始yolov5模型训练
- YOLOv8问题分析和解决