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    fn next(&mut self) -> Option<Self::Item> {
185        loop {
186            match self.prev.take() {
187                Some(Run(c, len)) => {
188                    // Just return stored run
189                    return Some(Run(c, len));
190                }
191                Some(Norun(idx, len)) => {
192                    // Let's see if we need to continue norun
193                    match self.runiter.next() {
194                        Some(Norun(_, len1)) => {
195                            // norun continues
196                            let clen = len + len1; // combined length
197                            match clen.cmp(&NORUN_MAX_LEN) {
198                                Ordering::Equal => return Some(Norun(idx, clen)),
199                                Ordering::Greater => {
200                                    // combined norun exceeds maximum length. store extra part of norun
201                                    self.prev =
202                                        Some(Norun(idx + NORUN_MAX_LEN, clen - NORUN_MAX_LEN));
203                                    // then return maximal norun
204                                    return Some(Norun(idx, NORUN_MAX_LEN));
205                                }
206                                Ordering::Less => {
207                                    // len + len1 < NORUN_MAX_LEN
208                                    self.prev = Some(Norun(idx, len + len1));
209                                    // combine and continue loop
210                                }
211                            }
212                        }
213                        Some(Run(c, len1)) => {
214                            // Run encountered. Store it
215                            self.prev = Some(Run(c, len1));
216                            return Some(Norun(idx, len)); // and return combined norun
217                        }
218                        None => {
219                            // End of sequence
220                            return Some(Norun(idx, len)); // return combined norun
221                        }
222                    }
223                } // End match self.prev.take() == Some(NoRun())
224                None => {
225                    // No norun to combine
226                    match self.runiter.next() {
227                        Some(Norun(idx, len)) => {
228                            self.prev = Some(Norun(idx, len));
229                            // store for combine and continue the loop
230                        }
231                        Some(Run(c, len)) => {
232                            // Some run. Just return it
233                            return Some(Run(c, len));
234                        }
235                        None => {
236                            // That's all, folks
237                            return None;
238                        }
239                    }
240                } // End match self.prev.take() == None
241            } // End match
242        } // End loop
243    }
244}
245
246// Appends RLE compressed ```data``` to ```rle```
247fn rle_compress(data: &[u8], rle: &mut Vec<u8>) {
248    rle.clear();
249    if data.is_empty() {
250        rle.push(0); // Technically correct. It means read next 0 bytes.
251        return;
252    }
253    // Task: split data into chunks of repeating (max 127) and non-repeating bytes (max 128)
254    // Prepend non-repeating chunk with its length
255    // Replace repeating byte with (run length + 128) and the byte
256    for rnr in NorunCombineIterator::new(data) {
257        match rnr {
258            Run(c, len) => {
259                assert!(len <= 127);
260                rle.push(128u8 + len as u8);
261                rle.push(c);
262            }
263            Norun(idx, len) => {
264                assert!(len <= 128);
265                rle.push(len as u8);
266                rle.extend_from_slice(&data[idx..idx + len]);
267            }
268        }
269    }
270}
271
272fn write_rgbe8<W: Write>(w: &mut W, v: Rgbe8Pixel) -> Result<()> {
273    w.write_all(&[v.c[0], v.c[1], v.c[2], v.e])
274}
275
276/// Converts ```Rgb<f32>``` into ```Rgbe8Pixel```
277pub(crate) fn to_rgbe8(pix: Rgb<f32>) -> Rgbe8Pixel {
278    let pix = pix.0;
279    let mx = f32::max(pix[0], f32::max(pix[1], pix[2]));
280    if mx <= 0.0 {
281        Rgbe8Pixel { c: [0, 0, 0], e: 0 }
282    } else {
283        // let (frac, exp) = mx.frexp(); // unstable yet
284        let exp = mx.log2().floor() as i32 + 1;
285        let mul = f32::powi(2.0, exp);
286        let mut conv = [0u8; 3];
287        for (cv, &sv) in conv.iter_mut().zip(pix.iter()) {
288            *cv = f32::trunc(sv / mul * 256.0) as u8;
289        }
290        Rgbe8Pixel {
291            c: conv,
292            e: (exp + 128) as u8,
293        }
294    }
295}
296
297#[test]
298fn to_rgbe8_test() {
299    use crate::codecs::hdr::rgbe8;
300    let test_cases = vec![rgbe8(0, 0, 0, 0), rgbe8(1, 1, 128, 128)];
301    for &pix in &test_cases {
302        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
303    }
304    for mc in 128..255 {
305        // TODO: use inclusive range when stable
306        let pix = rgbe8(mc, mc, mc, 100);
307        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
308        let pix = rgbe8(mc, 0, mc, 130);
309        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
310        let pix = rgbe8(0, 0, mc, 140);
311        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
312        let pix = rgbe8(1, 0, mc, 150);
313        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
314        let pix = rgbe8(1, mc, 10, 128);
315        assert_eq!(pix, to_rgbe8(pix.to_hdr()));
316        for c in 0..255 {
317            // Radiance HDR seems to be pre IEEE 754.
318            // exponent can be -128 (represented as 0u8), so some colors cannot be represented in normalized f32
319            // Let's exclude exponent value of -128 (0u8) from testing
320            let pix = rgbe8(1, mc, c, if c == 0 { 1 } else { c });
321            assert_eq!(pix, to_rgbe8(pix.to_hdr()));
322        }
323    }
324    fn relative_dist(a: Rgb<f32>, b: Rgb<f32>) -> f32 {
325        // maximal difference divided by maximal value
326        let max_diff =
327            a.0.iter()
328                .zip(b.0.iter())
329                .fold(0.0, |diff, (&a, &b)| f32::max(diff, (a - b).abs()));
330        let max_val =
331            a.0.iter()
332                .chain(b.0.iter())
333                .fold(0.0, |maxv, &a| f32::max(maxv, a));
334        if max_val == 0.0 {
335            0.0
336        } else {
337            max_diff / max_val
338        }
339    }
340    let test_values = vec![
341        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,
342        300_000.0,
343    ];
344    for &r in &test_values {
345        for &g in &test_values {
346            for &b in &test_values {
347                let c1 = Rgb([r, g, b]);
348                let c2 = to_rgbe8(c1).to_hdr();
349                let rel_dist = relative_dist(c1, c2);
350                // Maximal value is normalized to the range 128..256, thus we have 1/128 precision
351                assert!(
352                    rel_dist <= 1.0 / 128.0,
353                    "Relative distance ({}) exceeds 1/128 for {:?} and {:?}",
354                    rel_dist,
355                    c1,
356                    c2
357                );
358            }
359        }
360    }
361}
362
363#[test]
364fn runiterator_test() {
365    let data = [];
366    let mut run_iter = RunIterator::new(&data[..]);
367    assert_eq!(run_iter.next(), None);
368    let data = [5];
369    let mut run_iter = RunIterator::new(&data[..]);
370    assert_eq!(run_iter.next(), Some(Norun(0, 1)));
371    assert_eq!(run_iter.next(), None);
372    let data = [1, 1];
373    let mut run_iter = RunIterator::new(&data[..]);
374    assert_eq!(run_iter.next(), Some(Norun(0, 2)));
375    assert_eq!(run_iter.next(), None);
376    let data = [0, 0, 0];
377    let mut run_iter = RunIterator::new(&data[..]);
378    assert_eq!(run_iter.next(), Some(Run(0u8, 3)));
379    assert_eq!(run_iter.next(), None);
380    let data = [0, 0, 1, 1];
381    let mut run_iter = RunIterator::new(&data[..]);
382    assert_eq!(run_iter.next(), Some(Norun(0, 2)));
383    assert_eq!(run_iter.next(), Some(Norun(2, 2)));
384    assert_eq!(run_iter.next(), None);
385    let data = [0, 0, 0, 1, 1];
386    let mut run_iter = RunIterator::new(&data[..]);
387    assert_eq!(run_iter.next(), Some(Run(0u8, 3)));
388    assert_eq!(run_iter.next(), Some(Norun(3, 2)));
389    assert_eq!(run_iter.next(), None);
390    let data = [1, 2, 2, 2];
391    let mut run_iter = RunIterator::new(&data[..]);
392    assert_eq!(run_iter.next(), Some(Norun(0, 1)));
393    assert_eq!(run_iter.next(), Some(Run(2u8, 3)));
394    assert_eq!(run_iter.next(), None);
395    let data = [1, 1, 2, 2, 2];
396    let mut run_iter = RunIterator::new(&data[..]);
397    assert_eq!(run_iter.next(), Some(Norun(0, 2)));
398    assert_eq!(run_iter.next(), Some(Run(2u8, 3)));
399    assert_eq!(run_iter.next(), None);
400    let data = [2; 128];
401    let mut run_iter = RunIterator::new(&data[..]);
402    assert_eq!(run_iter.next(), Some(Run(2u8, 127)));
403    assert_eq!(run_iter.next(), Some(Norun(127, 1)));
404    assert_eq!(run_iter.next(), None);
405    let data = [2; 129];
406    let mut run_iter = RunIterator::new(&data[..]);
407    assert_eq!(run_iter.next(), Some(Run(2u8, 127)));
408    assert_eq!(run_iter.next(), Some(Norun(127, 2)));
409    assert_eq!(run_iter.next(), None);
410    let data = [2; 130];
411    let mut run_iter = RunIterator::new(&data[..]);
412    assert_eq!(run_iter.next(), Some(Run(2u8, 127)));
413    assert_eq!(run_iter.next(), Some(Run(2u8, 3)));
414    assert_eq!(run_iter.next(), None);
415}
416
417#[test]
418fn noruncombine_test() {
419    fn a<T>(mut v: Vec<T>, mut other: Vec<T>) -> Vec<T> {
420        v.append(&mut other);
421        v
422    }
423
424    let v = [];
425    let mut rsi = NorunCombineIterator::new(&v[..]);
426    assert_eq!(rsi.next(), None);
427
428    let v = [1];
429    let mut rsi = NorunCombineIterator::new(&v[..]);
430    assert_eq!(rsi.next(), Some(Norun(0, 1)));
431    assert_eq!(rsi.next(), None);
432
433    let v = [2, 2];
434    let mut rsi = NorunCombineIterator::new(&v[..]);
435    assert_eq!(rsi.next(), Some(Norun(0, 2)));
436    assert_eq!(rsi.next(), None);
437
438    let v = [3, 3, 3];
439    let mut rsi = NorunCombineIterator::new(&v[..]);
440    assert_eq!(rsi.next(), Some(Run(3, 3)));
441    assert_eq!(rsi.next(), None);
442
443    let v = [4, 4, 3, 3, 3];
444    let mut rsi = NorunCombineIterator::new(&v[..]);
445    assert_eq!(rsi.next(), Some(Norun(0, 2)));
446    assert_eq!(rsi.next(), Some(Run(3, 3)));
447    assert_eq!(rsi.next(), None);
448
449    let v = vec![40; 400];
450    let mut rsi = NorunCombineIterator::new(&v[..]);
451    assert_eq!(rsi.next(), Some(Run(40, 127)));
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, 19)));
455    assert_eq!(rsi.next(), None);
456
457    let v = a(a(vec![5; 3], vec![6; 129]), vec![7, 3, 7, 10, 255]);
458    let mut rsi = NorunCombineIterator::new(&v[..]);
459    assert_eq!(rsi.next(), Some(Run(5, 3)));
460    assert_eq!(rsi.next(), Some(Run(6, 127)));
461    assert_eq!(rsi.next(), Some(Norun(130, 7)));
462    assert_eq!(rsi.next(), None);
463
464    let v = a(a(vec![5; 2], vec![6; 129]), vec![7, 3, 7, 7, 255]);
465    let mut rsi = NorunCombineIterator::new(&v[..]);
466    assert_eq!(rsi.next(), Some(Norun(0, 2)));
467    assert_eq!(rsi.next(), Some(Run(6, 127)));
468    assert_eq!(rsi.next(), Some(Norun(129, 7)));
469    assert_eq!(rsi.next(), None);
470
471    let v: Vec<_> = std::iter::repeat(())
472        .flat_map(|_| (0..2))
473        .take(257)
474        .collect();
475    let mut rsi = NorunCombineIterator::new(&v[..]);
476    assert_eq!(rsi.next(), Some(Norun(0, 128)));
477    assert_eq!(rsi.next(), Some(Norun(128, 128)));
478    assert_eq!(rsi.next(), Some(Norun(256, 1)));
479    assert_eq!(rsi.next(), None);
480}
image/codecs/hdr/encoder.rs

image/codecs/hdr/
encoder.rs
