clearml dataset parser

parser.py

#!/usr/bin/env python

import argparse
import json
import sys

def parse_args():
    parser = argparse.ArgumentParser(description='ClearML dataset parser')
    parser.add_argument('--file', '-f', type=str, required=True, help='dataset file name')
    parser.add_argument('--output', '-o', type=str, required=True, help='output file name')
    return parser.parse_args()

def parse_json(filename: str) -> tuple[list[str], list[list[list[float]]]]:
    with open(filename, 'r') as f:
        raw_data = json.load(f)
    
    image_list: list[str] = []
    rois_list: list[list[list[float]]] = []
    for r in raw_data:
        image_filename = r['sources'][0]['uri'].split('/')[-1]
        image_list.append(image_filename)
        if len(r['rois']) == 0:
            continue
        
        rois: list[list[float]] = []
        for roi in r['rois']:
            poly: list[float] = roi['poly']
            rois.append(poly)

        rois_list.append(rois)
    return (image_list, rois_list)

def write_file(filename: str, data: tuple[list[str], list[list[list[float]]]]):
    with open(filename, 'w') as f:
        for (filename, rois) in zip(data[0], data[1]):
            if len(rois) == 0:
                continue
            for roi in rois: # [x1,y1,x2,y2,x3,y3,x4,y4]
                x, y, width, height = roi[0], roi[1], roi[2] - roi[0], roi[7] - roi[1]
                line = f'{filename}\t{x}\t{y}\t{width}\t{height}\n'
                f.write(line)



if __name__ == '__main__':
    args = parse_args()
    print(args)
    data = parse_json(args.file)
    write_file(args.output, data)