image/codecs/hdr/
encoder.rs

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

image/codecs/hdr/
encoder.rs
