import json
import os
from pycocotools.coco import COCO
from tqdm import tqdm
import concurrent.futures
import re

if __name__ == '__main__':
    root_path = 'datasets/lvis'
    splits = ['train', 'val']
    for split in splits:
        print(f'Processing {split}.')
        annotation_path = os.path.join(root_path, f'lvis_v1_{split}.json')
        save_path = os.path.join(root_path, f'lvis_instance_psalm.json')

        with open(annotation_path) as f:
            data = json.load(f)
        categories = data['categories']
        categories_path = os.path.join(root_path, 'lvis_categories.json')
        with open(categories_path, 'w') as f:
            json.dump(categories, f, indent=2)

        annotation_map = {}
        for anno in data['annotations']:
            image_id = anno['image_id']
            if image_id not in annotation_map:
                annotation_map[image_id] = []
            annotation_map[image_id].append(anno)

        lvis_anno = []
        pattern = re.compile(r'.*/((?:train|val)\d+/\d+\.jpg)')
        new_img_id = 0
        for img in tqdm(data['images']):
            if img['id'] not in annotation_map:
                continue
            match = pattern.search(img['coco_url'])
            if match:
                image = match.group(1)
            else:
                image = img['coco_url']
            lvis_anno.append(
                {
                    'image': image,
                    'image_info': img,
                    'new_img_id': new_img_id,
                    'anns': annotation_map[img['id']]
                }
            )
            new_img_id += 1

        with open(save_path, 'w') as f:
            json.dump(lvis_anno, f, indent=2)
        print(f'saving at {save_path}.')