Yolov5数据集转EfficientDet数据集格式

最新推荐文章于 2024-07-03 15:24:12 发布

Franklin0607

最新推荐文章于 2024-07-03 15:24:12 发布

阅读量495

点赞数

文章标签： python

本文链接：https://blog.csdn.net/weixin_46403336/article/details/115511577

版权

Yolov5数据集转COCO数据集格式

Yolo V5数据集结构

-- train
	-- images
	-- labels
-- valid
	-- images
	-- labels
data.yaml

data.yaml

train: ../../train/images
val: ../Mask_Wearing/valid/images

nc: 2
names: ['0', '1']

EfficientDet (coco) 数据集结构

-- annotations
	instances_valid.json
	instances_train.json
-- train
-- valid

instances_train.json

{    
	"info": {info}, 
	"licenses": [{license}],
	"categories":[{category}],
	"images": [{image}],
	"annotations": [{annotation}], 
}

Yolov5数据集转EfficientDet数据集格式

import os
import json
import cv2
import random
import time

# EfficientDet所需数据集最后储存位置
coco_format_save_path = './datasets/jam/annotations/instances_train.json'
# yolov5 类别文件，一行一个类
yolo_format_classes_path = './datasets/jam/jam.names'
# yolo格式的注释文件
yolo_format_annotation_path = './datasets/jam/train/labels/'
# 图片文件
yolo_image_path = './datasets/jam/train/images/'

#Yolo注释文件及图片文件名
annotations_list = os.listdir(yolo_format_annotation_path)
images_list = os.listdir(yolo_image_path)
with open(yolo_format_classes_path,'r') as f1:
    lines1 = f1.readlines()
categories = []
for j,label in enumerate(lines1):
    label = label.strip()
    categories.append({'id':(j+1),'name':label,'supercategory': 'None'})  # EfficientDet 的 Category_id 从1开始

write_json_context = dict()
write_json_context['info'] = {'description': '', 'url': '', 'version': '', 'year': 2021, 'contributor': '',
                              'date_created': '2021-04-02'}
write_json_context['licenses'] = [{'id': 1, 'name': None, 'url': None}]
write_json_context['categories'] = categories
write_json_context['images'] = []
write_json_context['annotations'] = []
num_bboxes = 0
for k in range(len(images_list)):
    img_context = {}
    img_path = yolo_image_path+images_list[k]
    img_name = os.path.basename(img_path)
    height,width = cv2.imread(img_path).shape[:2]
    img_context['file_name'] = img_name
    img_context['height'] = height
    img_context['width'] = width
    img_context['date_captured'] = '2021-04-02'
    img_context['id'] = k
    img_context['license'] = 1
    img_context['coco_url'] =''
    img_context['flickr_url'] = ''
    write_json_context['images'].append(img_context)    
    with open(yolo_format_annotation_path+annotations_list[k],'r') as f2:
        lines2 = f2.readlines()
    for i,line in enumerate(lines2):
        line = line.split(' ')        
        bbox_dict = {}
        class_id, x, y, w, h = int(line[0]),float(line[1]),float(line[2]),float(line[3]),float(line[4])
        xmin = int(x*width-w*width/2)
        ymin = int(y*height-h*height/2)
        xmax = int(x*width+w*width/2)
        ymax = int(y*height+h*height/2)
        bbox_dict['id'] = num_bboxes
        bbox_dict['image_id'] = k
        bbox_dict['category_id'] = class_id+1  # EfficientDet 的 Category_id 从1开始
        h,w = abs(ymax-ymin),abs(xmax-xmin)
        bbox_dict['area']  = h * w
        bbox_dict['bbox'] = [xmin,ymin,w,h]
        bbox_dict['segmentation'] = []
        bbox_dict['iscrowd'] = 0
        write_json_context['annotations'].append(bbox_dict)
        num_bboxes+=1
        
with open(coco_format_save_path,'w') as fw:
    json.dump(write_json_context,fw)