import json
import os
import numpy as np
import random

ann_file = './Datasets/VG/vg_category_200_labels_index.json'
json_data = json.load(open(ann_file, "r"))
print(len(json_data))

data_dir = "../Datasets/VG/"
data = []

for item in json_data.items():
    key = item[0]
    label = item[1]
    
    if len(label) == 0:
        print('no label')
        print(key, label)
        continue
    
    if max(label) >= 200:
        print('More than 200 labels:')
        print(key, label)
        continue
    
    img_path = os.path.join(os.path.join(data_dir, 'VG_100K/'), key)
    target = np.zeros((200, ), dtype=np.int)
    target[label] = 1
    data.append({"target": target.tolist(), "img_path": img_path})

json.dump(data, open(os.path.join(data_dir, "VG200.json"), "w"))

num_test = 10000
seed = 1
np.random.seed(seed)  # Numpy module.
random.seed(seed)  # Python random module.

test_idx = np.random.choice(len(data), num_test, replace=False)
test_idx = np.sort(test_idx)
train_data, test_data = [], []
cur = 0
for i in range(len(data)):
    if cur < num_test and i == test_idx[cur]:
        test_data.append(data[i])
        cur += 1
    else:
        train_data.append(data[i])
print(len(data), len(train_data), len(test_data))

json.dump(train_data, open(os.path.join(data_dir, "train_vg200.json"), "w"))
json.dump(test_data, open(os.path.join(data_dir, "test_vg200.json"), "w"))
    