mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
Add comprehensive photo restoration tool that chains multiple AI models: - Scratch/tear/spot detection via morphological analysis (top-hat/black-hat transforms) - Damage inpainting via LaMa ONNX model (reuses existing infrastructure) - Face enhancement via CodeFormer ONNX (~377MB, from facefusion/models-3.0.0) - Noise reduction via OpenCV NLMeans in LAB color space - Optional B&W auto-colorization via DDColor (reuses existing model) Settings: 3 restoration modes (Light/Auto/Heavy), individual feature toggles for scratch removal, face enhancement (with fidelity slider), denoising (with strength slider), and auto-colorize. Before/after comparison view. Handles HEIC, HEIF, and all standard formats. Batch processing supported. No new Python dependencies - reuses onnxruntime, cv2, mediapipe, PIL. 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
8071fe61c5
commit
6a43cc1b77
@@ -33,6 +33,7 @@ import { registerRedEyeRemoval } from "./red-eye-removal.js";
|
||||
import { registerRemoveBackground } from "./remove-background.js";
|
||||
import { registerReplaceColor } from "./replace-color.js";
|
||||
import { registerResize } from "./resize.js";
|
||||
import { registerRestorePhoto } from "./restore-photo.js";
|
||||
import { registerRotate } from "./rotate.js";
|
||||
import { registerSharpening } from "./sharpening.js";
|
||||
import { registerSmartCrop } from "./smart-crop.js";
|
||||
@@ -138,6 +139,7 @@ export async function registerToolRoutes(app: FastifyInstance): Promise<void> {
|
||||
{ id: "enhance-faces", register: registerEnhanceFaces },
|
||||
{ id: "noise-removal", register: registerNoiseRemoval },
|
||||
{ id: "red-eye-removal", register: registerRedEyeRemoval },
|
||||
{ id: "restore-photo", register: registerRestorePhoto },
|
||||
];
|
||||
|
||||
let skipped = 0;
|
||||
|
||||
@@ -0,0 +1,214 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { writeFile } from "node:fs/promises";
|
||||
import { basename, join } from "node:path";
|
||||
import { restorePhoto } from "@stirling-image/ai";
|
||||
import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify";
|
||||
import sharp from "sharp";
|
||||
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 { resolveOutputFormat } from "../../lib/output-format.js";
|
||||
import { createWorkspace } from "../../lib/workspace.js";
|
||||
import { updateSingleFileProgress } from "../progress.js";
|
||||
import { registerToolProcessFn } from "../tool-factory.js";
|
||||
|
||||
const settingsSchema = z.object({
|
||||
mode: z.enum(["auto", "light", "heavy"]).default("auto"),
|
||||
scratchRemoval: z.boolean().default(true),
|
||||
faceEnhancement: z.boolean().default(true),
|
||||
fidelity: z.number().min(0).max(1).default(0.7),
|
||||
denoise: z.boolean().default(true),
|
||||
denoiseStrength: z.number().min(0).max(100).default(40),
|
||||
colorize: z.boolean().default(false),
|
||||
});
|
||||
|
||||
/**
|
||||
* AI photo restoration route.
|
||||
* Multi-step pipeline: scratch repair, face enhancement, denoising,
|
||||
* optional colorization.
|
||||
*/
|
||||
export function registerRestorePhoto(app: FastifyInstance) {
|
||||
app.post("/api/v1/tools/restore-photo", 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 = basename(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 settings = settingsSchema.parse(settingsRaw ? JSON.parse(settingsRaw) : {});
|
||||
|
||||
request.log.info(
|
||||
{ toolId: "restore-photo", imageSize: fileBuffer.length, mode: settings.mode },
|
||||
"Starting photo restoration",
|
||||
);
|
||||
|
||||
// Decode HEIC/HEIF input
|
||||
if (validation.format === "heif") {
|
||||
fileBuffer = await decodeHeic(fileBuffer);
|
||||
}
|
||||
|
||||
// Auto-orient to fix EXIF rotation
|
||||
fileBuffer = await autoOrient(fileBuffer);
|
||||
|
||||
const jobId = randomUUID();
|
||||
const workspacePath = await createWorkspace(jobId);
|
||||
|
||||
// Save input
|
||||
const inputPath = join(workspacePath, "input", filename);
|
||||
await writeFile(inputPath, fileBuffer);
|
||||
|
||||
// Progress callback
|
||||
const jobIdForProgress = clientJobId;
|
||||
const onProgress = jobIdForProgress
|
||||
? (percent: number, stage: string) => {
|
||||
updateSingleFileProgress({
|
||||
jobId: jobIdForProgress,
|
||||
phase: "processing",
|
||||
stage,
|
||||
percent,
|
||||
});
|
||||
}
|
||||
: undefined;
|
||||
|
||||
// Process with Python sidecar
|
||||
const result = await restorePhoto(
|
||||
fileBuffer,
|
||||
join(workspacePath, "output"),
|
||||
{
|
||||
mode: settings.mode,
|
||||
scratchRemoval: settings.scratchRemoval,
|
||||
faceEnhancement: settings.faceEnhancement,
|
||||
fidelity: settings.fidelity,
|
||||
denoise: settings.denoise,
|
||||
denoiseStrength: settings.denoiseStrength,
|
||||
colorize: settings.colorize,
|
||||
},
|
||||
onProgress,
|
||||
);
|
||||
|
||||
// Resolve output format to match input
|
||||
const outputFormat = await resolveOutputFormat(fileBuffer, filename);
|
||||
let outputBuffer = result.buffer;
|
||||
|
||||
// Convert from PNG (Python output) to target format
|
||||
if (outputFormat.format !== "png") {
|
||||
outputBuffer = await sharp(result.buffer)
|
||||
.toFormat(outputFormat.format, { quality: outputFormat.quality })
|
||||
.toBuffer();
|
||||
}
|
||||
|
||||
// Save output
|
||||
const ext = outputFormat.format === "jpeg" ? "jpg" : outputFormat.format;
|
||||
const outputFilename = `${filename.replace(/\.[^.]+$/, "")}_restored.${ext}`;
|
||||
const outputPath = join(workspacePath, "output", outputFilename);
|
||||
await writeFile(outputPath, outputBuffer);
|
||||
|
||||
// Generate browser-compatible preview for non-previewable formats
|
||||
const BROWSER_PREVIEWABLE = new Set(["png", "jpg", "jpeg", "webp", "gif", "avif", "bmp"]);
|
||||
let previewUrl: string | undefined;
|
||||
if (!BROWSER_PREVIEWABLE.has(ext)) {
|
||||
try {
|
||||
const previewBuffer = await sharp(outputBuffer).webp({ quality: 80 }).toBuffer();
|
||||
const previewPath = join(workspacePath, "output", "preview.webp");
|
||||
await writeFile(previewPath, previewBuffer);
|
||||
previewUrl = `/api/v1/download/${jobId}/preview.webp`;
|
||||
} catch {
|
||||
// Non-fatal
|
||||
}
|
||||
}
|
||||
|
||||
if (clientJobId) {
|
||||
updateSingleFileProgress({
|
||||
jobId: clientJobId,
|
||||
phase: "complete",
|
||||
percent: 100,
|
||||
});
|
||||
}
|
||||
|
||||
return reply.send({
|
||||
jobId,
|
||||
downloadUrl: `/api/v1/download/${jobId}/${encodeURIComponent(outputFilename)}`,
|
||||
previewUrl,
|
||||
originalSize: fileBuffer.length,
|
||||
processedSize: outputBuffer.length,
|
||||
width: result.width,
|
||||
height: result.height,
|
||||
steps: result.steps,
|
||||
scratchCoverage: result.scratchCoverage,
|
||||
facesEnhanced: result.facesEnhanced,
|
||||
isGrayscale: result.isGrayscale,
|
||||
colorized: result.colorized,
|
||||
});
|
||||
} catch (err) {
|
||||
request.log.error({ err, toolId: "restore-photo" }, "Photo restoration failed");
|
||||
return reply.status(422).send({
|
||||
error: "Photo restoration failed",
|
||||
details: err instanceof Error ? err.message : "Unknown error",
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
// Register in the pipeline/batch registry
|
||||
registerToolProcessFn({
|
||||
toolId: "restore-photo",
|
||||
settingsSchema: z.object({
|
||||
mode: z.enum(["auto", "light", "heavy"]).default("auto"),
|
||||
scratchRemoval: z.boolean().default(true),
|
||||
faceEnhancement: z.boolean().default(true),
|
||||
fidelity: z.number().min(0).max(1).default(0.7),
|
||||
denoise: z.boolean().default(true),
|
||||
denoiseStrength: z.number().min(0).max(100).default(40),
|
||||
colorize: z.boolean().default(false),
|
||||
}),
|
||||
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 restorePhoto(orientedBuffer, join(workspacePath, "output"), {
|
||||
mode: s.mode,
|
||||
scratchRemoval: s.scratchRemoval,
|
||||
faceEnhancement: s.faceEnhancement,
|
||||
fidelity: s.fidelity,
|
||||
denoise: s.denoise,
|
||||
denoiseStrength: s.denoiseStrength,
|
||||
colorize: s.colorize,
|
||||
});
|
||||
const outputFilename = `${filename.replace(/\.[^.]+$/, "")}_restored.png`;
|
||||
return { buffer: result.buffer, filename: outputFilename, contentType: "image/png" };
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,288 @@
|
||||
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 Mode = "auto" | "light" | "heavy";
|
||||
|
||||
const MODES: { id: Mode; label: string; desc: string }[] = [
|
||||
{ id: "light", label: "Light", desc: "Gentle touch, preserves details" },
|
||||
{ id: "auto", label: "Auto", desc: "Balanced restoration" },
|
||||
{ id: "heavy", label: "Heavy", desc: "Aggressive repair for severe damage" },
|
||||
];
|
||||
|
||||
export interface RestorePhotoControlsProps {
|
||||
settings?: Record<string, unknown>;
|
||||
onChange?: (settings: Record<string, unknown>) => void;
|
||||
}
|
||||
|
||||
export function RestorePhotoControls({
|
||||
settings: initialSettings,
|
||||
onChange,
|
||||
}: RestorePhotoControlsProps) {
|
||||
const [mode, setMode] = useState<Mode>("auto");
|
||||
const [scratchRemoval, setScratchRemoval] = useState(true);
|
||||
const [faceEnhancement, setFaceEnhancement] = useState(true);
|
||||
const [fidelity, setFidelity] = useState(70);
|
||||
const [denoise, setDenoise] = useState(true);
|
||||
const [denoiseStrength, setDenoiseStrength] = useState(40);
|
||||
const [colorize, setColorize] = useState(false);
|
||||
|
||||
// One-time init from pipeline settings
|
||||
const initializedRef = useRef(false);
|
||||
useEffect(() => {
|
||||
if (!initialSettings || initializedRef.current) return;
|
||||
initializedRef.current = true;
|
||||
if (initialSettings.mode != null) setMode(initialSettings.mode as Mode);
|
||||
if (initialSettings.scratchRemoval != null)
|
||||
setScratchRemoval(Boolean(initialSettings.scratchRemoval));
|
||||
if (initialSettings.faceEnhancement != null)
|
||||
setFaceEnhancement(Boolean(initialSettings.faceEnhancement));
|
||||
if (initialSettings.fidelity != null) setFidelity(Number(initialSettings.fidelity) * 100);
|
||||
if (initialSettings.denoise != null) setDenoise(Boolean(initialSettings.denoise));
|
||||
if (initialSettings.denoiseStrength != null)
|
||||
setDenoiseStrength(Number(initialSettings.denoiseStrength));
|
||||
if (initialSettings.colorize != null) setColorize(Boolean(initialSettings.colorize));
|
||||
}, [initialSettings]);
|
||||
|
||||
// Emit settings on change
|
||||
const onChangeRef = useRef(onChange);
|
||||
useEffect(() => {
|
||||
onChangeRef.current = onChange;
|
||||
});
|
||||
|
||||
useEffect(() => {
|
||||
onChangeRef.current?.({
|
||||
mode,
|
||||
scratchRemoval,
|
||||
faceEnhancement,
|
||||
fidelity: fidelity / 100,
|
||||
denoise,
|
||||
denoiseStrength,
|
||||
colorize,
|
||||
});
|
||||
}, [mode, scratchRemoval, faceEnhancement, fidelity, denoise, denoiseStrength, colorize]);
|
||||
|
||||
const activeMode = MODES.find((m) => m.id === mode);
|
||||
|
||||
return (
|
||||
<div className="space-y-4">
|
||||
{/* Mode selector */}
|
||||
<div>
|
||||
<p className="text-xs text-muted-foreground mb-1">Restoration Mode</p>
|
||||
<div className="grid grid-cols-3 gap-1">
|
||||
{MODES.map((m) => (
|
||||
<button
|
||||
key={m.id}
|
||||
type="button"
|
||||
onClick={() => setMode(m.id)}
|
||||
className={`flex flex-col items-center gap-0.5 text-xs py-2 rounded transition-colors ${
|
||||
mode === m.id
|
||||
? "bg-primary text-primary-foreground"
|
||||
: "bg-muted text-muted-foreground hover:bg-muted/80"
|
||||
}`}
|
||||
>
|
||||
{m.label}
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
{activeMode && <p className="text-[10px] text-muted-foreground mt-1">{activeMode.desc}</p>}
|
||||
</div>
|
||||
|
||||
<div className="border-t border-border pt-3" />
|
||||
|
||||
{/* Feature toggles */}
|
||||
<div className="space-y-3">
|
||||
{/* Scratch Removal */}
|
||||
<label className="flex items-center justify-between cursor-pointer">
|
||||
<div>
|
||||
<p className="text-sm font-medium">Scratch Removal</p>
|
||||
<p className="text-[10px] text-muted-foreground">
|
||||
Detect and repair scratches, tears, spots
|
||||
</p>
|
||||
</div>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={scratchRemoval}
|
||||
onChange={(e) => setScratchRemoval(e.target.checked)}
|
||||
className="h-4 w-4 rounded border-muted-foreground accent-primary"
|
||||
/>
|
||||
</label>
|
||||
|
||||
{/* Face Enhancement */}
|
||||
<label className="flex items-center justify-between cursor-pointer">
|
||||
<div>
|
||||
<p className="text-sm font-medium">Face Enhancement</p>
|
||||
<p className="text-[10px] text-muted-foreground">
|
||||
Restore degraded faces with CodeFormer AI
|
||||
</p>
|
||||
</div>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={faceEnhancement}
|
||||
onChange={(e) => setFaceEnhancement(e.target.checked)}
|
||||
className="h-4 w-4 rounded border-muted-foreground accent-primary"
|
||||
/>
|
||||
</label>
|
||||
|
||||
{/* Fidelity slider (only when face enhancement is on) */}
|
||||
{faceEnhancement && (
|
||||
<div className="pl-2 border-l-2 border-primary/20">
|
||||
<div className="flex justify-between items-center">
|
||||
<p className="text-xs text-muted-foreground">Face Fidelity</p>
|
||||
<span className="text-xs font-mono tabular-nums">{fidelity}%</span>
|
||||
</div>
|
||||
<input
|
||||
type="range"
|
||||
min={0}
|
||||
max={100}
|
||||
step={5}
|
||||
value={fidelity}
|
||||
onChange={(e) => setFidelity(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>Enhanced</span>
|
||||
<span>Faithful</span>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Noise Reduction */}
|
||||
<label className="flex items-center justify-between cursor-pointer">
|
||||
<div>
|
||||
<p className="text-sm font-medium">Noise Reduction</p>
|
||||
<p className="text-[10px] text-muted-foreground">
|
||||
Remove grain and noise from old photos
|
||||
</p>
|
||||
</div>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={denoise}
|
||||
onChange={(e) => setDenoise(e.target.checked)}
|
||||
className="h-4 w-4 rounded border-muted-foreground accent-primary"
|
||||
/>
|
||||
</label>
|
||||
|
||||
{/* Denoise strength slider */}
|
||||
{denoise && (
|
||||
<div className="pl-2 border-l-2 border-primary/20">
|
||||
<div className="flex justify-between items-center">
|
||||
<p className="text-xs text-muted-foreground">Denoise Strength</p>
|
||||
<span className="text-xs font-mono tabular-nums">{denoiseStrength}</span>
|
||||
</div>
|
||||
<input
|
||||
type="range"
|
||||
min={0}
|
||||
max={100}
|
||||
step={5}
|
||||
value={denoiseStrength}
|
||||
onChange={(e) => setDenoiseStrength(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>Strong</span>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="border-t border-border pt-3" />
|
||||
|
||||
{/* Auto-Colorize */}
|
||||
<label className="flex items-center justify-between cursor-pointer">
|
||||
<div>
|
||||
<p className="text-sm font-medium">Auto-Colorize</p>
|
||||
<p className="text-[10px] text-muted-foreground">
|
||||
Add color to B&W photos using DDColor AI
|
||||
</p>
|
||||
</div>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={colorize}
|
||||
onChange={(e) => setColorize(e.target.checked)}
|
||||
className="h-4 w-4 rounded border-muted-foreground accent-primary"
|
||||
/>
|
||||
</label>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
export function RestorePhotoSettings() {
|
||||
const { files } = useFileStore();
|
||||
const {
|
||||
processFiles,
|
||||
processAllFiles,
|
||||
processing,
|
||||
error,
|
||||
downloadUrl,
|
||||
originalSize,
|
||||
processedSize,
|
||||
progress,
|
||||
} = useToolProcessor("restore-photo");
|
||||
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;
|
||||
|
||||
return (
|
||||
<div className="space-y-4">
|
||||
<RestorePhotoControls onChange={setSettings} />
|
||||
|
||||
{/* 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>Restored: {(processedSize / 1024).toFixed(1)} KB</p>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Process button / progress */}
|
||||
{processing ? (
|
||||
<ProgressCard
|
||||
active={processing}
|
||||
phase={progress.phase === "idle" ? "uploading" : progress.phase}
|
||||
label={hasMultiple ? `Restoring ${files.length} photos` : "Restoring photo"}
|
||||
percent={progress.percent}
|
||||
elapsed={progress.elapsed}
|
||||
/>
|
||||
) : (
|
||||
<button
|
||||
type="button"
|
||||
data-testid="restore-photo-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 ? `Restore Photos (${files.length})` : "Restore Photo"}
|
||||
</button>
|
||||
)}
|
||||
|
||||
{/* Download (single file) */}
|
||||
{!hasMultiple && downloadUrl && (
|
||||
<a
|
||||
href={downloadUrl}
|
||||
download
|
||||
data-testid="restore-photo-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>
|
||||
);
|
||||
}
|
||||
@@ -15,6 +15,7 @@ const TOOL_SUGGESTIONS: Record<string, string[]> = {
|
||||
"watermark-image": ["compress", "convert"],
|
||||
"text-overlay": ["compress", "convert"],
|
||||
colorize: ["adjust-colors", "image-enhancement", "upscale", "compress"],
|
||||
"restore-photo": ["colorize", "upscale", "image-enhancement", "adjust-colors"],
|
||||
sharpening: ["adjust-colors", "compress", "convert", "resize"],
|
||||
border: ["compress", "convert", "resize"],
|
||||
};
|
||||
|
||||
@@ -264,6 +264,11 @@ const RedEyeRemovalSettings = lazy(() =>
|
||||
default: m.RedEyeRemovalSettings,
|
||||
})),
|
||||
);
|
||||
const RestorePhotoSettings = lazy(() =>
|
||||
import("@/components/tools/restore-photo-settings").then((m) => ({
|
||||
default: m.RestorePhotoSettings,
|
||||
})),
|
||||
);
|
||||
|
||||
// ── Color tool wrapper ─────────────────────────────────────────────
|
||||
// Color tools share a single component but differ by toolId.
|
||||
@@ -396,6 +401,7 @@ export const toolRegistry = new Map<string, ToolRegistryEntry>([
|
||||
["colorize", { displayMode: "before-after", Settings: ColorizeSettings }],
|
||||
["noise-removal", { displayMode: "before-after", Settings: NoiseRemovalSettings }],
|
||||
["red-eye-removal", { displayMode: "before-after", Settings: RedEyeRemovalSettings }],
|
||||
["restore-photo", { displayMode: "before-after", Settings: RestorePhotoSettings }],
|
||||
]);
|
||||
|
||||
export function getToolRegistryEntry(toolId: string): ToolRegistryEntry | undefined {
|
||||
|
||||
@@ -53,6 +53,10 @@ SCUNET_MODEL_URL = (
|
||||
SCUNET_MODEL_PATH = os.path.join(SCUNET_MODEL_DIR, "scunet_color_real_psnr.pth")
|
||||
SCUNET_MIN_SIZE = 3_000_000 # ~4 MB
|
||||
|
||||
CODEFORMER_MODEL_DIR = "/opt/models/codeformer"
|
||||
CODEFORMER_ONNX_PATH = os.path.join(CODEFORMER_MODEL_DIR, "codeformer.onnx")
|
||||
CODEFORMER_MIN_SIZE = 100_000_000 # ~377 MB
|
||||
|
||||
NAFNET_MODEL_DIR = "/opt/models/nafnet"
|
||||
NAFNET_MODEL_URL = (
|
||||
"https://huggingface.co/mikestealth/nafnet-models/resolve/main/"
|
||||
@@ -219,6 +223,35 @@ def download_ddcolor_model():
|
||||
|
||||
|
||||
|
||||
def download_codeformer_model():
|
||||
"""Download CodeFormer ONNX model for AI face restoration.
|
||||
|
||||
Uses the pre-converted ONNX model from HuggingFace (facefusion repo)
|
||||
for direct inference via onnxruntime without needing PyTorch.
|
||||
"""
|
||||
print("=== Downloading CodeFormer ONNX model ===")
|
||||
os.makedirs(CODEFORMER_MODEL_DIR, exist_ok=True)
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
print(" Downloading CodeFormer ONNX from HuggingFace...")
|
||||
downloaded_path = hf_hub_download(
|
||||
repo_id="facefusion/models-3.0.0",
|
||||
filename="codeformer.onnx",
|
||||
local_dir=CODEFORMER_MODEL_DIR,
|
||||
)
|
||||
|
||||
actual_path = os.path.join(CODEFORMER_MODEL_DIR, "codeformer.onnx")
|
||||
if not os.path.exists(actual_path) and os.path.exists(downloaded_path):
|
||||
os.rename(downloaded_path, actual_path)
|
||||
|
||||
size = os.path.getsize(actual_path)
|
||||
assert size > CODEFORMER_MIN_SIZE, (
|
||||
f"CodeFormer model too small: {size} bytes (expected > {CODEFORMER_MIN_SIZE})"
|
||||
)
|
||||
print(f" CodeFormer ONNX model ready ({size / 1_000_000:.1f} MB)\n")
|
||||
|
||||
|
||||
def download_paddleocr_models():
|
||||
"""Pre-download PaddleOCR PP-OCRv5 model weights from HuggingFace.
|
||||
|
||||
@@ -358,6 +391,15 @@ def smoke_test():
|
||||
)
|
||||
print(" DDColor ONNX model file verified")
|
||||
|
||||
# CodeFormer ONNX model must exist
|
||||
assert os.path.exists(CODEFORMER_ONNX_PATH), (
|
||||
f"CodeFormer model missing: {CODEFORMER_ONNX_PATH}"
|
||||
)
|
||||
assert os.path.getsize(CODEFORMER_ONNX_PATH) > CODEFORMER_MIN_SIZE, (
|
||||
"CodeFormer model file is too small"
|
||||
)
|
||||
print(" CodeFormer 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
|
||||
@@ -392,6 +434,7 @@ def main():
|
||||
download_gfpgan_model()
|
||||
download_codeformer_model()
|
||||
download_ddcolor_model()
|
||||
download_codeformer_model()
|
||||
download_paddleocr_models()
|
||||
download_paddleocr_vl_model()
|
||||
download_scunet_model()
|
||||
|
||||
@@ -0,0 +1,597 @@
|
||||
"""AI photo restoration pipeline.
|
||||
|
||||
Multi-step pipeline for restoring old and damaged photos:
|
||||
1. Scratch & damage detection (morphological analysis)
|
||||
2. Damage inpainting (LaMa ONNX)
|
||||
3. Face enhancement (CodeFormer ONNX)
|
||||
4. Noise reduction (OpenCV NLMeans)
|
||||
5. Optional B&W colorization (DDColor ONNX)
|
||||
"""
|
||||
import sys
|
||||
import json
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import Image
|
||||
|
||||
|
||||
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 ───────────────────────────────────────────────────────
|
||||
|
||||
LAMA_MODEL_DIR = os.environ.get("LAMA_MODEL_DIR", "/opt/models/lama")
|
||||
LAMA_MODEL_PATH = os.path.join(LAMA_MODEL_DIR, "lama_fp32.onnx")
|
||||
LAMA_LOCAL_CACHE = os.path.join(os.path.expanduser("~"), ".cache", "stirling-image", "lama")
|
||||
LAMA_LOCAL_PATH = os.path.join(LAMA_LOCAL_CACHE, "lama_fp32.onnx")
|
||||
|
||||
CODEFORMER_MODEL_DIR = os.environ.get("CODEFORMER_MODEL_DIR", "/opt/models/codeformer")
|
||||
CODEFORMER_MODEL_PATH = os.path.join(CODEFORMER_MODEL_DIR, "codeformer.onnx")
|
||||
CODEFORMER_LOCAL_CACHE = os.path.join(
|
||||
os.path.expanduser("~"), ".cache", "stirling-image", "codeformer"
|
||||
)
|
||||
CODEFORMER_LOCAL_PATH = os.path.join(CODEFORMER_LOCAL_CACHE, "codeformer.onnx")
|
||||
|
||||
DDCOLOR_MODEL_PATH = os.environ.get(
|
||||
"DDCOLOR_MODEL_PATH", "/opt/models/ddcolor/ddcolor.onnx"
|
||||
)
|
||||
|
||||
LAMA_MODEL_SIZE = 512
|
||||
CODEFORMER_SIZE = 512
|
||||
|
||||
|
||||
# ── Scratch detection ─────────────────────────────────────────────────
|
||||
|
||||
def detect_scratches(img_bgr, sensitivity="medium"):
|
||||
"""Detect scratches, tears, and spots using morphological analysis.
|
||||
|
||||
Uses top-hat and black-hat transforms with oriented line kernels
|
||||
to find both bright and dark linear structures at multiple scales
|
||||
and angles. Returns a binary mask (255 = damage, 0 = clean).
|
||||
"""
|
||||
gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# CLAHE for local contrast enhancement to reveal faint scratches
|
||||
clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
|
||||
enhanced = clahe.apply(gray)
|
||||
|
||||
# Sensitivity controls the detection threshold
|
||||
thresh_map = {"light": 170, "medium": 130, "heavy": 90}
|
||||
thresh = thresh_map.get(sensitivity, 130)
|
||||
|
||||
h, w = gray.shape
|
||||
base_dim = min(h, w)
|
||||
|
||||
# Scale kernel sizes to image resolution
|
||||
kernel_sizes = [
|
||||
max(9, base_dim // 80),
|
||||
max(15, base_dim // 50),
|
||||
max(25, base_dim // 30),
|
||||
]
|
||||
|
||||
mask = np.zeros_like(gray)
|
||||
|
||||
for ksize in kernel_sizes:
|
||||
ksize = ksize | 1 # ensure odd
|
||||
for angle in [0, 45, 90, 135]:
|
||||
kernel = _make_line_kernel(ksize, angle)
|
||||
|
||||
# Black-hat: detects dark structures (dark scratches on light areas)
|
||||
blackhat = cv2.morphologyEx(enhanced, cv2.MORPH_BLACKHAT, kernel)
|
||||
# Top-hat: detects bright structures (light scratches on dark areas)
|
||||
tophat = cv2.morphologyEx(enhanced, cv2.MORPH_TOPHAT, kernel)
|
||||
|
||||
combined = cv2.add(blackhat, tophat)
|
||||
_, binary = cv2.threshold(combined, thresh, 255, cv2.THRESH_BINARY)
|
||||
mask = cv2.bitwise_or(mask, binary)
|
||||
|
||||
# Clean up: remove isolated noise pixels
|
||||
kernel_open = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
|
||||
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_open)
|
||||
|
||||
# Connect nearby scratch segments
|
||||
kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
|
||||
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_close)
|
||||
|
||||
# Dilate to include scratch edges for cleaner inpainting
|
||||
kernel_dilate = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
|
||||
mask = cv2.dilate(mask, kernel_dilate, iterations=1)
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
def _make_line_kernel(size, angle):
|
||||
"""Create an oriented line structuring element."""
|
||||
kernel = np.zeros((size, size), np.uint8)
|
||||
mid = size // 2
|
||||
if angle == 0:
|
||||
kernel[mid, :] = 1
|
||||
elif angle == 90:
|
||||
kernel[:, mid] = 1
|
||||
elif angle == 45:
|
||||
for i in range(size):
|
||||
kernel[i, i] = 1
|
||||
elif angle == 135:
|
||||
for i in range(size):
|
||||
kernel[i, size - 1 - i] = 1
|
||||
return kernel
|
||||
|
||||
|
||||
# ── LaMa inpainting ──────────────────────────────────────────────────
|
||||
|
||||
def _get_lama_path():
|
||||
"""Resolve LaMa model path, downloading if needed."""
|
||||
if os.path.exists(LAMA_MODEL_PATH):
|
||||
return LAMA_MODEL_PATH
|
||||
if os.path.exists(LAMA_LOCAL_PATH):
|
||||
return LAMA_LOCAL_PATH
|
||||
# Auto-download for local dev
|
||||
os.makedirs(LAMA_LOCAL_CACHE, exist_ok=True)
|
||||
import urllib.request
|
||||
url = "https://huggingface.co/Carve/LaMa-ONNX/resolve/main/lama_fp32.onnx"
|
||||
urllib.request.urlretrieve(url, LAMA_LOCAL_PATH)
|
||||
return LAMA_LOCAL_PATH
|
||||
|
||||
|
||||
def inpaint_damage(img_bgr, mask):
|
||||
"""Inpaint damaged areas using LaMa ONNX model.
|
||||
|
||||
Args:
|
||||
img_bgr: Input BGR image as numpy array.
|
||||
mask: Binary mask (255 = damage to repair, 0 = keep).
|
||||
|
||||
Returns:
|
||||
Restored BGR image with damage inpainted.
|
||||
"""
|
||||
import onnxruntime as ort
|
||||
|
||||
model_path = _get_lama_path()
|
||||
providers = ["CPUExecutionProvider"]
|
||||
if "CUDAExecutionProvider" in ort.get_available_providers():
|
||||
providers.insert(0, "CUDAExecutionProvider")
|
||||
|
||||
session = ort.InferenceSession(model_path, providers=providers)
|
||||
|
||||
orig_h, orig_w = img_bgr.shape[:2]
|
||||
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# Preprocess image: resize to 512x512, normalize to [0,1], NCHW
|
||||
img_resized = cv2.resize(img_rgb, (LAMA_MODEL_SIZE, LAMA_MODEL_SIZE))
|
||||
img_input = img_resized.astype(np.float32) / 255.0
|
||||
img_input = np.transpose(img_input, (2, 0, 1))[np.newaxis, ...] # (1,3,512,512)
|
||||
|
||||
# Preprocess mask: resize to 512x512, binary, NCHW
|
||||
mask_resized = cv2.resize(mask, (LAMA_MODEL_SIZE, LAMA_MODEL_SIZE),
|
||||
interpolation=cv2.INTER_NEAREST)
|
||||
mask_binary = (mask_resized > 127).astype(np.float32)
|
||||
mask_input = mask_binary[np.newaxis, np.newaxis, ...] # (1,1,512,512)
|
||||
|
||||
# Run inference
|
||||
outputs = session.run(None, {"image": img_input, "mask": mask_input})
|
||||
result = outputs[0][0] # (3, 512, 512)
|
||||
result = np.transpose(result, (1, 2, 0)) # (512, 512, 3)
|
||||
result = np.clip(result, 0, 255).astype(np.uint8)
|
||||
|
||||
# Resize inpainted result back to original dimensions
|
||||
inpainted = cv2.resize(result, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
|
||||
|
||||
# Feathered composite: preserve quality outside mask, blend at edges
|
||||
mask_full = mask.astype(np.float32) / 255.0
|
||||
feather_r = max(3, min(orig_w, orig_h) // 200)
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (feather_r, feather_r))
|
||||
dilated = cv2.dilate(mask_full, kernel, iterations=1)
|
||||
blur_size = feather_r * 2 + 1
|
||||
alpha = cv2.GaussianBlur(dilated, (blur_size, blur_size), 0)
|
||||
alpha = np.clip(alpha, 0.0, 1.0)[:, :, np.newaxis]
|
||||
|
||||
# Composite in RGB space, then convert back to BGR
|
||||
inpainted_rgb = inpainted
|
||||
original_rgb = img_rgb
|
||||
composited = (original_rgb.astype(np.float32) * (1.0 - alpha) +
|
||||
inpainted_rgb.astype(np.float32) * alpha)
|
||||
composited = np.clip(composited, 0, 255).astype(np.uint8)
|
||||
|
||||
return cv2.cvtColor(composited, cv2.COLOR_RGB2BGR)
|
||||
|
||||
|
||||
# ── CodeFormer face enhancement ──────────────────────────────────────
|
||||
|
||||
def _get_codeformer_path():
|
||||
"""Resolve CodeFormer ONNX model path, downloading if needed."""
|
||||
if os.path.exists(CODEFORMER_MODEL_PATH):
|
||||
return CODEFORMER_MODEL_PATH
|
||||
if os.path.exists(CODEFORMER_LOCAL_PATH):
|
||||
return CODEFORMER_LOCAL_PATH
|
||||
|
||||
# Auto-download for local dev
|
||||
os.makedirs(CODEFORMER_LOCAL_CACHE, exist_ok=True)
|
||||
emit_progress(35, "Downloading CodeFormer model")
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download(
|
||||
repo_id="facefusion/models-3.0.0",
|
||||
filename="codeformer.onnx",
|
||||
local_dir=CODEFORMER_LOCAL_CACHE,
|
||||
)
|
||||
return CODEFORMER_LOCAL_PATH
|
||||
|
||||
|
||||
def enhance_faces(img_bgr, fidelity=0.7):
|
||||
"""Enhance faces in the image using CodeFormer ONNX.
|
||||
|
||||
1. Detect faces with MediaPipe
|
||||
2. Crop each face with generous padding
|
||||
3. Run CodeFormer ONNX inference
|
||||
4. Paste enhanced face back with feathered blending
|
||||
|
||||
Args:
|
||||
img_bgr: Input BGR image.
|
||||
fidelity: 0.0 = aggressive enhancement, 1.0 = faithful to original.
|
||||
|
||||
Returns:
|
||||
Tuple of (enhanced BGR image, number of faces found).
|
||||
"""
|
||||
import mediapipe as mp
|
||||
import onnxruntime as ort
|
||||
|
||||
# Detect faces
|
||||
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
|
||||
ih, iw = img_bgr.shape[:2]
|
||||
|
||||
mp_face = mp.solutions.face_detection
|
||||
detections = []
|
||||
for model_sel in [0, 1]:
|
||||
detector = mp_face.FaceDetection(
|
||||
model_selection=model_sel, min_detection_confidence=0.4
|
||||
)
|
||||
results = detector.process(img_rgb)
|
||||
detector.close()
|
||||
if results.detections:
|
||||
detections = results.detections
|
||||
break
|
||||
|
||||
if not detections:
|
||||
return img_bgr, 0
|
||||
|
||||
# Load CodeFormer model
|
||||
model_path = _get_codeformer_path()
|
||||
providers = ["CPUExecutionProvider"]
|
||||
if "CUDAExecutionProvider" in ort.get_available_providers():
|
||||
providers.insert(0, "CUDAExecutionProvider")
|
||||
|
||||
session = ort.InferenceSession(model_path, providers=providers)
|
||||
input_names = [inp.name for inp in session.get_inputs()]
|
||||
|
||||
result = img_bgr.copy()
|
||||
faces_enhanced = 0
|
||||
|
||||
for detection in detections:
|
||||
bbox = detection.location_data.relative_bounding_box
|
||||
# Convert relative coords to absolute
|
||||
x = int(bbox.xmin * iw)
|
||||
y = int(bbox.ymin * ih)
|
||||
w = int(bbox.width * iw)
|
||||
h = int(bbox.height * ih)
|
||||
|
||||
# Skip very small faces (under 48px) - enhancement won't help
|
||||
if w < 48 or h < 48:
|
||||
continue
|
||||
|
||||
# Expand bounding box by ~80% for hair, forehead, chin
|
||||
pad_x = int(w * 0.8)
|
||||
pad_y = int(h * 0.8)
|
||||
x1 = max(0, x - pad_x)
|
||||
y1 = max(0, y - pad_y)
|
||||
x2 = min(iw, x + w + pad_x)
|
||||
y2 = min(ih, y + h + pad_y)
|
||||
|
||||
# Crop face region
|
||||
face_crop = img_bgr[y1:y2, x1:x2].copy()
|
||||
crop_h, crop_w = face_crop.shape[:2]
|
||||
|
||||
# Resize to 512x512 for CodeFormer
|
||||
face_resized = cv2.resize(face_crop, (CODEFORMER_SIZE, CODEFORMER_SIZE),
|
||||
interpolation=cv2.INTER_LANCZOS4)
|
||||
|
||||
# Preprocess: BGR -> RGB, normalize to [-1, 1]
|
||||
face_rgb = cv2.cvtColor(face_resized, cv2.COLOR_BGR2RGB)
|
||||
face_input = face_rgb.astype(np.float32) / 255.0
|
||||
face_input = (face_input - 0.5) / 0.5
|
||||
face_input = np.transpose(face_input, (2, 0, 1))
|
||||
face_input = np.expand_dims(face_input, 0) # (1, 3, 512, 512)
|
||||
|
||||
# Build model inputs
|
||||
model_inputs = {}
|
||||
for name in input_names:
|
||||
if name == "input":
|
||||
model_inputs[name] = face_input.astype(np.float32)
|
||||
elif name == "weight":
|
||||
model_inputs[name] = np.array([fidelity]).astype(np.float64)
|
||||
|
||||
# Run inference
|
||||
try:
|
||||
output = session.run(None, model_inputs)[0][0] # (3, 512, 512)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
# Postprocess: [-1, 1] -> [0, 255], RGB -> BGR
|
||||
output = np.clip(output, -1, 1)
|
||||
output = (output + 1) / 2
|
||||
output = np.transpose(output, (1, 2, 0)) # (512, 512, 3)
|
||||
output = (output * 255.0).clip(0, 255).astype(np.uint8)
|
||||
output_bgr = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
|
||||
|
||||
# Resize back to original crop size
|
||||
enhanced_crop = cv2.resize(output_bgr, (crop_w, crop_h),
|
||||
interpolation=cv2.INTER_LANCZOS4)
|
||||
|
||||
# Create feathered elliptical mask for smooth blending
|
||||
blend_mask = np.zeros((crop_h, crop_w), dtype=np.float32)
|
||||
center = (crop_w // 2, crop_h // 2)
|
||||
axes = (int(crop_w * 0.42), int(crop_h * 0.45))
|
||||
cv2.ellipse(blend_mask, center, axes, 0, 0, 360, 1.0, -1)
|
||||
|
||||
# Feather the mask edges
|
||||
blur_r = max(5, min(crop_w, crop_h) // 8) | 1
|
||||
blend_mask = cv2.GaussianBlur(blend_mask, (blur_r, blur_r), 0)
|
||||
blend_mask = blend_mask[:, :, np.newaxis]
|
||||
|
||||
# Blend enhanced face into result
|
||||
face_region = result[y1:y2, x1:x2].astype(np.float32)
|
||||
blended = face_region * (1.0 - blend_mask) + enhanced_crop.astype(np.float32) * blend_mask
|
||||
result[y1:y2, x1:x2] = np.clip(blended, 0, 255).astype(np.uint8)
|
||||
faces_enhanced += 1
|
||||
|
||||
return result, faces_enhanced
|
||||
|
||||
|
||||
# ── Denoising ─────────────────────────────────────────────────────────
|
||||
|
||||
def denoise_image(img_bgr, strength=40):
|
||||
"""Apply noise reduction using Non-Local Means in LAB color space.
|
||||
|
||||
Processes luminance and chrominance channels independently
|
||||
for better color preservation.
|
||||
|
||||
Args:
|
||||
img_bgr: Input BGR image.
|
||||
strength: 0-100, higher = more aggressive denoising.
|
||||
|
||||
Returns:
|
||||
Denoised BGR image.
|
||||
"""
|
||||
if strength <= 0:
|
||||
return img_bgr
|
||||
|
||||
# Map 0-100 to NLMeans filter strength
|
||||
h = 3 + (strength / 100) * 12 # 3-15
|
||||
|
||||
lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
|
||||
l_ch, a_ch, b_ch = cv2.split(lab)
|
||||
|
||||
# Denoise luminance channel
|
||||
l_ch = cv2.fastNlMeansDenoising(l_ch, None, h, 7, 21)
|
||||
|
||||
# Lightly denoise color channels to remove chroma noise
|
||||
color_h = h * 0.5
|
||||
if color_h > 1:
|
||||
a_ch = cv2.fastNlMeansDenoising(a_ch, None, color_h, 7, 21)
|
||||
b_ch = cv2.fastNlMeansDenoising(b_ch, None, color_h, 7, 21)
|
||||
|
||||
result = cv2.merge([l_ch, a_ch, b_ch])
|
||||
return cv2.cvtColor(result, cv2.COLOR_LAB2BGR)
|
||||
|
||||
|
||||
# ── B&W detection ────────────────────────────────────────────────────
|
||||
|
||||
def is_grayscale(img_bgr):
|
||||
"""Detect if an image is grayscale/B&W.
|
||||
|
||||
Checks if color channels are nearly identical by measuring
|
||||
the standard deviation of channel differences.
|
||||
"""
|
||||
if len(img_bgr.shape) == 2:
|
||||
return True
|
||||
if img_bgr.shape[2] == 1:
|
||||
return True
|
||||
|
||||
b, g, r = cv2.split(img_bgr)
|
||||
diff_rg = np.abs(r.astype(np.float32) - g.astype(np.float32)).mean()
|
||||
diff_rb = np.abs(r.astype(np.float32) - b.astype(np.float32)).mean()
|
||||
diff_gb = np.abs(g.astype(np.float32) - b.astype(np.float32)).mean()
|
||||
avg_diff = (diff_rg + diff_rb + diff_gb) / 3
|
||||
|
||||
return bool(avg_diff < 5.0)
|
||||
|
||||
|
||||
# ── DDColor colorization ─────────────────────────────────────────────
|
||||
|
||||
def colorize_bw(img_bgr, intensity=0.85):
|
||||
"""Colorize a B&W image using DDColor ONNX.
|
||||
|
||||
Reuses the DDColor model that the colorize tool already downloads.
|
||||
"""
|
||||
import onnxruntime as ort
|
||||
|
||||
if not os.path.exists(DDCOLOR_MODEL_PATH):
|
||||
return img_bgr, False
|
||||
|
||||
providers = ["CPUExecutionProvider"]
|
||||
try:
|
||||
from gpu import gpu_available
|
||||
if gpu_available():
|
||||
providers.insert(0, "CUDAExecutionProvider")
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
session = ort.InferenceSession(DDCOLOR_MODEL_PATH, providers=providers)
|
||||
input_name = session.get_inputs()[0].name
|
||||
input_shape = session.get_inputs()[0].shape
|
||||
model_size = (
|
||||
input_shape[2]
|
||||
if len(input_shape) == 4 and isinstance(input_shape[2], int)
|
||||
else 512
|
||||
)
|
||||
|
||||
orig_h, orig_w = img_bgr.shape[:2]
|
||||
img_lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
|
||||
orig_l = img_lab[:, :, 0].astype(np.float32)
|
||||
|
||||
# Prepare input
|
||||
img_resized = cv2.resize(img_bgr, (model_size, model_size))
|
||||
img_float = img_resized.astype(np.float32) / 255.0
|
||||
img_nchw = np.transpose(img_float, (2, 0, 1))[np.newaxis, ...]
|
||||
|
||||
output = session.run(None, {input_name: img_nchw})[0]
|
||||
ab_pred = output[0] # (2, model_size, model_size)
|
||||
|
||||
# Resize ab channels back to original
|
||||
ab_resized = np.zeros((2, orig_h, orig_w), dtype=np.float32)
|
||||
for i in range(2):
|
||||
ab_resized[i] = cv2.resize(ab_pred[i], (orig_w, orig_h))
|
||||
|
||||
ab_a = np.clip(ab_resized[0], -128, 127)
|
||||
ab_b = np.clip(ab_resized[1], -128, 127)
|
||||
|
||||
# Apply intensity blending
|
||||
if intensity < 1.0:
|
||||
orig_a = img_lab[:, :, 1].astype(np.float32) - 128.0
|
||||
orig_b = img_lab[:, :, 2].astype(np.float32) - 128.0
|
||||
ab_a = orig_a * (1 - intensity) + ab_a * intensity
|
||||
ab_b = orig_b * (1 - intensity) + ab_b * intensity
|
||||
|
||||
result_lab = np.zeros((orig_h, orig_w, 3), dtype=np.uint8)
|
||||
result_lab[:, :, 0] = np.clip(orig_l, 0, 255).astype(np.uint8)
|
||||
result_lab[:, :, 1] = np.clip(ab_a + 128.0, 0, 255).astype(np.uint8)
|
||||
result_lab[:, :, 2] = np.clip(ab_b + 128.0, 0, 255).astype(np.uint8)
|
||||
|
||||
return cv2.cvtColor(result_lab, cv2.COLOR_LAB2BGR), True
|
||||
|
||||
|
||||
# ── Main pipeline ─────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
input_path = sys.argv[1]
|
||||
output_path = sys.argv[2]
|
||||
settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
|
||||
|
||||
mode = settings.get("mode", "auto")
|
||||
scratch_removal = settings.get("scratchRemoval", True)
|
||||
face_enhancement = settings.get("faceEnhancement", True)
|
||||
fidelity = float(settings.get("fidelity", 0.7))
|
||||
do_denoise = settings.get("denoise", True)
|
||||
denoise_strength = float(settings.get("denoiseStrength", 40))
|
||||
do_colorize = settings.get("colorize", False)
|
||||
|
||||
# Mode presets override individual settings
|
||||
if mode == "light":
|
||||
scratch_sensitivity = "light"
|
||||
if denoise_strength > 30:
|
||||
denoise_strength = 30
|
||||
elif mode == "heavy":
|
||||
scratch_sensitivity = "heavy"
|
||||
if denoise_strength < 60:
|
||||
denoise_strength = 60
|
||||
else:
|
||||
scratch_sensitivity = "medium"
|
||||
|
||||
try:
|
||||
emit_progress(5, "Opening image")
|
||||
img_bgr = cv2.imread(input_path, cv2.IMREAD_COLOR)
|
||||
if img_bgr is None:
|
||||
pil_img = Image.open(input_path).convert("RGB")
|
||||
img_bgr = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
|
||||
|
||||
orig_h, orig_w = img_bgr.shape[:2]
|
||||
result = img_bgr.copy()
|
||||
steps_applied = []
|
||||
|
||||
# ── Step 1: Analyze photo ────────────────────────────────
|
||||
emit_progress(8, "Analyzing photo")
|
||||
bw_detected = is_grayscale(img_bgr)
|
||||
scratch_mask = None
|
||||
scratch_coverage = 0.0
|
||||
|
||||
# ── Step 2: Scratch detection & inpainting ───────────────
|
||||
if scratch_removal:
|
||||
emit_progress(10, "Detecting damage")
|
||||
scratch_mask = detect_scratches(result, scratch_sensitivity)
|
||||
scratch_pixels = np.count_nonzero(scratch_mask)
|
||||
total_pixels = scratch_mask.shape[0] * scratch_mask.shape[1]
|
||||
scratch_coverage = float(scratch_pixels / total_pixels)
|
||||
|
||||
if scratch_coverage > 0.001: # At least 0.1% coverage
|
||||
emit_progress(15, f"Repairing damage ({scratch_coverage:.1%} affected)")
|
||||
result = inpaint_damage(result, scratch_mask)
|
||||
steps_applied.append("scratch_removal")
|
||||
emit_progress(30, "Damage repaired")
|
||||
else:
|
||||
emit_progress(15, "No significant damage detected")
|
||||
else:
|
||||
emit_progress(15, "Scratch removal disabled")
|
||||
|
||||
# ── Step 3: Face enhancement ─────────────────────────────
|
||||
faces_found = 0
|
||||
if face_enhancement:
|
||||
emit_progress(35, "Detecting faces")
|
||||
try:
|
||||
result, faces_found = enhance_faces(result, fidelity)
|
||||
if faces_found > 0:
|
||||
steps_applied.append("face_enhancement")
|
||||
emit_progress(65, f"Enhanced {faces_found} face{'s' if faces_found != 1 else ''}")
|
||||
else:
|
||||
emit_progress(65, "No faces detected")
|
||||
except Exception as e:
|
||||
emit_progress(65, f"Face enhancement skipped: {str(e)[:40]}")
|
||||
else:
|
||||
emit_progress(65, "Face enhancement disabled")
|
||||
|
||||
# ── Step 4: Noise reduction ──────────────────────────────
|
||||
if do_denoise and denoise_strength > 0:
|
||||
emit_progress(70, "Reducing noise")
|
||||
result = denoise_image(result, denoise_strength)
|
||||
steps_applied.append("denoise")
|
||||
emit_progress(80, "Noise reduced")
|
||||
else:
|
||||
emit_progress(80, "Denoising disabled")
|
||||
|
||||
# ── Step 5: Colorization ─────────────────────────────────
|
||||
colorized = False
|
||||
if do_colorize and bw_detected:
|
||||
emit_progress(82, "Colorizing B&W photo")
|
||||
try:
|
||||
result, colorized = colorize_bw(result, intensity=0.85)
|
||||
if colorized:
|
||||
steps_applied.append("colorize")
|
||||
emit_progress(92, "Colorization complete")
|
||||
else:
|
||||
emit_progress(92, "Colorization model not available")
|
||||
except Exception as e:
|
||||
emit_progress(92, f"Colorization skipped: {str(e)[:40]}")
|
||||
else:
|
||||
emit_progress(92, "Colorization skipped")
|
||||
|
||||
# ── Save result ──────────────────────────────────────────
|
||||
emit_progress(95, "Saving result")
|
||||
cv2.imwrite(output_path, result)
|
||||
|
||||
print(json.dumps({
|
||||
"success": True,
|
||||
"width": orig_w,
|
||||
"height": orig_h,
|
||||
"steps": steps_applied,
|
||||
"scratchCoverage": round(scratch_coverage * 100, 2),
|
||||
"facesEnhanced": faces_found,
|
||||
"isGrayscale": bw_detected,
|
||||
"colorized": colorized,
|
||||
"output_path": output_path,
|
||||
}))
|
||||
|
||||
except Exception as e:
|
||||
print(json.dumps({"success": False, "error": str(e)}))
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -8,5 +8,6 @@ export { inpaint } from "./inpainting.js";
|
||||
export { noiseRemoval } from "./noise-removal.js";
|
||||
export { extractText } from "./ocr.js";
|
||||
export { removeRedEye } from "./red-eye-removal.js";
|
||||
export { restorePhoto } from "./restoration.js";
|
||||
export { seamCarve } from "./seam-carving.js";
|
||||
export { upscale } from "./upscaling.js";
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
import { readFile, writeFile } from "node:fs/promises";
|
||||
import { join } from "node:path";
|
||||
import { type ProgressCallback, runPythonWithProgress } from "./bridge.js";
|
||||
|
||||
export interface RestorePhotoOptions {
|
||||
mode?: string;
|
||||
scratchRemoval?: boolean;
|
||||
faceEnhancement?: boolean;
|
||||
fidelity?: number;
|
||||
denoise?: boolean;
|
||||
denoiseStrength?: number;
|
||||
colorize?: boolean;
|
||||
}
|
||||
|
||||
export interface RestorePhotoResult {
|
||||
buffer: Buffer;
|
||||
width: number;
|
||||
height: number;
|
||||
steps: string[];
|
||||
scratchCoverage: number;
|
||||
facesEnhanced: number;
|
||||
isGrayscale: boolean;
|
||||
colorized: boolean;
|
||||
}
|
||||
|
||||
export async function restorePhoto(
|
||||
inputBuffer: Buffer,
|
||||
outputDir: string,
|
||||
options: RestorePhotoOptions = {},
|
||||
onProgress?: ProgressCallback,
|
||||
): Promise<RestorePhotoResult> {
|
||||
const inputPath = join(outputDir, "input_restore.png");
|
||||
const outputPath = join(outputDir, "output_restore.png");
|
||||
|
||||
await writeFile(inputPath, inputBuffer);
|
||||
const { stdout } = await runPythonWithProgress(
|
||||
"restore.py",
|
||||
[inputPath, outputPath, JSON.stringify(options)],
|
||||
{ onProgress },
|
||||
);
|
||||
|
||||
const result = JSON.parse(stdout);
|
||||
if (!result.success) {
|
||||
throw new Error(result.error || "Photo restoration failed");
|
||||
}
|
||||
|
||||
const actualOutputPath = result.output_path || outputPath;
|
||||
const buffer = await readFile(actualOutputPath);
|
||||
return {
|
||||
buffer,
|
||||
width: result.width,
|
||||
height: result.height,
|
||||
steps: result.steps ?? [],
|
||||
scratchCoverage: result.scratchCoverage ?? 0,
|
||||
facesEnhanced: result.facesEnhanced ?? 0,
|
||||
isGrayscale: result.isGrayscale ?? false,
|
||||
colorized: result.colorized ?? false,
|
||||
};
|
||||
}
|
||||
@@ -209,6 +209,14 @@ export const TOOLS: Tool[] = [
|
||||
icon: "Eye",
|
||||
route: "/red-eye-removal",
|
||||
},
|
||||
{
|
||||
id: "restore-photo",
|
||||
name: "Photo Restoration",
|
||||
description: "Fix scratches, tears, and damage on old photos with AI",
|
||||
category: "ai",
|
||||
icon: "Undo2",
|
||||
route: "/restore-photo",
|
||||
},
|
||||
// Watermark & Overlay
|
||||
{
|
||||
id: "watermark-text",
|
||||
@@ -432,4 +440,5 @@ export const PYTHON_SIDECAR_TOOLS = [
|
||||
"enhance-faces",
|
||||
"noise-removal",
|
||||
"red-eye-removal",
|
||||
"restore-photo",
|
||||
] as const;
|
||||
|
||||
@@ -92,6 +92,10 @@ export const en = {
|
||||
name: "Red Eye Removal",
|
||||
description: "AI-powered red eye detection and correction for flash photos",
|
||||
},
|
||||
"restore-photo": {
|
||||
name: "Photo Restoration",
|
||||
description: "Fix scratches, tears, and damage on old photos with AI",
|
||||
},
|
||||
"content-aware-resize": {
|
||||
name: "Content-Aware Resize",
|
||||
description: "Intelligently resize images while preserving important content",
|
||||
|
||||
Reference in New Issue
Block a user