SRCNN_Pytorch/utils.py at master · thangvubk/SRCNN_Pytorch · GitHub

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
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import os
import glob
import scipy.misc
import scipy.ndimage
import numpy as np

def get_image_paths(dir_path):
    """
    Get image paths from dir
    """
    dir_path = os.path.join(os.getcwd(), dir_path)
    image_paths = glob.glob(os.path.join(dir_path, '*.bmp'))
    return image_paths

def imread(image_path, is_grayscale=True):
    """
    Read image from path
    """
    if is_grayscale:
        return scipy.misc.imread(image_path, flatten=True, mode='YCbCr').astype(np.float)
    else:
        return scipy.misc.imread(image_path, mode='YCbCr').astype(np.float)

def mod_crop(image, scale=3):
    """
    modulo crop the image with the scale
    """
    h, w = image.shape[0], image.shape[1]
    h = h - np.mod(h, scale)
    w = w - np.mod(w, scale)
    return image[:h, :w]

def process_image(image, scale=3):
    """
    Normalize
    Aplly bicubic interpolation

    Args:
        file_path: path to single image
    Return:
        - inp: image applied bicubic interpolation (low resolution)
        - label: image with original resolution (high resolution)
    """
    label = mod_crop(image, scale)
    label = label/255 #normalize

    inp = scipy.ndimage.interpolation.zoom(label, (1./scale), prefilter=False)
    inp = scipy.ndimage.interpolation.zoom(inp, (scale/1.), prefilter=False)

    return inp, label

def slicing_image(inp, label, I, L, stride):
    """
    Slice inp image (low resolution) to obtain multiple sub-image of size (kernel x kernel)
    with the step being stride

    Args:
        - inp: input image (low resolution)
        - label: image of high resolution
        - I: filter size of to slice inp
        - L: filter size to slice label
        - stride: step size to move filter
    Return
        - sub_inputs: sub-images from inp
        - sub_labels: sub_images from label
    """
    sub_inputs = []
    sub_labels = []
    h, w = inp.shape[0], inp.shape[1]
    offset = abs(I - L)//2 #offset to identify where to slice label

    for hh in range(0, h-I+1, stride):
        for ww in range(0, w-I+1, stride):
            sub_input = inp[hh:hh+I, ww:ww+I]
            sub_label = label[hh+offset:hh+offset+L, ww+offset:ww+offset+L]

            # Make channel value
            sub_input = sub_input.reshape(I, I, 1)
            sub_label = sub_label.reshape(L, L, 1)

            sub_inputs.append(sub_input)
            sub_labels.append(sub_label)
    return sub_inputs, sub_labels

def make_input(config):
    """
    Produce image and label from config

    Args:
        Config: config to produce image and label

    Return:
        - inputs: data for the model
        - labels: groud true for the model
    """
    inputs = []
    labels = []

    dir_path = config['dir_path']
    scale = config['scale']
    is_grayscale = config['is_gray']
    I = config['input_size']
    L = config['label_size']
    stride = config['stride']

    image_paths = get_image_paths(dir_path)

    for path in image_paths:
        image = imread(path, is_grayscale)
        inp, label = process_image(image, scale)
        sub_inputs, sub_labels = slicing_image(inp, label, I, L, stride)
        inputs += sub_inputs
        labels += sub_labels

    inputs = np.asarray(inputs) # shape (N, I, I, 1)
    labels = np.asarray(labels) # shape (N, L, L, 1)
    return inputs, labels

# testing
if __name__ == "__main__":
    paths = get_image_paths('Train')
    image = imread(paths[0])
    print(image)
    inp, label = process_image(image)
    scipy.misc.imsave('inp.png', inp)
    scipy.misc.imsave('label.png', label)
    print(inp)
    print(label)