diff --git a/apps/api/src/routes/tools/index.ts b/apps/api/src/routes/tools/index.ts index cc96952e..fc0ccf8e 100644 --- a/apps/api/src/routes/tools/index.ts +++ b/apps/api/src/routes/tools/index.ts @@ -24,6 +24,7 @@ import { registerGifTools } from "./gif-tools.js"; import { registerImageEnhancement } from "./image-enhancement.js"; import { registerImageToPdf } from "./image-to-pdf.js"; import { registerInfo } from "./info.js"; +import { registerNoiseRemoval } from "./noise-removal.js"; import { registerOcr } from "./ocr.js"; import { registerPdfToImage } from "./pdf-to-image.js"; import { registerQrGenerate } from "./qr-generate.js"; @@ -132,6 +133,7 @@ export async function registerToolRoutes(app: FastifyInstance): Promise { { id: "image-enhancement", register: registerImageEnhancement }, { id: "content-aware-resize", register: registerContentAwareResize }, { id: "colorize", register: registerColorize }, + { id: "noise-removal", register: registerNoiseRemoval }, ]; let skipped = 0; diff --git a/apps/api/src/routes/tools/noise-removal.ts b/apps/api/src/routes/tools/noise-removal.ts new file mode 100644 index 00000000..4cf3199a --- /dev/null +++ b/apps/api/src/routes/tools/noise-removal.ts @@ -0,0 +1,183 @@ +import { randomUUID } from "node:crypto"; +import { join } from "node:path"; +import { noiseRemoval } from "@stirling-image/ai"; +import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify"; +import { z } from "zod"; +import { autoOrient } from "../../lib/auto-orient.js"; +import { validateImageBuffer } from "../../lib/file-validation.js"; +import { decodeHeic } from "../../lib/heic-converter.js"; +import { createWorkspace } from "../../lib/workspace.js"; +import { updateSingleFileProgress } from "../progress.js"; +import { registerToolProcessFn } from "../tool-factory.js"; + +const settingsSchema = z.object({ + tier: z.enum(["quick", "balanced", "quality", "maximum"]).default("balanced"), + strength: z.union([z.number(), z.string()]).transform(Number).default(50), + detailPreservation: z.union([z.number(), z.string()]).transform(Number).default(50), + colorNoise: z.union([z.number(), z.string()]).transform(Number).default(30), + format: z.enum(["original", "png", "jpeg", "webp"]).default("original"), + quality: z.union([z.number(), z.string()]).transform(Number).default(90), +}); + +/** + * AI noise removal route. + * Uses the Python sidecar for multi-tier denoising. + */ +export function registerNoiseRemoval(app: FastifyInstance) { + app.post("/api/v1/tools/noise-removal", async (request: FastifyRequest, reply: FastifyReply) => { + let fileBuffer: Buffer | null = null; + let filename = "image"; + let settingsRaw: string | null = null; + let clientJobId: string | null = null; + + try { + const parts = request.parts(); + for await (const part of parts) { + if (part.type === "file") { + const chunks: Buffer[] = []; + for await (const chunk of part.file) { + chunks.push(chunk); + } + fileBuffer = Buffer.concat(chunks); + filename = part.filename ?? "image"; + } else if (part.fieldname === "settings") { + settingsRaw = part.value as string; + } else if (part.fieldname === "clientJobId") { + clientJobId = part.value as string; + } + } + } catch (err) { + return reply.status(400).send({ + error: "Failed to parse multipart request", + details: err instanceof Error ? err.message : String(err), + }); + } + + if (!fileBuffer || fileBuffer.length === 0) { + return reply.status(400).send({ error: "No image file provided" }); + } + + const validation = await validateImageBuffer(fileBuffer); + if (!validation.valid) { + return reply.status(400).send({ error: `Invalid image: ${validation.reason}` }); + } + + try { + const parsed = settingsSchema.parse(settingsRaw ? JSON.parse(settingsRaw) : {}); + request.log.info( + { toolId: "noise-removal", imageSize: fileBuffer.length, tier: parsed.tier }, + "Starting noise removal", + ); + + // Decode HEIC/HEIF input via system decoder + if (validation.format === "heif") { + fileBuffer = await decodeHeic(fileBuffer); + } + + // Auto-orient to fix EXIF rotation before processing + fileBuffer = await autoOrient(fileBuffer); + + const jobId = randomUUID(); + const workspacePath = await createWorkspace(jobId); + + // Progress callback + const jobIdForProgress = clientJobId; + const onProgress = jobIdForProgress + ? (percent: number, stage: string) => { + updateSingleFileProgress({ + jobId: jobIdForProgress, + phase: "processing", + stage, + percent, + }); + } + : undefined; + + const result = await noiseRemoval( + fileBuffer, + join(workspacePath, "output"), + { + tier: parsed.tier, + strength: parsed.strength, + detailPreservation: parsed.detailPreservation, + colorNoise: parsed.colorNoise, + format: parsed.format, + quality: parsed.quality, + }, + onProgress, + ); + + if (clientJobId) { + updateSingleFileProgress({ + jobId: clientJobId, + phase: "complete", + percent: 100, + }); + } + + const CONTENT_TYPES: Record = { + png: "image/png", + jpeg: "image/jpeg", + jpg: "image/jpeg", + webp: "image/webp", + }; + const contentType = CONTENT_TYPES[result.format] || "image/png"; + const ext = result.format === "jpeg" ? "jpg" : result.format; + const outputFilename = `${filename.replace(/\.[^.]+$/, "")}_denoised.${ext}`; + + return reply + .header("Content-Type", contentType) + .header("Content-Disposition", `attachment; filename="${outputFilename}"`) + .header("X-Image-Width", String(result.width)) + .header("X-Image-Height", String(result.height)) + .send(result.buffer); + } catch (err) { + request.log.error({ err, toolId: "noise-removal" }, "Noise removal failed"); + return reply.status(422).send({ + error: "Noise removal failed", + details: err instanceof Error ? err.message : "Unknown error", + }); + } + }); + + // Register in the pipeline/batch registry so this tool can be used + // as a step in automation pipelines (without progress callbacks). + registerToolProcessFn({ + toolId: "noise-removal", + settingsSchema: z.object({ + tier: z.enum(["quick", "balanced", "quality", "maximum"]).default("balanced"), + strength: z.union([z.number(), z.string()]).transform(Number).default(50), + detailPreservation: z.union([z.number(), z.string()]).transform(Number).default(50), + colorNoise: z.union([z.number(), z.string()]).transform(Number).default(30), + format: z.enum(["original", "png", "jpeg", "webp"]).default("original"), + quality: z.union([z.number(), z.string()]).transform(Number).default(90), + }), + process: async (inputBuffer, settings, filename) => { + const s = settings as z.infer; + const orientedBuffer = await autoOrient(inputBuffer); + const jobId = randomUUID(); + const workspacePath = await createWorkspace(jobId); + const result = await noiseRemoval(orientedBuffer, join(workspacePath, "output"), { + tier: s.tier, + strength: s.strength, + detailPreservation: s.detailPreservation, + colorNoise: s.colorNoise, + format: s.format, + quality: s.quality, + }); + const ext = result.format === "jpeg" ? "jpg" : result.format; + const outputFilename = `${filename.replace(/\.[^.]+$/, "")}_denoised.${ext}`; + const CONTENT_TYPES: Record = { + png: "image/png", + jpeg: "image/jpeg", + jpg: "image/jpeg", + webp: "image/webp", + }; + return { + buffer: result.buffer, + filename: outputFilename, + contentType: CONTENT_TYPES[result.format] || "image/png", + }; + }, + }); +} diff --git a/apps/web/src/components/tools/noise-removal-settings.tsx b/apps/web/src/components/tools/noise-removal-settings.tsx new file mode 100644 index 00000000..73df7793 --- /dev/null +++ b/apps/web/src/components/tools/noise-removal-settings.tsx @@ -0,0 +1,290 @@ +import { Download } from "lucide-react"; +import { useEffect, useRef, useState } from "react"; +import { ProgressCard } from "@/components/common/progress-card"; +import { useToolProcessor } from "@/hooks/use-tool-processor"; +import { useFileStore } from "@/stores/file-store"; + +type Tier = "quick" | "balanced" | "quality" | "maximum"; + +const TIERS: { id: Tier; label: string; desc: string }[] = [ + { id: "quick", label: "Quick", desc: "Fast, lightweight" }, + { id: "balanced", label: "Balanced", desc: "Good quality, moderate speed" }, + { id: "quality", label: "Quality", desc: "AI-powered, slow" }, + { id: "maximum", label: "Maximum", desc: "Best AI model, slowest" }, +]; + +const LOSSY_FORMATS = new Set(["jpeg", "webp"]); + +export interface NoiseRemovalControlsProps { + settings?: Record; + onChange?: (settings: Record) => void; +} + +export function NoiseRemovalControls({ + settings: initialSettings, + onChange, +}: NoiseRemovalControlsProps) { + const [tier, setTier] = useState("balanced"); + const [strength, setStrength] = useState(50); + const [detailPreservation, setDetailPreservation] = useState(50); + const [colorNoise, setColorNoise] = useState(30); + const [outputFormat, setOutputFormat] = useState<"original" | "png" | "jpeg" | "webp">( + "original", + ); + const [quality, setQuality] = useState(90); + + // One-time init from pipeline settings + const initializedRef = useRef(false); + useEffect(() => { + if (!initialSettings || initializedRef.current) return; + initializedRef.current = true; + if (initialSettings.tier != null) setTier(initialSettings.tier as Tier); + if (initialSettings.strength != null) setStrength(Number(initialSettings.strength)); + if (initialSettings.detailPreservation != null) + setDetailPreservation(Number(initialSettings.detailPreservation)); + if (initialSettings.colorNoise != null) setColorNoise(Number(initialSettings.colorNoise)); + if (initialSettings.format != null) + setOutputFormat(initialSettings.format as "original" | "png" | "jpeg" | "webp"); + if (initialSettings.quality != null) setQuality(Number(initialSettings.quality)); + }, [initialSettings]); + + // Emit settings on change + const onChangeRef = useRef(onChange); + useEffect(() => { + onChangeRef.current = onChange; + }); + + useEffect(() => { + onChangeRef.current?.({ + tier, + strength, + detailPreservation, + colorNoise, + format: outputFormat, + quality, + }); + }, [tier, strength, detailPreservation, colorNoise, outputFormat, quality]); + + const tabClass = (active: boolean) => + `flex-1 text-xs py-1.5 rounded ${active ? "bg-primary text-primary-foreground" : "bg-muted text-muted-foreground hover:bg-muted/80"}`; + + const activeTier = TIERS.find((t) => t.id === tier); + + return ( +
+ {/* Tier selector */} +
+

