ai_image/codecs/bmp/
encoder.rs

1use alloc::string::ToString;
2use byteorder_lite::{LittleEndian, WriteBytesExt};
3use no_std_io::io::{self, Write};
4
5use crate::error::{
6    EncodingError, ImageError, ImageFormatHint, ImageResult, ParameterError, ParameterErrorKind,
7    UnsupportedError, UnsupportedErrorKind,
8};
9use crate::{DynamicImage, ExtendedColorType, ImageEncoder, ImageFormat};
10
11const BITMAPFILEHEADER_SIZE: u32 = 14;
12const BITMAPINFOHEADER_SIZE: u32 = 40;
13const BITMAPV4HEADER_SIZE: u32 = 108;
14
15/// The representation of a BMP encoder.
16pub struct BmpEncoder<'a, W: 'a> {
17    writer: &'a mut W,
18}
19
20impl<'a, W: Write + 'a> BmpEncoder<'a, W> {
21    /// Create a new encoder that writes its output to ```w```.
22    pub fn new(w: &'a mut W) -> Self {
23        BmpEncoder { writer: w }
24    }
25
26    /// Encodes the image `image` that has dimensions `width` and `height` and `ExtendedColorType` `c`.
27    ///
28    /// # Panics
29    ///
30    /// Panics if `width * height * c.bytes_per_pixel() != image.len()`.
31    #[track_caller]
32    pub fn encode(
33        &mut self,
34        image: &[u8],
35        width: u32,
36        height: u32,
37        c: ExtendedColorType,
38    ) -> ImageResult<()> {
39        self.encode_with_palette(image, width, height, c, None)
40    }
41
42    /// Same as `encode`, but allow a palette to be passed in. The `palette` is ignored for color
43    /// types other than Luma/Luma-with-alpha.
44    ///
45    /// # Panics
46    ///
47    /// Panics if `width * height * c.bytes_per_pixel() != image.len()`.
48    #[track_caller]
49    pub fn encode_with_palette(
50        &mut self,
51        image: &[u8],
52        width: u32,
53        height: u32,
54        color_type: ExtendedColorType,
55        palette: Option<&[[u8; 3]]>,
56    ) -> ImageResult<()> {
57        if palette.is_some()
58            && color_type != ExtendedColorType::L8
59            && color_type != ExtendedColorType::La8
60        {
61            return Err(ImageError::Parameter(ParameterError::from_kind(
62                ParameterErrorKind::Generic(
63                    "Palette given which must only be used with L8 or La8 color types".to_string(),
64                ),
65            )));
66        }
67
68        let expected_buffer_len = color_type.buffer_size(width, height);
69        assert_eq!(
70            expected_buffer_len,
71            image.len() as u64,
72            "Invalid buffer length: expected {expected_buffer_len} got {} for {width}x{height} image",
73            image.len(),
74        );
75
76        let bmp_header_size = BITMAPFILEHEADER_SIZE;
77
78        let (dib_header_size, written_pixel_size, palette_color_count) =
79            written_pixel_info(color_type, palette)?;
80
81        let (padded_row, image_size) = width
82            .checked_mul(written_pixel_size)
83            // each row must be padded to a multiple of 4 bytes
84            .and_then(|v| v.checked_next_multiple_of(4))
85            .and_then(|v| {
86                let image_bytes = v.checked_mul(height)?;
87                Some((v, image_bytes))
88            })
89            .ok_or_else(|| {
90                ImageError::Parameter(ParameterError::from_kind(
91                    ParameterErrorKind::DimensionMismatch,
92                ))
93            })?;
94
95        let row_padding = padded_row - width * written_pixel_size;
96
97        // all palette colors are BGRA
98        let palette_size = palette_color_count.checked_mul(4).ok_or_else(|| {
99            ImageError::Encoding(EncodingError::new(
100                ImageFormatHint::Exact(ImageFormat::Bmp),
101                "calculated palette size larger than 2^32",
102            ))
103        })?;
104
105        let file_size = bmp_header_size
106            .checked_add(dib_header_size)
107            .and_then(|v| v.checked_add(palette_size))
108            .and_then(|v| v.checked_add(image_size))
109            .ok_or_else(|| {
110                ImageError::Encoding(EncodingError::new(
111                    ImageFormatHint::Exact(ImageFormat::Bmp),
112                    "calculated BMP header size larger than 2^32",
113                ))
114            })?;
115
116        let image_data_offset = bmp_header_size
117            .checked_add(dib_header_size)
118            .and_then(|v| v.checked_add(palette_size))
119            .ok_or_else(|| {
120                ImageError::Encoding(EncodingError::new(
121                    ImageFormatHint::Exact(ImageFormat::Bmp),
122                    "calculated BMP size larger than 2^32",
123                ))
124            })?;
125
126        // write BMP header
127        self.writer.write_u8(b'B')?;
128        self.writer.write_u8(b'M')?;
129        self.writer.write_u32::<LittleEndian>(file_size)?; // file size
130        self.writer.write_u16::<LittleEndian>(0)?; // reserved 1
131        self.writer.write_u16::<LittleEndian>(0)?; // reserved 2
132        self.writer.write_u32::<LittleEndian>(image_data_offset)?; // image data offset
133
134        // write DIB header
135        self.writer.write_u32::<LittleEndian>(dib_header_size)?;
136        self.writer.write_i32::<LittleEndian>(width as i32)?;
137        self.writer.write_i32::<LittleEndian>(height as i32)?;
138        self.writer.write_u16::<LittleEndian>(1)?; // color planes
139        self.writer
140            .write_u16::<LittleEndian>((written_pixel_size * 8) as u16)?; // bits per pixel
141        if dib_header_size >= BITMAPV4HEADER_SIZE {
142            // Assume BGRA32
143            self.writer.write_u32::<LittleEndian>(3)?; // compression method - bitfields
144        } else {
145            self.writer.write_u32::<LittleEndian>(0)?; // compression method - no compression
146        }
147        self.writer.write_u32::<LittleEndian>(image_size)?;
148        self.writer.write_i32::<LittleEndian>(0)?; // horizontal ppm
149        self.writer.write_i32::<LittleEndian>(0)?; // vertical ppm
150        self.writer.write_u32::<LittleEndian>(palette_color_count)?;
151        self.writer.write_u32::<LittleEndian>(0)?; // all colors are important
152        if dib_header_size >= BITMAPV4HEADER_SIZE {
153            // Assume BGRA32
154            self.writer.write_u32::<LittleEndian>(0xff << 16)?; // red mask
155            self.writer.write_u32::<LittleEndian>(0xff << 8)?; // green mask
156            self.writer.write_u32::<LittleEndian>(0xff)?; // blue mask
157            self.writer.write_u32::<LittleEndian>(0xff << 24)?; // alpha mask
158            self.writer.write_u32::<LittleEndian>(0x7352_4742)?; // colorspace - sRGB
159
160            // endpoints (3x3) and gamma (3)
161            for _ in 0..12 {
162                self.writer.write_u32::<LittleEndian>(0)?;
163            }
164        }
165
166        // write image data
167        match color_type {
168            ExtendedColorType::Rgb8 => self.encode_rgb(image, width, height, row_padding, 3)?,
169            ExtendedColorType::Rgba8 => self.encode_rgba(image, width, height, row_padding, 4)?,
170            ExtendedColorType::L8 => {
171                self.encode_gray(image, width, height, row_padding, 1, palette)?;
172            }
173            ExtendedColorType::La8 => {
174                self.encode_gray(image, width, height, row_padding, 2, palette)?;
175            }
176            _ => {
177                return Err(ImageError::Unsupported(
178                    UnsupportedError::from_format_and_kind(
179                        ImageFormat::Bmp.into(),
180                        UnsupportedErrorKind::Color(color_type),
181                    ),
182                ));
183            }
184        }
185
186        Ok(())
187    }
188
189    fn encode_rgb(
190        &mut self,
191        image: &[u8],
192        width: u32,
193        height: u32,
194        row_padding: u32,
195        bytes_per_pixel: u32,
196    ) -> io::Result<()> {
197        let width = width as usize;
198        let height = height as usize;
199        let x_stride = bytes_per_pixel as usize;
200        let y_stride = width * x_stride;
201        for row in (0..height).rev() {
202            // from the bottom up
203            let row_start = row * y_stride;
204            for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
205                let r = px[0];
206                let g = px[1];
207                let b = px[2];
208                // written as BGR
209                self.writer.write_all(&[b, g, r])?;
210            }
211            self.write_row_pad(row_padding)?;
212        }
213
214        Ok(())
215    }
216
217    fn encode_rgba(
218        &mut self,
219        image: &[u8],
220        width: u32,
221        height: u32,
222        row_padding: u32,
223        bytes_per_pixel: u32,
224    ) -> io::Result<()> {
225        let width = width as usize;
226        let height = height as usize;
227        let x_stride = bytes_per_pixel as usize;
228        let y_stride = width * x_stride;
229        for row in (0..height).rev() {
230            // from the bottom up
231            let row_start = row * y_stride;
232            for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
233                let r = px[0];
234                let g = px[1];
235                let b = px[2];
236                let a = px[3];
237                // written as BGRA
238                self.writer.write_all(&[b, g, r, a])?;
239            }
240            self.write_row_pad(row_padding)?;
241        }
242
243        Ok(())
244    }
245
246    fn encode_gray(
247        &mut self,
248        image: &[u8],
249        width: u32,
250        height: u32,
251        row_padding: u32,
252        bytes_per_pixel: u32,
253        palette: Option<&[[u8; 3]]>,
254    ) -> io::Result<()> {
255        // write grayscale palette
256        if let Some(palette) = palette {
257            for item in palette {
258                // each color is written as BGRA, where A is always 0
259                self.writer.write_all(&[item[2], item[1], item[0], 0])?;
260            }
261        } else {
262            for val in 0u8..=255 {
263                // each color is written as BGRA, where A is always 0 and since only grayscale is being written, B = G = R = index
264                self.writer.write_all(&[val, val, val, 0])?;
265            }
266        }
267
268        // write image data
269        let x_stride = bytes_per_pixel;
270        let y_stride = width * x_stride;
271        for row in (0..height).rev() {
272            // from the bottom up
273            let row_start = row * y_stride;
274
275            // color value is equal to the palette index
276            if x_stride == 1 {
277                // improve performance by writing the whole row at once
278                self.writer
279                    .write_all(&image[row_start as usize..][..y_stride as usize])?;
280            } else {
281                for col in 0..width {
282                    let pixel_start = (row_start + (col * x_stride)) as usize;
283                    self.writer.write_u8(image[pixel_start])?;
284                    // alpha is never written as it's not widely supported
285                }
286            }
287
288            self.write_row_pad(row_padding)?;
289        }
290
291        Ok(())
292    }
293
294    fn write_row_pad(&mut self, row_pad_size: u32) -> io::Result<()> {
295        for _ in 0..row_pad_size {
296            self.writer.write_u8(0)?;
297        }
298
299        Ok(())
300    }
301}
302
303impl<W: Write> ImageEncoder for BmpEncoder<'_, W> {
304    #[track_caller]
305    fn write_image(
306        mut self,
307        buf: &[u8],
308        width: u32,
309        height: u32,
310        color_type: ExtendedColorType,
311    ) -> ImageResult<()> {
312        self.encode(buf, width, height, color_type)
313    }
314
315    fn make_compatible_img(
316        &self,
317        _: crate::io::encoder::MethodSealedToImage,
318        img: &DynamicImage,
319    ) -> Option<DynamicImage> {
320        crate::io::encoder::dynimage_conversion_8bit(img)
321    }
322}
323
324/// Returns a tuple representing: (dib header size, written pixel size, palette color count).
325fn written_pixel_info(
326    c: ExtendedColorType,
327    palette: Option<&[[u8; 3]]>,
328) -> Result<(u32, u32, u32), ImageError> {
329    let (header, color_bytes, palette_count) = match c {
330        ExtendedColorType::Rgb8 => (BITMAPINFOHEADER_SIZE, 3, Some(0)),
331        ExtendedColorType::Rgba8 => (BITMAPV4HEADER_SIZE, 4, Some(0)),
332        ExtendedColorType::L8 => (
333            BITMAPINFOHEADER_SIZE,
334            1,
335            u32::try_from(palette.map(|p| p.len()).unwrap_or(256)).ok(),
336        ),
337        ExtendedColorType::La8 => (
338            BITMAPINFOHEADER_SIZE,
339            1,
340            u32::try_from(palette.map(|p| p.len()).unwrap_or(256)).ok(),
341        ),
342        _ => {
343            return Err(ImageError::Unsupported(
344                UnsupportedError::from_format_and_kind(
345                    ImageFormat::Bmp.into(),
346                    UnsupportedErrorKind::Color(c),
347                ),
348            ));
349        }
350    };
351
352    let palette_count = palette_count.ok_or_else(|| {
353        ImageError::Encoding(EncodingError::new(
354            ImageFormatHint::Exact(ImageFormat::Bmp),
355            "calculated palette size larger than 2^32",
356        ))
357    })?;
358
359    Ok((header, color_bytes, palette_count))
360}
361
362#[cfg(test)]
363mod tests {
364    use super::super::BmpDecoder;
365    use super::BmpEncoder;
366
367    use crate::ExtendedColorType;
368    use crate::ImageDecoder as _;
369    use no_std_io::io::Cursor;
370
371    fn round_trip_image(image: &[u8], width: u32, height: u32, c: ExtendedColorType) -> Vec<u8> {
372        let mut encoded_data = Vec::new();
373        {
374            let mut encoder = BmpEncoder::new(&mut encoded_data);
375            encoder
376                .encode(image, width, height, c)
377                .expect("could not encode image");
378        }
379
380        let decoder = BmpDecoder::new(Cursor::new(&encoded_data)).expect("failed to decode");
381
382        let mut buf = vec![0; decoder.total_bytes() as usize];
383        decoder.read_image(&mut buf).expect("failed to decode");
384        buf
385    }
386
387    #[test]
388    fn round_trip_single_pixel_rgb() {
389        let image = [255u8, 0, 0]; // single red pixel
390        let decoded = round_trip_image(&image, 1, 1, ExtendedColorType::Rgb8);
391        assert_eq!(3, decoded.len());
392        assert_eq!(255, decoded[0]);
393        assert_eq!(0, decoded[1]);
394        assert_eq!(0, decoded[2]);
395    }
396
397    #[test]
398    #[cfg(target_pointer_width = "64")]
399    fn huge_files_return_error() {
400        let mut encoded_data = Vec::new();
401        let image = vec![0u8; 3 * 40_000 * 40_000]; // 40_000x40_000 pixels, 3 bytes per pixel, allocated on the heap
402        let mut encoder = BmpEncoder::new(&mut encoded_data);
403        let result = encoder.encode(&image, 40_000, 40_000, ExtendedColorType::Rgb8);
404        assert!(result.is_err());
405    }
406
407    #[test]
408    fn round_trip_single_pixel_rgba() {
409        let image = [1, 2, 3, 4];
410        let decoded = round_trip_image(&image, 1, 1, ExtendedColorType::Rgba8);
411        assert_eq!(&decoded[..], &image[..]);
412    }
413
414    #[test]
415    fn round_trip_3px_rgb() {
416        let image = [0u8; 3 * 3 * 3]; // 3x3 pixels, 3 bytes per pixel
417        let _decoded = round_trip_image(&image, 3, 3, ExtendedColorType::Rgb8);
418    }
419
420    #[test]
421    fn round_trip_gray() {
422        let image = [0u8, 1, 2]; // 3 pixels
423        let decoded = round_trip_image(&image, 3, 1, ExtendedColorType::L8);
424        // should be read back as 3 RGB pixels
425        assert_eq!(9, decoded.len());
426        assert_eq!(0, decoded[0]);
427        assert_eq!(0, decoded[1]);
428        assert_eq!(0, decoded[2]);
429        assert_eq!(1, decoded[3]);
430        assert_eq!(1, decoded[4]);
431        assert_eq!(1, decoded[5]);
432        assert_eq!(2, decoded[6]);
433        assert_eq!(2, decoded[7]);
434        assert_eq!(2, decoded[8]);
435    }
436
437    #[test]
438    fn round_trip_graya() {
439        let image = [0u8, 0, 1, 0, 2, 0]; // 3 pixels, each with an alpha channel
440        let decoded = round_trip_image(&image, 1, 3, ExtendedColorType::La8);
441        // should be read back as 3 RGB pixels
442        assert_eq!(9, decoded.len());
443        assert_eq!(0, decoded[0]);
444        assert_eq!(0, decoded[1]);
445        assert_eq!(0, decoded[2]);
446        assert_eq!(1, decoded[3]);
447        assert_eq!(1, decoded[4]);
448        assert_eq!(1, decoded[5]);
449        assert_eq!(2, decoded[6]);
450        assert_eq!(2, decoded[7]);
451        assert_eq!(2, decoded[8]);
452    }
453
454    #[test]
455    fn regression_issue_2604() {
456        let mut image = vec![];
457        let mut encoder = BmpEncoder::new(&mut image);
458        encoder
459            .encode(&[], 1 << 31, 0, ExtendedColorType::Rgb8)
460            .unwrap_err();
461    }
462}
ai_image/codecs/bmp/encoder.rs

ai_image/codecs/bmp/
encoder.rs
