1. 拉取代码创建环境
- 执行
nvidia-smi
验证cuda环境是否可用; - 拉取官方代码;
- clone官方代码仓库
git clone https://github.com/WongKinYiu/yolov7
; - 从main分支切换到u7分支
cd yolov7 && git checkout 44f30af0daccb1a3baecc5d80eae22948516c579
;(YOLO_v5的所有视觉任务在同一个代码仓库中,YOLO_v7不同的分支对应不同的任务,分割任务对应u7
)
- clone官方代码仓库
- 安装环境
pip install --upgrade pip && pip install -r requirements.txt
;
2. 用预训练模型推理
- 下载与训练权重;
wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-seg.pt
;(更换代理会出现下载的权重只有2k,可以选择取u7
分支下手动下载) - 用
predict.py
推理;python segment/predict.py --weights <权重文件路径> --source <图片/视频路径> --name xxx
3. 准备数据集
labelimg/labelme/Roboflow,此处选用Roboflow;
- 创建项目
- 通过web控件和本地API上传图片;
- 标注
手动标注固然稳,使用基于SAM
的Smart Polygon
是真的香; - 生成数据集
创建数据集(不同版本) ~> 切分 ~> 前处理 ~> 增强
- 导出数据集
Roboflow提供了训练平台(只是免费的次数有限);
4. 训练
下面以实力分割为例,更换任务类型方法如下;
训练代码如下;
python segment/train.py --batch-size 16 --img-size 640 --epochs 10 --data {dataset.location}/data.yaml --weights $WEIGHTS_PATH --device 0 --name custom