ai_image/io/
free_functions.rs

1use alloc::boxed::Box;
2use alloc::vec::Vec;
3use core::iter;
4use core::mem::size_of;
5#[cfg(feature = "std")]
6use no_std_io::io::BufRead;
7use no_std_io::io::{self, Seek, Write};
8#[cfg(feature = "std")]
9use std::fs::File;
10#[cfg(feature = "std")]
11use std::path::Path;
12
13#[cfg(all(
14    not(feature = "std"),
15    any(
16        feature = "png",
17        feature = "jpeg",
18        feature = "gif",
19        feature = "bmp",
20        feature = "ico",
21        feature = "pnm",
22        feature = "tga",
23        feature = "tiff",
24        feature = "exr",
25        feature = "avif",
26        feature = "qoi",
27        feature = "webp",
28        feature = "hdr",
29        feature = "ff"
30    )
31))]
32use crate::codecs::*;
33use crate::io::encoder::ImageEncoderBoxed;
34#[cfg(feature = "std")]
35use crate::{codecs::*, ExtendedColorType, ImageReader};
36
37use crate::error::{
38    ImageError, ImageFormatHint, ImageResult, LimitError, LimitErrorKind, ParameterError,
39    ParameterErrorKind, UnsupportedError, UnsupportedErrorKind,
40};
41#[cfg(feature = "std")]
42use crate::DynamicImage;
43use crate::{ImageDecoder, ImageFormat};
44
45/// Create a new image from a Reader.
46///
47/// Assumes the reader is already buffered. For optimal performance,
48/// consider wrapping the reader with a `BufReader::new()`.
49///
50/// Try [`ImageReader`] for more advanced uses.
51#[cfg(feature = "std")]
52pub fn load<R: BufRead + Seek>(r: R, format: ImageFormat) -> ImageResult<DynamicImage> {
53    let mut reader = ImageReader::new(r);
54    reader.set_format(format);
55    reader.decode()
56}
57
58/// Saves the supplied buffer to a file at the path specified.
59///
60/// The image format is derived from the file extension. The buffer is assumed to have the correct
61/// format according to the specified color type. This will lead to corrupted files if the buffer
62/// contains malformed data.
63#[cfg(feature = "std")]
64pub fn save_buffer(
65    path: impl AsRef<Path>,
66    buf: &[u8],
67    width: u32,
68    height: u32,
69    color: impl Into<ExtendedColorType>,
70) -> ImageResult<()> {
71    let format = ImageFormat::from_path(path.as_ref())?;
72    save_buffer_with_format(path, buf, width, height, color, format)
73}
74
75/// Saves the supplied buffer to a file given the path and desired format.
76///
77/// The buffer is assumed to have the correct format according to the specified color type. This
78/// will lead to corrupted files if the buffer contains malformed data.
79#[cfg(feature = "std")]
80pub fn save_buffer_with_format(
81    path: impl AsRef<Path>,
82    buf: &[u8],
83    width: u32,
84    height: u32,
85    color: impl Into<ExtendedColorType>,
86    format: ImageFormat,
87) -> ImageResult<()> {
88    let buffered_file_write = &mut io::BufWriter::new(File::create(path)?); // always seekable
89    let encoder = encoder_for_format(format, buffered_file_write)?;
90    encoder.write_image(buf, width, height, color.into())
91}
92
93pub(crate) fn encoder_for_format<'a, W: Write + Seek>(
94    format: ImageFormat,
95    buffered_write: &'a mut W,
96) -> ImageResult<Box<dyn ImageEncoderBoxed + 'a>> {
97    let _ = &buffered_write;
98    match format {
99        #[cfg(feature = "png")]
100        ImageFormat::Png => Ok(Box::new(png::PngEncoder::new(buffered_write))),
101        #[cfg(feature = "jpeg")]
102        ImageFormat::Jpeg => Ok(Box::new(jpeg::JpegEncoder::new(buffered_write))),
103        #[cfg(feature = "pnm")]
104        ImageFormat::Pnm => Ok(Box::new(pnm::PnmEncoder::new(buffered_write))),
105        #[cfg(feature = "gif")]
106        ImageFormat::Gif => Ok(Box::new(gif::GifEncoder::new(buffered_write))),
107        #[cfg(feature = "ico")]
108        ImageFormat::Ico => Ok(Box::new(ico::IcoEncoder::new(buffered_write))),
109        #[cfg(feature = "bmp")]
110        ImageFormat::Bmp => Ok(Box::new(bmp::BmpEncoder::new(buffered_write))),
111        #[cfg(feature = "ff")]
112        ImageFormat::Farbfeld => Ok(Box::new(farbfeld::FarbfeldEncoder::new(buffered_write))),
113        #[cfg(feature = "tga")]
114        ImageFormat::Tga => Ok(Box::new(tga::TgaEncoder::new(buffered_write))),
115        #[cfg(feature = "exr")]
116        ImageFormat::OpenExr => Ok(Box::new(openexr::OpenExrEncoder::new(buffered_write))),
117        #[cfg(feature = "tiff")]
118        ImageFormat::Tiff => Ok(Box::new(tiff::TiffEncoder::new(buffered_write))),
119        #[cfg(feature = "avif")]
120        ImageFormat::Avif => Ok(Box::new(avif::AvifEncoder::new(buffered_write))),
121        #[cfg(feature = "qoi")]
122        ImageFormat::Qoi => Ok(Box::new(qoi::QoiEncoder::new(buffered_write))),
123        #[cfg(feature = "webp")]
124        ImageFormat::WebP => Ok(Box::new(webp::WebPEncoder::new_lossless(buffered_write))),
125        #[cfg(feature = "hdr")]
126        ImageFormat::Hdr => Ok(Box::new(hdr::HdrEncoder::new(buffered_write))),
127        _ => Err(ImageError::Unsupported(
128            UnsupportedError::from_format_and_kind(
129                ImageFormatHint::Unknown,
130                UnsupportedErrorKind::Format(ImageFormatHint::Name(alloc::format!("{format:?}"))),
131            ),
132        )),
133    }
134}
135
136static MAGIC_BYTES: [(&[u8], &[u8], ImageFormat); 22] = [
137    (b"\x89PNG\r\n\x1a\n", b"", ImageFormat::Png),
138    (&[0xff, 0xd8, 0xff], b"", ImageFormat::Jpeg),
139    (b"GIF89a", b"", ImageFormat::Gif),
140    (b"GIF87a", b"", ImageFormat::Gif),
141    (
142        b"RIFF\0\0\0\0WEBP",
143        b"\xFF\xFF\xFF\xFF\0\0\0\0",
144        ImageFormat::WebP,
145    ),
146    (b"MM\x00*", b"", ImageFormat::Tiff),
147    (b"II*\x00", b"", ImageFormat::Tiff),
148    (b"DDS ", b"", ImageFormat::Dds),
149    (b"BM", b"", ImageFormat::Bmp),
150    (&[0, 0, 1, 0], b"", ImageFormat::Ico),
151    (b"#?RADIANCE", b"", ImageFormat::Hdr),
152    (b"\0\0\0\0ftypavif", b"\xFF\xFF\0\0", ImageFormat::Avif),
153    (&[0x76, 0x2f, 0x31, 0x01], b"", ImageFormat::OpenExr), // = &exr::meta::magic_number::BYTES
154    (b"qoif", b"", ImageFormat::Qoi),
155    (b"P1", b"", ImageFormat::Pnm),
156    (b"P2", b"", ImageFormat::Pnm),
157    (b"P3", b"", ImageFormat::Pnm),
158    (b"P4", b"", ImageFormat::Pnm),
159    (b"P5", b"", ImageFormat::Pnm),
160    (b"P6", b"", ImageFormat::Pnm),
161    (b"P7", b"", ImageFormat::Pnm),
162    (b"farbfeld", b"", ImageFormat::Farbfeld),
163];
164
165/// Guess image format from memory block
166///
167/// Makes an educated guess about the image format based on the Magic Bytes at the beginning.
168/// TGA is not supported by this function.
169/// This is not to be trusted on the validity of the whole memory block
170pub fn guess_format(buffer: &[u8]) -> ImageResult<ImageFormat> {
171    match guess_format_impl(buffer) {
172        Some(format) => Ok(format),
173        None => Err(ImageError::Unsupported(ImageFormatHint::Unknown.into())),
174    }
175}
176
177pub(crate) fn guess_format_impl(buffer: &[u8]) -> Option<ImageFormat> {
178    for &(signature, mask, format) in &MAGIC_BYTES {
179        if mask.is_empty() {
180            if buffer.starts_with(signature) {
181                return Some(format);
182            }
183        } else if buffer.len() >= signature.len()
184            && buffer
185                .iter()
186                .zip(signature.iter())
187                .zip(mask.iter().chain(iter::repeat(&0xFF)))
188                .all(|((&byte, &sig), &mask)| byte & mask == sig)
189        {
190            return Some(format);
191        }
192    }
193
194    None
195}
196
197/// Decodes a specific region of the image, represented by the rectangle
198/// starting from ```x``` and ```y``` and having ```length``` and ```width```
199#[allow(dead_code)]
200#[allow(clippy::too_many_arguments)]
201pub(crate) fn load_rect<D, F1, F2, E>(
202    x: u32,
203    y: u32,
204    width: u32,
205    height: u32,
206    buf: &mut [u8],
207    row_pitch: usize,
208    decoder: &mut D,
209    scanline_bytes: usize,
210    mut seek_scanline: F1,
211    mut read_scanline: F2,
212) -> ImageResult<()>
213where
214    D: ImageDecoder,
215    F1: FnMut(&mut D, u64) -> io::Result<()>,
216    F2: FnMut(&mut D, &mut [u8]) -> Result<(), E>,
217    ImageError: From<E>,
218{
219    let scanline_bytes = u64::try_from(scanline_bytes).unwrap();
220    let row_pitch = u64::try_from(row_pitch).unwrap();
221
222    let (x, y, width, height) = (
223        u64::from(x),
224        u64::from(y),
225        u64::from(width),
226        u64::from(height),
227    );
228    let dimensions = decoder.dimensions();
229    let bytes_per_pixel = u64::from(decoder.color_type().bytes_per_pixel());
230    let row_bytes = bytes_per_pixel * u64::from(dimensions.0);
231    let total_bytes = width * height * bytes_per_pixel;
232
233    assert!(
234        buf.len() >= usize::try_from(total_bytes).unwrap_or(usize::MAX),
235        "output buffer too short\n expected `{}`, provided `{}`",
236        total_bytes,
237        buf.len()
238    );
239
240    let mut current_scanline = 0;
241    let mut tmp = Vec::new();
242    let mut tmp_scanline = None;
243
244    {
245        // Read a range of the image starting from byte number `start` and continuing until byte
246        // number `end`. Updates `current_scanline` and `bytes_read` appropriately.
247        let mut read_image_range =
248            |mut start: u64, end: u64, mut output: &mut [u8]| -> ImageResult<()> {
249                // If the first scanline we need is already stored in the temporary buffer, then handle
250                // it first.
251                let target_scanline = start / scanline_bytes;
252                if tmp_scanline == Some(target_scanline) {
253                    let position = target_scanline * scanline_bytes;
254                    let offset = start.saturating_sub(position);
255                    let len = (end - start)
256                        .min(scanline_bytes - offset)
257                        .min(end - position);
258
259                    output
260                        .write_all(&tmp[offset as usize..][..len as usize])
261                        .unwrap();
262                    start += len;
263
264                    if start == end {
265                        return Ok(());
266                    }
267                }
268
269                let target_scanline = start / scanline_bytes;
270                if target_scanline != current_scanline {
271                    seek_scanline(decoder, target_scanline)?;
272                    current_scanline = target_scanline;
273                }
274
275                let mut position = current_scanline * scanline_bytes;
276                while position < end {
277                    if position >= start && end - position >= scanline_bytes {
278                        read_scanline(decoder, &mut output[..(scanline_bytes as usize)])?;
279                        output = &mut output[scanline_bytes as usize..];
280                    } else {
281                        tmp.resize(scanline_bytes as usize, 0u8);
282                        read_scanline(decoder, &mut tmp)?;
283                        tmp_scanline = Some(current_scanline);
284
285                        let offset = start.saturating_sub(position);
286                        let len = (end - start)
287                            .min(scanline_bytes - offset)
288                            .min(end - position);
289
290                        output
291                            .write_all(&tmp[offset as usize..][..len as usize])
292                            .unwrap();
293                    }
294
295                    current_scanline += 1;
296                    position += scanline_bytes;
297                }
298                Ok(())
299            };
300
301        if x + width > u64::from(dimensions.0)
302            || y + height > u64::from(dimensions.1)
303            || width == 0
304            || height == 0
305        {
306            return Err(ImageError::Parameter(ParameterError::from_kind(
307                ParameterErrorKind::DimensionMismatch,
308            )));
309        }
310        if scanline_bytes > usize::MAX as u64 {
311            return Err(ImageError::Limits(LimitError::from_kind(
312                LimitErrorKind::InsufficientMemory,
313            )));
314        }
315
316        if x == 0 && width == u64::from(dimensions.0) && row_pitch == row_bytes {
317            let start = x * bytes_per_pixel + y * row_bytes;
318            let end = (x + width) * bytes_per_pixel + (y + height - 1) * row_bytes;
319            read_image_range(start, end, buf)?;
320        } else {
321            for (output_slice, row) in buf.chunks_mut(row_pitch as usize).zip(y..(y + height)) {
322                let start = x * bytes_per_pixel + row * row_bytes;
323                let end = (x + width) * bytes_per_pixel + row * row_bytes;
324                read_image_range(start, end, output_slice)?;
325            }
326        }
327    }
328
329    // Seek back to the start
330    Ok(seek_scanline(decoder, 0)?)
331}
332
333/// Reads all of the bytes of a decoder into a Vec<T>. No particular alignment
334/// of the output buffer is guaranteed.
335///
336/// Panics if there isn't enough memory to decode the image.
337pub(crate) fn decoder_to_vec<T>(decoder: impl ImageDecoder) -> ImageResult<Vec<T>>
338where
339    T: crate::traits::Primitive + bytemuck::Pod,
340{
341    let total_bytes = usize::try_from(decoder.total_bytes());
342    if total_bytes.is_err() || total_bytes.unwrap() > isize::MAX as usize {
343        return Err(ImageError::Limits(LimitError::from_kind(
344            LimitErrorKind::InsufficientMemory,
345        )));
346    }
347
348    let mut buf = alloc::vec![num_traits::Zero::zero(); total_bytes.unwrap() / size_of::<T>()];
349    decoder.read_image(bytemuck::cast_slice_mut(buf.as_mut_slice()))?;
350    Ok(buf)
351}
352
353#[cfg(test)]
354mod tests {
355    use crate::ColorType;
356    use std::io;
357
358    use super::{load_rect, ImageDecoder, ImageResult};
359
360    #[test]
361    fn test_load_rect() {
362        struct MockDecoder {
363            scanline_number: u64,
364            scanline_bytes: u64,
365        }
366        impl ImageDecoder for MockDecoder {
367            fn dimensions(&self) -> (u32, u32) {
368                (5, 5)
369            }
370            fn color_type(&self) -> ColorType {
371                ColorType::L8
372            }
373            fn read_image(self, _buf: &mut [u8]) -> ImageResult<()> {
374                unimplemented!()
375            }
376            fn read_image_boxed(self: Box<Self>, buf: &mut [u8]) -> ImageResult<()> {
377                (*self).read_image(buf)
378            }
379        }
380
381        const DATA: [u8; 25] = [
382            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
383            24,
384        ];
385
386        fn seek_scanline(m: &mut MockDecoder, n: u64) -> io::Result<()> {
387            m.scanline_number = n;
388            Ok(())
389        }
390        fn read_scanline(m: &mut MockDecoder, buf: &mut [u8]) -> io::Result<()> {
391            let bytes_read = m.scanline_number * m.scanline_bytes;
392            if bytes_read >= 25 {
393                return Ok(());
394            }
395
396            let len = m.scanline_bytes.min(25 - bytes_read);
397            buf[..(len as usize)].copy_from_slice(&DATA[(bytes_read as usize)..][..(len as usize)]);
398            m.scanline_number += 1;
399            Ok(())
400        }
401
402        for scanline_bytes in 1..30 {
403            let mut output = [0u8; 26];
404
405            load_rect(
406                0,
407                0,
408                5,
409                5,
410                &mut output,
411                5,
412                &mut MockDecoder {
413                    scanline_number: 0,
414                    scanline_bytes,
415                },
416                scanline_bytes as usize,
417                seek_scanline,
418                read_scanline,
419            )
420            .unwrap();
421            assert_eq!(output[0..25], DATA);
422            assert_eq!(output[25], 0);
423
424            output = [0u8; 26];
425            load_rect(
426                3,
427                2,
428                1,
429                1,
430                &mut output,
431                1,
432                &mut MockDecoder {
433                    scanline_number: 0,
434                    scanline_bytes,
435                },
436                scanline_bytes as usize,
437                seek_scanline,
438                read_scanline,
439            )
440            .unwrap();
441            assert_eq!(output[0..2], [13, 0]);
442
443            output = [0u8; 26];
444            load_rect(
445                3,
446                2,
447                2,
448                2,
449                &mut output,
450                2,
451                &mut MockDecoder {
452                    scanline_number: 0,
453                    scanline_bytes,
454                },
455                scanline_bytes as usize,
456                seek_scanline,
457                read_scanline,
458            )
459            .unwrap();
460            assert_eq!(output[0..5], [13, 14, 18, 19, 0]);
461
462            output = [0u8; 26];
463            load_rect(
464                1,
465                1,
466                2,
467                4,
468                &mut output,
469                2,
470                &mut MockDecoder {
471                    scanline_number: 0,
472                    scanline_bytes,
473                },
474                scanline_bytes as usize,
475                seek_scanline,
476                read_scanline,
477            )
478            .unwrap();
479            assert_eq!(output[0..9], [6, 7, 11, 12, 16, 17, 21, 22, 0]);
480        }
481    }
482
483    #[test]
484    fn test_load_rect_single_scanline() {
485        const DATA: [u8; 25] = [
486            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
487            24,
488        ];
489
490        struct MockDecoder;
491        impl ImageDecoder for MockDecoder {
492            fn dimensions(&self) -> (u32, u32) {
493                (5, 5)
494            }
495            fn color_type(&self) -> ColorType {
496                ColorType::L8
497            }
498            fn read_image(self, _buf: &mut [u8]) -> ImageResult<()> {
499                unimplemented!()
500            }
501            fn read_image_boxed(self: Box<Self>, buf: &mut [u8]) -> ImageResult<()> {
502                (*self).read_image(buf)
503            }
504        }
505
506        // Ensure that seek scanline is called only once.
507        let mut seeks = 0;
508        let seek_scanline = |_d: &mut MockDecoder, n: u64| -> io::Result<()> {
509            seeks += 1;
510            assert_eq!(n, 0);
511            assert_eq!(seeks, 1);
512            Ok(())
513        };
514
515        fn read_scanline(_m: &mut MockDecoder, buf: &mut [u8]) -> io::Result<()> {
516            buf.copy_from_slice(&DATA);
517            Ok(())
518        }
519
520        let mut output = [0; 26];
521        load_rect(
522            1,
523            1,
524            2,
525            4,
526            &mut output,
527            2,
528            &mut MockDecoder,
529            DATA.len(),
530            seek_scanline,
531            read_scanline,
532        )
533        .unwrap();
534        assert_eq!(output[0..9], [6, 7, 11, 12, 16, 17, 21, 22, 0]);
535    }
536}
ai_image/io/free_functions.rs

ai_image/io/
free_functions.rs
