corridorkey-green-screen

What this skill does

# CorridorKey Green Screen Keying

> Skill by [ara.so](https://ara.so) — Daily 2026 Skills collection.

CorridorKey is a neural network that solves the color *unmixing* problem in green screen footage. For every pixel — including semi-transparent ones from motion blur, hair, or out-of-focus edges — it predicts the true straight (un-premultiplied) foreground color and a clean linear alpha channel. It reads/writes 16-bit and 32-bit EXR files for VFX pipeline integration.

## How It Works

Two inputs required per frame:
1. **RGB green screen image** — sRGB or linear gamma, sRGB/REC709 gamut
2. **Alpha Hint** — rough coarse B&W mask (doesn't need to be precise)

The model fills in fine detail from the hint; it's trained on blurry/eroded masks.

## Installation

### Prerequisites
- [uv](https://docs.astral.sh/uv/) package manager (handles Python automatically)
- NVIDIA GPU with CUDA 12.8+ drivers (for GPU), or Apple M1+ (for MLX), or CPU fallback

### Windows
```bat
# Double-click or run from terminal:
Install_CorridorKey_Windows.bat

# Optional heavy modules:
Install_GVM_Windows.bat
Install_VideoMaMa_Windows.bat
```

### Linux / macOS
```bash
# Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh

# Install dependencies — pick one:
uv sync                  # CPU / Apple MPS (universal)
uv sync --extra cuda     # NVIDIA GPU (Linux/Windows)
uv sync --extra mlx      # Apple Silicon MLX

# Download required model (~300MB)
mkdir -p CorridorKeyModule/checkpoints
# Place downloaded CorridorKey_v1.0.pth as:
# CorridorKeyModule/checkpoints/CorridorKey.pth
```

Model download: https://huggingface.co/nikopueringer/CorridorKey_v1.0/resolve/main/CorridorKey_v1.0.pth

### Optional Alpha Hint Generators
```bash
# GVM (automatic, ~80GB VRAM, good for people)
uv run hf download geyongtao/gvm --local-dir gvm_core/weights

# VideoMaMa (requires mask hint, <24GB VRAM with community tweaks)
uv run hf download SammyLim/VideoMaMa \
  --local-dir VideoMaMaInferenceModule/checkpoints/VideoMaMa

uv run hf download stabilityai/stable-video-diffusion-img2vid-xt \
  --local-dir VideoMaMaInferenceModule/checkpoints/stable-video-diffusion-img2vid-xt \
  --include "feature_extractor/*" "image_encoder/*" "vae/*" "model_index.json"
```

## Key CLI Commands

```bash
# Run inference on prepared clips
uv run python main.py run_inference --device cuda
uv run python main.py run_inference --device cpu
uv run python main.py run_inference --device mps   # Apple Silicon

# List available clips/shots
uv run python main.py list

# Interactive setup wizard
uv run python main.py wizard
uv run python main.py wizard --win_path /path/to/ClipsForInference
```

## Docker (Linux + NVIDIA GPU)

```bash
# Build
docker build -t corridorkey:latest .

# Run inference
docker run --rm -it --gpus all \
  -e OPENCV_IO_ENABLE_OPENEXR=1 \
  -v "$(pwd)/ClipsForInference:/app/ClipsForInference" \
  -v "$(pwd)/Output:/app/Output" \
  -v "$(pwd)/CorridorKeyModule/checkpoints:/app/CorridorKeyModule/checkpoints" \
  corridorkey:latest run_inference --device cuda

# Docker Compose
docker compose build
docker compose --profile gpu run --rm corridorkey run_inference --device cuda
docker compose --profile gpu run --rm corridorkey list

# Pin to specific GPU on multi-GPU systems
NVIDIA_VISIBLE_DEVICES=0 docker compose --profile gpu run --rm corridorkey run_inference --device cuda
```

## Directory Structure

```
CorridorKey/
├── ClipsForInference/          # Input shots go here
│   └── my_shot/
│       ├── frames/             # Green screen RGB frames (PNG/EXR)
│       ├── alpha_hints/        # Coarse alpha masks (grayscale)
│       └── VideoMamaMaskHint/  # Optional: hand-drawn hints for VideoMaMa
├── Output/                     # Processed results
│   └── my_shot/
│       ├── foreground/         # Straight RGBA EXR frames
│       └── alpha/              # Linear alpha channel frames
├── CorridorKeyModule/
│   └── checkpoints/
│       └── CorridorKey.pth     # Required model weights
├── gvm_core/weights/           # Optional GVM weights
└── VideoMaMaInferenceModule/
    └── checkpoints/            # Optional VideoMaMa weights
```

## Python Usage Examples

### Basic Inference Pipeline
```python
import torch
from pathlib import Path
from CorridorKeyModule.model import CorridorKeyModel  # adjust to actual module path
from CorridorKeyModule.inference import run_inference

# Load model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CorridorKeyModel()
model.load_state_dict(torch.load("CorridorKeyModule/checkpoints/CorridorKey.pth"))
model.to(device)
model.eval()

# Run inference on a shot folder
run_inference(
    shot_dir=Path("ClipsForInference/my_shot"),
    output_dir=Path("Output/my_shot"),
    device=device,
)
```

### Reading/Writing EXR Files
```python
import cv2
import numpy as np
import os

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"

# Read a 32-bit linear EXR frame
frame = cv2.imread("frame_0001.exr", cv2.IMREAD_UNCHANGED | cv2.IMREAD_ANYCOLOR)
# frame is float32, linear light, BGR channel order

# Convert BGR -> RGB for processing
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

# Write output EXR (straight RGBA)
# Assume `foreground` is float32 HxWx4 (RGBA, linear, straight alpha)
foreground_bgra = cv2.cvtColor(foreground, cv2.COLOR_RGBA2BGRA)
cv2.imwrite("output_0001.exr", foreground_bgra.astype(np.float32))
```

### Generating a Coarse Alpha Hint with OpenCV
```python
import cv2
import numpy as np

def generate_chroma_key_hint(image_bgr: np.ndarray, erode_px: int = 5) -> np.ndarray:
    """
    Quick-and-dirty green screen hint for CorridorKey input.
    Returns grayscale mask (0=background, 255=foreground).
    """
    hsv = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2HSV)

    # Tune these ranges for your specific green screen
    lower_green = np.array([35, 50, 50])
    upper_green = np.array([85, 255, 255])

    green_mask = cv2.inRange(hsv, lower_green, upper_green)
    foreground_mask = cv2.bitwise_not(green_mask)

    # Erode to pull mask away from edges (CorridorKey handles edge detail)
    kernel = np.ones((erode_px, erode_px), np.uint8)
    eroded = cv2.erode(foreground_mask, kernel, iterations=2)

    # Optional: slight blur to soften hint
    blurred = cv2.GaussianBlur(eroded, (15, 15), 5)
    return blurred


# Usage
frame = cv2.imread("greenscreen_frame.png")
hint = generate_chroma_key_hint(frame, erode_px=8)
cv2.imwrite("alpha_hint.png", hint)
```

### Batch Processing Frames
```python
from pathlib import Path
import cv2
import numpy as np
import os

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"

def prepare_shot_folder(
    raw_frames_dir: Path,
    output_shot_dir: Path,
    hint_generator_fn=None
):
    """
    Prepares a CorridorKey shot folder from raw green screen frames.
    """
    frames_out = output_shot_dir / "frames"
    hints_out = output_shot_dir / "alpha_hints"
    frames_out.mkdir(parents=True, exist_ok=True)
    hints_out.mkdir(parents=True, exist_ok=True)

    frame_paths = sorted(raw_frames_dir.glob("*.png")) + \
                  sorted(raw_frames_dir.glob("*.exr"))

    for frame_path in frame_paths:
        frame = cv2.imread(str(frame_path), cv2.IMREAD_UNCHANGED | cv2.IMREAD_ANYCOLOR)

        # Copy frame
        cv2.imwrite(str(frames_out / frame_path.name), frame)

        # Generate hint
        if hint_generator_fn:
            hint = hint_generator_fn(frame)
        else:
            hint = generate_chroma_key_hint(frame)

        hint_name = frame_path.stem + ".png"
        cv2.imwrite(str(hints_out / hint_name), hint)

    print(f"Prepared {len(frame_paths)} frames in {output_shot_dir}")


prepare_shot_folder(
    raw_frames_dir=Path("raw_footage/shot_01"),
    output_shot_dir=Path("ClipsForInference/shot_01"),
)
```

### Using clip_manager.py Alpha Hint Generators
```python
# GVM (automatic — no extra input needed)
from clip_manager import generate_alpha_hints_gvm

generate_alpha_hints_gvm(
    shot_dir="ClipsForInference/my_shot",