Spaces:
Runtime error
Runtime error
File size: 3,784 Bytes
e9b996f |
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 |
import os
import numpy as np
import torch
from SwinIR.models.network_swinir import SwinIR as net
ROOT_PATH = os.path.dirname(__file__)
class SwinIRDemo:
def __init__(self):
self.scale = 4
self.window_size = 8
self.tile = 800
self.tile_overlap = 32
self.device = 'cuda'
model_path = os.path.join(ROOT_PATH, 'weight/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth')
self.model = self.model_init(model_path)
def model_init(self, model_path):
model = net(upscale=self.scale, in_chans=3, img_size=64, window_size=8,
img_range=1., depths=[6, 6, 6, 6, 6, 6], embed_dim=180, num_heads=[6, 6, 6, 6, 6, 6],
mlp_ratio=2, upsampler='nearest+conv', resi_connection='1conv')
param_key_g = 'params_ema'
pretrained_model = torch.load(model_path)
model.load_state_dict(
pretrained_model[param_key_g] if param_key_g in pretrained_model.keys() else pretrained_model,
strict=True)
model.eval()
model = model.to(self.device)
return model
def img_preprocess(self, img_PIL, device, window_size):
# imgname, img_lq, img_gt = get_image_pair(args, path) # image to HWC-BGR, float32
# img_lq = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
# img_lq = img_PIL.convert('BGR')
img_lq = np.asarray(img_PIL)
img_lq = img_lq / 255
img_lq = np.transpose(img_lq[:, :, [0, 1, 2]], (2, 0, 1)) # HCW-BGR to CHW-RGB
img_lq = torch.from_numpy(img_lq).float().unsqueeze(0).to(device) # CHW-RGB to NCHW-RGB
# pad input image to be a multiple of window_size
_, _, h_old, w_old = img_lq.size()
h_pad = (h_old // window_size + 1) * window_size - h_old
w_pad = (w_old // window_size + 1) * window_size - w_old
img_lq = torch.cat([img_lq, torch.flip(img_lq, [2])], 2)[:, :, :h_old + h_pad, :]
img_lq = torch.cat([img_lq, torch.flip(img_lq, [3])], 3)[:, :, :, :w_old + w_pad]
return img_lq, h_old, w_old
def test(self, img_lq):
b, c, h, w = img_lq.size()
tile = min(self.tile, h, w)
assert tile % self.window_size == 0, "tile size should be a multiple of window_size"
sf = self.scale
stride = tile - self.tile_overlap
h_idx_list = list(range(0, h - tile, stride)) + [h - tile]
w_idx_list = list(range(0, w - tile, stride)) + [w - tile]
E = torch.zeros(b, c, h * sf, w * sf).type_as(img_lq)
W = torch.zeros_like(E)
for h_idx in h_idx_list:
for w_idx in w_idx_list:
in_patch = img_lq[..., h_idx:h_idx + tile, w_idx:w_idx + tile]
out_patch = self.model(in_patch)
out_patch_mask = torch.ones_like(out_patch)
E[..., h_idx * sf:(h_idx + tile) * sf, w_idx * sf:(w_idx + tile) * sf].add_(out_patch)
W[..., h_idx * sf:(h_idx + tile) * sf, w_idx * sf:(w_idx + tile) * sf].add_(out_patch_mask)
output = E.div_(W)
return output
def infer(self, img_lq):
img_lq, h_old, w_old = self.img_preprocess(img_lq, self.device, self.window_size)
with torch.no_grad():
output = self.test(img_lq)
output = output[..., :h_old * self.scale, :w_old * self.scale]
output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
if output.ndim == 3:
output = np.transpose(output[[0, 1, 2], :, :], (1, 2, 0)) # CHW-RGB to HCW-BGR
output = (output * 255.0).round().astype(np.uint8)
return output
|