Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

8 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

CorridorKey for Nuke

A TensorRT-accelerated Nuke plugin for CorridorKey — the AI green screen keyer by Niko Pueringer / Corridor Digital that produces physically accurate unmixed foreground color and alpha.

This plugin runs the CorridorKey neural network directly inside Nuke as a native node (TRTCorridorKey), powered by NVIDIA TensorRT for real-time inference on the GPU.

What It Does

CorridorKey takes a green screen plate and a coarse alpha hint, and produces a clean straight foreground color and linear alpha matte — preserving hair, motion blur, and semi-transparent edges that traditional keyers destroy.

This plugin wraps that model into a two-input Nuke node:

  • Input 0 (plate): RGB green screen footage
  • Input 1 (mask): Coarse alpha hint (rough chroma key or AI roto)
  • Output: RGBA — unmixed FG color (sRGB) in RGB + linear alpha in A

The TensorRT engine runs at ~300ms per frame at 2048×2048 FP16 on an RTX A5000 (24 GB).

Requirements

  • Linux x86_64
  • NVIDIA GPU with 24+ GB VRAM (tested on RTX A5000)
  • CUDA 13.1 (or 12.9)
  • TensorRT 10.15.1
  • Nuke 17.0 (NDK)
  • ~128 GB+ system RAM for the ONNX export step

Build Guide

The full pipeline is: .pth → ONNX → TensorRT engine → Nuke plugin. Each step only needs to be done once.

Step 1: Export CorridorKey to ONNX

Clone the original CorridorKey repo and set up its environment:

git clone https://github.com/nikopueringer/CorridorKey.git
cd CorridorKey
uv sync
uv pip install onnx onnxruntime onnxscript

Download the model weights (~300 MB) into the weights/ folder:

mkdir -p weights
wget -O weights/CorridorKey_v1.0.pth \
    https://huggingface.co/nikopueringer/CorridorKey_v1.0/resolve/main/CorridorKey_v1.0.pth

Copy export_corridorkey_onnx.py from this repo into the CorridorKey directory, then export:

source .venv/bin/activate
uv run python export_corridorkey_onnx.py --img-size 2048 --no-verify

RAM note: The ONNX trace at 2048×2048 requires ~140 GB of memory. If you have 128 GB RAM, add temporary swap:

sudo fallocate -l 64G /swapfile
sudo chmod 600 /swapfile
sudo mkswap /swapfile
sudo swapon /swapfile
sync && sudo sh -c 'echo 3 > /proc/sys/vm/drop_caches'
uv run python export_corridorkey_onnx.py --img-size 2048 --no-verify
sudo swapoff /swapfile
sudo rm /swapfile

Alternatively, export at 1024 (--img-size 1024) which needs much less RAM. The model still works well — the Nuke plugin resizes the input to model resolution anyway.

This produces weights/CorridorKey_v1.0.onnx.

Step 2: Build the TensorRT Engine

Download the pre-built TensorRT 10.15.1 and extract it to /opt/TensorRT-10.15.1.29 (or wherever you prefer).

export LD_LIBRARY_PATH=/opt/TensorRT-10.15.1.29/lib:/usr/local/cuda-13.1/lib64:$LD_LIBRARY_PATH

/opt/TensorRT-10.15.1.29/bin/trtexec \
    --onnx=weights/CorridorKey_v1.0.onnx \
    --saveEngine=weights/CorridorKey_v1.0_fp16.engine \
    --fp16 \
    --memPoolSize=workspace:2G

FP16 is recommended — CorridorKey already uses float16 autocast during inference. This takes about 5 minutes and produces a ~198 MB engine file.

Step 3: Build the Nuke Plugin

export PATH=/usr/local/cuda-13.1/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-13.1/lib64:$LD_LIBRARY_PATH

rm -rf build && mkdir build && cd build
cmake .. \
    -DNUKE_ROOT=/opt/Nuke17.0v1 \
    -DTENSORRT_ROOT=/opt/TensorRT-10.15.1.29
make -j8

Copy TRTCorridorKey.so to your Nuke plugin path (e.g. ~/.nuke/).

Usage in Nuke

  1. Create a TRTCorridorKey node (found under AI menu)
  2. Connect your green screen plate to the plate input
  3. Connect a coarse alpha hint to the mask input (a rough Keylight/Primatte key or AI roto works. Shuffle to Alpha to RGBA)
  4. Point the Engine File knob to your .engine file
  5. Set Resolution to match your engine (2048×2048 or 1024×1024)

Knobs

Knob Description
Engine File Path to the TensorRT .engine file
Resolution Must match the ONNX export resolution
Output RGBA (FG + Alpha), Alpha Only, or FG Only
Invert Matte Inverts the predicted alpha
GPU Device GPU index for multi-GPU systems

Output Details

  • RGB channels: Straight (unpremultiplied) foreground color in sRGB
  • Alpha channel: Linear alpha matte

To composite properly in Nuke, convert the FG from sRGB to linear, then premultiply:

TRTCorridorKey → Colorspace (sRGB to linear) → Premult → Merge (over)

Model Details

The CorridorKey engine uses:

  • Input: [1, 4, 2048, 2048] — ImageNet-normalized RGB concatenated with the raw alpha hint
  • Outputs: alpha [1, 1, 2048, 2048] and fg [1, 3, 2048, 2048], both post-sigmoid
  • Architecture: Hiera Base Plus backbone (from timm) with dual decoder heads and a CNN refiner
  • Native resolution: 2048×2048 (positional embeddings are baked at export resolution)

Files in This Repo

├── TRTCorridorKey.cpp          # Nuke plugin source
├── CMakeLists.txt              # Build configuration
├── export_corridorkey_onnx.py  # PyTorch → ONNX export script
├── test_trt_engine.py          # Standalone TRT engine test
├── LICENSE                     # License
└── README.md                   # This file

Credits

License

This plugin code is released under the MIT License. See LICENSE for details.

The CorridorKey model weights and architecture are licensed under CC BY-NC-SA 4.0 by Corridor Digital. You must comply with their license when using the model. See the CorridorKey repository for full terms.

About

CorridorKey for Nuke

Resources

Stars

35 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages