marduk191andClaude Opus 4.5
b4428dabc1
Features: - NAFNet denoising (SIDD-trained models) - NAFNet deblurring (GoPro/REDS-trained models) - NAFSSR stereo image super-resolution (2x/4x) - Tiled processing for large images - Automatic model caching Nodes: - NAFNet Load Model - NAFNet Restore (generic with tiling support) - NAFNet Denoise (SIDD) - NAFNet Deblur (GoPro) - NAFSSR Stereo Super-Resolution Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
276 lines
9.5 KiB
Python
276 lines
9.5 KiB
Python
# ------------------------------------------------------------------------
|
|
# Copyright (c) 2022 megvii-model. All Rights Reserved.
|
|
# ------------------------------------------------------------------------
|
|
# NAFNet: Simple Baselines for Image Restoration
|
|
# https://github.com/megvii-research/NAFNet
|
|
# ------------------------------------------------------------------------
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class LayerNormFunction(torch.autograd.Function):
|
|
@staticmethod
|
|
def forward(ctx, x, weight, bias, eps):
|
|
ctx.eps = eps
|
|
N, C, H, W = x.size()
|
|
mu = x.mean(1, keepdim=True)
|
|
var = (x - mu).pow(2).mean(1, keepdim=True)
|
|
y = (x - mu) / (var + eps).sqrt()
|
|
ctx.save_for_backward(y, var, weight)
|
|
y = weight.view(1, C, 1, 1) * y + bias.view(1, C, 1, 1)
|
|
return y
|
|
|
|
@staticmethod
|
|
def backward(ctx, grad_output):
|
|
eps = ctx.eps
|
|
N, C, H, W = grad_output.size()
|
|
y, var, weight = ctx.saved_tensors
|
|
g = grad_output * weight.view(1, C, 1, 1)
|
|
mean_g = g.mean(dim=1, keepdim=True)
|
|
mean_gy = (g * y).mean(dim=1, keepdim=True)
|
|
gx = 1. / torch.sqrt(var + eps) * (g - y * mean_gy - mean_g)
|
|
return gx, (grad_output * y).sum(dim=3).sum(dim=2).sum(dim=0), grad_output.sum(dim=3).sum(dim=2).sum(dim=0), None
|
|
|
|
|
|
class LayerNorm2d(nn.Module):
|
|
def __init__(self, channels, eps=1e-6):
|
|
super(LayerNorm2d, self).__init__()
|
|
self.register_parameter('weight', nn.Parameter(torch.ones(channels)))
|
|
self.register_parameter('bias', nn.Parameter(torch.zeros(channels)))
|
|
self.eps = eps
|
|
|
|
def forward(self, x):
|
|
return LayerNormFunction.apply(x, self.weight, self.bias, self.eps)
|
|
|
|
|
|
class SimpleGate(nn.Module):
|
|
def forward(self, x):
|
|
x1, x2 = x.chunk(2, dim=1)
|
|
return x1 * x2
|
|
|
|
|
|
class NAFBlock(nn.Module):
|
|
def __init__(self, c, DW_Expand=2, FFN_Expand=2, drop_out_rate=0.):
|
|
super().__init__()
|
|
dw_channel = c * DW_Expand
|
|
self.conv1 = nn.Conv2d(c, dw_channel, 1, 1, 0, bias=True)
|
|
self.conv2 = nn.Conv2d(dw_channel, dw_channel, 3, 1, 1, groups=dw_channel, bias=True)
|
|
self.conv3 = nn.Conv2d(dw_channel // 2, c, 1, 1, 0, bias=True)
|
|
|
|
# Simplified Channel Attention
|
|
self.sca = nn.Sequential(
|
|
nn.AdaptiveAvgPool2d(1),
|
|
nn.Conv2d(dw_channel // 2, dw_channel // 2, 1, 1, 0, bias=True),
|
|
)
|
|
|
|
self.sg = SimpleGate()
|
|
|
|
ffn_channel = FFN_Expand * c
|
|
self.conv4 = nn.Conv2d(c, ffn_channel, 1, 1, 0, bias=True)
|
|
self.conv5 = nn.Conv2d(ffn_channel // 2, c, 1, 1, 0, bias=True)
|
|
|
|
self.norm1 = LayerNorm2d(c)
|
|
self.norm2 = LayerNorm2d(c)
|
|
|
|
self.dropout1 = nn.Dropout(drop_out_rate) if drop_out_rate > 0. else nn.Identity()
|
|
self.dropout2 = nn.Dropout(drop_out_rate) if drop_out_rate > 0. else nn.Identity()
|
|
|
|
self.beta = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
|
|
self.gamma = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
|
|
|
|
def forward(self, inp):
|
|
x = inp
|
|
x = self.norm1(x)
|
|
x = self.conv1(x)
|
|
x = self.conv2(x)
|
|
x = self.sg(x)
|
|
x = x * self.sca(x)
|
|
x = self.conv3(x)
|
|
x = self.dropout1(x)
|
|
|
|
y = inp + x * self.beta
|
|
|
|
x = self.conv4(self.norm2(y))
|
|
x = self.sg(x)
|
|
x = self.conv5(x)
|
|
x = self.dropout2(x)
|
|
|
|
return y + x * self.gamma
|
|
|
|
|
|
class NAFNet(nn.Module):
|
|
def __init__(self, img_channel=3, width=16, middle_blk_num=1, enc_blk_nums=[], dec_blk_nums=[]):
|
|
super().__init__()
|
|
|
|
self.intro = nn.Conv2d(img_channel, width, 3, 1, 1, bias=True)
|
|
self.ending = nn.Conv2d(width, img_channel, 3, 1, 1, bias=True)
|
|
|
|
self.encoders = nn.ModuleList()
|
|
self.decoders = nn.ModuleList()
|
|
self.middle_blks = nn.ModuleList()
|
|
self.ups = nn.ModuleList()
|
|
self.downs = nn.ModuleList()
|
|
|
|
chan = width
|
|
for num in enc_blk_nums:
|
|
self.encoders.append(
|
|
nn.Sequential(*[NAFBlock(chan) for _ in range(num)])
|
|
)
|
|
self.downs.append(
|
|
nn.Conv2d(chan, 2 * chan, 2, 2)
|
|
)
|
|
chan = chan * 2
|
|
|
|
self.middle_blks = nn.Sequential(*[NAFBlock(chan) for _ in range(middle_blk_num)])
|
|
|
|
for num in dec_blk_nums:
|
|
self.ups.append(
|
|
nn.Sequential(
|
|
nn.Conv2d(chan, chan * 2, 1, bias=False),
|
|
nn.PixelShuffle(2)
|
|
)
|
|
)
|
|
chan = chan // 2
|
|
self.decoders.append(
|
|
nn.Sequential(*[NAFBlock(chan) for _ in range(num)])
|
|
)
|
|
|
|
self.padder_size = 2 ** len(self.encoders)
|
|
|
|
def forward(self, inp):
|
|
B, C, H, W = inp.shape
|
|
inp = self.check_image_size(inp)
|
|
|
|
x = self.intro(inp)
|
|
|
|
encs = []
|
|
for encoder, down in zip(self.encoders, self.downs):
|
|
x = encoder(x)
|
|
encs.append(x)
|
|
x = down(x)
|
|
|
|
x = self.middle_blks(x)
|
|
|
|
for decoder, up, enc_skip in zip(self.decoders, self.ups, encs[::-1]):
|
|
x = up(x)
|
|
x = x + enc_skip
|
|
x = decoder(x)
|
|
|
|
x = self.ending(x)
|
|
x = x + inp
|
|
|
|
return x[:, :, :H, :W]
|
|
|
|
def check_image_size(self, x):
|
|
_, _, h, w = x.size()
|
|
mod_pad_h = (self.padder_size - h % self.padder_size) % self.padder_size
|
|
mod_pad_w = (self.padder_size - w % self.padder_size) % self.padder_size
|
|
x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h))
|
|
return x
|
|
|
|
|
|
# ------------------------------------------------------------------------
|
|
# NAFSSR: Stereo Image Super-Resolution Using NAFNet
|
|
# ------------------------------------------------------------------------
|
|
|
|
class SCAM(nn.Module):
|
|
"""Stereo Cross Attention Module"""
|
|
def __init__(self, c):
|
|
super().__init__()
|
|
self.scale = c ** -0.5
|
|
self.norm_l = LayerNorm2d(c)
|
|
self.norm_r = LayerNorm2d(c)
|
|
self.l_proj1 = nn.Conv2d(c, c, kernel_size=1, stride=1, padding=0)
|
|
self.r_proj1 = nn.Conv2d(c, c, kernel_size=1, stride=1, padding=0)
|
|
self.beta = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
|
|
self.gamma = nn.Parameter(torch.zeros((1, c, 1, 1)), requires_grad=True)
|
|
self.l_proj2 = nn.Conv2d(c, c, kernel_size=1, stride=1, padding=0)
|
|
self.r_proj2 = nn.Conv2d(c, c, kernel_size=1, stride=1, padding=0)
|
|
|
|
def forward(self, x_l, x_r):
|
|
Q_l = self.l_proj1(self.norm_l(x_l)).permute(0, 2, 3, 1)
|
|
Q_r_T = self.r_proj1(self.norm_r(x_r)).permute(0, 2, 1, 3)
|
|
V_l = self.l_proj2(x_l).permute(0, 2, 3, 1)
|
|
V_r = self.r_proj2(x_r).permute(0, 2, 3, 1)
|
|
attention = torch.matmul(Q_l, Q_r_T) * self.scale
|
|
F_r2l = torch.matmul(torch.softmax(attention, dim=-1), V_r)
|
|
F_l2r = torch.matmul(torch.softmax(attention.permute(0, 1, 3, 2), dim=-1), V_l)
|
|
F_r2l = F_r2l.permute(0, 3, 1, 2) * self.beta
|
|
F_l2r = F_l2r.permute(0, 3, 1, 2) * self.gamma
|
|
return x_l + F_r2l, x_r + F_l2r
|
|
|
|
|
|
class DropPath(nn.Module):
|
|
def __init__(self, drop_rate, module):
|
|
super().__init__()
|
|
self.drop_rate = drop_rate
|
|
self.module = module
|
|
|
|
def forward(self, *feats):
|
|
if self.training and np.random.rand() < self.drop_rate:
|
|
return feats
|
|
new_feats = self.module(*feats)
|
|
factor = 1. / (1 - self.drop_rate) if self.training else 1.
|
|
if self.training and factor != 1.:
|
|
new_feats = tuple([x + factor * (new_x - x) for x, new_x in zip(feats, new_feats)])
|
|
return new_feats
|
|
|
|
|
|
class NAFBlockSR(nn.Module):
|
|
"""NAFBlock for Super-Resolution"""
|
|
def __init__(self, c, fusion=False, drop_out_rate=0.):
|
|
super().__init__()
|
|
self.blk = NAFBlock(c, drop_out_rate=drop_out_rate)
|
|
self.fusion = SCAM(c) if fusion else None
|
|
|
|
def forward(self, *feats):
|
|
feats = tuple([self.blk(x) for x in feats])
|
|
if self.fusion:
|
|
feats = self.fusion(*feats)
|
|
return feats
|
|
|
|
|
|
class MySequential(nn.Sequential):
|
|
def forward(self, *inputs):
|
|
for module in self._modules.values():
|
|
if isinstance(inputs, tuple):
|
|
inputs = module(*inputs)
|
|
else:
|
|
inputs = module(inputs)
|
|
return inputs
|
|
|
|
|
|
class NAFNetSR(nn.Module):
|
|
"""NAFNet for Super-Resolution"""
|
|
def __init__(self, up_scale=4, width=48, num_blks=16, img_channel=3,
|
|
drop_path_rate=0., drop_out_rate=0., fusion_from=-1, fusion_to=-1, dual=False):
|
|
super().__init__()
|
|
self.dual = dual
|
|
self.intro = nn.Conv2d(img_channel, width, kernel_size=3, padding=1, stride=1, bias=True)
|
|
self.body = MySequential(
|
|
*[DropPath(
|
|
drop_path_rate,
|
|
NAFBlockSR(width, fusion=(fusion_from <= i and i <= fusion_to), drop_out_rate=drop_out_rate)
|
|
) for i in range(num_blks)]
|
|
)
|
|
self.up = nn.Sequential(
|
|
nn.Conv2d(width, img_channel * up_scale ** 2, kernel_size=3, padding=1, stride=1, bias=True),
|
|
nn.PixelShuffle(up_scale)
|
|
)
|
|
self.up_scale = up_scale
|
|
|
|
def forward(self, inp):
|
|
inp_hr = F.interpolate(inp, scale_factor=self.up_scale, mode='bilinear', align_corners=False)
|
|
if self.dual:
|
|
inp = inp.chunk(2, dim=1)
|
|
else:
|
|
inp = (inp,)
|
|
feats = [self.intro(x) for x in inp]
|
|
feats = self.body(*feats)
|
|
out = torch.cat([self.up(x) for x in feats], dim=1)
|
|
out = out + inp_hr
|
|
return out
|