def __init__(self, root, config, hr_shape): self.device = config.device self.image_paths = sorted(glob.glob(root + "/*.*")) self.image_size = config.image_size self.scale_factor = config.scale_factor hr_height, hr_width = hr_shape K, P = load_kernels(file_path='kernels/', scale_factor=self.scale_factor) #K = kernels -> K.shape = (15,15,1,358) #P = Matriz de projeçao do PCA --> P.shape = (15,225) self.randkern = Kernels(K, P)
def __init__(self, root, config=None): """Initialize image paths and preprocessing module.""" self.image_paths = [] for ext in TYPES: self.image_paths.extend(glob.glob(os.path.join(root, ext))) self.image_size = config.image_size self.scale_factor = config.scale_factor K, P = load_kernels(file_path='kernels/', scale_factor=self.scale_factor) self.randkern = Kernels(K, P)
def __init__(self, root, config=None): """Initialize image paths and preprocessing module.""" self.device = config.device self.image_paths = [] for ext in TYPES: self.image_paths.extend(glob.glob(os.path.join(root, ext))) self.image_size = config.image_size # LR image size self.scale_factor = config.scale_factor K, P = load_kernels(file_path='kernels/', scale_factor=self.scale_factor) #K = kernels -> K.shape = (15,15,1,358) #P = Matriz de projeçao do PCA --> P.shape = (15,225) self.randkern = Kernels(K, P)
class ImageFolder(data.Dataset): """Custom Dataset compatible with prebuilt DataLoader.""" def __init__(self, root, config=None): """Initialize image paths and preprocessing module.""" self.image_paths = [] for ext in TYPES: self.image_paths.extend(glob.glob(os.path.join(root, ext))) self.image_size = config.image_size self.scale_factor = config.scale_factor K, P = load_kernels(file_path='kernels/', scale_factor=self.scale_factor) self.randkern = Kernels(K, P) def __getitem__(self, index): """Read an image from a file and preprocesses it and returns.""" image_path = self.image_paths[index] image = Image.open(image_path).convert('RGB') # target (high-resolution image) transform = transforms.RandomCrop(self.image_size * self.scale_factor) hr_image = transform(image) # input (low-resolution image) transform = transforms.Compose([ transforms.Lambda(lambda x: self.randkern.RandomBlur(x)), transforms.Resize((self.image_size, self.image_size)), transforms.Lambda(lambda x: Scaling(x)), transforms.Lambda(lambda x: self.randkern.ConcatDegraInfo(x)) ]) lr_image = transform(hr_image) transform = transforms.ToTensor() lr_image, hr_image = transform(lr_image), transform(hr_image) return lr_image.to(torch.float64), hr_image.to(torch.float64) def __len__(self): """Return the total number of image files.""" return len(self.image_paths)
def test(self): #receives single image --> can be easily modified to handle multiple images 'Takes single LR image as input. Returns LR image + (models approx) HR image concatenated' 'image location must be given by flag --test_image_path' self.model.eval() step = self.start_step + 1 # if not loading trained start = 0 lr_image = Image.open(self.test_image_path) lr_image_size = lr_image.size[0] #CONSIDER RGB IMAGE from utils import Kernels, load_kernels K, P = load_kernels(file_path='kernels/', scale_factor=2) randkern = Kernels(K, P) # LR_image_scaled + LR_residual_scaled (CONCAT) ---> TO TORCH lr_image_scaled = Scaling(lr_image) lr_image_with_kernel = randkern.ConcatDegraInfo(lr_image_scaled) lr_image_with_kernel = torch.from_numpy(lr_image_with_kernel).float().to(self.device) # NUMPY to TORCH # LR_image to torch lr_image_scaled = torch.from_numpy(lr_image_scaled).float().to(self.device) # NUMPY to TORCH #Transpose - Permute since for model we need input with channels first lr_image_scaled = lr_image_scaled.permute(2,0,1) lr_image_with_kernel = lr_image_with_kernel.permute(2,0,1) lr_image_with_kernel = lr_image_with_kernel.unsqueeze(0) #just add one dimension (index on batch) lr_image_scaled = lr_image_scaled.unsqueeze(0) lr_image, x = lr_image_scaled.to(torch.float64), lr_image_with_kernel.to(torch.float64) lr_image, x = lr_image.to(self.device), x.to(self.device) x = x.to(torch.float64) reconst = self.model(x) tmp1 = lr_image.data.cpu().numpy().transpose(0,2,3,1)*255 image_list = [np.array(Image.fromarray(tmp1.astype(np.uint8)[i]).resize((128,128), Image.BICUBIC)) \ for i in range(self.data_loader.batch_size)] image_hr_bicubic= np.stack(image_list) image_hr_bicubic_single = np.squeeze(image_hr_bicubic) #return this ^ image_hr_bicubic = image_hr_bicubic.transpose(0,3,1,2) image_hr_bicubic = Scaling(image_hr_bicubic) image_hr_bicubic = torch.from_numpy(image_hr_bicubic).double().to(self.device) # NUMPY to TORCH hr_image_hat = reconst hr_image_hat_np = hr_image_hat.data.cpu().numpy() hr_image_hat_np_scaled = hr_image_hat_np #just to try different types of scaling. It already comes scaled hr_image_hat_np_scaled = np.squeeze(hr_image_hat_np_scaled).transpose((1, 2, 0)) hr_image_hat_np_png = (hr_image_hat_np_scaled*255).astype(np.uint8) #return this ^ #Saving Image Bicubic and HR Image Hat Image.fromarray(image_hr_bicubic_single).save('./results/HR_bicub_images/'+ os.path.basename(self.test_image_path)+'_hr_bic_{}.png'.format(step)) Image.fromarray(hr_image_hat_np_png).save('./results/HR_HAT_images/'+ os.path.basename(self.test_image_path)+'_hr_hat_{}.png'.format(step)) #Create Grid hr_image_hat_np_scaled = hr_image_hat_np #It's already scaled (comes out of model scaled) hr_image_hat_torch = torch.from_numpy(hr_image_hat_np_scaled).double().to(self.device) # NUMPY to TORCH pairs = torch.cat((image_hr_bicubic.data, \ hr_image_hat_torch.data), dim=3) grid = make_grid(pairs, 1) tmp = np.squeeze(grid.cpu().numpy().transpose((1, 2, 0))) tmp = (255 * tmp).astype(np.uint8) random_number = np.random.rand(1)[0] Image.fromarray(tmp).save('./results/grids/'+ os.path.basename(self.test_image_path).split('.')[0]+'_grid_{}.png'.format(step))
class ImageFolder(data.Dataset): def __init__(self, root, config=None): self.device = config.device self.image_paths = sorted(glob.glob(root + "/*.*")) self.image_size = config.image_size self.scale_factor = config.scale_factor hr_height, hr_width = config.image_size * config.scale_factor, config.image_size * config.scale_factor K, P = load_kernels(file_path='kernels/', scale_factor=self.scale_factor) #K = kernels -> K.shape = (15,15,1,358) #P = Matriz de projeçao do PCA --> P.shape = (15,225) self.randkern = Kernels(K, P) # Transforms for low resolution images and high resolution images def __getitem__(self, index): """Read an image from a file and preprocesses it and returns.""" image_path = self.image_paths[index % len(self.image_paths)] image = Image.open(image_path) if np.array(image).shape == 4: image = np.array(image)[:, :, :3] image = Image.fromarray(image) # target (high-resolution image) # Random Crop: #transform = transforms.RandomCrop(self.image_size * self.scale_factor) #hr_image = transform(image) # Face Crop: face_width = face_height = 128 j = (image.size[0] - face_width) // 2 i = (image.size[1] - face_height) // 2 image = image.crop([j, i, j + face_width, i + face_height]) hr_image = image # HR_image --> [0,1] --> torch hr_image_scaled = Scaling(hr_image) hr_image_scaled = torch.from_numpy(hr_image_scaled).float().to( self.device) # NUMPY to TORCH # get HR residuals --> [-1,1] --> torch transform_HR = transforms.Compose([ #random blur transforms.Lambda(lambda x: self.randkern.RandomBlur(x)), #downscale BICUBIC pro tamanho LR transforms.Resize((self.image_size, self.image_size), Image.BICUBIC), #upscale BICUBIC pro tamanho HR transforms.Resize((self.image_size * self.scale_factor, self.image_size * self.scale_factor), Image.BICUBIC) ]) hr_image_hat = transform_HR(hr_image) hr_residual = np.array(hr_image).astype(float) - np.array( hr_image_hat).astype(float) hr_residual_scaled = Scaling(hr_residual) hr_residual_scaled = torch.from_numpy(hr_residual_scaled).float().to( self.device) # NUMPY to TORCH # get LR_IMAGE --> [0,1] transform_to_lr = transforms.Compose([ transforms.Lambda(lambda x: self.randkern.RandomBlur(x)), transforms.Resize((self.image_size, self.image_size), Image.BICUBIC) ]) lr_image = transform_to_lr(hr_image) lr_image_scaled = Scaling(lr_image) # get LR_RESIDUAL --> [-1,1] transform_to_vlr_bicub = transforms.Compose([ transforms.Lambda( lambda x: self.randkern.RandomBlur(x)), #random blur transforms.Lambda(lambda x: bicub_downscale(x, self.scale_factor)), transforms.Resize((self.image_size, self.image_size), Image.BICUBIC) #upscale pro tamanho LR ]) transform_to_vlr_bilin = transforms.Compose([ transforms.Lambda( lambda x: self.randkern.RandomBlur(x)), #random blur transforms.Lambda(lambda x: bilin_downscale(x, self.scale_factor)), transforms.Resize((self.image_size, self.image_size), Image.BICUBIC) #upscale pro tamanho LR ]) transform_to_vlr_nears = transforms.Compose([ transforms.Lambda( lambda x: self.randkern.RandomBlur(x)), #random blur transforms.Lambda( lambda x: nearest_downscale(x, self.scale_factor)), transforms.Resize((self.image_size, self.image_size), Image.BICUBIC) #upscale pro tamanho LR ]) lr_image_hat_bicub = transform_to_vlr_bicub(lr_image) lr_residual_bicub = np.array(lr_image).astype( np.float32) - np.array(lr_image_hat_bicub).astype(np.float32) lr_residual_bicub_scaled = Scaling(lr_residual_bicub) lr_image_hat_bilin = transform_to_vlr_bilin(lr_image) lr_residual_bilin = np.array(lr_image).astype( np.float32) - np.array(lr_image_hat_bilin).astype(np.float32) lr_residual_bilin_scaled = Scaling(lr_residual_bilin) lr_image_hat_nears = transform_to_vlr_nears(lr_image) lr_residual_nears = np.array(lr_image).astype( np.float32) - np.array(lr_image_hat_nears).astype(np.float32) lr_residual_nears_scaled = Scaling(lr_residual_nears) # LR_image_scaled + kernel_proj + LR_residual_scaled (CONCAT) ---> TO TORCH lr_image_without_kernel = lr_image_scaled #self.randkern.ConcatDegraInfo(lr_image_scaled) lr_image_with_resid = np.concatenate((lr_image_without_kernel,\ lr_residual_bicub_scaled,\ lr_residual_bilin_scaled,\ lr_residual_nears_scaled), axis=-1) lr_image_with_resid = torch.from_numpy(lr_image_with_resid).float().to( self.device) # NUMPY to TORCH # LR_image to torch lr_image_scaled = torch.from_numpy(lr_image_scaled).float().to( self.device) # NUMPY to TORCH #Transpose - Permute since for model we need input with channels first lr_image_scaled = lr_image_scaled.permute(2, 0, 1) hr_image_scaled = hr_image_scaled.permute(2, 0, 1) lr_image_with_resid = lr_image_with_resid.permute(2, 0, 1) hr_residual_scaled = hr_residual_scaled.permute(2, 0, 1) return image_path,\ lr_image_scaled,\ hr_image_scaled, \ lr_image_with_resid, \ hr_residual_scaled def __len__(self): """Return the total number of image files.""" return len(self.image_paths)
def test(self): #receives single image --> can be easily modified to handle multiple images 'Takes single LR image as input. Returns LR image + (models approx) HR image concatenated' 'image location must be given by flag --test_image_path' self.model.eval() step = self.start_step + 1 # if not loading trained start = 0 lr_image = Image.open(self.test_image_path) lr_image_size = lr_image.size[0] #CONSIDER RGB IMAGE from utils import Kernels, load_kernels K, P = load_kernels(file_path='kernels/', scale_factor=2) randkern = Kernels(K, P) # get LR_RESIDUAL --> [-1,1] transform_to_vlr = transforms.Compose([ transforms.Lambda(lambda x: randkern.RandomBlur(x)), #random blur transforms.Lambda(lambda x: random_downscale(x,self.scale_factor)), #random downscale transforms.Resize((lr_image_size, lr_image_size), Image.BICUBIC) #upscale pro tamanho LR ]) lr_image_hat = transform_to_vlr(lr_image) lr_residual = np.array(lr_image).astype(np.float32) - np.array(lr_image_hat).astype(np.float32) lr_residual_scaled = Scaling(lr_residual) # LR_image_scaled + LR_residual_scaled (CONCAT) ---> TO TORCH #lr_image_with_kernel = self.randkern.ConcatDegraInfo(lr_image_scaled) #lr_image_with_resid = np.concatenate((lr_image_with_kernel, lr_residual_scaled), axis=-1) lr_image_scaled = Scaling(lr_image) lr_image_with_resid = np.concatenate((lr_image_scaled, lr_residual_scaled), axis=-1) lr_image_with_resid = torch.from_numpy(lr_image_with_resid).float().to(self.device) # NUMPY to TORCH # LR_image to torch lr_image_scaled = torch.from_numpy(lr_image_scaled).float().to(self.device) # NUMPY to TORCH #Transpose - Permute since for model we need input with channels first lr_image_scaled = lr_image_scaled.permute(2,0,1) lr_image_with_resid = lr_image_with_resid.permute(2,0,1) lr_image_with_resid = lr_image_with_resid.unsqueeze(0) #just add one dimension (index on batch) lr_image_scaled = lr_image_scaled.unsqueeze(0) lr_image, x = lr_image_scaled, lr_image_with_resid lr_image, x = lr_image.to(self.device), x.to(self.device) reconst = self.model(x) tmp1 = lr_image.data.cpu().numpy().transpose(0,2,3,1)*255 image_list = [np.array(Image.fromarray(tmp1.astype(np.uint8)[i]).resize((128,128), Image.BICUBIC)) \ for i in range(self.data_loader.batch_size)] image_hr_bicubic= np.stack(image_list) image_hr_bicubic_single = np.squeeze(image_hr_bicubic) #return this ^ image_hr_bicubic = image_hr_bicubic.transpose(0,3,1,2) image_hr_bicubic = Scaling(image_hr_bicubic) image_hr_bicubic = torch.from_numpy(image_hr_bicubic).float().to(self.device) # NUMPY to TORCH hr_image_hat = reconst + image_hr_bicubic hr_image_hat_np = hr_image_hat.data.cpu().numpy() hr_image_hat_np_scaled = Scaling01(hr_image_hat_np) hr_image_hat_np_scaled = np.squeeze(hr_image_hat_np_scaled).transpose((1, 2, 0)) hr_image_hat_np_png = (hr_image_hat_np_scaled*255).astype(np.uint8) #return this ^ #Saving Image Bicubic and HR Image Hat Image.fromarray(image_hr_bicubic_single).save('./results/HR_bicub_images/'+ os.path.basename(self.test_image_path)+'_hr_bic_{}.png'.format(step)) Image.fromarray(hr_image_hat_np_png).save('./results/HR_HAT_images/'+ os.path.basename(self.test_image_path)+'_hr_hat_{}.png'.format(step)) #Create Grid hr_image_hat_np_scaled = Scaling01(hr_image_hat_np) hr_image_hat_torch = torch.from_numpy(hr_image_hat_np_scaled).float().to(self.device) # NUMPY to TORCH pairs = torch.cat((image_hr_bicubic.data, \ hr_image_hat_torch.data), dim=3) grid = make_grid(pairs, 1) tmp = np.squeeze(grid.cpu().numpy().transpose((1, 2, 0))) tmp = (255 * tmp).astype(np.uint8) random_number = np.random.rand(1)[0] Image.fromarray(tmp).save('./results/grids/'+ os.path.basename(self.test_image_path).split('.')[0]+'_grid_{}.png'.format(step))
class ImageDataset(Dataset): def __init__(self, root, config, hr_shape): self.device = config.device self.image_paths = sorted(glob.glob(root + "/*.*")) self.image_size = config.image_size self.scale_factor = config.scale_factor hr_height, hr_width = hr_shape K, P = load_kernels(file_path='kernels/', scale_factor=self.scale_factor) #K = kernels -> K.shape = (15,15,1,358) #P = Matriz de projeçao do PCA --> P.shape = (15,225) self.randkern = Kernels(K, P) # Transforms for low resolution images and high resolution images def __getitem__(self, index): image_path = self.image_paths[index % len(self.image_paths)] image = Image.open(image_path) if np.array(image).shape==4: image = np.array(image)[:,:,:3] image = Image.fromarray(image) face_width = face_height = 128 ######## HARDCODED HR.shape = 128 ############### j = (image.size[0] - face_width) // 2 i = (image.size[1] - face_height) // 2 image = image.crop([j, i, j + face_width, i + face_height]) hr_image = image # HR_image --> [0,1] --> torch hr_image_scaled = Scaling(hr_image) hr_image_scaled = torch.from_numpy(hr_image_scaled).float().to(self.device) # NUMPY to TORCH # get HR residuals --> [-1,1] --> torch transform_HR = transforms.Compose([ #random blur transforms.Lambda(lambda x: self.randkern.RandomBlur(x)), #downscale BICUBIC pro tamanho LR transforms.Resize((self.image_size, self.image_size), Image.BICUBIC), #upscale BICUBIC pro tamanho HR transforms.Resize((self.image_size*self.scale_factor, self.image_size*self.scale_factor), Image.BICUBIC) ]) hr_image_hat = transform_HR(hr_image) hr_residual = np.array(hr_image).astype(float) - np.array(hr_image_hat).astype(float) hr_residual_scaled = Scaling(hr_residual) hr_residual_scaled = torch.from_numpy(hr_residual_scaled).float().to(self.device) # NUMPY to TORCH # get LR_RESIDUAL --> [-1,1] transform_to_vlr = transforms.Compose([ transforms.Lambda(lambda x: self.randkern.RandomBlur(x)), #random blur transforms.Lambda(lambda x: random_downscale(x,self.scale_factor)), #random downscale transforms.Resize((self.image_size, self.image_size), Image.BICUBIC) #upscale pro tamanho LR ]) lr_image_hat = transform_to_vlr(lr_image) lr_residual = np.array(lr_image).astype(np.float32) - np.array(lr_image_hat).astype(np.float32) lr_residual_scaled = Scaling(lr_residual) # LR_image_scaled + LR_residual_scaled (CONCAT) ---> TO TORCH lr_image_without_kernel = lr_image_scaled #self.randkern.ConcatDegraInfo(lr_image_scaled) lr_image_with_resid = np.concatenate((lr_image_without_kernel, lr_residual_scaled), axis=-1) lr_image_with_resid = torch.from_numpy(lr_image_with_resid).float().to(self.device) # NUMPY to TORCH # LR_image to torch lr_image_scaled = torch.from_numpy(lr_image_scaled).float().to(self.device) # NUMPY to TORCH #Transpose - Permute since for model we need input with channels first lr_image_scaled = lr_image_scaled.permute(2,0,1) hr_image_scaled = hr_image_scaled.permute(2,0,1) lr_image_with_resid = lr_image_with_resid.permute(2,0,1) hr_residual_scaled = hr_residual_scaled.permute(2,0,1) return {'image_path': image_path,\ 'lr_image_scaled':lr_image_scaled.to(torch.float64),\ 'hr_image_scaled':hr_image_scaled.to(torch.float64), \ 'lr_image_with_resid':lr_image_with_resid.to(torch.float64), \ 'hr_residual_scaled':hr_residual_scaled.to(torch.float64)} def __len__(self): return len(self.image_paths)