diff --git a/docs/specs/2026-04-15-restore-photo-design.md b/docs/specs/2026-04-15-restore-photo-design.md new file mode 100644 index 00000000..fdd09c52 --- /dev/null +++ b/docs/specs/2026-04-15-restore-photo-design.md @@ -0,0 +1,207 @@ +# Restore-Photo: State-of-the-Art Redesign + +**Date:** 2026-04-15 +**Status:** Approved +**Scope:** Full rewrite of the restore-photo pipeline and settings UI + +--- + +## Overview + +Replace the current restore-photo pipeline with a single, always-best-quality pipeline backed by state-of-the-art open-source models running fully locally inside the Docker container. No presets, no quality tiers — the pipeline always runs at maximum quality. GPU is used when available; CPU is always a supported fallback. Users control *which* steps run via toggles, not *how well*. + +--- + +## Models + +All models are ONNX (or TFLite for BlazeFace) and downloaded at Docker build time via `hf_hub_download`. If a model file is missing at runtime, the affected step is skipped gracefully with a logged warning. + +| Model | Purpose | Format | Approx. Size | HuggingFace Source | +|---|---|---|---|---| +| `lama_fp32.onnx` | Inpainting fallback (when SD model absent) | ONNX fp32 | ~200 MB | Carve/LaMa-ONNX | +| `codeformer.onnx` | Face restoration (primary) | ONNX | ~170 MB | facefusion/models-3.0.0 | +| `gfpgan_1.4.onnx` | Face restoration fallback | ONNX | ~330 MB | facefusion/models-3.0.0 | +| `nafnet_denoise.onnx` | Neural denoising | ONNX | ~68 MB | Export from megvii-research/NAFNet official weights (NAFNet-SIDD-width32) | +| `ddcolor.onnx` | Colorization (shared with colorize tool) | ONNX | ~450 MB | existing | +| `ct2_colorization.onnx` | Higher-quality colorization | ONNX | ~850 MB | Export from shuchenweng/CT2 (ICCV 2023) official weights | +| `realesrgan_x2.onnx` | Super-resolution 2× (shared with upscale tool) | ONNX | ~65 MB | existing | +| `realesrgan_x4.onnx` | Super-resolution 4× (shared with upscale tool) | ONNX | ~65 MB | existing | +| `mediapipe_blazeface.tflite` | Face detection | TFLite float16 | ~2 MB | Google MediaPipe (existing) | +| `sd_inpaint_fp16.onnx` | Diffusion inpainting (primary) | ONNX fp16 | ~1.7 GB | runwayml/stable-diffusion-inpainting quantized | + +**Total new model footprint:** ~3.9 GB. Models shared with other tools are not re-downloaded. + +--- + +## Processing Pipeline + +Stages run sequentially. Each stage is a self-contained function: takes image + config, returns modified image. GPU provider priority: `CUDAExecutionProvider` → `CoreMLExecutionProvider` → `CPUExecutionProvider`. + +``` +Input image + │ + ▼ +[1] Pre-processing + - EXIF auto-orient + - Detect grayscale (is_grayscale flag) — mean channel diff < 5.0 + - Compute adaptive tile size from image resolution + │ + ▼ +[2] Damage Detection (always runs to produce mask for step 3) + - 8-angle morphological line kernels (0°, 22.5°, 45°, 67.5°, 90°, 112.5°, 135°, 157.5°) + - Laplacian-based stain/fade/water-damage detection layer (new) + - Combines top-hat + black-hat transforms at full resolution + - Output: damage_mask at full resolution + │ + ▼ +[3] Inpainting (if scratchRemoval=true AND damage coverage ≥ 0.1%) + Primary: SD-inpainting ONNX fp16 + - Overlapping tile-stitch (tile=512, stride=384) + - Poisson-blend seam correction + Fallback (if SD model absent): LaMa tiled + - tile=768, stride=512, feathered Gaussian blend + │ + ▼ +[4] Face Enhancement (if faceEnhancement=true) + - MediaPipe BlazeFace detection (unchanged) + - Skip faces < 48px + - Crop with 80% expansion for context + - CodeFormer at 512×512 (primary) + - GFPGAN at 512×512 as fallback if CodeFormer output SSIM vs. original face crop < 0.6 + - Paste-back with face-aware feathered elliptical mask + color-transfer correction + │ + ▼ +[5] Denoising (if denoise=true) + - NAFNet ONNX, tiled at full resolution + - denoiseStrength (0–100) maps to NAFNet sigma parameter + - Falls back to NLMeans if NAFNet model absent + │ + ▼ +[6] Super Resolution (if superResolution=true) + - Real-ESRGAN x4 if input < 2 MP + - Real-ESRGAN x2 if input ≥ 2 MP + - Tiled inference with seam blending (reuse upscale tool implementation) + │ + ▼ +[7] Colorization (if colorize=true) + - CT2 tiled (tile=512, stride=384) + - Falls back to DDColor tiled if CT2 model absent + - colorizeStrength (0–100) controls LAB a/b blend intensity + - Preserves original luminance channel + │ + ▼ +Output image +``` + +**Progress milestones** (stderr JSON `{"progress": N, "stage": "..."}`): + +| % | Stage | +|---|---| +| 5 | Pre-processing | +| 10 | Damage detection | +| 25 | Inpainting | +| 45 | Face enhancement | +| 60 | Denoising | +| 75 | Super resolution | +| 90 | Colorization | +| 95 | Saving output | +| 100 | Complete (set by API route) | + +--- + +## API Changes + +### Removed +- `mode: "auto" | "light" | "heavy"` — removed entirely + +### Unchanged +- `scratchRemoval: boolean` (default `true`) +- `faceEnhancement: boolean` (default `true`) +- `fidelity: number` 0–1 (default `0.7`) — CodeFormer fidelity +- `denoise: boolean` (default `true`) +- `denoiseStrength: number` 0–100 (default `40`) +- `colorize: boolean` (default `false`) + +### Added +- `superResolution: boolean` (default `true`) — allow user to disable SR if they want original dimensions +- `colorizeStrength: number` 0–100 (default `85`) — was hardcoded `0.85` internally + +### Timeout +Bridge timeout raised from 300s → 1200s to accommodate SD-inpainting on CPU for large images. + +--- + +## Frontend Changes (`restore-photo-settings.tsx`) + +- Remove 3-column mode button group (`auto / light / heavy`) +- Add **Super Resolution** checkbox (default on) +- Add **Colorize Strength** slider (0–100, step 5) — shown when `colorize` is on, same pattern as `fidelity` slider +- `denoiseStrength` slider remains; now controls NAFNet sigma rather than NLMeans h + +No other UI changes. Settings panel stays compact. + +--- + +## i18n Additions (`packages/shared/src/i18n/en.ts`) + +```ts +superResolution: "Super Resolution", +colorizeStrength: "Colorize Strength", +``` + +Remove: `restoreMode`, `restoreModeAuto`, `restoreModeLight`, `restoreModeHeavy` (if present). + +--- + +## Error Handling + +| Scenario | Behavior | +|---|---| +| SD model missing | Fall back to LaMa tiled; log warning | +| CT2 model missing | Fall back to DDColor tiled; log warning | +| NAFNet model missing | Fall back to NLMeans; log warning | +| GFPGAN model missing | Skip GFPGAN fallback; CodeFormer only | +| No damage detected (< 0.1% coverage) | Skip inpainting entirely | +| Tiling dimension mismatch | Fall back to 512px single-pass for that stage; log warning | +| CUDA unavailable | Silently use CPU provider | +| SD-inpainting CPU timeout | Bridge timeout is 1200s; if exceeded, API returns 504 with clear message | + +--- + +## Testing + +### Unit (Vitest) +- API Zod schema accepts `superResolution` and `colorizeStrength` +- API schema rejects old `mode` field +- `colorizeStrength` default is 85 + +### Integration (Vitest) +- POST restore-photo with fixture image → assert output dimensions ≥ input (SR ran) +- POST with `superResolution: false` → assert output dimensions equal input +- POST with `colorize: true` → assert output is not identical to input + +### Python (pytest) +- Each pipeline stage tested in isolation with small fixture image: + - Damage detection returns a binary mask with same H×W as input + - LaMa tiled returns same-shape image + - NAFNet tiled returns same-shape image + - Real-ESRGAN x2 returns 2× dimensions + - CT2 tiled returns same-shape image + +### E2E (Playwright) +- Upload fixture old photo → verify before-after slider renders with non-blank result +- Verify Super Resolution checkbox appears in settings panel + +--- + +## File Paths Affected + +| File | Change | +|---|---| +| `packages/ai/python/restore.py` | Full rewrite of pipeline | +| `packages/ai/python/download_models.py` | Add NAFNet, CT2, SD-inpainting, GFPGAN download entries | +| `packages/ai/src/restoration.ts` | Update `RestoreOptions` interface | +| `apps/api/src/routes/tools/restore-photo.ts` | Update Zod schema, raise timeout | +| `apps/web/src/components/tools/restore-photo-settings.tsx` | Remove mode group, add SR + colorizeStrength controls | +| `packages/shared/src/i18n/en.ts` | Add/remove i18n keys | +| `packages/shared/src/constants.ts` | Update tool timeout constant if applicable |