External/SnapOtter
1
0
Fork 0
mirror of https://github.com/snapotter-hq/SnapOtter.git synced 2026-08-03 07:46:42 +02:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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
2026-07-11 22:27:35 +08:00
committed by GitHub
parent cb5db59f77
commit 380419dd06
2 changed files with 272 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ai/python/inpaint.py
+136 -70
View File
@@ -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)