combs_media/
images.rs

1//! Image preprocessing (SigLIP / Idefics3 single-image path).
2
3use crate::{MediaError, Result};
4
5/// A preprocessed image: planar CHW f32, normalized, ready for
6/// `Tensor::from_data(TensorData::new(data, [1, c, h, w]))`.
7#[derive(Debug, Clone)]
8pub struct PixelBatch {
9    /// Image width (pixels, after resize+pad).
10    pub width: usize,
11    /// Image height (pixels, after resize+pad).
12    pub height: usize,
13    /// Channel count (3 = RGB).
14    pub channels: usize,
15    /// Planar CHW data, normalized with the preprocessor's mean/std.
16    pub data: Vec<f32>,
17}
18
19impl PixelBatch {
20    /// `[batch=1, channels, height, width]` tensor shape.
21    pub fn shape(&self) -> [usize; 4] {
22        [1, self.channels, self.height, self.width]
23    }
24}
25
26/// Turns encoded image bytes (PNG/JPEG/WebP) into normalized pixel batches.
27pub trait ImagePreprocessor: Send + Sync {
28    /// Decodes and preprocesses one image.
29    fn preprocess(&self, bytes: &[u8]) -> Result<PixelBatch>;
30}
31
32/// SigLIP / Idefics3 single-image preprocessing (SmolVLM-256M/500M):
33/// RGB → resize longest edge to `image_size` (aspect preserved, bilinear) →
34/// pad to a square with 0.5 → rescale 1/255 → normalize mean/std 0.5.
35/// (Padding at the normalization mean maps to 0 after normalization.)
36pub struct SiglipPreprocessor {
37    image_size: usize,
38    mean: [f32; 3],
39    std: [f32; 3],
40}
41
42impl SiglipPreprocessor {
43    /// Preprocessor for a square `image_size` (SmolVLM: 512).
44    pub fn new(image_size: usize) -> Self {
45        SiglipPreprocessor {
46            image_size,
47            mean: [0.5, 0.5, 0.5],
48            std: [0.5, 0.5, 0.5],
49        }
50    }
51}
52
53impl ImagePreprocessor for SiglipPreprocessor {
54    fn preprocess(&self, bytes: &[u8]) -> Result<PixelBatch> {
55        let img = image::load_from_memory(bytes)
56            .map_err(|e| MediaError::Decode(e.to_string()))?
57            .to_rgb8();
58        let (w, h) = (img.width() as usize, img.height() as usize);
59        if w == 0 || h == 0 {
60            return Err(MediaError::BadShape("empty image".to_string()));
61        }
62
63        // Longest edge -> image_size, aspect preserved.
64        let size = self.image_size as f64;
65        let scale = size / w.max(h) as f64;
66        let new_w = ((w as f64 * scale).round() as usize).max(1);
67        let new_h = ((h as f64 * scale).round() as usize).max(1);
68        let resized = image::imageops::resize(
69            &img,
70            new_w as u32,
71            new_h as u32,
72            image::imageops::FilterType::Triangle,
73        );
74
75        // Pad to square with 0.5 (≈128/255), top-left anchored.
76        let pad_byte = 128u8;
77        let mut canvas = image::RgbImage::from_pixel(
78            self.image_size as u32,
79            self.image_size as u32,
80            image::Rgb([pad_byte, pad_byte, pad_byte]),
81        );
82        image::imageops::overlay(&mut canvas, &resized, 0, 0);
83
84        // Planar CHW, rescale 1/255, normalize.
85        let n = self.image_size * self.image_size;
86        let mut data = vec![0f32; 3 * n];
87        for (x, y, pixel) in canvas.enumerate_pixels() {
88            let idx = y as usize * self.image_size + x as usize;
89            for c in 0..3 {
90                let v = pixel[c] as f32 / 255.0;
91                data[c * n + idx] = (v - self.mean[c]) / self.std[c];
92            }
93        }
94
95        Ok(PixelBatch {
96            width: self.image_size,
97            height: self.image_size,
98            channels: 3,
99            data,
100        })
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    fn png_bytes(w: u32, h: u32, rgb: [u8; 3]) -> Vec<u8> {
109        let img = image::RgbImage::from_pixel(w, h, image::Rgb(rgb));
110        let mut buf = std::io::Cursor::new(Vec::new());
111        img.write_to(&mut buf, image::ImageFormat::Png).unwrap();
112        buf.into_inner()
113    }
114
115    #[test]
116    fn output_shape_and_range() {
117        let pp = SiglipPreprocessor::new(512);
118        let out = pp.preprocess(&png_bytes(64, 32, [255, 0, 0])).unwrap();
119        assert_eq!(out.shape(), [1, 3, 512, 512]);
120        assert_eq!(out.data.len(), 3 * 512 * 512);
121        let (mut lo, mut hi) = (f32::INFINITY, f32::NEG_INFINITY);
122        for &v in &out.data {
123            lo = lo.min(v);
124            hi = hi.max(v);
125        }
126        assert!(lo >= -1.0 && hi <= 1.0, "normalized range [{lo}, {hi}]");
127        // A pure-red image must saturate the R channel somewhere near +1.
128        let n = 512 * 512;
129        let r_max = out.data[..n].iter().copied().fold(f32::NEG_INFINITY, f32::max);
130        assert!(r_max > 0.9, "red channel max {r_max}");
131    }
132
133    #[test]
134    fn padding_is_zero_after_normalization() {
135        let pp = SiglipPreprocessor::new(512);
136        // Wide image: bottom half is padding.
137        let out = pp.preprocess(&png_bytes(64, 16, [0, 255, 0])).unwrap();
138        let n = 512 * 512;
139        // Sample a pixel deep in the padded region (bottom-right).
140        let pad_idx = n - 1;
141        for c in 0..3 {
142            assert!(
143                out.data[c * n + pad_idx].abs() < 0.02,
144                "pad should normalize to ~0, got {}",
145                out.data[c * n + pad_idx]
146            );
147        }
148    }
149
150    #[test]
151    fn rejects_garbage() {
152        let pp = SiglipPreprocessor::new(512);
153        assert!(pp.preprocess(b"not an image").is_err());
154    }
155}
combs_media/images.rs

combs_media/
images.rs
