Fix pink halo: finish the matte (despeckle/choke/feather) before un-mixing color; expose matte shrink

app.py

@@ -74,6 +74,7 @@ def forge(
     softness,
     despill,
     despeckle,
+    shrink,
     playback_fps,
     sprite_height,
     request: gr.Request,
@@ -148,6 +149,7 @@ def forge(
             softness=float(softness),
             despill=float(despill),
             despeckle=int(despeckle),
+            shrink=float(shrink),
         )
         key_hex = "#%02X%02X%02X" % tuple((key_rgb * 255).astype(int))
 
@@ -278,6 +280,10 @@ with gr.Blocks(title="Sprite Forge") as demo:
                 softness = gr.Slider(0.01, 0.25, value=0.09, step=0.01, label="Edge softness")
                 despill = gr.Slider(0, 1, value=1.0, step=0.05, label="Despill")
                 despeckle = gr.Slider(0, 4, value=2, step=1, label="Despeckle")
+                shrink = gr.Slider(
+                    0, 0.5, value=0.2, step=0.05, label="Matte shrink",
+                    info="Chokes the matte inward. The clip is 4:2:0, so chroma bleeds ~1px past the real edge.",
+                )
 
         demo_frame = gr.HTML(IDLE_PANEL, label="Playable")
         status = gr.Markdown()
@@ -296,7 +302,7 @@ with gr.Blocks(title="Sprite Forge") as demo:
         forge,
         [
             subject, uploaded, style, image_backend, canvas, selected_moves, duration, steps, seed,
-            tolerance, softness, despill, despeckle, playback_fps, sprite_height,
+            tolerance, softness, despill, despeckle, shrink, playback_fps, sprite_height,
         ],
         [demo_frame, character_preview, clip_gallery, sheet_image, sprite_previews, downloads, status],
     )

spritekit.py

@@ -87,70 +87,87 @@ def _smoothstep(edge0: float, edge1: float, x: np.ndarray) -> np.ndarray:
     return t * t * (3.0 - 2.0 * t)
 
 
