diff --git a/apps/api/src/routes/tools/colorize.ts b/apps/api/src/routes/tools/colorize.ts new file mode 100644 index 00000000..72bb6136 --- /dev/null +++ b/apps/api/src/routes/tools/colorize.ts @@ -0,0 +1,183 @@ +import { randomUUID } from "node:crypto"; +import { writeFile } from "node:fs/promises"; +import { basename, join } from "node:path"; +import { colorize } 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"; + +/** + * AI photo colorization route. + * Converts B&W / grayscale photos to full color using DDColor, + * with OpenCV DNN fallback. + */ +export function registerColorize(app: FastifyInstance) { + app.post("/api/v1/tools/colorize", 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 intensity = Math.min(1, Math.max(0, Number(settings.intensity) || 1.0)); + const model = settings.model || "auto"; + + request.log.info( + { toolId: "colorize", imageSize: fileBuffer.length, intensity, model }, + "Starting colorization", + ); + + // 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 colorize( + fileBuffer, + join(workspacePath, "output"), + { intensity, model }, + 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(/\.[^.]+$/, "")}_colorized.${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, + method: result.method, + }); + } catch (err) { + request.log.error({ err, toolId: "colorize" }, "Colorization failed"); + return reply.status(422).send({ + error: "Colorization failed", + details: err instanceof Error ? err.message : "Unknown error", + }); + } + }); + + // Register in the pipeline/batch registry + registerToolProcessFn({ + toolId: "colorize", + settingsSchema: z.object({ + intensity: z.number().min(0).max(1).default(1.0), + model: z.enum(["auto", "ddcolor", "opencv"]).default("auto"), + }), + process: async (inputBuffer, settings, filename) => { + const orientedBuffer = await autoOrient(inputBuffer); + const jobId = randomUUID(); + const workspacePath = await createWorkspace(jobId); + const result = await colorize(orientedBuffer, join(workspacePath, "output"), { + intensity: (settings as { intensity?: number }).intensity ?? 1.0, + model: (settings as { model?: string }).model ?? "auto", + }); + const outputFilename = `${filename.replace(/\.[^.]+$/, "")}_colorized.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 6dc90130..cc96952e 100644 --- a/apps/api/src/routes/tools/index.ts +++ b/apps/api/src/routes/tools/index.ts @@ -9,6 +9,7 @@ import { registerBulkRename } from "./bulk-rename.js"; import { registerCollage } from "./collage.js"; import { registerColorAdjustments } from "./color-adjustments.js"; import { registerColorPalette } from "./color-palette.js"; +import { registerColorize } from "./colorize.js"; import { registerCompare } from "./compare.js"; import { registerCompose } from "./compose.js"; import { registerCompress } from "./compress.js"; @@ -130,6 +131,7 @@ export async function registerToolRoutes(app: FastifyInstance): Promise { { id: "smart-crop", register: registerSmartCrop }, { id: "image-enhancement", register: registerImageEnhancement }, { id: "content-aware-resize", register: registerContentAwareResize }, + { id: "colorize", register: registerColorize }, ]; let skipped = 0; diff --git a/apps/web/src/components/tools/colorize-settings.tsx b/apps/web/src/components/tools/colorize-settings.tsx new file mode 100644 index 00000000..3d659647 --- /dev/null +++ b/apps/web/src/components/tools/colorize-settings.tsx @@ -0,0 +1,154 @@ +import { Download } from "lucide-react"; +import { useState } from "react"; +import { ProgressCard } from "@/components/common/progress-card"; +import { useToolProcessor } from "@/hooks/use-tool-processor"; +import { useFileStore } from "@/stores/file-store"; + +type Model = "auto" | "ddcolor" | "opencv"; + +const MODEL_OPTIONS: { value: Model; label: string; desc: string }[] = [ + { value: "auto", label: "Auto", desc: "Best available" }, + { value: "ddcolor", label: "DDColor", desc: "SOTA deep learning" }, + { value: "opencv", label: "Classic", desc: "Fast, lightweight" }, +]; + +export function ColorizeSettings() { + const { files } = useFileStore(); + const { + processFiles, + processAllFiles, + processing, + error, + downloadUrl, + originalSize, + processedSize, + progress, + } = useToolProcessor("colorize"); + + const [model, setModel] = useState("auto"); + const [intensity, setIntensity] = useState(100); + + const hasFile = files.length > 0; + const hasMultiple = files.length > 1; + + const handleProcess = () => { + const settings = { + model, + intensity: intensity / 100, + }; + if (hasMultiple) { + processAllFiles(files, settings); + } else { + processFiles(files, settings); + } + }; + + const handleSubmit = (e: React.FormEvent) => { + e.preventDefault(); + if (hasFile && !processing) handleProcess(); + }; + + return ( +
+ {/* Model selector */} + AI Model +
+ {MODEL_OPTIONS.map((opt) => ( + + ))} +
+ + {/* Color intensity */} + Color Intensity +
+
+ + {intensity === 0 + ? "Grayscale" + : intensity < 50 + ? "Subtle" + : intensity < 80 + ? "Natural" + : "Vivid"} + + + {intensity}% + +
+ setIntensity(Number(e.target.value))} + className="w-full mt-0.5" + /> +

