imagegen_bridge_artifacts/
chroma.rs

1//! Deterministic chroma-key removal with bounded image decoding.
2
3use std::io::Cursor;
4
5use image::{ImageFormat, ImageReader, Limits};
6use imagegen_bridge_core::{BridgeError, ErrorCode, OutputFormat};
7
8use crate::{ImageLimits, ImageMetadata, inspect_image};
9
10const KEY_DOMINANCE_THRESHOLD: i16 = 16;
11const ALPHA_NOISE_FLOOR: u8 = 8;
12const ALPHA_SCALE: u64 = 65_535;
13
14/// RGB color used as a generated solid background.
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16pub struct ChromaKey(pub u8, pub u8, pub u8);
17
18impl ChromaKey {
19    /// Parses `#RRGGBB` or `RRGGBB`.
20    pub fn parse(value: &str) -> Result<Self, BridgeError> {
21        let value = value.strip_prefix('#').unwrap_or(value);
22        if value.len() != 6 || !value.bytes().all(|byte| byte.is_ascii_hexdigit()) {
23            return Err(chroma_error(
24                "chroma key must be a hex RGB color like #00ff00",
25            ));
26        }
27        let channel = |offset| {
28            u8::from_str_radix(&value[offset..offset + 2], 16)
29                .map_err(|_| chroma_error("chroma key contains an invalid RGB channel"))
30        };
31        Ok(Self(channel(0)?, channel(2)?, channel(4)?))
32    }
33
34    /// Lowercase CSS-compatible hexadecimal representation.
35    #[must_use]
36    pub fn hex(self) -> String {
37        format!("#{:02x}{:02x}{:02x}", self.0, self.1, self.2)
38    }
39}
40
41/// Controls for a soft chroma matte.
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub struct ChromaKeyOptions {
44    /// Known key color requested from the image model.
45    pub key: ChromaKey,
46    /// Distance at or below which a pixel becomes fully transparent.
47    pub transparent_threshold: u8,
48    /// Distance at or above which a pixel becomes fully opaque.
49    pub opaque_threshold: u8,
50    /// Remove key-colored spill from partially transparent edges.
51    pub despill: bool,
52}
53
54impl ChromaKeyOptions {
55    fn validate(self) -> Result<(), BridgeError> {
56        if self.transparent_threshold >= self.opaque_threshold {
57            return Err(chroma_error(
58                "transparent threshold must be lower than opaque threshold",
59            ));
60        }
61        Ok(())
62    }
63}
64
65impl Default for ChromaKeyOptions {
66    fn default() -> Self {
67        Self {
68            key: ChromaKey(0, 255, 0),
69            transparent_threshold: 12,
70            opaque_threshold: 96,
71            despill: true,
72        }
73    }
74}
75
76/// Alpha validation counters returned after background removal.
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub struct AlphaSummary {
79    /// Total decoded pixels.
80    pub total_pixels: u64,
81    /// Fully transparent pixels.
82    pub transparent_pixels: u64,
83    /// Partially transparent antialiasing pixels.
84    pub partial_pixels: u64,
85    /// Fully opaque pixels.
86    pub opaque_pixels: u64,
87    /// Transparent pixels among the four corners.
88    pub transparent_corners: u8,
89}
90
91/// Re-encoded alpha image and its independently verified metadata.
92#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct ChromaKeyResult {
94    /// PNG or WebP bytes with an alpha channel.
95    pub bytes: Vec<u8>,
96    /// Metadata verified after encoding.
97    pub metadata: ImageMetadata,
98    /// Matte validation counters.
99    pub alpha: AlphaSummary,
100}
101
102/// Removes a known solid chroma background and validates a usable alpha matte.
103pub fn remove_chroma_key(
104    bytes: &[u8],
105    output_format: OutputFormat,
106    options: ChromaKeyOptions,
107    limits: ImageLimits,
108) -> Result<ChromaKeyResult, BridgeError> {
109    options.validate()?;
110    if !matches!(output_format, OutputFormat::Png | OutputFormat::Webp) {
111        return Err(chroma_error(
112            "transparent output requires PNG or WebP encoding",
113        ));
114    }
115    let source = inspect_image(bytes, limits)?;
116    let decoder_format = image_format(source.format);
117    let mut reader = ImageReader::with_format(Cursor::new(bytes), decoder_format);
118    let mut decode_limits = Limits::default();
119    decode_limits.max_image_width = Some(limits.max_edge);
120    decode_limits.max_image_height = Some(limits.max_edge);
121    decode_limits.max_alloc = Some(limits.max_decode_alloc);
122    reader.limits(decode_limits);
123    let mut rgba = reader
124        .decode()
125        .map_err(|_| chroma_error("image could not be decoded for background removal"))?
126        .to_rgba8();
127
128    let key = [options.key.0, options.key.1, options.key.2];
129    for pixel in rgba.pixels_mut() {
130        let rgb = [pixel[0], pixel[1], pixel[2]];
131        let distance = channel_distance(rgb, key);
132        let key_like = looks_key_colored(rgb, key, distance);
133        let matte_alpha = if key_like {
134            soft_alpha(
135                distance,
136                options.transparent_threshold,
137                options.opaque_threshold,
138            )
139            .min(dominance_alpha(rgb, key))
140        } else {
141            255
142        };
143        let mut alpha = u8::try_from((u16::from(matte_alpha) * u16::from(pixel[3]) + 127) / 255)
144            .unwrap_or(u8::MAX);
145        if alpha <= ALPHA_NOISE_FLOOR {
146            alpha = 0;
147        }
148        if alpha == 0 {
149            *pixel = image::Rgba([0, 0, 0, 0]);
150            continue;
151        }
152        if options.despill && key_like {
153            let cleaned = cleanup_spill(rgb, key, alpha);
154            pixel[0] = cleaned[0];
155            pixel[1] = cleaned[1];
156            pixel[2] = cleaned[2];
157        }
158        pixel[3] = alpha;
159    }
160
161    let alpha = summarize_alpha(&rgba);
162    validate_alpha(alpha, true)?;
163    let mut encoded = Cursor::new(Vec::new());
164    image::DynamicImage::ImageRgba8(rgba)
165        .write_to(&mut encoded, image_format(output_format))
166        .map_err(|_| chroma_error("transparent image could not be encoded"))?;
167    let bytes = encoded.into_inner();
168    let metadata = inspect_image(&bytes, limits)?;
169    Ok(ChromaKeyResult {
170        bytes,
171        metadata,
172        alpha,
173    })
174}
175
176/// Samples the median RGB color from a bounded band around the image border.
177pub fn detect_border_chroma_key(
178    bytes: &[u8],
179    limits: ImageLimits,
180) -> Result<ChromaKey, BridgeError> {
181    let source = inspect_image(bytes, limits)?;
182    let mut reader = ImageReader::with_format(Cursor::new(bytes), image_format(source.format));
183    let mut decode_limits = Limits::default();
184    decode_limits.max_image_width = Some(limits.max_edge);
185    decode_limits.max_image_height = Some(limits.max_edge);
186    decode_limits.max_alloc = Some(limits.max_decode_alloc);
187    reader.limits(decode_limits);
188    let rgba = reader
189        .decode()
190        .map_err(|_| chroma_error("image could not be decoded for key detection"))?
191        .to_rgba8();
192    let band = rgba.width().min(rgba.height()).clamp(1, 6);
193    let step = (rgba.width().min(rgba.height()) / 256).max(1);
194    let mut red = Vec::new();
195    let mut green = Vec::new();
196    let mut blue = Vec::new();
197    let mut sample = |pixel: &image::Rgba<u8>| {
198        red.push(pixel[0]);
199        green.push(pixel[1]);
200        blue.push(pixel[2]);
201    };
202    for x in (0..rgba.width()).step_by(step as usize) {
203        for offset in 0..band {
204            sample(rgba.get_pixel(x, offset));
205            sample(rgba.get_pixel(x, rgba.height() - 1 - offset));
206        }
207    }
208    for y in (0..rgba.height()).step_by(step as usize) {
209        for offset in 0..band {
210            sample(rgba.get_pixel(offset, y));
211            sample(rgba.get_pixel(rgba.width() - 1 - offset, y));
212        }
213    }
214    red.sort_unstable();
215    green.sort_unstable();
216    blue.sort_unstable();
217    let middle = red.len() / 2;
218    if red.is_empty() {
219        return Err(chroma_error(
220            "image border did not contain sampleable pixels",
221        ));
222    }
223    Ok(ChromaKey(red[middle], green[middle], blue[middle]))
224}
225
226/// Verifies that an encoded image contains a plausible transparent-background matte.
227pub fn inspect_transparent_alpha(
228    bytes: &[u8],
229    limits: ImageLimits,
230) -> Result<AlphaSummary, BridgeError> {
231    let source = inspect_image(bytes, limits)?;
232    let mut reader = ImageReader::with_format(Cursor::new(bytes), image_format(source.format));
233    let mut decode_limits = Limits::default();
234    decode_limits.max_image_width = Some(limits.max_edge);
235    decode_limits.max_image_height = Some(limits.max_edge);
236    decode_limits.max_alloc = Some(limits.max_decode_alloc);
237    reader.limits(decode_limits);
238    let rgba = reader
239        .decode()
240        .map_err(|_| chroma_error("image could not be decoded for alpha validation"))?
241        .to_rgba8();
242    let summary = summarize_alpha(&rgba);
243    validate_alpha(summary, false)?;
244    Ok(summary)
245}
246
247fn image_format(format: OutputFormat) -> ImageFormat {
248    match format {
249        OutputFormat::Png => ImageFormat::Png,
250        OutputFormat::Jpeg => ImageFormat::Jpeg,
251        OutputFormat::Webp => ImageFormat::WebP,
252    }
253}
254
255fn channel_distance(rgb: [u8; 3], key: [u8; 3]) -> u8 {
256    rgb.into_iter()
257        .zip(key)
258        .map(|(left, right)| left.abs_diff(right))
259        .max()
260        .unwrap_or(0)
261}
262
263fn spill_channels(key: [u8; 3]) -> [bool; 3] {
264    let maximum = key.into_iter().max().unwrap_or(0);
265    if maximum < 128 {
266        return [false; 3];
267    }
268    key.map(|value| value >= maximum.saturating_sub(16) && value >= 128)
269}
270
271fn key_dominance(rgb: [u8; 3], key: [u8; 3]) -> i16 {
272    let spill = spill_channels(key);
273    if !spill.iter().any(|value| *value) {
274        return 0;
275    }
276    let key_strength = rgb
277        .iter()
278        .enumerate()
279        .filter(|(index, _)| spill[*index])
280        .map(|(_, value)| *value)
281        .min()
282        .unwrap_or(0);
283    let non_key_strength = rgb
284        .iter()
285        .enumerate()
286        .filter(|(index, _)| !spill[*index])
287        .map(|(_, value)| *value)
288        .max()
289        .unwrap_or(0);
290    i16::from(key_strength) - i16::from(non_key_strength)
291}
292
293fn looks_key_colored(rgb: [u8; 3], key: [u8; 3], distance: u8) -> bool {
294    distance <= 32 || key_dominance(rgb, key) >= KEY_DOMINANCE_THRESHOLD
295}
296
297fn soft_alpha(distance: u8, transparent: u8, opaque: u8) -> u8 {
298    if distance <= transparent {
299        return 0;
300    }
301    if distance >= opaque {
302        return 255;
303    }
304    let ratio = u64::from(distance - transparent) * ALPHA_SCALE / u64::from(opaque - transparent);
305    let smooth = ratio * ratio * (3 * ALPHA_SCALE - 2 * ratio) / (ALPHA_SCALE * ALPHA_SCALE);
306    u8::try_from((smooth * 255 + ALPHA_SCALE / 2) / ALPHA_SCALE).unwrap_or(u8::MAX)
307}
308
309fn dominance_alpha(rgb: [u8; 3], key: [u8; 3]) -> u8 {
310    let dominance = key_dominance(rgb, key);
311    if dominance <= 0 {
312        return 255;
313    }
314    let spill = spill_channels(key);
315    let non_key_strength = rgb
316        .iter()
317        .enumerate()
318        .filter(|(index, _)| !spill[*index])
319        .map(|(_, value)| *value)
320        .max()
321        .unwrap_or(0);
322    let denominator =
323        (i16::from(key.into_iter().max().unwrap_or(0)) - i16::from(non_key_strength)).max(1);
324    let remaining = denominator.saturating_sub(dominance.min(denominator));
325    u8::try_from((i32::from(remaining) * 255 + i32::from(denominator) / 2) / i32::from(denominator))
326        .unwrap_or(u8::MAX)
327}
328
329fn cleanup_spill(rgb: [u8; 3], key: [u8; 3], alpha: u8) -> [u8; 3] {
330    if alpha >= 252 {
331        return rgb;
332    }
333    let spill = spill_channels(key);
334    let anchor = rgb
335        .iter()
336        .enumerate()
337        .filter(|(index, _)| !spill[*index])
338        .map(|(_, value)| *value)
339        .max()
340        .unwrap_or(0)
341        .saturating_sub(1);
342    let mut output = rgb;
343    for (index, value) in output.iter_mut().enumerate() {
344        if spill[index] {
345            *value = (*value).min(anchor);
346        }
347    }
348    output
349}
350
351fn summarize_alpha(image: &image::RgbaImage) -> AlphaSummary {
352    let mut transparent_pixels = 0_u64;
353    let mut partial_pixels = 0_u64;
354    let mut opaque_pixels = 0_u64;
355    for pixel in image.pixels() {
356        match pixel[3] {
357            0 => transparent_pixels += 1,
358            255 => opaque_pixels += 1,
359            _ => partial_pixels += 1,
360        }
361    }
362    let maximum_x = image.width().saturating_sub(1);
363    let maximum_y = image.height().saturating_sub(1);
364    let transparent_corners = u8::try_from(
365        [
366            (0, 0),
367            (maximum_x, 0),
368            (0, maximum_y),
369            (maximum_x, maximum_y),
370        ]
371        .into_iter()
372        .filter(|(x, y)| image.get_pixel(*x, *y)[3] == 0)
373        .count(),
374    )
375    .unwrap_or(4);
376    AlphaSummary {
377        total_pixels: u64::from(image.width()) * u64::from(image.height()),
378        transparent_pixels,
379        partial_pixels,
380        opaque_pixels,
381        transparent_corners,
382    }
383}
384
385fn validate_alpha(
386    summary: AlphaSummary,
387    require_transparent_corners: bool,
388) -> Result<(), BridgeError> {
389    let minimum_region = (summary.total_pixels / 1_000).max(1);
390    if summary.transparent_pixels < minimum_region {
391        return Err(
392            chroma_error("background removal produced no meaningful transparent region")
393                .with_detail("transparent_pixels", summary.transparent_pixels)
394                .with_detail("total_pixels", summary.total_pixels),
395        );
396    }
397    if summary.opaque_pixels + summary.partial_pixels < minimum_region {
398        return Err(chroma_error(
399            "background removal erased the entire visible subject",
400        ));
401    }
402    if require_transparent_corners && summary.transparent_corners < 2 {
403        return Err(chroma_error(
404            "background removal left an opaque border around the generated subject",
405        )
406        .with_detail("transparent_corners", summary.transparent_corners));
407    }
408    Ok(())
409}
410
411fn chroma_error(message: impl Into<String>) -> BridgeError {
412    BridgeError::new(ErrorCode::Artifact, message).with_detail("stage", "transparent_background")
413}
414
415#[cfg(test)]
416mod tests {
417    #![allow(clippy::unwrap_used)]
418
419    use super::*;
420
421    fn keyed_fixture() -> Vec<u8> {
422        let mut image = image::RgbaImage::from_pixel(20, 20, image::Rgba([0, 255, 0, 255]));
423        for y in 5..15 {
424            for x in 5..15 {
425                image.put_pixel(x, y, image::Rgba([220, 30, 20, 255]));
426            }
427        }
428        let mut bytes = Cursor::new(Vec::new());
429        image::DynamicImage::ImageRgba8(image)
430            .write_to(&mut bytes, ImageFormat::Png)
431            .unwrap();
432        bytes.into_inner()
433    }
434
435    #[test]
436    fn removes_key_preserves_subject_and_normalizes_transparent_rgb() {
437        let result = remove_chroma_key(
438            &keyed_fixture(),
439            OutputFormat::Png,
440            ChromaKeyOptions::default(),
441            ImageLimits::default(),
442        )
443        .unwrap();
444        assert_eq!(result.alpha.transparent_pixels, 300);
445        assert_eq!(result.alpha.opaque_pixels, 100);
446        assert_eq!(result.alpha.transparent_corners, 4);
447        let decoded = image::load_from_memory(&result.bytes).unwrap().to_rgba8();
448        assert_eq!(decoded.get_pixel(0, 0).0, [0, 0, 0, 0]);
449        assert_eq!(decoded.get_pixel(10, 10).0, [220, 30, 20, 255]);
450    }
451
452    #[test]
453    fn rejects_invalid_thresholds_and_non_keyed_outputs() {
454        let invalid = ChromaKeyOptions {
455            transparent_threshold: 96,
456            opaque_threshold: 12,
457            ..ChromaKeyOptions::default()
458        };
459        assert!(
460            remove_chroma_key(
461                &keyed_fixture(),
462                OutputFormat::Png,
463                invalid,
464                ImageLimits::default()
465            )
466            .is_err()
467        );
468
469        let white = image::RgbaImage::from_pixel(20, 20, image::Rgba([255, 255, 255, 255]));
470        let mut bytes = Cursor::new(Vec::new());
471        image::DynamicImage::ImageRgba8(white)
472            .write_to(&mut bytes, ImageFormat::Png)
473            .unwrap();
474        assert!(
475            remove_chroma_key(
476                &bytes.into_inner(),
477                OutputFormat::Png,
478                ChromaKeyOptions::default(),
479                ImageLimits::default()
480            )
481            .is_err()
482        );
483    }
484}
imagegen_bridge_artifacts/chroma.rs

imagegen_bridge_artifacts/
chroma.rs
