Spaces:
Sleeping
Sleeping
Upload 10 files
Browse files- RealESRGAN/__init__.py +1 -0
- RealESRGAN/arch_utils.py +197 -0
- RealESRGAN/model.py +93 -0
- RealESRGAN/rrdbnet_arch.py +121 -0
- RealESRGAN/utils.py +135 -0
- app.py +75 -0
- requirements.txt +17 -0
- weights/RealESRGAN_x2.pth +3 -0
- weights/RealESRGAN_x4.pth +3 -0
- weights/RealESRGAN_x8.pth +3 -0
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
|