diff --git a/dataloader/load.py b/dataloader/load.py index 9133d6d9175a23dd4e5e72f40c9b5988293d648c..6beb5474a72e58867d6611954f0ea4b10dc7b079 100644 --- a/dataloader/load.py +++ b/dataloader/load.py @@ -20,69 +20,53 @@ class UnconditionalDataset(Dataset): """ ### Create DataFrame - #file_list = [] - #for root, dirs, files in os.walk(fpath, topdown=False): - # for name in sorted(files): - # file_list.append(os.path.join(root, name)) - # - #df = pd.DataFrame({"Filepath":file_list},) - #self.df = df[df["Filepath"].str.endswith(ext)] + file_list = [] + for root, dirs, files in os.walk(fpath, topdown=False): + for name in sorted(files): + file_list.append(os.path.join(root, name)) + + df = pd.DataFrame({"Filepath":file_list},) + self.df = df[df["Filepath"].str.endswith(ext)] if skip_first_n: self.df = self.df[skip_first_n:] - print(fpath) - if train: - fpath = os.path.join(fpath, 'train') + if train: + df_train = self.df.sample(frac=frac,random_state=2) + self.df = df_train else: - fpath = os.path.join(fpath, 'valid') - - file_list =[] - for root, dirs, files in os.walk(fpath, topdown=False): - for name in sorted(files): - file_list.append(os.path.join(root, name)) - df = pd.DataFrame({"Filepath":file_list},) - self.df = df[df["Filepath"].str.endswith(ext)] + df_train = self.df.sample(frac=frac,random_state=2) + df_test = df.drop(df_train.index) + self.df = df_test if transform: - # for training 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(), + 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), - transforms.Normalize(mean=(0.5,0.5,0.5),std=(0.5,0.5,0.5))]) + transforms.CenterCrop(img_size),transforms.RandomHorizontalFlip(p=0.5)]) - transform_randomcrop = transforms.Compose([transforms.ToTensor(), - transforms.Resize(intermediate_size, antialias=True), - transforms.RandomCrop(img_size),transforms.RandomHorizontalFlip(p=0.5), - transforms.Normalize(mean=(0.5,0.5,0.5),std=(0.5,0.5,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 : - # for evaluation + else : self.transform = transforms.Compose([transforms.ToTensor(), - transforms.Lambda(lambda x: (x * 255).type(torch.uint8)), transforms.Resize(img_size)]) - - if train==False: - # for testing - self.transform = transforms.Compose([transforms.ToTensor(), - transforms.Resize(intermediate_size, antialias=True), - transforms.Normalize(mean=(0.5,0.5,0.5),std=(0.5,0.5,0.5))]) - + 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