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)
# 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 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)
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)