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/// Binary PPM (`P6`, maxval ≤ 255). `ffmpeg -i x.mp4 frames/%06d.ppm` is the
118/// cheapest way to hand a video to the frame-directory source, and the
119/// `image` crate is built here without its PNM decoder.
120fn decode_ppm(bytes: &[u8]) -> Result<RgbFrame, String> {
121    let mut pos = 2usize;
122    let mut fields = [0usize; 3];
123    for field in &mut fields {
124        loop {
125            while pos < bytes.len() && bytes[pos].is_ascii_whitespace() {
126                pos += 1;
127            }
128            if pos < bytes.len() && bytes[pos] == b'#' {
129                while pos < bytes.len() && bytes[pos] != b'\n' {
130                    pos += 1;
131                }
132                continue;
133            }
134            break;
135        }
136        let start = pos;
137        while pos < bytes.len() && bytes[pos].is_ascii_digit() {
138            pos += 1;
139        }
140        if start == pos {
141            return Err("PPM header is malformed".into());
142        }
143        *field = std::str::from_utf8(&bytes[start..pos])
144            .ok()
145            .and_then(|s| s.parse().ok())
146            .ok_or_else(|| "PPM header number is malformed".to_string())?;
147    }
148    // Exactly one whitespace byte separates the header from the raster.
149    pos += 1;
150    let [w, h, maxval] = fields;
151    if maxval == 0 || maxval > 255 {
152        return Err(format!("PPM maxval {maxval} is not supported (8-bit only)"));
153    }
154    let n = w
155        .checked_mul(h)
156        .and_then(|n| n.checked_mul(3))
157        .ok_or_else(|| "PPM size overflow".to_string())?;
158    let raster = bytes
159        .get(pos..pos + n)
160        .ok_or_else(|| format!("PPM raster is truncated: need {n} bytes"))?;
161    let data = if maxval == 255 {
162        raster.to_vec()
163    } else {
164        raster
165            .iter()
166            .map(|&v| ((v as u32 * 255 + maxval as u32 / 2) / maxval as u32).min(255) as u8)
167            .collect()
168    };
169    RgbFrame::new(w, h, data)
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175
176    #[test]
177    fn ppm_round_trip() {
178        let mut bytes = b"P6\n# c\n2 1\n255\n".to_vec();
179        bytes.extend_from_slice(&[1, 2, 3, 4, 5, 6]);
180        let f = decode_rgb(&bytes).unwrap();
181        assert_eq!((f.width, f.height), (2, 1));
182        assert_eq!(f.data, vec![1, 2, 3, 4, 5, 6]);
183    }
184
185    #[test]
186    fn data_url_decodes() {
187        let rec = serde_json::json!({"url": "data:image/png;base64,AAEC"});
188        assert_eq!(load_image_bytes(&rec).unwrap(), vec![0, 1, 2]);
189    }
190}
cortiq_engine/media.rs

cortiq_engine/
media.rs
