mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
SnapOtterandGitHub
380419dd06
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.
137 lines
5.0 KiB
Python
137 lines
5.0 KiB
Python
"""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
|