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
相关推荐
红色的山茶花1 天前
YOLOv6-4.0部分代码阅读笔记-engine.py
笔记·yolo
小哥谈2 天前
注意力机制篇 | YOLO11改进 | 即插即用的高效多尺度注意力模块EMA
人工智能·深度学习·神经网络·yolo·目标检测·机器学习·计算机视觉
静静AI学堂2 天前
Yolo11改进策略:上采样改进|CARAFE,轻量级上采样|即插即用|附改进方法+代码
人工智能·yolo·目标跟踪
小张贼嚣张2 天前
yolov8涨点系列之引入CBAM注意力机制
yolo·目标检测
辛勤的程序猿3 天前
YOLO即插即用---CPCA
yolo
王哈哈^_^3 天前
【数据集】【YOLO】【目标检测】树木倒塌识别数据集 9957 张,YOLO道路树木断裂识别算法实战训练教程!
人工智能·深度学习·算法·yolo·目标检测·计算机视觉·pyqt
深度学习lover3 天前
<项目代码>YOLOv7 草莓叶片病害识别<目标检测>
人工智能·yolo·目标检测·计算机视觉·草莓叶片病害识别
oceanao3 天前
半自动标注yolo训练数据并逆向为Labelimg标注文件
yolo
王哈哈^_^3 天前
【数据集】【YOLO】【目标检测】航拍船只识别数据集 3550 张,YOLO航拍水面船只识别算法实战训练教程!
人工智能·深度学习·算法·yolo·目标检测·计算机视觉·pyqt
王哈哈^_^4 天前
【数据集】【YOLO】【目标检测】摔跤识别数据集 5097 张,YOLO行人摔倒识别算法实战训练教程!
人工智能·深度学习·算法·yolo·目标检测·计算机视觉·pyqt