+ Lower values produce more muted, vintage-style colors. +

+
+ + {error &&

{error}

} + + {originalSize != null && processedSize != null && ( +
+

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

+

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

+
+ )} + + {processing ? ( + + ) : ( + + )} + + {!hasMultiple && downloadUrl && ( + + + Download + + )} + + ); +} + +function SectionLabel({ children }: { children: React.ReactNode }) { + return ( +

+ {children} +

+ ); +} diff --git a/apps/web/src/components/tools/pipeline-step-summary.ts b/apps/web/src/components/tools/pipeline-step-summary.ts index d092f1b8..1732940e 100644 --- a/apps/web/src/components/tools/pipeline-step-summary.ts +++ b/apps/web/src/components/tools/pipeline-step-summary.ts @@ -53,6 +53,10 @@ export function getSettingsSummary(toolId: string, settings: Record = { "watermark-text": ["compress", "convert"], "watermark-image": ["compress", "convert"], "text-overlay": ["compress", "convert"], + colorize: ["adjust-colors", "image-enhancement", "upscale", "compress"], 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 f4bb30b1..fffe5097 100644 --- a/apps/web/src/lib/tool-registry.tsx +++ b/apps/web/src/lib/tool-registry.tsx @@ -244,6 +244,11 @@ const ImageEnhancementSettings = lazy(() => default: m.ImageEnhancementSettings, })), ); +const ColorizeSettings = lazy(() => + import("@/components/tools/colorize-settings").then((m) => ({ + default: m.ColorizeSettings, + })), +); // ── Color tool wrapper ───────────────────────────────────────────── // Color tools share a single component but differ by toolId. @@ -372,6 +377,7 @@ export const toolRegistry = new Map([ Settings: ImageEnhancementSettings as never, }, ], + ["colorize", { displayMode: "before-after", Settings: ColorizeSettings }], ]); export function getToolRegistryEntry(toolId: string): ToolRegistryEntry | undefined { diff --git a/docker/download_models.py b/docker/download_models.py index 271292b7..21c620cc 100644 --- a/docker/download_models.py +++ b/docker/download_models.py @@ -32,6 +32,14 @@ GFPGAN_MODEL_URL = ( GFPGAN_MODEL_PATH = os.path.join(GFPGAN_MODEL_DIR, "GFPGANv1.3.pth") GFPGAN_MIN_SIZE = 300_000_000 # ~332 MB +DDCOLOR_MODEL_DIR = "/opt/models/ddcolor" +DDCOLOR_MODEL_URL = ( + "https://huggingface.co/piddnad/DDColor-models/resolve/main/ddcolor_paper_tiny.pth" +) +DDCOLOR_ONNX_PATH = os.path.join(DDCOLOR_MODEL_DIR, "ddcolor.onnx") +DDCOLOR_MIN_SIZE = 50_000_000 # ~220 MB ONNX + + REMBG_MODELS = [ "u2net", "isnet-general-use", @@ -145,6 +153,37 @@ def download_gfpgan_model(): print(f" GFPGANv1.3.pth downloaded ({size / 1_000_000:.1f} MB)\n") +def download_ddcolor_model(): + """Download pre-exported DDColor ONNX model for AI photo colorization. + + Uses the pre-converted ONNX model from HuggingFace (facefusion repo) + for direct inference via onnxruntime without needing PyTorch. + """ + print("=== Downloading DDColor ONNX model ===") + os.makedirs(DDCOLOR_MODEL_DIR, exist_ok=True) + + from huggingface_hub import hf_hub_download + + print(" Downloading DDColor ONNX from HuggingFace...") + downloaded_path = hf_hub_download( + repo_id="facefusion/models-3.0.0", + filename="ddcolor.onnx", + local_dir=DDCOLOR_MODEL_DIR, + ) + + # huggingface_hub downloads to local_dir/filename + actual_path = os.path.join(DDCOLOR_MODEL_DIR, "ddcolor.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 > DDCOLOR_MIN_SIZE, ( + f"DDColor model too small: {size} bytes (expected > {DDCOLOR_MIN_SIZE})" + ) + print(f" DDColor ONNX model ready ({size / 1_000_000:.1f} MB)\n") + + + def download_paddleocr_models(): """Pre-download PaddleOCR PP-OCRv5 model weights from HuggingFace. @@ -242,6 +281,15 @@ def smoke_test(): ) print(" GFPGAN model file verified") + # DDColor ONNX model must exist + assert os.path.exists(DDCOLOR_ONNX_PATH), ( + f"DDColor model missing: {DDCOLOR_ONNX_PATH}" + ) + assert os.path.getsize(DDCOLOR_ONNX_PATH) > DDCOLOR_MIN_SIZE, ( + "DDColor model file is too small" + ) + print(" DDColor ONNX model file verified") + # PaddleOCR model directories must exist for repo_id in PADDLEOCR_MODELS: model_name = repo_id.split("/", 1)[1] @@ -264,6 +312,7 @@ def main(): download_rembg_models() download_realesrgan_model() download_gfpgan_model() + download_ddcolor_model() download_paddleocr_models() download_paddleocr_vl_model() verify_mediapipe() diff --git a/packages/ai/python/colorize.py b/packages/ai/python/colorize.py new file mode 100644 index 00000000..bb5a828e --- /dev/null +++ b/packages/ai/python/colorize.py @@ -0,0 +1,231 @@ +"""AI photo colorization using DDColor ONNX model. + +Converts grayscale / black-and-white photos to full color using the DDColor +dual-decoder architecture. Falls back to a lightweight OpenCV DNN colorizer +when the DDColor model is unavailable. +""" +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) + + +DDCOLOR_MODEL_PATH = os.environ.get( + "DDCOLOR_MODEL_PATH", + "/opt/models/ddcolor/ddcolor.onnx", +) + +# OpenCV DNN fallback model paths (lightweight ~17 MB) +OPENCV_PROTO_PATH = os.environ.get( + "OPENCV_COLORIZE_PROTO", + "/opt/models/colorize-opencv/colorization_deploy_v2.prototxt", +) +OPENCV_MODEL_PATH = os.environ.get( + "OPENCV_COLORIZE_MODEL", + "/opt/models/colorize-opencv/colorization_release_v2.caffemodel", +) +OPENCV_POINTS_PATH = os.environ.get( + "OPENCV_COLORIZE_POINTS", + "/opt/models/colorize-opencv/pts_in_hull.npy", +) + + +def colorize_ddcolor(img_bgr, intensity): + """Colorize using DDColor ONNX model.""" + import onnxruntime as ort + + emit_progress(15, "Loading DDColor model") + + providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] + try: + from gpu import gpu_available + if not gpu_available(): + providers = ["CPUExecutionProvider"] + except ImportError: + providers = ["CPUExecutionProvider"] + + session = ort.InferenceSession(DDCOLOR_MODEL_PATH, providers=providers) + input_name = session.get_inputs()[0].name + input_shape = session.get_inputs()[0].shape + # Dynamic dims are strings ('w', 'h'), so default to 512 if not int + model_size = input_shape[2] if len(input_shape) == 4 and isinstance(input_shape[2], int) else 512 + + emit_progress(25, "Preprocessing image") + + orig_h, orig_w = img_bgr.shape[:2] + + # Convert to Lab, extract L channel + img_lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB) + orig_l = img_lab[:, :, 0].astype(np.float32) + + # Prepare input: resize, normalize to [0, 1], NCHW format + 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)) + img_nchw = np.expand_dims(img_nchw, axis=0) + + emit_progress(40, "Running AI colorization") + + # Run inference - model outputs predicted ab channels + output = session.run(None, {input_name: img_nchw})[0] + + emit_progress(75, "Post-processing colors") + + # Output shape: (1, 2, H, W) - predicted ab channels + ab_pred = output[0] # (2, model_size, model_size) + + # Resize ab channels back to original dimensions + 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)) + + # Model outputs ab values already in Lab scale (roughly -50 to +70) + 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: + # Blend with original ab channels (grayscale has ab near 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 + + # Reconstruct Lab image + 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) + + # Convert back to BGR + result_bgr = cv2.cvtColor(result_lab, cv2.COLOR_LAB2BGR) + return result_bgr, "ddcolor" + + +def colorize_opencv(img_bgr, intensity): + """Fallback colorization using lightweight OpenCV DNN model (Zhang et al.).""" + emit_progress(15, "Loading OpenCV colorizer") + + net = cv2.dnn.readNetFromCaffe(OPENCV_PROTO_PATH, OPENCV_MODEL_PATH) + pts = np.load(OPENCV_POINTS_PATH).transpose().reshape(2, 313, 1, 1) + + # Set cluster centers as 1x1 convolution kernel + net.getLayer(net.getLayerId("class8_ab")).blobs = [pts.astype(np.float32)] + net.getLayer(net.getLayerId("conv8_313_rh")).blobs = [ + np.full([1, 313], 2.606, dtype=np.float32) + ] + + emit_progress(25, "Preprocessing image") + + orig_h, orig_w = img_bgr.shape[:2] + + # Convert to Lab + img_lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB) + orig_l = img_lab[:, :, 0].astype(np.float32) + + # Resize L channel and normalize for the network + l_resized = cv2.resize(orig_l, (224, 224)) + l_resized -= 50 # Mean subtraction + + emit_progress(40, "Running colorization") + + net.setInput(cv2.dnn.blobFromImage(l_resized)) + ab_out = net.forward()[0] # (2, 56, 56) + + emit_progress(75, "Post-processing colors") + + # Resize ab to original size + ab_a = cv2.resize(ab_out[0], (orig_w, orig_h)) + ab_b = cv2.resize(ab_out[1], (orig_w, orig_h)) + + # Apply intensity + 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 + + # Reconstruct + 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) + + result_bgr = cv2.cvtColor(result_lab, cv2.COLOR_LAB2BGR) + return result_bgr, "opencv" + + +def main(): + input_path = sys.argv[1] + output_path = sys.argv[2] + settings = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {} + + intensity = float(settings.get("intensity", 1.0)) + model_choice = settings.get("model", "auto") + + try: + emit_progress(5, "Opening image") + img_bgr = cv2.imread(input_path, cv2.IMREAD_COLOR) + if img_bgr is None: + # Try with Pillow for formats OpenCV can't read + 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_bgr = None + method = "unknown" + + # Try DDColor first + if model_choice in ("auto", "ddcolor"): + try: + if os.path.exists(DDCOLOR_MODEL_PATH): + result_bgr, method = colorize_ddcolor(img_bgr, intensity) + elif model_choice == "ddcolor": + emit_progress(10, "DDColor model not found, using fallback") + except Exception as e: + if model_choice == "ddcolor": + emit_progress(10, f"DDColor failed: {str(e)[:50]}") + result_bgr = None + + # Try OpenCV fallback + if result_bgr is None and model_choice in ("auto", "opencv"): + try: + if os.path.exists(OPENCV_PROTO_PATH) and os.path.exists(OPENCV_MODEL_PATH): + result_bgr, method = colorize_opencv(img_bgr, intensity) + except Exception: + result_bgr = None + + if result_bgr is None: + print(json.dumps({ + "success": False, + "error": "No colorization model available. Install DDColor or OpenCV models.", + })) + sys.exit(1) + + emit_progress(90, "Saving result") + + # Save output as PNG (API route handles format conversion) + cv2.imwrite(output_path, result_bgr) + + print(json.dumps({ + "success": True, + "width": orig_w, + "height": orig_h, + "method": method, + "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/colorization.ts b/packages/ai/src/colorization.ts new file mode 100644 index 00000000..4849f746 --- /dev/null +++ b/packages/ai/src/colorization.ts @@ -0,0 +1,46 @@ +import { readFile, writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import { type ProgressCallback, runPythonWithProgress } from "./bridge.js"; + +export interface ColorizeOptions { + intensity?: number; + model?: string; +} + +export interface ColorizeResult { + buffer: Buffer; + width: number; + height: number; + method: string; +} + +export async function colorize( + inputBuffer: Buffer, + outputDir: string, + options: ColorizeOptions = {}, + onProgress?: ProgressCallback, +): Promise { + const inputPath = join(outputDir, "input_colorize.png"); + const outputPath = join(outputDir, "output_colorize.png"); + + await writeFile(inputPath, inputBuffer); + const { stdout } = await runPythonWithProgress( + "colorize.py", + [inputPath, outputPath, JSON.stringify(options)], + { onProgress }, + ); + + const result = JSON.parse(stdout); + if (!result.success) { + throw new Error(result.error || "Colorization failed"); + } + + const actualOutputPath = result.output_path || outputPath; + const buffer = await readFile(actualOutputPath); + return { + buffer, + width: result.width, + height: result.height, + method: result.method ?? "unknown", + }; +} diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 3924106c..ba0ec0d1 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -1,5 +1,6 @@ export { removeBackground } from "./background-removal.js"; 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 { inpaint } from "./inpainting.js"; diff --git a/packages/shared/src/constants.ts b/packages/shared/src/constants.ts index 6be516f8..19210004 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: "colorize", + name: "AI Colorization", + description: "Convert B&W photos to full color with AI", + category: "ai", + icon: "Palette", + route: "/colorize", + }, // Watermark & Overlay { id: "watermark-text", @@ -396,4 +404,5 @@ export const PYTHON_SIDECAR_TOOLS = [ "blur-faces", "erase-object", "ocr", + "colorize", ] as const; diff --git a/packages/shared/src/i18n/en.ts b/packages/shared/src/i18n/en.ts index b2a65366..ff3f85c0 100644 --- a/packages/shared/src/i18n/en.ts +++ b/packages/shared/src/i18n/en.ts @@ -84,6 +84,10 @@ export const en = { name: "Content-Aware Resize", description: "Intelligently resize images while preserving important content", }, + colorize: { + name: "AI Colorization", + description: "Convert black & white photos to full color using AI deep learning models", + }, "watermark-text": { name: "Text Watermark", description: "Add text watermark overlay" }, "watermark-image": { name: "Image Watermark", description: "Overlay a logo as watermark" }, "text-overlay": { name: "Text Overlay", description: "Add styled text to images" },