rpi_tools/
image_processing.rs

1//! Image processing utilities — resize, convert, and transform images.
2//!
3//! Mirrors `packages/coding-agent/src/utils/image-*.ts`. Provides cross-platform
4//! image processing for the agent's vision capabilities.
5//!
6//! # Features
7//!
8//! - **Resize**: Scale images while maintaining aspect ratio
9//! - **Convert**: Convert between image formats (PNG, JPEG, WebP)
10//! - **EXIF handling**: Strip EXIF data for privacy
11//! - **Optimization**: Reduce file size for API transmission
12//!
13//! # Example
14//!
15//! ```rust,no_run
16//! use rpi_tools::image_processing::{process_image, ImageProcessingOptions, OutputFormat};
17//!
18//! # fn example(input_bytes: &[u8]) -> Result<(), String> {
19//! let options = ImageProcessingOptions {
20//!     max_width: Some(1024),
21//!     max_height: Some(1024),
22//!     format: Some(OutputFormat::Jpeg { quality: 85 }),
23//!     strip_exif: true,
24//! };
25//!
26//! let processed = process_image(input_bytes, options)?;
27//! # Ok(())
28//! # }
29//! ```
30
31use image::{DynamicImage, ImageFormat};
32use std::io::Cursor;
33
34/// Pi's default inline limits: 2000 pixels per side and 4.5 MiB of base64.
35pub const INLINE_IMAGE_MAX_SIDE: u32 = 2000;
36pub const INLINE_IMAGE_MAX_BASE64_BYTES: usize = 9 * 1024 * 1024 / 2;
37
38/// Decode, orient, resize and encode an image for model input. Small supported
39/// images retain their original bytes; oversized images fall back to JPEG.
40pub fn prepare_inline_image(bytes: &[u8]) -> Result<crate::tools::read::ProcessedImage, String> {
41    use base64::{engine::general_purpose::STANDARD, Engine};
42    use image::ImageDecoder;
43    let reader = image::ImageReader::new(Cursor::new(bytes))
44        .with_guessed_format()
45        .map_err(|e| e.to_string())?;
46    let mut decoder = reader.into_decoder().map_err(|e| e.to_string())?;
47    let orientation = decoder.orientation().map_err(|e| e.to_string())?;
48    let mut img = DynamicImage::from_decoder(decoder).map_err(|e| e.to_string())?;
49    img.apply_orientation(orientation);
50    let original = (img.width(), img.height());
51    let mime = crate::image::detect_supported_image_mime_type(bytes);
52    let base64_len = bytes.len().div_ceil(3) * 4;
53    if matches!(
54        mime,
55        Some("image/png" | "image/jpeg" | "image/gif" | "image/webp")
56    ) && original.0 <= INLINE_IMAGE_MAX_SIDE
57        && original.1 <= INLINE_IMAGE_MAX_SIDE
58        && base64_len < INLINE_IMAGE_MAX_BASE64_BYTES
59        && orientation == image::metadata::Orientation::NoTransforms
60    {
61        return Ok(crate::tools::read::ProcessedImage {
62            data: STANDARD.encode(bytes),
63            mime_type: mime.unwrap().into(),
64            hints: vec![],
65        });
66    }
67    img = resize_image(
68        img,
69        Some(INLINE_IMAGE_MAX_SIDE),
70        Some(INLINE_IMAGE_MAX_SIDE),
71    );
72    for _ in 0..8 {
73        let png = encode_image(&img, Some(OutputFormat::Png))?;
74        let mut candidate = (png, "image/png");
75        if candidate.0.len().div_ceil(3) * 4 >= INLINE_IMAGE_MAX_BASE64_BYTES {
76            // JPEG cannot encode alpha; convert explicitly before encoding.
77            let mut rgb = img.to_rgb8();
78            let rgba = img.to_rgba8();
79            for (target, source) in rgb.pixels_mut().zip(rgba.pixels()) {
80                let alpha = source[3] as u16;
81                for channel in 0..3 {
82                    target[channel] =
83                        ((source[channel] as u16 * alpha + 255 * (255 - alpha)) / 255) as u8;
84                }
85            }
86            let rgb = DynamicImage::ImageRgb8(rgb);
87            for quality in [85, 70, 50] {
88                let jpeg = encode_image(&rgb, Some(OutputFormat::Jpeg { quality }))?;
89                if jpeg.len() < candidate.0.len() {
90                    candidate = (jpeg, "image/jpeg");
91                }
92                if candidate.0.len().div_ceil(3) * 4 < INLINE_IMAGE_MAX_BASE64_BYTES {
93                    break;
94                }
95            }
96        }
97        if candidate.0.len().div_ceil(3) * 4 < INLINE_IMAGE_MAX_BASE64_BYTES {
98            let mut hints = vec![];
99            if let Some(mime) = mime {
100                if mime != candidate.1 {
101                    hints.push(format!("[Image converted from {mime} to {}.]", candidate.1));
102                }
103            }
104            if original != (img.width(), img.height()) {
105                hints.push(format!(
106                    "[Image resized from {}x{} to {}x{}.]",
107                    original.0,
108                    original.1,
109                    img.width(),
110                    img.height()
111                ));
112            }
113            return Ok(crate::tools::read::ProcessedImage {
114                data: STANDARD.encode(candidate.0),
115                mime_type: candidate.1.into(),
116                hints,
117            });
118        }
119        let bounds = ((img.width() * 3 / 4).max(1), (img.height() * 3 / 4).max(1));
120        img = resize_image(img, Some(bounds.0), Some(bounds.1));
121    }
122    Err("Image could not be resized below the inline image size limit".into())
123}
124
125/// Shared processor for the built-in read tool.
126pub struct InlineImageProcessor;
127
128/// Normalize image blocks while retaining order and surfacing omissions as
129/// text, rather than sending an invalid image to a provider.
130pub fn normalize_image_blocks(blocks: &mut Vec<rpi_ai::types::Content>) {
131    use rpi_ai::types::Content;
132    let mut output = Vec::with_capacity(blocks.len());
133    for block in std::mem::take(blocks) {
134        if let Content::Image(mut img) = block {
135            use base64::{engine::general_purpose::STANDARD, Engine};
136            match STANDARD
137                .decode(&img.data)
138                .map_err(|e| e.to_string())
139                .and_then(|bytes| prepare_inline_image(&bytes))
140            {
141                Ok(processed) => {
142                    img.data = processed.data;
143                    img.mime_type = processed.mime_type;
144                    output.push(Content::Image(img));
145                    output.extend(processed.hints.into_iter().map(Content::text));
146                }
147                Err(error) => output.push(Content::text(format!("[Image omitted: {error}.]"))),
148            }
149        } else {
150            output.push(block);
151        }
152    }
153    *blocks = output;
154}
155
156impl crate::tools::read::ReadImageProcessor for InlineImageProcessor {
157    fn process(
158        &self,
159        bytes: &[u8],
160        _mime_type: &str,
161        auto_resize: bool,
162    ) -> Result<crate::tools::read::ProcessedImage, String> {
163        if auto_resize {
164            return prepare_inline_image(bytes);
165        }
166        use base64::{engine::general_purpose::STANDARD, Engine};
167        let mime = crate::image::detect_supported_image_mime_type(bytes);
168        let (bytes, mime) = if matches!(
169            mime,
170            Some("image/png" | "image/jpeg" | "image/gif" | "image/webp")
171        ) {
172            (bytes.to_vec(), mime.unwrap())
173        } else {
174            (
175                process_image(bytes, ImageProcessingOptions::default())?,
176                "image/png",
177            )
178        };
179        Ok(crate::tools::read::ProcessedImage {
180            data: STANDARD.encode(bytes),
181            mime_type: mime.into(),
182            hints: vec![],
183        })
184    }
185}
186
187/// Output format for processed images.
188#[derive(Debug, Clone, Copy)]
189pub enum OutputFormat {
190    /// PNG format (lossless)
191    Png,
192    /// JPEG format with quality (1-100)
193    Jpeg { quality: u8 },
194    /// WebP format with quality (1-100)
195    WebP { quality: u8 },
196}
197
198/// Options for image processing.
199#[derive(Debug, Clone)]
200pub struct ImageProcessingOptions {
201    /// Maximum width in pixels. If set, image is scaled down to fit.
202    pub max_width: Option<u32>,
203    /// Maximum height in pixels. If set, image is scaled down to fit.
204    pub max_height: Option<u32>,
205    /// Output format. If `None`, preserves the input format.
206    pub format: Option<OutputFormat>,
207    /// Strip EXIF metadata (privacy). Default: `true`.
208    pub strip_exif: bool,
209}
210
211impl Default for ImageProcessingOptions {
212    fn default() -> Self {
213        Self {
214            max_width: None,
215            max_height: None,
216            format: None,
217            strip_exif: true,
218        }
219    }
220}
221
222/// Process an image according to the given options.
223///
224/// # Arguments
225///
226/// * `input` - Raw image bytes (PNG, JPEG, WebP, etc.)
227/// * `options` - Processing options (resize, format, EXIF stripping)
228///
229/// # Returns
230///
231/// * `Ok(Vec<u8>)` - Processed image bytes
232/// * `Err(String)` - Processing failed
233///
234/// # Example
235///
236/// ```rust,no_run
237/// use rpi_tools::image_processing::{process_image, ImageProcessingOptions, OutputFormat};
238///
239/// # fn example(raw_bytes: &[u8]) -> Result<(), String> {
240/// let options = ImageProcessingOptions {
241///     max_width: Some(800),
242///     max_height: Some(600),
243///     format: Some(OutputFormat::Jpeg { quality: 85 }),
244///     strip_exif: true,
245/// };
246///
247/// let processed = process_image(raw_bytes, options)?;
248/// # Ok(())
249/// # }
250/// ```
251pub fn process_image(input: &[u8], options: ImageProcessingOptions) -> Result<Vec<u8>, String> {
252    // Load the image
253    let mut img =
254        image::load_from_memory(input).map_err(|e| format!("Failed to load image: {}", e))?;
255
256    // Resize if needed
257    if options.max_width.is_some() || options.max_height.is_some() {
258        img = resize_image(img, options.max_width, options.max_height);
259    }
260
261    // Encode to the desired format
262    let output_bytes = encode_image(&img, options.format)?;
263
264    Ok(output_bytes)
265}
266
267/// Resize an image to fit within the given bounds while maintaining aspect ratio.
268fn resize_image(
269    img: DynamicImage,
270    max_width: Option<u32>,
271    max_height: Option<u32>,
272) -> DynamicImage {
273    let (orig_width, orig_height) = (img.width(), img.height());
274
275    // Calculate the scaling factor
276    let width_scale = max_width
277        .map(|mw| mw as f64 / orig_width as f64)
278        .unwrap_or(1.0);
279    let height_scale = max_height
280        .map(|mh| mh as f64 / orig_height as f64)
281        .unwrap_or(1.0);
282
283    let scale = width_scale.min(height_scale).min(1.0); // Never upscale
284
285    if scale >= 1.0 {
286        // No need to resize
287        return img;
288    }
289
290    let new_width = (orig_width as f64 * scale).round() as u32;
291    let new_height = (orig_height as f64 * scale).round() as u32;
292
293    img.resize(new_width, new_height, image::imageops::FilterType::Lanczos3)
294}
295
296/// Encode an image to the specified format.
297fn encode_image(img: &DynamicImage, format: Option<OutputFormat>) -> Result<Vec<u8>, String> {
298    let mut buffer = Vec::new();
299    let mut cursor = Cursor::new(&mut buffer);
300
301    match format {
302        Some(OutputFormat::Png) => {
303            img.write_to(&mut cursor, ImageFormat::Png)
304                .map_err(|e| format!("Failed to encode PNG: {}", e))?;
305        }
306        Some(OutputFormat::Jpeg { quality }) => {
307            // For JPEG with quality, we need to use the encoder directly
308            let encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut cursor, quality);
309            img.write_with_encoder(encoder)
310                .map_err(|e| format!("Failed to encode JPEG: {}", e))?;
311        }
312        Some(OutputFormat::WebP { quality }) => {
313            // WebP encoding requires the webp feature
314            #[cfg(feature = "webp")]
315            {
316                let encoder = image::codecs::webp::WebPEncoder::new_with_quality(
317                    &mut cursor,
318                    image::codecs::webp::WebPQuality::lossy(quality),
319                );
320                img.write_with_encoder(encoder)
321                    .map_err(|e| format!("Failed to encode WebP: {}", e))?;
322            }
323            #[cfg(not(feature = "webp"))]
324            {
325                let _ = quality;
326                return Err("WebP encoding requires the 'webp' feature".to_string());
327            }
328        }
329        None => {
330            // Default to PNG
331            img.write_to(&mut cursor, ImageFormat::Png)
332                .map_err(|e| format!("Failed to encode image: {}", e))?;
333        }
334    }
335
336    Ok(buffer)
337}
338
339/// Detect the MIME type of an image from its bytes.
340///
341/// # Arguments
342///
343/// * `data` - Raw image bytes
344///
345/// # Returns
346///
347/// * `Some(&str)` - MIME type (e.g., "image/png", "image/jpeg")
348/// * `None` - Unknown or unsupported format
349pub fn detect_image_mime(data: &[u8]) -> Option<&'static str> {
350    // Check magic bytes
351    if data.len() < 8 {
352        return None;
353    }
354
355    // PNG: 89 50 4E 47 0D 0A 1A 0A
356    if data.starts_with(&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]) {
357        return Some("image/png");
358    }
359
360    // JPEG: FF D8 FF
361    if data.starts_with(&[0xFF, 0xD8, 0xFF]) {
362        return Some("image/jpeg");
363    }
364
365    // GIF: 47 49 46 38
366    if data.starts_with(&[0x47, 0x49, 0x46, 0x38]) {
367        return Some("image/gif");
368    }
369
370    // WebP: 52 49 46 46 ... 57 45 42 50
371    if data.len() >= 12
372        && data.starts_with(&[0x52, 0x49, 0x46, 0x46])
373        && data[8..12].starts_with(&[0x57, 0x45, 0x42, 0x50])
374    {
375        return Some("image/webp");
376    }
377
378    // BMP: 42 4D
379    if data.starts_with(&[0x42, 0x4D]) {
380        return Some("image/bmp");
381    }
382
383    None
384}
385
386/// Get image dimensions without fully decoding the image.
387///
388/// # Arguments
389///
390/// * `data` - Raw image bytes
391///
392/// # Returns
393///
394/// * `Ok((width, height))` - Image dimensions in pixels
395/// * `Err(String)` - Failed to read dimensions
396pub fn get_image_dimensions(data: &[u8]) -> Result<(u32, u32), String> {
397    let img = image::load_from_memory(data).map_err(|e| format!("Failed to load image: {}", e))?;
398    Ok((img.width(), img.height()))
399}
400
401#[cfg(test)]
402mod tests {
403    use super::*;
404
405    #[test]
406    fn inline_images_keep_small_bytes_and_resize_large_images() {
407        use base64::{engine::general_purpose::STANDARD, Engine};
408        let small = encode_image(&DynamicImage::new_rgba8(8, 8), Some(OutputFormat::Png)).unwrap();
409        let processed = prepare_inline_image(&small).unwrap();
410        assert_eq!(STANDARD.decode(processed.data).unwrap(), small);
411        let large = encode_image(
412            &DynamicImage::new_rgba8(2400, 1200),
413            Some(OutputFormat::Png),
414        )
415        .unwrap();
416        let processed = prepare_inline_image(&large).unwrap();
417        assert!(processed.data.len() < INLINE_IMAGE_MAX_BASE64_BYTES);
418        let resized = image::load_from_memory(&STANDARD.decode(processed.data).unwrap()).unwrap();
419        assert_eq!((resized.width(), resized.height()), (2000, 1000));
420        assert!(processed.hints[0].contains("2400x1200 to 2000x1000"));
421    }
422
423    #[test]
424    fn inline_bmp_is_converted_and_invalid_images_are_reported() {
425        let mut bmp = Cursor::new(Vec::new());
426        DynamicImage::new_rgb8(4, 4)
427            .write_to(&mut bmp, ImageFormat::Bmp)
428            .unwrap();
429        let processed = prepare_inline_image(&bmp.into_inner()).unwrap();
430        assert_eq!(processed.mime_type, "image/png");
431        let mut blocks = vec![rpi_ai::types::Content::Image(rpi_ai::types::ImageContent {
432            kind: rpi_ai::types::ImageContentType,
433            data: "invalid".into(),
434            mime_type: "image/png".into(),
435        })];
436        normalize_image_blocks(&mut blocks);
437        assert!(
438            matches!(&blocks[0], rpi_ai::types::Content::Text(text) if text.text.contains("Image omitted"))
439        );
440    }
441
442    #[test]
443    fn inline_images_reduce_encoded_payload_even_below_dimension_limit() {
444        let mut rgba = image::RgbaImage::new(1200, 1200);
445        let mut seed = 1u32;
446        for pixel in rgba.pixels_mut() {
447            for channel in &mut pixel.0 {
448                seed ^= seed << 13;
449                seed ^= seed >> 17;
450                seed ^= seed << 5;
451                *channel = seed as u8;
452            }
453        }
454        let png = encode_image(&DynamicImage::ImageRgba8(rgba), Some(OutputFormat::Png)).unwrap();
455        assert!(png.len().div_ceil(3) * 4 > INLINE_IMAGE_MAX_BASE64_BYTES);
456        let processed = prepare_inline_image(&png).unwrap();
457        assert!(processed.data.len() < INLINE_IMAGE_MAX_BASE64_BYTES);
458        assert_eq!(processed.mime_type, "image/jpeg");
459    }
460
461    #[test]
462    fn test_detect_image_mime_png() {
463        let png_header = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
464        assert_eq!(detect_image_mime(&png_header), Some("image/png"));
465    }
466
467    #[test]
468    fn test_detect_image_mime_jpeg() {
469        let jpeg_header = vec![0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10, 0x4A, 0x46];
470        assert_eq!(detect_image_mime(&jpeg_header), Some("image/jpeg"));
471    }
472
473    #[test]
474    fn test_detect_image_mime_unknown() {
475        let unknown = vec![0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07];
476        assert_eq!(detect_image_mime(&unknown), None);
477    }
478
479    #[test]
480    fn test_detect_image_mime_too_short() {
481        let short = vec![0x89, 0x50];
482        assert_eq!(detect_image_mime(&short), None);
483    }
484}
rpi_tools/image_processing.rs

rpi_tools/
image_processing.rs
