from cProfile import label
import os
import json
from PIL import Image
from tqdm import tqdm

import sys
sys.path.append("..")

from data_generate.split_generator import SplitGenerator

if __name__ == "__main__":
    # load the config file
    config_name = 'prompt.json'
    jsonfile = open(os.path.join('../config/5-shot', config_name))
    config = json.loads(jsonfile.read())

    dest_dir = os.path.join(os.path.join(config['data_dir'], '5-shot'), 'vggflowers')
    label_dir = os.path.join(os.path.join(os.path.join(config['data_dir'], 'raw'), 'vggflowers_raw'), "variant.txt")
    raw_data_dir = os.path.join(os.path.join(os.path.join(config['data_dir'], 'raw'), 'vggflowers_raw'), 'data')

    # organise images into class folders, since it has been done
    if False:
        prepare_images_and_classes(
            image_folder_dir=raw_data_dir,
            label_only_file_dir=os.path.join(label_dir, 'variants.txt'),
            image_label_file_dir_list = [
                os.path.join(label_dir, 'images_variant_trainval.txt'), os.path.join(label_dir, 'images_variant_test.txt')
            ],
            crop_bottom_pixel=20
        )

    #* split dataset
    split_aircraft = SplitGenerator(data_dir=raw_data_dir, dest_dir=dest_dir, split_dir=None)
    split_aircraft.split_train_val_test(nclass_train=64, nclass_val=16, nclass_test=20, save_split_npy=True)