diff --git a/dataloader/load.py b/dataloader/load.py
index a57ea04a6e0c438f9ba654874a3952c81b63f781..93d033b4f1d6ff30ae0bc8abde1aebcaaa894c49 100644
--- a/dataloader/load.py
+++ b/dataloader/load.py
@@ -4,11 +4,11 @@ from torchvision import transforms
import os
from PIL import Image
import pandas as pd
+import numpy as np
class UnconditionalDataset(Dataset):
def __init__(self,fpath,img_size,train,frac =0.8,skip_first_n = 0,ext = ".png",transform=True ):
"""
- Customized to datasets where all images are within a folder and the filenames are sorted by likeliehood.(Landscape Dataset)
Args:
fpath (string): Path to the folder where images are stored
img_size (int): size of output image img_size=height=width
@@ -42,17 +42,30 @@ class UnconditionalDataset(Dataset):
self.df = df_test
if transform:
- self.transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=(0.5,0.5,0.5),std=(0.5,0.5,0.5)),
- transforms.Resize(img_size,antialias=True),transforms.RandomHorizontalFlip(p=0.5)])
+ intermediate_size = 150
+ theta = np.pi/4 -np.arccos(intermediate_size/(np.sqrt(2)*img_size)) #Check dataloading.ipynb in analysis-depot for more details
+
+ transform_rotate = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=(0.5,0.5,0.5),std=(0.5,0.5,0.5)),
+ transforms.Resize(intermediate_size,antialias=True),
+ transforms.RandomRotation(theta/np.pi*180,interpolation=transforms.InterpolationMode.BILINEAR),
+ transforms.CenterCrop(img_size),transforms.RandomHorizontalFlip(p=0.5)])
+
+ transform_randomcrop = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=(0.5,0.5,0.5),std=(0.5,0.5,0.5)),
+ transforms.Resize(intermediate_size),transforms.RandomCrop(img_size),transforms.RandomHorizontalFlip(p=0.5)])
+
+ self.transform = transforms.RandomChoice([transform_rotate,transform_randomcrop])
+ else :
+ self.transform = transforms.Compose([transforms.ToTensor(),
+ transforms.Resize(img_size)])
def __len__(self):
return len(self.df)
def __getitem__(self,idx):
- path = self.df.iloc[idx].Filepath
+ path = self.df.iloc[idx].Filepaths
img = Image.open(path)
return self.transform(img),0
-
+
def tensor2PIL(self,img):
back2pil = transforms.Compose([transforms.Normalize(mean=(-1,-1,-1),std=(2,2,2)),transforms.ToPILImage()])
return back2pil(img)
\ No newline at end of file