Palette keying: explain each pixel as one flat color composited over the backdrop; guarantees no key color survives

app.py

@@ -22,6 +22,8 @@ import time
 import traceback
 
 import gradio as gr
+import numpy as np
+from PIL import Image
 
 import backends
 import moves as move_defs
@@ -87,6 +89,7 @@ def forge(
     despill,
     despeckle,
     shrink,
+    palette_size,
     playback_fps,
     sprite_height,
     request: gr.Request,
@@ -211,6 +214,25 @@ def forge(
     yield PENDING, character, gallery, None, [], None, f"{len(clips)} clips generated. Cutting to pose extremes…"
 
     # ---- 3. key, cut to extremes, pack ----
+    # One palette for the whole character, read before any keying. Shared across every move so the sheet cannot
+    # drift in color between rows, and seeded from the character still because that is a PNG — the only view of the
+    # art that has not been through H.264 chroma subsampling.
+    palette = None
+    if int(palette_size) > 0:
+        progress(0.72, desc="Reading the character's palette")
+        samples = []
+        still = np.asarray(Image.open(character).convert("RGB"), np.float32) / 255.0
+        samples.append(sk.sample_interior(still, sk.detect_key_color(still), 6000))
+        for clip in clips:
+            frames, _fps = sk.read_video(clip["video"])
+            clip_key = np.median(np.stack([sk.detect_key_color(f) for f in frames]), axis=0)
+            for j in {0, len(frames) // 2, len(frames) - 1}:
+                samples.append(sk.sample_interior(frames[j], clip_key, 2000, seed=j))
+            del frames
+        stacked = np.concatenate([s for s in samples if len(s)]) if any(len(s) for s in samples) else None
+        if stacked is not None and len(stacked):
+            palette = sk.extract_palette(stacked, int(palette_size))
+
     packed, previews, key_hex = [], [], None
     for i, clip in enumerate(clips):
         progress(0.75 + 0.2 * i / len(clips), desc=f"Keying {clip['label']}")
@@ -222,6 +244,7 @@ def forge(
             despill=float(despill),
             despeckle=int(despeckle),
             shrink=float(shrink),
+            palette=palette,
         )
         key_hex = "#%02X%02X%02X" % tuple((key_rgb * 255).astype(int))
 
@@ -343,6 +366,10 @@ with gr.Blocks(title="Sprite Forge") as demo:
                     with gr.Row():
                         despill = gr.Slider(0, 1, value=1.0, step=0.05, label="Despill")
                         despeckle = gr.Slider(0, 4, value=2, step=1, label="Despeckle")
+                        palette_size = gr.Slider(
+                            0, 32, value=16, step=2, label="Palette colors",
+                            info="Flat-art keying: 0 falls back to plain chroma keying.",
+                        )
                         shrink = gr.Slider(
                             0, 0.5, value=0.3, step=0.05, label="Matte shrink",
                             info="Chokes the matte inward against 4:2:0 chroma bleed.",
@@ -364,7 +391,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, shrink, playback_fps, sprite_height,
+            tolerance, softness, despill, despeckle, shrink, palette_size, playback_fps, sprite_height,
         ],
         [demo_frame, character_preview, clip_gallery, sheet_image, sprite_previews, downloads, status],
     )

spritekit.py

@@ -181,6 +181,107 @@ def extend_edges(rgb: np.ndarray, alpha: np.ndarray, radius: int = 3) -> np.ndar
     return (rgb * weight + extended * (1.0 - weight)).astype(np.float32)
 
 
+def _kmeans(points: np.ndarray, k: int, iterations: int = 12, seed: int = 0) -> np.ndarray:
+    """Small k-means with k-means++ seeding. Enough for a dozen flat colors out of a few thousand samples."""
+    rng = np.random.default_rng(seed)
+    centers = [points[rng.integers(len(points))]]
+    nearest = ((points - centers[0]) ** 2).sum(1)
+    for _ in range(k - 1):
+        total = float(nearest.sum())
+        if total < 1e-12:
+            break
+        centers.append(points[rng.choice(len(points), p=nearest / total)])
+        nearest = np.minimum(nearest, ((points - centers[-1]) ** 2).sum(1))
+    current = np.stack(centers)
+    for _ in range(iterations):
+        labels = np.argmin(((points[:, None, :] - current[None]) ** 2).sum(2), axis=1)
+        for i in range(len(current)):
+            member = labels == i
+            if member.any():
+                current[i] = points[member].mean(0)
+    return current.astype(np.float32)
+
+
+def sample_interior(
+    frame: np.ndarray,
+    key_rgb: np.ndarray,
+    limit: int = 4000,
+    seed: int = 0,
+) -> np.ndarray:
+    """Confidently-interior pixels of one frame, for palette extraction.
+
+    Eroded hard — only pixels whose whole neighbourhood is opaque qualify. The palette must describe the character's
+    real flat colors, and a single edge pixel dragged in by 4:2:0 bleed would add a magenta-tinted entry that then
+    licenses magenta everywhere.
+    """
+    alpha = matte(frame, key_rgb)
+    interior = _box_blur(alpha[..., None], 3)[..., 0] > 0.995
+    points = frame[interior]
+    if len(points) == 0:
+        return points.reshape(0, 3)
+    if len(points) > limit:
+        points = points[np.random.default_rng(seed).choice(len(points), limit, replace=False)]
+    return points
+
+
+def extract_palette(samples: np.ndarray, size: int = 16) -> np.ndarray:
+    """The character's flat palette, as ``(k, 3)`` float RGB."""
+    if len(samples) < size:
+        return samples.reshape(-1, 3).astype(np.float32)
+    return _kmeans(samples.astype(np.float32), size)
+
+
+def palette_matte(
+    frame: np.ndarray,
+    palette: np.ndarray,
+    key_rgb: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Explain every pixel as *one palette color composited over the backdrop*. Returns ``(alpha, rgb)``.
+
+    This is the only step here that can *guarantee* the backdrop is gone, and it earns that by exploiting something
+    the earlier chroma keying ignored: the source art is flat, so the character is made of a dozen known colors and
+    nothing in between. Every in-between color in the decoded clip is therefore an artifact of H.264 — a blend of a
+    real color with the backdrop — and is exactly what a composite equation describes.
+
+    For each palette color ``P`` the coverage that best explains the pixel is a projection onto the segment from the
+    backdrop to that color::
+
+        a = clamp(<observed - key, P - key> / |P - key|^2, 0, 1)
+
+    and the residual says how well it explains it. Take the palette color with the smallest residual.
+
+    Two properties follow, and they are the point:
+
+    * **A pure-backdrop pixel is explained by ``a = 0`` with zero residual, for every palette entry.** It becomes
+      fully transparent no matter what its chroma got smeared into — which is precisely the rim that survived
+      choking, un-mixing and edge extension, all of which are identities at ``a == 1``.
+    * **The output color is always a palette color**, never a blend, so no amount of backdrop can be left in it.
+
+    It also removes the failure mode of `suppress_residual`: a character that genuinely *is* near the key color has
+    that color in its palette, so it matches at ``a = 1`` and is preserved rather than eroded.
+    """
+    offset = frame - key_rgb
+    directions = palette.astype(np.float32) - key_rgb
+    lengths = np.maximum((directions ** 2).sum(1), 1e-8)
+
+    best_residual = np.full(frame.shape[:2], np.inf, np.float32)
+    best_alpha = np.zeros(frame.shape[:2], np.float32)
+    best_color = np.zeros_like(frame)
+
+    # Looped over the palette rather than broadcast: a (H, W, k, 3) tensor is hundreds of MB at 768^2 and this Space
+    # runs on 2 vCPU / 16 GB.
+    for i in range(len(palette)):
+        direction = directions[i]
+        alpha = np.clip((offset @ direction) / lengths[i], 0.0, 1.0)
+        residual = np.linalg.norm(offset - alpha[..., None] * direction, axis=-1)
+        better = residual < best_residual
+        best_residual = np.where(better, residual, best_residual)
+        best_alpha = np.where(better, alpha, best_alpha)
+        best_color = np.where(better[..., None], palette[i], best_color)
+
+    return best_alpha.astype(np.float32), best_color.astype(np.float32)
+
+
 def suppress_residual(
     rgb: np.ndarray,
     alpha: np.ndarray,
@@ -252,6 +353,7 @@ def key_frame(
     despeckle: int = 2,
     shrink: float = 0.2,
     feather: float = 0.6,
+    palette: np.ndarray | None = None,
 ) -> np.ndarray:
     """Returns RGBA float 0..1.
 
@@ -266,6 +368,17 @@ def key_frame(
 
     So: recover color against the truth, then hand out the tidied matte.
     """
+    if palette is not None and len(palette):
+        raw, matched = palette_matte(frame, palette, key_rgb)
+        refined = refine_matte(raw, despeckle, 0.0, feather)
+        # Snap only where the composite is partial. A fully opaque pixel is interior and keeps its original color,
+        # so soft shading and small details survive; the transition band takes the palette color, which is what
+        # makes a backdrop-colored edge impossible. No choke is applied — palette matting derives coverage from the
+        # full RGB rather than from chroma alone, so it does not need one.
+        keep = _smoothstep(0.90, 1.0, refined)[..., None]
+        rgb = matched * (1.0 - keep) + frame * keep
+        return np.concatenate([rgb, refined[..., None]], axis=-1).astype(np.float32)
+
     raw = matte(frame, key_rgb, tolerance, softness)
     refined = refine_matte(raw, despeckle, shrink, feather)
     rgb = unmix(frame, raw, key_rgb, despill)
@@ -286,7 +399,8 @@ def key_clip(
     despeckle: int = 2,
     shrink: float = 0.2,
     feather: float = 0.6,
-) -> tuple[list[np.ndarray], np.ndarray]:
+    palette: np.ndarray | None = None,
+) -> tuple[list[np.ndarray], np.ndarray, list[int]]:
     """Key every frame against the clip's **median** backdrop, and report which frames that backdrop still describes.
 
     MiniMax-H3 is asked to hold the backdrop flat and unchanged and mostly does — but not always. A run of this
@@ -310,7 +424,9 @@ def key_clip(
     # A well-behaved clip drifts by a couple of levels; the magenta->green slide was a distance of ~0.8. Anything
     # past this is a different backdrop, not noise.
     reliable = [i for i in range(len(frames)) if float(np.linalg.norm(measured[i] - key_rgb)) < 0.15]
-    keyed = [key_frame(f, key_rgb, tolerance, softness, despill, despeckle, shrink, feather) for f in frames]
+    keyed = [
+        key_frame(f, key_rgb, tolerance, softness, despill, despeckle, shrink, feather, palette) for f in frames
+    ]
     return keyed, key_rgb, (reliable or list(range(len(frames))))