Spaces:
Sleeping
Sleeping
| import os | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import Dataset | |
| from torchvision.transforms import Compose, Resize, ToTensor | |
| import imageio | |
| from tqdm import tqdm | |
| class pix2pixDataset(Dataset): | |
| def __init__(self, dataset="maps", data_dir="/projects/ml4science/datasets_pix2pix/", split="train", normalize=True, transforms=None, preload=False, image_size=256, direction="BtoA"): | |
| self.datadir = os.path.join(data_dir, dataset) | |
| self.img_name_list_path = os.path.join(data_dir, dataset, split) | |
| if not os.path.exists(self.datadir): | |
| print(f'Dataset directory {self.datadir} does not exists') | |
| self.normalize=normalize | |
| self.image_name_list = os.listdir(self.img_name_list_path) | |
| self.preload = preload | |
| self.direction = direction | |
| if transforms is None: | |
| self.transforms = Compose([ | |
| ToTensor(), # Convert to torch tensor | |
| Resize((image_size, image_size), antialias=False), # Resize to 256x256 | |
| ]) | |
| else: | |
| self.transforms = transforms | |
| if self.preload: | |
| self.x_list, self.y_list= (), () | |
| for name in tqdm(self.image_name_list): | |
| x, y = self.load_every(name) | |
| self.x_list = self.x_list + (x,) | |
| self.y_list = self.y_list + (y,) | |
| self.x_list = torch.stack(self.x_list, 0) | |
| self.y_list = torch.stack(self.y_list, 0) | |
| print(f"{split} dataset preloaded!") | |
| def load_every(self, name): | |
| img_array = np.asarray(imageio.imread(os.path.join(self.img_name_list_path, name))) | |
| img_H, img_W = img_array.shape[0], img_array.shape[1] | |
| if self.normalize: | |
| img_array = self.normalize_fn(img_array) | |
| x_img, y_img = img_array[:,:img_W//2, :], img_array[:, img_W//2:, :] | |
| x_img, y_img = self.transforms(x_img), self.transforms(y_img) # Apply the resize transform | |
| return x_img.float(), y_img.float() | |
| def normalize_fn(self, x): | |
| return (x/255. -0.5)*2 | |
| def unnormalize_fn(self, x): | |
| return ((x/2 + 0.5) * 255).int().clamp(0, 255) #since these are images | |
| def __getitem__(self, index): # getitem should return x0, x1, y (where y is the class label for class conditional generation) | |
| class_cond = None | |
| if self.preload: | |
| x_img, y_img = self.x_list[index], self.y_list[index] | |
| else: | |
| name = self.image_name_list[index] | |
| x_img, y_img = self.load_every(name) | |
| # if self.direction == "BtoA": | |
| # return x_img, y_img, class_cond | |
| # elif self.direction == "AtoB": | |
| # return y_img, x_img, class_cond | |
| batch ={ | |
| "image1":x_img, | |
| "image2":y_img, | |
| } | |
| return batch | |
| def __len__(self): | |
| return len(self.image_name_list) | |
| class FishDataset(Dataset): | |
| def __init__(self, data_dir="/projects/ml4science/FishDiffusion/", split="train", normalize=True, transforms=None, preload=False, image_size=128): | |
| self.datadir = os.path.join(data_dir) | |
| self.img_name_list_path = os.path.join(data_dir, split) | |
| if not os.path.exists(self.datadir): | |
| print(f'Dataset directory {self.datadir} does not exists') | |
| self.normalize=normalize | |
| self.image_name_list = os.listdir(self.img_name_list_path) | |
| self.preload = preload | |
| if transforms is None: | |
| # self.transforms = Compose([ | |
| # ToTensor(), # Convert to torch tensor | |
| # Resize((image_size, image_size), antialias=False), # Resize to 256x256 | |
| # ]) | |
| self.transforms = Compose([ | |
| ToTensor(), # Convert to torch tensor | |
| ]) | |
| else: | |
| self.transforms = transforms | |
| if self.preload: | |
| self.x_list, self.y_list, self.class_id = (), (), [] | |
| for name in tqdm(self.image_name_list): | |
| x, y = self.load_every(name) | |
| cls_id = int(name.split("_")[-1][:-4]) | |
| self.x_list = self.x_list + (x,) | |
| self.y_list = self.y_list + (y,) | |
| self.class_id.append(cls_id) | |
| self.x_list = torch.stack(self.x_list, 0) | |
| self.y_list = torch.stack(self.y_list, 0) | |
| self.class_id = torch.tensor(self.class_id) | |
| print(f"{split} dataset preloaded!") | |
| def load_every(self, name): | |
| img_array = np.asarray(imageio.imread(os.path.join(self.img_name_list_path, name))) | |
| img_H, img_W = img_array.shape[0], img_array.shape[1] | |
| if self.normalize: | |
| img_array = self.normalize_fn(img_array) | |
| x_img, y_img = img_array[:,:img_W//2, :], img_array[:, img_W//2:, :] | |
| x_img, y_img = self.transforms(x_img), self.transforms(y_img) # Apply the resize transform | |
| return x_img.float(), y_img.float() | |
| def normalize_fn(self, x): | |
| return (x/255. -0.5)*2 | |
| def unnormalize_fn(self, x): | |
| return ((x/2 + 0.5) * 255).int().clamp(0, 255) #since these are images | |
| def __getitem__(self, index): # getitem should return x0, x1, y (where y is the class label for class conditional generation) | |
| if self.preload: | |
| x_img, y_img, class_id = self.x_list[index], self.y_list[index], self.class_id[index] | |
| else: | |
| name = self.image_name_list[index] | |
| class_id = torch.tensor(int(name.split("_")[-1][:-4])) | |
| x_img, y_img = self.load_every(name) | |
| return x_img, y_img, class_id | |
| def __len__(self): | |
| return len(self.image_name_list) |