Denoising Tier

+
+ {TIERS.map((t) => ( + + ))} +
+ {activeTier &&

{activeTier.desc}

} +
+ +
+ + {/* Strength slider */} +
+
+

Strength

+ {strength} +
+ setStrength(Number(e.target.value))} + className="w-full h-1.5 rounded-full appearance-none bg-muted accent-primary" + /> +
+ Subtle + Aggressive +
+
+ + {/* Detail Preservation slider */} +
+
+

Detail Preservation

+ {detailPreservation} +
+ setDetailPreservation(Number(e.target.value))} + className="w-full h-1.5 rounded-full appearance-none bg-muted accent-primary" + /> +
+ Smooth + Sharp +
+
+ + {/* Color Noise slider */} +
+
+

Color Noise

+ {colorNoise} +
+ setColorNoise(Number(e.target.value))} + className="w-full h-1.5 rounded-full appearance-none bg-muted accent-primary" + /> +
+ Off + Heavy +
+
+ +
+ + {/* Output format */} +
+

Output Format

+
+ {(["original", "png", "jpeg", "webp"] as const).map((f) => ( + + ))} +
+
+ + {/* Quality slider (lossy formats only) */} + {LOSSY_FORMATS.has(outputFormat) && ( +
+
+

Quality

+ {quality} +
+ setQuality(Number(e.target.value))} + className="w-full h-1.5 rounded-full appearance-none bg-muted accent-primary" + /> +
+ )} +
+ ); +} + +export function NoiseRemovalSettings() { + const { files, entries } = useFileStore(); + const { + processFiles, + processAllFiles, + processing, + error, + downloadUrl, + originalSize, + processedSize, + progress, + } = useToolProcessor("noise-removal"); + const [settings, setSettings] = useState>({}); + + const handleProcess = () => { + if (files.length > 1) { + processAllFiles(files, settings); + } else { + processFiles(files, settings); + } + }; + + const hasFile = files.length > 0; + const hasMultiple = files.length > 1; + + // Warn about GIF + AI tiers + const isGif = entries.some((e) => e.file.type === "image/gif"); + const isAiTier = settings.tier === "quality" || settings.tier === "maximum"; + + return ( +
+ + + {/* GIF + AI tier warning */} + {isGif && isAiTier && ( +

+ AI denoising on GIF files processes only the first frame. For animated GIFs, use the Quick + or Balanced tier. +

+ )} + + {/* Error */} + {error &&

{error}

} + + {/* Size info */} + {originalSize != null && processedSize != null && ( +
+

Original: {(originalSize / 1024).toFixed(1)} KB

+

Denoised: {(processedSize / 1024).toFixed(1)} KB

+
+ )} + + {/* Process button / progress */} + {processing ? ( + + ) : ( + + )} + + {/* Download (single file - batch uses Download All ZIP in tool-page) */} + {!hasMultiple && downloadUrl && ( + + + Download + + )} +
+ ); +} diff --git a/apps/web/src/components/tools/pipeline-step-settings.tsx b/apps/web/src/components/tools/pipeline-step-settings.tsx index 12bb1e71..16dda3e4 100644 --- a/apps/web/src/components/tools/pipeline-step-settings.tsx +++ b/apps/web/src/components/tools/pipeline-step-settings.tsx @@ -5,6 +5,7 @@ import { CompressControls } from "./compress-settings"; import { ConvertControls } from "./convert-settings"; import { CropControls } from "./crop-settings"; import { GifToolsControls } from "./gif-tools-settings"; +import { NoiseRemovalControls } from "./noise-removal-settings"; import { RemoveBgControls } from "./remove-bg-settings"; import { ReplaceColorControls } from "./replace-color-settings"; import { ResizeControls } from "./resize-settings"; @@ -44,6 +45,8 @@ export function PipelineStepSettings({ toolId, settings, onChange }: PipelineSte if (toolId === "blur-faces") return ; if (toolId === "remove-background") return ; + if (toolId === "noise-removal") + return ; if (COLOR_TOOL_IDS.has(toolId)) return ; diff --git a/apps/web/src/lib/tool-registry.tsx b/apps/web/src/lib/tool-registry.tsx index fffe5097..4929e047 100644 --- a/apps/web/src/lib/tool-registry.tsx +++ b/apps/web/src/lib/tool-registry.tsx @@ -249,6 +249,11 @@ const ColorizeSettings = lazy(() => default: m.ColorizeSettings, })), ); +const NoiseRemovalSettings = lazy(() => + import("@/components/tools/noise-removal-settings").then((m) => ({ + default: m.NoiseRemovalSettings, + })), +); // ── Color tool wrapper ───────────────────────────────────────────── // Color tools share a single component but differ by toolId. @@ -378,6 +383,7 @@ export const toolRegistry = new Map([ }, ], ["colorize", { displayMode: "before-after", Settings: ColorizeSettings }], + ["noise-removal", { displayMode: "before-after", Settings: NoiseRemovalSettings }], ]); export function getToolRegistryEntry(toolId: string): ToolRegistryEntry | undefined { diff --git a/docker/download_models.py b/docker/download_models.py index 21c620cc..5f9d7de9 100644 --- a/docker/download_models.py +++ b/docker/download_models.py @@ -39,6 +39,20 @@ DDCOLOR_MODEL_URL = ( DDCOLOR_ONNX_PATH = os.path.join(DDCOLOR_MODEL_DIR, "ddcolor.onnx") DDCOLOR_MIN_SIZE = 50_000_000 # ~220 MB ONNX +SCUNET_MODEL_DIR = "/opt/models/scunet" +SCUNET_MODEL_URL = ( + "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_psnr.pth" +) +SCUNET_MODEL_PATH = os.path.join(SCUNET_MODEL_DIR, "scunet_color_real_psnr.pth") +SCUNET_MIN_SIZE = 3_000_000 # ~4 MB + +NAFNET_MODEL_DIR = "/opt/models/nafnet" +NAFNET_MODEL_URL = ( + "https://huggingface.co/mikestealth/nafnet-models/resolve/main/" + "NAFNet-SIDD-width64.pth" +) +NAFNET_MODEL_PATH = os.path.join(NAFNET_MODEL_DIR, "NAFNet-SIDD-width64.pth") +NAFNET_MIN_SIZE = 60_000_000 # ~67 MB REMBG_MODELS = [ "u2net", @@ -218,6 +232,30 @@ def download_paddleocr_vl_model(): print(f" {model_name} ready\n") +def download_scunet_model(): + """Download SCUNet real-noise denoising model.""" + print(f"Downloading SCUNet model to {SCUNET_MODEL_PATH}...") + os.makedirs(SCUNET_MODEL_DIR, exist_ok=True) + urllib.request.urlretrieve(SCUNET_MODEL_URL, SCUNET_MODEL_PATH) + size = os.path.getsize(SCUNET_MODEL_PATH) + assert size > SCUNET_MIN_SIZE, ( + f"SCUNet model too small: {size} bytes (expected >{SCUNET_MIN_SIZE})" + ) + print(f" SCUNet model downloaded: {size:,} bytes") + + +def download_nafnet_model(): + """Download NAFNet SIDD width-64 denoising model.""" + print(f"Downloading NAFNet model to {NAFNET_MODEL_PATH}...") + os.makedirs(NAFNET_MODEL_DIR, exist_ok=True) + urllib.request.urlretrieve(NAFNET_MODEL_URL, NAFNET_MODEL_PATH) + size = os.path.getsize(NAFNET_MODEL_PATH) + assert size > NAFNET_MIN_SIZE, ( + f"NAFNet model too small: {size} bytes (expected >{NAFNET_MIN_SIZE})" + ) + print(f" NAFNet model downloaded: {size:,} bytes") + + def verify_mediapipe(): """Verify MediaPipe face detection models are bundled in the wheel.""" print("=== Verifying MediaPipe models ===") @@ -290,6 +328,16 @@ def smoke_test(): ) print(" DDColor ONNX model file verified") + # SCUNet model file must exist + assert os.path.exists(SCUNET_MODEL_PATH), f"SCUNet model not found: {SCUNET_MODEL_PATH}" + assert os.path.getsize(SCUNET_MODEL_PATH) > SCUNET_MIN_SIZE + print(" SCUNet model file verified") + + # NAFNet model file must exist + assert os.path.exists(NAFNET_MODEL_PATH), f"NAFNet model not found: {NAFNET_MODEL_PATH}" + assert os.path.getsize(NAFNET_MODEL_PATH) > NAFNET_MIN_SIZE + print(" NAFNet model file verified") + # PaddleOCR model directories must exist for repo_id in PADDLEOCR_MODELS: model_name = repo_id.split("/", 1)[1] @@ -315,6 +363,8 @@ def main(): download_ddcolor_model() download_paddleocr_models() download_paddleocr_vl_model() + download_scunet_model() + download_nafnet_model() verify_mediapipe() smoke_test() print("All models downloaded and verified.") diff --git a/packages/ai/python/models/__init__.py b/packages/ai/python/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/packages/ai/python/models/nafnet_arch.py b/packages/ai/python/models/nafnet_arch.py new file mode 100644 index 00000000..e0f0abe6 --- /dev/null +++ b/packages/ai/python/models/nafnet_arch.py @@ -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 diff --git a/packages/ai/python/models/scunet_arch.py b/packages/ai/python/models/scunet_arch.py new file mode 100644 index 00000000..26411e19 --- /dev/null +++ b/packages/ai/python/models/scunet_arch.py @@ -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 diff --git a/packages/ai/python/noise_removal.py b/packages/ai/python/noise_removal.py new file mode 100644 index 00000000..d92d21fd --- /dev/null +++ b/packages/ai/python/noise_removal.py @@ -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() diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index ba0ec0d1..880ba6ba 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -4,6 +4,7 @@ export { colorize } from "./colorization.js"; export type { DetectFacesResult, FaceRegion } from "./face-detection.js"; export { blurFaces, detectFaces } from "./face-detection.js"; export { inpaint } from "./inpainting.js"; +export { noiseRemoval } from "./noise-removal.js"; export { extractText } from "./ocr.js"; export { seamCarve } from "./seam-carving.js"; export { upscale } from "./upscaling.js"; diff --git a/packages/ai/src/noise-removal.ts b/packages/ai/src/noise-removal.ts new file mode 100644 index 00000000..e199af1e --- /dev/null +++ b/packages/ai/src/noise-removal.ts @@ -0,0 +1,52 @@ +import { readFile, writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import { type ProgressCallback, runPythonWithProgress } from "./bridge.js"; + +export interface NoiseRemovalOptions { + tier?: string; + strength?: number; + detailPreservation?: number; + colorNoise?: number; + format?: string; + quality?: number; +} + +export interface NoiseRemovalResult { + buffer: Buffer; + width: number; + height: number; + format: string; + tier: string; +} + +export async function noiseRemoval( + inputBuffer: Buffer, + outputDir: string, + options: NoiseRemovalOptions = {}, + onProgress?: ProgressCallback, +): Promise { + const inputPath = join(outputDir, "input_denoise.png"); + const outputPath = join(outputDir, "output_denoise.png"); + + await writeFile(inputPath, inputBuffer); + const { stdout } = await runPythonWithProgress( + "noise_removal.py", + [inputPath, outputPath, JSON.stringify(options)], + { onProgress }, + ); + + const result = JSON.parse(stdout); + if (!result.success) { + throw new Error(result.error || "Noise removal failed"); + } + + const actualOutputPath = result.output_path || outputPath; + const buffer = await readFile(actualOutputPath); + return { + buffer, + width: result.width, + height: result.height, + format: result.format ?? "png", + tier: result.tier ?? options.tier ?? "balanced", + }; +} diff --git a/packages/shared/src/constants.ts b/packages/shared/src/constants.ts index 19210004..a7ec073f 100644 --- a/packages/shared/src/constants.ts +++ b/packages/shared/src/constants.ts @@ -185,6 +185,14 @@ export const TOOLS: Tool[] = [ icon: "Palette", route: "/colorize", }, + { + id: "noise-removal", + name: "Noise Removal", + description: "AI-powered noise and grain removal", + category: "ai", + icon: "Sparkles", + route: "/noise-removal", + }, // Watermark & Overlay { id: "watermark-text", @@ -405,4 +413,5 @@ export const PYTHON_SIDECAR_TOOLS = [ "erase-object", "ocr", "colorize", + "noise-removal", ] as const; diff --git a/packages/shared/src/i18n/en.ts b/packages/shared/src/i18n/en.ts index ff3f85c0..5673790b 100644 --- a/packages/shared/src/i18n/en.ts +++ b/packages/shared/src/i18n/en.ts @@ -80,6 +80,10 @@ export const en = { description: "One-click auto-improve with smart exposure, contrast, color, and sharpness correction", }, + "noise-removal": { + name: "Noise Removal", + description: "AI-powered noise and grain removal", + }, "content-aware-resize": { name: "Content-Aware Resize", description: "Intelligently resize images while preserving important content", diff --git a/tests/e2e/tools-all.spec.ts b/tests/e2e/tools-all.spec.ts index 128cc0b2..ea382069 100644 --- a/tests/e2e/tools-all.spec.ts +++ b/tests/e2e/tools-all.spec.ts @@ -40,6 +40,7 @@ const TOOLS_WITH_DROPZONE = [ { id: "svg-to-raster", name: "SVG to Raster" }, { id: "vectorize", name: "Image to SVG" }, { id: "gif-tools", name: "GIF" }, + { id: "noise-removal", name: "Noise Removal" }, ]; const TOOLS_WITHOUT_DROPZONE = [{ id: "qr-generate", name: "QR Code" }]; diff --git a/tests/integration/api.test.ts b/tests/integration/api.test.ts index 38965f1a..41d6a22a 100644 --- a/tests/integration/api.test.ts +++ b/tests/integration/api.test.ts @@ -2720,6 +2720,7 @@ describe("Pipeline", () => { expect(toolIds).toContain("remove-background"); expect(toolIds).toContain("upscale"); expect(toolIds).toContain("blur-faces"); + expect(toolIds).toContain("noise-removal"); }); it("excludes tools that are not pipeline-compatible", async () => { @@ -4056,3 +4057,112 @@ describe("Image Enhancement", () => { expect(res.statusCode).toBe(400); }); }); + +// ═══════════════════════════════════════════════════════════════════════════ +// NOISE REMOVAL +// ═══════════════════════════════════════════════════════════════════════════ +describe("Noise Removal", () => { + it("POST /api/v1/tools/noise-removal processes with quick tier", async () => { + const { body: payload, contentType } = createMultipartPayload([ + { name: "file", filename: "noisy.png", contentType: "image/png", content: PNG_200x150 }, + { + name: "settings", + content: JSON.stringify({ tier: "quick", strength: 50 }), + }, + ]); + const res = await app.inject({ + method: "POST", + url: "/api/v1/tools/noise-removal", + headers: { + authorization: `Bearer ${adminToken}`, + "content-type": contentType, + }, + payload, + }); + // Either succeeds or Python sidecar unavailable in test env + expect(res.statusCode).not.toBe(400); + expect([200, 422]).toContain(res.statusCode); + if (res.statusCode === 200) { + expect(res.headers["content-type"]).toMatch(/^image\//); + expect(res.headers["content-disposition"]).toMatch(/attachment/); + } + }); + + it("returns 400 when no file is provided", async () => { + const { body: payload, contentType } = createMultipartPayload([ + { + name: "settings", + content: JSON.stringify({ tier: "quick", strength: 50 }), + }, + ]); + const res = await app.inject({ + method: "POST", + url: "/api/v1/tools/noise-removal", + headers: { + authorization: `Bearer ${adminToken}`, + "content-type": contentType, + }, + payload, + }); + expect(res.statusCode).toBe(400); + }); + + it("returns 400 for empty file", async () => { + const { body: payload, contentType } = createMultipartPayload([ + { name: "file", filename: "empty.png", contentType: "image/png", content: Buffer.alloc(0) }, + { + name: "settings", + content: JSON.stringify({ tier: "quick" }), + }, + ]); + const res = await app.inject({ + method: "POST", + url: "/api/v1/tools/noise-removal", + headers: { + authorization: `Bearer ${adminToken}`, + "content-type": contentType, + }, + payload, + }); + expect(res.statusCode).toBe(400); + }); + + it("accepts JPEG input", async () => { + const { body: payload, contentType } = createMultipartPayload([ + { name: "file", filename: "noisy.jpg", contentType: "image/jpeg", content: JPG_100x100 }, + { + name: "settings", + content: JSON.stringify({ tier: "quick", strength: 30 }), + }, + ]); + const res = await app.inject({ + method: "POST", + url: "/api/v1/tools/noise-removal", + headers: { + authorization: `Bearer ${adminToken}`, + "content-type": contentType, + }, + payload, + }); + expect(res.statusCode).not.toBe(400); + expect([200, 422]).toContain(res.statusCode); + }); + + it("uses default settings when none provided", async () => { + const { body: payload, contentType } = createMultipartPayload([ + { name: "file", filename: "noisy.png", contentType: "image/png", content: PNG_200x150 }, + ]); + const res = await app.inject({ + method: "POST", + url: "/api/v1/tools/noise-removal", + headers: { + authorization: `Bearer ${adminToken}`, + "content-type": contentType, + }, + payload, + }); + // Defaults should be accepted without validation errors + expect(res.statusCode).not.toBe(400); + expect([200, 422]).toContain(res.statusCode); + }); +});