mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
feat(noise-removal): SOTA noise removal with 4 quality tiers (#57)
* feat(noise-removal): register tool in shared constants and i18n * feat(noise-removal): add SCUNet and NAFNet model architectures * feat(noise-removal): add Python denoising engine with 4 quality tiers * feat(noise-removal): add TypeScript bridge for Python sidecar * feat(noise-removal): add frontend settings with 4-tier selector * feat(noise-removal): register in tool registry and pipeline * feat(noise-removal): add Fastify API route with Zod validation * feat(noise-removal): add SCUNet and NAFNet model downloads to Docker build * test(noise-removal): add to e2e tool page rendering tests * test(noise-removal): add integration tests for API endpoint * style: fix biome formatting and import ordering * fix(noise-removal): use correct model download URLs NAFNet model is hosted on HuggingFace, not GitHub releases. Also align SCUNet URL to use the KAIR releases (same as Docker build). * fix(noise-removal): remove emojis from tier selector, simplify labels Drop emoji icons from Quick/Balanced/Quality/Maximum buttons. Replace technical algorithm names with plain descriptions users can understand. --------- Co-authored-by: stirling-image <stirling-image@users.noreply.github.com>
This commit is contained in:
stirling-image
co-authored by
stirling-image
stirling-image
parent
61794dca2d
commit
dfffc0a8cc
@@ -0,0 +1,213 @@
|
||||
# NAFNet: Nonlinear Activation Free Network for image restoration
|
||||
# Adapted from https://github.com/megvii-research/NAFNet (MIT License)
|
||||
# Original paper: "Simple Baselines for Image Restoration"
|
||||
# Authors: Liangyu Chen, Xiaojie Chu, Xiangyu Zhang, Jian Sun
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class LayerNormFunction(torch.autograd.Function):
|
||||
"""Custom autograd function for efficient 2D layer normalization."""
|
||||
|
||||
@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.0 / 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):
|
||||
"""Channel-wise layer normalization for 2D feature maps."""
|
||||
|
||||
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):
|
||||
"""Split channels in half and multiply - a simple gating mechanism."""
|
||||
|
||||
def forward(self, x):
|
||||
x1, x2 = x.chunk(2, dim=1)
|
||||
return x1 * x2
|
||||
|
||||
|
||||
class NAFBlock(nn.Module):
|
||||
"""NAFNet building block with simplified channel attention and SimpleGate."""
|
||||
|
||||
def __init__(self, c, DW_Expand=2, FFN_Expand=2, drop_out_rate=0.0):
|
||||
super().__init__()
|
||||
dw_channel = c * DW_Expand
|
||||
self.conv1 = nn.Conv2d(
|
||||
in_channels=c, out_channels=dw_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True,
|
||||
)
|
||||
self.conv2 = nn.Conv2d(
|
||||
in_channels=dw_channel, out_channels=dw_channel, kernel_size=3, padding=1, stride=1,
|
||||
groups=dw_channel, bias=True,
|
||||
)
|
||||
self.conv3 = nn.Conv2d(
|
||||
in_channels=dw_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True,
|
||||
)
|
||||
|
||||
# Simplified Channel Attention
|
||||
self.sca = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d(1),
|
||||
nn.Conv2d(
|
||||
in_channels=dw_channel // 2, out_channels=dw_channel // 2,
|
||||
kernel_size=1, padding=0, stride=1, groups=1, bias=True,
|
||||
),
|
||||
)
|
||||
|
||||
self.sg = SimpleGate()
|
||||
|
||||
ffn_channel = FFN_Expand * c
|
||||
self.conv4 = nn.Conv2d(
|
||||
in_channels=c, out_channels=ffn_channel, kernel_size=1, padding=0, stride=1, groups=1, bias=True,
|
||||
)
|
||||
self.conv5 = nn.Conv2d(
|
||||
in_channels=ffn_channel // 2, out_channels=c, kernel_size=1, padding=0, stride=1, groups=1, bias=True,
|
||||
)
|
||||
|
||||
self.norm1 = LayerNorm2d(c)
|
||||
self.norm2 = LayerNorm2d(c)
|
||||
|
||||
self.dropout1 = nn.Dropout(drop_out_rate) if drop_out_rate > 0.0 else nn.Identity()
|
||||
self.dropout2 = nn.Dropout(drop_out_rate) if drop_out_rate > 0.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):
|
||||
"""NAFNet: Nonlinear Activation Free Network for image restoration.
|
||||
|
||||
Encoder-decoder architecture with skip connections, using NAFBlock
|
||||
as the basic building block.
|
||||
|
||||
Args:
|
||||
img_channel: Number of input/output image channels. Default: 3.
|
||||
width: Base channel width. Default: 64.
|
||||
middle_blk_num: Number of NAFBlocks in the bottleneck. Default: 12.
|
||||
enc_blk_nums: Number of NAFBlocks per encoder stage. Default: [2,2,4,8].
|
||||
dec_blk_nums: Number of NAFBlocks per decoder stage. Default: [2,2,2,2].
|
||||
"""
|
||||
|
||||
def __init__(self, img_channel=3, width=64, middle_blk_num=12, enc_blk_nums=[2, 2, 4, 8],
|
||||
dec_blk_nums=[2, 2, 2, 2]):
|
||||
super().__init__()
|
||||
|
||||
self.intro = nn.Conv2d(
|
||||
in_channels=img_channel, out_channels=width, kernel_size=3, padding=1, stride=1, groups=1, bias=True,
|
||||
)
|
||||
self.ending = nn.Conv2d(
|
||||
in_channels=width, out_channels=img_channel, kernel_size=3, padding=1, stride=1, groups=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
|
||||
@@ -0,0 +1,305 @@
|
||||
# SCUNet: Swin-Conv-UNet for blind image denoising
|
||||
# Adapted from https://github.com/cszn/SCUNet (MIT License)
|
||||
# Original paper: "Practical Blind Image Denoising via Swin-Conv-UNet and Data Synthesis"
|
||||
# Authors: Kai Zhang, Yawei Li, Jingyun Liang, Jiezhang Cao, Yulun Zhang,
|
||||
# Hao Tang, Deng-Ping Fan, Radu Timofte, Luc Van Gool
|
||||
|
||||
import math
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from einops.layers.torch import Rearrange
|
||||
|
||||
|
||||
def _trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
|
||||
"""Truncated normal initialization (inline to avoid timm dependency)."""
|
||||
with torch.no_grad():
|
||||
l = (1.0 + math.erf((a - mean) / (std * math.sqrt(2.0)))) / 2.0
|
||||
u = (1.0 + math.erf((b - mean) / (std * math.sqrt(2.0)))) / 2.0
|
||||
tensor.uniform_(2 * l - 1, 2 * u - 1)
|
||||
tensor.erfinv_()
|
||||
tensor.mul_(std * math.sqrt(2.0))
|
||||
tensor.add_(mean)
|
||||
tensor.clamp_(min=a, max=b)
|
||||
return tensor
|
||||
|
||||
|
||||
class DropPath(nn.Module):
|
||||
"""Stochastic depth (drop path) for regularization."""
|
||||
|
||||
def __init__(self, drop_prob=0.0):
|
||||
super().__init__()
|
||||
self.drop_prob = drop_prob
|
||||
|
||||
def forward(self, x):
|
||||
if self.drop_prob == 0.0 or not self.training:
|
||||
return x
|
||||
keep_prob = 1 - self.drop_prob
|
||||
shape = (x.shape[0],) + (1,) * (x.ndim - 1)
|
||||
random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device)
|
||||
random_tensor = torch.floor_(random_tensor + keep_prob)
|
||||
return x.div(keep_prob) * random_tensor
|
||||
|
||||
|
||||
class WMSA(nn.Module):
|
||||
"""Window Multi-head 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=0.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 the attention mask for shifted window 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):
|
||||
"""Swin Transformer 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 <= 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.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):
|
||||
"""Combined Swin Transformer and Convolution 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 <= 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: Swin-Conv-UNet for blind image denoising.
|
||||
|
||||
Args:
|
||||
in_nc: Number of input channels. Default: 3.
|
||||
config: Number of ConvTransBlocks at each stage. Default: [4,4,4,4,4,4,4].
|
||||
dim: Base channel dimension. Default: 64.
|
||||
drop_path_rate: Stochastic depth rate. Default: 0.0.
|
||||
input_resolution: Expected input spatial resolution. Default: 256.
|
||||
"""
|
||||
|
||||
def __init__(self, in_nc=3, config=[4, 4, 4, 4, 4, 4, 4], 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
|
||||
|
||||
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(config))]
|
||||
|
||||
self.m_head = [nn.Conv2d(in_nc, dim, 3, 1, 1, bias=False)]
|
||||
|
||||
begin = 0
|
||||
self.m_down1 = [
|
||||
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 = [
|
||||
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 = [
|
||||
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 = [
|
||||
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.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.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.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.Conv2d(dim, in_nc, 3, 1, 1, bias=False)]
|
||||
|
||||
self.m_head = nn.Sequential(*self.m_head)
|
||||
self.m_down1 = nn.Sequential(*self.m_down1)
|
||||
self.m_down2 = nn.Sequential(*self.m_down2)
|
||||
self.m_down3 = nn.Sequential(*self.m_down3)
|
||||
self.m_body = nn.Sequential(*self.m_body)
|
||||
self.m_up3 = nn.Sequential(*self.m_up3)
|
||||
self.m_up2 = nn.Sequential(*self.m_up2)
|
||||
self.m_up1 = nn.Sequential(*self.m_up1)
|
||||
self.m_tail = nn.Sequential(*self.m_tail)
|
||||
|
||||
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
|
||||
@@ -0,0 +1,575 @@
|
||||
"""Image noise removal with 4 quality tiers: quick, balanced, quality, maximum."""
|
||||
import sys
|
||||
import json
|
||||
import os
|
||||
|
||||
|
||||
def emit_progress(percent, stage):
|
||||
"""Emit structured progress to stderr for bridge.ts to capture."""
|
||||
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
||||
|
||||
|
||||
# Model paths - Docker locations as defaults, with env var overrides
|
||||
SCUNET_MODEL_PATH = os.environ.get(
|
||||
"SCUNET_MODEL_PATH",
|
||||
"/opt/models/scunet/scunet_color_real_psnr.pth",
|
||||
)
|
||||
|
||||
NAFNET_MODEL_PATH = os.environ.get(
|
||||
"NAFNET_MODEL_PATH",
|
||||
"/opt/models/nafnet/NAFNet-SIDD-width64.pth",
|
||||
)
|
||||
|
||||
# Local cache for dev installs
|
||||
_CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "stirling-image", "models")
|
||||
|
||||
# GitHub release URLs for auto-download
|
||||
SCUNET_URL = "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_psnr.pth"
|
||||
NAFNET_URL = "https://huggingface.co/mikestealth/nafnet-models/resolve/main/NAFNet-SIDD-width64.pth"
|
||||
|
||||
|
||||
def _get_model_path(env_path, filename, url):
|
||||
"""Resolve model path: env/Docker first, then local cache, then download."""
|
||||
if os.path.exists(env_path):
|
||||
return env_path
|
||||
|
||||
local_path = os.path.join(_CACHE_DIR, filename)
|
||||
if os.path.exists(local_path):
|
||||
return local_path
|
||||
|
||||
# Auto-download
|
||||
emit_progress(10, f"Downloading {filename}")
|
||||
os.makedirs(_CACHE_DIR, exist_ok=True)
|
||||
import urllib.request
|
||||
urllib.request.urlretrieve(url, local_path)
|
||||
return local_path
|
||||
|
||||
|
||||
def denoise_quick(img_array, strength, detail, color_noise):
|
||||
"""Bilateral filter denoising - fast, good for mild noise.
|
||||
|
||||
Processes in LAB color space to independently handle luminance
|
||||
and chrominance noise.
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
# Map strength 0-100 to bilateral filter params
|
||||
d = int(3 + (strength / 100) * 12) # diameter: 3-15
|
||||
sigma_base_color = 20 + (strength / 100) * 130 # 20-150
|
||||
sigma_base_space = 20 + (strength / 100) * 130 # 20-150
|
||||
|
||||
# Detail preservation reduces sigma values
|
||||
detail_factor = 1.0 - (detail / 100) * 0.7 # 1.0 down to 0.3
|
||||
sigma_color = sigma_base_color * detail_factor
|
||||
sigma_space = sigma_base_space * detail_factor
|
||||
|
||||
is_gray = len(img_array.shape) == 2 or (
|
||||
len(img_array.shape) == 3 and img_array.shape[2] == 1
|
||||
)
|
||||
|
||||
if is_gray:
|
||||
gray = img_array if len(img_array.shape) == 2 else img_array[:, :, 0]
|
||||
result = cv2.bilateralFilter(gray, d, sigma_color, sigma_space)
|
||||
if len(img_array.shape) == 3:
|
||||
result = result[:, :, np.newaxis]
|
||||
return result
|
||||
|
||||
# Convert to LAB for split luminance/chrominance processing
|
||||
lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB)
|
||||
l_ch, a_ch, b_ch = cv2.split(lab)
|
||||
|
||||
# Denoise L (luminance) channel
|
||||
l_ch = cv2.bilateralFilter(l_ch, d, sigma_color, sigma_space)
|
||||
|
||||
# Optionally denoise A/B (color) channels based on color_noise param
|
||||
if color_noise > 0:
|
||||
color_factor = color_noise / 100
|
||||
color_sigma = sigma_color * color_factor * 0.7
|
||||
color_d = max(3, int(d * 0.7))
|
||||
a_ch = cv2.bilateralFilter(a_ch, color_d, color_sigma, sigma_space * 0.5)
|
||||
b_ch = cv2.bilateralFilter(b_ch, color_d, color_sigma, sigma_space * 0.5)
|
||||
|
||||
result = cv2.merge([l_ch, a_ch, b_ch])
|
||||
return cv2.cvtColor(result, cv2.COLOR_LAB2RGB)
|
||||
|
||||
|
||||
def denoise_balanced(img_array, strength, detail, color_noise):
|
||||
"""Non-Local Means denoising with LAB split - good balance of speed and quality.
|
||||
|
||||
NLMeans compares patches across the image for more accurate noise
|
||||
estimation than bilateral filtering.
|
||||
"""
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
# Map strength 0-100 to filter strength h: 3-20
|
||||
h = 3 + (strength / 100) * 17
|
||||
|
||||
# Detail preservation controls search window size: 21 down to 11
|
||||
search_window = 21 - int((detail / 100) * 10)
|
||||
template_window = 7
|
||||
|
||||
is_gray = len(img_array.shape) == 2 or (
|
||||
len(img_array.shape) == 3 and img_array.shape[2] == 1
|
||||
)
|
||||
|
||||
if is_gray:
|
||||
gray = img_array if len(img_array.shape) == 2 else img_array[:, :, 0]
|
||||
result = cv2.fastNlMeansDenoising(gray, None, h, template_window, search_window)
|
||||
if len(img_array.shape) == 3:
|
||||
result = result[:, :, np.newaxis]
|
||||
return result
|
||||
|
||||
# Process in LAB space
|
||||
lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB)
|
||||
l_ch, a_ch, b_ch = cv2.split(lab)
|
||||
|
||||
# Denoise L channel with NLMeans
|
||||
l_ch = cv2.fastNlMeansDenoising(l_ch, None, h, template_window, search_window)
|
||||
|
||||
# Optionally denoise color channels
|
||||
if color_noise > 0:
|
||||
color_h = h * (color_noise / 100) * 0.6
|
||||
if color_h > 1:
|
||||
a_ch = cv2.fastNlMeansDenoising(a_ch, None, color_h, template_window, search_window)
|
||||
b_ch = cv2.fastNlMeansDenoising(b_ch, None, color_h, template_window, search_window)
|
||||
|
||||
result = cv2.merge([l_ch, a_ch, b_ch])
|
||||
return cv2.cvtColor(result, cv2.COLOR_LAB2RGB)
|
||||
|
||||
|
||||
def _run_ai_denoise(model, img_array, strength, detail, color_noise, device):
|
||||
"""Shared inference helper for AI-based denoise tiers (SCUNet, NAFNet).
|
||||
|
||||
Handles tensor conversion, tiling for large images, strength blending,
|
||||
detail preservation, and optional color noise post-processing.
|
||||
"""
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
original = img_array.copy()
|
||||
h, w = img_array.shape[:2]
|
||||
|
||||
# Convert to float32 tensor [0,1] in NCHW format
|
||||
tensor = torch.from_numpy(img_array.astype(np.float32) / 255.0)
|
||||
tensor = tensor.permute(2, 0, 1).unsqueeze(0) # HWC -> NCHW
|
||||
tensor = tensor.to(device)
|
||||
|
||||
emit_progress(50, "Running AI denoising")
|
||||
|
||||
with torch.inference_mode():
|
||||
# Decide whether to use tiling (for images > 2048px on either side)
|
||||
if h > 2048 or w > 2048:
|
||||
result_tensor = _tile_process(model, tensor, tile_size=512, overlap=32, device=device)
|
||||
else:
|
||||
# Pad to multiple of 8 for model compatibility
|
||||
pad_h = (8 - h % 8) % 8
|
||||
pad_w = (8 - w % 8) % 8
|
||||
if pad_h > 0 or pad_w > 0:
|
||||
tensor = F.pad(tensor, (0, pad_w, 0, pad_h), mode="reflect")
|
||||
|
||||
result_tensor = model(tensor)
|
||||
|
||||
# Remove padding
|
||||
if pad_h > 0 or pad_w > 0:
|
||||
result_tensor = result_tensor[:, :, :h, :w]
|
||||
|
||||
emit_progress(70, "Post-processing")
|
||||
|
||||
# Convert back to numpy uint8
|
||||
result = result_tensor.squeeze(0).permute(1, 2, 0).cpu().clamp(0, 1).numpy()
|
||||
result = (result * 255).astype(np.uint8)
|
||||
|
||||
# Blend with original based on strength (0 = no change, 100 = full denoise)
|
||||
blend = strength / 100.0
|
||||
result = (original.astype(np.float32) * (1 - blend) + result.astype(np.float32) * blend)
|
||||
result = np.clip(result, 0, 255).astype(np.uint8)
|
||||
|
||||
# Detail preservation: extract high-frequency from original, add back
|
||||
if detail > 0:
|
||||
detail_scale = detail / 100.0
|
||||
# Blur original to get low-frequency component
|
||||
kernel_size = 5
|
||||
blurred = cv2.GaussianBlur(
|
||||
original.astype(np.float32),
|
||||
(kernel_size, kernel_size),
|
||||
0,
|
||||
)
|
||||
# High-frequency = original - blurred
|
||||
high_freq = original.astype(np.float32) - blurred
|
||||
# Add high frequency back to result, scaled by detail
|
||||
result = result.astype(np.float32) + high_freq * detail_scale
|
||||
result = np.clip(result, 0, 255).astype(np.uint8)
|
||||
|
||||
# Color noise post-processing: denoise A/B channels in LAB with NLMeans
|
||||
if color_noise > 0:
|
||||
color_h = 3 + (color_noise / 100) * 12
|
||||
lab = cv2.cvtColor(result, cv2.COLOR_RGB2LAB)
|
||||
l_ch, a_ch, b_ch = cv2.split(lab)
|
||||
a_ch = cv2.fastNlMeansDenoising(a_ch, None, color_h, 7, 21)
|
||||
b_ch = cv2.fastNlMeansDenoising(b_ch, None, color_h, 7, 21)
|
||||
lab = cv2.merge([l_ch, a_ch, b_ch])
|
||||
result = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _tile_process(model, tensor, tile_size=512, overlap=32, device="cpu"):
|
||||
"""Process large images in overlapping tiles to avoid OOM.
|
||||
|
||||
Tiles are blended at overlapping edges using linear ramps
|
||||
for seamless results.
|
||||
"""
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
_, c, h, w = tensor.shape
|
||||
step = tile_size - overlap
|
||||
|
||||
# Allocate output and weight map for blending
|
||||
output = torch.zeros_like(tensor)
|
||||
weight = torch.zeros((1, 1, h, w), device=device)
|
||||
|
||||
# Create blending weight ramp for overlap regions
|
||||
ramp = torch.ones((1, 1, tile_size, tile_size), device=device)
|
||||
if overlap > 0:
|
||||
for i in range(overlap):
|
||||
factor = (i + 1) / (overlap + 1)
|
||||
ramp[:, :, i, :] *= factor # top edge
|
||||
ramp[:, :, -1 - i, :] *= factor # bottom edge
|
||||
ramp[:, :, :, i] *= factor # left edge
|
||||
ramp[:, :, :, -1 - i] *= factor # right edge
|
||||
|
||||
tiles_y = max(1, (h - overlap + step - 1) // step)
|
||||
tiles_x = max(1, (w - overlap + step - 1) // step)
|
||||
total_tiles = tiles_y * tiles_x
|
||||
tile_count = 0
|
||||
|
||||
for y in range(0, h, step):
|
||||
for x in range(0, w, step):
|
||||
y_end = min(y + tile_size, h)
|
||||
x_end = min(x + tile_size, w)
|
||||
y_start = max(0, y_end - tile_size)
|
||||
x_start = max(0, x_end - tile_size)
|
||||
|
||||
tile = tensor[:, :, y_start:y_end, x_start:x_end]
|
||||
|
||||
# Pad tile if smaller than tile_size
|
||||
th, tw = tile.shape[2], tile.shape[3]
|
||||
pad_h = tile_size - th
|
||||
pad_w = tile_size - tw
|
||||
if pad_h > 0 or pad_w > 0:
|
||||
tile = F.pad(tile, (0, pad_w, 0, pad_h), mode="reflect")
|
||||
|
||||
result_tile = model(tile)
|
||||
|
||||
# Remove padding
|
||||
if pad_h > 0 or pad_w > 0:
|
||||
result_tile = result_tile[:, :, :th, :tw]
|
||||
|
||||
# Trim ramp to actual tile size
|
||||
tile_ramp = ramp[:, :, :th, :tw]
|
||||
|
||||
output[:, :, y_start:y_end, x_start:x_end] += result_tile * tile_ramp
|
||||
weight[:, :, y_start:y_end, x_start:x_end] += tile_ramp
|
||||
|
||||
tile_count += 1
|
||||
pct = 50 + int(20 * tile_count / total_tiles)
|
||||
emit_progress(pct, f"Processing tile {tile_count}/{total_tiles}")
|
||||
|
||||
# Normalize by weight
|
||||
weight = torch.clamp(weight, min=1e-6)
|
||||
output = output / weight
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def denoise_quality(img_array, strength, detail, color_noise, model_path):
|
||||
"""SCUNet-based denoising - high quality, slower.
|
||||
|
||||
Uses the Swin-Conv-UNet architecture trained on real-world noise.
|
||||
"""
|
||||
import torch
|
||||
from gpu import gpu_available
|
||||
|
||||
emit_progress(15, "Loading SCUNet model")
|
||||
|
||||
# Redirect stdout during model loading/inference
|
||||
stdout_fd = os.dup(1)
|
||||
os.dup2(2, 1)
|
||||
|
||||
try:
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "models"))
|
||||
from scunet_arch import SCUNet
|
||||
|
||||
use_gpu = gpu_available()
|
||||
device = torch.device("cuda" if use_gpu else "cpu")
|
||||
|
||||
model = SCUNet(in_nc=3, config=[4, 4, 4, 4, 4, 4, 4], dim=64)
|
||||
|
||||
resolved_path = _get_model_path(model_path, "scunet_color_real_psnr.pth", SCUNET_URL)
|
||||
checkpoint = torch.load(resolved_path, map_location=device, weights_only=True)
|
||||
model.load_state_dict(checkpoint)
|
||||
|
||||
model = model.to(device)
|
||||
model.eval()
|
||||
|
||||
emit_progress(30, "SCUNet model loaded")
|
||||
|
||||
result = _run_ai_denoise(model, img_array, strength, detail, color_noise, device)
|
||||
|
||||
# Free model from memory
|
||||
del model
|
||||
if use_gpu:
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return result
|
||||
finally:
|
||||
os.dup2(stdout_fd, 1)
|
||||
os.close(stdout_fd)
|
||||
|
||||
|
||||
def denoise_maximum(img_array, strength, detail, color_noise, model_path):
|
||||
"""NAFNet-based denoising - maximum quality, slowest.
|
||||
|
||||
Uses the Nonlinear Activation Free Network architecture for
|
||||
state-of-the-art image restoration.
|
||||
"""
|
||||
import torch
|
||||
from gpu import gpu_available
|
||||
|
||||
emit_progress(15, "Loading NAFNet model")
|
||||
|
||||
# Redirect stdout during model loading/inference
|
||||
stdout_fd = os.dup(1)
|
||||
os.dup2(2, 1)
|
||||
|
||||
try:
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "models"))
|
||||
from nafnet_arch import NAFNet
|
||||
|
||||
use_gpu = gpu_available()
|
||||
device = torch.device("cuda" if use_gpu else "cpu")
|
||||
|
||||
model = NAFNet(
|
||||
img_channel=3,
|
||||
width=64,
|
||||
middle_blk_num=12,
|
||||
enc_blk_nums=[2, 2, 4, 8],
|
||||
dec_blk_nums=[2, 2, 2, 2],
|
||||
)
|
||||
|
||||
resolved_path = _get_model_path(model_path, "NAFNet-SIDD-width64.pth", NAFNET_URL)
|
||||
checkpoint = torch.load(resolved_path, map_location=device, weights_only=True)
|
||||
|
||||
# NAFNet checkpoints may wrap state_dict under "params" key
|
||||
if "params" in checkpoint:
|
||||
checkpoint = checkpoint["params"]
|
||||
|
||||
model.load_state_dict(checkpoint)
|
||||
|
||||
model = model.to(device)
|
||||
model.eval()
|
||||
|
||||
emit_progress(30, "NAFNet model loaded")
|
||||
|
||||
result = _run_ai_denoise(model, img_array, strength, detail, color_noise, device)
|
||||
|
||||
# Free model from memory
|
||||
del model
|
||||
if use_gpu:
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return result
|
||||
finally:
|
||||
os.dup2(stdout_fd, 1)
|
||||
os.close(stdout_fd)
|
||||
|
||||
|
||||
def _process_single_image(img_array, settings, tier, strength, detail, color_noise):
|
||||
"""Run the appropriate denoise tier on a single image array (RGB uint8)."""
|
||||
if tier == "quick":
|
||||
return denoise_quick(img_array, strength, detail, color_noise)
|
||||
elif tier == "balanced":
|
||||
return denoise_balanced(img_array, strength, detail, color_noise)
|
||||
elif tier == "quality":
|
||||
return denoise_quality(img_array, strength, detail, color_noise, SCUNET_MODEL_PATH)
|
||||
elif tier == "maximum":
|
||||
return denoise_maximum(img_array, strength, detail, color_noise, NAFNET_MODEL_PATH)
|
||||
else:
|
||||
raise ValueError(f"Unknown tier: {tier}")
|
||||
|
||||
|
||||
def main():
|
||||
input_path = sys.argv[1]
|
||||
output_path = sys.argv[2]
|
||||
settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
|
||||
|
||||
tier = settings.get("tier", "balanced")
|
||||
strength = float(settings.get("strength", 50))
|
||||
detail = float(settings.get("detailPreservation", 50))
|
||||
color_noise = float(settings.get("colorNoise", 30))
|
||||
output_format = settings.get("format", "original")
|
||||
quality = int(settings.get("quality", 90))
|
||||
|
||||
try:
|
||||
emit_progress(5, "Opening image")
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
import cv2
|
||||
except ImportError:
|
||||
print(json.dumps({
|
||||
"success": False,
|
||||
"error": "OpenCV is not installed. Install with: pip install opencv-python-headless",
|
||||
}))
|
||||
sys.exit(1)
|
||||
|
||||
img = Image.open(input_path)
|
||||
original_format = img.format or "PNG"
|
||||
is_animated = getattr(img, "is_animated", False) and getattr(img, "n_frames", 1) > 1
|
||||
is_gif = original_format.upper() == "GIF"
|
||||
|
||||
# Determine output format
|
||||
if output_format == "original":
|
||||
# HEIC/HEIF -> PNG (Pillow can read but not always write these)
|
||||
if original_format.upper() in ("HEIC", "HEIF"):
|
||||
fmt = "png"
|
||||
elif is_gif and is_animated:
|
||||
fmt = "gif"
|
||||
else:
|
||||
fmt = original_format.lower()
|
||||
else:
|
||||
fmt = output_format.lower()
|
||||
|
||||
# Resolve output path extension
|
||||
ext_map = {
|
||||
"jpeg": ".jpg",
|
||||
"jpg": ".jpg",
|
||||
"png": ".png",
|
||||
"webp": ".webp",
|
||||
"tiff": ".tiff",
|
||||
"gif": ".gif",
|
||||
}
|
||||
base_path = output_path.rsplit(".", 1)[0]
|
||||
final_path = base_path + ext_map.get(fmt, ".png")
|
||||
|
||||
# Handle animated GIFs
|
||||
if is_animated and is_gif:
|
||||
ai_tier = tier in ("quality", "maximum")
|
||||
|
||||
if ai_tier:
|
||||
# AI tiers: process first frame only, save as static image
|
||||
emit_progress(10, "Processing first frame (AI mode)")
|
||||
frame = img.convert("RGB")
|
||||
img_array = np.array(frame)
|
||||
result_array = _process_single_image(img_array, settings, tier, strength, detail, color_noise)
|
||||
result = Image.fromarray(result_array)
|
||||
|
||||
# Override to static format
|
||||
if fmt == "gif":
|
||||
fmt = "png"
|
||||
final_path = base_path + ".png"
|
||||
else:
|
||||
# Classical tiers: process frame-by-frame
|
||||
frames = []
|
||||
durations = []
|
||||
n_frames = img.n_frames
|
||||
|
||||
for i in range(n_frames):
|
||||
img.seek(i)
|
||||
frame = img.convert("RGB")
|
||||
img_array = np.array(frame)
|
||||
|
||||
pct = 10 + int(80 * i / n_frames)
|
||||
emit_progress(pct, f"Denoising frame {i + 1}/{n_frames}")
|
||||
|
||||
result_array = _process_single_image(img_array, settings, tier, strength, detail, color_noise)
|
||||
result_frame = Image.fromarray(result_array)
|
||||
frames.append(result_frame)
|
||||
durations.append(img.info.get("duration", 100))
|
||||
|
||||
emit_progress(92, "Saving animated GIF")
|
||||
frames[0].save(
|
||||
final_path,
|
||||
save_all=True,
|
||||
append_images=frames[1:],
|
||||
duration=durations,
|
||||
loop=img.info.get("loop", 0),
|
||||
optimize=True,
|
||||
)
|
||||
|
||||
actual_w, actual_h = frames[0].size
|
||||
print(json.dumps({
|
||||
"success": True,
|
||||
"tier": tier,
|
||||
"width": actual_w,
|
||||
"height": actual_h,
|
||||
"frames": n_frames,
|
||||
"output_path": final_path,
|
||||
"format": "gif",
|
||||
}))
|
||||
return
|
||||
else:
|
||||
# Static image processing
|
||||
emit_progress(10, f"Denoising with {tier} tier")
|
||||
|
||||
# Convert to RGB for processing (handle RGBA, palette, grayscale)
|
||||
has_alpha = img.mode in ("RGBA", "LA", "PA")
|
||||
alpha_channel = None
|
||||
if has_alpha:
|
||||
alpha_channel = np.array(img.convert("RGBA"))[:, :, 3]
|
||||
img_rgb = img.convert("RGB")
|
||||
elif img.mode in ("L", "1"):
|
||||
img_rgb = img
|
||||
else:
|
||||
img_rgb = img.convert("RGB")
|
||||
|
||||
img_array = np.array(img_rgb)
|
||||
result_array = _process_single_image(img_array, settings, tier, strength, detail, color_noise)
|
||||
result = Image.fromarray(result_array)
|
||||
|
||||
# Re-attach alpha channel if present
|
||||
if alpha_channel is not None:
|
||||
result_rgba = result.convert("RGBA")
|
||||
r, g, b, _ = result_rgba.split()
|
||||
result = Image.merge("RGBA", (r, g, b, Image.fromarray(alpha_channel)))
|
||||
|
||||
# Save with format-specific options
|
||||
emit_progress(92, "Saving result")
|
||||
save_kwargs = {}
|
||||
|
||||
if fmt in ("jpeg", "jpg"):
|
||||
result = result.convert("RGB")
|
||||
save_kwargs["quality"] = quality
|
||||
save_kwargs["optimize"] = True
|
||||
elif fmt == "webp":
|
||||
save_kwargs["quality"] = quality
|
||||
elif fmt == "tiff":
|
||||
save_kwargs["compression"] = "tiff_lzw"
|
||||
elif fmt == "gif":
|
||||
result = result.convert("P", palette=Image.ADAPTIVE, colors=256)
|
||||
|
||||
result.save(final_path, **save_kwargs)
|
||||
|
||||
actual_w, actual_h = result.size
|
||||
print(json.dumps({
|
||||
"success": True,
|
||||
"tier": tier,
|
||||
"width": actual_w,
|
||||
"height": actual_h,
|
||||
"output_path": final_path,
|
||||
"format": fmt,
|
||||
}))
|
||||
|
||||
except Exception as e:
|
||||
print(json.dumps({"success": False, "error": str(e)}))
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user