diff --git a/packages/ai/python/outpaint.py b/packages/ai/python/outpaint.py new file mode 100644 index 00000000..37de83c5 --- /dev/null +++ b/packages/ai/python/outpaint.py @@ -0,0 +1,272 @@ +"""Enhanced LaMa outpainting with 6-step preprocessing pipeline.""" +import sys +import os +import json + + +MODEL_SIZE = 512 +BAND_SIZE = 128 +MASK_DILATE_PX = 18 +SEAM_STRIP_PX = 24 + + +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) + + +# Resolve the LaMa ONNX model path. +_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models") +LAMA_MODEL_DIR = os.environ.get("LAMA_MODEL_DIR", os.path.join(_MODELS_BASE, "lama")) +LAMA_MODEL_PATH = os.path.join(LAMA_MODEL_DIR, "lama_fp32.onnx") +LAMA_LOCAL_CACHE = os.path.join(os.path.expanduser("~"), ".cache", "snapotter", "lama") +LAMA_LOCAL_PATH = os.path.join(LAMA_LOCAL_CACHE, "lama_fp32.onnx") +LAMA_HF_URL = "https://huggingface.co/Carve/LaMa-ONNX/resolve/main/lama_fp32.onnx" + + +def _get_model_path(): + """Return path to the LaMa ONNX model, 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 + + emit_progress(5, "Downloading LaMa model") + os.makedirs(LAMA_LOCAL_CACHE, exist_ok=True) + import urllib.request + urllib.request.urlretrieve(LAMA_HF_URL, LAMA_LOCAL_PATH) + return LAMA_LOCAL_PATH + + +def _run_lama(session, canvas, mask, feather_radius=5): + """Run a single LaMa inference pass with feathered compositing.""" + import cv2 + import numpy as np + + h, w = canvas.shape[:2] + + # Skip if mask is empty + if np.sum(mask) == 0: + return canvas + + # Resize to model input size + img_resized = cv2.resize(canvas, (MODEL_SIZE, MODEL_SIZE), interpolation=cv2.INTER_AREA) + mask_resized = cv2.resize(mask, (MODEL_SIZE, MODEL_SIZE), interpolation=cv2.INTER_NEAREST) + + # Preprocess image: HWC -> NCHW float32 [0,1] + img_input = np.transpose(img_resized, (2, 0, 1)).astype(np.float32) / 255.0 + img_input = img_input[np.newaxis, ...] # (1, 3, 512, 512) + + # Preprocess mask: HW -> N1HW float32 binary + 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}) + + # Postprocess: (1, 3, 512, 512) -> (512, 512, 3) uint8 + raw_output = outputs[0][0] # (3, 512, 512) + raw_output = np.transpose(raw_output, (1, 2, 0)) # (512, 512, 3) + raw_output = np.clip(raw_output, 0, 255).astype(np.uint8) + + # Resize back to original dimensions + inpainted = cv2.resize(raw_output, (w, h), interpolation=cv2.INTER_LANCZOS4) + + # Feathered blending: smooth transition at mask edges + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (feather_radius, feather_radius)) + dilated = cv2.dilate(mask.astype(np.uint8), kernel, iterations=1) + blur_size = feather_radius * 2 + 1 + alpha = cv2.GaussianBlur(dilated.astype(np.float32) / 255.0, (blur_size, blur_size), 0) + alpha = np.clip(alpha, 0.0, 1.0) + alpha_3ch = alpha[:, :, np.newaxis] + + result = (canvas.astype(np.float32) * (1.0 - alpha_3ch) + + inpainted.astype(np.float32) * alpha_3ch) + return np.clip(result, 0, 255).astype(np.uint8) + + +def _progressive_outpaint(session, canvas, mask): + """Process mask in concentric bands from original edge outward.""" + import cv2 + import numpy as np + + remaining = mask.copy() + current_canvas = canvas.copy() + + # Count total bands for progress scaling + total_white = np.sum(remaining > 127) + if total_white == 0: + return current_canvas + + band_index = 0 + total_bands = 0 + temp = remaining.copy() + while np.sum(temp > 127) > 0: + kernel = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, (BAND_SIZE * 2 + 1, BAND_SIZE * 2 + 1) + ) + eroded = cv2.erode(temp, kernel, iterations=1) + temp = eroded + total_bands += 1 + + if total_bands == 0: + total_bands = 1 + + while np.sum(remaining > 127) > 0: + # Erode remaining mask to peel off outermost band + kernel = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, (BAND_SIZE * 2 + 1, BAND_SIZE * 2 + 1) + ) + eroded = cv2.erode(remaining, kernel, iterations=1) + + # Current band = remaining minus eroded + band = cv2.subtract(remaining, eroded) + + # Run LaMa on this band + current_canvas = _run_lama(session, current_canvas, band) + + remaining = eroded + band_index += 1 + + # Scale progress between 30% and 75% + progress = 30 + int(45 * band_index / total_bands) + emit_progress(min(progress, 75), f"AI outpainting band {band_index}/{total_bands}") + + return current_canvas + + +def main(): + input_path = sys.argv[1] + output_path = sys.argv[2] + extend_top = int(sys.argv[3]) + extend_right = int(sys.argv[4]) + extend_bottom = int(sys.argv[5]) + extend_left = int(sys.argv[6]) + + try: + emit_progress(5, "Preparing") + from PIL import Image + import numpy as np + + try: + import cv2 + import onnxruntime + except ImportError as e: + print(json.dumps({ + "success": False, + "error": f"Missing dependency: {e}. Requires opencv-python-headless and onnxruntime.", + })) + sys.exit(1) + + emit_progress(8, "Loading model") + model_path = _get_model_path() + + from gpu import safe_onnx_session + session, _device = safe_onnx_session(model_path) + + emit_progress(15, "Loading image") + img = Image.open(input_path).convert("RGB") + orig_w, orig_h = img.size + img_array = np.array(img) + + # Step 1: Mirror pre-fill -- extend canvas with reflected content + emit_progress(20, "Extending canvas") + canvas = cv2.copyMakeBorder( + img_array, + extend_top, extend_bottom, extend_left, extend_right, + cv2.BORDER_REFLECT_101, + ) + new_h, new_w = canvas.shape[:2] + + # Step 2: Create mask -- white for extended regions, black for original + mask = np.zeros((new_h, new_w), dtype=np.uint8) + mask[:extend_top, :] = 255 + mask[extend_top + orig_h:, :] = 255 + mask[:, :extend_left] = 255 + mask[:, extend_left + orig_w:] = 255 + + # Dilate mask into original area for overlap + dilate_kernel = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, (MASK_DILATE_PX * 2 + 1, MASK_DILATE_PX * 2 + 1) + ) + mask = cv2.dilate(mask, dilate_kernel, iterations=1) + + # Step 3: Telea pre-inpainting for gradient hints + emit_progress(25, "Pre-filling gradients") + canvas = cv2.inpaint(canvas, mask, 3, cv2.INPAINT_TELEA) + + # Step 4: Progressive LaMa outpainting in concentric bands + canvas = _progressive_outpaint(session, canvas, mask) + + # Step 5: Seam refinement -- second LaMa pass on thin boundary strip + emit_progress(80, "Refining seams") + seam_mask = np.zeros((new_h, new_w), dtype=np.uint8) + + # Create thin strip along original image boundary + inner_kernel = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, (SEAM_STRIP_PX + 1, SEAM_STRIP_PX + 1) + ) + outer_kernel = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, (SEAM_STRIP_PX * 2 + 1, SEAM_STRIP_PX * 2 + 1) + ) + + # Original region mask (before dilation) + orig_mask = np.zeros((new_h, new_w), dtype=np.uint8) + orig_mask[extend_top:extend_top + orig_h, extend_left:extend_left + orig_w] = 255 + + inner_edge = cv2.erode(orig_mask, inner_kernel, iterations=1) + outer_edge = cv2.dilate(orig_mask, outer_kernel, iterations=1) + seam_mask = cv2.subtract(outer_edge, inner_edge) + + canvas = _run_lama(session, canvas, seam_mask) + + # Step 6: Poisson blending -- paste untouched original back + emit_progress(90, "Blending") + + # Extract the region where original will go + center_x = extend_left + orig_w // 2 + center_y = extend_top + orig_h // 2 + + try: + # Poisson seamless clone for smooth boundary + blended = cv2.seamlessClone( + img_array, canvas, np.ones_like(img_array[:, :, 0]) * 255, + (center_x, center_y), cv2.NORMAL_CLONE, + ) + except Exception: + # Fall back to alpha blending if Poisson fails + blended = canvas.copy() + # Create gradient alpha mask for the boundary + blend_mask = np.zeros((new_h, new_w), dtype=np.float32) + blend_mask[extend_top:extend_top + orig_h, extend_left:extend_left + orig_w] = 1.0 + # Blur edges for smooth transition + blend_mask = cv2.GaussianBlur(blend_mask, (31, 31), 0) + alpha_3ch = blend_mask[:, :, np.newaxis] + + # Paste original into canvas position + paste = canvas.copy() + paste[extend_top:extend_top + orig_h, extend_left:extend_left + orig_w] = img_array + + blended = (canvas.astype(np.float32) * (1.0 - alpha_3ch) + + paste.astype(np.float32) * alpha_3ch) + blended = np.clip(blended, 0, 255).astype(np.uint8) + + emit_progress(95, "Saving") + Image.fromarray(blended).save(output_path) + + print(json.dumps({ + "success": True, + "method": "lama-enhanced-outpaint", + "originalWidth": orig_w, + "originalHeight": orig_h, + "newWidth": new_w, + "newHeight": new_h, + })) + + except Exception as e: + print(json.dumps({"success": False, "error": str(e)})) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/packages/shared/src/constants.ts b/packages/shared/src/constants.ts index 13bab078..16f6ab79 100644 --- a/packages/shared/src/constants.ts +++ b/packages/shared/src/constants.ts @@ -249,6 +249,14 @@ export const TOOLS: Tool[] = [ icon: "Scaling", route: "/content-aware-resize", }, + { + id: "content-aware-crop", + name: "Content-Aware Crop", + description: "Extend canvas beyond image bounds with AI-powered fill", + category: "ai", + icon: "Expand", + route: "/content-aware-crop", + }, { id: "transparency-fixer", name: "PNG Transparency Fixer", @@ -1240,4 +1248,5 @@ export const PYTHON_SIDECAR_TOOLS = [ "restore-photo", "passport-photo", "transparency-fixer", + "content-aware-crop", ] as const; diff --git a/packages/shared/src/features.ts b/packages/shared/src/features.ts index c6be7c38..1148c1f9 100644 --- a/packages/shared/src/features.ts +++ b/packages/shared/src/features.ts @@ -40,7 +40,7 @@ export const FEATURE_BUNDLES: Record = { name: "Object Eraser & Colorize", description: "Erase objects from photos and colorize B&W images", estimatedSize: "1-2 GB", - enablesTools: ["erase-object", "colorize"], + enablesTools: ["erase-object", "colorize", "content-aware-crop"], }, "upscale-enhance": { id: "upscale-enhance", diff --git a/packages/shared/src/i18n/en.ts b/packages/shared/src/i18n/en.ts index 40c9d18a..32104ad1 100644 --- a/packages/shared/src/i18n/en.ts +++ b/packages/shared/src/i18n/en.ts @@ -119,6 +119,10 @@ export const en = { name: "Content-Aware Resize", description: "Intelligently resize images while preserving important content", }, + "content-aware-crop": { + name: "Content-Aware Crop", + description: "Extend canvas with AI-powered fill", + }, colorize: { name: "AI Colorization", description: "Convert black & white photos to full color using AI deep learning models",