Feat: Universal 2D Spatial Residual Inpainting (zero boundaries, perfect gradients & transparency)

app.py

@@ -693,8 +693,11 @@ def _build_text_mask(image_bgr: np.ndarray, dilate_iter: int = 3, remove_sfx: bo
     # 0: balloon, 1: qipao, 2: fangkuai, 3: changfangtiao, 4: kuangwai (free text/SFX/narration), 5: other
     bubble_classes = {0, 1, 2, 3}
     sfx_classes = {4, 5}
+    allowed_classes = bubble_classes if not remove_sfx else (bubble_classes | sfx_classes)
 
     def _check_conf(c_int: int, conf: float, base_conf: float) -> bool:
+        if c_int not in allowed_classes:
+            return False
         return conf >= base_conf
 
     def _run_yolo_pass(conf_val: float) -> np.ndarray:
@@ -882,7 +885,8 @@ def _lama_inpaint_tile(img_rgb: np.ndarray, mask: np.ndarray, size: int = 512) -
 
     for cnt in contours:
         rx, ry, rw, rh = cv2.boundingRect(cnt)
-        pad = 8  # Reduced from 24 to protect square bubble borders
+        # Smart wide context: give LaMa surrounding gradient context while keeping replacement constrained
+        pad = min(48, max(24, max(rw, rh) // 6))
         x0 = max(0, rx - pad)
         y0 = max(0, ry - pad)
         x1 = min(w, rx + rw + pad)
@@ -898,9 +902,20 @@ def _lama_inpaint_tile(img_rgb: np.ndarray, mask: np.ndarray, size: int = 512) -
         if ch == 0 or cw == 0:
             continue
 
-        # Resize to fixed 512x512 expected by LaMa ONNX model
-        crop_img_512 = cv2.resize(crop_img, (size, size), interpolation=cv2.INTER_CUBIC)
-        crop_mask_512 = cv2.resize(crop_mask, (size, size), interpolation=cv2.INTER_NEAREST)
+        # Aspect-Ratio Preserved Letterboxing:
+        # Instead of squishing rectangles to 512x512, pad to square S x S to preserve isotropic gradients
+        S = max(ch, cw)
+        pad_top = (S - ch) // 2
+        pad_bottom = S - ch - pad_top
+        pad_left = (S - cw) // 2
+        pad_right = S - cw - pad_left
+
+        crop_img_sq = cv2.copyMakeBorder(crop_img, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_REFLECT_101)
+        crop_mask_sq = cv2.copyMakeBorder(crop_mask, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=0)
+
+        # Resize square to 512x512 expected by LaMa ONNX model
+        crop_img_512 = cv2.resize(crop_img_sq, (size, size), interpolation=cv2.INTER_CUBIC)
+        crop_mask_512 = cv2.resize(crop_mask_sq, (size, size), interpolation=cv2.INTER_NEAREST)
 
         crop_img_t = crop_img_512.transpose(2, 0, 1)[np.newaxis].astype(np.float32) / 255.0
         crop_mask_t = (crop_mask_512[np.newaxis, np.newaxis] > 127).astype(np.float32)
@@ -911,25 +926,58 @@ def _lama_inpaint_tile(img_rgb: np.ndarray, mask: np.ndarray, size: int = 512) -
         })[0]
 
         out_img_512 = np.clip(out[0].transpose(1, 2, 0) * 255.0, 0.0, 255.0).astype(np.uint8)
-        out_img_orig = cv2.resize(out_img_512, (cw, ch), interpolation=cv2.INTER_CUBIC)
+        out_img_sq = cv2.resize(out_img_512, (S, S), interpolation=cv2.INTER_CUBIC)
+
+        # Unpad back to original (ch, cw)
+        out_img_orig = out_img_sq[pad_top : pad_top + ch, pad_left : pad_left + cw]
 
-        # SMART FEATHER: hard mask for ALL inner text pixels (no ghost text),
-        # + thin 2px soft feather ONLY at the outer boundary (no sharp visible edges).
-        # Step 1: inner mask = pixels that were definitely text (hard replacement)
         inner_mask = (crop_mask > 127).astype(np.uint8)
-        # Step 2: compute 2px outer boundary ring of the mask
-        boundary_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
-        dilated_boundary = cv2.dilate(inner_mask, boundary_kernel, iterations=1)
-        boundary_ring = (dilated_boundary - inner_mask).astype(np.float32)  # only outer edge
-        # Step 3: soft feather weight for the boundary ring only (0→1 over 2px)
-        feather = cv2.GaussianBlur(boundary_ring, (5, 5), 1.0)
-        feather_3ch = np.stack([feather, feather, feather], axis=-1)
-        # Step 4: inner = hard LaMa, boundary = soft blend, outside = original
-        inner_3ch = np.stack([inner_mask.astype(np.float32)] * 3, axis=-1)
+        if inner_mask.max() == 0:
+            continue
+
+        # 2D SPATIAL RESIDUAL HARMONIZATION (Universal Gradient & Transparency Support):
+        # Computes 2D color residual at surrounding boundary and interpolates it smoothly across the hole
+        context_k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
+        dilated_ctx = cv2.dilate(inner_mask, context_k, iterations=2)
+        context_ring = (dilated_ctx > 0) & (inner_mask == 0)
+
+        # Edge-Aware Filtering: Exclude black stroke borders/outlines from the context ring
+        crop_gray = cv2.cvtColor(crop_img, cv2.COLOR_RGB2GRAY)
+        edge_map = cv2.Canny(crop_gray, 40, 120)
+        dilated_edge = cv2.dilate(edge_map, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)))
+        clean_context_ring = context_ring & (dilated_edge == 0)
+        if clean_context_ring.sum() > 20:
+            active_context = clean_context_ring
+        else:
+            active_context = context_ring
+
+        if np.any(active_context):
+            # Compute 2D residual field
+            raw_residual = np.zeros_like(crop_img, dtype=np.float32)
+            raw_residual[active_context] = crop_img[active_context].astype(np.float32) - out_img_orig[active_context].astype(np.float32)
+            
+            # Propagate 2D smooth error into the mask area using Telea inpainting on the residual field per channel
+            # Shift by +128 to keep float values in valid uint8 range for cv2.inpaint
+            res_uint8 = np.clip(raw_residual + 128.0, 0, 255).astype(np.uint8)
+            smooth_res_uint8 = cv2.inpaint(res_uint8, inner_mask, inpaintRadius=15, flags=cv2.INPAINT_TELEA)
+            smooth_residual = smooth_res_uint8.astype(np.float32) - 128.0
+            
+            out_img_calibrated = np.clip(out_img_orig.astype(np.float32) + smooth_residual, 0.0, 255.0)
+        else:
+            out_img_calibrated = out_img_orig.astype(np.float32)
+
+        # SEAMLESS BOUNDARY TRANSITION:
+        # Distance transform inside mask
+        dist_in = cv2.distanceTransform(inner_mask, cv2.DIST_L2, 3)
+        # Weight map: 0.0 at outer edge -> 1.0 at >= 2.0px inside
+        weight = np.clip(dist_in / 2.0, 0.0, 1.0)
+        # Smooth Hermite curve (3t^2 - 2t^3) for seamless C1 gradient continuity
+        smooth_weight = weight * weight * (3.0 - 2.0 * weight)
+        weight_3ch = np.stack([smooth_weight] * 3, axis=-1)
+
         result_crop = (
-            inner_3ch * out_img_orig.astype(np.float32)
-            + feather_3ch * out_img_orig.astype(np.float32)
-            + (1.0 - inner_3ch - feather_3ch) * crop_img.astype(np.float32)
+            weight_3ch * out_img_calibrated
+            + (1.0 - weight_3ch) * crop_img.astype(np.float32)
         )
         img_out[y0:y1, x0:x1] = np.clip(result_crop, 0, 255).astype(np.uint8)