From a7f7a8b677c5faafa4fd4efa5876d79e3f06cae6 Mon Sep 17 00:00:00 2001 From: FuouM Date: Fri, 2 Aug 2024 22:31:25 +0700 Subject: [PATCH] run, but doesnt seem to work that well --- README.md | 11 + __init__.py | 8 + asvr_module/ADISTS_allstage.py | 191 +++++ asvr_module/SIREN.py | 70 ++ asvr_module/asvr_model.py | 88 ++ asvr_module/blocks.py | 758 ++++++++++++++++++ asvr_module/components.py | 120 +++ asvr_module/cuf_siren.py | 253 ++++++ asvr_module/init.py | 69 ++ asvr_module/network_vidue_refsr_cuf3_siren.py | 224 ++++++ .../pytorch_pwc/correlation/correlation.py | 435 ++++++++++ asvr_module/pytorch_pwc/extract_flow.py | 32 + asvr_module/pytorch_pwc/pwc.py | 242 ++++++ .../refsrrnn_adists_fgda_only_future.py | 214 +++++ asvr_module/warp.py | 33 + asvr_run.py | 47 ++ constants.py | 3 +- 17 files changed, 2797 insertions(+), 1 deletion(-) create mode 100644 asvr_module/ADISTS_allstage.py create mode 100644 asvr_module/SIREN.py create mode 100644 asvr_module/asvr_model.py create mode 100644 asvr_module/blocks.py create mode 100644 asvr_module/components.py create mode 100644 asvr_module/cuf_siren.py create mode 100644 asvr_module/init.py create mode 100644 asvr_module/network_vidue_refsr_cuf3_siren.py create mode 100644 asvr_module/pytorch_pwc/correlation/correlation.py create mode 100644 asvr_module/pytorch_pwc/extract_flow.py create mode 100644 asvr_module/pytorch_pwc/pwc.py create mode 100644 asvr_module/refsrrnn_adists_fgda_only_future.py create mode 100644 asvr_module/warp.py create mode 100644 asvr_run.py diff --git a/README.md b/README.md index 0ef2865..b6c4d2f 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ Click name to jump to workflow 2. [**RealViformer**](#realviformer) - Paper: [Investigating Attention for Real-World Video Super-Resolution](https://github.com/Yuehan717/RealViformer) 3. [**ProPIH**](#propih). Paper: [Progressive Painterly Image Harmonization from Low-level Styles to High-level Styles](https://github.com/bcmi/ProPIH-Painterly-Image-Harmonization) 4. [**CoLIE**](#colie). Paper: [Fast Context-Based Low-Light Image Enhancement via Neural Implicit Representations](https://github.com/ctom2/colie) +5. [**ST-AVSR**](#st-avsr). Paper: [Arbitrary-Scale Video Super-Resolution with Structural and Textural Priors](https://github.com/shangwei5/ST-AVSR) ## Workflows @@ -52,3 +53,13 @@ No model needed to be downloaded. Lower `loss_mean` seems to result in brighter [workflow_colie_lowlight.json](workflow/workflow_colie_lowlight.json) ![colie_lowlight](workflow/colie_lowlight.png) + +### ST-AVSR + +You will need to install [mmcv](https://github.com/open-mmlab/mmcv), which can be a real pain. Refer to https://mmcv.readthedocs.io/en/latest/get_started/installation.html. Cupy is also required. This is tested on: + +``` +mmcv==2.2.0 +torch=2.3.0+cu121 +cupy==13.2.0 +``` \ No newline at end of file diff --git a/__init__.py b/__init__.py index 2bdd7e6..7c02db8 100644 --- a/__init__.py +++ b/__init__.py @@ -9,6 +9,7 @@ NODE_CLASS_MAPPINGS = { "CoLIE_LowLight_Enhance": CoLIE_LowLight_Enhance, } + NODE_DISPLAY_NAME_MAPPINGS = { "WFEN": "WFEN Face Restore", "RealViFormerSR": "RealViFormer Video SR", @@ -16,5 +17,12 @@ NODE_DISPLAY_NAME_MAPPINGS = { "CoLIE_LowLight_Enhance": "CoLIE LowLight Enhance", } +try: + from .asvr_run import ASVR_VideoSR + NODE_CLASS_MAPPINGS["ASVR_VideoSR"] = ASVR_VideoSR + NODE_DISPLAY_NAME_MAPPINGS["ASVR_VideoSR"] = "ASVR Video SR" +except ImportError as e: + print(f"Failed to load ASVR_VideoSR. This will not affect other nodes. {e}") + pass __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/asvr_module/ADISTS_allstage.py b/asvr_module/ADISTS_allstage.py new file mode 100644 index 0000000..9db0634 --- /dev/null +++ b/asvr_module/ADISTS_allstage.py @@ -0,0 +1,191 @@ +import torch +import torch.nn as nn +import numpy as np +from torchvision import models +import torch.nn.functional as F + +class Downsample(nn.Module): + def __init__(self, filter_size=5, stride=2, channels=None, pad_off=0): + super(Downsample, self).__init__() + self.padding = (filter_size - 2) // 2 + self.stride = stride + self.channels = channels + a = np.hanning(filter_size)[1:-1] + g = torch.Tensor(a[:, None] * a[None, :]) + g = g / torch.sum(g) + self.register_buffer('filter', g[None, None, :, :].repeat((self.channels, 1, 1, 1))) + + def forward(self, input): + input = input ** 2 + out = F.conv2d(input, self.filter, stride=self.stride, padding=self.padding, groups=input.shape[1]) + return (out + 1e-12).sqrt() + + +class ADISTS(torch.nn.Module): + def __init__(self, window_size=21): + super(ADISTS, self).__init__() + vgg_pretrained_features = models.vgg16(pretrained=True).features + self.stage1 = torch.nn.Sequential() + self.stage2 = torch.nn.Sequential() + self.stage3 = torch.nn.Sequential() + self.stage4 = torch.nn.Sequential() + self.stage5 = torch.nn.Sequential() + for x in range(0, 4): + self.stage1.add_module(str(x), vgg_pretrained_features[x]) + self.stage2.add_module(str(4), Downsample(channels=64)) + for x in range(5, 9): + self.stage2.add_module(str(x), vgg_pretrained_features[x]) + self.stage3.add_module(str(9), Downsample(channels=128)) + for x in range(10, 16): + self.stage3.add_module(str(x), vgg_pretrained_features[x]) + self.stage4.add_module(str(16), Downsample(channels=256)) + for x in range(17, 23): + self.stage4.add_module(str(x), vgg_pretrained_features[x]) + self.stage5.add_module(str(23), Downsample(channels=512)) + for x in range(24, 30): + self.stage5.add_module(str(x), vgg_pretrained_features[x]) + + for param in self.parameters(): + param.requires_grad = False + + self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, -1, 1, 1)) + self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1, -1, 1, 1)) + + # self.chns = [3, 64] #[3, 64, 128, 256, 512, 512] + # self.windows = nn.ParameterList() + # self.window_size = window_size + # for k in range(len(self.chns)): + # self.windows.append(self.create_window(self.window_size, self.window_size / 3, self.chns[k])) + + def forward_once(self, x): + h_list = [] + h = (x - self.mean) / self.std + h = self.stage1(h) + _, _, H, W = h.size() + h_list.append(h) + h = self.stage2(h) + h_list.append(F.interpolate(h, (H, W), mode="bilinear")) + h = self.stage3(h) + h_list.append(F.interpolate(h, (H, W), mode="bilinear")) + h = self.stage4(h) + h_list.append(F.interpolate(h, (H, W), mode="bilinear")) + h = self.stage5(h) + h_list.append(F.interpolate(h, (H, W), mode="bilinear")) + + return torch.cat(h_list, dim=1) #B, 64+12+256+512+512, H, W + + def std_(self, img, window_size=3): + assert window_size % 2 == 1 + pad = window_size // 2 + + # calculate std on the mean image of the color channels + # img = torch.mean(img, dim=1, keepdim=True) + N, C, H, W = img.shape + img = nn.functional.pad(img, [pad] * 4, mode='reflect') + img = nn.functional.unfold(img, kernel_size=window_size) + img = img.view(N, C, window_size * window_size, H, W) + img = img - torch.mean(img, dim=2, keepdim=True) + # print(img.shape) + img = img * img + img = torch.mean(img, dim=2, keepdim=True) + # print(img.shape) + img = torch.sqrt(img) + # print(img.shape) + img = img.squeeze(2) + return img + + def generate_alpha(self,input, lower=1, upper=5): + N, C, H, W = input.shape + ratio = input.new_ones((N, C, H, W)) * 0.5 + input_std = self.std_(input) + ratio[input_std < lower] = torch.sigmoid((input_std - lower))[input_std < lower] + ratio[input_std > upper] = torch.sigmoid((input_std - upper))[input_std > upper] + ratio = ratio #.detach() + + return ratio + + def forward(self, x): + + # with torch.no_grad(): + feats_x = self.forward_once(x) + + # print(feats_x.shape) + # feats_x = self.generate_alpha(feats_x*255) + # pad = nn.ReflectionPad2d(0) + # ps_x = self.compute_prob(feats_x) + # # print(ps_x[1].shape) + # weight = [] + # for k in range(0, len(self.chns)): + # weight.append(self.entropy(feats_x[k])) + # weight = torch.concat(weight, dim=1) + # # print(weight.shape) + # + # weight = weight / weight.sum(dim=(1, 2), keepdim=True) + # weight_mean = weight.mean(dim=(1, 2), keepdim=True) + # weight_std = torch.sqrt(((weight - weight_mean) ** 2).mean(dim=(1, 2), keepdim=True)) + # weight = weight.clamp(min=weight_mean - 0.5 * weight_std, max=weight_mean + 0.5 * weight_std) + # weight = weight / weight.sum(dim=(1, 2), keepdim=True) + # weight_list = torch.split(weight, self.chns, dim=1) + # # print(weight_list[1].shape) + # + # # for k in range(len(self.chns) - 1, -1, -1): + # feat_x = F.normalize(feats_x[1], dim=(2, 3)) + # try: + # x_mean = F.conv2d(pad(feat_x), self.windows[1], stride=1, padding=0, groups=self.chns[1]) + # # y_mean = F.conv2d(pad(feat_y), self.windows[k], stride=1, padding=0, groups=self.chns[k]) + # x_var = F.conv2d(pad(feat_x ** 2), self.windows[1], stride=1, padding=0, + # groups=self.chns[1]) - x_mean ** 2 + # # y_var = F.conv2d(pad(feat_y ** 2), self.windows[k], stride=1, padding=0, + # # groups=self.chns[k]) - y_mean ** 2 + # xy_cov = F.conv2d(pad(feat_x * feat_x), self.windows[1], stride=1, padding=0, groups=self.chns[1]) \ + # - x_mean * x_mean + # except: + # x_mean = feat_x.mean([2, 3], keepdim=True) + # # y_mean = feat_y.mean([2, 3], keepdim=True) + # x_var = ((feat_x - x_mean) ** 2).mean([2, 3], keepdim=True) + # # y_var = ((feat_y - y_mean) ** 2).mean([2, 3], keepdim=True) + # xy_cov = (feat_x * feat_x).mean([2, 3], keepdim=True) - x_mean * x_mean + # + # T = (2 * x_mean * x_mean + 1e-6) / (x_mean ** 2 + x_mean ** 2 + 1e-6) + # S = (2 * xy_cov + 1e-6) / (x_var + x_var + 1e-6) + # S = F.interpolate(S, size=(ps_x[1].shape[-2], ps_x[1].shape[-1]), mode='bilinear', align_corners=True) + # T = F.interpolate(T, size=(ps_x[1].shape[-2], ps_x[1].shape[-1]), mode='bilinear', align_corners=True) + # # print(S.shape, T.shape, ps_x[1].shape) + # ps = ps_x[1] #.expand(x_mean.shape[0], x_mean.shape[1], -1, -1) + # pt = 1 - ps + # D_map = (pt * T + ps * S) * weight_list[1].unsqueeze(3) + # print(D_map.mean([2, 3])) + return feats_x + +# def prepare_image(image, resize=True): +# if resize and min(image.size) > 256: +# image = transforms.functional.resize(image, 256) +# image = transforms.ToTensor()(image) +# return image.unsqueeze(0) +# +# +# if __name__ == '__main__': +# from PIL import Image +# import argparse +# +# parser = argparse.ArgumentParser() +# parser.add_argument('--ref', type=str, default='/data1/shangwei/dataset/video/REDS/val/val_sharp_bicubic/X4/002/00000002.png') +# parser.add_argument('--dist', type=str, default='/data1/shangwei/dataset/video/REDS/val/val_sharp_bicubic/X4/002/00000003.png') +# args = parser.parse_args() +# +# ref = prepare_image(Image.open(args.ref).convert("RGB"), resize=False) +# dist = prepare_image(Image.open(args.dist).convert("RGB"), resize=False) +# assert ref.shape == dist.shape +# # print(ref.shape) +# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +# model = ADISTS().to(device) +# ref = ref.to(device) +# dist = dist.to(device) +# map = model(ref).mean(1) +# # print(score.item()) +# print(map.shape) +# # print(map.min(), map.max()) +# save = map.cpu().numpy().squeeze() * 255. +# save = save.astype(np.uint8) +# cv2.imwrite("/data1/shangwei/dataset/video/REDS/val/002_00002_lr_ratio_w21_prob_flat255.png", save) + diff --git a/asvr_module/SIREN.py b/asvr_module/SIREN.py new file mode 100644 index 0000000..41fb761 --- /dev/null +++ b/asvr_module/SIREN.py @@ -0,0 +1,70 @@ +import torch +from torch import nn +import numpy as np + +class SineLayer(nn.Module): + # See paper sec. 3.2, final paragraph, and supplement Sec. 1.5 for discussion of omega_0. + + # If is_first=True, omega_0 is a frequency factor which simply multiplies the activations before the + # nonlinearity. Different signals may require different omega_0 in the first layer - this is a + # hyperparameter. + + # If is_first=False, then the weights will be divided by omega_0 so as to keep the magnitude of + # activations constant, but boost gradients to the weight matrix (see supplement Sec. 1.5) + + def __init__(self, in_features, out_features, bias=True, + is_first=False, omega_0=30): + super().__init__() + self.omega_0 = omega_0 + self.is_first = is_first + + self.in_features = in_features + self.linear = nn.Linear(in_features, out_features, bias=bias) + + self.init_weights() + + def init_weights(self): + with torch.no_grad(): + if self.is_first: + self.linear.weight.uniform_(-1 / self.in_features, + 1 / self.in_features) + else: + self.linear.weight.uniform_(-np.sqrt(6 / self.in_features) / self.omega_0, + np.sqrt(6 / self.in_features) / self.omega_0) + + def forward(self, input): + return torch.sin(self.omega_0 * self.linear(input)) + + + +class Siren(nn.Module): + def __init__(self, in_features, hidden_features, hidden_layers, out_features, outermost_linear=False, + first_omega_0=30, hidden_omega_0=30.): + super().__init__() + + self.net = [] + self.net.append(SineLayer(in_features, hidden_features[0], + is_first=True, omega_0=first_omega_0)) + + for i in range(hidden_layers): + self.net.append(SineLayer(hidden_features[i], hidden_features[i + 1], + is_first=False, omega_0=hidden_omega_0)) + + if outermost_linear: + final_linear = nn.Linear(hidden_features[-1], out_features) + + with torch.no_grad(): + final_linear.weight.uniform_(-np.sqrt(6 / hidden_features[-1]) / hidden_omega_0, + np.sqrt(6 / hidden_features[-1]) / hidden_omega_0) + + self.net.append(final_linear) + else: + self.net.append(SineLayer(hidden_features[-1], out_features, + is_first=False, omega_0=hidden_omega_0)) + + self.net = nn.Sequential(*self.net) + + def forward(self, x): + shape = x.shape[:-1] + output = self.net(x.view(-1, x.shape[-1])) + return output.view(*shape, -1) diff --git a/asvr_module/asvr_model.py b/asvr_module/asvr_model.py new file mode 100644 index 0000000..0bb877f --- /dev/null +++ b/asvr_module/asvr_model.py @@ -0,0 +1,88 @@ +import torch +# import tqdm +import torch.nn.functional as F + +from .refsrrnn_adists_fgda_only_future import RefsrRNN + + +class ASVR_Config: + def __init__( + self, model_path: str, space_scale=(8, 8), border: bool = False + ) -> None: + self.space_scale = space_scale + self.model_path = model_path + self.border = border + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +class ASVR: + def __init__(self, cfg: ASVR_Config) -> None: + self.net = RefsrRNN() + self.net.load_state_dict(torch.load(cfg.model_path)) + self.net = self.net.to(cfg.device) + self.net.eval() + self.scale = cfg.space_scale + self.device = cfg.device + + def inference(self, frames: torch.Tensor): + # frames: [B, H, W, C] + print(f"{self.scale=}") + print(f"{frames.shape=}") + frames = frames.to(self.device) + with torch.no_grad(): + scale_h = torch.ones(1).to(self.device) * (1 / float(self.scale[0])) + scale_w = torch.ones(1).to(self.device) * (1 / float(self.scale[1])) + + hs, hw = 1.0 / scale_h, 1.0 / scale_w + hs, hw = hs.unsqueeze(-1), hw.unsqueeze(-1) + kernel = None + + b, h, w, c = frames.shape + + hr_coord = make_coord((h, w)).unsqueeze(0).to(self.device) + + cell = torch.ones(2).unsqueeze(0).to(self.device) + cell[:, 0] *= 2.0 / h + cell[:, 1] *= 2.0 / w + # [B, C, H, W] + frames = frames.permute(0, 3, 1, 2) + + frames_scaled: list[torch.Tensor] = [] + for frame in frames: + frames_scaled.append( + F.interpolate( + frame.unsqueeze(0), + (int(scale_h[0].item() * h), int(scale_w[0].item() * w)), + mode="bicubic", + ) + ) + + if kernel is None: + res = torch.zeros( + ( + 1, + self.net.num_channels, + frames_scaled[0].shape[-2], + frames_scaled[0].shape[-1], + ), + device=frames_scaled[0].device, + ) + kernel = self.net.kernel_predict(res, hr_coord, cell) + + output = self.net.test_forward(frames_scaled, kernel, hr_coord).squeeze(0) #T,C,H,W + output = output.permute(0, 2, 3, 1) + return output + + +def make_coord(shape): + """Make coordinates at grid centers.""" + coord_seqs = [] + for i, n in enumerate(shape): + # v0, v1 = -1, 1 + + r = 1 / n + seq = -1 + r + (2 * r) * torch.arange(n).float() + coord_seqs.append(seq) + # ret = torch.stack(torch.meshgrid(coord_seqs, indexing='ij'), dim=-1) + ret = torch.stack(torch.meshgrid(coord_seqs, indexing='ij'), dim=-1) + return ret diff --git a/asvr_module/blocks.py b/asvr_module/blocks.py new file mode 100644 index 0000000..9a523b2 --- /dev/null +++ b/asvr_module/blocks.py @@ -0,0 +1,758 @@ +import torch.nn as nn +import torch +import torch.nn.functional as F +############################### +# common +############################### +class _NonLocalBlockND(nn.Module): + def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True): + super(_NonLocalBlockND, self).__init__() + + assert dimension in [1, 2, 3] + + self.dimension = dimension + self.sub_sample = sub_sample + + self.in_channels = in_channels + self.inter_channels = inter_channels + + if self.inter_channels is None: + self.inter_channels = in_channels // 2 + if self.inter_channels == 0: + self.inter_channels = 1 + + if dimension == 3: + conv_nd = nn.Conv3d + max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2)) + bn = nn.BatchNorm3d + elif dimension == 2: + conv_nd = nn.Conv2d + max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2)) + bn = nn.BatchNorm2d + else: + conv_nd = nn.Conv1d + max_pool_layer = nn.MaxPool1d(kernel_size=(2)) + bn = nn.BatchNorm1d + + self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels, + kernel_size=1, stride=1, padding=0) + + if bn_layer: + self.W = nn.Sequential( + conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels, + kernel_size=1, stride=1, padding=0), + bn(self.in_channels) + ) + nn.init.constant_(self.W[1].weight, 0) + nn.init.constant_(self.W[1].bias, 0) + else: + self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels, + kernel_size=1, stride=1, padding=0) + nn.init.constant_(self.W.weight, 0) + nn.init.constant_(self.W.bias, 0) + + self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels, + kernel_size=1, stride=1, padding=0) + + self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels, + kernel_size=1, stride=1, padding=0) + + # self.fusion = conv_nd(in_channels=self.inter_channels, out_channels=self.inter_channels, + # kernel_size=1, stride=1, padding=0) + + if sub_sample: + self.g = nn.Sequential(self.g, max_pool_layer) + self.phi = nn.Sequential(self.phi, max_pool_layer) + + self.softmax = nn.Softmax(dim=2) + + def forward(self, x, eve): + ''' + :param x: (b, c, t, h, w) + :return: + ''' + batch_size = x.size(0) + + g_x = self.g(x).view(batch_size, self.inter_channels, -1) #b, c, h*w/4 + g_x = g_x.permute(0, 2, 1) + + theta_eve = self.theta(eve).view(batch_size, self.inter_channels, -1) #b, c, h*w + theta_eve = theta_eve.permute(0, 2, 1) + phi_x = self.phi(x).view(batch_size, self.inter_channels, -1) #b, c, h*w/4 + f = torch.matmul(theta_eve, phi_x) # (b, h*w, c) X (b, c, h*w/4) = (b, h*w, h*w/4) + N = f.size(-1) + f_div_C = f / N + f_div_C = self.softmax(f_div_C) + + y = torch.matmul(f_div_C, g_x) # (b, h*w, h*w/4) X (b, h*w/4, c) = (b, h*w, c) + y = y.permute(0, 2, 1).contiguous() + y = y.view(batch_size, self.inter_channels, *x.size()[2:]) # b,c,h,w + W_y = self.W(y) + z = W_y + x + + return z + + +class NONLocalBlock2D(_NonLocalBlockND): + def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=False): + super(NONLocalBlock2D, self).__init__(in_channels, + inter_channels=inter_channels, + dimension=2, sub_sample=sub_sample, + bn_layer=bn_layer) + + +def get_same_padding(kernel_size, dilation): + kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1) + padding = (kernel_size - 1) // 2 + return padding + +class ConvLSTM(nn.Module): + def __init__(self, inp_dim, oup_dim, kernel, dilation): + super().__init__() + pad_x = int(dilation * (kernel - 1) / 2) + self.conv_xf = nn.Conv2d(inp_dim, oup_dim, kernel, padding=pad_x, dilation=dilation) + self.conv_xi = nn.Conv2d(inp_dim, oup_dim, kernel, padding=pad_x, dilation=dilation) + self.conv_xo = nn.Conv2d(inp_dim, oup_dim, kernel, padding=pad_x, dilation=dilation) + self.conv_xj = nn.Conv2d(inp_dim, oup_dim, kernel, padding=pad_x, dilation=dilation) + + pad_h = int((kernel - 1) / 2) + self.conv_hf = nn.Conv2d(oup_dim, oup_dim, kernel, padding=pad_h) + self.conv_hi = nn.Conv2d(oup_dim, oup_dim, kernel, padding=pad_h) + self.conv_ho = nn.Conv2d(oup_dim, oup_dim, kernel, padding=pad_h) + self.conv_hj = nn.Conv2d(oup_dim, oup_dim, kernel, padding=pad_h) + + self.se = SEBlock(oup_dim, 4) + self.relu = nn.LeakyReLU(0.2) + + def forward(self, x, pair=None): + if pair is None: + i = F.sigmoid(self.conv_xi(x)) + o = F.sigmoid(self.conv_xo(x)) + j = F.tanh(self.conv_xj(x)) + c = i * j + h = o * c + else: + h, c = pair + f = F.sigmoid(self.conv_xf(x) + self.conv_hf(h)) + i = F.sigmoid(self.conv_xi(x) + self.conv_hi(h)) + o = F.sigmoid(self.conv_xo(x) + self.conv_ho(h)) + j = F.tanh(self.conv_xj(x) + self.conv_hj(h)) + c = f * c + i * j + h = o * F.tanh(c) + + h = self.relu(self.se(h)) + return h, [h, c] + +class SEBlock(nn.Module): + def __init__(self, input_dim, reduction): + super().__init__() + mid = int(input_dim / reduction) + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(input_dim, mid), + nn.ReLU(inplace=True), + nn.Linear(mid, input_dim), + nn.Sigmoid() + ) + + def forward(self, x): + b, c, _, _ = x.size() + y = self.avg_pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + +############################### +# ResNet +############################### +# class ResBlock(nn.Module): +# def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bn=False, se=True): +# super(ResBlock, self).__init__() +# self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=kernel_size, stride=stride, +# padding=get_same_padding(kernel_size, dilation), dilation=dilation) +# self.conv2 = nn.Conv2d(planes, planes, kernel_size=kernel_size, stride=1, +# padding=get_same_padding(kernel_size, dilation), dilation=dilation) +# self.relu = nn.LeakyReLU(0.1, inplace=True) +# if se: +# self.se = SEBlock(planes, 4) +# if bn: +# self.bn1 = nn.InstanceNorm2d(planes) +# self.bn2 = nn.InstanceNorm2d(planes) +# self.se_ = se +# self.bn_ = bn +# self.res_translate = None +# if not inplanes == planes or not stride == 1: +# self.res_translate = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride) +# +# def forward(self, x): +# residual = x +# if self.bn_: +# out = self.relu(self.bn1(self.conv1(x))) +# else: +# out = self.relu(self.conv1(x)) +# if self.bn_: +# out = self.bn2(self.conv2(out)) +# else: +# out = self.conv2(out) +# if self.se_: +# out = self.se(out) +# if self.res_translate is not None: +# residual = self.res_translate(residual) +# out += residual +# +# return out + +class CoralLayer(torch.nn.Module): + """ Implements CORAL layer described in + + Cao, Mirjalili, and Raschka (2020) + *Rank Consistent Ordinal Regression for Neural Networks + with Application to Age Estimation* + Pattern Recognition Letters, https://doi.org/10.1016/j.patrec.2020.11.008 + + Parameters + ----------- + size_in : int + Number of input features for the inputs to the forward method, which + are expected to have shape=(num_examples, num_features). + + num_classes : int + Number of classes in the dataset. + + + """ + def __init__(self, size_in, num_classes): + super().__init__() + self.size_in, self.size_out = size_in, 1 + + self.coral_weights = torch.nn.Linear(self.size_in, 1, bias=False) + self.coral_bias = torch.nn.Parameter( + torch.zeros(num_classes-1).float()) + + def forward(self, x): + """ + Computes forward pass. + + Parameters + ----------- + x : torch.tensor, shape=(num_examples, num_features) + Input features. + + Returns + ----------- + logits : torch.tensor, shape=(num_examples, num_classes-1) + """ + return self.coral_weights(x) + self.coral_bias + +class ResBlock(nn.Module): + def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bn=False, se=True): + super(ResBlock, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=kernel_size, stride=stride, + padding=get_same_padding(kernel_size, dilation), dilation=dilation) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=kernel_size, stride=1, + padding=get_same_padding(kernel_size, dilation), dilation=dilation) + self.relu = nn.LeakyReLU(0.1, inplace=True) + if se: + self.se = SEBlock(planes, 4) + if bn: + self.bn1 = nn.BatchNorm2d(planes) + self.bn2 = nn.BatchNorm2d(planes) + self.se_ = se + self.bn_ = bn + self.res_translate = None + if not inplanes == planes or not stride == 1: + self.res_translate = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride) + + def forward(self, x): + residual = x + if self.bn_: + out = self.relu(self.bn1(self.conv1(x))) + else: + out = self.relu(self.conv1(x)) + if self.bn_: + out = self.bn2(self.conv2(out)) + else: + out = self.conv2(out) + if self.se_: + out = self.se(out) + if self.res_translate is not None: + residual = self.res_translate(residual) + out += residual + + return out + +class GResBlock(nn.Module): + def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1,group=1): + super(GResBlock, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=kernel_size, stride=stride, + padding=get_same_padding(kernel_size, dilation), dilation=dilation, groups=group) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=kernel_size, stride=1, + padding=get_same_padding(kernel_size, dilation), dilation=dilation, groups=group) + self.relu = nn.LeakyReLU(0.1, inplace=True) + self.conv3 = nn.Conv2d(planes, planes, kernel_size=kernel_size, stride=1, + padding=get_same_padding(kernel_size, dilation), dilation=dilation, groups=group) + self.res_translate = None + if not inplanes == planes or not stride == 1: + self.res_translate = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride) + + def forward(self, x): + residual = x + + out = self.relu(self.conv1(x)) + out = self.relu(self.conv2(out)) + out = self.conv3(out) + if self.res_translate is not None: + residual = self.res_translate(residual) + out += residual + + return out + +class SKConv(nn.Module): + def __init__(self, features, WH, M, G, r, stride=1, L=32): + """ Constructor + Args: + features: input channel dimensionality. + WH: input spatial dimensionality, used for GAP kernel size. + M: the number of branchs. + G: num of convolution groups. + r: the radio for compute d, the length of z. + stride: stride, default 1. + L: the minimum dim of the vector z in paper, default 32. + """ + super(SKConv, self).__init__() + d = max(int(features / r), L) + self.M = M + self.features = features + self.convs = nn.ModuleList([]) + for i in range(M): + self.convs.append(nn.Sequential( + nn.Conv2d(features, features, kernel_size=3 + i * 2, stride=stride, padding=1 + i, groups=G), + # nn.BatchNorm2d(features), + nn.LeakyReLU(0.1) + )) + # self.gap = nn.AvgPool2d(int(WH/stride)) + self.fc = nn.Linear(features, d) + self.fcs = nn.ModuleList([]) + for i in range(M): + self.fcs.append( + nn.Linear(d, features) + ) + self.softmax = nn.Softmax(dim=1) + + def forward(self, x): + for i, conv in enumerate(self.convs): + fea = conv(x).unsqueeze_(dim=1) + if i == 0: + feas = fea + else: + feas = torch.cat([feas, fea], dim=1) + fea_U = torch.sum(feas, dim=1) + # fea_s = self.gap(fea_U).squeeze_() + fea_s = fea_U.mean(-1).mean(-1) + fea_z = self.fc(fea_s) + for i, fc in enumerate(self.fcs): + vector = fc(fea_z).unsqueeze_(dim=1) + if i == 0: + attention_vectors = vector + else: + attention_vectors = torch.cat([attention_vectors, vector], dim=1) + attention_vectors = self.softmax(attention_vectors) + attention_vectors = attention_vectors.unsqueeze(-1).unsqueeze(-1) + fea_v = (feas * attention_vectors).sum(dim=1) + return fea_v + +class BasicConv(nn.Module): + def __init__(self, in_channel, out_channel, kernel_size, stride, bias=True, norm=False, relu=True, transpose=False): + super(BasicConv, self).__init__() + if bias and norm: + bias = False + + padding = kernel_size // 2 + layers = list() + if transpose: + padding = kernel_size // 2 -1 + layers.append(nn.ConvTranspose2d(in_channel, out_channel, kernel_size, padding=padding, stride=stride, bias=bias)) + else: + layers.append( + nn.Conv2d(in_channel, out_channel, kernel_size, padding=padding, stride=stride, bias=bias)) + if norm: + layers.append(nn.BatchNorm2d(out_channel)) + if relu: + layers.append(nn.ReLU(inplace=True)) + self.main = nn.Sequential(*layers) + + def forward(self, x): + return self.main(x) + +class SCM(nn.Module): + def __init__(self, out_plane): + super(SCM, self).__init__() + self.main = nn.Sequential( + BasicConv(3, out_plane//4, kernel_size=3, stride=1, relu=True), + BasicConv(out_plane // 4, out_plane // 2, kernel_size=1, stride=1, relu=True), + BasicConv(out_plane // 2, out_plane // 2, kernel_size=3, stride=1, relu=True), + BasicConv(out_plane // 2, out_plane-3, kernel_size=1, stride=1, relu=True) + ) + + self.conv = BasicConv(out_plane, out_plane, kernel_size=1, stride=1, relu=False) + + def forward(self, x): + x = torch.cat([x, self.main(x)], dim=1) + return self.conv(x) + +class SKUnit(nn.Module): + def __init__(self, in_features, out_features, WH, M, G, r, mid_features=None, stride=1, L=32): + """ Constructor + Args: + in_features: input channel dimensionality. + out_features: output channel dimensionality. + WH: input spatial dimensionality, used for GAP kernel size. + M: the number of branchs. + G: num of convolution groups. + r: the radio for compute d, the length of z. + mid_features: the channle dim of the middle conv with stride not 1, default out_features/2. + stride: stride. + L: the minimum dim of the vector z in paper. + """ + super(SKUnit, self).__init__() + if mid_features is None: + mid_features = int(out_features / 2) + self.feas = nn.Sequential( + nn.Conv2d(in_features, mid_features, 1, stride=1), + # nn.BatchNorm2d(mid_features), + SKConv(mid_features, WH, M, G, r, stride=stride, L=L), + # nn.BatchNorm2d(mid_features), + nn.Conv2d(mid_features, out_features, 1, stride=1), + # nn.BatchNorm2d(out_features) + ) + if in_features == out_features: # when dim not change, in could be added diectly to out + self.shortcut = nn.Sequential() + else: # when dim not change, in should also change dim to be added to out + self.shortcut = nn.Sequential( + nn.Conv2d(in_features, out_features, 1, stride=stride), + # nn.BatchNorm2d(out_features) + ) + + def forward(self, x): + fea = self.feas(x) + return fea + self.shortcut(x) + + +class ChannelAttention(nn.Module): + def __init__(self, in_planes, ratio=16): + super(ChannelAttention, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.max_pool = nn.AdaptiveMaxPool2d(1) + + self.fc = nn.Sequential(nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False), + nn.ReLU(), + nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + avg_out = self.fc(self.avg_pool(x)) + max_out = self.fc(self.max_pool(x)) + out = avg_out + max_out + return self.sigmoid(out) + + +class SpatialAttention(nn.Module): + def __init__(self, kernel_size=7): + super(SpatialAttention, self).__init__() + + self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + avg_out = torch.mean(x, dim=1, keepdim=True) + max_out, _ = torch.max(x, dim=1, keepdim=True) + x = torch.cat([avg_out, max_out], dim=1) + x = self.conv1(x) + return self.sigmoid(x) + +# class Unet(nn.Module): +# +# def __init__(self, in_dim=3*4, out_dim=3*3, block=2, feats=32, kernel_size=3): +# super(Unet, self).__init__() +# InBlock = [] +# +# InBlock.extend([nn.Sequential( +# nn.Conv2d(in_dim, feats, kernel_size=7, stride=1, +# padding=7 // 2), +# nn.LeakyReLU(0.1,inplace=True) +# )]) +# +# InBlock.extend([ResBlock(feats, feats, kernel_size=kernel_size, stride=1) +# for _ in range(block)]) +# +# # encoder1 +# Encoder_first = [nn.Sequential( +# nn.Conv2d(feats, feats*2, kernel_size=kernel_size, stride=2, padding=kernel_size // 2), +# nn.LeakyReLU(inplace=True) +# )] +# Encoder_first.extend([ResBlock(feats*2, feats*2, kernel_size=kernel_size, stride=1) +# for _ in range(block)]) +# # encoder2 +# Encoder_second = [nn.Sequential( +# nn.Conv2d(feats*2, feats*4, kernel_size=kernel_size, stride=2, padding=kernel_size // 2), +# nn.LeakyReLU(inplace=True) +# )] +# Encoder_second.extend([ResBlock(feats * 4, feats * 4, kernel_size=kernel_size, stride=1) +# for _ in range(block)]) +# +# # decoder2 +# Decoder_second = [ResBlock(feats * 4, feats * 4, kernel_size=kernel_size, stride=1) +# for _ in range(block)] +# Decoder_second.append(nn.Sequential( +# nn.ConvTranspose2d(feats * 4, feats * 2, kernel_size=3, stride=2, padding=1, output_padding=1), +# nn.LeakyReLU(inplace=True) +# )) +# # decoder1 +# Decoder_first = [ResBlock(feats * 2, feats * 2, kernel_size=kernel_size, stride=1) +# for _ in range(block)] +# Decoder_first.append(nn.Sequential( +# nn.ConvTranspose2d(feats * 2, feats, kernel_size=3, stride=2, padding=1, output_padding=1), +# nn.LeakyReLU(inplace=True) +# )) +# +# OutBlock = [ResBlock(feats, feats, kernel_size=kernel_size, stride=1) +# for _ in range(block)] +# OutBlock.append(nn.Conv2d(feats, out_dim, kernel_size=kernel_size, stride=1, padding=kernel_size // 2)) +# +# +# self.inBlock = nn.Sequential(*InBlock) +# self.encoder_first = nn.Sequential(*Encoder_first) +# self.encoder_second = nn.Sequential(*Encoder_second) +# self.decoder_second = nn.Sequential(*Decoder_second) +# self.decoder_first = nn.Sequential(*Decoder_first) +# self.outBlock = nn.Sequential(*OutBlock) +# +# def forward(self, x): +# +# x = torch.cat(x, dim=1) +# +# first_scale_inblock = self.inBlock(x) +# first_scale_encoder_first = self.encoder_first(first_scale_inblock) +# first_scale_encoder_second = self.encoder_second(first_scale_encoder_first) +# +# return first_scale_inblock, first_scale_encoder_first, first_scale_encoder_second + + +class DA_conv(nn.Module): + def __init__(self, channels_in, channels_out, kernel_size, reduction, kernel_dim): + super(DA_conv, self).__init__() + self.channels_out = channels_out + self.channels_in = channels_in + self.kernel_size = kernel_size + + # self.kernel = nn.Sequential( + # nn.Linear(kernel_dim, kernel_dim//2, bias=False), + # nn.LeakyReLU(0.1, True), + # nn.Linear(kernel_dim//2, channels_in * self.kernel_size * self.kernel_size, bias=False) #3 + # ) + # self.conv = nn.Conv2d(channels_in, channels_out, 1, padding=0) + self.ca = CA_layer(kernel_dim, channels_out, reduction) + + # self.relu = nn.LeakyReLU(0.1, True) + + def forward(self, x, vec): + ''' + :param x[0]: feature map: B * C * H * W + :param x[1]: degradation representation: B * C #4,256 + ''' + # print(x.size(), vec.size()) + # b, c, h, w = x.size() + # # print(x[0].size()) + # # branch 1 + # kernel = self.kernel(vec).view(-1, 1, self.kernel_size, self.kernel_size) #4,128*3*3 512,1,3,3 + # # print(kernel.size()) + # out = self.relu(F.conv2d(x.view(1, -1, h, w), kernel, groups=b*c, padding=(self.kernel_size-1)//2)) + # out = self.conv(out.view(b, -1, h, w)) + # print(out.size()) + # branch 2 + out = x + out = out + self.ca(x, vec) + + return out + + +class CA_layer(nn.Module): + def __init__(self, channels_in, channels_out, reduction): + super(CA_layer, self).__init__() + self.conv_du = nn.Sequential( + nn.Conv2d(channels_in, channels_in//reduction, 1, 1, 0, bias=False), + nn.LeakyReLU(0.1, True), + nn.Conv2d(channels_in // reduction, channels_out, 1, 1, 0, bias=False), + nn.Sigmoid() + ) + + def forward(self, x,vec): + ''' + :param x[0]: feature map: B * C * H * W + :param x[1]: degradation representation: B * C + ''' + att = self.conv_du(vec[:, :, None, None]) + + return x * att + +class ChanNorm(nn.Module): + def __init__(self, dim, eps = 1e-5): + super().__init__() + self.eps = eps + self.g = nn.Parameter(torch.ones(1, dim, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, dim, 1, 1)) + + def forward(self, x): + var = torch.var(x, dim = 1, unbiased = False, keepdim = True) + mean = torch.mean(x, dim = 1, keepdim = True) + return (x - mean) / (var + self.eps).sqrt() * self.g + self.b +# nn.InstanceNorm2d +# b=torch.tensor([ +# [ +# [[1,0,1],[0,1,0],[1,0,1]],[[1,0,1],[0,1,0],[1,0,1]] +# ], +# [ +# [[1, 0, 1], [0, 1, 0], [1, 0, 1]], [[1, 0, 1], [0, 1, 0], [1, 0, 1]] +# ] +# ]) + +class Conv2DMod(nn.Module): + def __init__(self, in_chan, out_chan, kernel, demod=True, stride=1, dilation=1, eps = 1e-8, **kwargs): + super().__init__() + self.filters = out_chan + self.demod = demod + self.kernel = kernel + self.stride = stride + self.dilation = dilation + self.weight = nn.Parameter(torch.randn((out_chan, in_chan, kernel, kernel))) + self.eps = eps + nn.init.kaiming_normal_(self.weight, a=0, mode='fan_in', nonlinearity='leaky_relu') + + def _get_same_padding(self, size, kernel, dilation, stride): + return ((size - 1) * (stride - 1) + dilation * (kernel - 1)) // 2 + + def forward(self, x, y): + b, c, h, w = x.shape + + w1 = y[:, None, :, None, None] + w2 = self.weight[None, :, :, :, :] + weights = w2 * (w1 + 1) + + if self.demod: + d = torch.rsqrt((weights ** 2).sum(dim=(2, 3, 4), keepdim=True) + self.eps) + weights = weights * d + + x = x.reshape(1, -1, h, w) + + _, _, *ws = weights.shape + weights = weights.reshape(b * self.filters, *ws) + + padding = self._get_same_padding(h, self.kernel, self.dilation, self.stride) + x = F.conv2d(x, weights, padding=padding, groups=b) + + x = x.reshape(-1, self.filters, h, w) + return x + +class Sty_layer(nn.Module): + def __init__(self, channels_in, channels_out, reduction): + super(Sty_layer, self).__init__() + self.norm = ChanNorm(channels_out) + self.conv_du = nn.Sequential( + nn.Linear(channels_in, channels_in//reduction, bias=False), + nn.LeakyReLU(0.1, True), + nn.Linear(channels_in // reduction, channels_out, bias=False), + # nn.Sigmoid() + ) + self.conv = Conv2DMod(channels_out, channels_out, 3) + + def forward(self, x, vec): + ''' + :param x[0]: feature map: B * C * H * W + :param x[1]: degradation representation: B * C + ''' + norm_x = self.norm(x) + sty = self.conv_du(vec) + out = self.conv(norm_x, sty) + + return out + x + + +class Ada_layer(nn.Module): + def __init__(self, channels_in, channels_out, reduction): + super(Ada_layer, self).__init__() + self.conv_alpha = nn.Sequential( + nn.Conv2d(channels_in, channels_in//reduction, 1, 1, 0, bias=False), + nn.LeakyReLU(0.1, True), + nn.Conv2d(channels_in // reduction, channels_out, 1, 1, 0, bias=False), + nn.Sigmoid() + ) + self.conv_belta = nn.Sequential( + nn.Conv2d(channels_in, channels_in // reduction, 1, 1, 0, bias=False), + nn.LeakyReLU(0.1, True), + nn.Conv2d(channels_in // reduction, channels_out, 1, 1, 0, bias=False) + ) + + def forward(self, x,vec): + ''' + :param x[0]: feature map: B * C * H * W + :param x[1]: degradation representation: B * C + ''' + alpha = self.conv_alpha(vec[:, :, None, None]) + belta = self.conv_belta(vec[:, :, None, None]) + + return x * alpha + belta + +class DAB(nn.Module): + def __init__(self, n_feat, kernel_size, kernel_dim, reduction=8): + super(DAB, self).__init__() + + self.da_conv1 = DA_conv(n_feat, n_feat, kernel_size, reduction, kernel_dim) + self.da_conv2 = DA_conv(n_feat, n_feat, kernel_size, reduction, kernel_dim) + self.conv1 = nn.Conv2d(n_feat, n_feat, kernel_size, padding=(kernel_size//2)) + self.conv2 = nn.Conv2d(n_feat, n_feat, kernel_size, padding=(kernel_size//2)) + + self.relu = nn.LeakyReLU(0.1, True) + + def forward(self, inp): + ''' + :param x[0]: feature map: B * C * H * W + :param x[1]: degradation representation: B * C + ''' + x, vec = inp + out = self.relu(self.da_conv1(x,vec)) + out = self.relu(self.conv1(out)) + out = self.relu(self.da_conv2(out, vec)) + out = self.conv2(out) + x + + return out,vec + +class SFT_Layer(nn.Module): + def __init__(self, nf=64, para=10): + super(SFT_Layer, self).__init__() + self.mul_conv1 = nn.Conv2d(para + nf, nf//2, kernel_size=3, stride=1, padding=1) + self.mul_leaky = nn.LeakyReLU(0.2) + self.mul_conv2 = nn.Conv2d(nf//2, nf, kernel_size=3, stride=1, padding=1) + + self.add_conv1 = nn.Conv2d(para + nf, nf//2, kernel_size=3, stride=1, padding=1) + self.add_leaky = nn.LeakyReLU(0.2) + self.add_conv2 = nn.Conv2d(nf//2, nf, kernel_size=3, stride=1, padding=1) + + def forward(self, feature_maps, para_maps): + cat_input = torch.cat((feature_maps, para_maps), dim=1) + mul = torch.sigmoid(self.mul_conv2(self.mul_leaky(self.mul_conv1(cat_input)))) + add = self.add_conv2(self.add_leaky(self.add_conv1(cat_input))) + return feature_maps * mul + add + + +class SFT_Residual_Block(nn.Module): + def __init__(self, nf=64, para=10): + super(SFT_Residual_Block, self).__init__() + self.sft1 = SFT_Layer(nf=nf, para=para) + self.sft2 = SFT_Layer(nf=nf, para=para) + self.conv1 = nn.Conv2d(in_channels=nf, out_channels=nf, kernel_size=3, stride=1, padding=1, bias=True) + self.conv2 = nn.Conv2d(in_channels=nf, out_channels=nf, kernel_size=3, stride=1, padding=1, bias=True) + + def forward(self, feature_maps, para_maps): + fea1 = F.relu(self.sft1(feature_maps, para_maps)) + fea2 = F.relu(self.sft2(self.conv1(fea1), para_maps)) + fea3 = self.conv2(fea2) + return torch.add(feature_maps, fea3) diff --git a/asvr_module/components.py b/asvr_module/components.py new file mode 100644 index 0000000..d1981da --- /dev/null +++ b/asvr_module/components.py @@ -0,0 +1,120 @@ +import functools +from .init import init_fn +import torch +import torch.nn as nn +import torch.nn.functional as F + +class ResBlock(nn.Module): + def __init__(self, in_channels, mid_channels, out_channels): + super(ResBlock, self).__init__() + self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=mid_channels, kernel_size=3, stride=1, padding=1) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x): + output = self.conv2(self.relu(self.conv1(x))) + output = torch.add(output, x) + return output + +class ResBlocks(nn.Module): + def __init__(self, input_channels, num_resblocks, num_channels): + super(ResBlocks, self).__init__() + self.input_channels = input_channels + self.first_conv = nn.Conv2d(in_channels=self.input_channels, out_channels=num_channels, kernel_size=3, stride=1, padding=1) + + modules = [] + for _ in range(num_resblocks): + modules.append(ResBlock(in_channels=num_channels, mid_channels=num_channels, out_channels=num_channels)) + self.resblocks = nn.Sequential(*modules) + + fn = functools.partial(init_fn, init_type='kaiming_normal', init_bn_type='uniform', gain=0.2) + self.apply(fn) + + def forward(self, h): + shallow_feature = self.first_conv(h) + new_h = self.resblocks(shallow_feature) + return new_h + +class Conv2DMod(nn.Module): + def __init__(self, in_chan, out_chan, kernel, style_ch=82, demod=True, stride=1, padding=1, dilation=1, eps = 1e-8, **kwargs): + super().__init__() + self.filters = out_chan + self.demod = demod + self.kernel = kernel + self.stride = stride + self.padding = padding + self.dilation = dilation + self.weight = nn.Parameter(torch.randn((out_chan, in_chan, kernel, kernel))) + self.eps = eps + self.to_style = nn.Linear(style_ch, in_chan) + nn.init.kaiming_normal_(self.weight, a=0, mode='fan_in', nonlinearity='leaky_relu') + + # def _get_same_padding(self, size, kernel, dilation, stride): + # return ((size - 1) * (stride - 1) + dilation * (kernel - 1)) // 2 + + def forward(self, x, scale): + b, c, h, w = x.shape + + y = self.to_style(scale) + w1 = y[:, None, :, None, None] + # print('w1--', w1.shape) + w2 = self.weight[None, :, :, :, :] + # print('w2--',w2.shape) + weights = w2 * (w1 + 1) + + if self.demod: + d = torch.rsqrt((weights ** 2).sum(dim=(2, 3, 4), keepdim=True) + self.eps) + weights = weights * d + + x = x.reshape(1, -1, h, w) + + _, _, *ws = weights.shape + weights = weights.reshape(b * self.filters, *ws) + + # padding = self._get_same_padding(h, self.kernel, self.dilation, self.stride) + x = F.conv2d(x, weights, padding=self.padding, groups=b) + # print(x.shape) + + x = x.reshape(-1, self.filters, h, w) + return x + +class ResBlocks_mod(nn.Module): + def __init__(self, input_channels, num_resblocks, num_channels): + super(ResBlocks_mod, self).__init__() + self.input_channels = input_channels + self.first_conv = nn.Conv2d(in_channels=self.input_channels, out_channels=num_channels, kernel_size=3, stride=1, padding=1) + self.num_resblocks = num_resblocks + modules = [] + for _ in range(num_resblocks): + modules.append(ResBlock(in_channels=num_channels, mid_channels=num_channels, out_channels=num_channels)) + self.resblocks = nn.Sequential(*modules) + fn = functools.partial(init_fn, init_type='kaiming_normal', init_bn_type='uniform', gain=0.2) + self.apply(fn) + + modc = [] + for _ in range(num_resblocks // 5): + modc.append(Conv2DMod(num_channels, num_channels, 3)) + self.mbconv = nn.Sequential(*modc) + + def forward(self, h, scale): + new_h = self.first_conv(h) + for i in range(self.num_resblocks//5): + new_h = self.mbconv[i](new_h, scale) + new_h = self.resblocks[i*5:(i+1)*5](new_h) + return new_h + +class D(nn.Module): + def __init__(self, in_channels, mid_channels, out_channels): + super(D, self).__init__() + layers = [] + layers.append(nn.Conv2d(in_channels=in_channels, out_channels=mid_channels, kernel_size=3, stride=1, padding=1)) + layers.append(nn.ReLU()) + layers.append(nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)) + self.convs = nn.Sequential(*layers) + + fn = functools.partial(init_fn, init_type='kaiming_normal', init_bn_type='uniform', gain=0.2) + self.apply(fn) + + def forward(self, x): + x = self.convs(x) + return x diff --git a/asvr_module/cuf_siren.py b/asvr_module/cuf_siren.py new file mode 100644 index 0000000..1c82e97 --- /dev/null +++ b/asvr_module/cuf_siren.py @@ -0,0 +1,253 @@ +import numpy as np + +import torch +import torch.fft +import torch.nn as nn +import torch.nn.functional as F +from .SIREN import Siren + +def make_coord(shape, ranges=None, flatten=True): + """ Make coordinates at grid centers. + coord_x = -1+(2*i+1)/W + coord_y = -1+(2*i+1)/H + normalize to (-1, 1) + """ + coord_seqs = [] + for i, n in enumerate(shape): + if ranges is None: + v0, v1 = -1, 1 + else: + v0, v1 = ranges[i] + r = (v1 - v0) / (2 * n) + seq = v0 + r + (2 * r) * torch.arange(n).float() + coord_seqs.append(seq) + ret = torch.stack(torch.meshgrid(*coord_seqs, indexing='ij'), dim=-1) + if flatten: + ret = ret.view(-1, ret.shape[-1]) + return ret + +class PositionEncoder(nn.Module): + def __init__( + self, + posenc_type='sinusoid', + complex_transform=False, + posenc_scale=10, + gauss_scale=1, + in_dims=2, + enc_dims=64, + hidden_dims=64, + out_dims=32, + head=8, + gamma=1 + ): + super().__init__() + + self.posenc_type = posenc_type + self.complex_transform = complex_transform + self.posenc_scale = posenc_scale + self.gauss_scale = gauss_scale + self.out_dims = out_dims + self.in_dims = in_dims + self.enc_dims = enc_dims + self.hidden_dims = hidden_dims + self.head = head + self.gamma = gamma + + self.define_parameter() + + def define_parameter(self): + if self.posenc_type == 'sinusoid' or self.posenc_type == 'ipe': + self.b_vals = 2.**torch.linspace( + 0, self.posenc_scale, self.enc_dims // 4 + ) - 1 # -1 -> (2 * pi) + self.b_vals = torch.stack([self.b_vals, torch.zeros_like(self.b_vals)], dim=-1) + self.b_vals = torch.cat([self.b_vals, torch.roll(self.b_vals, 1, -1)], dim=0) + self.a_vals = torch.ones(self.b_vals.shape[0]) + self.proj = nn.Linear(self.enc_dims, self.out_dims) #self.head) + + elif self.posenc_type == 'learn': + self.Wr = nn.Linear(self.in_dims, self.hidden_dims // 2, bias=False) + self.mlp = nn.Sequential( + nn.LayerNorm(self.hidden_dims, eps=1e-6), + nn.Linear(self.hidden_dims, self.hidden_dims), + nn.GELU(), + nn.LayerNorm(self.hidden_dims, eps=1e-6), + nn.Linear(self.hidden_dims, self.enc_dims) + ) + self.proj = nn.Sequential(nn.GELU(), nn.Linear(self.enc_dims, self.head)) + self.init_weight() + + elif self.posenc_type == 'dpb': + self.mlp = nn.Sequential( + nn.Linear(2, self.hidden_dims), + nn.LayerNorm(self.hidden_dims, eps=1e-6), + nn.ReLU(), + nn.Linear(self.hidden_dims, self.hidden_dims), + nn.LayerNorm(self.hidden_dims, eps=1e-6), + nn.ReLU(), + nn.Linear(self.hidden_dims, self.enc_dims) + ) + self.proj = nn.Sequential( + nn.LayerNorm(self.hidden_dims, eps=1e-6), + nn.ReLU(), + nn.Linear(self.enc_dims, self.head) + ) + + def init_weight(self): + nn.init.normal_(self.Wr.weight.data, mean=0, std=self.gamma**-2) + + def forward(self, positions, cells=None): + + if self.posenc_type is None: + return positions + + if self.posenc_type == 'sinusoid' or self.posenc_type == 'ipe': + self.b_vals = self.b_vals.to(positions.device) #.cuda() + self.a_vals = self.a_vals.to(positions.device) #.cuda() + + # b, q, 1, c (x -> c/2, y -> c/2) + sin_part = self.a_vals * torch.sin( + torch.matmul(positions, self.b_vals.transpose(-2, -1)) + ) + cos_part = self.a_vals * torch.cos( + torch.matmul(positions, self.b_vals.transpose(-2, -1)) + ) + + if self.posenc_type == 'ipe': + # b, q, 2 + cell = cells.clone() + cell_part = torch.sinc( + torch.matmul((1 / np.pi * cell), self.b_vals.transpose(-2, -1)) + ) + + sin_part = sin_part * cell_part + cos_part = cos_part * cell_part + + if self.complex_transform: + pos_enocoding = torch.view_as_complex(torch.stack([cos_part, sin_part], dim=-1)) + else: + pos_enocoding = torch.cat([sin_part, cos_part], dim=-1) + # pos_bias = self.proj(pos_enocoding) + + elif self.posenc_type == 'learn': + projected_pos = self.Wr(positions) + + sin_part = torch.sin(projected_pos) + cos_part = torch.cos(projected_pos) + + if self.complex_transform: + pos_enocoding = 1 / np.sqrt(self.hidden_dims) * torch.view_as_complex( + torch.stack([cos_part, sin_part], dim=-1) + ) + else: + pos_enocoding = 1 / np.sqrt(self.hidden_dims + ) * torch.cat([sin_part, cos_part], dim=-1) + pos_enocoding = self.mlp(pos_enocoding) + + elif self.posenc_type == 'dpb': + pos_enocoding = self.mlp(positions) + + pos_bias = None if self.complex_transform else self.proj(pos_enocoding) + + return pos_enocoding, pos_bias + +class KerP(nn.Module): + def __init__( + self, + feat_dim, + sq_factor=4, + r=1, + ): + super().__init__() + + # self.dim = base_dim + self.r = r + self.pb_encoder1 = PositionEncoder() + self.pb_encoder2 = PositionEncoder() + self.pb_encoder3 = PositionEncoder() + self.r_area = (2 * self.r + 1)**2 + imnet_in_dim = 32 * 3 + self.imnets = Siren(in_features=imnet_in_dim, out_features=feat_dim, hidden_features=[feat_dim//sq_factor, feat_dim//sq_factor, feat_dim//sq_factor], + hidden_layers=2, outermost_linear=True) #MLP(imnet_in_dim, feat_dim) + + def forward(self, feat, sample_coord, cell): + bs = feat.shape[0] + r = self.r + coord_lr = make_coord(feat.shape[-2:], flatten=False).to(sample_coord.device).permute(2, 0, 1).contiguous(). \ + unsqueeze(0).expand(bs, 2, *feat.shape[-2:]) + + # b, 2, h_in, w_in -> b, 2, h_t, w_t -> b, h_t, w_t, 2 -> b, q, 2 + sample_coord_k = F.grid_sample( + coord_lr, sample_coord.flip(-1), mode='nearest', align_corners=False + ).permute(0, 2, 3, 1).contiguous().view(bs, -1, 2).contiguous() + + rel_coord = sample_coord.view(bs, -1, 2).contiguous() - sample_coord_k + rel_coord[..., 0] *= feat.shape[-2] + rel_coord[..., 1] *= feat.shape[-1] + + dh = torch.linspace(-r, r, 2 * r + 1).to(sample_coord.device) # * rh + dw = torch.linspace(-r, r, 2 * r + 1).to(sample_coord.device) # * rw + # 1, r_area, 2 + delta = torch.stack(torch.meshgrid(dh, dw, indexing='ij'), axis=-1).view(1, -1, 2).repeat(bs, 1, 1).contiguous() + + # b, 2 -> b, q, 2 + rel_cell = cell.unsqueeze(1).repeat(1, sample_coord_k.shape[1], 1) + rel_cell[..., 0] *= feat.shape[-2] + rel_cell[..., 1] *= feat.shape[-1] + + _, pb1 = self.pb_encoder1(rel_cell) # b, q, 32 + _, pb2 = self.pb_encoder2(delta) # b, r_area, 32 + _, pb3 = self.pb_encoder3(rel_coord) # b, q, 32 + # print(pb1.shape, pb2.shape,pb3.shape) + pb1 = pb1.unsqueeze(1).repeat(1, self.r_area, 1, 1).view(-1, sample_coord_k.shape[1], 32) + pb2 = pb2.unsqueeze(2).repeat(1, 1, sample_coord_k.shape[1], 1).view(-1, sample_coord_k.shape[1], + 32) # b*r_area, q, 32 + pb3 = pb3.unsqueeze(1).repeat(1, self.r_area, 1, 1).view(-1, sample_coord_k.shape[1], 32) + # print(pb1.shape, pb2.shape, pb3.shape) + feat_in = torch.cat([pb1, pb2, pb3], dim=-1) + + pred = self.imnets(feat_in) # b*r_area, q, feat_dim + pred = pred.view(bs, self.r_area, sample_coord_k.shape[1], feat.shape[1]).contiguous().permute(0, 3, 1, + 2).contiguous() + kernel = pred.view(bs, feat.shape[1] * self.r_area, -1).contiguous() + return kernel + +class CUF(nn.Module): + def __init__( + self, + feat_dim, + r=1, + ): + super().__init__() + + self.r = r + self.r_area = (2 * self.r + 1)**2 + self.unfold = torch.nn.Unfold((2 * r + 1, 2 * r + 1), 1, 1) + self.imnets = nn.Sequential( + nn.Conv2d(feat_dim, feat_dim, 1), + nn.LeakyReLU(negative_slope=0.1, inplace=True), + nn.Conv2d(feat_dim, 3, 3, 1, 1) + ) + + def forward(self, feat, kernel, sample_coord, inp): + ''' + feat: B, C, H_in, W_in + smaple_coord: B, H_target, W_target, 2 + cell: B, 2 + inp: B, C, H_in, W_in + ''' + + bs, H_t, W_t, _ = sample_coord.shape + + feat_unf = self.unfold(feat) # B, C*r_area, H_in*W_in + feat_unf = F.grid_sample(feat_unf.view(bs, -1, feat.shape[-2], feat.shape[-1]).contiguous(), sample_coord.flip(-1), mode='bilinear', align_corners=False).view(bs,feat.shape[1]*self.r_area,-1).contiguous() + + pred = feat_unf * kernel # B, C*r_area, H_target * W_target + + Fold = torch.nn.Fold(sample_coord.shape[1:3], (2 * self.r + 1, 2 * self.r + 1), 1, 1) + pred_ = Fold(pred) + pred = self.imnets(pred_) + + pred_f = pred + F.grid_sample(inp, sample_coord.flip(-1), mode='bilinear', padding_mode='border', align_corners=False) + return pred_f + diff --git a/asvr_module/init.py b/asvr_module/init.py new file mode 100644 index 0000000..2b3cc84 --- /dev/null +++ b/asvr_module/init.py @@ -0,0 +1,69 @@ +""" +# -------------------------------------------- +# weights initialization +# -------------------------------------------- +""" +from torch.nn import init + +""" +# Kai Zhang, https://github.com/cszn/KAIR +# +# Args: +# init_type: +# normal; normal; xavier_normal; xavier_uniform; +# kaiming_normal; kaiming_uniform; orthogonal +# init_bn_type: +# uniform; constant +# gain: +# 0.2 +""" + +def init_fn(m, init_type='kaiming_normal', init_bn_type='uniform', gain=0.2): + classname = m.__class__.__name__ + + if classname.find('Conv') != -1 or classname.find('Linear') != -1: + + if init_type == 'normal': + init.normal_(m.weight.data, 0, 0.1) + m.weight.data.clamp_(-1, 1).mul_(gain) + + elif init_type == 'uniform': + init.uniform_(m.weight.data, -0.2, 0.2) + m.weight.data.mul_(gain) + + elif init_type == 'xavier_normal': + init.xavier_normal_(m.weight.data, gain=gain) + m.weight.data.clamp_(-1, 1) + + elif init_type == 'xavier_uniform': + init.xavier_uniform_(m.weight.data, gain=gain) + + elif init_type == 'kaiming_normal': + init.kaiming_normal_(m.weight.data, a=0, mode='fan_in', nonlinearity='relu') + m.weight.data.clamp_(-1, 1).mul_(gain) + + elif init_type == 'kaiming_uniform': + init.kaiming_uniform_(m.weight.data, a=0, mode='fan_in', nonlinearity='relu') + m.weight.data.mul_(gain) + + elif init_type == 'orthogonal': + init.orthogonal_(m.weight.data, gain=gain) + + else: + raise NotImplementedError('Initialization method [{:s}] is not implemented'.format(init_type)) + + if m.bias is not None: + m.bias.data.zero_() + + elif classname.find('BatchNorm2d') != -1: + + if init_bn_type == 'uniform': # preferred + if m.affine: + init.uniform_(m.weight.data, 0.1, 1.0) + init.constant_(m.bias.data, 0.0) + elif init_bn_type == 'constant': + if m.affine: + init.constant_(m.weight.data, 1.0) + init.constant_(m.bias.data, 0.0) + else: + raise NotImplementedError('Initialization method [{:s}] is not implemented'.format(init_bn_type)) \ No newline at end of file diff --git a/asvr_module/network_vidue_refsr_cuf3_siren.py b/asvr_module/network_vidue_refsr_cuf3_siren.py new file mode 100644 index 0000000..e33cabe --- /dev/null +++ b/asvr_module/network_vidue_refsr_cuf3_siren.py @@ -0,0 +1,224 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from . import blocks +from mmcv.ops import ModulatedDeformConv2d, modulated_deform_conv2d + +def make_coord(shape): + """ Make coordinates at grid centers. + """ + coord_seqs = [] + for i, n in enumerate(shape): + # v0, v1 = -1, 1 + + r = 1 / n + seq = -1 + r + (2 * r) * torch.arange(n).float() + coord_seqs.append(seq) + # ret = torch.stack(torch.meshgrid(coord_seqs, indexing='ij'), dim=-1) + ret = torch.stack(torch.meshgrid(coord_seqs), dim=-1) + return ret + +class DynAgg(ModulatedDeformConv2d): + ''' + Use other features to generate offsets and masks. + Intialized the offset with precomputed non-local offset. + ''' + + def __init__(self, + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation=1, + groups=1, + deform_groups=1, + extra_offset_mask=True): + super().__init__(in_channels, out_channels, kernel_size, stride, + padding, dilation, groups, deform_groups) + self.extra_offset_mask = extra_offset_mask + channels_ = self.deform_groups * 3 * self.kernel_size[ + 0] * self.kernel_size[1] + self.conv_offset_mask = nn.Conv2d( + self.in_channels, + channels_, + kernel_size=self.kernel_size, + stride=self.stride, + padding=self.padding, + bias=True) + self.init_offset() + + def init_offset(self): + self.conv_offset_mask.weight.data.zero_() + self.conv_offset_mask.bias.data.zero_() + + def forward(self, x, pre_offset): + ''' + Args: + pre_offset: precomputed_offset. Size: [b, 2, h, w] + ''' + # pre_offset + + if self.extra_offset_mask: + # x = [input, features] + out = self.conv_offset_mask(x[1]) + x = x[0] + else: + out = self.conv_offset_mask(x) + o1, o2, mask = torch.chunk(out, 3, dim=1) + offset_ = torch.cat((o1, o2), dim=1) + + offset = offset_ + pre_offset.flip(1).repeat(1, offset_.size(1) // 2, 1, 1) + + # # repeat pre_offset along dim1, shape: [b, 9*groups, h, w, 2] + # pre_offset = pre_offset.repeat([1, self.deform_groups, 1, 1, 1]) + # # the order of offset is [y, x, y, x, ..., y, x] + # pre_offset_reorder = torch.zeros_like(offset) + # # add pre_offset on y-axis + # pre_offset_reorder[:, 0::2, :, :] = pre_offset[:, :, :, :, 1] + # # add pre_offset on x-axis + # pre_offset_reorder[:, 1::2, :, :] = pre_offset[:, :, :, :, 0] + # offset = offset + pre_offset_reorder + # print(offset.size()) + mask = torch.sigmoid(mask) + + offset_mean = torch.mean(torch.abs(offset_)) + if offset_mean > 100: + print('Offset mean is {}, larger than 100.'.format(offset_mean)) + return modulated_deform_conv2d(x, offset, mask, self.weight, self.bias, + self.stride, self.padding, self.dilation, + self.groups, self.deform_groups) + + +class MRAPAFusion(nn.Module): + """Temporal Spatial Attention (TSA) fusion module. It is used in EDVRNet. + Args: + nf (int): Number of the channels of middle features. + Default: 64. + ref_nf (int): Number of the channels of middle features. + Default: 256. + """ + + def __init__(self, + nf=64, + ref_nf=256): + super().__init__() + + # multi-ref attention (before fusion conv) + self.patch_size = 3 + channels = ref_nf + self.conv_emb1 = nn.Sequential( + nn.Conv2d(nf, channels, 1), + nn.PReLU()) + self.conv_emb2 = nn.Sequential( + nn.Conv2d(ref_nf, channels, + self.patch_size, 1, self.patch_size // 2), + nn.PReLU()) + self.conv_ass = nn.Conv2d(ref_nf, channels * 2, + self.patch_size, 1, self.patch_size // 2) + self.scale = channels ** -0.5 + self.feat_fusion = nn.Conv2d( + nf + channels * 2, nf, 1) + + # spatial attention (after fusion conv) + self.spatial_attn = nn.Conv2d( + nf + channels * 2, channels * 2, 1) + self.spatial_attn_mul1 = nn.Conv2d( + channels * 2, channels * 2, 3, padding=1) + self.spatial_attn_mul2 = nn.Conv2d( + channels * 2, channels * 2, 3, padding=1) + self.spatial_attn_add1 = nn.Conv2d( + channels * 2, channels * 2, 3, padding=1) + self.spatial_attn_add2 = nn.Conv2d( + channels * 2, channels * 2, 3, padding=1) + + self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True) + + def spatial_padding(self, feats): + _, _, h, w = feats.size() + pad_h = (4 - h % 4) % 4 + pad_w = (4 - w % 4) % 4 + feats = F.pad(feats, [0, pad_w, 0, pad_h], mode='reflect') + return feats + + def forward(self, target, refs): + n, _, h_input, w_input = target.size() + t = len(refs) + + target = self.spatial_padding(target) + refs = torch.stack(refs, dim=1).flatten(0, 1) + refs = self.spatial_padding(refs) + # multi-ref attention + embedding_target = self.conv_emb1(target) * self.scale # (n, c, h, w) + embedding_target = embedding_target.permute(0, 2, 3, 1).contiguous().unsqueeze(3) # (n, h, w, 1, c) + embedding_target = embedding_target.contiguous().flatten(0, 2) # (n*h*w, 1, c) + emb = self.conv_emb2(refs).unflatten(0, (n, t)) # (n, t, c, h, w) + emb = emb.permute(0, 3, 4, 2, 1).contiguous() # (n, h, w, c, t) + emb = emb.flatten(0, 2) # (n*h*w, c, t) + ass = self.conv_ass(refs).unflatten(0, (n, t)) # (n, t, c*2, h, w) + ass = ass.permute(0, 3, 4, 1, 2).contiguous() # (n, h, w, t, c*2) + ass = ass.flatten(0, 2) # (n*h*w, t, c*2) + + corr_prob = torch.matmul(embedding_target, emb) # (n*h*w, 1, t) + corr_prob = F.softmax(corr_prob, dim=2) + refs = torch.matmul(corr_prob, ass).squeeze(1) # (n*h*w, c*2) + refs = refs.unflatten(0, (n, *target.shape[-2:])) # (n, h, w, c*2) + refs = refs.permute(0, 3, 1, 2).contiguous() # (n, c*2, h, w) + + del embedding_target, emb, ass, corr_prob + + # spatial attention + attn = self.lrelu(self.spatial_attn(torch.cat([target, refs], dim=1))) + attn_mul = self.spatial_attn_mul2(self.lrelu(self.spatial_attn_mul1(attn))) + attn_add = self.spatial_attn_add2(self.lrelu(self.spatial_attn_add1(attn))) + attn_mul = torch.sigmoid(attn_mul) + + # after initialization, * 2 makes (attn_mul * 2) to be close to 1. + refs = refs * attn_mul * 2 + attn_add + + # fusion + feat = self.lrelu(self.feat_fusion(torch.cat([target, refs], dim=1))) + return feat[:, :, :h_input, :w_input] + +class SynBlock(nn.Module): + def __init__(self, nf, ks, groups=8): + super(SynBlock, self).__init__() + self.xN = 5 + self.offset_conv1 = nn.Conv2d( + nf*2, nf, 3, 1, 1, bias=True) # concat for diff + self.offset_conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True) + self.dyn_agg = DynAgg( + nf, + nf, + 3, + stride=1, + padding=1, + dilation=1, + deform_groups=groups, + extra_offset_mask=True) + self.head = MRAPAFusion(nf=nf, ref_nf=nf) + self.body = blocks.ResBlock(nf, nf, kernel_size=ks, stride=1, se=False) + self.tail = nn.Sequential( + nn.Conv2d(nf, nf, kernel_size=3, stride=1, padding=1)) + self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True) + + def forward(self, x, pre_offset_list, img_ref_feat_list, size): #fea, flow_list, vec, out_size): # + + swapped_feat_list = [] + for pre_offset, img_ref_feat in zip(pre_offset_list, img_ref_feat_list): + _, _, H, W = pre_offset.shape + TH, TW = size + scale = int(TW / W) + # print(scale, TW, W) + pre_offset = F.interpolate(pre_offset, size=size, mode='bilinear') * scale + offset = torch.cat([x, img_ref_feat], 1) + offset = self.lrelu(self.offset_conv1(offset)) + offset = self.lrelu(self.offset_conv2(offset)) + swapped_feat = self.lrelu( + self.dyn_agg([img_ref_feat, offset], pre_offset)) + swapped_feat_list.append(swapped_feat) + + h = self.head(x, swapped_feat_list) + h = self.body(h) + x + x = self.tail(h) + return x diff --git a/asvr_module/pytorch_pwc/correlation/correlation.py b/asvr_module/pytorch_pwc/correlation/correlation.py new file mode 100644 index 0000000..2c592cc --- /dev/null +++ b/asvr_module/pytorch_pwc/correlation/correlation.py @@ -0,0 +1,435 @@ +#!/usr/bin/env python + +import torch + +import cupy +import re + +kernel_Correlation_rearrange = ''' + extern "C" __global__ void kernel_Correlation_rearrange( + const int n, + const float* input, + float* output + ) { + int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; + + if (intIndex >= n) { + return; + } + + int intSample = blockIdx.z; + int intChannel = blockIdx.y; + + float fltValue = input[(((intSample * SIZE_1(input)) + intChannel) * SIZE_2(input) * SIZE_3(input)) + intIndex]; + + __syncthreads(); + + int intPaddedY = (intIndex / SIZE_3(input)) + 4; + int intPaddedX = (intIndex % SIZE_3(input)) + 4; + int intRearrange = ((SIZE_3(input) + 8) * intPaddedY) + intPaddedX; + + output[(((intSample * SIZE_1(output) * SIZE_2(output)) + intRearrange) * SIZE_1(input)) + intChannel] = fltValue; + } +''' + +kernel_Correlation_updateOutput = ''' + extern "C" __global__ void kernel_Correlation_updateOutput( + const int n, + const float* rbot0, + const float* rbot1, + float* top + ) { + extern __shared__ char patch_data_char[]; + + float *patch_data = (float *)patch_data_char; + + // First (upper left) position of kernel upper-left corner in current center position of neighborhood in image 1 + int x1 = blockIdx.x + 4; + int y1 = blockIdx.y + 4; + int item = blockIdx.z; + int ch_off = threadIdx.x; + + // Load 3D patch into shared shared memory + for (int j = 0; j < 1; j++) { // HEIGHT + for (int i = 0; i < 1; i++) { // WIDTH + int ji_off = (j + i) * SIZE_3(rbot0); + for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS + int idx1 = ((item * SIZE_1(rbot0) + y1+j) * SIZE_2(rbot0) + x1+i) * SIZE_3(rbot0) + ch; + int idxPatchData = ji_off + ch; + patch_data[idxPatchData] = rbot0[idx1]; + } + } + } + + __syncthreads(); + + __shared__ float sum[32]; + + // Compute correlation + for (int top_channel = 0; top_channel < SIZE_1(top); top_channel++) { + sum[ch_off] = 0; + + int s2o = top_channel % 9 - 4; + int s2p = top_channel / 9 - 4; + + for (int j = 0; j < 1; j++) { // HEIGHT + for (int i = 0; i < 1; i++) { // WIDTH + int ji_off = (j + i) * SIZE_3(rbot0); + for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS + int x2 = x1 + s2o; + int y2 = y1 + s2p; + + int idxPatchData = ji_off + ch; + int idx2 = ((item * SIZE_1(rbot0) + y2+j) * SIZE_2(rbot0) + x2+i) * SIZE_3(rbot0) + ch; + + sum[ch_off] += patch_data[idxPatchData] * rbot1[idx2]; + } + } + } + + __syncthreads(); + + if (ch_off == 0) { + float total_sum = 0; + for (int idx = 0; idx < 32; idx++) { + total_sum += sum[idx]; + } + const int sumelems = SIZE_3(rbot0); + const int index = ((top_channel*SIZE_2(top) + blockIdx.y)*SIZE_3(top))+blockIdx.x; + top[index + item*SIZE_1(top)*SIZE_2(top)*SIZE_3(top)] = total_sum / (float)sumelems; + } + } + } +''' + +kernel_Correlation_updateGradFirst = ''' + #define ROUND_OFF 50000 + + extern "C" __global__ void kernel_Correlation_updateGradFirst( + const int n, + const int intSample, + const float* rbot0, + const float* rbot1, + const float* gradOutput, + float* gradFirst, + float* gradSecond + ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { + int n = intIndex % SIZE_1(gradFirst); // channels + int l = (intIndex / SIZE_1(gradFirst)) % SIZE_3(gradFirst) + 4; // w-pos + int m = (intIndex / SIZE_1(gradFirst) / SIZE_3(gradFirst)) % SIZE_2(gradFirst) + 4; // h-pos + + // round_off is a trick to enable integer division with ceil, even for negative numbers + // We use a large offset, for the inner part not to become negative. + const int round_off = ROUND_OFF; + const int round_off_s1 = round_off; + + // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior: + int xmin = (l - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4) + int ymin = (m - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4) + + // Same here: + int xmax = (l - 4 + round_off_s1) - round_off; // floor (l - 4) + int ymax = (m - 4 + round_off_s1) - round_off; // floor (m - 4) + + float sum = 0; + if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) { + xmin = max(0,xmin); + xmax = min(SIZE_3(gradOutput)-1,xmax); + + ymin = max(0,ymin); + ymax = min(SIZE_2(gradOutput)-1,ymax); + + for (int p = -4; p <= 4; p++) { + for (int o = -4; o <= 4; o++) { + // Get rbot1 data: + int s2o = o; + int s2p = p; + int idxbot1 = ((intSample * SIZE_1(rbot0) + (m+s2p)) * SIZE_2(rbot0) + (l+s2o)) * SIZE_3(rbot0) + n; + float bot1tmp = rbot1[idxbot1]; // rbot1[l+s2o,m+s2p,n] + + // Index offset for gradOutput in following loops: + int op = (p+4) * 9 + (o+4); // index[o,p] + int idxopoffset = (intSample * SIZE_1(gradOutput) + op); + + for (int y = ymin; y <= ymax; y++) { + for (int x = xmin; x <= xmax; x++) { + int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p] + sum += gradOutput[idxgradOutput] * bot1tmp; + } + } + } + } + } + const int sumelems = SIZE_1(gradFirst); + const int bot0index = ((n * SIZE_2(gradFirst)) + (m-4)) * SIZE_3(gradFirst) + (l-4); + gradFirst[bot0index + intSample*SIZE_1(gradFirst)*SIZE_2(gradFirst)*SIZE_3(gradFirst)] = sum / (float)sumelems; + } } +''' + +kernel_Correlation_updateGradSecond = ''' + #define ROUND_OFF 50000 + + extern "C" __global__ void kernel_Correlation_updateGradSecond( + const int n, + const int intSample, + const float* rbot0, + const float* rbot1, + const float* gradOutput, + float* gradFirst, + float* gradSecond + ) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) { + int n = intIndex % SIZE_1(gradSecond); // channels + int l = (intIndex / SIZE_1(gradSecond)) % SIZE_3(gradSecond) + 4; // w-pos + int m = (intIndex / SIZE_1(gradSecond) / SIZE_3(gradSecond)) % SIZE_2(gradSecond) + 4; // h-pos + + // round_off is a trick to enable integer division with ceil, even for negative numbers + // We use a large offset, for the inner part not to become negative. + const int round_off = ROUND_OFF; + const int round_off_s1 = round_off; + + float sum = 0; + for (int p = -4; p <= 4; p++) { + for (int o = -4; o <= 4; o++) { + int s2o = o; + int s2p = p; + + //Get X,Y ranges and clamp + // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior: + int xmin = (l - 4 - s2o + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o) + int ymin = (m - 4 - s2p + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o) + + // Same here: + int xmax = (l - 4 - s2o + round_off_s1) - round_off; // floor (l - 4 - s2o) + int ymax = (m - 4 - s2p + round_off_s1) - round_off; // floor (m - 4 - s2p) + + if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) { + xmin = max(0,xmin); + xmax = min(SIZE_3(gradOutput)-1,xmax); + + ymin = max(0,ymin); + ymax = min(SIZE_2(gradOutput)-1,ymax); + + // Get rbot0 data: + int idxbot0 = ((intSample * SIZE_1(rbot0) + (m-s2p)) * SIZE_2(rbot0) + (l-s2o)) * SIZE_3(rbot0) + n; + float bot0tmp = rbot0[idxbot0]; // rbot1[l+s2o,m+s2p,n] + + // Index offset for gradOutput in following loops: + int op = (p+4) * 9 + (o+4); // index[o,p] + int idxopoffset = (intSample * SIZE_1(gradOutput) + op); + + for (int y = ymin; y <= ymax; y++) { + for (int x = xmin; x <= xmax; x++) { + int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p] + sum += gradOutput[idxgradOutput] * bot0tmp; + } + } + } + } + } + const int sumelems = SIZE_1(gradSecond); + const int bot1index = ((n * SIZE_2(gradSecond)) + (m-4)) * SIZE_3(gradSecond) + (l-4); + gradSecond[bot1index + intSample*SIZE_1(gradSecond)*SIZE_2(gradSecond)*SIZE_3(gradSecond)] = sum / (float)sumelems; + } } +''' + + +def cupy_kernel(strFunction, objVariables): + strKernel = globals()[strFunction] + + while True: + objMatch = re.search('(SIZE_)([0-4])(\()([^\)]*)(\))', strKernel) + + if objMatch is None: + break + # end + + intArg = int(objMatch.group(2)) + + strTensor = objMatch.group(4) + intSizes = objVariables[strTensor].size() + + strKernel = strKernel.replace(objMatch.group(), str(intSizes[intArg])) + # end + + while True: + objMatch = re.search('(VALUE_)([0-4])(\()([^\)]+)(\))', strKernel) + + if objMatch is None: + break + # end + + intArgs = int(objMatch.group(2)) + strArgs = objMatch.group(4).split(',') + + strTensor = strArgs[0] + intStrides = objVariables[strTensor].stride() + strIndex = ['((' + strArgs[intArg + 1].replace('{', '(').replace('}', ')').strip() + ')*' + str( + intStrides[intArg]) + ')' for intArg in range(intArgs)] + + strKernel = strKernel.replace(objMatch.group(0), strTensor + '[' + str.join('+', strIndex) + ']') + # end + + return strKernel + + +# end + +# @cupy.memoize(for_each_device=True) +# def cupy_launch(strFunction, strKernel): +# return cupy.cuda.compile_with_cache(strKernel).get_function(strFunction) + +compiled_modules = {} + +def get_device_id(): + return cupy.cuda.Device().id + +def get_or_compile_module(device_id, strKernel): + if device_id not in compiled_modules: + compiled_modules[device_id] = {} + + if strKernel not in compiled_modules[device_id]: + module = cupy.RawModule(code=strKernel) + compiled_modules[device_id][strKernel] = module + + return compiled_modules[device_id][strKernel] + +def cupy_launch(strFunction, strKernel): + device_id = get_device_id() + module = get_or_compile_module(device_id, strKernel) + function = module.get_function(strFunction) + return function + +# end + +class _FunctionCorrelation(torch.autograd.Function): + @staticmethod + def forward(self, first, second): + rbot0 = first.new_zeros([first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1]]) + rbot1 = first.new_zeros([first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1]]) + + self.save_for_backward(first, second, rbot0, rbot1) + + assert (first.is_contiguous() == True) + assert (second.is_contiguous() == True) + + output = first.new_zeros([first.shape[0], 81, first.shape[2], first.shape[3]]) + + if first.is_cuda == True: + n = first.shape[2] * first.shape[3] + cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', { + 'input': first, + 'output': rbot0 + }))( + grid=tuple([int((n + 16 - 1) / 16), first.shape[1], first.shape[0]]), + block=tuple([16, 1, 1]), + args=[n, first.data_ptr(), rbot0.data_ptr()] + ) + + n = second.shape[2] * second.shape[3] + cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', { + 'input': second, + 'output': rbot1 + }))( + grid=tuple([int((n + 16 - 1) / 16), second.shape[1], second.shape[0]]), + block=tuple([16, 1, 1]), + args=[n, second.data_ptr(), rbot1.data_ptr()] + ) + + n = output.shape[1] * output.shape[2] * output.shape[3] + cupy_launch('kernel_Correlation_updateOutput', cupy_kernel('kernel_Correlation_updateOutput', { + 'rbot0': rbot0, + 'rbot1': rbot1, + 'top': output + }))( + grid=tuple([output.shape[3], output.shape[2], output.shape[0]]), + block=tuple([32, 1, 1]), + shared_mem=first.shape[1] * 4, + args=[n, rbot0.data_ptr(), rbot1.data_ptr(), output.data_ptr()] + ) + + elif first.is_cuda == False: + raise NotImplementedError() + + # end + + return output + + # end + + @staticmethod + def backward(self, gradOutput): + first, second, rbot0, rbot1 = self.saved_tensors + + assert (gradOutput.is_contiguous() == True) + + gradFirst = first.new_zeros([first.shape[0], first.shape[1], first.shape[2], first.shape[3]]) if \ + self.needs_input_grad[0] == True else None + gradSecond = first.new_zeros([first.shape[0], first.shape[1], first.shape[2], first.shape[3]]) if \ + self.needs_input_grad[1] == True else None + + if first.is_cuda == True: + if gradFirst is not None: + for intSample in range(first.shape[0]): + n = first.shape[1] * first.shape[2] * first.shape[3] + cupy_launch('kernel_Correlation_updateGradFirst', + cupy_kernel('kernel_Correlation_updateGradFirst', { + 'rbot0': rbot0, + 'rbot1': rbot1, + 'gradOutput': gradOutput, + 'gradFirst': gradFirst, + 'gradSecond': None + }))( + grid=tuple([int((n + 512 - 1) / 512), 1, 1]), + block=tuple([512, 1, 1]), + args=[n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(), + gradFirst.data_ptr(), None] + ) + # end + # end + + if gradSecond is not None: + for intSample in range(first.shape[0]): + n = first.shape[1] * first.shape[2] * first.shape[3] + cupy_launch('kernel_Correlation_updateGradSecond', + cupy_kernel('kernel_Correlation_updateGradSecond', { + 'rbot0': rbot0, + 'rbot1': rbot1, + 'gradOutput': gradOutput, + 'gradFirst': None, + 'gradSecond': gradSecond + }))( + grid=tuple([int((n + 512 - 1) / 512), 1, 1]), + block=tuple([512, 1, 1]), + args=[n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(), None, + gradSecond.data_ptr()] + ) + # end + # end + + elif first.is_cuda == False: + raise NotImplementedError() + + # end + + return gradFirst, gradSecond + + +# end +# end + +def FunctionCorrelation(tenFirst, tenSecond): + return _FunctionCorrelation.apply(tenFirst, tenSecond) + + +# end + +class ModuleCorrelation(torch.nn.Module): + def __init__(self): + super(ModuleCorrelation, self).__init__() + + # end + + def forward(self, tenFirst, tenSecond): + return _FunctionCorrelation.apply(tenFirst, tenSecond) +# end +# end \ No newline at end of file diff --git a/asvr_module/pytorch_pwc/extract_flow.py b/asvr_module/pytorch_pwc/extract_flow.py new file mode 100644 index 0000000..95ebe75 --- /dev/null +++ b/asvr_module/pytorch_pwc/extract_flow.py @@ -0,0 +1,32 @@ +import math +import torch + +# im1_torch, im2_torch in shape (N, C, H, W) +def extract_flow_torch(model, im1_torch, im2_torch): + # interpolate image, make new_H, mew_W divide by 64 + assert im1_torch.shape == im2_torch.shape + N, C, H, W = im1_torch.shape + device = im1_torch.device + new_H = int(math.floor(math.ceil(H / 64.0) * 64.0)) + new_W = int(math.floor(math.ceil(W / 64.0) * 64.0)) + im1_torch = torch.nn.functional.interpolate(input=im1_torch, size=(new_H, new_W), mode='bilinear', + align_corners=False) + im2_torch = torch.nn.functional.interpolate(input=im2_torch, size=(new_H, new_W), mode='bilinear', + align_corners=False) + model.eval() + with torch.no_grad(): + flo12 = model(im1_torch, im2_torch) + flo12 = 20.0 * torch.nn.functional.interpolate(input=flo12, size=(H, W), mode='bilinear', + align_corners=False) + flo12[:, 0, :, :] *= float(W) / float(new_W) + flo12[:, 1, :, :] *= float(H) / float(new_H) + return flo12 + +# im1_np, im2_np in shape (C, H, W) +def extract_flow_np(model, im1_np, im2_np): + im1_torch = torch.from_numpy(im1_np).unsqueeze(0).to(torch.device('cuda')) + im2_torch = torch.from_numpy(im2_np).unsqueeze(0).to(torch.device('cuda')) + flo12_torch = extract_flow_torch(model, im1_torch, im2_torch) + flo12_np = flo12_torch.detach().cpu().squeeze(0).numpy() + return flo12_np + diff --git a/asvr_module/pytorch_pwc/pwc.py b/asvr_module/pytorch_pwc/pwc.py new file mode 100644 index 0000000..4c7a94c --- /dev/null +++ b/asvr_module/pytorch_pwc/pwc.py @@ -0,0 +1,242 @@ +from .correlation import correlation # the custom cost volume layer +import torch + +########################################################## + +assert(int(str('').join(torch.__version__.split('.')[0:2])) >= 13) # requires at least pytorch version 1.3.0 + +# torch.set_grad_enabled(False) + +torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance + +arguments_strModel = 'default' + +backwarp_tenGrid = {} +backwarp_tenPartial = {} + +def backwarp(tenInput, tenFlow): + if (str(tenFlow.shape)+str(tenFlow.device)) not in backwarp_tenGrid: + tenHor = torch.linspace(-1.0 + (1.0 / tenFlow.shape[3]), 1.0 - (1.0 / tenFlow.shape[3]), tenFlow.shape[3]).view(1, 1, 1, -1).expand(-1, -1, tenFlow.shape[2], -1) + tenVer = torch.linspace(-1.0 + (1.0 / tenFlow.shape[2]), 1.0 - (1.0 / tenFlow.shape[2]), tenFlow.shape[2]).view(1, 1, -1, 1).expand(-1, -1, -1, tenFlow.shape[3]) + + backwarp_tenGrid[str(tenFlow.shape) + str(tenFlow.device)] = torch.cat([ tenHor, tenVer ], 1).to(tenFlow.device) + # end + + if (str(tenFlow.shape)+str(tenFlow.device)) not in backwarp_tenPartial: + backwarp_tenPartial[str(tenFlow.shape)+str(tenFlow.device)] = tenFlow.new_ones([ tenFlow.shape[0], 1, tenFlow.shape[2], tenFlow.shape[3] ]) + # end + + tenFlow = torch.cat([ tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0) ], 1) + tenInput = torch.cat([ tenInput, backwarp_tenPartial[str(tenFlow.shape)+str(tenFlow.device)] ], 1) + + tenOutput = torch.nn.functional.grid_sample(input=tenInput, grid=(backwarp_tenGrid[str(tenFlow.shape) + str(tenFlow.device)] + tenFlow).permute(0, 2, 3, 1), mode='bilinear', padding_mode='zeros', align_corners=False) + + tenMask = tenOutput[:, -1:, :, :]; tenMask[tenMask > 0.999] = 1.0; tenMask[tenMask < 1.0] = 0.0 + + return tenOutput[:, :-1, :, :] * tenMask +# end + +########################################################## + +class PWCNet(torch.nn.Module): + def __init__(self): + super(PWCNet, self).__init__() + + class Extractor(torch.nn.Module): + def __init__(self): + super(Extractor, self).__init__() + + self.netOne = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1) + ) + + self.netTwo = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=2, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1) + ) + + self.netThr = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1) + ) + + self.netFou = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=64, out_channels=96, kernel_size=3, stride=2, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1) + ) + + self.netFiv = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=96, out_channels=128, kernel_size=3, stride=2, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1) + ) + + self.netSix = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=128, out_channels=196, kernel_size=3, stride=2, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1) + ) + # end + + def forward(self, tenInput): + tenOne = self.netOne(tenInput) + tenTwo = self.netTwo(tenOne) + tenThr = self.netThr(tenTwo) + tenFou = self.netFou(tenThr) + tenFiv = self.netFiv(tenFou) + tenSix = self.netSix(tenFiv) + + return [ tenOne, tenTwo, tenThr, tenFou, tenFiv, tenSix ] + # end + # end + + class Decoder(torch.nn.Module): + def __init__(self, intLevel): + super(Decoder, self).__init__() + + intPrevious = [ None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None ][intLevel + 1] + intCurrent = [ None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None ][intLevel + 0] + + if intLevel < 6: self.netUpflow = torch.nn.ConvTranspose2d(in_channels=2, out_channels=2, kernel_size=4, stride=2, padding=1) + if intLevel < 6: self.netUpfeat = torch.nn.ConvTranspose2d(in_channels=intPrevious + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=4, stride=2, padding=1) + if intLevel < 6: self.fltBackwarp = [ None, None, None, 5.0, 2.5, 1.25, 0.625, None ][intLevel + 1] + + self.netOne = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=intCurrent, out_channels=128, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1) + ) + + self.netTwo = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=intCurrent + 128, out_channels=128, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1) + ) + + self.netThr = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=intCurrent + 128 + 128, out_channels=96, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1) + ) + + self.netFou = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96, out_channels=64, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1) + ) + + self.netFiv = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64, out_channels=32, kernel_size=3, stride=1, padding=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1) + ) + + self.netSix = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=3, stride=1, padding=1) + ) + # end + + def forward(self, tenFirst, tenSecond, objPrevious): + tenFlow = None + tenFeat = None + + if objPrevious is None: + tenFlow = None + tenFeat = None + tenVolume = torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tenFirst=tenFirst, tenSecond=tenSecond), negative_slope=0.1, inplace=False) + + tenFeat = torch.cat([ tenVolume ], 1) + + elif objPrevious is not None: + tenFlow = self.netUpflow(objPrevious['tenFlow']) + tenFeat = self.netUpfeat(objPrevious['tenFeat']) + + tenVolume = torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tenFirst=tenFirst, tenSecond=backwarp(tenInput=tenSecond, tenFlow=tenFlow * self.fltBackwarp)), negative_slope=0.1, inplace=False) + + tenFeat = torch.cat([ tenVolume, tenFirst, tenFlow, tenFeat ], 1) + + # end + + tenFeat = torch.cat([ self.netOne(tenFeat), tenFeat ], 1) + tenFeat = torch.cat([ self.netTwo(tenFeat), tenFeat ], 1) + tenFeat = torch.cat([ self.netThr(tenFeat), tenFeat ], 1) + tenFeat = torch.cat([ self.netFou(tenFeat), tenFeat ], 1) + tenFeat = torch.cat([ self.netFiv(tenFeat), tenFeat ], 1) + + tenFlow = self.netSix(tenFeat) + + return { + 'tenFlow': tenFlow, + 'tenFeat': tenFeat + } + # end + # end + + class Refiner(torch.nn.Module): + def __init__(self): + super(Refiner, self).__init__() + + self.netMain = torch.nn.Sequential( + torch.nn.Conv2d(in_channels=81 + 32 + 2 + 2 + 128 + 128 + 96 + 64 + 32, out_channels=128, kernel_size=3, stride=1, padding=1, dilation=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=2, dilation=2), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=4, dilation=4), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=128, out_channels=96, kernel_size=3, stride=1, padding=8, dilation=8), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=96, out_channels=64, kernel_size=3, stride=1, padding=16, dilation=16), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1, dilation=1), + torch.nn.LeakyReLU(inplace=False, negative_slope=0.1), + torch.nn.Conv2d(in_channels=32, out_channels=2, kernel_size=3, stride=1, padding=1, dilation=1) + ) + # end + + def forward(self, tenInput): + return self.netMain(tenInput) + # end + # end + + self.netExtractor = Extractor() + + self.netTwo = Decoder(2) + self.netThr = Decoder(3) + self.netFou = Decoder(4) + self.netFiv = Decoder(5) + self.netSix = Decoder(6) + + self.netRefiner = Refiner() + + self.load_state_dict({ strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.hub.load_state_dict_from_url(url='http://content.sniklaus.com/github/pytorch-pwc/network-' + arguments_strModel + '.pytorch', file_name='pwc-' + arguments_strModel).items() }) + # end + + def forward(self, tenFirst, tenSecond): + tenFirst = self.netExtractor(tenFirst) + tenSecond = self.netExtractor(tenSecond) + + objEstimate = self.netSix(tenFirst[-1], tenSecond[-1], None) + objEstimate = self.netFiv(tenFirst[-2], tenSecond[-2], objEstimate) + objEstimate = self.netFou(tenFirst[-3], tenSecond[-3], objEstimate) + objEstimate = self.netThr(tenFirst[-4], tenSecond[-4], objEstimate) + objEstimate = self.netTwo(tenFirst[-5], tenSecond[-5], objEstimate) + + return objEstimate['tenFlow'] + self.netRefiner(objEstimate['tenFeat']) \ No newline at end of file diff --git a/asvr_module/refsrrnn_adists_fgda_only_future.py b/asvr_module/refsrrnn_adists_fgda_only_future.py new file mode 100644 index 0000000..a7e63a1 --- /dev/null +++ b/asvr_module/refsrrnn_adists_fgda_only_future.py @@ -0,0 +1,214 @@ +from .components import ResBlocks, D +from .pytorch_pwc.extract_flow import extract_flow_torch +from .pytorch_pwc.pwc import PWCNet +import torch +import torch.nn as nn +from .warp import warp +import torch.nn.functional as F +from .cuf_siren import KerP, CUF +# # from models.refsrrnn_fgda import SynBlock +from .network_vidue_refsr_cuf3_siren import SynBlock +from .ADISTS_allstage import ADISTS + +def std(img, window_size=3): + assert window_size % 2 == 1 + pad = window_size // 2 + + # calculate std on the mean image of the color channels + # img = torch.mean(img, dim=1, keepdim=True) + N, C, H, W = img.shape + img = nn.functional.pad(img, [pad] * 4, mode='reflect') + img = nn.functional.unfold(img, kernel_size=window_size) + img = img.view(N, C, window_size * window_size, H, W) + img = img - torch.mean(img, dim=2, keepdim=True) + # print(img.shape) + img = img * img + img = torch.mean(img, dim=2, keepdim=True) + # print(img.shape) + img = torch.sqrt(img) + # print(img.shape) + img = img.squeeze(2) + return img + +def generate_alpha(input, lower=1, upper=5): + N, C, H, W = input.shape + ratio = input.new_ones((N, C, H, W)) * 0.5 + input_std = std(input) + ratio[input_std < lower] = torch.sigmoid((input_std - lower))[input_std < lower] + ratio[input_std > upper] = torch.sigmoid((input_std - upper))[input_std > upper] + ratio = ratio.detach() + + return ratio + +class RefsrRNN(nn.Module): + def __init__(self, img_channels=3, num_resblocks=15, num_channels=64, count=2): + super(RefsrRNN, self).__init__() + self.num_channels = num_channels + self.pwcnet = PWCNet() + self.forward_rnn = ResBlocks(input_channels=img_channels + num_channels*2, num_resblocks=num_resblocks, num_channels=num_channels) + self.forward_rnn2 = ResBlocks(input_channels=img_channels + num_channels*2, num_resblocks=num_resblocks, num_channels=num_channels) + self.d = D(in_channels=num_channels * 2, mid_channels=num_channels * 2, out_channels=num_channels) + self.kernel_predict = KerP(num_channels) + self.upsample = CUF(num_channels) + self.predict = SynBlock(num_channels, ks=3) + self.count = count + self.adists = ADISTS() + self.multi_prior_fusion = nn.Conv2d(64+128+256+512+512,num_channels,1,1,0) + for param in self.pwcnet.parameters(): + param.requires_grad = False + for param in self.adists.parameters(): + param.requires_grad = False + + # def trainable_parameters(self): + # return [{'params':self.forward_rnn.parameters()}, {'params':self.d.parameters()}, {'params':self.kernel_predict.parameters()}, {'params':self.upsample.parameters()}] + + def forward(self, images, scale, coord, cell, isTrain=True): + seqn_not_pad = torch.stack(images, dim=1) + N, T, C, H, W = seqn_not_pad.shape + seqn_not_pad_ = F.interpolate(seqn_not_pad.view(-1, C, H, W).contiguous(), size=(coord.shape[-3:-1]), mode='bicubic') + seqn_not_pad_ = seqn_not_pad_.view(N, T, C, coord.shape[-3], coord.shape[-2]).contiguous() + seqdn = torch.empty_like(seqn_not_pad_) + del seqn_not_pad_ + + # border_queue = [] + # size_h, size_w = int(H * self.border_ratio), int(W * self.border_ratio) + + # reflect pad seqn and noise_level_map + seqn = torch.empty((N, T + self.count, C, H, W), device=seqn_not_pad.device) + seqn[:, :T] = seqn_not_pad + for i in range(self.count): + seqn[:, T + i] = seqn_not_pad[:, T - 2 - i] + + # flatten_map = generate_alpha(seqn.reshape(-1, C, H, W).contiguous()*255.).view(N, T+self.count*2, C, H, W).contiguous() + # flatten_map = self.adists(seqn.reshape(-1, C, H, W).contiguous()).view(N, T+self.count, self.num_channels, H, W).contiguous() + + inp1, inp2 = seqn[:, 1:].reshape(-1, C, H, W).contiguous(), seqn[:, :-1].reshape(-1, C, H, W).contiguous() + flow = extract_flow_torch(self.pwcnet, inp1, inp2) + flow_queue = flow.view(N, T + self.count-1, 2, H, W).contiguous() + + hidden_list = [] + prior_list = [] + init_forward_h = torch.zeros((N, self.num_channels, H, W), device=seqn.device) + flatten_map = self.multi_prior_fusion(self.adists(seqn[:, 0])) + h_ = self.forward_rnn(torch.cat((seqn[:, 0], flatten_map, init_forward_h), dim=1)) + hidden_list.append(h_) + prior_list.append(flatten_map) + arb_up_kernel = self.kernel_predict(h_, coord, cell) + for j in range(1, self.count): + h_, _ = warp(h_, flow_queue[:, j-1]) + flatten_map = self.multi_prior_fusion(self.adists(seqn[:, j])) + h_ = self.forward_rnn(torch.cat((seqn[:, j], flatten_map, h_), dim=1)) + hidden_list.append(h_) + prior_list.append(flatten_map) + + for i in range(self.count, T + self.count): + # if i == self.count: + # for j in range(i, i+self.count): + # h_, _ = warp(h_, flow_queue[:, j - 1]) + # h_ = self.forward_rnn(torch.cat((seqn[:, j], flatten_map[:, j], h_), dim=1)) + # hidden_list.append(h_) + # else: + # #print(hidden_list[-1].shape) + # #print(flow_queue[:, i + self.count - 1].shape) + h_, _ = warp(hidden_list[-1], flow_queue[:, i-1]) + flatten_map = self.multi_prior_fusion(self.adists(seqn[:, i])) + h_ = self.forward_rnn(torch.cat((seqn[:, i], flatten_map, h_), dim=1)) + hidden_list.append(h_) + prior_list.append(flatten_map) + if i > self.count: + hidden_list.pop(0) + prior_list.pop(0) + + assert len(hidden_list) == self.count+1 + + refsr_flow_list = [] + for j in range(1, self.count+1): + refsr_flow_list.append(extract_flow_torch(self.pwcnet, seqn[:, i-self.count], seqn[:, i-self.count + j])) + + assert len(refsr_flow_list) == self.count + + hidden = hidden_list[0] + fusion_h = self.predict(hidden, refsr_flow_list, hidden_list[1:], hidden.size()[-2:]) + + h = self.forward_rnn2(torch.cat((seqn[:, i-self.count], prior_list[0], fusion_h), dim=1)) + res = self.d(torch.cat([h, hidden], dim=1)) + seqdn[:, i-self.count] = self.upsample(res, arb_up_kernel, coord, seqn[:, i-self.count]) + + + return seqdn + + def test_forward(self, images, arb_up_kernel, coord): + seqn_not_pad = torch.stack(images, dim=1) + print(f"{seqn_not_pad.shape=}") + N, T, C, H, W = seqn_not_pad.shape + seqn_not_pad_ = F.interpolate(seqn_not_pad.view(-1, C, H, W).contiguous(), size=(coord.shape[-3:-1]), mode='bicubic') + seqn_not_pad_ = seqn_not_pad_.view(N, T, C, coord.shape[-3], coord.shape[-2]).contiguous() + seqdn = torch.empty_like(seqn_not_pad_) + del seqn_not_pad_ + + # border_queue = [] + # size_h, size_w = int(H * self.border_ratio), int(W * self.border_ratio) + + # reflect pad seqn and noise_level_map + seqn = torch.empty((N, T + self.count, C, H, W), device=seqn_not_pad.device) + seqn[:, :T] = seqn_not_pad + for i in range(self.count): + seqn[:, T + i] = seqn_not_pad[:, T - 2 - i] + + # flatten_map = generate_alpha(seqn.reshape(-1, C, H, W).contiguous()*255.).view(N, T+self.count*2, C, H, W).contiguous() + # flatten_map = self.adists(seqn.reshape(-1, C, H, W).contiguous()).view(N, T+self.count, self.num_channels, H, W).contiguous() + + inp1, inp2 = seqn[:, 1:].reshape(-1, C, H, W).contiguous(), seqn[:, :-1].reshape(-1, C, H, W).contiguous() + flow = extract_flow_torch(self.pwcnet, inp1, inp2) + flow_queue = flow.view(N, T + self.count-1, 2, H, W).contiguous() + + hidden_list = [] + prior_list = [] + init_forward_h = torch.zeros((N, self.num_channels, H, W), device=seqn.device) + flatten_map = self.multi_prior_fusion(self.adists(seqn[:, 0])) + h_ = self.forward_rnn(torch.cat((seqn[:, 0], flatten_map, init_forward_h), dim=1)) + hidden_list.append(h_) + prior_list.append(flatten_map) + #arb_up_kernel = self.kernel_predict(h_, coord, cell) + for j in range(1, self.count): + h_, _ = warp(h_, flow_queue[:, j-1]) + flatten_map = self.multi_prior_fusion(self.adists(seqn[:, j])) + h_ = self.forward_rnn(torch.cat((seqn[:, j], flatten_map, h_), dim=1)) + hidden_list.append(h_) + prior_list.append(flatten_map) + + for i in range(self.count, T + self.count): + # if i == self.count: + # for j in range(i, i+self.count): + # h_, _ = warp(h_, flow_queue[:, j - 1]) + # h_ = self.forward_rnn(torch.cat((seqn[:, j], flatten_map[:, j], h_), dim=1)) + # hidden_list.append(h_) + # else: + # #print(hidden_list[-1].shape) + # #print(flow_queue[:, i + self.count - 1].shape) + h_, _ = warp(hidden_list[-1], flow_queue[:, i-1]) + flatten_map = self.multi_prior_fusion(self.adists(seqn[:, i])) + h_ = self.forward_rnn(torch.cat((seqn[:, i], flatten_map, h_), dim=1)) + hidden_list.append(h_) + prior_list.append(flatten_map) + if i > self.count: + hidden_list.pop(0) + prior_list.pop(0) + + assert len(hidden_list) == self.count+1 + + refsr_flow_list = [] + for j in range(1, self.count+1): + refsr_flow_list.append(extract_flow_torch(self.pwcnet, seqn[:, i-self.count], seqn[:, i-self.count + j])) + + assert len(refsr_flow_list) == self.count + + hidden = hidden_list[0] + fusion_h = self.predict(hidden, refsr_flow_list, hidden_list[1:], hidden.size()[-2:]) + + h = self.forward_rnn2(torch.cat((seqn[:, i-self.count], prior_list[0], fusion_h), dim=1)) + res = self.d(torch.cat([h, hidden], dim=1)) + seqdn[:, i-self.count] = self.upsample(res, arb_up_kernel, coord, seqn[:, i-self.count]) + print(f"{seqdn.shape=}") + + return seqdn diff --git a/asvr_module/warp.py b/asvr_module/warp.py new file mode 100644 index 0000000..087fcf2 --- /dev/null +++ b/asvr_module/warp.py @@ -0,0 +1,33 @@ +import torch + +def warp(x, flo): + ''' + warp an image/tensor (im2) back to im1, according to the optical flow + x: [B, C, H, W] (im2) + flo: [B, 2, H, W] (flow) + ''' + B, C, H, W = x.size() + # mesh grid + xx = torch.arange(0, W, device=x.device).view(1, -1).repeat(H, 1) + yy = torch.arange(0, H, device=x.device).view(-1, 1).repeat(1, W) + xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1) + yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1) + grid = torch.cat((xx, yy), 1).float() + + if x.is_cuda: + grid = grid.to(x.device) + vgrid = torch.autograd.Variable(grid) + flo + + # scale grid to [-1, 1] + vgrid[:, 0, :, :] = 2.0 * vgrid[:, 0, :, :].clone() / max(W - 1, 1) - 1.0 + vgrid[:, 1, :, :] = 2.0 * vgrid[:, 1, :, :].clone() / max(H - 1, 1) - 1.0 + + vgrid = vgrid.permute(0, 2, 3, 1) + output = torch.nn.functional.grid_sample(x, vgrid, align_corners=True) + mask = torch.autograd.Variable(torch.ones((B, C, H, W), device=x.device)) + mask = torch.nn.functional.grid_sample(mask, vgrid, align_corners=True) + + mask[mask < 0.9999] = 0 + mask[mask > 0] = 1 + + return output * mask, mask \ No newline at end of file diff --git a/asvr_run.py b/asvr_run.py new file mode 100644 index 0000000..a10f385 --- /dev/null +++ b/asvr_run.py @@ -0,0 +1,47 @@ + +from pathlib import Path + +import torch +import torch.nn.functional as F +import tqdm +from comfy.utils import ProgressBar + +from .constants import ASVR_MODEL_PATH +from .asvr_module.asvr_model import ASVR_Config, ASVR + +base_dir = Path(__file__).resolve().parent + +class ASVR_VideoSR: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "src_video": ("IMAGE",), + "height_scale": ( + "FLOAT", + {"default": 8.0, "min": 0.1, "step": 0.01}, + ), + "width_scale": ( + "FLOAT", + {"default": 8.0, "min": 0.1, "step": 0.01}, + ), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("res_video",) + FUNCTION = "todo" + CATEGORY = "FM_nodes" + + def todo( + self, + src_video: torch.Tensor, + height_scale: float, + width_scale: float + ): + asvr_runner = ASVR(ASVR_Config( + f"{base_dir}/{ASVR_MODEL_PATH}", + space_scale=(height_scale, width_scale) + )) + output = asvr_runner.inference(src_video) + return (output, ) diff --git a/constants.py b/constants.py index 8f31860..58b3d96 100644 --- a/constants.py +++ b/constants.py @@ -1,4 +1,5 @@ WFEN_MODEL_PATH = "models/wfen/WFEN.pth" REALVIFORMER_MODEL_PATH = "models/realviformer/weights.pth" PROPIH_VGG_MODEL_PATH = "models/propih/vgg_normalised.pth" -PROPIH_G_MODEL_PATH = "models/propih/latest_net_G.pth" \ No newline at end of file +PROPIH_G_MODEL_PATH = "models/propih/latest_net_G.pth" +ASVR_MODEL_PATH = "models/asvr/refsrrnn_cuf_siren_adists_allstage_only_future_t2.pth" \ No newline at end of file