ai_image/codecs/hdr/
encoder.rs

1use alloc::{format, vec, vec::Vec};
2use core::cmp::Ordering;
3use no_std_io::io::{Result, Write};
4#[cfg(not(feature = "std"))]
5use num_traits::Float as _;
6
7use crate::codecs::hdr::{rgbe8, Rgbe8Pixel, SIGNATURE};
8use crate::color::Rgb;
9use crate::error::{ImageResult, UnsupportedError, UnsupportedErrorKind};
10use crate::{ExtendedColorType, ImageEncoder, ImageError, ImageFormat};
11
12/// Radiance HDR encoder
13pub struct HdrEncoder<W: Write> {
14    w: W,
15}
16
17impl<W: Write> ImageEncoder for HdrEncoder<W> {
18    fn write_image(
19        self,
20        unaligned_bytes: &[u8],
21        width: u32,
22        height: u32,
23        color_type: ExtendedColorType,
24    ) -> ImageResult<()> {
25        match color_type {
26            ExtendedColorType::Rgb32F => {
27                let bytes_per_pixel = color_type.bits_per_pixel() as usize / 8;
28                let rgbe_pixels = unaligned_bytes
29                    .chunks_exact(bytes_per_pixel)
30                    .map(|bytes| to_rgbe8(Rgb::<f32>(bytemuck::pod_read_unaligned(bytes))));
31
32                // the length will be checked inside encode_pixels
33                self.encode_pixels(rgbe_pixels, width as usize, height as usize)
34            }
35            _ => Err(ImageError::Unsupported(
36                UnsupportedError::from_format_and_kind(
37                    ImageFormat::Hdr.into(),
38                    UnsupportedErrorKind::Color(color_type),
39                ),
40            )),
41        }
42    }
43}
44
45impl<W: Write> HdrEncoder<W> {
46    /// Creates encoder
47    pub fn new(w: W) -> HdrEncoder<W> {
48        HdrEncoder { w }
49    }
50
51    /// Encodes the image ```rgb```
52    /// that has dimensions ```width``` and ```height```
53    pub fn encode(self, rgb: &[Rgb<f32>], width: usize, height: usize) -> ImageResult<()> {
54        self.encode_pixels(rgb.iter().map(|&rgb| to_rgbe8(rgb)), width, height)
55    }
56
57    /// Encodes the image ```flattened_rgbe_pixels```
58    /// that has dimensions ```width``` and ```height```.
59    /// The callback must return the color for the given flattened index of the pixel (row major).
60    fn encode_pixels(
61        mut self,
62        mut flattened_rgbe_pixels: impl ExactSizeIterator<Item = Rgbe8Pixel>,
63        width: usize,
64        height: usize,
65    ) -> ImageResult<()> {
66        assert!(
67            flattened_rgbe_pixels.len() >= width * height,
68            "not enough pixels provided"
69        ); // bonus: this might elide some bounds checks
70
71        let w = &mut self.w;
72        w.write_all(SIGNATURE)?;
73        w.write_all(b"\n")?;
74        w.write_all(b"# Rust HDR encoder\n")?;
75        w.write_all(b"FORMAT=32-bit_rle_rgbe\n\n")?;
76        w.write_all(format!("-Y {height} +X {width}\n").as_bytes())?;
77
78        if !(8..=32_768).contains(&width) {
79            for pixel in flattened_rgbe_pixels {
80                write_rgbe8(w, pixel)?;
81            }
82        } else {
83            // new RLE marker contains scanline width
84            let marker = rgbe8(2, 2, (width / 256) as u8, (width % 256) as u8);
85            // buffers for encoded pixels
86            let mut bufr = vec![0; width];
87            let mut bufg = vec![0; width];
88            let mut bufb = vec![0; width];
89            let mut bufe = vec![0; width];
90            let mut rle_buf = vec![0; width];
91            for _scanline_index in 0..height {
92                assert!(flattened_rgbe_pixels.len() >= width); // may reduce the bound checks
93
94                for ((((r, g), b), e), pixel) in bufr
95                    .iter_mut()
96                    .zip(bufg.iter_mut())
97                    .zip(bufb.iter_mut())
98                    .zip(bufe.iter_mut())
99                    .zip(&mut flattened_rgbe_pixels)
100                {
101                    *r = pixel.c[0];
102                    *g = pixel.c[1];
103                    *b = pixel.c[2];
104                    *e = pixel.e;
105                }
106
107                write_rgbe8(w, marker)?; // New RLE encoding marker
108                rle_buf.clear();
109                rle_compress(&bufr[..], &mut rle_buf);
110                w.write_all(&rle_buf[..])?;
111                rle_buf.clear();
112                rle_compress(&bufg[..], &mut rle_buf);
113                w.write_all(&rle_buf[..])?;
114                rle_buf.clear();
115                rle_compress(&bufb[..], &mut rle_buf);
116                w.write_all(&rle_buf[..])?;
117                rle_buf.clear();
118                rle_compress(&bufe[..], &mut rle_buf);
119                w.write_all(&rle_buf[..])?;
120            }
121        }
122        Ok(())
123    }
124}
125
126#[derive(Debug, PartialEq, Eq)]
127enum RunOrNot {
128    Run(u8, usize),
129    Norun(usize, usize),
130}
131
132use self::RunOrNot::{Norun, Run};
133
134const RUN_MAX_LEN: usize = 127;
135const NORUN_MAX_LEN: usize = 128;
136
137struct RunIterator<'a> {
138    data: &'a [u8],
139    curidx: usize,
140}
141
142impl<'a> RunIterator<'a> {
143    fn new(data: &'a [u8]) -> RunIterator<'a> {
144        RunIterator { data, curidx: 0 }
145    }
146}
147
148impl Iterator for RunIterator<'_> {
149    type Item = RunOrNot;
150
151    fn next(&mut self) -> Option<Self::Item> {
152        if self.curidx == self.data.len() {
153            None
154        } else {
155            let cv = self.data[self.curidx];
156            let crun = self.data[self.curidx..]
157                .iter()
158                .take_while(|&&v| v == cv)
159                .take(RUN_MAX_LEN)
160                .count();
161            let ret = if crun > 2 {
162                Run(cv, crun)
163            } else {
164                Norun(self.curidx, crun)
165            };
166            self.curidx += crun;
167            Some(ret)
168        }
169    }
170}
171
172struct NorunCombineIterator<'a> {
173    runiter: RunIterator<'a>,
174    prev: Option<RunOrNot>,
175}
176
177impl<'a> NorunCombineIterator<'a> {
178    fn new(data: &'a [u8]) -> NorunCombineIterator<'a> {
179        NorunCombineIterator {
180            runiter: RunIterator::new(data),
181            prev: None,
182        }
183    }
184}
185
186// Combines sequential noruns produced by RunIterator
187impl Iterator for NorunCombineIterator<'_> {
188    type Item = RunOrNot;
189
190    fn next(&mut self) -> Option<Self::Item> {
191        loop {
192            match self.prev.take() {
193                Some(Run(c, len)) => {
194                    // Just return stored run
195                    return Some(Run(c, len));
196                }
197                Some(Norun(idx, len)) => {
198                    // Let's see if we need to continue norun
199                    match self.runiter.next() {
200                        Some(Norun(_, len1)) => {
201                            // norun continues
202                            let clen = len + len1; // combined length
203                            match clen.cmp(&NORUN_MAX_LEN) {
204                                Ordering::Equal => return Some(Norun(idx, clen)),
205                                Ordering::Greater => {
206                                    // combined norun exceeds maximum length. store extra part of norun
207                                    self.prev =
208                                        Some(Norun(idx + NORUN_MAX_LEN, clen - NORUN_MAX_LEN));
209                                    // then return maximal norun
210                                    return Some(Norun(idx, NORUN_MAX_LEN));
211                                }
212                                Ordering::Less => {
213                                    // len + len1 < NORUN_MAX_LEN
214                                    self.prev = Some(Norun(idx, len + len1));
215                                    // combine and continue loop
216                                }
217                            }
218                        }
219                        Some(Run(c, len1)) => {
220                            // Run encountered. Store it
221                            self.prev = Some(Run(c, len1));
222                            return Some(Norun(idx, len)); // and return combined norun
223                        }
224                        None => {
225                            // End of sequence
226                            return Some(Norun(idx, len)); // return combined norun
227                        }
228                    }
229                } // End match self.prev.take() == Some(NoRun())
230                None => {
231                    // No norun to combine
232                    match self.runiter.next() {
233                        Some(Norun(idx, len)) => {
234                            self.prev = Some(Norun(idx, len));
235                            // store for combine and continue the loop
236                        }
237                        Some(Run(c, len)) => {
238                            // Some run. Just return it
239                            return Some(Run(c, len));
240                        }
241                        None => {
242                            // That's all, folks
243                            return None;
244                        }
245                    }
246                } // End match self.prev.take() == None
247            } // End match
248        } // End loop
249    }
250}
251
252// Appends RLE compressed ```data``` to ```rle```
253fn rle_compress(data: &[u8], rle: &mut Vec<u8>) {
254    rle.clear();
255    if data.is_empty() {
256        rle.push(0); // Technically correct. It means read next 0 bytes.
257        return;
258    }
259    // Task: split data into chunks of repeating (max 127) and non-repeating bytes (max 128)
260    // Prepend non-repeating chunk with its length
261    // Replace repeating byte with (run length + 128) and the byte
262    for rnr in NorunCombineIterator::new(data) {
263        match rnr {
264            Run(c, len) => {
265                assert!(len <= 127);
266                rle.push(128u8 + len as u8);
267                rle.push(c);
268            }
269            Norun(idx, len) => {
270                assert!(len <= 128);
271                rle.push(len as u8);
272                rle.extend_from_slice(&data[idx..idx + len]);
273            }
274        }
275    }
276}
277
278fn write_rgbe8<W: Write>(w: &mut W, v: Rgbe8Pixel) -> Result<()> {
279    w.write_all(&[v.c[0], v.c[1], v.c[2], v.e])
280}
281
282/// Converts ```Rgb<f32>``` into ```Rgbe8Pixel```
283pub(crate) fn to_rgbe8(pix: Rgb<f32>) -> Rgbe8Pixel {
284    let pix = pix.0;
285    let mx = f32::max(pix[0], f32::max(pix[1], pix[2]));
286    if mx <= 0.0 {
287        Rgbe8Pixel { c: [0, 0, 0], e: 0 }
288    } else {
289        // let (frac, exp) = mx.frexp(); // unstable yet
290        let exp = mx.log2().floor() as i32 + 1;
291        let mul = 2.0f32.powi(exp);
292        let mut conv = [0u8; 3];
293        for (cv, &sv) in conv.iter_mut().zip(pix.iter()) {
294            *cv = (sv / mul * 256.0).trunc() as u8;
295        }
296        Rgbe8Pixel {
297            c: conv,
298            e: (exp + 128) as u8,
299        }
300    }
301}
302
303#[test]
304fn to_rgbe8_test() {
305    use crate::codecs::hdr::rgbe8;
306    let test_cases = vec![rgbe8(0, 0, 0, 0), rgbe8(1, 1, 128, 128)];
307    for &pix in &test_cases {
308        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
309    }
310    for mc in 128..255 {
311        // TODO: use inclusive range when stable
312        let pix = rgbe8(mc, mc, mc, 100);
313        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
314        let pix = rgbe8(mc, 0, mc, 130);
315        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
316        let pix = rgbe8(0, 0, mc, 140);
317        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
318        let pix = rgbe8(1, 0, mc, 150);
319        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
320        let pix = rgbe8(1, mc, 10, 128);
321        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
322        for c in 0..255 {
323            // Radiance HDR seems to be pre IEEE 754.
324            // exponent can be -128 (represented as 0u8), so some colors cannot be represented in normalized f32
325            // Let's exclude exponent value of -128 (0u8) from testing
326            let pix = rgbe8(1, mc, c, if c == 0 { 1 } else { c });
327            assert_eq!(pix, to_rgbe8(pix.to_hdr()));
328        }
329    }
330    fn relative_dist(a: Rgb<f32>, b: Rgb<f32>) -> f32 {
331        // maximal difference divided by maximal value
332        let max_diff =
333            a.0.iter()
334                .zip(b.0.iter())
335                .fold(0.0, |diff, (&a, &b)| f32::max(diff, (a - b).abs()));
336        let max_val =
337            a.0.iter()
338                .chain(b.0.iter())
339                .fold(0.0, |maxv, &a| f32::max(maxv, a));
340        if max_val == 0.0 {
341            0.0
342        } else {
343            max_diff / max_val
344        }
345    }
346    let test_values = vec![
347        0.000_001, 0.000_02, 0.000_3, 0.004, 0.05, 0.6, 7.0, 80.0, 900.0, 1_000.0, 20_000.0,
348        300_000.0,
349    ];
350    for &r in &test_values {
351        for &g in &test_values {
352            for &b in &test_values {
353                let c1 = Rgb([r, g, b]);
354                let c2 = to_rgbe8(c1).to_hdr();
355                let rel_dist = relative_dist(c1, c2);
356                // Maximal value is normalized to the range 128..256, thus we have 1/128 precision
357                assert!(
358                    rel_dist <= 1.0 / 128.0,
359                    "Relative distance ({rel_dist}) exceeds 1/128 for {c1:?} and {c2:?}"
360                );
361            }
362        }
363    }
364}
365
366#[test]
367fn runiterator_test() {
368    let data = [];
369    let mut run_iter = RunIterator::new(&data[..]);
370    assert_eq!(run_iter.next(), None);
371    let data = [5];
372    let mut run_iter = RunIterator::new(&data[..]);
373    assert_eq!(run_iter.next(), Some(Norun(0, 1)));
374    assert_eq!(run_iter.next(), None);
375    let data = [1, 1];
376    let mut run_iter = RunIterator::new(&data[..]);
377    assert_eq!(run_iter.next(), Some(Norun(0, 2)));
378    assert_eq!(run_iter.next(), None);
379    let data = [0, 0, 0];
380    let mut run_iter = RunIterator::new(&data[..]);
381    assert_eq!(run_iter.next(), Some(Run(0u8, 3)));
382    assert_eq!(run_iter.next(), None);
383    let data = [0, 0, 1, 1];
384    let mut run_iter = RunIterator::new(&data[..]);
385    assert_eq!(run_iter.next(), Some(Norun(0, 2)));
386    assert_eq!(run_iter.next(), Some(Norun(2, 2)));
387    assert_eq!(run_iter.next(), None);
388    let data = [0, 0, 0, 1, 1];
389    let mut run_iter = RunIterator::new(&data[..]);
390    assert_eq!(run_iter.next(), Some(Run(0u8, 3)));
391    assert_eq!(run_iter.next(), Some(Norun(3, 2)));
392    assert_eq!(run_iter.next(), None);
393    let data = [1, 2, 2, 2];
394    let mut run_iter = RunIterator::new(&data[..]);
395    assert_eq!(run_iter.next(), Some(Norun(0, 1)));
396    assert_eq!(run_iter.next(), Some(Run(2u8, 3)));
397    assert_eq!(run_iter.next(), None);
398    let data = [1, 1, 2, 2, 2];
399    let mut run_iter = RunIterator::new(&data[..]);
400    assert_eq!(run_iter.next(), Some(Norun(0, 2)));
401    assert_eq!(run_iter.next(), Some(Run(2u8, 3)));
402    assert_eq!(run_iter.next(), None);
403    let data = [2; 128];
404    let mut run_iter = RunIterator::new(&data[..]);
405    assert_eq!(run_iter.next(), Some(Run(2u8, 127)));
406    assert_eq!(run_iter.next(), Some(Norun(127, 1)));
407    assert_eq!(run_iter.next(), None);
408    let data = [2; 129];
409    let mut run_iter = RunIterator::new(&data[..]);
410    assert_eq!(run_iter.next(), Some(Run(2u8, 127)));
411    assert_eq!(run_iter.next(), Some(Norun(127, 2)));
412    assert_eq!(run_iter.next(), None);
413    let data = [2; 130];
414    let mut run_iter = RunIterator::new(&data[..]);
415    assert_eq!(run_iter.next(), Some(Run(2u8, 127)));
416    assert_eq!(run_iter.next(), Some(Run(2u8, 3)));
417    assert_eq!(run_iter.next(), None);
418}
419
420#[test]
421fn noruncombine_test() {
422    fn a<T>(mut v: Vec<T>, mut other: Vec<T>) -> Vec<T> {
423        v.append(&mut other);
424        v
425    }
426
427    let v = [];
428    let mut rsi = NorunCombineIterator::new(&v[..]);
429    assert_eq!(rsi.next(), None);
430
431    let v = [1];
432    let mut rsi = NorunCombineIterator::new(&v[..]);
433    assert_eq!(rsi.next(), Some(Norun(0, 1)));
434    assert_eq!(rsi.next(), None);
435
436    let v = [2, 2];
437    let mut rsi = NorunCombineIterator::new(&v[..]);
438    assert_eq!(rsi.next(), Some(Norun(0, 2)));
439    assert_eq!(rsi.next(), None);
440
441    let v = [3, 3, 3];
442    let mut rsi = NorunCombineIterator::new(&v[..]);
443    assert_eq!(rsi.next(), Some(Run(3, 3)));
444    assert_eq!(rsi.next(), None);
445
446    let v = [4, 4, 3, 3, 3];
447    let mut rsi = NorunCombineIterator::new(&v[..]);
448    assert_eq!(rsi.next(), Some(Norun(0, 2)));
449    assert_eq!(rsi.next(), Some(Run(3, 3)));
450    assert_eq!(rsi.next(), None);
451
452    let v = vec![40; 400];
453    let mut rsi = NorunCombineIterator::new(&v[..]);
454    assert_eq!(rsi.next(), Some(Run(40, 127)));
455    assert_eq!(rsi.next(), Some(Run(40, 127)));
456    assert_eq!(rsi.next(), Some(Run(40, 127)));
457    assert_eq!(rsi.next(), Some(Run(40, 19)));
458    assert_eq!(rsi.next(), None);
459
460    let v = a(a(vec![5; 3], vec![6; 129]), vec![7, 3, 7, 10, 255]);
461    let mut rsi = NorunCombineIterator::new(&v[..]);
462    assert_eq!(rsi.next(), Some(Run(5, 3)));
463    assert_eq!(rsi.next(), Some(Run(6, 127)));
464    assert_eq!(rsi.next(), Some(Norun(130, 7)));
465    assert_eq!(rsi.next(), None);
466
467    let v = a(a(vec![5; 2], vec![6; 129]), vec![7, 3, 7, 7, 255]);
468    let mut rsi = NorunCombineIterator::new(&v[..]);
469    assert_eq!(rsi.next(), Some(Norun(0, 2)));
470    assert_eq!(rsi.next(), Some(Run(6, 127)));
471    assert_eq!(rsi.next(), Some(Norun(129, 7)));
472    assert_eq!(rsi.next(), None);
473
474    let v: Vec<_> = core::iter::repeat(())
475        .flat_map(|()| 0..2)
476        .take(257)
477        .collect();
478    let mut rsi = NorunCombineIterator::new(&v[..]);
479    assert_eq!(rsi.next(), Some(Norun(0, 128)));
480    assert_eq!(rsi.next(), Some(Norun(128, 128)));
481    assert_eq!(rsi.next(), Some(Norun(256, 1)));
482    assert_eq!(rsi.next(), None);
483}
ai_image/codecs/hdr/encoder.rs

ai_image/codecs/hdr/
encoder.rs
