fix(erase-object): crop-based HD inpainting to remove ghosting and blur (#501)

Dilate the mask, crop a padded box around it, run LaMa on the crop at 512, and composite back cleanly. Fixes the ghost remnants (#491) and sharpens small/medium-object fills in high-res images (#141 core). Same model, still offline, no new bundle. Closes #491.

packages/ai/python/inpaint.py

@@ -22,6 +22,112 @@ LAMA_HF_URL = "https://huggingface.co/Carve/LaMa-ONNX/resolve/main/lama_fp32.onn
 # The ONNX model expects 512x512 fixed input.
 MODEL_SIZE = 512
 
+# Crop-and-composite tuning. Defaults validated against the real LaMa model on
+# small/medium objects in high-res images; safe to tune.
+MIN_IMAGE_DIM = 8  # below this, inpainting is meaningless; return the original
+WHOLE_FRAME_RATIO = 0.95  # crop this fraction of the frame -> just use the whole frame
+DILATE_FRAC = 0.04  # mask dilation as a fraction of the mask bbox diagonal
+DILATE_MIN = 6
+DILATE_MAX = 96
+MARGIN_FRAC = 0.5  # context margin around the dilated mask, fraction of its max side
+MARGIN_MIN = 32
+
+
+def _mask_bbox(mask_bin):
+    """Return (x0, y0, x1, y1) tight bounds of nonzero pixels, or None if empty."""
+    import numpy as np
+
+    ys, xs = np.where(mask_bin > 0)
+    if xs.size == 0:
+        return None
+    return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1
+
+
+def dilate_mask(mask_bin, d):
+    """Grow a binary mask by d px with an elliptical kernel. d<=0 is a no-op copy."""
+    import cv2
+
+    if d <= 0:
+        return mask_bin.copy()
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * d + 1, 2 * d + 1))
+    return cv2.dilate(mask_bin, kernel, iterations=1)
+
+
+def compute_crop_box(mask_dil, image_shape, margin_frac, margin_min):
+    """Bounding box of the dilated mask, expanded by a context margin, clamped."""
+    h, w = image_shape[:2]
+    x0, y0, x1, y1 = _mask_bbox(mask_dil)
+    margin = int(max(margin_min, round(margin_frac * max(x1 - x0, y1 - y0))))
+    return (
+        max(0, x0 - margin),
+        max(0, y0 - margin),
+        min(w, x1 + margin),
+        min(h, y1 + margin),
+    )
+
+
+def composite(original, inpainted_crop, mask_dil_native, crop_box, feather):
+    """Blend the inpainted crop into a copy of the full-res original.
+
+    Only the crop_box region is written, and within it only where the feathered
+    dilated mask has alpha > 0. Pixels beyond the feather stay byte-identical.
+    """
+    import cv2
+    import numpy as np
+
+    x0, y0, x1, y1 = crop_box
+    result = original.copy()
+    ksize = 2 * feather + 1
+    alpha = cv2.GaussianBlur((mask_dil_native > 0).astype(np.float32), (ksize, ksize), 0)
+    alpha = np.clip(alpha, 0.0, 1.0)[:, :, np.newaxis]
+    region = result[y0:y1, x0:x1].astype(np.float32)
+    blended = region * (1.0 - alpha) + inpainted_crop.astype(np.float32) * alpha
+    result[y0:y1, x0:x1] = np.clip(blended, 0, 255).astype(np.uint8)
+    return result
+
+
+def inpaint_array(img_array, mask_array, run_model, progress=None):
+    """Crop-and-composite inpainting orchestrator (model-agnostic).
+
+    img_array: HxWx3 uint8 RGB. mask_array: HxW uint8 (white = erase).
+    run_model(crop_img, crop_mask) -> inpainted crop, same HxWx3 size as crop_img.
+    progress(percent, stage) optional; called for the SSE progress UI.
+    """
+    import numpy as np
+
+    def _p(percent, stage):
+        if progress:
+            progress(percent, stage)
+
+    h, w = img_array.shape[:2]
+    mask_bin = (mask_array > 127).astype(np.uint8) * 255
+
+    # Guards: nothing to erase, or an image too small to inpaint meaningfully.
+    if int(mask_bin.max()) == 0 or min(h, w) < MIN_IMAGE_DIM:
+        return img_array.copy()
+
+    _p(30, "Preprocessing")
+    x0, y0, x1, y1 = _mask_bbox(mask_bin)
+    diag = float(np.hypot(x1 - x0, y1 - y0))
+    d = int(np.clip(round(DILATE_FRAC * diag), DILATE_MIN, DILATE_MAX))
+    mask_dil = dilate_mask(mask_bin, d)
+
+    bx0, by0, bx1, by1 = compute_crop_box(mask_dil, img_array.shape, MARGIN_FRAC, MARGIN_MIN)
+    if (bx1 - bx0) * (by1 - by0) >= WHOLE_FRAME_RATIO * w * h:
+        bx0, by0, bx1, by1 = 0, 0, w, h
+
+    crop_img = img_array[by0:by1, bx0:bx1]
+    crop_mask = mask_dil[by0:by1, bx0:bx1]
+
+    _p(40, "Erasing objects")
+    inpainted_crop = run_model(crop_img, crop_mask)
+
+    _p(75, "Compositing")
+    side = max(bx1 - bx0, by1 - by0)
+    feather = int(np.clip(round(0.01 * side), 2, 12))
+    feather = min(feather, max(1, d // 2))
+    return composite(img_array, inpainted_crop, crop_mask, (bx0, by0, bx1, by1), feather)
+
 
 def _get_model_path():
     """Return path to the LaMa ONNX model, downloading only if allowed."""
@@ -39,53 +145,32 @@ def _get_model_path():
     return LAMA_LOCAL_PATH
 
 
-def _preprocess_image(img_array):
+def _make_run_model(session):
-    """Convert HWC uint8 RGB image to NCHW float32 [0,1] at MODEL_SIZE."""
+    """Build a run_model(crop_img, crop_mask) that runs LaMa at its fixed 512x512.
-    import cv2
-    import numpy as np
 
-    resized = cv2.resize(img_array, (MODEL_SIZE, MODEL_SIZE), interpolation=cv2.INTER_AREA)
+    Small crops are upscaled to 512 (INTER_LINEAR), large crops downscaled
-    # HWC -> CHW, normalize to [0, 1], add batch dim
+    (INTER_AREA); the result is resized back to the native crop size. Preserves
-    chw = np.transpose(resized, (2, 0, 1)).astype(np.float32) / 255.0
+    the model's I/O contract: image in as float32 [0,1] NCHW, output in [0,255].
-    return chw[np.newaxis, ...]  # (1, 3, 512, 512)
-
-
-def _preprocess_mask(mask_array):
-    """Convert HW uint8 grayscale mask to NC(1)HW float32 binary at MODEL_SIZE."""
-    import cv2
-    import numpy as np
-
-    resized = cv2.resize(mask_array, (MODEL_SIZE, MODEL_SIZE), interpolation=cv2.INTER_NEAREST)
-    # Threshold to binary 0/1
-    binary = (resized > 127).astype(np.float32)
-    return binary[np.newaxis, np.newaxis, ...]  # (1, 1, 512, 512)
-
-
-def _feathered_composite(original, inpainted, mask, feather_radius=5):
-    """Composite inpainted region into original using a feathered mask.
-
-    This preserves full quality in non-masked areas and smoothly blends
-    the inpainted region at the boundary.
     """
     import cv2
     import numpy as np
 
-    # Dilate mask slightly for smoother transition
+    def run_model(crop_img, crop_mask):
-    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (feather_radius, feather_radius))
+        h, w = crop_img.shape[:2]
-    dilated = cv2.dilate(mask.astype(np.uint8), kernel, iterations=1)
+        interp = cv2.INTER_AREA if (w > MODEL_SIZE or h > MODEL_SIZE) else cv2.INTER_LINEAR
+        img_resized = cv2.resize(crop_img, (MODEL_SIZE, MODEL_SIZE), interpolation=interp)
+        mask_resized = cv2.resize(
+            crop_mask, (MODEL_SIZE, MODEL_SIZE), interpolation=cv2.INTER_NEAREST
+        )
 
-    # Gaussian blur the dilated mask for feathering
+        img_in = np.transpose(img_resized, (2, 0, 1)).astype(np.float32)[np.newaxis] / 255.0
-    blur_size = feather_radius * 2 + 1
+        mask_in = (mask_resized > 127).astype(np.float32)[np.newaxis, np.newaxis]
-    alpha = cv2.GaussianBlur(dilated.astype(np.float32), (blur_size, blur_size), 0)
-    alpha = np.clip(alpha, 0.0, 1.0)
 
-    # Expand alpha to 3 channels
+        out = session.run(None, {"image": img_in, "mask": mask_in})[0][0]
-    alpha_3ch = alpha[:, :, np.newaxis]
+        out = np.clip(np.transpose(out, (1, 2, 0)), 0, 255).astype(np.uint8)
+        return cv2.resize(out, (w, h), interpolation=cv2.INTER_LANCZOS4)
 
-    # Composite: original * (1 - alpha) + inpainted * alpha
+    return run_model
-    result = (original.astype(np.float32) * (1.0 - alpha_3ch) +
-              inpainted.astype(np.float32) * alpha_3ch)
-    return np.clip(result, 0, 255).astype(np.uint8)
 
 
 def main():
@@ -100,10 +185,14 @@ def main():
 
         try:
             import cv2
-            import onnxruntime
+            import onnxruntime  # noqa: F401
         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."
+            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}",
@@ -119,44 +208,21 @@ def main():
         emit_progress(20, "Loading images")
         img = Image.open(input_path).convert("RGB")
         mask = Image.open(mask_path).convert("L")
-
-        orig_w, orig_h = img.size
         img_array = np.array(img)
         mask_array = np.array(mask)
 
-        # Resize mask to match image if needed
+        # Resize mask to match the image if the client sent a different size.
         if mask_array.shape[:2] != img_array.shape[:2]:
             mask_array = cv2.resize(
-                mask_array, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST
+                mask_array,
+                (img_array.shape[1], img_array.shape[0]),
+                interpolation=cv2.INTER_NEAREST,
             )
 
-        # Threshold mask to binary
+        result = inpaint_array(
-        _, mask_binary = cv2.threshold(mask_array, 127, 255, cv2.THRESH_BINARY)
+            img_array, mask_array, _make_run_model(session), progress=emit_progress
-
-        emit_progress(30, "Preprocessing")
-        img_input = _preprocess_image(img_array)
-        mask_input = _preprocess_mask(mask_binary)
-
-        emit_progress(40, "Erasing objects")
-        outputs = session.run(
-            None,
-            {"image": img_input, "mask": mask_input},
         )
 
-        emit_progress(75, "Compositing")
-        # Output shape: (1, 3, 512, 512) with values in [0, 255]
-        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 inpainted result back to original dimensions
-        inpainted_full = cv2.resize(raw_output, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
-
-        # Feathered composite: preserve quality outside mask, blend at edges
-        mask_full = mask_binary.astype(np.float32) / 255.0
-        feather_r = max(3, min(orig_w, orig_h) // 200)
-        result = _feathered_composite(img_array, inpainted_full, mask_full, feather_r)
-
         emit_progress(90, "Saving")
         Image.fromarray(result).save(output_path)
 

packages/ai/python/tests/test_inpaint_geometry.py

@@ -0,0 +1,136 @@
+"""Pure-geometry unit tests for the crop-and-composite inpaint pipeline.
+
+No ONNX model needed: the pipeline's model step is injected as a fake. Skips
+cleanly where the AI env (numpy/cv2) is absent, as on CI integration shards.
+"""
+import os
+import sys
+
+import pytest
+
+np = pytest.importorskip("numpy")
+cv2 = pytest.importorskip("cv2")
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+import inpaint  # noqa: E402
+
+
+def test_dilate_mask_grows_the_region():
+    mask = np.zeros((100, 100), np.uint8)
+    mask[40:60, 40:60] = 255
+    before = int((mask > 0).sum())
+    grown = inpaint.dilate_mask(mask, 5)
+    assert int((grown > 0).sum()) > before
+    # original masked pixels stay masked
+    assert np.all(grown[mask > 0] == 255)
+
+
+def test_dilate_mask_zero_is_noop():
+    mask = np.zeros((20, 20), np.uint8)
+    mask[5:10, 5:10] = 255
+    assert np.array_equal(inpaint.dilate_mask(mask, 0), mask)
+
+
+def test_mask_bbox_none_for_empty():
+    assert inpaint._mask_bbox(np.zeros((10, 10), np.uint8)) is None
+
+
+def test_mask_bbox_tight_bounds():
+    mask = np.zeros((100, 100), np.uint8)
+    mask[30:41, 20:51] = 255  # rows 30..40, cols 20..50
+    assert inpaint._mask_bbox(mask) == (20, 30, 51, 41)
+
+
+def test_crop_box_small_for_small_mask_in_large_image():
+    # The HD property: a small object in a big frame yields a small crop.
+    mask = np.zeros((1800, 2600), np.uint8)
+    cv2.circle(mask, (1400, 1000), 60, 255, -1)
+    mdil = inpaint.dilate_mask(mask, 8)
+    box = inpaint.compute_crop_box(mdil, mask.shape, inpaint.MARGIN_FRAC, inpaint.MARGIN_MIN)
+    area = (box[2] - box[0]) * (box[3] - box[1])
+    assert area < 0.05 * (2600 * 1800)
+
+
+def test_crop_box_clamps_to_frame_for_full_mask():
+    mask = np.full((300, 400), 255, np.uint8)
+    mdil = inpaint.dilate_mask(mask, 8)
+    box = inpaint.compute_crop_box(mdil, mask.shape, inpaint.MARGIN_FRAC, inpaint.MARGIN_MIN)
+    assert box == (0, 0, 400, 300)
+
+
+def test_composite_is_byte_identical_outside_the_feathered_mask():
+    rng = np.random.RandomState(0)
+    original = rng.randint(0, 256, (200, 200, 3), np.uint8)
+    box = (50, 50, 150, 150)
+    # a small dilated mask in the middle of the crop
+    mask_dil_native = np.zeros((100, 100), np.uint8)
+    cv2.circle(mask_dil_native, (50, 50), 20, 255, -1)
+    inpainted_crop = np.full((100, 100, 3), 255, np.uint8)  # obvious fill
+    out = inpaint.composite(original, inpainted_crop, mask_dil_native, box, feather=3)
+    # far from the mask -> untouched original, exactly
+    assert np.array_equal(out[0:40, 0:40], original[0:40, 0:40])
+    # mask core -> replaced by the fill
+    assert out[100, 100, 0] > 200
+
+
+def _scene(w=600, h=400, cx=300, cy=200, r=60):
+    """A teal disc (anti-aliased edge) on a smooth gradient. Returns img, mask, bg."""
+    yy, xx = np.mgrid[0:h, 0:w]
+    bg = np.stack(
+        [120 + 60 * xx / w, 100 + 50 * yy / h, 150 - 40 * xx / w], axis=-1
+    ).astype(np.uint8)
+    dist = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2)
+    alpha = np.clip((r - dist + 1.5) / 3.0, 0, 1)[..., None]
+    obj = np.array([20, 200, 180], np.float32)
+    img = (bg * (1 - alpha) + obj[None, None, :] * alpha).astype(np.uint8)
+    mask = (dist <= r).astype(np.uint8) * 255
+    return img, mask, bg
+
+
+def _fake_run_model(crop_img, crop_mask):
+    """Stand-in for LaMa: fill masked pixels with the mean of the unmasked crop."""
+    out = crop_img.copy()
+    m = crop_mask > 127
+    if m.any() and (~m).any():
+        out[m] = crop_img[~m].reshape(-1, 3).mean(axis=0).astype(np.uint8)
+    return out
+
+
+def test_inpaint_array_removes_the_object_ghost():
+    img, mask, bg = _scene()
+    out = inpaint.inpaint_array(img, mask, _fake_run_model)
+    core = mask > 127
+    err_before = np.abs(img[core].astype(np.float32) - bg[core].astype(np.float32)).mean()
+    err_after = np.abs(out[core].astype(np.float32) - bg[core].astype(np.float32)).mean()
+    # the object was very different from bg; after erasing, the core is close to bg
+    assert err_after < err_before * 0.5
+    assert err_after < 25
+    # and the fill is NOT the object's green (200) -> no ghost
+    assert out[core][:, 1].mean() < 175
+
+
+def test_inpaint_array_byte_identical_far_from_mask():
+    img, mask, _ = _scene()
+    out = inpaint.inpaint_array(img, mask, _fake_run_model)
+    yy, xx = np.mgrid[0:400, 0:600]
+    far = np.sqrt((xx - 300) ** 2 + (yy - 200) ** 2) > 120
+    assert np.array_equal(out[far], img[far])
+
+
+def test_inpaint_array_empty_mask_returns_original():
+    img, _, _ = _scene()
+    mask = np.zeros(img.shape[:2], np.uint8)
+    assert np.array_equal(inpaint.inpaint_array(img, mask, _fake_run_model), img)
+
+
+def test_inpaint_array_tiny_image_returns_original():
+    img = np.random.RandomState(2).randint(0, 256, (4, 4, 3), np.uint8)
+    mask = np.full((4, 4), 255, np.uint8)
+    assert np.array_equal(inpaint.inpaint_array(img, mask, _fake_run_model), img)
+
+
+def test_inpaint_array_full_mask_does_not_crash():
+    img = np.random.RandomState(3).randint(0, 256, (120, 160, 3), np.uint8)
+    mask = np.full((120, 160), 255, np.uint8)
+    out = inpaint.inpaint_array(img, mask, _fake_run_model)
+    assert out.shape == img.shape