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
722 lines
30 KiB
Python
722 lines
30 KiB
Python
"""
|
|
SCUNet (Swin-Conv-UNet) Implementation for Mixed Degradation Restoration
|
|
State-of-the-art for handling film grain, digital noise, JPEG compression, and blur
|
|
|
|
Based on the official SCUNet implementation:
|
|
https://github.com/cszn/SCUNet
|
|
|
|
Memory-optimized version to prevent system crashes
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from typing import Tuple, Optional
|
|
import math
|
|
from einops import rearrange
|
|
from einops.layers.torch import Rearrange, Reduce
|
|
|
|
# Import memory utilities
|
|
try:
|
|
from .memory_utils import MemoryManager, SafeModelLoader, safe_torch_operation
|
|
except ImportError:
|
|
# Fallback if memory_utils not available
|
|
class MemoryManager:
|
|
@staticmethod
|
|
def safe_device_selection(device='auto', min_gpu_memory_gb=4.0):
|
|
if device == 'cpu' or not torch.cuda.is_available():
|
|
return 'cpu'
|
|
return 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
|
|
@staticmethod
|
|
def cleanup_memory():
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
@staticmethod
|
|
def optimize_tile_size(image_shape, max_memory_gb=2.0, min_tile_size=64, max_tile_size=1024):
|
|
return min_tile_size
|
|
|
|
class SafeModelLoader:
|
|
def __init__(self, device='auto', cleanup_on_exit=True):
|
|
self.device = device
|
|
def __enter__(self):
|
|
return self.device
|
|
def __exit__(self, *args):
|
|
pass
|
|
|
|
def safe_torch_operation(func):
|
|
return func
|
|
|
|
def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
|
|
"""Truncated normal initialization (simplified version)"""
|
|
with torch.no_grad():
|
|
tensor.normal_(mean, std)
|
|
tensor.clamp_(min=a, max=b)
|
|
return tensor
|
|
|
|
class DropPath(nn.Module):
|
|
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
|
|
def __init__(self, drop_prob=None):
|
|
super(DropPath, self).__init__()
|
|
self.drop_prob = drop_prob
|
|
|
|
def forward(self, x):
|
|
if self.drop_prob == 0. or not self.training:
|
|
return x
|
|
keep_prob = 1 - self.drop_prob
|
|
shape = (x.shape[0],) + (1,) * (x.ndim - 1)
|
|
random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
|
|
random_tensor.floor_()
|
|
output = x.div(keep_prob) * random_tensor
|
|
return output
|
|
|
|
class WMSA(nn.Module):
|
|
"""Self-attention module in Swin Transformer"""
|
|
def __init__(self, input_dim, output_dim, head_dim, window_size, type):
|
|
super(WMSA, self).__init__()
|
|
self.input_dim = input_dim
|
|
self.output_dim = output_dim
|
|
self.head_dim = head_dim
|
|
self.scale = self.head_dim ** -0.5
|
|
self.n_heads = input_dim // head_dim
|
|
self.window_size = window_size
|
|
self.type = type
|
|
|
|
self.embedding_layer = nn.Linear(self.input_dim, 3 * self.input_dim, bias=True)
|
|
self.relative_position_params = nn.Parameter(torch.zeros((2 * window_size - 1)*(2 * window_size - 1), self.n_heads))
|
|
self.linear = nn.Linear(self.input_dim, self.output_dim)
|
|
|
|
trunc_normal_(self.relative_position_params, std=.02)
|
|
self.relative_position_params = torch.nn.Parameter(self.relative_position_params.view(2*window_size-1, 2*window_size-1, self.n_heads).transpose(1,2).transpose(0,1))
|
|
|
|
def generate_mask(self, h, w, p, shift):
|
|
"""Generate mask for SW-MSA"""
|
|
attn_mask = torch.zeros(h, w, p, p, p, p, dtype=torch.bool, device=self.relative_position_params.device)
|
|
if self.type == 'W':
|
|
return attn_mask
|
|
|
|
s = p - shift
|
|
attn_mask[-1, :, :s, :, s:, :] = True
|
|
attn_mask[-1, :, s:, :, :s, :] = True
|
|
attn_mask[:, -1, :, :s, :, s:] = True
|
|
attn_mask[:, -1, :, s:, :, :s] = True
|
|
attn_mask = rearrange(attn_mask, 'w1 w2 p1 p2 p3 p4 -> 1 1 (w1 w2) (p1 p2) (p3 p4)')
|
|
return attn_mask
|
|
|
|
def forward(self, x):
|
|
if self.type != 'W':
|
|
x = torch.roll(x, shifts=(-(self.window_size//2), -(self.window_size//2)), dims=(1,2))
|
|
|
|
x = rearrange(x, 'b (w1 p1) (w2 p2) c -> b w1 w2 p1 p2 c', p1=self.window_size, p2=self.window_size)
|
|
h_windows = x.size(1)
|
|
w_windows = x.size(2)
|
|
|
|
x = rearrange(x, 'b w1 w2 p1 p2 c -> b (w1 w2) (p1 p2) c', p1=self.window_size, p2=self.window_size)
|
|
qkv = self.embedding_layer(x)
|
|
q, k, v = rearrange(qkv, 'b nw np (threeh c) -> threeh b nw np c', c=self.head_dim).chunk(3, dim=0)
|
|
|
|
sim = torch.einsum('hbwpc,hbwqc->hbwpq', q, k) * self.scale
|
|
sim = sim + rearrange(self.relative_embedding(), 'h p q -> h 1 1 p q')
|
|
|
|
if self.type != 'W':
|
|
attn_mask = self.generate_mask(h_windows, w_windows, self.window_size, shift=self.window_size//2)
|
|
sim = sim.masked_fill_(attn_mask, float("-inf"))
|
|
|
|
probs = nn.functional.softmax(sim, dim=-1)
|
|
output = torch.einsum('hbwij,hbwjc->hbwic', probs, v)
|
|
output = rearrange(output, 'h b w p c -> b w p (h c)')
|
|
output = self.linear(output)
|
|
output = rearrange(output, 'b (w1 w2) (p1 p2) c -> b (w1 p1) (w2 p2) c', w1=h_windows, p1=self.window_size)
|
|
|
|
if self.type != 'W':
|
|
output = torch.roll(output, shifts=(self.window_size//2, self.window_size//2), dims=(1,2))
|
|
|
|
return output
|
|
|
|
def relative_embedding(self):
|
|
cord = torch.tensor(np.array([[i, j] for i in range(self.window_size) for j in range(self.window_size)]))
|
|
relation = cord[:, None, :] - cord[None, :, :] + self.window_size - 1
|
|
return self.relative_position_params[:, relation[:,:,0].long(), relation[:,:,1].long()]
|
|
|
|
class Block(nn.Module):
|
|
"""SwinTransformer Block"""
|
|
def __init__(self, input_dim, output_dim, head_dim, window_size, drop_path, type='W', input_resolution=None):
|
|
super(Block, self).__init__()
|
|
self.input_dim = input_dim
|
|
self.output_dim = output_dim
|
|
assert type in ['W', 'SW']
|
|
self.type = type
|
|
|
|
if input_resolution is not None and input_resolution <= window_size:
|
|
self.type = 'W'
|
|
|
|
self.ln1 = nn.LayerNorm(input_dim)
|
|
self.msa = WMSA(input_dim, input_dim, head_dim, window_size, self.type)
|
|
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
|
|
self.ln2 = nn.LayerNorm(input_dim)
|
|
self.mlp = nn.Sequential(
|
|
nn.Linear(input_dim, 4 * input_dim),
|
|
nn.GELU(),
|
|
nn.Linear(4 * input_dim, output_dim),
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = x + self.drop_path(self.msa(self.ln1(x)))
|
|
x = x + self.drop_path(self.mlp(self.ln2(x)))
|
|
return x
|
|
|
|
class ConvTransBlock(nn.Module):
|
|
"""SwinTransformer and Conv Block"""
|
|
def __init__(self, conv_dim, trans_dim, head_dim, window_size, drop_path, type='W', input_resolution=None):
|
|
super(ConvTransBlock, self).__init__()
|
|
self.conv_dim = conv_dim
|
|
self.trans_dim = trans_dim
|
|
self.head_dim = head_dim
|
|
self.window_size = window_size
|
|
self.drop_path = drop_path
|
|
self.type = type
|
|
self.input_resolution = input_resolution
|
|
|
|
assert self.type in ['W', 'SW']
|
|
if self.input_resolution is not None and self.input_resolution <= self.window_size:
|
|
self.type = 'W'
|
|
|
|
self.trans_block = Block(self.trans_dim, self.trans_dim, self.head_dim,
|
|
self.window_size, self.drop_path, self.type, self.input_resolution)
|
|
self.conv1_1 = nn.Conv2d(self.conv_dim + self.trans_dim, self.conv_dim + self.trans_dim, 1, 1, 0, bias=True)
|
|
self.conv1_2 = nn.Conv2d(self.conv_dim + self.trans_dim, self.conv_dim + self.trans_dim, 1, 1, 0, bias=True)
|
|
|
|
self.conv_block = nn.Sequential(
|
|
nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False),
|
|
nn.ReLU(True),
|
|
nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False)
|
|
)
|
|
|
|
def forward(self, x):
|
|
conv_x, trans_x = torch.split(self.conv1_1(x), (self.conv_dim, self.trans_dim), dim=1)
|
|
conv_x = self.conv_block(conv_x) + conv_x
|
|
trans_x = Rearrange('b c h w -> b h w c')(trans_x)
|
|
trans_x = self.trans_block(trans_x)
|
|
trans_x = Rearrange('b h w c -> b c h w')(trans_x)
|
|
res = self.conv1_2(torch.cat((conv_x, trans_x), dim=1))
|
|
x = x + res
|
|
return x
|
|
|
|
class SCUNet(nn.Module):
|
|
"""SCUNet model with official architecture"""
|
|
def __init__(self, in_nc=3, config=[2,2,2,2,2,2,2], dim=64, drop_path_rate=0.0, input_resolution=256):
|
|
super(SCUNet, self).__init__()
|
|
self.config = config
|
|
self.dim = dim
|
|
self.head_dim = 32
|
|
self.window_size = 8
|
|
|
|
# Drop path rate for each layer
|
|
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(config))]
|
|
|
|
self.m_head = nn.Sequential(nn.Conv2d(in_nc, dim, 3, 1, 1, bias=False))
|
|
|
|
begin = 0
|
|
self.m_down1 = nn.Sequential(
|
|
*[ConvTransBlock(dim//2, dim//2, self.head_dim, self.window_size, dpr[i+begin],
|
|
'W' if not i%2 else 'SW', input_resolution) for i in range(config[0])],
|
|
nn.Conv2d(dim, 2*dim, 2, 2, 0, bias=False)
|
|
)
|
|
|
|
begin += config[0]
|
|
self.m_down2 = nn.Sequential(
|
|
*[ConvTransBlock(dim, dim, self.head_dim, self.window_size, dpr[i+begin],
|
|
'W' if not i%2 else 'SW', input_resolution//2) for i in range(config[1])],
|
|
nn.Conv2d(2*dim, 4*dim, 2, 2, 0, bias=False)
|
|
)
|
|
|
|
begin += config[1]
|
|
self.m_down3 = nn.Sequential(
|
|
*[ConvTransBlock(2*dim, 2*dim, self.head_dim, self.window_size, dpr[i+begin],
|
|
'W' if not i%2 else 'SW', input_resolution//4) for i in range(config[2])],
|
|
nn.Conv2d(4*dim, 8*dim, 2, 2, 0, bias=False)
|
|
)
|
|
|
|
begin += config[2]
|
|
self.m_body = nn.Sequential(
|
|
*[ConvTransBlock(4*dim, 4*dim, self.head_dim, self.window_size, dpr[i+begin],
|
|
'W' if not i%2 else 'SW', input_resolution//8) for i in range(config[3])]
|
|
)
|
|
|
|
begin += config[3]
|
|
self.m_up3 = nn.Sequential(
|
|
nn.ConvTranspose2d(8*dim, 4*dim, 2, 2, 0, bias=False),
|
|
*[ConvTransBlock(2*dim, 2*dim, self.head_dim, self.window_size, dpr[i+begin],
|
|
'W' if not i%2 else 'SW', input_resolution//4) for i in range(config[4])]
|
|
)
|
|
|
|
begin += config[4]
|
|
self.m_up2 = nn.Sequential(
|
|
nn.ConvTranspose2d(4*dim, 2*dim, 2, 2, 0, bias=False),
|
|
*[ConvTransBlock(dim, dim, self.head_dim, self.window_size, dpr[i+begin],
|
|
'W' if not i%2 else 'SW', input_resolution//2) for i in range(config[5])]
|
|
)
|
|
|
|
begin += config[5]
|
|
self.m_up1 = nn.Sequential(
|
|
nn.ConvTranspose2d(2*dim, dim, 2, 2, 0, bias=False),
|
|
*[ConvTransBlock(dim//2, dim//2, self.head_dim, self.window_size, dpr[i+begin],
|
|
'W' if not i%2 else 'SW', input_resolution) for i in range(config[6])]
|
|
)
|
|
|
|
self.m_tail = nn.Sequential(nn.Conv2d(dim, in_nc, 3, 1, 1, bias=False))
|
|
|
|
def forward(self, x0):
|
|
h, w = x0.size()[-2:]
|
|
paddingBottom = int(np.ceil(h/64)*64-h)
|
|
paddingRight = int(np.ceil(w/64)*64-w)
|
|
x0 = nn.ReplicationPad2d((0, paddingRight, 0, paddingBottom))(x0)
|
|
|
|
x1 = self.m_head(x0)
|
|
x2 = self.m_down1(x1)
|
|
x3 = self.m_down2(x2)
|
|
x4 = self.m_down3(x3)
|
|
x = self.m_body(x4)
|
|
x = self.m_up3(x + x4)
|
|
x = self.m_up2(x + x3)
|
|
x = self.m_up1(x + x2)
|
|
x = self.m_tail(x + x1)
|
|
|
|
x = x[..., :h, :w]
|
|
return x
|
|
|
|
class WindowAttention(nn.Module):
|
|
"""Window based multi-head self attention (W-MSA) module"""
|
|
def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0.):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.window_size = window_size
|
|
self.num_heads = num_heads
|
|
head_dim = dim // num_heads
|
|
self.scale = head_dim ** -0.5
|
|
|
|
# Define relative position bias table
|
|
self.relative_position_bias_table = nn.Parameter(
|
|
torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))
|
|
|
|
# Get pair-wise relative position index
|
|
coords_h = torch.arange(self.window_size[0])
|
|
coords_w = torch.arange(self.window_size[1])
|
|
coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij'))
|
|
coords_flatten = torch.flatten(coords, 1)
|
|
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
|
|
relative_coords = relative_coords.permute(1, 2, 0).contiguous()
|
|
relative_coords[:, :, 0] += self.window_size[0] - 1
|
|
relative_coords[:, :, 1] += self.window_size[1] - 1
|
|
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
|
|
relative_position_index = relative_coords.sum(-1)
|
|
self.register_buffer("relative_position_index", relative_position_index)
|
|
|
|
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
|
|
self.attn_drop = nn.Dropout(attn_drop)
|
|
self.proj = nn.Linear(dim, dim)
|
|
self.proj_drop = nn.Dropout(proj_drop)
|
|
|
|
def forward(self, x, mask=None):
|
|
B_, N, C = x.shape
|
|
qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
|
q, k, v = qkv[0], qkv[1], qkv[2]
|
|
|
|
q = q * self.scale
|
|
attn = (q @ k.transpose(-2, -1))
|
|
|
|
relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
|
|
self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)
|
|
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
|
|
attn = attn + relative_position_bias.unsqueeze(0)
|
|
|
|
if mask is not None:
|
|
nW = mask.shape[0]
|
|
attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
|
|
attn = attn.view(-1, self.num_heads, N, N)
|
|
attn = F.softmax(attn, dim=-1)
|
|
else:
|
|
attn = F.softmax(attn, dim=-1)
|
|
|
|
attn = self.attn_drop(attn)
|
|
|
|
x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
|
|
x = self.proj(x)
|
|
x = self.proj_drop(x)
|
|
return x
|
|
|
|
class SwinTransformerBlock(nn.Module):
|
|
"""Swin Transformer Block"""
|
|
def __init__(self, dim, num_heads, window_size=7, shift_size=0, mlp_ratio=4.):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.num_heads = num_heads
|
|
self.window_size = window_size
|
|
self.shift_size = shift_size
|
|
self.mlp_ratio = mlp_ratio
|
|
|
|
self.norm1 = nn.LayerNorm(dim)
|
|
self.attn = WindowAttention(dim, window_size=(window_size, window_size), num_heads=num_heads)
|
|
|
|
self.norm2 = nn.LayerNorm(dim)
|
|
mlp_hidden_dim = int(dim * mlp_ratio)
|
|
self.mlp = nn.Sequential(
|
|
nn.Linear(dim, mlp_hidden_dim),
|
|
nn.GELU(),
|
|
nn.Linear(mlp_hidden_dim, dim)
|
|
)
|
|
|
|
def forward(self, x, H, W):
|
|
B, L, C = x.shape
|
|
|
|
shortcut = x
|
|
x = self.norm1(x)
|
|
x = x.view(B, H, W, C)
|
|
|
|
# Pad feature maps to multiples of window size
|
|
pad_l = pad_t = 0
|
|
pad_r = (self.window_size - W % self.window_size) % self.window_size
|
|
pad_b = (self.window_size - H % self.window_size) % self.window_size
|
|
x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
|
|
_, Hp, Wp, _ = x.shape
|
|
|
|
# Window partition
|
|
x_windows = window_partition(x, self.window_size)
|
|
x_windows = x_windows.view(-1, self.window_size * self.window_size, C)
|
|
|
|
# W-MSA/SW-MSA
|
|
attn_windows = self.attn(x_windows)
|
|
|
|
# Merge windows
|
|
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
|
|
x = window_reverse(attn_windows, self.window_size, Hp, Wp)
|
|
|
|
# Remove padding
|
|
if pad_r > 0 or pad_b > 0:
|
|
x = x[:, :H, :W, :].contiguous()
|
|
|
|
x = x.view(B, H * W, C)
|
|
|
|
# FFN
|
|
x = shortcut + x
|
|
x = x + self.mlp(self.norm2(x))
|
|
|
|
return x
|
|
|
|
def window_partition(x, window_size):
|
|
"""Partition into non-overlapping windows"""
|
|
B, H, W, C = x.shape
|
|
x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
|
|
windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
|
|
return windows
|
|
|
|
def window_reverse(windows, window_size, H, W):
|
|
"""Reverse window partition"""
|
|
B = int(windows.shape[0] / (H * W / window_size / window_size))
|
|
x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
|
|
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
|
|
return x
|
|
|
|
class SCUNetProcessor:
|
|
"""Memory-safe processor for SCUNet mixed degradation restoration"""
|
|
|
|
def __init__(self, model_name: str = "scunet_color_real_psnr", device: str = 'auto',
|
|
lightweight: bool = True):
|
|
# Safe device selection with memory checking
|
|
self.device = MemoryManager.safe_device_selection(device, min_gpu_memory_gb=3.0)
|
|
self.model = None
|
|
self.model_name = model_name
|
|
|
|
# Determine model path based on model name
|
|
models_dir = Path(__file__).parent.parent / "models"
|
|
model_path = models_dir / f"{model_name}.pth"
|
|
|
|
if not model_path.exists():
|
|
print(f"Warning: Model {model_name}.pth not found in {models_dir}")
|
|
print("Available models:")
|
|
for model_file in models_dir.glob("scunet_*.pth"):
|
|
print(f" - {model_file.name}")
|
|
# Try to find any SCUNet model as fallback
|
|
fallback_models = list(models_dir.glob("scunet_*.pth"))
|
|
if fallback_models:
|
|
model_path = fallback_models[0]
|
|
print(f"Using fallback model: {model_path.name}")
|
|
else:
|
|
print("No SCUNet models found. Download models using: python download_models.py")
|
|
|
|
self.model_path = str(model_path) if model_path.exists() else None
|
|
self.lightweight = lightweight
|
|
|
|
# Determine input channels based on model name
|
|
if "gray" in model_name.lower():
|
|
self.in_channels = 1
|
|
else:
|
|
self.in_channels = 3
|
|
|
|
print(f"SCUNet initializing: {model_name} on {self.device} (channels: {self.in_channels})")
|
|
self._load_model()
|
|
|
|
def _load_model(self):
|
|
"""Load SCUNet model with memory-safe parameters"""
|
|
with SafeModelLoader(self.device) as device:
|
|
try:
|
|
# Use official SCUNet architecture with full configuration
|
|
# to match the pretrained model
|
|
config = [4, 4, 4, 4, 4, 4, 4] # Official SCUNet config
|
|
dim = 64
|
|
print(f"Using official SCUNet configuration for {self.model_name}")
|
|
|
|
# Initialize model with correct input channels
|
|
self.model = SCUNet(
|
|
in_nc=self.in_channels,
|
|
config=config,
|
|
dim=dim,
|
|
drop_path_rate=0.0,
|
|
input_resolution=256
|
|
)
|
|
|
|
# Load pretrained weights if available
|
|
if self.model_path and Path(self.model_path).exists():
|
|
checkpoint = torch.load(self.model_path, map_location='cpu')
|
|
|
|
# Handle different checkpoint formats
|
|
if 'params' in checkpoint:
|
|
state_dict = checkpoint['params']
|
|
elif 'params_ema' in checkpoint:
|
|
state_dict = checkpoint['params_ema']
|
|
elif 'model' in checkpoint:
|
|
state_dict = checkpoint['model']
|
|
else:
|
|
state_dict = checkpoint
|
|
|
|
# Remove any module prefix if present
|
|
new_state_dict = {}
|
|
for key, value in state_dict.items():
|
|
if key.startswith('module.'):
|
|
new_key = key[7:] # Remove 'module.' prefix
|
|
else:
|
|
new_key = key
|
|
new_state_dict[new_key] = value
|
|
|
|
# Load with strict=False to handle architecture differences
|
|
missing_keys, unexpected_keys = self.model.load_state_dict(new_state_dict, strict=False)
|
|
|
|
if missing_keys:
|
|
print(f"Missing keys: {len(missing_keys)}")
|
|
if unexpected_keys:
|
|
print(f"Unexpected keys: {len(unexpected_keys)}")
|
|
|
|
print(f"Loaded SCUNet model: {self.model_name}")
|
|
else:
|
|
print(f"No pretrained model found for {self.model_name}. Using randomly initialized weights.")
|
|
|
|
self.model.to(device)
|
|
self.model.eval()
|
|
|
|
# Clear memory after loading
|
|
MemoryManager.cleanup_memory()
|
|
|
|
except Exception as e:
|
|
print(f"Error loading SCUNet model: {e}")
|
|
print("Model loading failed - check architecture compatibility")
|
|
raise e
|
|
|
|
@safe_torch_operation
|
|
def process_image(self, image: np.ndarray, tile_size: int = None, overlap: int = 16) -> np.ndarray:
|
|
"""
|
|
Process image with SCUNet (memory-safe version)
|
|
|
|
Args:
|
|
image: Input image as numpy array (H, W, C)
|
|
tile_size: Size of processing tiles (auto-calculated if None)
|
|
overlap: Overlap between tiles
|
|
"""
|
|
if self.model is None:
|
|
raise ValueError("Model not loaded")
|
|
|
|
try:
|
|
if image.ndim == 4:
|
|
if image.shape[0] != 1:
|
|
raise ValueError("SCUNet processor expects a single image (batch size 1)")
|
|
image = image[0]
|
|
|
|
# Auto-calculate tile size based on available memory and image size
|
|
if tile_size is None:
|
|
tile_size = MemoryManager.optimize_tile_size(
|
|
image.shape,
|
|
max_memory_gb=2.0 if self.device == 'cuda' else 1.0,
|
|
min_tile_size=64,
|
|
max_tile_size=512 if self.device == 'cuda' else 256
|
|
)
|
|
print(f"Auto-selected tile size: {tile_size}")
|
|
|
|
# Convert to tensor
|
|
input_tensor = self._preprocess(image)
|
|
|
|
# Use tiling for any reasonably sized image to be safe
|
|
max_direct_size = 256 if self.device == 'cuda' else 128
|
|
if input_tensor.shape[2] > max_direct_size or input_tensor.shape[3] > max_direct_size:
|
|
output_tensor = self._process_with_tiles(input_tensor, tile_size, overlap)
|
|
else:
|
|
with torch.no_grad():
|
|
output_tensor = self.model(input_tensor)
|
|
|
|
# Convert back to numpy
|
|
result = self._postprocess(output_tensor)
|
|
|
|
# Clean up memory
|
|
del input_tensor, output_tensor
|
|
MemoryManager.cleanup_memory()
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
# Clean up on error
|
|
MemoryManager.cleanup_memory()
|
|
print(f"SCUNet processing error: {e}")
|
|
return image # Return original image on error
|
|
|
|
def _preprocess(self, image: np.ndarray) -> torch.Tensor:
|
|
"""Preprocess image for SCUNet"""
|
|
# Convert to float and normalize
|
|
if image.dtype == np.uint8:
|
|
image = image.astype(np.float32) / 255.0
|
|
elif image.dtype == np.uint16:
|
|
image = image.astype(np.float32) / 65535.0
|
|
|
|
# Handle channel dimension
|
|
if len(image.shape) == 2:
|
|
# Grayscale image without channel dimension
|
|
image = np.expand_dims(image, axis=2)
|
|
|
|
# Ensure correct number of channels
|
|
if self.in_channels == 1 and image.shape[2] == 3:
|
|
# Convert RGB to grayscale
|
|
image = np.mean(image, axis=2, keepdims=True)
|
|
elif self.in_channels == 3 and image.shape[2] == 1:
|
|
# Convert grayscale to RGB
|
|
image = np.repeat(image, 3, axis=2)
|
|
|
|
# Convert to tensor format (C, H, W)
|
|
# FIXED: Ensure image has correct shape before transpose
|
|
if len(image.shape) != 3:
|
|
raise ValueError(f"Image must have 3 dimensions (H,W,C), got shape {image.shape}")
|
|
|
|
# FIXED: Ensure contiguous array before transpose
|
|
image = np.ascontiguousarray(image)
|
|
image_tensor = torch.from_numpy(image.transpose(2, 0, 1)).float()
|
|
|
|
# Add batch dimension
|
|
image_tensor = image_tensor.unsqueeze(0)
|
|
|
|
return image_tensor.to(self.device)
|
|
|
|
def _postprocess(self, tensor: torch.Tensor) -> np.ndarray:
|
|
"""Convert model output back to image"""
|
|
# Remove batch dimension and convert to numpy
|
|
image = tensor.squeeze(0).cpu().detach().numpy()
|
|
|
|
# FIXED: Ensure contiguous array before transpose
|
|
image = np.ascontiguousarray(image)
|
|
|
|
# Convert from (C, H, W) to (H, W, C)
|
|
image = image.transpose(1, 2, 0)
|
|
|
|
# Handle single channel output
|
|
if image.shape[2] == 1:
|
|
# For grayscale output, keep as single channel
|
|
pass
|
|
|
|
# Clip values to valid range [0, 1]
|
|
image = np.clip(image, 0.0, 1.0)
|
|
|
|
return image
|
|
|
|
def _process_with_tiles(self, input_tensor: torch.Tensor, tile_size: int, overlap: int) -> torch.Tensor:
|
|
"""Process large image with overlapping tiles (memory-optimized)"""
|
|
B, C, H, W = input_tensor.shape
|
|
|
|
# Calculate tile positions
|
|
stride = tile_size - overlap
|
|
tiles_h = (H - overlap) // stride + (1 if (H - overlap) % stride > 0 else 0)
|
|
tiles_w = (W - overlap) // stride + (1 if (W - overlap) % stride > 0 else 0)
|
|
|
|
# Use smaller chunks to avoid memory issues
|
|
output_tensor = torch.zeros_like(input_tensor)
|
|
weight_tensor = torch.zeros((B, 1, H, W), device=input_tensor.device)
|
|
|
|
# Process tiles in batches to reduce memory usage
|
|
batch_size = 4 if self.device == 'cuda' else 8
|
|
|
|
for i in range(0, tiles_h, batch_size):
|
|
for j in range(0, tiles_w, batch_size):
|
|
# Process batch of tiles
|
|
batch_end_i = min(i + batch_size, tiles_h)
|
|
batch_end_j = min(j + batch_size, tiles_w)
|
|
|
|
for ii in range(i, batch_end_i):
|
|
for jj in range(j, batch_end_j):
|
|
# Calculate tile boundaries
|
|
start_h = ii * stride
|
|
end_h = min(start_h + tile_size, H)
|
|
start_w = jj * stride
|
|
end_w = min(start_w + tile_size, W)
|
|
|
|
# Extract tile
|
|
tile = input_tensor[:, :, start_h:end_h, start_w:end_w]
|
|
|
|
# Pad tile if necessary
|
|
pad_h = tile_size - (end_h - start_h)
|
|
pad_w = tile_size - (end_w - start_w)
|
|
if pad_h > 0 or pad_w > 0:
|
|
tile = F.pad(tile, (0, pad_w, 0, pad_h), mode='reflect')
|
|
|
|
# Process tile
|
|
with torch.no_grad():
|
|
processed_tile = self.model(tile)
|
|
|
|
# Remove padding
|
|
if pad_h > 0 or pad_w > 0:
|
|
processed_tile = processed_tile[:, :, :end_h-start_h, :end_w-start_w]
|
|
|
|
# Add to output with weighting
|
|
output_tensor[:, :, start_h:end_h, start_w:end_w] += processed_tile
|
|
weight_tensor[:, :, start_h:end_h, start_w:end_w] += 1
|
|
|
|
# Clean up tile memory
|
|
del tile, processed_tile
|
|
|
|
# Clean GPU cache after each batch
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
# Normalize by weights
|
|
output_tensor = torch.where(weight_tensor > 0, output_tensor / weight_tensor, output_tensor)
|
|
|
|
return output_tensor
|
|
|
|
# Test function
|
|
def test_scunet_processor():
|
|
"""Test SCUNet processor with memory safety"""
|
|
# Create small test image to avoid memory issues
|
|
test_image = np.random.rand(128, 128, 3).astype(np.float32)
|
|
test_image = (test_image * 255).astype(np.uint8)
|
|
|
|
# Test processor with safe settings
|
|
processor = SCUNetProcessor(device='cpu', lightweight=True)
|
|
result = processor.process_image(test_image, tile_size=64)
|
|
|
|
print(f"Input shape: {test_image.shape}")
|
|
print(f"Output shape: {result.shape}")
|
|
print(f"Input range: [{test_image.min()}, {test_image.max()}]")
|
|
print(f"Output range: [{result.min()}, {result.max()}]")
|
|
print("SCUNet test completed successfully!")
|
|
|
|
return result
|
|
|
|
if __name__ == "__main__":
|
|
test_scunet_processor()
|