From 380419dd0603693ad542082bd6197e6505f9b1e9 Mon Sep 17 00:00:00 2001 From: SnapOtter Date: Sat, 11 Jul 2026 22:27:35 +0800 Subject: [PATCH] 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 | 206 ++++++++++++------ .../ai/python/tests/test_inpaint_geometry.py | 136 ++++++++++++ 2 files changed, 272 insertions(+), 70 deletions(-) create mode 100644 packages/ai/python/tests/test_inpaint_geometry.py diff --git a/packages/ai/python/inpaint.py b/packages/ai/python/inpaint.py index 72af5e0d..99213d4c 100644 --- a/packages/ai/python/inpaint.py +++ b/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): - """Convert HWC uint8 RGB image to NCHW float32 [0,1] at MODEL_SIZE.""" - import cv2 - import numpy as np +def _make_run_model(session): + """Build a run_model(crop_img, crop_mask) that runs LaMa at its fixed 512x512. - resized = cv2.resize(img_array, (MODEL_SIZE, MODEL_SIZE), interpolation=cv2.INTER_AREA) - # HWC -> CHW, normalize to [0, 1], add batch dim - chw = np.transpose(resized, (2, 0, 1)).astype(np.float32) / 255.0 - 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. + Small crops are upscaled to 512 (INTER_LINEAR), large crops downscaled + (INTER_AREA); the result is resized back to the native crop size. Preserves + the model's I/O contract: image in as float32 [0,1] NCHW, output in [0,255]. """ import cv2 import numpy as np - # Dilate mask slightly for smoother transition - kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (feather_radius, feather_radius)) - dilated = cv2.dilate(mask.astype(np.uint8), kernel, iterations=1) + def run_model(crop_img, crop_mask): + h, w = crop_img.shape[:2] + 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 - blur_size = feather_radius * 2 + 1 - alpha = cv2.GaussianBlur(dilated.astype(np.float32), (blur_size, blur_size), 0) - alpha = np.clip(alpha, 0.0, 1.0) + img_in = np.transpose(img_resized, (2, 0, 1)).astype(np.float32)[np.newaxis] / 255.0 + mask_in = (mask_resized > 127).astype(np.float32)[np.newaxis, np.newaxis] - # Expand alpha to 3 channels - alpha_3ch = alpha[:, :, np.newaxis] + out = session.run(None, {"image": img_in, "mask": mask_in})[0][0] + 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 - result = (original.astype(np.float32) * (1.0 - alpha_3ch) + - inpainted.astype(np.float32) * alpha_3ch) - return np.clip(result, 0, 255).astype(np.uint8) + return run_model 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 - _, mask_binary = cv2.threshold(mask_array, 127, 255, cv2.THRESH_BINARY) - - 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}, + result = inpaint_array( + img_array, mask_array, _make_run_model(session), progress=emit_progress ) - 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) diff --git a/packages/ai/python/tests/test_inpaint_geometry.py b/packages/ai/python/tests/test_inpaint_geometry.py new file mode 100644 index 00000000..8ec87e3d --- /dev/null +++ b/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