diff --git a/apps/api/src/routes/tools/index.ts b/apps/api/src/routes/tools/index.ts index cdc5601a..a3651a34 100644 --- a/apps/api/src/routes/tools/index.ts +++ b/apps/api/src/routes/tools/index.ts @@ -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 { { id: "enhance-faces", register: registerEnhanceFaces }, { id: "noise-removal", register: registerNoiseRemoval }, { id: "red-eye-removal", register: registerRedEyeRemoval }, + { id: "restore-photo", register: registerRestorePhoto }, ]; let skipped = 0; diff --git a/apps/api/src/routes/tools/restore-photo.ts b/apps/api/src/routes/tools/restore-photo.ts new file mode 100644 index 00000000..f5c8fce4 --- /dev/null +++ b/apps/api/src/routes/tools/restore-photo.ts @@ -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; + 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" }; + }, + }); +} diff --git a/apps/web/src/components/tools/restore-photo-settings.tsx b/apps/web/src/components/tools/restore-photo-settings.tsx new file mode 100644 index 00000000..172c8860 --- /dev/null +++ b/apps/web/src/components/tools/restore-photo-settings.tsx @@ -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; + onChange?: (settings: Record) => void; +} + +export function RestorePhotoControls({ + settings: initialSettings, + onChange, +}: RestorePhotoControlsProps) { + const [mode, setMode] = useState("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 ( +
+ {/* Mode selector */} +
+

Restoration Mode

+
+ {MODES.map((m) => ( + + ))} +
+ {activeMode &&

{activeMode.desc}

} +
+ +
+ + {/* Feature toggles */} +
+ {/* Scratch Removal */} + + + {/* Face Enhancement */} + + + {/* Fidelity slider (only when face enhancement is on) */} + {faceEnhancement && ( +
+
+

Face Fidelity

+ {fidelity}% +
+ setFidelity(Number(e.target.value))} + className="w-full h-1.5 rounded-full appearance-none bg-muted accent-primary" + /> +
+ Enhanced + Faithful +
+
+ )} + + {/* Noise Reduction */} + + + {/* Denoise strength slider */} + {denoise && ( +
+
+

Denoise Strength

+ {denoiseStrength} +
+ setDenoiseStrength(Number(e.target.value))} + className="w-full h-1.5 rounded-full appearance-none bg-muted accent-primary" + /> +
+ Subtle + Strong +
+
+ )} + +
+ + {/* Auto-Colorize */} + +
+
+ ); +} + +export function RestorePhotoSettings() { + const { files } = useFileStore(); + const { + processFiles, + processAllFiles, + processing, + error, + downloadUrl, + originalSize, + processedSize, + progress, + } = useToolProcessor("restore-photo"); + const [settings, setSettings] = useState>({}); + + const handleProcess = () => { + if (files.length > 1) { + processAllFiles(files, settings); + } else { + processFiles(files, settings); + } + }; + + const hasFile = files.length > 0; + const hasMultiple = files.length > 1; + + return ( +
+ + + {/* Error */} + {error &&

{error}

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

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

+

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

+
+ )} + + {/* Process button / progress */} + {processing ? ( + + ) : ( + + )} + + {/* Download (single file) */} + {!hasMultiple && downloadUrl && ( + + + Download + + )} +
+ ); +} diff --git a/apps/web/src/lib/suggested-tools.ts b/apps/web/src/lib/suggested-tools.ts index 74917197..44b6e86d 100644 --- a/apps/web/src/lib/suggested-tools.ts +++ b/apps/web/src/lib/suggested-tools.ts @@ -15,6 +15,7 @@ const TOOL_SUGGESTIONS: Record = { "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"], }; diff --git a/apps/web/src/lib/tool-registry.tsx b/apps/web/src/lib/tool-registry.tsx index 02c9b1a8..31aebe64 100644 --- a/apps/web/src/lib/tool-registry.tsx +++ b/apps/web/src/lib/tool-registry.tsx @@ -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([ ["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 { diff --git a/docker/download_models.py b/docker/download_models.py index fad29169..cfa2af6c 100644 --- a/docker/download_models.py +++ b/docker/download_models.py @@ -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() diff --git a/packages/ai/python/restore.py b/packages/ai/python/restore.py new file mode 100644 index 00000000..eb9d6f19 --- /dev/null +++ b/packages/ai/python/restore.py @@ -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() diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 5594ac47..a5174788 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -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"; diff --git a/packages/ai/src/restoration.ts b/packages/ai/src/restoration.ts new file mode 100644 index 00000000..74568215 --- /dev/null +++ b/packages/ai/src/restoration.ts @@ -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 { + 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, + }; +} diff --git a/packages/shared/src/constants.ts b/packages/shared/src/constants.ts index 22e92ec9..63b303e3 100644 --- a/packages/shared/src/constants.ts +++ b/packages/shared/src/constants.ts @@ -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; diff --git a/packages/shared/src/i18n/en.ts b/packages/shared/src/i18n/en.ts index 16d6c574..7327dcd1 100644 --- a/packages/shared/src/i18n/en.ts +++ b/packages/shared/src/i18n/en.ts @@ -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",