use crate::color::rgb_to_hsv;
#[derive(Debug, Clone, Copy)]
pub struct ChromaKeyConfig {
pub target_hue: f32,
pub hue_tolerance: f32,
pub min_saturation: f32,
pub min_value: f32,
pub spill_suppression: f32,
}
impl Default for ChromaKeyConfig {
fn default() -> Self {
Self {
target_hue: 120.0, hue_tolerance: 30.0,
min_saturation: 0.2,
min_value: 0.2,
spill_suppression: 0.5,
}
}
}
impl ChromaKeyConfig {
pub fn green() -> Self {
Self::default()
}
pub fn blue() -> Self {
Self {
target_hue: 240.0,
..Self::default()
}
}
}
pub fn apply_chroma_key(row: &mut [u8], config: &ChromaKeyConfig) {
for pixel in row.chunks_exact_mut(4) {
let hsv = rgb_to_hsv(pixel[0], pixel[1], pixel[2]);
let hue_diff = hue_distance(hsv.h, config.target_hue);
if hue_diff < config.hue_tolerance
&& hsv.s >= config.min_saturation
&& hsv.v >= config.min_value
{
let hue_strength = 1.0 - (hue_diff / config.hue_tolerance);
let sat_strength =
((hsv.s - config.min_saturation) / (1.0 - config.min_saturation)).clamp(0.0, 1.0);
let match_strength = hue_strength * sat_strength;
let new_alpha = (pixel[3] as f32 * (1.0 - match_strength))
.round()
.clamp(0.0, 255.0) as u8;
pixel[3] = new_alpha;
if config.spill_suppression > 0.0 {
let gray = crate::color::rgba_to_gray(pixel[0], pixel[1], pixel[2]) as f32;
let factor = 1.0 - (match_strength * config.spill_suppression);
for c in &mut pixel[..3] {
let v = gray + (*c as f32 - gray) * factor;
*c = v.round().clamp(0.0, 255.0) as u8;
}
}
}
}
}
fn hue_distance(h1: f32, h2: f32) -> f32 {
let diff = (h1 - h2).abs();
if diff > 180.0 {
360.0 - diff
} else {
diff
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pure_green_keyed() {
let mut row = [0, 255, 0, 255]; apply_chroma_key(&mut row, &ChromaKeyConfig::green());
assert!(
row[3] < 30,
"Pure green should be mostly transparent, got alpha={}",
row[3]
);
}
#[test]
fn test_red_not_keyed() {
let mut row = [255, 0, 0, 255];
apply_chroma_key(&mut row, &ChromaKeyConfig::green());
assert_eq!(row[3], 255, "Red should not be keyed");
}
#[test]
fn test_blue_screen() {
let mut row = [0, 0, 255, 255]; apply_chroma_key(&mut row, &ChromaKeyConfig::blue());
assert!(
row[3] < 30,
"Pure blue should be mostly transparent with blue config"
);
}
#[test]
fn test_dark_green_not_keyed() {
let mut row = [0, 20, 0, 255];
apply_chroma_key(&mut row, &ChromaKeyConfig::green());
assert_eq!(
row[3], 255,
"Dark green should not be keyed (below min_value)"
);
}
#[test]
fn test_hue_distance() {
assert!((hue_distance(10.0, 350.0) - 20.0).abs() < 1e-5);
assert!((hue_distance(0.0, 180.0) - 180.0).abs() < 1e-5);
assert!((hue_distance(120.0, 120.0)).abs() < 1e-5);
}
}