from dataLoader import *
import torchvision.transforms as tf
import SimpleITK as sitk
import os 


transform = tf.Compose([
    tf.ToTensor(),  # Convert image to tensor
])

mapping_files_bert = {
    # 'TotalSegmentor': '/home/data/Github/data/data_gen_def/DATASETS_processed/TotalSegmentorCT_MRI/nifti_mappings.json',
    # 'MSD': '/home/jachin/data/Github/data/data_gen_def/DATASETS_processed/MSD_processed/nifti_mappings_updated.json',
    'CancerImageArchive': '/home/data/Github/OmniMorph/Dataloader/nifty_mappings/CIA_mappings.json',
}

if __name__ == "__main__":
    # dataset = OminiDataset_v1(transform=None)
    # datasetp = OminiDataset_paired(transform=None)
    # dataset = OminiDataset_paired_inf(transform=None)
    # dataset = OminiDataset_inference_w_all(transform=None)
    # dataset = OminiDataset_bertembd(transform=None,mapping_files=mapping_files_bert)
    dataset = OminiDataset(transform=None)
    
    
    
    
    # print(dataset.get_keys_dist())
    # print(len(dataset))
    # print(dataset.build_batch().shape)
    # exit()
    dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

    for i, data in enumerate(dataloader):
        print(data[1])
        exit()

    
    # print(dataset.get_ALLdata())