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

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
相关推荐
Suyuoa3 小时前
附录2-pytorch yolov5目标检测
python·深度学习·yolo
红色的山茶花21 小时前
YOLOv8-ultralytics-8.2.103部分代码阅读笔记-block.py
笔记·深度学习·yolo
unix2linux1 天前
YOLO v5 Series - Image & Video Storage ( Openresty + Lua)
yolo·lua·openresty
菠菠萝宝1 天前
【YOLOv8】安卓端部署-1-项目介绍
android·java·c++·yolo·目标检测·目标跟踪·kotlin
ZZZZ_Y_1 天前
YOLOv5指定标签框背景颜色和标签字
yolo
红色的山茶花2 天前
YOLOv8-ultralytics-8.2.103部分代码阅读笔记-conv.py
笔记·yolo
Eric.Lee20212 天前
数据集-目标检测系列- 花卉 鸡蛋花 检测数据集 frangipani >> DataBall
人工智能·python·yolo·目标检测·计算机视觉·鸡蛋花检查
阿_旭2 天前
【模型级联】YOLO-World与SAM2通过文本实现指定目标的零样本分割
yolo·yolo-world·sam2
CSBLOG2 天前
OpenCV、YOLO、VOC、COCO之间的关系和区别
人工智能·opencv·yolo
2zcode2 天前
基于YOLOv8深度学习的医学影像骨折检测诊断系统研究与实现(PyQt5界面+数据集+训练代码)
人工智能·深度学习·yolo