Spaces:
Runtime error
Runtime error
| import os | |
| from enum import IntEnum | |
| from pathlib import Path | |
| from tempfile import mktemp | |
| from typing import IO, Dict, Type | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from gradio import Interface, inputs, outputs | |
| DEVICE = "cpu" | |
| WEIGHTS_PATH = Path(__file__).parent / "weights" | |
| AVALIABLE_WEIGHTS = { | |
| basename: path | |
| for basename, ext in ( | |
| os.path.splitext(filename) for filename in os.listdir(WEIGHTS_PATH) | |
| ) | |
| if (path := WEIGHTS_PATH / (basename + ext)).is_file() and ext.endswith("pth") | |
| } | |
| class ScaleMode(IntEnum): | |
| up2x = 2 | |
| up3x = 3 | |
| up4x = 4 | |
| class TileMode(IntEnum): | |
| full = 0 | |
| half = 1 | |
| quarter = 2 | |
| ninth = 3 | |
| sixteenth = 4 | |
| class SEBlock(nn.Module): | |
| def __init__(self, in_channels, reduction=8, bias=False): | |
| super(SEBlock, self).__init__() | |
| self.conv1 = nn.Conv2d( | |
| in_channels, in_channels // reduction, 1, 1, 0, bias=bias | |
| ) | |
| self.conv2 = nn.Conv2d( | |
| in_channels // reduction, in_channels, 1, 1, 0, bias=bias | |
| ) | |
| def forward(self, x): | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| x0 = torch.mean(x.float(), dim=(2, 3), keepdim=True).half() | |
| else: | |
| x0 = torch.mean(x, dim=(2, 3), keepdim=True) | |
| x0 = self.conv1(x0) | |
| x0 = F.relu(x0, inplace=True) | |
| x0 = self.conv2(x0) | |
| x0 = torch.sigmoid(x0) | |
| x = torch.mul(x, x0) | |
| return x | |
| def forward_mean(self, x, x0): | |
| x0 = self.conv1(x0) | |
| x0 = F.relu(x0, inplace=True) | |
| x0 = self.conv2(x0) | |
| x0 = torch.sigmoid(x0) | |
| x = torch.mul(x, x0) | |
| return x | |
| class UNetConv(nn.Module): | |
| def __init__(self, in_channels, mid_channels, out_channels, se): | |
| super(UNetConv, self).__init__() | |
| self.conv = nn.Sequential( | |
| nn.Conv2d(in_channels, mid_channels, 3, 1, 0), | |
| nn.LeakyReLU(0.1, inplace=True), | |
| nn.Conv2d(mid_channels, out_channels, 3, 1, 0), | |
| nn.LeakyReLU(0.1, inplace=True), | |
| ) | |
| if se: | |
| self.seblock = SEBlock(out_channels, reduction=8, bias=True) | |
| else: | |
| self.seblock = None | |
| def forward(self, x): | |
| z = self.conv(x) | |
| if self.seblock is not None: | |
| z = self.seblock(z) | |
| return z | |
| class UNet1(nn.Module): | |
| def __init__(self, in_channels, out_channels, deconv): | |
| super(UNet1, self).__init__() | |
| self.conv1 = UNetConv(in_channels, 32, 64, se=False) | |
| self.conv1_down = nn.Conv2d(64, 64, 2, 2, 0) | |
| self.conv2 = UNetConv(64, 128, 64, se=True) | |
| self.conv2_up = nn.ConvTranspose2d(64, 64, 2, 2, 0) | |
| self.conv3 = nn.Conv2d(64, 64, 3, 1, 0) | |
| if deconv: | |
| self.conv_bottom = nn.ConvTranspose2d(64, out_channels, 4, 2, 3) | |
| else: | |
| self.conv_bottom = nn.Conv2d(64, out_channels, 3, 1, 0) | |
| for m in self.modules(): | |
| if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)): | |
| nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") | |
| elif isinstance(m, nn.Linear): | |
| nn.init.normal_(m.weight, 0, 0.01) | |
| if m.bias is not None: | |
| nn.init.constant_(m.bias, 0) | |
| def forward(self, x): | |
| x1 = self.conv1(x) | |
| x2 = self.conv1_down(x1) | |
| x2 = F.leaky_relu(x2, 0.1, inplace=True) | |
| x2 = self.conv2(x2) | |
| x2 = self.conv2_up(x2) | |
| x2 = F.leaky_relu(x2, 0.1, inplace=True) | |
| x1 = F.pad(x1, (-4, -4, -4, -4)) | |
| x3 = self.conv3(x1 + x2) | |
| x3 = F.leaky_relu(x3, 0.1, inplace=True) | |
| z = self.conv_bottom(x3) | |
| return z | |
| def forward_a(self, x): | |
| x1 = self.conv1(x) | |
| x2 = self.conv1_down(x1) | |
| x2 = F.leaky_relu(x2, 0.1, inplace=True) | |
| x2 = self.conv2.conv(x2) | |
| return x1, x2 | |
| def forward_b(self, x1, x2): | |
| x2 = self.conv2_up(x2) | |
| x2 = F.leaky_relu(x2, 0.1, inplace=True) | |
| x1 = F.pad(x1, (-4, -4, -4, -4)) | |
| x3 = self.conv3(x1 + x2) | |
| x3 = F.leaky_relu(x3, 0.1, inplace=True) | |
| z = self.conv_bottom(x3) | |
| return z | |
| class UNet1x3(nn.Module): | |
| def __init__(self, in_channels, out_channels, deconv): | |
| super(UNet1x3, self).__init__() | |
| self.conv1 = UNetConv(in_channels, 32, 64, se=False) | |
| self.conv1_down = nn.Conv2d(64, 64, 2, 2, 0) | |
| self.conv2 = UNetConv(64, 128, 64, se=True) | |
| self.conv2_up = nn.ConvTranspose2d(64, 64, 2, 2, 0) | |
| self.conv3 = nn.Conv2d(64, 64, 3, 1, 0) | |
| if deconv: | |
| self.conv_bottom = nn.ConvTranspose2d(64, out_channels, 5, 3, 2) | |
| else: | |
| self.conv_bottom = nn.Conv2d(64, out_channels, 3, 1, 0) | |
| for m in self.modules(): | |
| if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)): | |
| nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") | |
| elif isinstance(m, nn.Linear): | |
| nn.init.normal_(m.weight, 0, 0.01) | |
| if m.bias is not None: | |
| nn.init.constant_(m.bias, 0) | |
| def forward(self, x): | |
| x1 = self.conv1(x) | |
| x2 = self.conv1_down(x1) | |
| x2 = F.leaky_relu(x2, 0.1, inplace=True) | |
| x2 = self.conv2(x2) | |
| x2 = self.conv2_up(x2) | |
| x2 = F.leaky_relu(x2, 0.1, inplace=True) | |
| x1 = F.pad(x1, (-4, -4, -4, -4)) | |
| x3 = self.conv3(x1 + x2) | |
| x3 = F.leaky_relu(x3, 0.1, inplace=True) | |
| z = self.conv_bottom(x3) | |
| return z | |
| def forward_a(self, x): | |
| x1 = self.conv1(x) | |
| x2 = self.conv1_down(x1) | |
| x2 = F.leaky_relu(x2, 0.1, inplace=True) | |
| x2 = self.conv2.conv(x2) | |
| return x1, x2 | |
| def forward_b(self, x1, x2): | |
| x2 = self.conv2_up(x2) | |
| x2 = F.leaky_relu(x2, 0.1, inplace=True) | |
| x1 = F.pad(x1, (-4, -4, -4, -4)) | |
| x3 = self.conv3(x1 + x2) | |
| x3 = F.leaky_relu(x3, 0.1, inplace=True) | |
| z = self.conv_bottom(x3) | |
| return z | |
| class UNet2(nn.Module): | |
| def __init__(self, in_channels, out_channels, deconv): | |
| super(UNet2, self).__init__() | |
| self.conv1 = UNetConv(in_channels, 32, 64, se=False) | |
| self.conv1_down = nn.Conv2d(64, 64, 2, 2, 0) | |
| self.conv2 = UNetConv(64, 64, 128, se=True) | |
| self.conv2_down = nn.Conv2d(128, 128, 2, 2, 0) | |
| self.conv3 = UNetConv(128, 256, 128, se=True) | |
| self.conv3_up = nn.ConvTranspose2d(128, 128, 2, 2, 0) | |
| self.conv4 = UNetConv(128, 64, 64, se=True) | |
| self.conv4_up = nn.ConvTranspose2d(64, 64, 2, 2, 0) | |
| self.conv5 = nn.Conv2d(64, 64, 3, 1, 0) | |
| if deconv: | |
| self.conv_bottom = nn.ConvTranspose2d(64, out_channels, 4, 2, 3) | |
| else: | |
| self.conv_bottom = nn.Conv2d(64, out_channels, 3, 1, 0) | |
| for m in self.modules(): | |
| if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)): | |
| nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") | |
| elif isinstance(m, nn.Linear): | |
| nn.init.normal_(m.weight, 0, 0.01) | |
| if m.bias is not None: | |
| nn.init.constant_(m.bias, 0) | |
| def forward(self, x): | |
| x1 = self.conv1(x) | |
| x2 = self.conv1_down(x1) | |
| x2 = F.leaky_relu(x2, 0.1, inplace=True) | |
| x2 = self.conv2(x2) | |
| x3 = self.conv2_down(x2) | |
| x3 = F.leaky_relu(x3, 0.1, inplace=True) | |
| x3 = self.conv3(x3) | |
| x3 = self.conv3_up(x3) | |
| x3 = F.leaky_relu(x3, 0.1, inplace=True) | |
| x2 = F.pad(x2, (-4, -4, -4, -4)) | |
| x4 = self.conv4(x2 + x3) | |
| x4 = self.conv4_up(x4) | |
| x4 = F.leaky_relu(x4, 0.1, inplace=True) | |
| x1 = F.pad(x1, (-16, -16, -16, -16)) | |
| x5 = self.conv5(x1 + x4) | |
| x5 = F.leaky_relu(x5, 0.1, inplace=True) | |
| z = self.conv_bottom(x5) | |
| return z | |
| def forward_a(self, x): # conv234结尾有se | |
| x1 = self.conv1(x) | |
| x2 = self.conv1_down(x1) | |
| x2 = F.leaky_relu(x2, 0.1, inplace=True) | |
| x2 = self.conv2.conv(x2) | |
| return x1, x2 | |
| def forward_b(self, x2): # conv234结尾有se | |
| x3 = self.conv2_down(x2) | |
| x3 = F.leaky_relu(x3, 0.1, inplace=True) | |
| x3 = self.conv3.conv(x3) | |
| return x3 | |
| def forward_c(self, x2, x3): # conv234结尾有se | |
| x3 = self.conv3_up(x3) | |
| x3 = F.leaky_relu(x3, 0.1, inplace=True) | |
| x2 = F.pad(x2, (-4, -4, -4, -4)) | |
| x4 = self.conv4.conv(x2 + x3) | |
| return x4 | |
| def forward_d(self, x1, x4): # conv234结尾有se | |
| x4 = self.conv4_up(x4) | |
| x4 = F.leaky_relu(x4, 0.1, inplace=True) | |
| x1 = F.pad(x1, (-16, -16, -16, -16)) | |
| x5 = self.conv5(x1 + x4) | |
| x5 = F.leaky_relu(x5, 0.1, inplace=True) | |
| z = self.conv_bottom(x5) | |
| return z | |
| class UpCunet2x(nn.Module): # 完美tile，全程无损 | |
| def __init__(self, in_channels=3, out_channels=3): | |
| super(UpCunet2x, self).__init__() | |
| self.unet1 = UNet1(in_channels, out_channels, deconv=True) | |
| self.unet2 = UNet2(in_channels, out_channels, deconv=False) | |
| def forward(self, x, tile_mode): # 1.7G | |
| n, c, h0, w0 = x.shape | |
| if tile_mode == 0: # 不tile | |
| ph = ((h0 - 1) // 2 + 1) * 2 | |
| pw = ((w0 - 1) // 2 + 1) * 2 | |
| x = F.pad(x, (18, 18 + pw - w0, 18, 18 + ph - h0), "reflect") # 需要保证被2整除 | |
| x = self.unet1.forward(x) | |
| x0 = self.unet2.forward(x) | |
| x1 = F.pad(x, (-20, -20, -20, -20)) | |
| x = torch.add(x0, x1) | |
| if w0 != pw or h0 != ph: | |
| x = x[:, :, : h0 * 2, : w0 * 2] | |
| return x | |
| elif tile_mode == 1: # 对长边减半 | |
| if w0 >= h0: | |
| crop_size_w = ((w0 - 1) // 4 * 4 + 4) // 2 # 减半后能被2整除，所以要先被4整除 | |
| crop_size_h = (h0 - 1) // 2 * 2 + 2 # 能被2整除 | |
| else: | |
| crop_size_h = ((h0 - 1) // 4 * 4 + 4) // 2 # 减半后能被2整除，所以要先被4整除 | |
| crop_size_w = (w0 - 1) // 2 * 2 + 2 # 能被2整除 | |
| crop_size = (crop_size_h, crop_size_w) # 6.6G | |
| elif tile_mode == 2: # hw都减半 | |
| crop_size = ( | |
| ((h0 - 1) // 4 * 4 + 4) // 2, | |
| ((w0 - 1) // 4 * 4 + 4) // 2, | |
| ) # 5.6G | |
| elif tile_mode == 3: # hw都三分之一 | |
| crop_size = ( | |
| ((h0 - 1) // 6 * 6 + 6) // 3, | |
| ((w0 - 1) // 6 * 6 + 6) // 3, | |
| ) # 4.2G | |
| elif tile_mode == 4: # hw都四分之一 | |
| crop_size = ( | |
| ((h0 - 1) // 8 * 8 + 8) // 4, | |
| ((w0 - 1) // 8 * 8 + 8) // 4, | |
| ) # 3.7G | |
| ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0] | |
| pw = ((w0 - 1) // crop_size[1] + 1) * crop_size[1] | |
| x = F.pad(x, (18, 18 + pw - w0, 18, 18 + ph - h0), "reflect") | |
| n, c, h, w = x.shape | |
| se_mean0 = torch.zeros((n, 64, 1, 1)).to(x.device) | |
| if "Half" in x.type(): | |
| se_mean0 = se_mean0.half() | |
| n_patch = 0 | |
| tmp_dict = {} | |
| opt_res_dict = {} | |
| for i in range(0, h - 36, crop_size[0]): | |
| tmp_dict[i] = {} | |
| for j in range(0, w - 36, crop_size[1]): | |
| x_crop = x[:, :, i : i + crop_size[0] + 36, j : j + crop_size[1] + 36] | |
| n, c1, h1, w1 = x_crop.shape | |
| tmp0, x_crop = self.unet1.forward_a(x_crop) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| x_crop.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(x_crop, dim=(2, 3), keepdim=True) | |
| se_mean0 += tmp_se_mean | |
| n_patch += 1 | |
| tmp_dict[i][j] = (tmp0, x_crop) | |
| se_mean0 /= n_patch | |
| se_mean1 = torch.zeros((n, 128, 1, 1)).to(x.device) # 64#128#128#64 | |
| if "Half" in x.type(): | |
| se_mean1 = se_mean1.half() | |
| for i in range(0, h - 36, crop_size[0]): | |
| for j in range(0, w - 36, crop_size[1]): | |
| tmp0, x_crop = tmp_dict[i][j] | |
| x_crop = self.unet1.conv2.seblock.forward_mean(x_crop, se_mean0) | |
| opt_unet1 = self.unet1.forward_b(tmp0, x_crop) | |
| tmp_x1, tmp_x2 = self.unet2.forward_a(opt_unet1) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| tmp_x2.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(tmp_x2, dim=(2, 3), keepdim=True) | |
| se_mean1 += tmp_se_mean | |
| tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2) | |
| se_mean1 /= n_patch | |
| se_mean0 = torch.zeros((n, 128, 1, 1)).to(x.device) # 64#128#128#64 | |
| if "Half" in x.type(): | |
| se_mean0 = se_mean0.half() | |
| for i in range(0, h - 36, crop_size[0]): | |
| for j in range(0, w - 36, crop_size[1]): | |
| opt_unet1, tmp_x1, tmp_x2 = tmp_dict[i][j] | |
| tmp_x2 = self.unet2.conv2.seblock.forward_mean(tmp_x2, se_mean1) | |
| tmp_x3 = self.unet2.forward_b(tmp_x2) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| tmp_x3.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(tmp_x3, dim=(2, 3), keepdim=True) | |
| se_mean0 += tmp_se_mean | |
| tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2, tmp_x3) | |
| se_mean0 /= n_patch | |
| se_mean1 = torch.zeros((n, 64, 1, 1)).to(x.device) # 64#128#128#64 | |
| if "Half" in x.type(): | |
| se_mean1 = se_mean1.half() | |
| for i in range(0, h - 36, crop_size[0]): | |
| for j in range(0, w - 36, crop_size[1]): | |
| opt_unet1, tmp_x1, tmp_x2, tmp_x3 = tmp_dict[i][j] | |
| tmp_x3 = self.unet2.conv3.seblock.forward_mean(tmp_x3, se_mean0) | |
| tmp_x4 = self.unet2.forward_c(tmp_x2, tmp_x3) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| tmp_x4.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(tmp_x4, dim=(2, 3), keepdim=True) | |
| se_mean1 += tmp_se_mean | |
| tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x4) | |
| se_mean1 /= n_patch | |
| for i in range(0, h - 36, crop_size[0]): | |
| opt_res_dict[i] = {} | |
| for j in range(0, w - 36, crop_size[1]): | |
| opt_unet1, tmp_x1, tmp_x4 = tmp_dict[i][j] | |
| tmp_x4 = self.unet2.conv4.seblock.forward_mean(tmp_x4, se_mean1) | |
| x0 = self.unet2.forward_d(tmp_x1, tmp_x4) | |
| x1 = F.pad(opt_unet1, (-20, -20, -20, -20)) | |
| x_crop = torch.add(x0, x1) # x0是unet2的最终输出 | |
| opt_res_dict[i][j] = x_crop | |
| del tmp_dict | |
| torch.cuda.empty_cache() | |
| res = torch.zeros((n, c, h * 2 - 72, w * 2 - 72)).to(x.device) | |
| if "Half" in x.type(): | |
| res = res.half() | |
| for i in range(0, h - 36, crop_size[0]): | |
| for j in range(0, w - 36, crop_size[1]): | |
| res[ | |
| :, :, i * 2 : i * 2 + h1 * 2 - 72, j * 2 : j * 2 + w1 * 2 - 72 | |
| ] = opt_res_dict[i][j] | |
| del opt_res_dict | |
| torch.cuda.empty_cache() | |
| if w0 != pw or h0 != ph: | |
| res = res[:, :, : h0 * 2, : w0 * 2] | |
| return res # | |
| class UpCunet3x(nn.Module): # 完美tile，全程无损 | |
| def __init__(self, in_channels=3, out_channels=3): | |
| super(UpCunet3x, self).__init__() | |
| self.unet1 = UNet1x3(in_channels, out_channels, deconv=True) | |
| self.unet2 = UNet2(in_channels, out_channels, deconv=False) | |
| def forward(self, x, tile_mode): # 1.7G | |
| n, c, h0, w0 = x.shape | |
| if tile_mode == 0: # 不tile | |
| ph = ((h0 - 1) // 4 + 1) * 4 | |
| pw = ((w0 - 1) // 4 + 1) * 4 | |
| x = F.pad(x, (14, 14 + pw - w0, 14, 14 + ph - h0), "reflect") # 需要保证被2整除 | |
| x = self.unet1.forward(x) | |
| x0 = self.unet2.forward(x) | |
| x1 = F.pad(x, (-20, -20, -20, -20)) | |
| x = torch.add(x0, x1) | |
| if w0 != pw or h0 != ph: | |
| x = x[:, :, : h0 * 3, : w0 * 3] | |
| return x | |
| elif tile_mode == 1: # 对长边减半 | |
| if w0 >= h0: | |
| crop_size_w = ((w0 - 1) // 8 * 8 + 8) // 2 # 减半后能被4整除，所以要先被8整除 | |
| crop_size_h = (h0 - 1) // 4 * 4 + 4 # 能被4整除 | |
| else: | |
| crop_size_h = ((h0 - 1) // 8 * 8 + 8) // 2 # 减半后能被4整除，所以要先被8整除 | |
| crop_size_w = (w0 - 1) // 4 * 4 + 4 # 能被4整除 | |
| crop_size = (crop_size_h, crop_size_w) # 6.6G | |
| elif tile_mode == 2: # hw都减半 | |
| crop_size = ( | |
| ((h0 - 1) // 8 * 8 + 8) // 2, | |
| ((w0 - 1) // 8 * 8 + 8) // 2, | |
| ) # 5.6G | |
| elif tile_mode == 3: # hw都三分之一 | |
| crop_size = ( | |
| ((h0 - 1) // 12 * 12 + 12) // 3, | |
| ((w0 - 1) // 12 * 12 + 12) // 3, | |
| ) # 4.2G | |
| elif tile_mode == 4: # hw都四分之一 | |
| crop_size = ( | |
| ((h0 - 1) // 16 * 16 + 16) // 4, | |
| ((w0 - 1) // 16 * 16 + 16) // 4, | |
| ) # 3.7G | |
| ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0] | |
| pw = ((w0 - 1) // crop_size[1] + 1) * crop_size[1] | |
| x = F.pad(x, (14, 14 + pw - w0, 14, 14 + ph - h0), "reflect") | |
| n, c, h, w = x.shape | |
| se_mean0 = torch.zeros((n, 64, 1, 1)).to(x.device) | |
| if "Half" in x.type(): | |
| se_mean0 = se_mean0.half() | |
| n_patch = 0 | |
| tmp_dict = {} | |
| opt_res_dict = {} | |
| for i in range(0, h - 28, crop_size[0]): | |
| tmp_dict[i] = {} | |
| for j in range(0, w - 28, crop_size[1]): | |
| x_crop = x[:, :, i : i + crop_size[0] + 28, j : j + crop_size[1] + 28] | |
| n, c1, h1, w1 = x_crop.shape | |
| tmp0, x_crop = self.unet1.forward_a(x_crop) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| x_crop.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(x_crop, dim=(2, 3), keepdim=True) | |
| se_mean0 += tmp_se_mean | |
| n_patch += 1 | |
| tmp_dict[i][j] = (tmp0, x_crop) | |
| se_mean0 /= n_patch | |
| se_mean1 = torch.zeros((n, 128, 1, 1)).to(x.device) # 64#128#128#64 | |
| if "Half" in x.type(): | |
| se_mean1 = se_mean1.half() | |
| for i in range(0, h - 28, crop_size[0]): | |
| for j in range(0, w - 28, crop_size[1]): | |
| tmp0, x_crop = tmp_dict[i][j] | |
| x_crop = self.unet1.conv2.seblock.forward_mean(x_crop, se_mean0) | |
| opt_unet1 = self.unet1.forward_b(tmp0, x_crop) | |
| tmp_x1, tmp_x2 = self.unet2.forward_a(opt_unet1) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| tmp_x2.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(tmp_x2, dim=(2, 3), keepdim=True) | |
| se_mean1 += tmp_se_mean | |
| tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2) | |
| se_mean1 /= n_patch | |
| se_mean0 = torch.zeros((n, 128, 1, 1)).to(x.device) # 64#128#128#64 | |
| if "Half" in x.type(): | |
| se_mean0 = se_mean0.half() | |
| for i in range(0, h - 28, crop_size[0]): | |
| for j in range(0, w - 28, crop_size[1]): | |
| opt_unet1, tmp_x1, tmp_x2 = tmp_dict[i][j] | |
| tmp_x2 = self.unet2.conv2.seblock.forward_mean(tmp_x2, se_mean1) | |
| tmp_x3 = self.unet2.forward_b(tmp_x2) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| tmp_x3.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(tmp_x3, dim=(2, 3), keepdim=True) | |
| se_mean0 += tmp_se_mean | |
| tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2, tmp_x3) | |
| se_mean0 /= n_patch | |
| se_mean1 = torch.zeros((n, 64, 1, 1)).to(x.device) # 64#128#128#64 | |
| if "Half" in x.type(): | |
| se_mean1 = se_mean1.half() | |
| for i in range(0, h - 28, crop_size[0]): | |
| for j in range(0, w - 28, crop_size[1]): | |
| opt_unet1, tmp_x1, tmp_x2, tmp_x3 = tmp_dict[i][j] | |
| tmp_x3 = self.unet2.conv3.seblock.forward_mean(tmp_x3, se_mean0) | |
| tmp_x4 = self.unet2.forward_c(tmp_x2, tmp_x3) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| tmp_x4.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(tmp_x4, dim=(2, 3), keepdim=True) | |
| se_mean1 += tmp_se_mean | |
| tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x4) | |
| se_mean1 /= n_patch | |
| for i in range(0, h - 28, crop_size[0]): | |
| opt_res_dict[i] = {} | |
| for j in range(0, w - 28, crop_size[1]): | |
| opt_unet1, tmp_x1, tmp_x4 = tmp_dict[i][j] | |
| tmp_x4 = self.unet2.conv4.seblock.forward_mean(tmp_x4, se_mean1) | |
| x0 = self.unet2.forward_d(tmp_x1, tmp_x4) | |
| x1 = F.pad(opt_unet1, (-20, -20, -20, -20)) | |
| x_crop = torch.add(x0, x1) # x0是unet2的最终输出 | |
| opt_res_dict[i][j] = x_crop # | |
| del tmp_dict | |
| torch.cuda.empty_cache() | |
| res = torch.zeros((n, c, h * 3 - 84, w * 3 - 84)).to(x.device) | |
| if "Half" in x.type(): | |
| res = res.half() | |
| for i in range(0, h - 28, crop_size[0]): | |
| for j in range(0, w - 28, crop_size[1]): | |
| res[ | |
| :, :, i * 3 : i * 3 + h1 * 3 - 84, j * 3 : j * 3 + w1 * 3 - 84 | |
| ] = opt_res_dict[i][j] | |
| del opt_res_dict | |
| torch.cuda.empty_cache() | |
| if w0 != pw or h0 != ph: | |
| res = res[:, :, : h0 * 3, : w0 * 3] | |
| return res | |
| class UpCunet4x(nn.Module): # 完美tile，全程无损 | |
| def __init__(self, in_channels=3, out_channels=3): | |
| super(UpCunet4x, self).__init__() | |
| self.unet1 = UNet1(in_channels, 64, deconv=True) | |
| self.unet2 = UNet2(64, 64, deconv=False) | |
| self.ps = nn.PixelShuffle(2) | |
| self.conv_final = nn.Conv2d(64, 12, 3, 1, padding=0, bias=True) | |
| def forward(self, x, tile_mode): | |
| n, c, h0, w0 = x.shape | |
| x00 = x | |
| if tile_mode == 0: # 不tile | |
| ph = ((h0 - 1) // 2 + 1) * 2 | |
| pw = ((w0 - 1) // 2 + 1) * 2 | |
| x = F.pad(x, (19, 19 + pw - w0, 19, 19 + ph - h0), "reflect") # 需要保证被2整除 | |
| x = self.unet1.forward(x) | |
| x0 = self.unet2.forward(x) | |
| x1 = F.pad(x, (-20, -20, -20, -20)) | |
| x = torch.add(x0, x1) | |
| x = self.conv_final(x) | |
| x = F.pad(x, (-1, -1, -1, -1)) | |
| x = self.ps(x) | |
| if w0 != pw or h0 != ph: | |
| x = x[:, :, : h0 * 4, : w0 * 4] | |
| x += F.interpolate(x00, scale_factor=4, mode="nearest") | |
| return x | |
| elif tile_mode == 1: # 对长边减半 | |
| if w0 >= h0: | |
| crop_size_w = ((w0 - 1) // 4 * 4 + 4) // 2 # 减半后能被2整除，所以要先被4整除 | |
| crop_size_h = (h0 - 1) // 2 * 2 + 2 # 能被2整除 | |
| else: | |
| crop_size_h = ((h0 - 1) // 4 * 4 + 4) // 2 # 减半后能被2整除，所以要先被4整除 | |
| crop_size_w = (w0 - 1) // 2 * 2 + 2 # 能被2整除 | |
| crop_size = (crop_size_h, crop_size_w) # 6.6G | |
| elif tile_mode == 2: # hw都减半 | |
| crop_size = ( | |
| ((h0 - 1) // 4 * 4 + 4) // 2, | |
| ((w0 - 1) // 4 * 4 + 4) // 2, | |
| ) # 5.6G | |
| elif tile_mode == 3: # hw都三分之一 | |
| crop_size = ( | |
| ((h0 - 1) // 6 * 6 + 6) // 3, | |
| ((w0 - 1) // 6 * 6 + 6) // 3, | |
| ) # 4.1G | |
| elif tile_mode == 4: # hw都四分之一 | |
| crop_size = ( | |
| ((h0 - 1) // 8 * 8 + 8) // 4, | |
| ((w0 - 1) // 8 * 8 + 8) // 4, | |
| ) # 3.7G | |
| ph = ((h0 - 1) // crop_size[0] + 1) * crop_size[0] | |
| pw = ((w0 - 1) // crop_size[1] + 1) * crop_size[1] | |
| x = F.pad(x, (19, 19 + pw - w0, 19, 19 + ph - h0), "reflect") | |
| n, c, h, w = x.shape | |
| se_mean0 = torch.zeros((n, 64, 1, 1)).to(x.device) | |
| if "Half" in x.type(): | |
| se_mean0 = se_mean0.half() | |
| n_patch = 0 | |
| tmp_dict = {} | |
| opt_res_dict = {} | |
| for i in range(0, h - 38, crop_size[0]): | |
| tmp_dict[i] = {} | |
| for j in range(0, w - 38, crop_size[1]): | |
| x_crop = x[:, :, i : i + crop_size[0] + 38, j : j + crop_size[1] + 38] | |
| n, c1, h1, w1 = x_crop.shape | |
| tmp0, x_crop = self.unet1.forward_a(x_crop) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| x_crop.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(x_crop, dim=(2, 3), keepdim=True) | |
| se_mean0 += tmp_se_mean | |
| n_patch += 1 | |
| tmp_dict[i][j] = (tmp0, x_crop) | |
| se_mean0 /= n_patch | |
| se_mean1 = torch.zeros((n, 128, 1, 1)).to(x.device) # 64#128#128#64 | |
| if "Half" in x.type(): | |
| se_mean1 = se_mean1.half() | |
| for i in range(0, h - 38, crop_size[0]): | |
| for j in range(0, w - 38, crop_size[1]): | |
| tmp0, x_crop = tmp_dict[i][j] | |
| x_crop = self.unet1.conv2.seblock.forward_mean(x_crop, se_mean0) | |
| opt_unet1 = self.unet1.forward_b(tmp0, x_crop) | |
| tmp_x1, tmp_x2 = self.unet2.forward_a(opt_unet1) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| tmp_x2.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(tmp_x2, dim=(2, 3), keepdim=True) | |
| se_mean1 += tmp_se_mean | |
| tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2) | |
| se_mean1 /= n_patch | |
| se_mean0 = torch.zeros((n, 128, 1, 1)).to(x.device) # 64#128#128#64 | |
| if "Half" in x.type(): | |
| se_mean0 = se_mean0.half() | |
| for i in range(0, h - 38, crop_size[0]): | |
| for j in range(0, w - 38, crop_size[1]): | |
| opt_unet1, tmp_x1, tmp_x2 = tmp_dict[i][j] | |
| tmp_x2 = self.unet2.conv2.seblock.forward_mean(tmp_x2, se_mean1) | |
| tmp_x3 = self.unet2.forward_b(tmp_x2) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| tmp_x3.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(tmp_x3, dim=(2, 3), keepdim=True) | |
| se_mean0 += tmp_se_mean | |
| tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x2, tmp_x3) | |
| se_mean0 /= n_patch | |
| se_mean1 = torch.zeros((n, 64, 1, 1)).to(x.device) # 64#128#128#64 | |
| if "Half" in x.type(): | |
| se_mean1 = se_mean1.half() | |
| for i in range(0, h - 38, crop_size[0]): | |
| for j in range(0, w - 38, crop_size[1]): | |
| opt_unet1, tmp_x1, tmp_x2, tmp_x3 = tmp_dict[i][j] | |
| tmp_x3 = self.unet2.conv3.seblock.forward_mean(tmp_x3, se_mean0) | |
| tmp_x4 = self.unet2.forward_c(tmp_x2, tmp_x3) | |
| if "Half" in x.type(): # torch.HalfTensor/torch.cuda.HalfTensor | |
| tmp_se_mean = torch.mean( | |
| tmp_x4.float(), dim=(2, 3), keepdim=True | |
| ).half() | |
| else: | |
| tmp_se_mean = torch.mean(tmp_x4, dim=(2, 3), keepdim=True) | |
| se_mean1 += tmp_se_mean | |
| tmp_dict[i][j] = (opt_unet1, tmp_x1, tmp_x4) | |
| se_mean1 /= n_patch | |
| for i in range(0, h - 38, crop_size[0]): | |
| opt_res_dict[i] = {} | |
| for j in range(0, w - 38, crop_size[1]): | |
| opt_unet1, tmp_x1, tmp_x4 = tmp_dict[i][j] | |
| tmp_x4 = self.unet2.conv4.seblock.forward_mean(tmp_x4, se_mean1) | |
| x0 = self.unet2.forward_d(tmp_x1, tmp_x4) | |
| x1 = F.pad(opt_unet1, (-20, -20, -20, -20)) | |
| x_crop = torch.add(x0, x1) # x0是unet2的最终输出 | |
| x_crop = self.conv_final(x_crop) | |
| x_crop = F.pad(x_crop, (-1, -1, -1, -1)) | |
| x_crop = self.ps(x_crop) | |
| opt_res_dict[i][j] = x_crop | |
| del tmp_dict | |
| torch.cuda.empty_cache() | |
| res = torch.zeros((n, c, h * 4 - 152, w * 4 - 152)).to(x.device) | |
| if "Half" in x.type(): | |
| res = res.half() | |
| for i in range(0, h - 38, crop_size[0]): | |
| for j in range(0, w - 38, crop_size[1]): | |
| # print(opt_res_dict[i][j].shape,res[:, :, i * 4:i * 4 + h1 * 4 - 144, j * 4:j * 4 + w1 * 4 - 144].shape) | |
| res[ | |
| :, :, i * 4 : i * 4 + h1 * 4 - 152, j * 4 : j * 4 + w1 * 4 - 152 | |
| ] = opt_res_dict[i][j] | |
| del opt_res_dict | |
| torch.cuda.empty_cache() | |
| if w0 != pw or h0 != ph: | |
| res = res[:, :, : h0 * 4, : w0 * 4] | |
| res += F.interpolate(x00, scale_factor=4, mode="nearest") | |
| return res # | |
| models: Dict[str, Type[nn.Module]] = { | |
| obj.__name__: obj | |
| for obj in globals().values() | |
| if isinstance(obj, type) and issubclass(obj, nn.Module) | |
| } | |
| class RealWaifuUpScaler: | |
| def __init__(self, scale: int, weight_path: str, half: bool, device: str): | |
| weight = torch.load(weight_path, map_location=device) | |
| self.model = models[f"UpCunet{scale}x"]() | |
| if half == True: | |
| self.model = self.model.half().to(device) | |
| else: | |
| self.model = self.model.to(device) | |
| self.model.load_state_dict(weight, strict=True) | |
| self.model.eval() | |
| self.half = half | |
| self.device = device | |
| def np2tensor(self, np_frame): | |
| if self.half == False: | |
| return ( | |
| torch.from_numpy(np.transpose(np_frame, (2, 0, 1))) | |
| .unsqueeze(0) | |
| .to(self.device) | |
| .float() | |
| / 255 | |
| ) | |
| else: | |
| return ( | |
| torch.from_numpy(np.transpose(np_frame, (2, 0, 1))) | |
| .unsqueeze(0) | |
| .to(self.device) | |
| .half() | |
| / 255 | |
| ) | |
| def tensor2np(self, tensor): | |
| if self.half == False: | |
| return np.transpose( | |
| (tensor.data.squeeze() * 255.0) | |
| .round() | |
| .clamp_(0, 255) | |
| .byte() | |
| .cpu() | |
| .numpy(), | |
| (1, 2, 0), | |
| ) | |
| else: | |
| return np.transpose( | |
| (tensor.data.squeeze().float() * 255.0) | |
| .round() | |
| .clamp_(0, 255) | |
| .byte() | |
| .cpu() | |
| .numpy(), | |
| (1, 2, 0), | |
| ) | |
| def __call__(self, frame, tile_mode): | |
| with torch.no_grad(): | |
| tensor = self.np2tensor(frame) | |
| result = self.tensor2np(self.model(tensor, tile_mode)) | |
| return result | |
| input_image = inputs.File(label="Input image") | |
| half_precision = inputs.Checkbox( | |
| label="Half precision (NOT work for CPU)", default=False | |
| ) | |
| model_weight = inputs.Dropdown(sorted(AVALIABLE_WEIGHTS), label="Choice model weight") | |
| tile_mode = inputs.Radio([mode.name for mode in TileMode], label="Output tile mode") | |
| output_image = outputs.Image(label="Output image preview") | |
| output_file = outputs.File(label="Output image file") | |
| def main(file: IO[bytes], half: bool, weight: str, tile: str): | |
| scale = next(mode.value for mode in ScaleMode if weight.startswith(mode.name)) | |
| upscaler = RealWaifuUpScaler( | |
| scale, weight_path=str(AVALIABLE_WEIGHTS[weight]), half=half, device=DEVICE | |
| ) | |
| frame = cv2.cvtColor(cv2.imread(file.name), cv2.COLOR_BGR2RGB) | |
| result = cv2.cvtColor(upscaler(frame, TileMode[tile]), cv2.COLOR_RGB2BGR) | |
| _, ext = os.path.splitext(file.name) | |
| tempfile = mktemp(suffix=ext) | |
| cv2.imwrite(tempfile, result) | |
| return tempfile, tempfile | |
| interface = Interface( | |
| main, | |
| inputs=[input_image, half_precision, model_weight, tile_mode], | |
| outputs=[output_image, output_file], | |
| ) | |
| interface.launch() | |