目录
前言
代码:https://github.com/dog-qiuqiu/FastestDet
算法介绍:https://zhuanlan.zhihu.com/p/536500269
训练环境:Win11
一、数据集构造
数据集格式与Darknet Yolo相同,每张图片对应一个txt标签文件。标签格式也是基于Darknet Yolo的数据集标签格式:“category cx cy wh”,其中category为类别下标,cx, cy为归一化标签框中心点的坐标,w, h为归一化标签框的宽度和高度,.txt标签文件内容示例如下:
2 0.2841796875 0.3815104166666667 0.068359375 0.257825 1 0.43505859375 0.55078125 0.0966796875 0.122395833333333
数据集文件结构如下:
其中train.txt和val.txt的格式为:
E:\cv_code\FastestDet-main\datasets\images\train\aug_0_10.jpg E:\cv_code\FastestDet-main\datasets\images\train\aug_0_100.jpg E:\cv_code\FastestDet-main\datasets\images\train\aug_0_101.jpg
生成train.txt和val.txt的代码示例:
import os
import glob
# Define the directory containing the dataset
train_dir = r'E:\cv_code\FastestDet-main\datasets\images\train'
val_dir = r'E:\cv_code\FastestDet-main\datasets\images\val'
# Get a list of all file paths in the dataset directory
train_files = glob.glob(os.path.join(train_dir, '*.jpg'))
val_files = glob.glob(os.path.join(val_dir, '*.jpg'))
# Write the file paths to a .txt file
with open('datasets/train.txt', 'w') as file:
for path in train_files:
file.write(path + '\n')
with open('datasets/val.txt', 'w') as file:
for path in val_files:
file.write(path + '\n')
二、配置文件设置
1..yaml配置文件介绍
以./configs/coco.yaml为例:
DATASET: TRAIN: 'E:\cv_code\FastestDet-main\datasets\train.txt' # train.txt VAL: 'E:\cv_code\FastestDet-main\datasets\val.txt' # val.txt NAMES: "configs/coco.names" # label name MODEL: NC: 4 # number of classes INPUT_WIDTH: 352 # input width INPUT_HEIGHT: 352 # input height TRAIN: LR: 0.001 # learning rate THRESH: 0.25 # threshold WARMUP: true # warmup BATCH_SIZE: 96 # batch size END_EPOCH: 10 # train epoch MILESTIONES: # Declining learning rate steps - 100 - 200 - 250
其中coco.names为你的数据集标签名
person bicycle car motorbike ....... toothbrush
2.开始训练
预训练模型为:shufflenetv2.pth
指定你的配置文件:
parser.add_argument('--yaml', type=str, default='configs/test.yaml', help='.yaml config')
开始训练:
python3 train.py
3.模型评估
python3 eval.py --yaml configs/coco.yaml --weight weights/weight_AP05:0.253207_280-epoch.pth
4.导出onnx进行测试
python3 test.py --yaml configs/coco.yaml --weight weights/weight_AP05:0.253207_280-epoch.pth --img data/3.jpg --onnx