plimb/codelor-color-1.2

Codelor Color 1.2

Codelor Color is a lightweight image segmentation model built specifically for chroma key: detecting and removing green from an image or video (green screens, leaves, fabric, green objects) and producing a clean alpha mask.

It takes an RGB image and a target color (green by default) and outputs a per-pixel binary mask.

---

Model at a glance

• Architecture: Lightweight Transformer with linear attention and depthwise convolutions, optimized for CPU.
• Parameters: 1,641,055
• File size: ~6.3 MB in safetensors format.
• Input: RGB image (any size) + RGB triplet of the target color.
• Output: probability mask (1 channel, same resolution as input).
• Use cases: green screen, video chroma key, green object removal, color-similarity segmentation.

---

Quick start

Installation

You can install the standalone architecture package directly from PyPI:

pip install codelor-color

Or install all the required dependencies to run this specific model:

pip install torch safetensors huggingface_hub pillow numpy

Load the model

You can use the standalone PyPI package, or load the architecture directly from this repository:

from codelor_color import CodelorColor
from safetensors.torch import load_file
import torch

# If loading from the HF repo directly:
# from modeling_codelor import load_from_pretrained
# model = load_from_pretrained("plimb/codelor-color-1.2", device="cpu")

# If using the PyPI package and local weights:
model = CodelorColor(dim=128, depth=8, heads=8, tied_after=4)
state_dict = load_file("model.safetensors")
model.load_state_dict(state_dict)
model.eval()

Predict a mask on an image

import torch
from PIL import Image
import numpy as np
from modeling_codelor import predict_mask

img = Image.open("photo.jpg").convert("RGB")
x = torch.from_numpy(np.array(img)).float().permute(2, 0, 1) / 255.0

mask, prob = predict_mask(model, x, color_rgb=(0.20, 0.85, 0.30), threshold=0.5)

print("Green surface detected: {:.1f}%".format(mask.mean() * 100))

Remove green from an image (RGBA output)

from modeling_codelor import remove_green

rgba = remove_green(x, model)
rgba_np = (rgba.permute(1, 2, 0).numpy() * 255).astype("uint8")

Image.fromarray(rgba_np, mode="RGBA").save("no_green.png")

---

Architecture

• A ConvNeXt-style convolutional stem with anti-aliased downsampling.
• 8 Transformer blocks with 2D RoPE and QK-Norm.
• A color cross-attention module injected at every block.
• A gated FPN decoder with deep supervision.
• An edge head (Sobel) to refine boundaries.
• LayerScale and DropPath.

---

Demo & Examples

Video Demo

Watch Codelor Color in action on a real video (green screen removal and transparent background):

Image Demo

---

Performance and Statistics

Model Metrics

Metric	Value
Parameters	1,641,055
File size	~6.3 MB
CPU latency (160x160)	~80-150 ms

Comparison with other approaches

Method	Parameters	IoU / mIoU	CPU latency	Notes
K-means (HSV)	-	0.642 (IoU)	very fast	Sensitive to lighting
U-Net	34.5M	0.794 (mIoU)	~290 ms	Heavy baseline
Attention U-Net	34.9M	0.845 (IoU)	~331 ms	Even heavier
BASNet	87.1M	0.561 (mIoU)	high	Very heavy
MobileSAM	10.1M	~0.80 (mIoU)	~407 ms	Promptable, heavier
Codelor Color 1.2	1.64M	to be filled	~80-150 ms	CPU-friendly

---

License

MIT.