Eric Rollei
d6dd6470cf
Features: - 40+ image processing algorithms - Traditional denoising: Wavelet, NLM, BM3D (GPU-accelerated) - AI models: DnCNN, SwinIR, SCUNet, NAFNet, Noise-DA - DeepInv diffusion denoising integration - Advanced film grain detection and processing - Frequency domain processing - Richardson-Lucy deconvolution - Wiener filter - External DeepInv service with FastAPI - Comprehensive documentation and guides
225 lines
7.4 KiB
Python
225 lines
7.4 KiB
Python
"""
|
|
Restormer network architecture.
|
|
|
|
Adapted from the official Restormer repository:
|
|
https://github.com/swz30/Restormer
|
|
|
|
Copyright (c) 2021 Shuhang Gu et al.
|
|
Released under the Apache License, Version 2.0.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class LayerNorm2d(nn.Module):
|
|
"""Channel-first LayerNorm that supports mixed precision."""
|
|
|
|
def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(1, num_channels, 1, 1))
|
|
self.bias = nn.Parameter(torch.zeros(1, num_channels, 1, 1))
|
|
self.eps = eps
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: # pragma: no cover - relies on torch ops
|
|
mean = x.mean((1, 2, 3), keepdim=True)
|
|
std = x.std((1, 2, 3), keepdim=True)
|
|
return (x - mean) / (std + self.eps) * self.weight + self.bias
|
|
|
|
|
|
class FeedForward(nn.Module):
|
|
"""Gated feed-forward network (GFN) used within transformer blocks."""
|
|
|
|
def __init__(self, dim: int, expansion_factor: float, bias: bool) -> None:
|
|
super().__init__()
|
|
hidden_dim = int(dim * expansion_factor)
|
|
self.project_in = nn.Conv2d(dim, hidden_dim * 2, kernel_size=1, bias=bias)
|
|
self.dwconv = nn.Conv2d(hidden_dim * 2, hidden_dim * 2, kernel_size=3, stride=1, padding=1, groups=hidden_dim * 2, bias=bias)
|
|
self.project_out = nn.Conv2d(hidden_dim, dim, kernel_size=1, bias=bias)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: # pragma: no cover
|
|
x = self.project_in(x)
|
|
x = self.dwconv(x)
|
|
x1, x2 = x.chunk(2, dim=1)
|
|
x = F.gelu(x1) * x2
|
|
x = self.project_out(x)
|
|
return x
|
|
|
|
|
|
class Attention(nn.Module):
|
|
"""Multi-Dconv head transposed attention (MDTA)."""
|
|
|
|
def __init__(self, dim: int, num_heads: int, bias: bool) -> None:
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
head_dim = dim // num_heads
|
|
self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
|
|
|
|
self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias)
|
|
self.qkv_dwconv = nn.Conv2d(dim * 3, dim * 3, kernel_size=3, stride=1, padding=1, groups=dim * 3, bias=bias)
|
|
self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: # pragma: no cover
|
|
b, c, h, w = x.shape
|
|
qkv = self.qkv_dwconv(self.qkv(x))
|
|
q, k, v = qkv.chunk(3, dim=1)
|
|
|
|
q = q.view(b, self.num_heads, c // self.num_heads, h * w)
|
|
k = k.view(b, self.num_heads, c // self.num_heads, h * w)
|
|
v = v.view(b, self.num_heads, c // self.num_heads, h * w)
|
|
|
|
q = F.normalize(q, dim=-1)
|
|
k = F.normalize(k, dim=-1)
|
|
attn = (q @ k.transpose(-2, -1)) * self.temperature
|
|
attn = attn.softmax(dim=-1)
|
|
|
|
out = attn @ v
|
|
out = out.view(b, c, h, w)
|
|
out = self.project_out(out)
|
|
return out
|
|
|
|
|
|
class TransformerBlock(nn.Module):
|
|
"""Core Restormer block with attention + feed-forward."""
|
|
|
|
def __init__(self, dim: int, num_heads: int, ffn_expansion_factor: float, bias: bool) -> None:
|
|
super().__init__()
|
|
self.norm1 = LayerNorm2d(dim)
|
|
self.attn = Attention(dim, num_heads, bias)
|
|
self.norm2 = LayerNorm2d(dim)
|
|
self.ffn = FeedForward(dim, ffn_expansion_factor, bias)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: # pragma: no cover
|
|
x = x + self.attn(self.norm1(x))
|
|
x = x + self.ffn(self.norm2(x))
|
|
return x
|
|
|
|
|
|
class OverlapPatchEmbed(nn.Module):
|
|
"""Stem convolution used to obtain overlapping patches."""
|
|
|
|
def __init__(self, in_chans: int, embed_dim: int, bias: bool) -> None:
|
|
super().__init__()
|
|
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=3, stride=1, padding=1, bias=bias)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: # pragma: no cover
|
|
return self.proj(x)
|
|
|
|
|
|
class Downsample(nn.Module):
|
|
def __init__(self, dim: int, bias: bool) -> None:
|
|
super().__init__()
|
|
self.body = nn.Sequential(
|
|
nn.Conv2d(dim, dim // 2, kernel_size=3, stride=1, padding=1, bias=bias),
|
|
nn.PixelUnshuffle(2),
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: # pragma: no cover
|
|
return self.body(x)
|
|
|
|
|
|
class Upsample(nn.Module):
|
|
def __init__(self, dim: int, bias: bool) -> None:
|
|
super().__init__()
|
|
self.body = nn.Sequential(
|
|
nn.Conv2d(dim, dim * 2, kernel_size=3, stride=1, padding=1, bias=bias),
|
|
nn.PixelShuffle(2),
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: # pragma: no cover
|
|
return self.body(x)
|
|
|
|
|
|
class Restormer(nn.Module):
|
|
"""Main Restormer network."""
|
|
|
|
def __init__(
|
|
self,
|
|
inp_channels: int = 3,
|
|
out_channels: int = 3,
|
|
dim: int = 48,
|
|
num_blocks: Optional[list[int]] = None,
|
|
num_refinement_blocks: int = 4,
|
|
heads: Optional[list[int]] = None,
|
|
ffn_expansion_factor: float = 2.66,
|
|
bias: bool = False,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
if num_blocks is None:
|
|
num_blocks = [4, 6, 6, 8]
|
|
if heads is None:
|
|
heads = [1, 2, 4, 8]
|
|
|
|
self.patch_embed = OverlapPatchEmbed(inp_channels, dim, bias)
|
|
|
|
# Encoder
|
|
self.encoder_levels = nn.ModuleList()
|
|
self.downsamples = nn.ModuleList()
|
|
|
|
dim_level = dim
|
|
for blocks, head in zip(num_blocks, heads):
|
|
level = nn.Sequential(
|
|
*[TransformerBlock(dim_level, head, ffn_expansion_factor, bias) for _ in range(blocks)]
|
|
)
|
|
self.encoder_levels.append(level)
|
|
self.downsamples.append(Downsample(dim_level, bias))
|
|
dim_level *= 2
|
|
|
|
# Bottleneck
|
|
self.latent = nn.Sequential(
|
|
*[TransformerBlock(dim_level, heads[-1] * 2, ffn_expansion_factor, bias) for _ in range(num_blocks[-1])]
|
|
)
|
|
|
|
# Decoder
|
|
self.upsamples = nn.ModuleList()
|
|
self.decoder_levels = nn.ModuleList()
|
|
for blocks, head in reversed(list(zip(num_blocks[:-1], heads[:-1]))):
|
|
dim_level //= 2
|
|
self.upsamples.append(Upsample(dim_level * 2, bias))
|
|
self.decoder_levels.append(
|
|
nn.Sequential(
|
|
*[TransformerBlock(dim_level, head, ffn_expansion_factor, bias) for _ in range(blocks)]
|
|
)
|
|
)
|
|
|
|
self.output = nn.Sequential(
|
|
LayerNorm2d(dim),
|
|
nn.Conv2d(dim, out_channels, kernel_size=3, stride=1, padding=1, bias=bias),
|
|
)
|
|
|
|
self.num_levels = len(num_blocks) - 1
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: # pragma: no cover
|
|
x = self.patch_embed(x)
|
|
|
|
features = []
|
|
for idx in range(self.num_levels):
|
|
x = self.encoder_levels[idx](x)
|
|
features.append(x)
|
|
x = self.downsamples[idx](x)
|
|
|
|
x = self.latent(x)
|
|
|
|
for idx in range(self.num_levels - 1, -1, -1):
|
|
x = self.upsamples[idx](x)
|
|
x = torch.cat([x, features[idx]], dim=1)
|
|
x = self.decoder_levels[self.num_levels - 1 - idx](x)
|
|
|
|
x = self.output(x)
|
|
return x
|
|
|
|
|
|
if __name__ == "__main__": # pragma: no cover
|
|
model = Restormer()
|
|
dummy = torch.randn(1, 3, 256, 256)
|
|
with torch.no_grad():
|
|
out = model(dummy)
|
|
print("Input:", dummy.shape, "Output:", out.shape)
|