mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
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.
This commit is contained in:
SnapOtter
+136
-70
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user