使用的版本为:yolov5-6.1
1、准备数据集
1.1 图像标注
使用labelimg或者CVAT等标注工具对图像样本进行逐一标注,标注格式选用YOLO(减少不必要的转格式麻烦),每一个图像文件xxxx.jpg标注完成后,都会生成对应的label文件xxxx.txt。labelimg标注示意:
1.2 数据集目录结构
最终的数据集目录结构采取如下方式:
--cat-dog-dataset
--images
--train
cat_1.jpg
cat_3.jpg
...
--val
cat_2.jpg
cat_19.jpg
...
--test
cat_13.jpg
cat_34.jpg
...
--labels
--train
cat_1.txt
cat_3.txt
...
--val
cat_2.txt
cat_19.txt
...
--test
cat_13.txt
cat_34.txt
...
数据集划分参考代码如下:
import os
import shutil
import random
random.seed(0)
def split_data(img_path, txt_path, new_file_path, train_rate, val_rate, test_rate):
directory = ["/images", "/labels", "/train", "/val", "/test"]
images_dir = new_file_path + directory[0]
labels_dir = new_file_path + directory[1]
if not os.path.exists(images_dir):
os.makedirs(images_dir)
if not os.path.exists(labels_dir):
os.makedirs(labels_dir)
for i in range(2, 5):
if not os.path.exists(images_dir + directory[i]):
os.makedirs(images_dir + directory[i])
if not os.path.exists(labels_dir + directory[i]):
os.makedirs(labels_dir + directory[i])
all_images = os.listdir(img_path)
all_labels = os.listdir(txt_path)
data = list(zip(all_images, all_labels))
random.shuffle(data)
shuffle_images, shuffle_labels = zip(*data)
train_images = shuffle_images[0: int(len(all_images) * train_rate)]
train_labels = shuffle_labels[0: int(len(all_labels) * train_rate)]
val_images = shuffle_images[int(len(all_images) * train_rate): int(len(all_images) * (train_rate + val_rate))]
val_labels = shuffle_labels[int(len(all_labels) * train_rate): int(len(all_labels) * (train_rate + val_rate))]
test_images = shuffle_images[int(len(all_images) * (train_rate + val_rate)):]
test_labels = shuffle_labels[int(len(all_labels) * (train_rate + val_rate)):]
for file in train_images:
file_prefix = file.split(".")[0]
origin_img_path = img_path + "/" + file_prefix + ".jpg"
new_img_path = images_dir + directory[2] + "/" + file_prefix + ".jpg"
origin_label_path = txt_path + "/" + file_prefix + ".txt"
new_label_path = labels_dir + directory[2] + "/" + file_prefix + ".txt"
shutil.copy(origin_img_path, new_img_path)
shutil.copy(origin_label_path, new_label_path)
for file in val_images:
file_prefix = file.split(".")[0]
origin_img_path = img_path + "/" + file_prefix + ".jpg"
new_img_path = images_dir + directory[3] + "/" + file_prefix + ".jpg"
origin_label_path = txt_path + "/" + file_prefix + ".txt"
new_label_path = labels_dir + directory[3] + "/" + file_prefix + ".txt"
shutil.copy(origin_img_path, new_img_path)
shutil.copy(origin_label_path, new_label_path)
for file in test_images:
file_prefix = file.split(".")[0]
origin_img_path = img_path + "/" + file_prefix + ".jpg"
new_img_path = images_dir + directory[4] + "/" + file_prefix + ".jpg"
origin_label_path = txt_path + "/" + file_prefix + ".txt"
new_label_path = labels_dir + directory[4] + "/" + file_prefix + ".txt"
shutil.copy(origin_img_path, new_img_path)
shutil.copy(origin_label_path, new_label_path)
if __name__ == '__main__':
img_path = "/home/temp/images"
txt_path = "/home/temp/labels"
new_file_path = "/home/dataset/cat-dog-dataset"
split_data(img_path, txt_path, new_file_path, 0.8, 0.1, 0.1)
以上代码实现了将数据集按照8:1:1的比例划分为训练集、验证集和测试集。
2、训练
2.1 创建数据集配置文件
yolov5-6.1源代码data目录下创建文件mydataset.yaml,并将data目录下的coco128.yaml文件内容粘贴进去,然后修改相应内容,如下:
# mydataset.yaml
path: /home/dataset/cat-dog-dataset # dataset root dir
train: images/train # train images(relative to 'path')
val: images/val # val images(relative to 'path')
test: # test images (optional)
# Classes
nc: 2 # number of classes
names: ['cat', 'dog'] # class names
2.2 修改模型配置文件
在models目录下选择一个yolov5模型,本人选择的是yolov5m,将其中的检测类别数nc修改为正确数值:
# yolov5m.yaml
nc: 80 # number of classes
depth_multiple: 0.67 # model depth multiple
width_multiple: 0.75 # layer channel multiple
2.3 下载预训练权重
yolov5m6.pt的预训练权重下载地址:yolov5m6_local.pt,提取密码:ge2y
将其下载后,放入yolov5-6.1目录下。
2.4 执行命令开始训练
python train.py --cfg models/yolov5m.yaml --data data/mydataset.yaml --weights yolov5m6_local.pt --epochs 100 --device 0,1,2
3、测试
python detect.py --weights runs/train/exp/weights/best.pt --source data/images --data data/mydataset.yaml --device 0