-def key_out(
+def matte(
     frame: np.ndarray,
     key_rgb: np.ndarray,
     tolerance: float = 0.10,
     softness: float = 0.09,
-    despill: float = 1.0,
 ) -> np.ndarray:
-    """Chroma-key one RGB frame (float 0..1) against ``key_rgb``. Returns RGBA float 0..1.
+    """The alpha channel alone, from the chroma distance to ``key_rgb``.
 
     ``tolerance`` is the chroma radius that is fully transparent and ``softness`` the width of the ramp out to fully
     opaque, both as distances in the (Cb, Cr) plane whose useful range is about 0..0.7.
     """
     ycc = to_ycc(frame)
-    key = to_ycc(key_rgb[None, None, :])[0, 0]
-
-    key_chroma = key[1:]
-    key_norm = float(np.linalg.norm(key_chroma))
-    if key_norm < 1e-4:
+    key_chroma = to_ycc(key_rgb[None, None, :])[0, 0][1:]
+    if float(np.linalg.norm(key_chroma)) < 1e-4:
         # A neutral backdrop has no chroma to key against; nothing here can separate it from a white glove.
-        return np.concatenate([frame, np.ones((*frame.shape[:2], 1), np.float32)], axis=-1)
-    key_dir = key_chroma / key_norm
-
+        return np.ones(frame.shape[:2], np.float32)
     delta = ycc[..., 1:] - key_chroma
-    alpha = _smoothstep(tolerance, tolerance + max(softness, 1e-4), np.linalg.norm(delta, axis=-1))
-
-    out = frame
-    if despill > 0:
-        # Un-mix the backdrop out of the edge instead of merely desaturating it.
-        #
-        # An edge pixel is a real alpha composite of the character over the backdrop,
-        # ``observed = a*F + (1-a)*key``, so the character's own color is recoverable exactly:
-        # ``F = (observed - (1-a)*key) / a``. Solving it removes the fringe rather than trading it for a grey rim,
-        # and it is an identity at ``a == 1``, so opaque interior pixels are provably untouched.
-        #
-        # That last property is why this replaced a chroma-projection despill. Projection subtracts the backdrop's
-        # chroma direction from *every* pixel, and a magenta key points almost exactly along +Cr, which every warm
-        # color shares — the first version of this turned an orange cat green. Scaling it by `1 - alpha` fixed the
-        # hue but then only removed a fraction of the spill on the edge, leaving the 1px magenta halo this does not.
-        safe = np.maximum(alpha, 1e-3)[..., None]
-        unmixed = np.clip((frame - (1.0 - alpha)[..., None] * key_rgb) / safe, 0.0, 1.0)
-        # Below a few percent coverage the division is mostly amplifying noise, and those pixels are invisible
-        # anyway — fade the correction out rather than letting it produce confetti.
-        trust = np.clip(alpha / 0.15, 0.0, 1.0)[..., None] * despill
-        out = frame * (1.0 - trust) + unmixed * trust
-
-    return np.concatenate([out, alpha[..., None]], axis=-1).astype(np.float32)
-
-
-def clean_matte(rgba: np.ndarray, despeckle: int = 2, feather: float = 0.6) -> np.ndarray:
-    """Drop isolated specks out of the matte and soften the remaining edge.
-
-    A median filter is what removes single-pixel noise without eating a whisker; an erode/dilate pair of the same
-    radius would take the whisker with it.
+    return _smoothstep(tolerance, tolerance + max(softness, 1e-4), np.linalg.norm(delta, axis=-1))
+
+
+def refine_matte(alpha: np.ndarray, despeckle: int = 2, shrink: float = 0.2, feather: float = 0.6) -> np.ndarray:
+    """Despeckle, choke and soften the matte — all of it before any color is recovered.
+
+    A median filter removes single-pixel noise without eating a whisker, which an erode/dilate pair of the same
+    radius would not.
+
+    ``shrink`` is the compositor's choke, and it is what actually kills the halo. The alpha ramp is derived from
+    chroma, and the clip arrives 4:2:0 subsampled, so chroma bleeds about a pixel past the real edge and a rim of
+    genuinely-backdrop pixels comes out faintly opaque. Rescaling the ramp so everything under ``shrink`` goes to
+    zero drops that rim.
+
+    Ordering matters more than any of the individual filters here: feathering *expands* alpha outward, and if it runs
+    after the color step it hands partial opacity to pixels that are still pure backdrop — a pink rim, reintroduced
+    by the very filter meant to tidy the edge. The matte is finished here and `unmix` is told about it.
     """
-    if despeckle <= 0 and feather <= 0:
-        return rgba
-    alpha = Image.fromarray((np.clip(rgba[..., 3], 0, 1) * 255).astype(np.uint8), "L")
     if despeckle > 0:
-        alpha = alpha.filter(ImageFilter.MedianFilter(size=despeckle * 2 + 1))
+        image = Image.fromarray((np.clip(alpha, 0, 1) * 255).astype(np.uint8), "L")
+        alpha = np.asarray(image.filter(ImageFilter.MedianFilter(size=despeckle * 2 + 1)), np.float32) / 255.0
+    if shrink > 0:
+        alpha = np.clip((alpha - shrink) / max(1.0 - shrink, 1e-3), 0.0, 1.0)
     if feather > 0:
-        alpha = alpha.filter(ImageFilter.GaussianBlur(radius=feather))
-    out = rgba.copy()
-    out[..., 3] = np.asarray(alpha, np.float32) / 255.0
-    return out
+        image = Image.fromarray((np.clip(alpha, 0, 1) * 255).astype(np.uint8), "L")
+        alpha = np.asarray(image.filter(ImageFilter.GaussianBlur(radius=feather)), np.float32) / 255.0
+    return alpha.astype(np.float32)
+
+
+def unmix(frame: np.ndarray, alpha: np.ndarray, key_rgb: np.ndarray, strength: float = 1.0) -> np.ndarray:
+    """Recover the character's own color from an edge pixel, given the finished matte.
+
+    An edge pixel is a real alpha composite of the character over the backdrop, ``observed = a*F + (1-a)*key``, so
+    the character's color comes back out algebraically: ``F = (observed - (1-a)*key) / a``. That removes the fringe
+    rather than trading it for a grey rim, and it is an identity at ``a == 1``, so opaque interior pixels are
+    provably untouched.
+
+    That last property is why this replaced a chroma-projection despill. Projection subtracts the backdrop's chroma
+    direction from *every* pixel, and a magenta key points almost exactly along +Cr — an axis every warm color
+    shares — so an early version of this turned an orange cat green.
+    """
+    if strength <= 0:
+        return frame
+    safe = np.maximum(alpha, 1e-3)[..., None]
+    unmixed = np.clip((frame - (1.0 - alpha)[..., None] * key_rgb) / safe, 0.0, 1.0)
+    # Below a few percent coverage the division mostly amplifies noise, and those pixels are invisible anyway —
+    # fade the correction out rather than letting it produce confetti.
+    trust = np.clip(alpha / 0.15, 0.0, 1.0)[..., None] * strength
+    return (frame * (1.0 - trust) + unmixed * trust).astype(np.float32)
+
+
+def key_frame(
+    frame: np.ndarray,
+    key_rgb: np.ndarray,
+    tolerance: float = 0.10,
+    softness: float = 0.09,
+    despill: float = 1.0,
+    despeckle: int = 2,
+    shrink: float = 0.2,
+    feather: float = 0.6,
+) -> np.ndarray:
+    """matte -> refine -> unmix, in that order. Returns RGBA float 0..1."""
+    alpha = refine_matte(matte(frame, key_rgb, tolerance, softness), despeckle, shrink, feather)
+    return np.concatenate([unmix(frame, alpha, key_rgb, despill), alpha[..., None]], axis=-1).astype(np.float32)
 
 
 def key_clip(
@@ -160,6 +177,7 @@ def key_clip(
     softness: float = 0.09,
     despill: float = 1.0,
     despeckle: int = 2,
+    shrink: float = 0.2,
     feather: float = 0.6,
 ) -> tuple[list[np.ndarray], np.ndarray]:
     """Key every frame of a clip against **one** color measured on the first frame.
@@ -169,9 +187,7 @@ def key_clip(
     """
     if key_rgb is None:
         key_rgb = detect_key_color(frames[0])
-    keyed = [
-        clean_matte(key_out(f, key_rgb, tolerance, softness, despill), despeckle, feather) for f in frames
-    ]
+    keyed = [key_frame(f, key_rgb, tolerance, softness, despill, despeckle, shrink, feather) for f in frames]
     return keyed, key_rgb