from data.data_reader import BIZCARD_LABEL_MAP, BizcardDataParser
import argparse
from pathlib import Path
import os
import json
import cv2
import numpy as np


def convert_bizcard_to_coco_format(image_dir, json_dir, id_list, out_dir, out_name):
    coco_json = {}
    images = []
    annotations = []
    categories = []

    for _, key in enumerate(BIZCARD_LABEL_MAP.keys()):
        categories.append({
            'id': BIZCARD_LABEL_MAP[key],
            'name': key
        })

    with open(id_list) as fp:
        ids = fp.readlines()

    for idx, file_id in enumerate(ids):
        file_id = Path(file_id.strip())
        print(idx, file_id)

        if not (image_dir / file_id).with_suffix('.jpg').exists():
            file_id = file_id.with_suffix('.jpeg')
        else:
            file_id = file_id.with_suffix('.jpg')

        height, width = cv2.imread(str(image_dir / file_id)).shape[:2]
        images.append({
            'file_name': str(file_id),
            'id': idx,
            'height': height,
            'width': width
        })

        try:
            gt = BizcardDataParser.parse_data(str((json_dir / file_id).with_suffix('.json')), str(image_dir / file_id))[
                0]
            for word in gt.words:
                anno = {
                    'id': len(annotations),
                    'image_id': idx,
                    'bbox': [word.bbox.min_x, word.bbox.min_y, (word.bbox.max_x - word.bbox.min_x),
                             (word.bbox.max_y - word.bbox.min_y)],
                    'segmentation': [word.bbox.val],
                    'category_id': word.label,
                    'iscrowd': 0,
                    'area': cv2.contourArea(np.reshape(word.bbox.val, [-1, 2]).astype(np.float32))
                }
                annotations.append(anno)
        except Exception as e:
            print(e)
            print(str(image_dir / file_id))

    coco_json['images'] = images
    coco_json['annotations'] = annotations
    coco_json['categories'] = categories
    with open(Path(out_dir, out_name), 'w') as f:
        json.dump(coco_json, f)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_dir', type=str)
    parser.add_argument('--gt_dir', type=str)
    parser.add_argument('--data_list', type=str)
    parser.add_argument('--out_dir', type=str)
    parser.add_argument('--out_name', type=str)
    args = parser.parse_args()

    if not Path(args.out_dir).exists():
        Path(args.out_dir).mkdir()

    convert_bizcard_to_coco_format(
        Path(args.img_dir),
        Path(args.gt_dir),
        args.data_list,
        args.out_dir,
        args.out_name)

[((69, 13, 69, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((16, 28, 16, 52), 'data.data_reader.BIZCARD_LABEL_MAP.keys', 'BIZCARD_LABEL_MAP.keys', ({}, {}), '()', False, 'from data.data_reader import BIZCARD_LABEL_MAP, BizcardDataParser\n'), ((65, 8, 65, 31), 'json.dump', 'json.dump', ({(65, 18, 65, 27): 'coco_json', (65, 29, 65, 30): 'f'}, {}), '(coco_json, f)', False, 'import json\n'), ((81, 8, 81, 26), 'pathlib.Path', 'Path', ({(81, 13, 81, 25): 'args.img_dir'}, {}), '(args.img_dir)', False, 'from pathlib import Path\n'), ((82, 8, 82, 25), 'pathlib.Path', 'Path', ({(82, 13, 82, 24): 'args.gt_dir'}, {}), '(args.gt_dir)', False, 'from pathlib import Path\n'), ((64, 14, 64, 37), 'pathlib.Path', 'Path', ({(64, 19, 64, 26): 'out_dir', (64, 28, 64, 36): 'out_name'}, {}), '(out_dir, out_name)', False, 'from pathlib import Path\n'), ((77, 11, 77, 29), 'pathlib.Path', 'Path', ({(77, 16, 77, 28): 'args.out_dir'}, {}), '(args.out_dir)', False, 'from pathlib import Path\n'), ((78, 8, 78, 26), 'pathlib.Path', 'Path', ({(78, 13, 78, 25): 'args.out_dir'}, {}), '(args.out_dir)', False, 'from pathlib import Path\n'), ((54, 44, 54, 78), 'numpy.reshape', 'np.reshape', ({(54, 55, 54, 68): 'word.bbox.val', (54, 70, 54, 77): '[-1, 2]'}, {}), '(word.bbox.val, [-1, 2])', True, 'import numpy as np\n')]