import os
import json
import argparse
import numpy as np
import glob
import cv2
from sklearn.model_selection import train_test_split
from labelme import utils
from tqdm import tqdm
import shutil

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--anno_dirs', type=str, nargs="+", default=[])
    parser.add_argument('--data_dir', type=str, default='')
    args = parser.parse_args()
    labelme_folds = args.anno_dirs
    json_list_path = []

    for labelme_path in labelme_folds:
        list_path = glob.glob(labelme_path + "/*.json")
        json_list_path.extend(list_path)

    print('reading...')
    print("images: %d" % len(json_list_path))

    for file in tqdm(json_list_path):
        json_path = file
        temp_json = {}
        img_name = ''
        with open(json_path, 'r', encoding='utf-8') as jp:
            info = json.load(jp)
            info['imagePath'] = os.path.basename(info['imagePath'])
            img_name = info['imagePath']
            info['imageData'] = None
            temp_json = info
        jp.close()
        img_path = file.replace('json', str(img_name).split('.')[1])
        try:
            shutil.copy(img_path, args.data_dir)
            json.dump(temp_json,
                      open(os.path.join(args.data_dir, os.path.basename(json_path)), 'w', encoding='utf-8'),
                      ensure_ascii=False, indent=4)
        except Exception as e:
            print(e)
            print('Wrong Image:', img_name)
            continue