Spaces:
Runtime error
Runtime error
Upload utils_sr.py
Browse files- 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], :]
|