SnapOtter/packages/ai/python/outpaint.py

"""Enhanced LaMa outpainting with 6-step preprocessing pipeline."""
import sys
import os
import json


MODEL_SIZE = 512

TIER_PARAMS = {
    "fast":     {"band_size": 192, "mask_dilate": 12, "seam_strip": 0,  "use_telea": False},
    "balanced": {"band_size": 128, "mask_dilate": 18, "seam_strip": 24, "use_telea": True},
    "high":     {"band_size": 72,  "mask_dilate": 24, "seam_strip": 36, "use_telea": True},
}


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 only if allowed."""
    if os.path.exists(LAMA_MODEL_PATH):
        return LAMA_MODEL_PATH
    if os.path.exists(LAMA_LOCAL_PATH):
        return LAMA_LOCAL_PATH

    from offline_guard import ensure_download_allowed
    ensure_download_allowed("LaMa inpainting model (lama_fp32.onnx)")
    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, band_size=128, progress_start=30, progress_end=75):
    """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

        progress_range = progress_end - progress_start
        progress = progress_start + int(progress_range * band_index / total_bands)
        emit_progress(min(progress, progress_end), 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])

    tier = sys.argv[7] if len(sys.argv) > 7 else "balanced"
    if tier not in TIER_PARAMS:
        tier = "balanced"
    params = TIER_PARAMS[tier]

    try:
        emit_progress(5, "Preparing")
        from PIL import Image
        import numpy as np

        try:
            import cv2
            import onnxruntime
        except ImportError as e:
            msg = str(e)
            hint = "Fix with: apt-get install -y libgl1" if "libGL" in msg else "Requires opencv-python-headless and onnxruntime."
            print(json.dumps({
                "success": False,
                "error": f"Missing dependency: {msg}. {hint}",
            }))
            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, (params["mask_dilate"] * 2 + 1, params["mask_dilate"] * 2 + 1)
        )
        mask = cv2.dilate(mask, dilate_kernel, iterations=1)

        # Step 3: Telea pre-inpainting for gradient hints
        if params["use_telea"]:
            emit_progress(25, "Pre-filling gradients")
            canvas = cv2.inpaint(canvas, mask, 3, cv2.INPAINT_TELEA)

        # Step 4: Progressive LaMa outpainting in concentric bands
        if params["use_telea"]:
            canvas = _progressive_outpaint(session, canvas, mask, params["band_size"], 30, 75)
        else:
            canvas = _progressive_outpaint(session, canvas, mask, params["band_size"], 20, 85)

        # Step 5: Seam refinement -- second LaMa pass on thin boundary strip
        seam_strip = params["seam_strip"]
        if seam_strip > 0:
            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 + 1, seam_strip + 1)
            )
            outer_kernel = cv2.getStructuringElement(
                cv2.MORPH_ELLIPSE, (seam_strip * 2 + 1, seam_strip * 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()