import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch

tr_tf = transforms.Compose([
                    transforms.RandomHorizontalFlip(),
                    transforms.Resize((224,224)),
                    transforms.CenterCrop((224,224)),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

te_tf = transforms.Compose([
                    transforms.Resize((224,224)),
                    transforms.ToTensor(),
                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])




tr_dataset = dataset.CocoCaptions(root = 'COCO_COCO_2014_Train_Images/train2014',
                        annFile = 'COCO_COCO_2014_Train_Val_annotations/annotations/captions_train2014.json',
                        transform=tr_tf)

te_dataset = dataset.CocoCaptions(root = 'COCO_COCO_2014_Val_Images/val2014',
                        annFile = 'COCO_COCO_2014_Train_Val_annotations/annotations/captions_val2014.json',
                        transform=te_tf)