import copy
import json
import argparse
import numpy
import glob
import os
from tqdm import tqdm

def Labelme2Label_Studio(labelme_dirs, project_id):
    all_label_list = []
    label_dict = {}
    data_dict = {}
    annotations_dict = {}
    result_dict = {}
    vale_dict = {}
    labelme_path = []

    for labelme_dir in labelme_dirs:
        list_path = glob.glob(labelme_dir + "/*.json")
        labelme_path.extend(list_path)
    # print(label_path)
    # 进行必要信息的json转换
    for label in tqdm(labelme_path):
        with open(label, 'r', encoding='utf-8') as f:
            info = json.load(f)
            if not info['shapes']:
                continue
            annotations_list = []  # 会把之前的也添加进来，必须重置
            result_list = []  # 同上
            width = info['imageWidth']
            height = info['imageHeight']
            result_dict['original_width'] = copy.deepcopy(width)
            result_dict['original_height'] = copy.deepcopy(height)
            result_dict['image_rotation'] = 0
            result_dict['from_name'] = 'label'
            result_dict['to_name'] = 'image'
            result_dict['type'] = 'rectanglelabels'
            result_dict['origin'] = 'manual'
            # 标注信息写入
            for shape in info['shapes']:
                # print(shape)
                # print(shape['points'][0][0])
                vale_dict['x'] = shape['points'][0][0] * 100 / copy.deepcopy(width)
                vale_dict['y'] = shape['points'][0][1] * 100 / copy.deepcopy(height)
                vale_dict['width'] = (shape['points'][1][0] - shape['points'][0][0]) * 100 / copy.deepcopy(width)
                vale_dict['height'] = (shape['points'][1][1] - shape['points'][0][1]) * 100 / copy.deepcopy(height)
                vale_dict['rotation'] = 0
                vale_dict['rectanglelabels'] = [copy.deepcopy(shape['label'])]
                result_dict['value'] = copy.deepcopy(vale_dict)
                result_list.append(copy.deepcopy(result_dict))
            annotations_dict['result'] = copy.deepcopy(result_list)
            if 'updated_by' in info:
                annotations_dict['updated_by'] = info['updated_by']
            annotations_list.append(copy.deepcopy(annotations_dict))
            label_dict['annotations'] = copy.deepcopy(annotations_list)
            data_dict['image'] = '/data/upload/' + str(project_id) + '/' + os.path.basename(info['imagePath'])
            label_dict['data'] = copy.deepcopy(data_dict)
            if 'updated_by' in info:
                label_dict['updated_by'] = info['updated_by']
            all_label_list.append(copy.deepcopy(label_dict))
        f.close()

    # print(all_label_list)
    if not all_label_list:
        print('无有效json数据')
    else:
        json.dump(all_label_list, open('project' + str(project_id) + '.json', 'w', encoding='utf-8'),
              ensure_ascii=False, indent=2)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--labelme_dirs', type=str, nargs="+", default=[])
    parser.add_argument('--project_id', type=int)
    args = parser.parse_args()

    Labelme2Label_Studio(args.labelme_dirs, args.project_id)