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