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: u8 = self.data[self.curidx];
151 let crun: usize = self.data[self.curidx..]
152 .iter()
153 .take_while(|&&v: u8| v == cv)
154 .take(RUN_MAX_LEN)
155 .count();
156 let ret: RunOrNot = 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: RunOrNot in NorunCombineIterator::new(data) {
258 match rnr {
259 Run(c: u8, len: usize) => {
260 assert!(len <= 127);
261 rle.push(128u8 + len as u8);
262 rle.push(c);
263 }
264 Norun(idx: usize, len: usize) => {
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: [f32; 3] = pix.0;
280 let mx: f32 = f32::max(self:pix[0], other:f32::max(self:pix[1], other: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: i32 = mx.log2().floor() as i32 + 1;
286 let mul: f32 = f32::powi(self:2.0, n:exp);
287 let mut conv: [u8; 3] = [0u8; 3];
288 for (cv: &mut u8, &sv: f32) in conv.iter_mut().zip(pix.iter()) {
289 *cv = f32::trunc(self: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}
482