yolov5训练coco数据集的部分类别

在测试yolov5系列不同类别的模型在各种加速卡上的精度和性能时，我们希望得到一个准确的评估结果。因此，本文从一个COCO数据集中创建一个子集，该子集仅包含特定的类别。具体来说，它首先从源数据集中读取JSON文件，然后过滤出所需的类别，并将它们保存到新的JSON文件中。接下来，它将所需的图像和标签复制到新的目标目录中。最后，它创建一个包含所有图像文件路径的文本文件，并更新数据集的YAML配置文件。以此为数据集，训练并测试模型，从而得到准确的评估结果。

创建容器

bash 复制代码

mkdir yolov5
cd yolov5
docker run -it --gpus all --name yolov5_dev -v $PWD:/home/ cuda_dev_image:v1.0 bash

准备yolov5环境

bash 复制代码

apt update
apt install git -y
git clone https://github.com/ultralytics/yolov5
cd yolov5
bash data/scripts/get_coco.sh --train --val
cd ../datasets/coco/
wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip
rm -rf annotations
unzip annotations_trainval2017.zip

定义需要训练的类别(coco-6.yaml)

yaml 复制代码

path: /home/dataset/coco 
train: train2017.txt 
val: val2017.txt
names:
  0: person
  1: bicycle
  2: car
  3: motorcycle
  4: airplane
  5: bus

根据coco-6.yaml中保留的类别,生成新的数据集

python 复制代码

#create_sub_coco_dataset.py

import json
import yaml
import sys
import os
import shutil
import tqdm

def MakeDirs(dir):
	if not os.path.exists(dir):
		os.makedirs(dir,True)

def create_sub_coco_dataset(data_yaml="coco-6.yaml",
							src_root_dir="../datasets/coco",
							dst_root_dir="class6/coco",
							folder="val2017"
							):

	MakeDirs(dst_root_dir+"/annotations/")
	MakeDirs(dst_root_dir+"/images/"+folder)
	MakeDirs(dst_root_dir+"/labels/"+folder)

	keep_names=[x+1 for x in yaml.safe_load(open(data_yaml).read())['names'].keys()]
	all_annotations=json.loads(open(src_root_dir+"/annotations/instances_{}.json".format(folder)).read())
	keep_categories=[x for x in all_annotations["categories"] if x["id"] in keep_names]
	keep_annotations=[x for x in all_annotations['annotations'] if x['category_id'] in keep_names]

	all_annotations['annotations']=keep_annotations
	all_annotations["categories"]=keep_categories
	
	if not os.path.exists(dst_root_dir+"/annotations/instances_{}.json".format(folder)):
		with open(dst_root_dir+"/annotations/instances_{}.json".format(folder), "w") as f:
			json.dump(all_annotations, f)
		
	filelist=set()
	for i in tqdm.tqdm(keep_annotations):
		img_src_path="/images/{}/{:012d}.jpg".format(folder,i["image_id"])
		label_src_path="/labels/{}/{:012d}.txt".format(folder,i["image_id"])
		if not os.path.exists(dst_root_dir+img_src_path):
			shutil.copy(src_root_dir+img_src_path, dst_root_dir+img_src_path)
		if not os.path.exists(dst_root_dir+label_src_path):
			keep_records=[x for x in open(src_root_dir+label_src_path,"r").readlines() if (int(x.strip().split(" ")[0])+1) in keep_names]
			with open(dst_root_dir+label_src_path,"w") as f:
				for r in keep_records:
					f.write(r)
		filelist.add("./images/{}/{:012d}.jpg\n".format(folder,i["image_id"]))
		
	with open(dst_root_dir+"/{}.txt".format(folder),"w") as f:
		for r in filelist:
			f.write(r)
			
	new_data_yaml=yaml.safe_load(open(data_yaml).read())
	new_data_yaml["path"]=dst_root_dir

	with open(dst_root_dir+"/coco.yaml", 'w') as f:
		f.write(yaml.dump(new_data_yaml, allow_unicode=True))

create_sub_coco_dataset(sys.argv[1],sys.argv[2],sys.argv[3],sys.argv[4])

生成新数据集

bash 复制代码

cd /home/yolov5
rm -rf class6
python create_sub_coco_dataset.py coco-6.yaml ../datasets/coco class6/coco train2017
python create_sub_coco_dataset.py coco-6.yaml ../datasets/coco class6/coco val2017

训练

bash 复制代码

python train.py --data class6/coco/coco.yaml \
				--weights '' --cfg models/yolov5m.yaml \
				--img 640 --workers 0 --device 0

测试

bash 复制代码

python val.py --weights best.pt --data class6/coco/coco.yaml \
					--img 640 --conf-thres 0.001 --iou-thres 0.6 \
					--workers 0 --device 0 --half --batch-size 1