This repository contains the implementation for an industry project exploring AI-assisted green screen spill suppression and video matting. The project focuses on correcting and refining residual green colour contamination in foreground subjects after chroma key compositing.
The project was developed as a staged pipeline. Stage 1 is the main implemented solution, while Stage 2 and Stage 3 are exploratory extensions.
Green screen compositing often produces residual green spill around foreground boundaries. Traditional keying tools can generate an alpha matte, but they may still leave colour contamination on the foreground. This project investigates whether a small neural network can use the original image and alpha information to predict a cleaner foreground result.
The final practical solution is based on SpillNet, a CNN-based model designed for foreground colour correction rather than full segmentation. The alpha matte is used as guidance so that the network can focus on despill and edge refinement.
stage1_spillnet/
model_spillnet.py
loss_spillnet.py
train_spillnet.py
eval_spillnet.ipynb
eval_spillnet_video.ipynb
stage2_sam2/
generate_alpha_sam2.py
eval_sam2_pipeline.ipynb
stage3_mattingnet/
model_mattingnet.py
loss_mattingnet.py
train_mattingnet.py
infer_mattingnet.py
eval_mattingnet.ipynb
dataset/
test/
dataset.py
postprocess.py
requirements.txt
| Stage | Module | Role |
|---|---|---|
| Stage 1 | SpillNet | Main supervised green spill suppression model |
| Stage 2 | SAM2 pipeline | Exploratory automatic alpha matte generation |
| Stage 3 | MattingNet | Exploratory dual-output matting and foreground prediction model |
Create a Python environment and install the project dependencies:
pip install -r requirements.txtIf running the SAM2-based experiment, install SAM2 separately:
git clone https://github.com/facebookresearch/sam2.git
cd sam2
pip install -e .Pretrained checkpoints are available via Google Drive: [https://drive.google.com/drive/folders/1koz_3R0dEurMPCgdub8Q5XF0ioF_Jspu?usp=drive_link]
The full training dataset is not included in this repository due to size constraints. A small set of sample test images is included under:
dataset/test/
The expected dataset structure for training is:
dataset/
image/
fg/
alpha/
test/
Where:
image/contains green screen input imagesfg/contains clean foreground target imagesalpha/contains alpha matte imagestest/contains sample images for inference
To train the main SpillNet model:
python stage1_spillnet/train_spillnet.pyTo train the exploratory MattingNet model:
python stage3_mattingnet/train_mattingnet.pyModel checkpoints are saved under:
checkpoints/
Checkpoints are not included in this repository because of file size limits.
To run inference with the exploratory MattingNet model:
python stage3_mattingnet/infer_mattingnet.pyThe script reads sample images from:
dataset/test/
The repository includes notebooks for visual inspection and comparison:
stage1_spillnet/eval_spillnet.ipynb
stage1_spillnet/eval_spillnet_video.ipynb
stage2_sam2/eval_sam2_pipeline.ipynb
stage3_mattingnet/eval_mattingnet.ipynb
These notebooks are used to compare input images, alpha mattes, model outputs, and target foreground images.
The current project is a prototype rather than a production-ready keying tool. The model depends on the quality of the provided alpha matte, and the training dataset is limited in scale and diversity. Difficult cases such as fine hair, motion blur, transparent objects, and strong green reflections remain challenging.
The SAM2 and MattingNet stages are exploratory experiments. They demonstrate possible future directions, but the main working solution of the project is Stage 1 SpillNet.
Third-party libraries and models used in this project include PyTorch, OpenCV, Pillow, Matplotlib, and Segment Anything Model 2 for the exploratory alpha generation stage.