
from model_zoo import get_model
from dataset_zoo import COCO_Order, Flickr30k_Order
import os
root_dir="/scratch/visual_gpt/"
COCO_ROOT = os.path.join(root_dir, 'coco2014')
FLICKR_ROOT = os.path.join(root_dir, 'flickr30k')
# splits = ['val', 'test']
# splits = [ 'test', 'val']
splits = ['val']
model, preprocess = get_model(model_name="openai-clip:ViT-B/32", device="cuda", root_dir=root_dir)

for split in splits:
    flickr_order_dataset = Flickr30k_Order(image_preprocess=preprocess, root_dir=FLICKR_ROOT, split=split, save_test_cases=True)
    coco_order_dataset = COCO_Order(image_preprocess=preprocess, download=True, root_dir=COCO_ROOT, split=split, save_test_cases=True) 