基于yolov3的知识蒸馏+剪枝+剪植1:coco数据集转yolo数据集

coco数据集2014下载地址:

链接: https://pan.baidu.com/s/1-oYvt2vpR_vzgvzGd7xmig 密码: r3bt

按照这个可以看到最终需要什么样子的格式,主要为:
基于yolov3的知识蒸馏+剪枝+剪植1:coco数据集转yolo数据集_第1张图片第一个是目标矩形框的类别标签外框坐标+长宽
第二个图片(train+test)
第三和四是图片检索列表

我这边现有的数据为 json格式的,因为原先的检测是centernet网络
基于yolov3的知识蒸馏+剪枝+剪植1:coco数据集转yolo数据集_第2张图片
里面的格式看这里:COCO数据集的标注格式

转换的代码如下:

#python
import json
from collections import defaultdict  #这个模块实现了特定目标的容器,以提供Python标准内建容器 dict、list、set、tuple 的替代选择
from pycocotools.coco import COCO
import numpy as np
import tqdm
import argparse

"""hyper parameters"""
json_file_path = '/home/psdz/桌面/YOLO/data/cocodj/annotations/instances_trainval2014.json'
images_dir_path = 'data/cocodj/images/train/'
output_file = '/home/psdz/桌面/YOLO/data/cocodj/train.txt'
name1='data/cocodj/images/train/'

"""load json file"""
name_box_id = defaultdict(list)   #defaultdict:字典的子类,提供了一个工厂函数,为字典查询提供了默认值
id_name = dict()   #Python内置了字典:dict的支持,dict全称dictionary,在其他语言中也称为map,字典是另 一种可变容器模型,且可存储任意类型对象。具有极快的查找速度。


def arg_parser():
    parser = argparse.ArgumentParser('code by rbj')
    parser.add_argument('--annotation_path', type=str,
                        default=json_file_path)
    #生成的txt文件保存的目录
    parser.add_argument('--save_base_path', type=str, default='/home/psdz/桌面/YOLO/data/cocodj/labels/train/')
    args = parser.parse_args()
    return args

args = arg_parser()
annotation_path = args.annotation_path
save_base_path = args.save_base_path
data_source = COCO(annotation_file=annotation_path)
catIds = data_source.getCatIds()
categories = data_source.loadCats(catIds)
categories.sort(key=lambda x: x['id'])
classes = {}
coco_labels = {}
coco_labels_inverse = {}

for c in categories:
    coco_labels[len(classes)] = c['id']
    coco_labels_inverse[c['id']] = len(classes)
    classes[c['name']] = len(classes)

img_ids = data_source.getImgIds()
for index, img_id in tqdm.tqdm(enumerate(img_ids), desc='change .json file to .txt file'):
    img_info = data_source.loadImgs(img_id)[0]
    file_name = img_info['file_name'].split('.')[0]
    height = img_info['height']
    width = img_info['width']

    save_path = save_base_path + file_name + '.txt'
    with open(save_path, mode='w') as fp:
        annotation_id = data_source.getAnnIds(img_id)
        boxes = np.zeros((0, 5))
        if len(annotation_id) == 0:
            fp.write('')
            continue
        annotations = data_source.loadAnns(annotation_id)
        lines = ''
        for annotation in annotations:
            box = annotation['bbox']
            # some annotations have basically no width / height, skip them
            if box[2] < 1 or box[3] < 1:
                continue
            #top_x,top_y,width,height---->cen_x,cen_y,width,height
            box[0] = round((box[0] + box[2] / 2) / width, 6)
            box[1] = round((box[1] + box[3] / 2) / height, 6)
            box[2] = round(box[2] / width, 6)
            box[3] = round(box[3] / height, 6)
            label = coco_labels_inverse[annotation['category_id']]
            lines = lines + str(label)
            for i in box:
                lines += ' ' + str(i)
            lines += '\n'
        fp.writelines(lines)
print('finish')

with open(json_file_path, encoding='utf-8') as f:
    data = json.load(f)             #https://blog.csdn.net/fireflychh/article/details/83040205
    annotations = data['annotations']   #选择这个词典,具体类别看上面的博客就可以了

data_source = COCO(annotation_file=json_file_path)   #读取的api initialize COCO api for instance annotations
for ant in annotations:
    id = ant['image_id']
    img_info = data_source.loadImgs(id)[0]
    file_name = img_info['file_name'].split('.')[0]
    name = name1+file_name+'.jpg'
    cat = ant['category_id']

    if cat >= 1 and cat <= 11:
        cat = cat - 1

    name_box_id[name].append([ant['bbox'], cat])
"""write to txt"""
with open(output_file, 'w') as f:
    for key in name_box_id.keys():
        f.write(key)
        f.write('\n')

可以看出里面有两个生成地址,因为我整合了一下大佬们的转换代码,一次性生成上面说的两种txt文件。
基于yolov3的知识蒸馏+剪枝+剪植1:coco数据集转yolo数据集_第3张图片这个是检索的txt
基于yolov3的知识蒸馏+剪枝+剪植1:coco数据集转yolo数据集_第4张图片这个是目标标记的txt

部分参考链接:

voc定位数据xml转coco数据集格式json

COCO数据集训练格式转换成YOLO格式

你可能感兴趣的:(目标检测,机器学习,深度学习,pytorch,神经网络)