Spaces:
Runtime error
Runtime error
| 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 | |