def __init__(self,
image_dir,
upscale_factor,
img_channels,
crop_size=-1,
augmentations=None,
resampling=64):
"""
- crop_size = -1 only when testing.
"""
super(SuperResDataset, self).__init__()
# Get image filenames.
self.image_filenames = [join(image_dir, x) for x in listdir(image_dir) \
if is_image_file(x)]
# Read all images.
self.images = []
for image_filename in self.image_filenames:
image = Image.open(image_filename).convert('RGB')
if image.size[0] < crop_size or image.size[1] < crop_size:
continue
if img_channels == 1:
image = _rgb2ycbcr(image)
image, _, _ = image.split()
self.images.append(image)
self.len = len(self.images)
self.upscale_factor = upscale_factor
self.img_channels = img_channels
self.crop_size = crop_size
self.augmentations = augmentations
self.resampling = resampling
res = {}
for i, f in enumerate(image_filenames):
# Read test image.
img = Image.open(f).convert('RGB')
width, height = img.size[0], img.size[1]
# Crop test image so that it has size that can be downsampled by the upscale factor.
pad_width = width % args.upscale_factor
pad_height = height % args.upscale_factor
width -= pad_width
height -= pad_height
img = img.crop((0, 0, width, height))
if args.img_channels == 1:
img = _rgb2ycbcr(img)
img_y = array(img.split()[0], dtype=np.float32) / 255.0
img_y = Image.fromarray(img_y, mode='F')
# Downsample to get low-res image.
lr_img = img_y.resize(
(width // args.upscale_factor, height // args.upscale_factor),
Image.BICUBIC)
# Achieve high-res using FSRCNN.
y = lr_img.copy()
y = ToTensor()(y).view(1, -1, y.size[1], y.size[0])
out_img_deep_y = model(y)[-1].detach().numpy().squeeze()
out_img_deep_y = out_img_deep_y.clip(0, 1)
if args.img_channels == 3:
model.load_state_dict(ckpt['model'])
res = {}
for i, f in enumerate(image_filenames):
# Read test image.
img = Image.open(f).convert('RGB')
width, height = img.size[0], img.size[1]
# Crop test image so that it has size that can be downsampled by the upscale factor.
pad_width = width % args.upscale_factor
pad_height = height % args.upscale_factor
width -= pad_width
height -= pad_height
img = img.crop((0, 0, width, height))
img = _rgb2ycbcr(img)
img_y = array(img.split()[0], dtype=np.float32) / 255.0
img_y = Image.fromarray(img_y, mode='F')
# Downsample to get low-res image.
img_y = img_y.resize(
(width // args.upscale_factor, height // args.upscale_factor),
Image.BICUBIC)
# Achive high-res using deep neural net.
img_y = img_y.resize((width, height), Image.BICUBIC)
img_y = ToTensor()(img_y).view(1, -1, img_y.size[1], img_y.size[0])
sr_y = model(img_y)[0].detach().numpy().squeeze()
sr_y = sr_y.clip(0, 1)
# Get image filename.