-
Notifications
You must be signed in to change notification settings - Fork 25
/
data_loader.py
66 lines (52 loc) · 2.33 KB
/
data_loader.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
import os
import glob
import torch
from torch.utils import data
from torchvision import transforms
from PIL import Image
import numpy as np
from utils import Kernels, load_kernels
TYPES = ('*.png', '*.jpg', '*.jpeg', '*.bmp')
torch.set_default_tensor_type(torch.DoubleTensor)
def Scaling(image):
return np.array(image) / 255.0
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 get_loader(image_path, config):
"""Create and return Dataloader."""
dataset = ImageFolder(image_path, config)
data_loader = data.DataLoader(dataset=dataset,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.num_workers)
return data_loader