cortiq_engine/
media.rs

1//! Media ingress shared by the vision front ends: fetching image bytes from
2//! the request shapes the CLI and the OpenAI server accept, and decoding them
3//! to 8-bit RGB.
4//!
5//! `load_image_bytes` used to live in `dsv41_vision`; it moved here unchanged
6//! so MiMo (and any later tower) reads images exactly the way DeepSeek-V4.1
7//! does. `dsv41_vision::load_image_bytes` re-exports it.
8
9use base64::Engine as _;
10use serde_json::Value;
11
12/// Load bytes from raw/base64 data, Anthropic source records, data URLs,
13/// HTTP(S), or local paths.  This follows the official loader's precedence.
14pub fn load_image_bytes(record: &Value) -> Result<Vec<u8>, String> {
15    let map = record
16        .as_object()
17        .ok_or_else(|| "image record must be an object".to_string())?;
18    if let Some(data) = map.get("data") {
19        if let Some(s) = data.as_str() {
20            return base64::engine::general_purpose::STANDARD
21                .decode(s)
22                .map_err(|e| format!("invalid base64 image data: {e}"));
23        }
24    }
25    if let Some(source) = map.get("source").and_then(Value::as_object) {
26        if let Some(data) = source.get("data").and_then(Value::as_str) {
27            return base64::engine::general_purpose::STANDARD
28                .decode(data)
29                .map_err(|e| format!("invalid base64 Anthropic image data: {e}"));
30        }
31        if let Some(url) = source.get("url").and_then(Value::as_str) {
32            return load_image_bytes(&serde_json::json!({"url": url}));
33        }
34    }
35    let url = map.get("url").and_then(Value::as_str).ok_or_else(|| {
36        format!(
37            "image record has no data/source/url (keys: {:?})",
38            map.keys()
39        )
40    })?;
41    if let Some((header, payload)) = url.split_once(',').filter(|(h, _)| h.starts_with("data:")) {
42        if !header.contains(";base64") {
43            return Err(format!("unsupported data URL encoding: {header}"));
44        }
45        return base64::engine::general_purpose::STANDARD
46            .decode(payload)
47            .map_err(|e| format!("invalid data URL image: {e}"));
48    }
49    if url.starts_with("http://") || url.starts_with("https://") {
50        let response = ureq::get(url)
51            .timeout(std::time::Duration::from_secs(30))
52            .call()
53            .map_err(|e| format!("image download failed: {e}"))?;
54        let mut reader = response.into_reader();
55        let mut bytes = Vec::new();
56        std::io::Read::read_to_end(&mut reader, &mut bytes)
57            .map_err(|e| format!("image download read failed: {e}"))?;
58        return Ok(bytes);
59    }
60    let path = url.strip_prefix("file://").unwrap_or(url);
61    std::fs::read(path).map_err(|e| format!("image path '{path}' could not be read: {e}"))
62}
63
64/// An 8-bit RGB raster, row-major `[height][width][3]`.
65#[derive(Clone, Debug, PartialEq, Eq)]
66pub struct RgbFrame {
67    pub width: usize,
68    pub height: usize,
69    pub data: Vec<u8>,
70}
71
72impl RgbFrame {
73    pub fn new(width: usize, height: usize, data: Vec<u8>) -> Result<Self, String> {
74        if width == 0 || height == 0 {
75            return Err(format!(
76                "image dimensions must be non-zero, got {width}x{height}"
77            ));
78        }
79        let want = width
80            .checked_mul(height)
81            .and_then(|n| n.checked_mul(3))
82            .ok_or_else(|| "image size overflow".to_string())?;
83        if data.len() != want {
84            return Err(format!(
85                "RGB buffer has {} bytes, expected {want} for {width}x{height}",
86                data.len()
87            ));
88        }
89        Ok(Self {
90            width,
91            height,
92            data,
93        })
94    }
95}
96
97/// Decode an encoded image (PNG, JPEG, GIF, WebP, or binary PPM `P6`) to
98/// RGB. Alpha is DROPPED, not composited: vLLM and sglang's PIL path both
99/// call `convert("RGB")`, which is what `to_rgb8` does. Grayscale is
100/// replicated to three channels, as PIL does.
101pub fn decode_rgb(bytes: &[u8]) -> Result<RgbFrame, String> {
102    if bytes.starts_with(b"P6") {
103        return decode_ppm(bytes);
104    }
105    let img = image::load_from_memory(bytes).map_err(|e| format!("image decode failed: {e}"))?;
106    let rgb = img.to_rgb8();
107    let (w, h) = (rgb.width() as usize, rgb.height() as usize);
108    RgbFrame::new(w, h, rgb.into_raw())
109}
110
111/// Read a file from disk and decode it with [`decode_rgb`].
112pub fn read_rgb(path: &std::path::Path) -> Result<RgbFrame, String> {
113    let bytes = std::fs::read(path).map_err(|e| format!("{}: {e}", path.display()))?;
114    decode_rgb(&bytes).map_err(|e| format!("{}: {e}", path.display()))
115}
116
117/// Write an RGB raster as PNG (lossless, 8-bit).
118pub fn save_png(path: &std::path::Path, frame: &RgbFrame) -> Result<(), String> {
119    let img = image::RgbImage::from_raw(frame.width as u32, frame.height as u32, frame.data.clone())
120        .ok_or_else(|| "RGB buffer does not match its dimensions".to_string())?;
121    img.save(path).map_err(|e| format!("png write failed: {e}"))
122}
123
124/// Binary PPM (`P6`, maxval ≤ 255). `ffmpeg -i x.mp4 frames/%06d.ppm` is the
125/// cheapest way to hand a video to the frame-directory source, and the
126/// `image` crate is built here without its PNM decoder.
127fn decode_ppm(bytes: &[u8]) -> Result<RgbFrame, String> {
128    let mut pos = 2usize;
129    let mut fields = [0usize; 3];
130    for field in &mut fields {
131        loop {
132            while pos < bytes.len() && bytes[pos].is_ascii_whitespace() {
133                pos += 1;
134            }
135            if pos < bytes.len() && bytes[pos] == b'#' {
136                while pos < bytes.len() && bytes[pos] != b'\n' {
137                    pos += 1;
138                }
139                continue;
140            }
141            break;
142        }
143        let start = pos;
144        while pos < bytes.len() && bytes[pos].is_ascii_digit() {
145            pos += 1;
146        }
147        if start == pos {
148            return Err("PPM header is malformed".into());
149        }
150        *field = std::str::from_utf8(&bytes[start..pos])
151            .ok()
152            .and_then(|s| s.parse().ok())
153            .ok_or_else(|| "PPM header number is malformed".to_string())?;
154    }
155    // Exactly one whitespace byte separates the header from the raster.
156    pos += 1;
157    let [w, h, maxval] = fields;
158    if maxval == 0 || maxval > 255 {
159        return Err(format!("PPM maxval {maxval} is not supported (8-bit only)"));
160    }
161    let n = w
162        .checked_mul(h)
163        .and_then(|n| n.checked_mul(3))
164        .ok_or_else(|| "PPM size overflow".to_string())?;
165    let raster = bytes
166        .get(pos..pos + n)
167        .ok_or_else(|| format!("PPM raster is truncated: need {n} bytes"))?;
168    let data = if maxval == 255 {
169        raster.to_vec()
170    } else {
171        raster
172            .iter()
173            .map(|&v| ((v as u32 * 255 + maxval as u32 / 2) / maxval as u32).min(255) as u8)
174            .collect()
175    };
176    RgbFrame::new(w, h, data)
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182
183    #[test]
184    fn ppm_round_trip() {
185        let mut bytes = b"P6\n# c\n2 1\n255\n".to_vec();
186        bytes.extend_from_slice(&[1, 2, 3, 4, 5, 6]);
187        let f = decode_rgb(&bytes).unwrap();
188        assert_eq!((f.width, f.height), (2, 1));
189        assert_eq!(f.data, vec![1, 2, 3, 4, 5, 6]);
190    }
191
192    #[test]
193    fn data_url_decodes() {
194        let rec = serde_json::json!({"url": "data:image/png;base64,AAEC"});
195        assert_eq!(load_image_bytes(&rec).unwrap(), vec![0, 1, 2]);
196    }
197}
cortiq_engine/media.rs

cortiq_engine/
media.rs
