mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
Add AI-powered photo colorization that converts B&W/grayscale images to full color using DDColor (ICCV 2023 dual-decoder architecture) via ONNX Runtime. Includes model selection (Auto/DDColor/Classic), adjustable color intensity, batch processing, before/after preview, and full HEIC/HEIF support. Co-authored-by: stirling-image <stirling-image@users.noreply.github.com>
This commit is contained in:
stirling-image
co-authored by
stirling-image
stirling-image
parent
58cdbe50b4
commit
c280076098
@@ -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" };
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -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<void> {
|
||||
{ id: "smart-crop", register: registerSmartCrop },
|
||||
{ id: "image-enhancement", register: registerImageEnhancement },
|
||||
{ id: "content-aware-resize", register: registerContentAwareResize },
|
||||
{ id: "colorize", register: registerColorize },
|
||||
];
|
||||
|
||||
let skipped = 0;
|
||||
|
||||
@@ -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<Model>("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 (
|
||||
<form onSubmit={handleSubmit} className="space-y-3">
|
||||
{/* Model selector */}
|
||||
<SectionLabel>AI Model</SectionLabel>
|
||||
<div className="grid grid-cols-3 gap-1">
|
||||
{MODEL_OPTIONS.map((opt) => (
|
||||
<button
|
||||
type="button"
|
||||
key={opt.value}
|
||||
onClick={() => setModel(opt.value)}
|
||||
className={`text-xs py-2 rounded transition-colors ${
|
||||
model === opt.value
|
||||
? "bg-primary text-primary-foreground"
|
||||
: "bg-muted text-muted-foreground hover:bg-primary/10"
|
||||
}`}
|
||||
>
|
||||
<span className="block font-medium">{opt.label}</span>
|
||||
<span className="block text-[10px] opacity-70">{opt.desc}</span>
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
|
||||
{/* Color intensity */}
|
||||
<SectionLabel>Color Intensity</SectionLabel>
|
||||
<div>
|
||||
<div className="flex justify-between items-center">
|
||||
<span className="text-xs text-muted-foreground">
|
||||
{intensity === 0
|
||||
? "Grayscale"
|
||||
: intensity < 50
|
||||
? "Subtle"
|
||||
: intensity < 80
|
||||
? "Natural"
|
||||
: "Vivid"}
|
||||
</span>
|
||||
<span className="text-xs font-mono text-foreground tabular-nums w-10 text-right">
|
||||
{intensity}%
|
||||
</span>
|
||||
</div>
|
||||
<input
|
||||
type="range"
|
||||
min={10}
|
||||
max={100}
|
||||
step={5}
|
||||
value={intensity}
|
||||
onChange={(e) => setIntensity(Number(e.target.value))}
|
||||
className="w-full mt-0.5"
|
||||
/>
|
||||
<p className="text-[10px] text-muted-foreground/60 mt-0.5">
|
||||
Lower values produce more muted, vintage-style colors.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
{error && <p className="text-xs text-red-500">{error}</p>}
|
||||
|
||||
{originalSize != null && processedSize != null && (
|
||||
<div className="text-xs text-muted-foreground space-y-0.5">
|
||||
<p>Original: {(originalSize / 1024).toFixed(1)} KB</p>
|
||||
<p>Colorized: {(processedSize / 1024).toFixed(1)} KB</p>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{processing ? (
|
||||
<ProgressCard
|
||||
active={processing}
|
||||
phase={progress.phase === "idle" ? "uploading" : progress.phase}
|
||||
label={hasMultiple ? `Colorizing ${files.length} images` : "Colorizing"}
|
||||
stage={progress.stage}
|
||||
percent={progress.percent}
|
||||
elapsed={progress.elapsed}
|
||||
/>
|
||||
) : (
|
||||
<button
|
||||
type="submit"
|
||||
data-testid="colorize-submit"
|
||||
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 ? `Colorize (${files.length} files)` : "Colorize"}
|
||||
</button>
|
||||
)}
|
||||
|
||||
{!hasMultiple && downloadUrl && (
|
||||
<a
|
||||
href={downloadUrl}
|
||||
download
|
||||
data-testid="colorize-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>
|
||||
)}
|
||||
</form>
|
||||
);
|
||||
}
|
||||
|
||||
function SectionLabel({ children }: { children: React.ReactNode }) {
|
||||
return (
|
||||
<p className="text-[11px] font-semibold uppercase tracking-wider text-muted-foreground/70 pt-1">
|
||||
{children}
|
||||
</p>
|
||||
);
|
||||
}
|
||||
@@ -53,6 +53,10 @@ export function getSettingsSummary(toolId: string, settings: Record<string, unkn
|
||||
if (settings.scale) return `${settings.scale}x`;
|
||||
return "";
|
||||
}
|
||||
case "colorize": {
|
||||
const pct = settings.intensity != null ? Math.round(Number(settings.intensity) * 100) : 100;
|
||||
return `${pct}% intensity`;
|
||||
}
|
||||
default:
|
||||
return "";
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@ const TOOL_SUGGESTIONS: Record<string, string[]> = {
|
||||
"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"],
|
||||
};
|
||||
|
||||
@@ -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<string, ToolRegistryEntry>([
|
||||
Settings: ImageEnhancementSettings as never,
|
||||
},
|
||||
],
|
||||
["colorize", { displayMode: "before-after", Settings: ColorizeSettings }],
|
||||
]);
|
||||
|
||||
export function getToolRegistryEntry(toolId: string): ToolRegistryEntry | undefined {
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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<ColorizeResult> {
|
||||
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",
|
||||
};
|
||||
}
|
||||
@@ -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";
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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" },
|
||||
|
||||
Reference in New Issue
Block a user