diff --git a/apps/api/src/routes/tools/enhance-faces.ts b/apps/api/src/routes/tools/enhance-faces.ts new file mode 100644 index 00000000..50462761 --- /dev/null +++ b/apps/api/src/routes/tools/enhance-faces.ts @@ -0,0 +1,174 @@ +import { randomUUID } from "node:crypto"; +import { writeFile } from "node:fs/promises"; +import { basename, join } from "node:path"; +import { enhanceFaces } 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 { createWorkspace } from "../../lib/workspace.js"; +import { updateSingleFileProgress } from "../progress.js"; +import { registerToolProcessFn } from "../tool-factory.js"; + +/** Face enhancement route using GFPGAN/CodeFormer. */ +export function registerEnhanceFaces(app: FastifyInstance) { + app.post("/api/v1/tools/enhance-faces", 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 = settingsRaw ? JSON.parse(settingsRaw) : {}; + const model = settings.model || "auto"; + const strength = Number(settings.strength) || 0.8; + const onlyCenterFace = Boolean(settings.onlyCenterFace); + const sensitivity = Number(settings.sensitivity) || 0.5; + request.log.info( + { toolId: "enhance-faces", imageSize: fileBuffer.length, model, strength }, + "Starting face enhancement", + ); + + // Decode HEIC/HEIF input via system decoder + if (validation.format === "heif") { + fileBuffer = await decodeHeic(fileBuffer); + } + + // Auto-orient to fix EXIF rotation before face detection + fileBuffer = await autoOrient(fileBuffer); + + const jobId = randomUUID(); + const workspacePath = await createWorkspace(jobId); + + // Save input + const inputPath = join(workspacePath, "input", filename); + await writeFile(inputPath, fileBuffer); + + // Process + const jobIdForProgress = clientJobId; + const onProgress = jobIdForProgress + ? (percent: number, stage: string) => { + updateSingleFileProgress({ + jobId: jobIdForProgress, + phase: "processing", + stage, + percent, + }); + } + : undefined; + + const result = await enhanceFaces( + fileBuffer, + join(workspacePath, "output"), + { model, strength, onlyCenterFace, sensitivity }, + onProgress, + ); + + // Save output + const outputFilename = `${filename.replace(/\.[^.]+$/, "")}_enhanced.png`; + const outputPath = join(workspacePath, "output", outputFilename); + await writeFile(outputPath, result.buffer); + + // Generate webp preview for the frontend + let previewUrl: string | undefined; + try { + const previewBuffer = await sharp(result.buffer).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 - frontend will show fallback + } + + 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: result.buffer.length, + facesDetected: result.facesDetected, + faces: result.faces, + model: result.model, + }); + } catch (err) { + request.log.error({ err, toolId: "enhance-faces" }, "Face enhancement failed"); + return reply.status(422).send({ + error: "Face enhancement failed", + details: err instanceof Error ? err.message : "Unknown error", + }); + } + }); + + // Register in the pipeline/batch registry so this tool can be used + // as a step in automation pipelines (without progress callbacks). + registerToolProcessFn({ + toolId: "enhance-faces", + settingsSchema: z.object({ + model: z.enum(["auto", "gfpgan", "codeformer"]).default("auto"), + strength: z.number().min(0).max(1).default(0.8), + onlyCenterFace: z.boolean().default(false), + sensitivity: z.number().min(0).max(1).default(0.5), + }), + process: async (inputBuffer, settings, filename) => { + const s = settings as { + model?: "auto" | "gfpgan" | "codeformer"; + strength?: number; + onlyCenterFace?: boolean; + sensitivity?: number; + }; + const orientedBuffer = await autoOrient(inputBuffer); + const jobId = randomUUID(); + const workspacePath = await createWorkspace(jobId); + const result = await enhanceFaces(orientedBuffer, join(workspacePath, "output"), { + model: s.model ?? "auto", + strength: s.strength ?? 0.8, + onlyCenterFace: s.onlyCenterFace ?? false, + sensitivity: s.sensitivity ?? 0.5, + }); + const outputFilename = `${filename.replace(/\.[^.]+$/, "")}_enhanced.png`; + return { buffer: result.buffer, filename: outputFilename, contentType: "image/png" }; + }, + }); +} diff --git a/apps/api/src/routes/tools/index.ts b/apps/api/src/routes/tools/index.ts index 0aad6fd8..cdc5601a 100644 --- a/apps/api/src/routes/tools/index.ts +++ b/apps/api/src/routes/tools/index.ts @@ -17,6 +17,7 @@ import { registerContentAwareResize } from "./content-aware-resize.js"; import { registerConvert } from "./convert.js"; import { registerCrop } from "./crop.js"; import { registerEditMetadata } from "./edit-metadata.js"; +import { registerEnhanceFaces } from "./enhance-faces.js"; import { registerEraseObject } from "./erase-object.js"; import { registerFavicon } from "./favicon.js"; import { registerFindDuplicates } from "./find-duplicates.js"; @@ -134,6 +135,7 @@ export async function registerToolRoutes(app: FastifyInstance): Promise { { id: "image-enhancement", register: registerImageEnhancement }, { id: "content-aware-resize", register: registerContentAwareResize }, { id: "colorize", register: registerColorize }, + { id: "enhance-faces", register: registerEnhanceFaces }, { id: "noise-removal", register: registerNoiseRemoval }, { id: "red-eye-removal", register: registerRedEyeRemoval }, ]; diff --git a/apps/web/src/components/tools/enhance-faces-settings.tsx b/apps/web/src/components/tools/enhance-faces-settings.tsx new file mode 100644 index 00000000..e2a9ee30 --- /dev/null +++ b/apps/web/src/components/tools/enhance-faces-settings.tsx @@ -0,0 +1,220 @@ +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"; + +const MODEL_OPTIONS = [ + { value: "gfpgan", label: "Fast" }, + { value: "auto", label: "Balanced" }, + { value: "codeformer", label: "Best" }, +] as const; + +export interface EnhanceFacesControlsProps { + settings?: Record; + onChange?: (settings: Record) => void; +} + +export function EnhanceFacesControls({ + settings: initialSettings, + onChange, +}: EnhanceFacesControlsProps) { + const [model, setModel] = useState<"gfpgan" | "auto" | "codeformer">("auto"); + const [strength, setStrength] = useState(80); + const [onlyCenterFace, setOnlyCenterFace] = useState(false); + const [sensitivity, setSensitivity] = useState(50); + + const initializedRef = useRef(false); + useEffect(() => { + if (!initialSettings || initializedRef.current) return; + initializedRef.current = true; + if (initialSettings.model != null) + setModel(initialSettings.model as "gfpgan" | "auto" | "codeformer"); + if (initialSettings.strength != null) setStrength(Number(initialSettings.strength) * 100); + if (initialSettings.onlyCenterFace != null) + setOnlyCenterFace(Boolean(initialSettings.onlyCenterFace)); + if (initialSettings.sensitivity != null) + setSensitivity(Number(initialSettings.sensitivity) * 100); + }, [initialSettings]); + + const onChangeRef = useRef(onChange); + useEffect(() => { + onChangeRef.current = onChange; + }); + + useEffect(() => { + onChangeRef.current?.({ + model, + strength: strength / 100, + onlyCenterFace, + sensitivity: sensitivity / 100, + }); + }, [model, strength, onlyCenterFace, sensitivity]); + + return ( +
+ {/* Quality */} +
+

Quality

+
+ {MODEL_OPTIONS.map(({ value, label }) => ( + + ))} +
+
+ + {/* Enhancement Strength */} +
+
+ + {strength}% +
+ setStrength(Number(e.target.value))} + className="w-full mt-1" + /> +
+ Subtle + Maximum +
+
+ + {/* Only enhance main face (only works with GFPGAN / Fast mode) */} + {model !== "codeformer" && ( +
+ +

+ For portraits - ignores background faces +

+
+ )} + + {/* Detection Sensitivity */} +
+
+ + {sensitivity}% +
+ setSensitivity(Number(e.target.value))} + className="w-full mt-1" + /> +
+ Fewer faces + More faces +
+
+
+ ); +} + +export function EnhanceFacesSettings() { + const { files } = useFileStore(); + const { + processFiles, + processAllFiles, + processing, + error, + downloadUrl, + originalSize, + processedSize, + progress, + } = useToolProcessor("enhance-faces"); + 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

+

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

+
+ )} + + {/* Process buttons / progress */} + {processing ? ( + + ) : ( + + )} + + {/* Download (single file - batch uses Download All ZIP in tool-page) */} + {!hasMultiple && downloadUrl && ( + + + Download + + )} +
+ ); +} diff --git a/apps/web/src/components/tools/pipeline-step-settings.tsx b/apps/web/src/components/tools/pipeline-step-settings.tsx index 16dda3e4..0675dda4 100644 --- a/apps/web/src/components/tools/pipeline-step-settings.tsx +++ b/apps/web/src/components/tools/pipeline-step-settings.tsx @@ -4,6 +4,7 @@ import { ColorControls } from "./color-settings"; import { CompressControls } from "./compress-settings"; import { ConvertControls } from "./convert-settings"; import { CropControls } from "./crop-settings"; +import { EnhanceFacesControls } from "./enhance-faces-settings"; import { GifToolsControls } from "./gif-tools-settings"; import { NoiseRemovalControls } from "./noise-removal-settings"; import { RemoveBgControls } from "./remove-bg-settings"; @@ -43,6 +44,8 @@ export function PipelineStepSettings({ toolId, settings, onChange }: PipelineSte if (toolId === "gif-tools") return ; if (toolId === "upscale") return ; if (toolId === "blur-faces") return ; + if (toolId === "enhance-faces") + return ; if (toolId === "remove-background") return ; if (toolId === "noise-removal") diff --git a/apps/web/src/lib/tool-registry.tsx b/apps/web/src/lib/tool-registry.tsx index 0c41e358..02c9b1a8 100644 --- a/apps/web/src/lib/tool-registry.tsx +++ b/apps/web/src/lib/tool-registry.tsx @@ -229,6 +229,11 @@ const BlurFacesSettings = lazy(() => default: m.BlurFacesSettings, })), ); +const EnhanceFacesSettings = lazy(() => + import("@/components/tools/enhance-faces-settings").then((m) => ({ + default: m.EnhanceFacesSettings, + })), +); const EraseObjectSettings = lazy(() => import("@/components/tools/erase-object-settings").then((m) => ({ default: m.EraseObjectSettings, @@ -371,6 +376,7 @@ export const toolRegistry = new Map([ ["upscale", { displayMode: "before-after", Settings: UpscaleSettings }], ["ocr", { displayMode: "before-after", Settings: OcrSettings }], ["blur-faces", { displayMode: "before-after", Settings: BlurFacesSettings }], + ["enhance-faces", { displayMode: "before-after", Settings: EnhanceFacesSettings }], [ "erase-object", { diff --git a/docker/Dockerfile b/docker/Dockerfile index d83d946d..fefa0a18 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -137,6 +137,12 @@ RUN if [ "$TARGETARCH" = "amd64" ]; then \ /opt/venv/bin/pip install mediapipe==0.10.18 \ ; fi +# CodeFormer face enhancement (install with --no-deps to avoid numpy 2.x conflict) +RUN /opt/venv/bin/pip install --no-deps codeformer-pip==0.0.4 lpips + +# Re-pin numpy to 1.26.4 in case any transitive dep upgraded it +RUN /opt/venv/bin/pip install numpy==1.26.4 + # Pre-download and verify all ML models # Note: on amd64, paddlepaddle-gpu can't import without the CUDA driver (only # available at runtime). The download script gracefully skips PaddleOCR model diff --git a/docker/download_models.py b/docker/download_models.py index 5f9d7de9..fad29169 100644 --- a/docker/download_models.py +++ b/docker/download_models.py @@ -32,6 +32,13 @@ GFPGAN_MODEL_URL = ( GFPGAN_MODEL_PATH = os.path.join(GFPGAN_MODEL_DIR, "GFPGANv1.3.pth") GFPGAN_MIN_SIZE = 300_000_000 # ~332 MB +CODEFORMER_MODEL_DIR = "/opt/models/codeformer" +CODEFORMER_MODEL_URL = ( + "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth" +) +CODEFORMER_MODEL_PATH = os.path.join(CODEFORMER_MODEL_DIR, "codeformer.pth") +CODEFORMER_MIN_SIZE = 350_000_000 # ~375 MB + DDCOLOR_MODEL_DIR = "/opt/models/ddcolor" DDCOLOR_MODEL_URL = ( "https://huggingface.co/piddnad/DDColor-models/resolve/main/ddcolor_paper_tiny.pth" @@ -167,6 +174,20 @@ def download_gfpgan_model(): print(f" GFPGANv1.3.pth downloaded ({size / 1_000_000:.1f} MB)\n") +def download_codeformer_model(): + """Download codeformer.pth pretrained weights for face enhancement.""" + print("=== Downloading CodeFormer model ===") + os.makedirs(CODEFORMER_MODEL_DIR, exist_ok=True) + print(f" Downloading from {CODEFORMER_MODEL_URL}...") + urllib.request.urlretrieve(CODEFORMER_MODEL_URL, CODEFORMER_MODEL_PATH) + + size = os.path.getsize(CODEFORMER_MODEL_PATH) + assert size > CODEFORMER_MIN_SIZE, ( + f"CodeFormer model too small: {size} bytes (expected > {CODEFORMER_MIN_SIZE})" + ) + print(f" codeformer.pth downloaded ({size / 1_000_000:.1f} MB)\n") + + def download_ddcolor_model(): """Download pre-exported DDColor ONNX model for AI photo colorization. @@ -319,6 +340,15 @@ def smoke_test(): ) print(" GFPGAN model file verified") + # CodeFormer model file must exist + assert os.path.exists(CODEFORMER_MODEL_PATH), ( + f"CodeFormer model missing: {CODEFORMER_MODEL_PATH}" + ) + assert os.path.getsize(CODEFORMER_MODEL_PATH) > CODEFORMER_MIN_SIZE, ( + "CodeFormer model file is too small" + ) + print(" CodeFormer model file verified") + # DDColor ONNX model must exist assert os.path.exists(DDCOLOR_ONNX_PATH), ( f"DDColor model missing: {DDCOLOR_ONNX_PATH}" @@ -360,6 +390,7 @@ def main(): download_rembg_models() download_realesrgan_model() download_gfpgan_model() + download_codeformer_model() download_ddcolor_model() download_paddleocr_models() download_paddleocr_vl_model() diff --git a/packages/ai/python/enhance_faces.py b/packages/ai/python/enhance_faces.py new file mode 100644 index 00000000..004ff2e1 --- /dev/null +++ b/packages/ai/python/enhance_faces.py @@ -0,0 +1,275 @@ +"""Face enhancement using GFPGAN or CodeFormer with MediaPipe detection.""" +import sys +import json +import os + +# Patch for basicsr compatibility with torchvision >= 0.18. +# torchvision removed transforms.functional_tensor, merging it into +# transforms.functional. basicsr still imports the old path, so we +# create a shim module to redirect the import. +try: + import torchvision.transforms.functional_tensor # noqa: F401 +except (ImportError, ModuleNotFoundError): + try: + import types + import torchvision.transforms.functional as _F + + _shim = types.ModuleType("torchvision.transforms.functional_tensor") + _shim.rgb_to_grayscale = _F.rgb_to_grayscale + sys.modules["torchvision.transforms.functional_tensor"] = _shim + except ImportError: + pass # torchvision not installed at all + + +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) + + +GFPGAN_MODEL_PATH = os.environ.get( + "GFPGAN_MODEL_PATH", + "/opt/models/gfpgan/GFPGANv1.3.pth", +) + +CODEFORMER_MODEL_PATH = os.environ.get( + "CODEFORMER_MODEL_PATH", + "/opt/models/codeformer/codeformer.pth", +) + + +def detect_faces_mediapipe(img_array, sensitivity): + """Detect faces using MediaPipe with dual-model approach. + + Returns a list of {x, y, w, h} dicts for each detected face. + """ + import mediapipe as mp + + min_confidence = max(0.1, 1.0 - sensitivity) + mp_face = mp.solutions.face_detection + + # Try short-range model first (model_selection=0, best for faces + # within ~2m which covers most photos), then fall back to + # full-range model (model_selection=1) for distant/group shots. + detections = [] + for model_sel in [0, 1]: + detector = mp_face.FaceDetection( + model_selection=model_sel, + min_detection_confidence=min_confidence, + ) + results = detector.process(img_array) + detector.close() + if results.detections: + detections = results.detections + break + + if not detections: + return [] + + ih, iw = img_array.shape[:2] + faces = [] + for detection in detections: + bbox = detection.location_data.relative_bounding_box + x = int(bbox.xmin * iw) + y = int(bbox.ymin * ih) + w = int(bbox.width * iw) + h = int(bbox.height * ih) + faces.append({"x": x, "y": y, "w": w, "h": h}) + + return faces + + +def enhance_with_gfpgan(img_array, only_center_face): + """Enhance faces using GFPGAN. Returns the enhanced image array.""" + from gfpgan import GFPGANer + + if not os.path.exists(GFPGAN_MODEL_PATH): + raise FileNotFoundError(f"GFPGAN model not found: {GFPGAN_MODEL_PATH}") + + enhancer = GFPGANer( + model_path=GFPGAN_MODEL_PATH, + upscale=1, + arch="clean", + channel_multiplier=2, + bg_upsampler=None, + ) + _, _, output = enhancer.enhance( + img_array, + has_aligned=False, + only_center_face=only_center_face, + paste_back=True, + ) + return output + + +def enhance_with_codeformer(img_array, fidelity_weight): + """Enhance faces using CodeFormer via codeformer-pip. + + The codeformer-pip package provides inference_app() which handles + face detection, alignment, restoration, and paste-back internally. + fidelity_weight controls quality vs fidelity (0 = quality, 1 = fidelity). + + NOTE: codeformer-pip's app.py runs heavy module-level initialization + (model downloads, GPU setup) on import. The Docker image must place + model weights where the package expects them, or set environment + variables so the download step succeeds. If the import or inference + fails, the auto model selection will fall back to GFPGAN. + """ + import numpy as np + + # Import may fail if codeformer-pip is not installed or if the + # module-level model loading fails (missing weights, no GPU, etc.) + from codeformer.app import inference_app + + # inference_app accepts a numpy array (BGR) or file path. + # It returns the restored image as a BGR numpy array. + # We pass our RGB array converted to BGR since OpenCV convention is used internally. + img_bgr = img_array[:, :, ::-1].copy() + restored_bgr = inference_app( + image=img_bgr, + background_enhance=False, + face_upsample=False, + upscale=1, + codeformer_fidelity=fidelity_weight, + ) + # Convert back to RGB + restored_rgb = restored_bgr[:, :, ::-1].copy() + return restored_rgb + + +def main(): + input_path = sys.argv[1] + output_path = sys.argv[2] + settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {} + + model_choice = settings.get("model", "auto") + strength = float(settings.get("strength", 0.8)) + only_center_face = settings.get("onlyCenterFace", False) + sensitivity = float(settings.get("sensitivity", 0.5)) + + try: + emit_progress(10, "Preparing") + from PIL import Image + import numpy as np + + img = Image.open(input_path).convert("RGB") + img_array = np.array(img) + + # Detect faces with MediaPipe + try: + emit_progress(20, "Scanning for faces") + faces = detect_faces_mediapipe(img_array, sensitivity) + except ImportError: + print( + json.dumps( + { + "success": False, + "error": "Face detection requires MediaPipe. Install with: pip install mediapipe", + } + ) + ) + sys.exit(1) + + num_faces = len(faces) + emit_progress(30, f"Found {num_faces} face{'s' if num_faces != 1 else ''}") + + # No faces found - save original unchanged + if num_faces == 0: + img.save(output_path) + print( + json.dumps( + { + "success": True, + "facesDetected": 0, + "faces": [], + "model": "none", + } + ) + ) + return + + emit_progress(40, "Loading AI model") + + # Redirect stdout to stderr for the ENTIRE AI pipeline. + # Libraries like basicsr, gfpgan, and torch print download + # progress and init messages to stdout which would corrupt + # our JSON result. + stdout_fd = os.dup(1) + os.dup2(2, 1) + + enhanced = None + model_used = None + + try: + if model_choice == "gfpgan": + enhanced = enhance_with_gfpgan(img_array, only_center_face) + model_used = "gfpgan" + + elif model_choice == "codeformer": + fidelity_weight = 1.0 - strength + enhanced = enhance_with_codeformer(img_array, fidelity_weight) + model_used = "codeformer" + + elif model_choice == "auto": + # Try CodeFormer first, fall back to GFPGAN. + # Catch broad Exception because codeformer-pip can fail in + # unexpected ways (AttributeError, TypeError, etc.) + try: + fidelity_weight = 1.0 - strength + enhanced = enhance_with_codeformer(img_array, fidelity_weight) + model_used = "codeformer" + except Exception: + enhanced = enhance_with_gfpgan(img_array, only_center_face) + model_used = "gfpgan" + + finally: + # Restore stdout after ALL AI processing + os.dup2(stdout_fd, 1) + os.close(stdout_fd) + + if enhanced is None: + raise RuntimeError("Face enhancement failed: no model available") + + emit_progress(85, "Enhancement complete") + + # Alpha blend result with original based on strength. + # For CodeFormer, strength is already applied via fidelity_weight, + # so skip the blend to avoid double-applying. + # For GFPGAN (which has no fidelity knob), blend with original. + if strength < 1.0 and model_used != "codeformer": + blended = ( + img_array.astype(np.float32) * (1.0 - strength) + + enhanced.astype(np.float32) * strength + ) + enhanced = np.clip(blended, 0, 255).astype(np.uint8) + + emit_progress(95, "Saving result") + Image.fromarray(enhanced).save(output_path) + + print( + json.dumps( + { + "success": True, + "facesDetected": num_faces, + "faces": faces, + "model": model_used, + } + ) + ) + + except ImportError: + print( + json.dumps( + { + "success": False, + "error": "Pillow is not installed. Install with: pip install Pillow", + } + ) + ) + sys.exit(1) + except Exception as e: + print(json.dumps({"success": False, "error": str(e)})) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/packages/ai/python/requirements-gpu.txt b/packages/ai/python/requirements-gpu.txt index 95fb1199..2f273bcb 100644 --- a/packages/ai/python/requirements-gpu.txt +++ b/packages/ai/python/requirements-gpu.txt @@ -7,3 +7,4 @@ onnxruntime-gpu==1.20.1 numpy==1.26.4 Pillow==11.1.0 opencv-python-headless==4.10.0.84 +codeformer-pip==0.0.4 diff --git a/packages/ai/python/requirements.txt b/packages/ai/python/requirements.txt index d449c378..ac7007c3 100644 --- a/packages/ai/python/requirements.txt +++ b/packages/ai/python/requirements.txt @@ -7,3 +7,4 @@ onnxruntime==1.20.1 numpy==1.26.4 Pillow==11.1.0 opencv-python-headless==4.10.0.84 +codeformer-pip==0.0.4 diff --git a/packages/ai/src/face-enhancement.ts b/packages/ai/src/face-enhancement.ts new file mode 100644 index 00000000..2e743894 --- /dev/null +++ b/packages/ai/src/face-enhancement.ts @@ -0,0 +1,47 @@ +import { readFile, writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import { type ProgressCallback, runPythonWithProgress } from "./bridge.js"; + +export interface EnhanceFacesOptions { + model?: "auto" | "gfpgan" | "codeformer"; + strength?: number; + onlyCenterFace?: boolean; + sensitivity?: number; +} + +export interface EnhanceFacesResult { + buffer: Buffer; + facesDetected: number; + faces: Array<{ x: number; y: number; w: number; h: number }>; + model: string; +} + +export async function enhanceFaces( + inputBuffer: Buffer, + outputDir: string, + options: EnhanceFacesOptions = {}, + onProgress?: ProgressCallback, +): Promise { + const inputPath = join(outputDir, "input_enhance_faces.png"); + const outputPath = join(outputDir, "output_enhance_faces.png"); + + await writeFile(inputPath, inputBuffer); + const { stdout } = await runPythonWithProgress( + "enhance_faces.py", + [inputPath, outputPath, JSON.stringify(options)], + { onProgress }, + ); + + const result = JSON.parse(stdout); + if (!result.success) { + throw new Error(result.error || "Face enhancement failed"); + } + + const buffer = await readFile(outputPath); + return { + buffer, + facesDetected: result.facesDetected, + faces: result.faces ?? [], + model: result.model ?? "unknown", + }; +} diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 9dfd65ce..5594ac47 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -3,6 +3,7 @@ export { isGpuAvailable, shutdownDispatcher } from "./bridge.js"; export { colorize } from "./colorization.js"; export type { DetectFacesResult, FaceRegion } from "./face-detection.js"; export { blurFaces, detectFaces } from "./face-detection.js"; +export { enhanceFaces } from "./face-enhancement.js"; export { inpaint } from "./inpainting.js"; export { noiseRemoval } from "./noise-removal.js"; export { extractText } from "./ocr.js"; diff --git a/packages/shared/src/constants.ts b/packages/shared/src/constants.ts index 612fcd98..22e92ec9 100644 --- a/packages/shared/src/constants.ts +++ b/packages/shared/src/constants.ts @@ -177,6 +177,14 @@ export const TOOLS: Tool[] = [ icon: "Sparkles", route: "/image-enhancement", }, + { + id: "enhance-faces", + name: "Face Enhancement", + description: "Restore and enhance faces with AI", + category: "ai", + icon: "ScanFace", + route: "/enhance-faces", + }, { id: "colorize", name: "AI Colorization", @@ -421,6 +429,7 @@ export const PYTHON_SIDECAR_TOOLS = [ "erase-object", "ocr", "colorize", + "enhance-faces", "noise-removal", "red-eye-removal", ] as const; diff --git a/packages/shared/src/i18n/en.ts b/packages/shared/src/i18n/en.ts index 30256a7b..16d6c574 100644 --- a/packages/shared/src/i18n/en.ts +++ b/packages/shared/src/i18n/en.ts @@ -71,6 +71,10 @@ export const en = { name: "Face / PII Blur", description: "Auto-detect and blur faces and sensitive info", }, + "enhance-faces": { + name: "Face Enhancement", + description: "Restore and enhance faces with AI", + }, "smart-crop": { name: "Smart Crop", description: "Smart subject, face, or trim-based cropping",