Edge-only green despill for chroma key footage

edge_despill.py

#!/usr/bin/env python3
"""
Edge-only green despill for chroma key footage.

Pipeline:
1. FFmpeg colorkey removes green background -> per-frame RGBA PNGs
2. For each frame: detect alpha edge band, despill green only within that band
3. Reassemble frames into MOV/WebM/GIF via FFmpeg
"""

import subprocess
import os
import sys
import glob
import numpy as np
from PIL import Image, ImageFilter

WORKDIR = "/tmp/edge_despill"
INPUT_GIF = "/Users/tizee/Desktop/IMG_6977.GIF"
OUTPUT_BASE = "/Users/tizee/Desktop/IMG_6977_transparent"
FPS = 20
COLORKEY_COLOR = "0x00CC25"
COLORKEY_SIMILARITY = 0.28
EDGE_WIDTH = 3  # pixels of edge band for despill


def extract_frames():
    """Extract frames with colorkey applied (no despill)."""
    frames_dir = os.path.join(WORKDIR, "frames")
    os.makedirs(frames_dir, exist_ok=True)
    subprocess.run([
        "ffmpeg", "-v", "warning", "-y",
        "-i", INPUT_GIF,
        "-vf", f"colorkey={COLORKEY_COLOR}:{COLORKEY_SIMILARITY}:0.0",
        os.path.join(frames_dir, "frame_%04d.png")
    ], check=True)
    return sorted(glob.glob(os.path.join(frames_dir, "frame_*.png")))


def edge_despill(frame_path):
    """Despill green only on edge pixels of the alpha mask."""
    img = np.array(Image.open(frame_path))
    if img.shape[2] < 4:
        return  # no alpha channel

    pil_img = Image.open(frame_path)
    alpha_channel = pil_img.split()[-1]

    alpha = img[:, :, 3]
    opaque_mask = alpha > 128

    # Use Pillow MaxFilter to dilate the transparent region inward
    # MaxFilter on inverted alpha = dilating transparent zone into opaque area
    inv_alpha = alpha_channel.point(lambda x: 255 - x)
    for _ in range(EDGE_WIDTH):
        inv_alpha = inv_alpha.filter(ImageFilter.MaxFilter(3))
    eroded_core = np.array(inv_alpha) > 128

    # Edge band = opaque pixels near the boundary of transparent pixels
    edge_band = opaque_mask & eroded_core

    # Despill: clamp green channel to max(red, blue) in edge band only
    r = img[:, :, 0].astype(np.float32)
    g = img[:, :, 1].astype(np.float32)
    b = img[:, :, 2].astype(np.float32)

    max_rb = np.maximum(r, b)
    # Only reduce green where it exceeds the max of red and blue
    g_clamped = np.where(edge_band & (g > max_rb), max_rb, g)

    img[:, :, 1] = g_clamped.astype(np.uint8)
    Image.fromarray(img).save(frame_path)


def assemble_outputs(frames_dir):
    """Reassemble processed frames into MOV, WebM, and GIF."""
    frame_pattern = os.path.join(frames_dir, "frame_%04d.png")

    # ProRes 4444 MOV
    subprocess.run([
        "ffmpeg", "-v", "warning", "-y",
        "-framerate", str(FPS),
        "-i", frame_pattern,
        "-c:v", "prores_ks", "-profile:v", "4444",
        "-pix_fmt", "yuva444p10le",
        f"{OUTPUT_BASE}.mov"
    ], check=True)

    # VP9 WebM
    subprocess.run([
        "ffmpeg", "-v", "warning", "-y",
        "-framerate", str(FPS),
        "-i", frame_pattern,
        "-c:v", "libvpx-vp9", "-pix_fmt", "yuva420p",
        "-auto-alt-ref", "0", "-b:v", "1M",
        f"{OUTPUT_BASE}.webm"
    ], check=True)

    # GIF with transparency
    subprocess.run([
        "ffmpeg", "-v", "warning", "-y",
        "-framerate", str(FPS),
        "-i", frame_pattern,
        "-vf", "split[s0][s1];[s0]palettegen=reserve_transparent=1[p];[s1][p]paletteuse=alpha_threshold=1",
        "-gifflags", "+transdiff",
        f"{OUTPUT_BASE}.gif"
    ], check=True)


def main():
    os.makedirs(WORKDIR, exist_ok=True)

    print("Extracting frames with colorkey...")
    frames = extract_frames()
    print(f"  {len(frames)} frames extracted")

    print("Applying edge-only despill...")
    for i, f in enumerate(frames):
        edge_despill(f)
        if (i + 1) % 50 == 0 or i == len(frames) - 1:
            print(f"  {i + 1}/{len(frames)}")

    print("Assembling outputs...")
    assemble_outputs(os.path.join(WORKDIR, "frames"))

    print("Done!")
    for ext in ["mov", "webm", "gif"]:
        path = f"{OUTPUT_BASE}.{ext}"
        size_mb = os.path.getsize(path) / (1024 * 1024)
        print(f"  {path}: {size_mb:.1f}MB")


if __name__ == "__main__":
    main()