Tzktz commited on
Commit
d761518
1 Parent(s): 499adea

Upload 10 files

Browse files
RealESRGAN/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ from .model import RealESRGAN
RealESRGAN/arch_utils.py ADDED
@@ -0,0 +1,197 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import torch
3
+ from torch import nn as nn
4
+ from torch.nn import functional as F
5
+ from torch.nn import init as init
6
+ from torch.nn.modules.batchnorm import _BatchNorm
7
+
8
+ @torch.no_grad()
9
+ def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs):
10
+ """Initialize network weights.
11
+
12
+ Args:
13
+ module_list (list[nn.Module] | nn.Module): Modules to be initialized.
14
+ scale (float): Scale initialized weights, especially for residual
15
+ blocks. Default: 1.
16
+ bias_fill (float): The value to fill bias. Default: 0
17
+ kwargs (dict): Other arguments for initialization function.
18
+ """
19
+ if not isinstance(module_list, list):
20
+ module_list = [module_list]
21
+ for module in module_list:
22
+ for m in module.modules():
23
+ if isinstance(m, nn.Conv2d):
24
+ init.kaiming_normal_(m.weight, **kwargs)
25
+ m.weight.data *= scale
26
+ if m.bias is not None:
27
+ m.bias.data.fill_(bias_fill)
28
+ elif isinstance(m, nn.Linear):
29
+ init.kaiming_normal_(m.weight, **kwargs)
30
+ m.weight.data *= scale
31
+ if m.bias is not None:
32
+ m.bias.data.fill_(bias_fill)
33
+ elif isinstance(m, _BatchNorm):
34
+ init.constant_(m.weight, 1)
35
+ if m.bias is not None:
36
+ m.bias.data.fill_(bias_fill)
37
+
38
+
39
+ def make_layer(basic_block, num_basic_block, **kwarg):
40
+ """Make layers by stacking the same blocks.
41
+
42
+ Args:
43
+ basic_block (nn.module): nn.module class for basic block.
44
+ num_basic_block (int): number of blocks.
45
+
46
+ Returns:
47
+ nn.Sequential: Stacked blocks in nn.Sequential.
48
+ """
49
+ layers = []
50
+ for _ in range(num_basic_block):
51
+ layers.append(basic_block(**kwarg))
52
+ return nn.Sequential(*layers)
53
+
54
+
55
+ class ResidualBlockNoBN(nn.Module):
56
+ """Residual block without BN.
57
+
58
+ It has a style of:
59
+ ---Conv-ReLU-Conv-+-
60
+ |________________|
61
+
62
+ Args:
63
+ num_feat (int): Channel number of intermediate features.
64
+ Default: 64.
65
+ res_scale (float): Residual scale. Default: 1.
66
+ pytorch_init (bool): If set to True, use pytorch default init,
67
+ otherwise, use default_init_weights. Default: False.
68
+ """
69
+
70
+ def __init__(self, num_feat=64, res_scale=1, pytorch_init=False):
71
+ super(ResidualBlockNoBN, self).__init__()
72
+ self.res_scale = res_scale
73
+ self.conv1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
74
+ self.conv2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=True)
75
+ self.relu = nn.ReLU(inplace=True)
76
+
77
+ if not pytorch_init:
78
+ default_init_weights([self.conv1, self.conv2], 0.1)
79
+
80
+ def forward(self, x):
81
+ identity = x
82
+ out = self.conv2(self.relu(self.conv1(x)))
83
+ return identity + out * self.res_scale
84
+
85
+
86
+ class Upsample(nn.Sequential):
87
+ """Upsample module.
88
+
89
+ Args:
90
+ scale (int): Scale factor. Supported scales: 2^n and 3.
91
+ num_feat (int): Channel number of intermediate features.
92
+ """
93
+
94
+ def __init__(self, scale, num_feat):
95
+ m = []
96
+ if (scale & (scale - 1)) == 0: # scale = 2^n
97
+ for _ in range(int(math.log(scale, 2))):
98
+ m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))
99
+ m.append(nn.PixelShuffle(2))
100
+ elif scale == 3:
101
+ m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))
102
+ m.append(nn.PixelShuffle(3))
103
+ else:
104
+ raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')
105
+ super(Upsample, self).__init__(*m)
106
+
107
+
108
+ def flow_warp(x, flow, interp_mode='bilinear', padding_mode='zeros', align_corners=True):
109
+ """Warp an image or feature map with optical flow.
110
+
111
+ Args:
112
+ x (Tensor): Tensor with size (n, c, h, w).
113
+ flow (Tensor): Tensor with size (n, h, w, 2), normal value.
114
+ interp_mode (str): 'nearest' or 'bilinear'. Default: 'bilinear'.
115
+ padding_mode (str): 'zeros' or 'border' or 'reflection'.
116
+ Default: 'zeros'.
117
+ align_corners (bool): Before pytorch 1.3, the default value is
118
+ align_corners=True. After pytorch 1.3, the default value is
119
+ align_corners=False. Here, we use the True as default.
120
+
121
+ Returns:
122
+ Tensor: Warped image or feature map.
123
+ """
124
+ assert x.size()[-2:] == flow.size()[1:3]
125
+ _, _, h, w = x.size()
126
+ # create mesh grid
127
+ grid_y, grid_x = torch.meshgrid(torch.arange(0, h).type_as(x), torch.arange(0, w).type_as(x))
128
+ grid = torch.stack((grid_x, grid_y), 2).float() # W(x), H(y), 2
129
+ grid.requires_grad = False
130
+
131
+ vgrid = grid + flow
132
+ # scale grid to [-1,1]
133
+ vgrid_x = 2.0 * vgrid[:, :, :, 0] / max(w - 1, 1) - 1.0
134
+ vgrid_y = 2.0 * vgrid[:, :, :, 1] / max(h - 1, 1) - 1.0
135
+ vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=3)
136
+ output = F.grid_sample(x, vgrid_scaled, mode=interp_mode, padding_mode=padding_mode, align_corners=align_corners)
137
+
138
+ # TODO, what if align_corners=False
139
+ return output
140
+
141
+
142
+ def resize_flow(flow, size_type, sizes, interp_mode='bilinear', align_corners=False):
143
+ """Resize a flow according to ratio or shape.
144
+
145
+ Args:
146
+ flow (Tensor): Precomputed flow. shape [N, 2, H, W].
147
+ size_type (str): 'ratio' or 'shape'.
148
+ sizes (list[int | float]): the ratio for resizing or the final output
149
+ shape.
150
+ 1) The order of ratio should be [ratio_h, ratio_w]. For
151
+ downsampling, the ratio should be smaller than 1.0 (i.e., ratio
152
+ < 1.0). For upsampling, the ratio should be larger than 1.0 (i.e.,
153
+ ratio > 1.0).
154
+ 2) The order of output_size should be [out_h, out_w].
155
+ interp_mode (str): The mode of interpolation for resizing.
156
+ Default: 'bilinear'.
157
+ align_corners (bool): Whether align corners. Default: False.
158
+
159
+ Returns:
160
+ Tensor: Resized flow.
161
+ """
162
+ _, _, flow_h, flow_w = flow.size()
163
+ if size_type == 'ratio':
164
+ output_h, output_w = int(flow_h * sizes[0]), int(flow_w * sizes[1])
165
+ elif size_type == 'shape':
166
+ output_h, output_w = sizes[0], sizes[1]
167
+ else:
168
+ raise ValueError(f'Size type should be ratio or shape, but got type {size_type}.')
169
+
170
+ input_flow = flow.clone()
171
+ ratio_h = output_h / flow_h
172
+ ratio_w = output_w / flow_w
173
+ input_flow[:, 0, :, :] *= ratio_w
174
+ input_flow[:, 1, :, :] *= ratio_h
175
+ resized_flow = F.interpolate(
176
+ input=input_flow, size=(output_h, output_w), mode=interp_mode, align_corners=align_corners)
177
+ return resized_flow
178
+
179
+
180
+ # TODO: may write a cpp file
181
+ def pixel_unshuffle(x, scale):
182
+ """ Pixel unshuffle.
183
+
184
+ Args:
185
+ x (Tensor): Input feature with shape (b, c, hh, hw).
186
+ scale (int): Downsample ratio.
187
+
188
+ Returns:
189
+ Tensor: the pixel unshuffled feature.
190
+ """
191
+ b, c, hh, hw = x.size()
192
+ out_channel = c * (scale**2)
193
+ assert hh % scale == 0 and hw % scale == 0
194
+ h = hh // scale
195
+ w = hw // scale
196
+ x_view = x.view(b, c, h, scale, w, scale)
197
+ return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)
RealESRGAN/model.py ADDED
@@ -0,0 +1,93 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ from torch.nn import functional as F
4
+ from PIL import Image
5
+ import numpy as np
6
+ import cv2
7
+ from huggingface_hub import hf_hub_url, hf_hub_download, cached_download
8
+
9
+ from .rrdbnet_arch import RRDBNet
10
+ from .utils import pad_reflect, split_image_into_overlapping_patches, stich_together, \
11
+ unpad_image
12
+
13
+ HF_MODELS = {
14
+ 2: dict(
15
+ repo_id='sberbank-ai/Real-ESRGAN',
16
+ filename='RealESRGAN_x2.pth',
17
+ ),
18
+ 4: dict(
19
+ repo_id='sberbank-ai/Real-ESRGAN',
20
+ filename='RealESRGAN_x4.pth',
21
+ ),
22
+ 8: dict(
23
+ repo_id='sberbank-ai/Real-ESRGAN',
24
+ filename='RealESRGAN_x8.pth',
25
+ ),
26
+ }
27
+
28
+
29
+ class RealESRGAN:
30
+ def __init__(self, device, scale=4):
31
+ self.device = device
32
+ self.scale = scale
33
+ self.model = RRDBNet(
34
+ num_in_ch=3, num_out_ch=3, num_feat=64,
35
+ num_block=23, num_grow_ch=32, scale=scale
36
+ )
37
+
38
+ def load_weights(self, model_path, download=True):
39
+ if not os.path.exists(model_path) and download:
40
+ assert self.scale in [2, 4, 8], 'You can download models only with scales: 2, 4, 8'
41
+ config = HF_MODELS[self.scale]
42
+ cache_dir = os.path.dirname(model_path)
43
+ local_filename = os.path.basename(model_path)
44
+ config_file_url = hf_hub_url(repo_id=config['repo_id'], filename=config['filename'])
45
+ htr = hf_hub_download(repo_id=config['repo_id'], cache_dir=cache_dir, local_dir=cache_dir,
46
+ filename=config['filename'])
47
+ print(htr)
48
+ # cached_download(config_file_url, cache_dir=cache_dir, force_filename=local_filename)
49
+ print('Weights downloaded to:', os.path.join(cache_dir, local_filename))
50
+
51
+ loadnet = torch.load(model_path)
52
+ if 'params' in loadnet:
53
+ self.model.load_state_dict(loadnet['params'], strict=True)
54
+ elif 'params_ema' in loadnet:
55
+ self.model.load_state_dict(loadnet['params_ema'], strict=True)
56
+ else:
57
+ self.model.load_state_dict(loadnet, strict=True)
58
+ self.model.eval()
59
+ self.model.to(self.device)
60
+
61
+ # @torch.cuda.amp.autocast()
62
+ def predict(self, lr_image, batch_size=4, patches_size=192,
63
+ padding=24, pad_size=15):
64
+ torch.autocast(device_type=self.device.type)
65
+ scale = self.scale
66
+ device = self.device
67
+ lr_image = np.array(lr_image)
68
+ lr_image = pad_reflect(lr_image, pad_size)
69
+
70
+ patches, p_shape = split_image_into_overlapping_patches(
71
+ lr_image, patch_size=patches_size, padding_size=padding
72
+ )
73
+ img = torch.FloatTensor(patches / 255).permute((0, 3, 1, 2)).to(device).detach()
74
+
75
+ with torch.no_grad():
76
+ res = self.model(img[0:batch_size])
77
+ for i in range(batch_size, img.shape[0], batch_size):
78
+ res = torch.cat((res, self.model(img[i:i + batch_size])), 0)
79
+
80
+ sr_image = res.permute((0, 2, 3, 1)).cpu().clamp_(0, 1)
81
+ np_sr_image = sr_image.numpy()
82
+
83
+ padded_size_scaled = tuple(np.multiply(p_shape[0:2], scale)) + (3,)
84
+ scaled_image_shape = tuple(np.multiply(lr_image.shape[0:2], scale)) + (3,)
85
+ np_sr_image = stich_together(
86
+ np_sr_image, padded_image_shape=padded_size_scaled,
87
+ target_shape=scaled_image_shape, padding_size=padding * scale
88
+ )
89
+ sr_img = (np_sr_image * 255).astype(np.uint8)
90
+ sr_img = unpad_image(sr_img, pad_size * scale)
91
+ sr_img = Image.fromarray(sr_img)
92
+
93
+ return sr_img
RealESRGAN/rrdbnet_arch.py ADDED
@@ -0,0 +1,121 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from torch import nn as nn
3
+ from torch.nn import functional as F
4
+
5
+ from .arch_utils import default_init_weights, make_layer, pixel_unshuffle
6
+
7
+
8
+ class ResidualDenseBlock(nn.Module):
9
+ """Residual Dense Block.
10
+
11
+ Used in RRDB block in ESRGAN.
12
+
13
+ Args:
14
+ num_feat (int): Channel number of intermediate features.
15
+ num_grow_ch (int): Channels for each growth.
16
+ """
17
+
18
+ def __init__(self, num_feat=64, num_grow_ch=32):
19
+ super(ResidualDenseBlock, self).__init__()
20
+ self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
21
+ self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
22
+ self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
23
+ self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
24
+ self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
25
+
26
+ self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
27
+
28
+ # initialization
29
+ default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
30
+
31
+ def forward(self, x):
32
+ x1 = self.lrelu(self.conv1(x))
33
+ x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
34
+ x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
35
+ x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
36
+ x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
37
+ # Emperically, we use 0.2 to scale the residual for better performance
38
+ return x5 * 0.2 + x
39
+
40
+
41
+ class RRDB(nn.Module):
42
+ """Residual in Residual Dense Block.
43
+
44
+ Used in RRDB-Net in ESRGAN.
45
+
46
+ Args:
47
+ num_feat (int): Channel number of intermediate features.
48
+ num_grow_ch (int): Channels for each growth.
49
+ """
50
+
51
+ def __init__(self, num_feat, num_grow_ch=32):
52
+ super(RRDB, self).__init__()
53
+ self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
54
+ self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
55
+ self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
56
+
57
+ def forward(self, x):
58
+ out = self.rdb1(x)
59
+ out = self.rdb2(out)
60
+ out = self.rdb3(out)
61
+ # Emperically, we use 0.2 to scale the residual for better performance
62
+ return out * 0.2 + x
63
+
64
+
65
+ class RRDBNet(nn.Module):
66
+ """Networks consisting of Residual in Residual Dense Block, which is used
67
+ in ESRGAN.
68
+
69
+ ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks.
70
+
71
+ We extend ESRGAN for scale x2 and scale x1.
72
+ Note: This is one option for scale 1, scale 2 in RRDBNet.
73
+ We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size
74
+ and enlarge the channel size before feeding inputs into the main ESRGAN architecture.
75
+
76
+ Args:
77
+ num_in_ch (int): Channel number of inputs.
78
+ num_out_ch (int): Channel number of outputs.
79
+ num_feat (int): Channel number of intermediate features.
80
+ Default: 64
81
+ num_block (int): Block number in the trunk network. Defaults: 23
82
+ num_grow_ch (int): Channels for each growth. Default: 32.
83
+ """
84
+
85
+ def __init__(self, num_in_ch, num_out_ch, scale=4, num_feat=64, num_block=23, num_grow_ch=32):
86
+ super(RRDBNet, self).__init__()
87
+ self.scale = scale
88
+ if scale == 2:
89
+ num_in_ch = num_in_ch * 4
90
+ elif scale == 1:
91
+ num_in_ch = num_in_ch * 16
92
+ self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
93
+ self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
94
+ self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
95
+ # upsample
96
+ self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
97
+ self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
98
+ if scale == 8:
99
+ self.conv_up3 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
100
+ self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
101
+ self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
102
+
103
+ self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
104
+
105
+ def forward(self, x):
106
+ if self.scale == 2:
107
+ feat = pixel_unshuffle(x, scale=2)
108
+ elif self.scale == 1:
109
+ feat = pixel_unshuffle(x, scale=4)
110
+ else:
111
+ feat = x
112
+ feat = self.conv_first(feat)
113
+ body_feat = self.conv_body(self.body(feat))
114
+ feat = feat + body_feat
115
+ # upsample
116
+ feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))
117
+ feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest')))
118
+ if self.scale == 8:
119
+ feat = self.lrelu(self.conv_up3(F.interpolate(feat, scale_factor=2, mode='nearest')))
120
+ out = self.conv_last(self.lrelu(self.conv_hr(feat)))
121
+ return out
RealESRGAN/utils.py ADDED
@@ -0,0 +1,135 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import torch
3
+ from PIL import Image
4
+ import os
5
+ import io
6
+
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
+
22
+ def unpad_image(image, pad_size):
23
+ return image[pad_size:-pad_size, pad_size:-pad_size, :]
24
+
25
+
26
+ def process_array(image_array, expand=True):
27
+ """ Process a 3-dimensional array into a scaled, 4 dimensional batch of size 1. """
28
+
29
+ image_batch = image_array / 255.0
30
+ if expand:
31
+ image_batch = np.expand_dims(image_batch, axis=0)
32
+ return image_batch
33
+
34
+
35
+ def process_output(output_tensor):
36
+ """ Transforms the 4-dimensional output tensor into a suitable image format. """
37
+
38
+ sr_img = output_tensor.clip(0, 1) * 255
39
+ sr_img = np.uint8(sr_img)
40
+ return sr_img
41
+
42
+
43
+ def pad_patch(image_patch, padding_size, channel_last=True):
44
+ """ Pads image_patch with with padding_size edge values. """
45
+
46
+ if channel_last:
47
+ return np.pad(
48
+ image_patch,
49
+ ((padding_size, padding_size), (padding_size, padding_size), (0, 0)),
50
+ 'edge',
51
+ )
52
+ else:
53
+ return np.pad(
54
+ image_patch,
55
+ ((0, 0), (padding_size, padding_size), (padding_size, padding_size)),
56
+ 'edge',
57
+ )
58
+
59
+
60
+ def unpad_patches(image_patches, padding_size):
61
+ return image_patches[:, padding_size:-padding_size, padding_size:-padding_size, :]
62
+
63
+
64
+ def split_image_into_overlapping_patches(image_array, patch_size, padding_size=2):
65
+ """ Splits the image into partially overlapping patches.
66
+ The patches overlap by padding_size pixels.
67
+ Pads the image twice:
68
+ - first to have a size multiple of the patch size,
69
+ - then to have equal padding at the borders.
70
+ Args:
71
+ image_array: numpy array of the input image.
72
+ patch_size: size of the patches from the original image (without padding).
73
+ padding_size: size of the overlapping area.
74
+ """
75
+
76
+ xmax, ymax, _ = image_array.shape
77
+ x_remainder = xmax % patch_size
78
+ y_remainder = ymax % patch_size
79
+
80
+ # modulo here is to avoid extending of patch_size instead of 0
81
+ x_extend = (patch_size - x_remainder) % patch_size
82
+ y_extend = (patch_size - y_remainder) % patch_size
83
+
84
+ # make sure the image is divisible into regular patches
85
+ extended_image = np.pad(image_array, ((0, x_extend), (0, y_extend), (0, 0)), 'edge')
86
+
87
+ # add padding around the image to simplify computations
88
+ padded_image = pad_patch(extended_image, padding_size, channel_last=True)
89
+
90
+ xmax, ymax, _ = padded_image.shape
91
+ patches = []
92
+
93
+ x_lefts = range(padding_size, xmax - padding_size, patch_size)
94
+ y_tops = range(padding_size, ymax - padding_size, patch_size)
95
+
96
+ for x in x_lefts:
97
+ for y in y_tops:
98
+ x_left = x - padding_size
99
+ y_top = y - padding_size
100
+ x_right = x + patch_size + padding_size
101
+ y_bottom = y + patch_size + padding_size
102
+ patch = padded_image[x_left:x_right, y_top:y_bottom, :]
103
+ patches.append(patch)
104
+
105
+ return np.array(patches), padded_image.shape
106
+
107
+
108
+ def stich_together(patches, padded_image_shape, target_shape, padding_size=4):
109
+ """ Reconstruct the image from overlapping patches.
110
+ After scaling, shapes and padding should be scaled too.
111
+ Args:
112
+ patches: patches obtained with split_image_into_overlapping_patches
113
+ padded_image_shape: shape of the padded image contructed in split_image_into_overlapping_patches
114
+ target_shape: shape of the final image
115
+ padding_size: size of the overlapping area.
116
+ """
117
+
118
+ xmax, ymax, _ = padded_image_shape
119
+ patches = unpad_patches(patches, padding_size)
120
+ patch_size = patches.shape[1]
121
+ n_patches_per_row = ymax // patch_size
122
+
123
+ complete_image = np.zeros((xmax, ymax, 3))
124
+
125
+ row = -1
126
+ col = 0
127
+ for i in range(len(patches)):
128
+ if i % n_patches_per_row == 0:
129
+ row += 1
130
+ col = 0
131
+ complete_image[
132
+ row * patch_size: (row + 1) * patch_size, col * patch_size: (col + 1) * patch_size, :
133
+ ] = patches[i]
134
+ col += 1
135
+ return complete_image[0: target_shape[0], 0: target_shape[1], :]
app.py ADDED
@@ -0,0 +1,75 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from PIL import Image
3
+ from RealESRGAN import RealESRGAN
4
+ import gradio as gr
5
+
6
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
7
+ model2 = RealESRGAN(device, scale=2)
8
+ model2.load_weights('weights/RealESRGAN_x2.pth', download=True)
9
+ model4 = RealESRGAN(device, scale=4)
10
+ model4.load_weights('weights/RealESRGAN_x4.pth', download=True)
11
+ model8 = RealESRGAN(device, scale=8)
12
+ model8.load_weights('weights/RealESRGAN_x8.pth', download=True)
13
+
14
+
15
+ def inference(image, size):
16
+ global model2
17
+ global model4
18
+ global model8
19
+ if image is None:
20
+ raise gr.Error("Image not uploaded")
21
+
22
+ width, height = image.size
23
+ if width >= 5000 or height >= 5000:
24
+ raise gr.Error("The image is too large.")
25
+
26
+ if torch.cuda.is_available():
27
+ torch.cuda.empty_cache()
28
+
29
+ if size == '2x':
30
+ try:
31
+ result = model2.predict(image.convert('RGB'))
32
+ except torch.cuda.OutOfMemoryError as e:
33
+ print(e)
34
+ model2 = RealESRGAN(device, scale=2)
35
+ model2.load_weights('weights/RealESRGAN_x2.pth', download=False)
36
+ result = model2.predict(image.convert('RGB'))
37
+ elif size == '4x':
38
+ try:
39
+ result = model4.predict(image.convert('RGB'))
40
+ except torch.cuda.OutOfMemoryError as e:
41
+ print(e)
42
+ model4 = RealESRGAN(device, scale=4)
43
+ model4.load_weights('weights/RealESRGAN_x4.pth', download=False)
44
+ result = model2.predict(image.convert('RGB'))
45
+ else:
46
+ try:
47
+ result = model8.predict(image.convert('RGB'))
48
+ except torch.cuda.OutOfMemoryError as e:
49
+ print(e)
50
+ model8 = RealESRGAN(device, scale=8)
51
+ model8.load_weights('weights/RealESRGAN_x8.pth', download=False)
52
+ result = model2.predict(image.convert('RGB'))
53
+
54
+ print(f"Image size ({device}): {size} ... OK")
55
+ return result
56
+
57
+
58
+ title = "Face Real ESRGAN UpScale: 2x 4x 8x"
59
+ description = "This is an unofficial demo for Real-ESRGAN. Scales the resolution of a photo. This model shows better results on faces compared to the original version.<br>Telegram BOT: https://t.me/restoration_photo_bot"
60
+ article = "<div style='text-align: center;'>Twitter <a href='https://twitter.com/DoEvent' target='_blank'>Max Skobeev</a> | <a href='https://huggingface.co/sberbank-ai/Real-ESRGAN' target='_blank'>Model card</a><div>"
61
+
62
+ gr.Interface(inference,
63
+ [gr.Image(type="pil"),
64
+ gr.Radio(['2x', '4x', '8x'],
65
+ type="value",
66
+ value='2x',
67
+ label='Resolution model')],
68
+ gr.Image(type="pil", label="Output"),
69
+ title=title,
70
+ description=description,
71
+ article=article,
72
+
73
+ allow_flagging='never',
74
+ cache_examples=False,
75
+ ).queue(api_open=False).launch(show_error=True, show_api=False)
requirements.txt ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ torch
2
+ numpy
3
+ opencv-python-headless
4
+ setuptools
5
+ Pillow
6
+ torchvision
7
+ addict
8
+ future
9
+ lmdb
10
+ pyyaml
11
+ requests
12
+ scikit-image
13
+ scipy
14
+ tb-nightly
15
+ tqdm
16
+ yapf
17
+ psutil
weights/RealESRGAN_x2.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c830d067d54fc767b9543a8432f36d91bc2de313584e8bbfe4ac26a47339e899
3
+ size 67061725
weights/RealESRGAN_x4.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:aa00f09ad753d88576b21ed977e97d634976377031b178acc3b5b238df463400
3
+ size 67040989
weights/RealESRGAN_x8.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8b72fb469d12f05a4770813d2603eb1b550f40df6fb8b37d6c7bc2db3d2bff5e
3
+ size 67189359