eiseg 转 yolov5数据格式

Posted 2022-02-05 东东就是我

# 转换json为txt，并按规则计算坐标

import os
import json

import pandas as pd

def find_annotations(annotations,image_id):
    # d = k: v for k, v in annotations.items() if v > 0
    result=[]
    for annotation in annotations:
        if annotation['id']==image_id:
            result.append(annotation)

    return result



def jsontotxt(json_path):

    txt_path = './txts/'
    imginfo = json.load(open(json_path))
    for image in imginfo['images']:
        image_id=image['id']
        image_width=image['width']
        image_height=image['height']
        image_name=image['file_name']
        annotations=find_annotations(imginfo['annotations'],image_id)
        fn = txt_path + image_name.replace('.jpg', '.txt')
        file = open(fn, 'a')
        for ap in annotations:  # 将坐标换算成yolov5的需求格式：相对坐标
            center_x=ap['bbox'][0]
            center_y=ap['bbox'][1]
            w=ap['bbox'][2]
            h=ap['bbox'][3]
            x = center_x / image_width
            y = center_y / image_height
            w = w / image_width
            h = h / image_height
            x = round(x, 8)
            y = round(y, 8)
            w = round(w, 8)
            h = round(h, 8)
            file.write(str(ap['category_id']-1) + ' ')
            file.write(str(x) + ' ')
            file.write(str(y) + ' ')
            file.write(str(w) + ' ')
            file.write(str(h) + '\\n')
        file.close()


if __name__ == "__main__":
    json_path='coco.json'
    jsontotxt(json_path)