feat: SOTA AI photo restoration with multi-step pipeline (#58) (#62)

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>

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<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;

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<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" };
+    },
+  });
+}

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<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>
+  );
+}

apps/web/src/lib/suggested-tools.ts

@@ -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"],
 };

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<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 {

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()

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()

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";

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<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,
+  };
+}

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;

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",