import argparse
import json
import os

from tqdm import tqdm


def work(fold_path, label):
    fold_list = os.listdir(fold_path)
    for fold in fold_list:
        json_list = os.listdir(fold_path + '/' + fold + '/images')
        print(fold)
        for json_path in tqdm(json_list):
            if json_path.endswith('.json'):
                file_name = fold_path + '/' + fold + '/images/' + json_path
                with open(str(file_name), 'r', encoding='utf-8') as fp:
                    json_data = json.load(fp)
                    shape_list = []
                    for shape in json_data['shapes']:
                        if shape['label'] != label:
                            shape_list.append(shape)
                    fp.close()
                    if len(shape_list) == 0:
                        # print('remove', file_name)
                        os.remove(file_name)
                        os.remove(str(file_name).replace('json', 'jpg'))
                    else:
                        os.remove(file_name)
                        json_data['shapes'] = shape_list
                        # print(file_name)
                        with open(str(file_name), 'w', encoding='utf-8') as fp1:
                            json.dump(json_data, fp1, sort_keys=False, indent=2)
                            fp1.close()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--fold', type=str, default='./fold')
    parser.add_argument('--label', type=str, default='Table')
    args = parser.parse_args()
    work(args.fold, args.label)