further augmentation added

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,14 +42,27 @@ 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