image/io/
free_functions.rs

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