目标检测和实例分割:COCO格式json直接转换成yolo格式txt

如果想使用yolo训练类似coco标注格式的数据集需要进行格式转换,将json格式转换为yolo的txt格式,下面是目标检测的

import os
import json
from tqdm import tqdm
import argparse

parser = argparse.ArgumentParser()
#这里根据自己的json文件位置，换成自己的就行
parser.add_argument('--json_path', default='/home/admin1/data/LVIS/lvis_v1_val.json',type=str, help="input: coco format(json)")
#这里设置.txt文件保存位置
parser.add_argument('--save_path', default='/home/admin1/data/LVIS/alltxt', type=str, help="specify where to save the output dir of labels")
arg = parser.parse_args()

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = box[0] + box[2] / 2.0
    y = box[1] + box[3] / 2.0
    w = box[2]
    h = box[3]
#round函数确定(xmin, ymin, xmax, ymax)的小数位数
    x = round(x * dw, 6)
    w = round(w * dw, 6)
    y = round(y * dh, 6)
    h = round(h * dh, 6)
    return (x, y, w, h)

if __name__ == '__main__':
    json_file =   arg.json_path # COCO Object Instance 类型的标注
    ana_txt_save_path = arg.save_path  # 保存的路径

    data = json.load(open(json_file, 'r'))
    if not os.path.exists(ana_txt_save_path):
        os.makedirs(ana_txt_save_path)

    id_map = {} # coco数据集的id不连续！重新映射一下再输出！
    with open(os.path.join(ana_txt_save_path, 'classes.txt'), 'w') as f:
        # 写入classes.txt
        for i, category in enumerate(data['categories']):
            f.write(f"{category['name']}\n")
            id_map[category['id']] = i
    # print(id_map)
    #这里需要根据自己的需要，更改写入图像相对路径的文件位置。
    list_file = open(os.path.join(ana_txt_save_path, 'train2017.txt'), 'w')
    for img in tqdm(data['images']):
        img_file_name = img['coco_url'] # 图片名称
        filename = img_file_name.split("/")[-1]
        img_width = img["width"]
        img_height = img["height"]
        img_id = img["id"]
        head, tail = os.path.splitext(filename)
        ana_txt_name = head + ".txt"  # 对应的txt名字，与jpg一致
        f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w')
        for ann in data['annotations']:
            # print(ann)
            if ann['image_id'] == img_id:
                # box = convert((img_width, img_height), ann["bbox"])
                segmentation = ann["segmentation"][0]  # 获取ann["segmentation"][0]

                for i in range(len(segmentation)):
                    if i % 2 == 0:  # 偶数项
                        segmentation[i] /= img_width
                    else:  # 奇数项
                        segmentation[i] /= img_height

                f_txt.write("%s %s\n" % (id_map[ann["category_id"]], ' '.join(map(str, segmentation))))
        f_txt.close()
        #将图片的相对路径写入train2017或val2017的路径
        # list_file.write('./images/train2017/%s.jpg\n' %(head))
    # list_file.close()

这是实例分割的格式转换

import os
import json
import shutil

def write_yolo_txt_file(txt_file_path,label_seg_x_y_list):
    if not os.path.exists(txt_file_path):
        with open(txt_file_path, "w") as file:
            for element in label_seg_x_y_list:
                file.write(str(element) + " ")
            file.write('\n')
    else:
        with open(txt_file_path, "a") as file:
            for element in label_seg_x_y_list:
                file.write(str(element) + " ")
            file.write('\n')

def read_json(in_json_path,img_dir,target_dir):
    with open(in_json_path, "r", encoding='utf-8') as f:
        # json.load数据到变量json_data
        json_data = json.load(f)

    # print(len(json_data['annotations']))
    # print(len(json_data['images']))
    # print(len(json_data['categories']))

    for annotation in json_data['annotations']: # 遍历标注数据信息
        # print(annotation)
        category_id = annotation['category_id']
        image_id = annotation['image_id']
        for image in json_data['images']: # 遍历图片相关信息
            if image['id'] == image_id:
                width = image['width'] # 图片宽
                height = image['height'] # 图片高
                img_file_name = image['coco_url'] # 图片名称
                img_file_name = img_file_name.split("/")[-1]
                # img_file_name=img_file_name[:-4]
                txt_file_name = img_file_name.split('.')[0] + '.txt' # 要保存的对应txt文件名
                break
        # print(width,height,img_file_name,txt_file_name)
        segmentation = annotation['segmentation'] # 图像分割点信息[[x1,y1,x2,y2,...,xn,yn]]
        seg_x_y_list = [i/width if num%2==0 else i/height for num,i in enumerate(segmentation[0])] # 归一化图像分割点信息
        label_seg_x_y_list = seg_x_y_list[:]
        label_seg_x_y_list.insert(0,category_id) # 图像类别与分割点信息[label,x1,y1,x2,y2,...,xn,yn]
        # print(label_seg_x_y_list)

        # 写txt文件
        txt_file_path = target_dir + txt_file_name
        # print(txt_file_path)
        write_yolo_txt_file(txt_file_path,label_seg_x_y_list)

        # 选出txt对应img文件
        img_file_path = img_dir + img_file_name
        # print(img_file_path)
        shutil.copy(img_file_path,target_dir)



if __name__=="__main__":
    img_dir = '/home/admin1/data/LVIS/val2017/'
    target_dir = '/home/admin1/data/LVIS/testset/'
    if not os.path.exists(target_dir):
        os.mkdir(target_dir)
    in_json_path = '/home/admin1/data/LVIS/lvis_v1_val.json'
    read_json(in_json_path,img_dir,target_dir)