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
@@ -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<void> {
|
||||
{ id: "image-enhancement", register: registerImageEnhancement },
|
||||
{ id: "content-aware-resize", register: registerContentAwareResize },
|
||||
{ id: "colorize", register: registerColorize },
|
||||
{ id: "noise-removal", register: registerNoiseRemoval },
|
||||
];
|
||||
|
||||
let skipped = 0;
|
||||
|
||||
@@ -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<string, string> = {
|
||||
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<typeof settingsSchema>;
|
||||
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<string, string> = {
|
||||
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",
|
||||
};
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -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<string, unknown>;
|
||||
onChange?: (settings: Record<string, unknown>) => void;
|
||||
}
|
||||
|
||||
export function NoiseRemovalControls({
|
||||
settings: initialSettings,
|
||||
onChange,
|
||||
}: NoiseRemovalControlsProps) {
|
||||
const [tier, setTier] = useState<Tier>("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 (
|
||||
<div className="space-y-4">
|
||||
{/* Tier selector */}
|
||||
<div>
|
||||
<p className="text-xs text-muted-foreground mb-1">Denoising Tier</p>
|
||||
<div className="grid grid-cols-4 gap-1">
|
||||
{TIERS.map((t) => (
|
||||
<button
|
||||
key={t.id}
|
||||
type="button"
|
||||
onClick={() => setTier(t.id)}
|
||||
className={`flex flex-col items-center gap-0.5 text-xs py-2 rounded transition-colors ${
|
||||
tier === t.id
|
||||
? "bg-primary text-primary-foreground"
|
||||
: "bg-muted text-muted-foreground hover:bg-muted/80"
|
||||
}`}
|
||||
>
|
||||
{t.label}
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
{activeTier && <p className="text-[10px] text-muted-foreground mt-1">{activeTier.desc}</p>}
|
||||
</div>
|
||||
|
||||
<div className="border-t border-border pt-3" />
|
||||
|
||||
{/* Strength slider */}
|
||||
<div>
|
||||
<div className="flex justify-between items-center">
|
||||
<p className="text-sm font-medium text-muted-foreground">Strength</p>
|
||||
<span className="text-sm font-mono tabular-nums font-medium">{strength}</span>
|
||||
</div>
|
||||
<input
|
||||
type="range"
|
||||
min={0}
|
||||
max={100}
|
||||
step={1}
|
||||
value={strength}
|
||||
onChange={(e) => setStrength(Number(e.target.value))}
|
||||
className="w-full h-1.5 rounded-full appearance-none bg-muted accent-primary"
|
||||
/>
|
||||
<div className="flex justify-between text-[10px] text-muted-foreground mt-0.5">
|
||||
<span>Subtle</span>
|
||||
<span>Aggressive</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Detail Preservation slider */}
|
||||
<div>
|
||||
<div className="flex justify-between items-center">
|
||||
<p className="text-sm font-medium text-muted-foreground">Detail Preservation</p>
|
||||
<span className="text-sm font-mono tabular-nums font-medium">{detailPreservation}</span>
|
||||
</div>
|
||||
<input
|
||||
type="range"
|
||||
min={0}
|
||||
max={100}
|
||||
step={1}
|
||||
value={detailPreservation}
|
||||
onChange={(e) => setDetailPreservation(Number(e.target.value))}
|
||||
className="w-full h-1.5 rounded-full appearance-none bg-muted accent-primary"
|
||||
/>
|
||||
<div className="flex justify-between text-[10px] text-muted-foreground mt-0.5">
|
||||
<span>Smooth</span>
|
||||
<span>Sharp</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Color Noise slider */}
|
||||
<div>
|
||||
<div className="flex justify-between items-center">
|
||||
<p className="text-sm font-medium text-muted-foreground">Color Noise</p>
|
||||
<span className="text-sm font-mono tabular-nums font-medium">{colorNoise}</span>
|
||||
</div>
|
||||
<input
|
||||
type="range"
|
||||
min={0}
|
||||
max={100}
|
||||
step={1}
|
||||
value={colorNoise}
|
||||
onChange={(e) => setColorNoise(Number(e.target.value))}
|
||||
className="w-full h-1.5 rounded-full appearance-none bg-muted accent-primary"
|
||||
/>
|
||||
<div className="flex justify-between text-[10px] text-muted-foreground mt-0.5">
|
||||
<span>Off</span>
|
||||
<span>Heavy</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="border-t border-border pt-3" />
|
||||
|
||||
{/* Output format */}
|
||||
<div>
|
||||
<p className="text-xs text-muted-foreground mb-1">Output Format</p>
|
||||
<div className="grid grid-cols-4 gap-1">
|
||||
{(["original", "png", "jpeg", "webp"] as const).map((f) => (
|
||||
<button
|
||||
key={f}
|
||||
type="button"
|
||||
onClick={() => setOutputFormat(f)}
|
||||
className={tabClass(outputFormat === f)}
|
||||
>
|
||||
{f === "original" ? "Original" : f.toUpperCase()}
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Quality slider (lossy formats only) */}
|
||||
{LOSSY_FORMATS.has(outputFormat) && (
|
||||
<div>
|
||||
<div className="flex justify-between items-center">
|
||||
<p className="text-sm font-medium text-muted-foreground">Quality</p>
|
||||
<span className="text-sm font-mono tabular-nums font-medium">{quality}</span>
|
||||
</div>
|
||||
<input
|
||||
type="range"
|
||||
min={1}
|
||||
max={100}
|
||||
step={1}
|
||||
value={quality}
|
||||
onChange={(e) => setQuality(Number(e.target.value))}
|
||||
className="w-full h-1.5 rounded-full appearance-none bg-muted accent-primary"
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
export function NoiseRemovalSettings() {
|
||||
const { files, entries } = useFileStore();
|
||||
const {
|
||||
processFiles,
|
||||
processAllFiles,
|
||||
processing,
|
||||
error,
|
||||
downloadUrl,
|
||||
originalSize,
|
||||
processedSize,
|
||||
progress,
|
||||
} = useToolProcessor("noise-removal");
|
||||
const [settings, setSettings] = useState<Record<string, unknown>>({});
|
||||
|
||||
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 (
|
||||
<div className="space-y-4">
|
||||
<NoiseRemovalControls onChange={setSettings} />
|
||||
|
||||
{/* GIF + AI tier warning */}
|
||||
{isGif && isAiTier && (
|
||||
<p className="text-xs text-amber-500">
|
||||
AI denoising on GIF files processes only the first frame. For animated GIFs, use the Quick
|
||||
or Balanced tier.
|
||||
</p>
|
||||
)}
|
||||
|
||||
{/* Error */}
|
||||
{error && <p className="text-xs text-red-500">{error}</p>}
|
||||
|
||||
{/* Size info */}
|
||||
{originalSize != null && processedSize != null && (
|
||||
<div className="text-xs text-muted-foreground space-y-0.5">
|
||||
<p>Original: {(originalSize / 1024).toFixed(1)} KB</p>
|
||||
<p>Denoised: {(processedSize / 1024).toFixed(1)} KB</p>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Process button / progress */}
|
||||
{processing ? (
|
||||
<ProgressCard
|
||||
active={processing}
|
||||
phase={progress.phase === "idle" ? "uploading" : progress.phase}
|
||||
label={hasMultiple ? `Removing noise from ${files.length} images` : "Removing noise"}
|
||||
percent={progress.percent}
|
||||
elapsed={progress.elapsed}
|
||||
/>
|
||||
) : (
|
||||
<button
|
||||
type="button"
|
||||
data-testid="noise-removal-submit"
|
||||
onClick={handleProcess}
|
||||
disabled={!hasFile || processing}
|
||||
className="w-full py-2.5 rounded-lg bg-primary text-primary-foreground font-medium disabled:opacity-50 disabled:cursor-not-allowed flex items-center justify-center gap-2"
|
||||
>
|
||||
{hasMultiple ? `Remove Noise (${files.length} files)` : "Remove Noise"}
|
||||
</button>
|
||||
)}
|
||||
|
||||
{/* Download (single file - batch uses Download All ZIP in tool-page) */}
|
||||
{!hasMultiple && downloadUrl && (
|
||||
<a
|
||||
href={downloadUrl}
|
||||
download
|
||||
data-testid="noise-removal-download"
|
||||
className="w-full py-2.5 rounded-lg border border-primary text-primary font-medium flex items-center justify-center gap-2 hover:bg-primary/5"
|
||||
>
|
||||
<Download className="h-4 w-4" />
|
||||
Download
|
||||
</a>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -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 <BlurFacesControls settings={settings} onChange={onChange} />;
|
||||
if (toolId === "remove-background")
|
||||
return <RemoveBgControls settings={settings} onChange={onChange} />;
|
||||
if (toolId === "noise-removal")
|
||||
return <NoiseRemovalControls settings={settings} onChange={onChange} />;
|
||||
if (COLOR_TOOL_IDS.has(toolId))
|
||||
return <ColorControls toolId={toolId} settings={settings} onChange={onChange} />;
|
||||
|
||||
|
||||
@@ -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<string, ToolRegistryEntry>([
|
||||
},
|
||||
],
|
||||
["colorize", { displayMode: "before-after", Settings: ColorizeSettings }],
|
||||
["noise-removal", { displayMode: "before-after", Settings: NoiseRemovalSettings }],
|
||||
]);
|
||||
|
||||
export function getToolRegistryEntry(toolId: string): ToolRegistryEntry | undefined {
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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()
|
||||
@@ -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";
|
||||
|
||||
@@ -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<NoiseRemovalResult> {
|
||||
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",
|
||||
};
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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" }];
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user