doevent commited on
Commit
53d095b
1 Parent(s): 1675c65

Upload utils_sr.py

Browse files
Files changed (1) hide show
  1. utils_sr.py +141 -0
utils_sr.py ADDED
@@ -0,0 +1,141 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import torch
3
+ from PIL import Image
4
+ import os
5
+ import io
6
+ import imageio
7
+
8
+ def pad_reflect(image, pad_size):
9
+ imsize = image.shape
10
+ height, width = imsize[:2]
11
+ new_img = np.zeros([height+pad_size*2, width+pad_size*2, imsize[2]]).astype(np.uint8)
12
+ new_img[pad_size:-pad_size, pad_size:-pad_size, :] = image
13
+
14
+ new_img[0:pad_size, pad_size:-pad_size, :] = np.flip(image[0:pad_size, :, :], axis=0) #top
15
+ new_img[-pad_size:, pad_size:-pad_size, :] = np.flip(image[-pad_size:, :, :], axis=0) #bottom
16
+ new_img[:, 0:pad_size, :] = np.flip(new_img[:, pad_size:pad_size*2, :], axis=1) #left
17
+ new_img[:, -pad_size:, :] = np.flip(new_img[:, -pad_size*2:-pad_size, :], axis=1) #right
18
+
19
+ return new_img
20
+
21
+ def unpad_image(image, pad_size):
22
+ return image[pad_size:-pad_size, pad_size:-pad_size, :]
23
+
24
+
25
+ def jpegBlur(im,q):
26
+ buf = io.BytesIO()
27
+ imageio.imwrite(buf,im,format='jpg',quality=q)
28
+ s = buf.getbuffer()
29
+ return imageio.imread(s,format='jpg')
30
+
31
+
32
+ def process_array(image_array, expand=True):
33
+ """ Process a 3-dimensional array into a scaled, 4 dimensional batch of size 1. """
34
+
35
+ image_batch = image_array / 255.0
36
+ if expand:
37
+ image_batch = np.expand_dims(image_batch, axis=0)
38
+ return image_batch
39
+
40
+
41
+ def process_output(output_tensor):
42
+ """ Transforms the 4-dimensional output tensor into a suitable image format. """
43
+
44
+ sr_img = output_tensor.clip(0, 1) * 255
45
+ sr_img = np.uint8(sr_img)
46
+ return sr_img
47
+
48
+
49
+ def pad_patch(image_patch, padding_size, channel_last=True):
50
+ """ Pads image_patch with with padding_size edge values. """
51
+
52
+ if channel_last:
53
+ return np.pad(
54
+ image_patch,
55
+ ((padding_size, padding_size), (padding_size, padding_size), (0, 0)),
56
+ 'edge',
57
+ )
58
+ else:
59
+ return np.pad(
60
+ image_patch,
61
+ ((0, 0), (padding_size, padding_size), (padding_size, padding_size)),
62
+ 'edge',
63
+ )
64
+
65
+
66
+ def unpad_patches(image_patches, padding_size):
67
+ return image_patches[:, padding_size:-padding_size, padding_size:-padding_size, :]
68
+
69
+
70
+ def split_image_into_overlapping_patches(image_array, patch_size, padding_size=2):
71
+ """ Splits the image into partially overlapping patches.
72
+ The patches overlap by padding_size pixels.
73
+ Pads the image twice:
74
+ - first to have a size multiple of the patch size,
75
+ - then to have equal padding at the borders.
76
+ Args:
77
+ image_array: numpy array of the input image.
78
+ patch_size: size of the patches from the original image (without padding).
79
+ padding_size: size of the overlapping area.
80
+ """
81
+
82
+ xmax, ymax, _ = image_array.shape
83
+ x_remainder = xmax % patch_size
84
+ y_remainder = ymax % patch_size
85
+
86
+ # modulo here is to avoid extending of patch_size instead of 0
87
+ x_extend = (patch_size - x_remainder) % patch_size
88
+ y_extend = (patch_size - y_remainder) % patch_size
89
+
90
+ # make sure the image is divisible into regular patches
91
+ extended_image = np.pad(image_array, ((0, x_extend), (0, y_extend), (0, 0)), 'edge')
92
+
93
+ # add padding around the image to simplify computations
94
+ padded_image = pad_patch(extended_image, padding_size, channel_last=True)
95
+
96
+ xmax, ymax, _ = padded_image.shape
97
+ patches = []
98
+
99
+ x_lefts = range(padding_size, xmax - padding_size, patch_size)
100
+ y_tops = range(padding_size, ymax - padding_size, patch_size)
101
+
102
+ for x in x_lefts:
103
+ for y in y_tops:
104
+ x_left = x - padding_size
105
+ y_top = y - padding_size
106
+ x_right = x + patch_size + padding_size
107
+ y_bottom = y + patch_size + padding_size
108
+ patch = padded_image[x_left:x_right, y_top:y_bottom, :]
109
+ patches.append(patch)
110
+
111
+ return np.array(patches), padded_image.shape
112
+
113
+
114
+ def stich_together(patches, padded_image_shape, target_shape, padding_size=4):
115
+ """ Reconstruct the image from overlapping patches.
116
+ After scaling, shapes and padding should be scaled too.
117
+ Args:
118
+ patches: patches obtained with split_image_into_overlapping_patches
119
+ padded_image_shape: shape of the padded image contructed in split_image_into_overlapping_patches
120
+ target_shape: shape of the final image
121
+ padding_size: size of the overlapping area.
122
+ """
123
+
124
+ xmax, ymax, _ = padded_image_shape
125
+ patches = unpad_patches(patches, padding_size)
126
+ patch_size = patches.shape[1]
127
+ n_patches_per_row = ymax // patch_size
128
+
129
+ complete_image = np.zeros((xmax, ymax, 3))
130
+
131
+ row = -1
132
+ col = 0
133
+ for i in range(len(patches)):
134
+ if i % n_patches_per_row == 0:
135
+ row += 1
136
+ col = 0
137
+ complete_image[
138
+ row * patch_size: (row + 1) * patch_size, col * patch_size: (col + 1) * patch_size,:
139
+ ] = patches[i]
140
+ col += 1
141
+ return complete_image[0: target_shape[0], 0: target_shape[1], :]