zune_hdr/
encoder.rs

1/*
2 * Copyright (c) 2023.
3 *
4 * This software is free software; You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
5 */
6
7//! Radiance HDR encoder
8
9use alloc::vec::Vec;
10use alloc::{format, vec};
11use std::collections::HashMap;
12
13use zune_core::bytestream::ZByteWriter;
14use zune_core::colorspace::ColorSpace;
15use zune_core::options::EncoderOptions;
16
17use crate::errors::HdrEncodeErrors;
18
19/// A simple HDR encoder
20pub struct HdrEncoder<'a> {
21    data:    &'a [f32],
22    headers: Option<&'a HashMap<String, String>>,
23    options: EncoderOptions
24}
25
26impl<'a> HdrEncoder<'a> {
27    /// Create a new HDR encoder context that can encode
28    /// the provided data
29    ///
30    /// # Arguments
31    ///  - `data`: Data to encode
32    ///  - `options`: Contains metadata for data, including width and height
33    pub fn new(data: &'a [f32], options: EncoderOptions) -> HdrEncoder<'a> {
34        Self {
35            data,
36            headers: None,
37            options
38        }
39    }
40    /// Add extra headers to be encoded  with the image
41    ///
42    /// This must be called before you call [`encode`](crate::encoder::HdrEncoder::encode)
43    /// otherwise it will have no effect.
44    ///
45    /// # Arguments:
46    /// - headers: A hashmap containing keys and values, the values will be encoded as key=value
47    /// in the hdr header before encoding
48    pub fn add_headers(&mut self, headers: &'a HashMap<String, String>) {
49        self.headers = Some(headers)
50    }
51
52    /// Calculate buffer with padding size needed for
53    /// encoding this into a vec
54    ///
55    /// This is a given upper limit and doesn't specifically mean that
56    /// your buffer should exactly be that size.
57    ///
58    /// The size of the output will depend on the nature of your data
59    pub fn expected_buffer_size(&self) -> Option<usize> {
60        self.options
61            .get_width()
62            .checked_mul(self.options.get_height())?
63            .checked_mul(4)?
64            .checked_add(1024)
65    }
66
67    /// Encode data in HDR format
68    ///
69    /// The encoder expects colorspace to be in RGB and the  length to match
70    ///
71    /// otherwise it's an error to try and encode
72    ///
73    /// The floating point data is expected to be normalized between 0.0 and 1.0, it will not be clipped
74    /// if not in this range
75    ///
76    /// # Returns
77    ///
78    /// The encoded data in case of no errors in case of errors returns error
79    ///
80    /// # See also
81    ///
82    /// - [encode_into](crate::encoder::HdrEncoder::encode_into): Encodes into a pre-allocated buffer
83    pub fn encode(&self) -> Result<Vec<u8>, HdrEncodeErrors> {
84        let expected = self
85            .options
86            .get_width()
87            .checked_mul(self.options.get_height())
88            .unwrap()
89            .checked_mul(3) // RGB
90            .unwrap();
91        let found = self.data.len();
92
93        if expected != found {
94            return Err(HdrEncodeErrors::WrongInputSize(expected, found));
95        }
96        if self.options.get_colorspace() != ColorSpace::RGB {
97            return Err(HdrEncodeErrors::UnsupportedColorspace(
98                self.options.get_colorspace()
99            ));
100        }
101        let size = self
102            .expected_buffer_size()
103            .ok_or_else(|| HdrEncodeErrors::Static("overflow detected"))?;
104
105        let mut out = vec![0_u8; size];
106
107        let pos = self.encode_into(&mut out)?;
108        out.truncate(pos);
109
110        Ok(out)
111    }
112
113    /// Encode into a pre-allocated buffer
114    ///
115    /// The size of output cannot be determined up until compression is over hence
116    /// we cannot be sure if the output buffer will be enough.
117    ///
118    /// The library provides [expected_buffer_size](crate::HdrEncoder::expected_buffer_size) which
119    /// guarantees that if your buffer is that big, encoding will always succeed
120    ///
121    /// # Arguments:
122    /// - out: The output buffer to write bytes into
123    ///
124    /// # Returns
125    /// - Ok(usize):  The number of bytes written into out
126    /// - Err(HdrEncodeErrors): An error if something occurred
127    ///
128    /// # Examples
129    /// Encode a black image of 10x10
130    ///```
131    /// use zune_core::bit_depth::BitDepth;
132    /// use zune_core::colorspace::ColorSpace;
133    /// use zune_core::options::EncoderOptions;
134    /// use zune_hdr::HdrEncoder;
135    /// let w = 10;
136    /// let h = 10;
137    /// let comp = 3;
138    /// let data = vec![0.0_f32;w*h*comp];
139    /// let opts = EncoderOptions::new(w,h,ColorSpace::RGB,BitDepth::Float32);
140    /// let encoder = HdrEncoder::new(&data,opts);
141    /// // create output buffer , this is the upper limit on it
142    /// let mut output = vec![0;encoder.expected_buffer_size().unwrap()];
143    /// let size = encoder.encode_into(&mut output).unwrap();
144    /// // trim it so that the buffer only contains encoded hdr
145    /// output.truncate(size);
146    ///```  
147    pub fn encode_into(&self, out: &mut [u8]) -> Result<usize, HdrEncodeErrors> {
148        let mut writer = ZByteWriter::new(out);
149
150        // write headers
151        {
152            writer.write_all(b"#?RADIANCE\n")?;
153            writer.write_all(b"SOFTWARE=zune-hdr\n")?;
154            if let Some(headers) = self.headers {
155                for (k, v) in headers {
156                    writer.write_all(format!("{}={}\n", k, v).as_bytes())?;
157                }
158            }
159            writer.write_all(b"FORMAT=32-bit_rle_rgbe\n\n")?;
160
161            // write lengths
162            let length_format = format!(
163                "-Y {} +X {}\n",
164                self.options.get_height(),
165                self.options.get_width()
166            );
167
168            writer.write_all(length_format.as_bytes())?;
169        }
170        let width = self.options.get_width();
171
172        let scanline_stride = width * 3;
173
174        let mut in_scanline = vec![0_u8; width * 4]; // RGBE
175
176        for scanline in self.data.chunks_exact(scanline_stride) {
177            if !(8..=0x7fff).contains(&width) {
178                for (pixels, out) in scanline
179                    .chunks_exact(3)
180                    .zip(in_scanline.chunks_exact_mut(4))
181                {
182                    float_to_rgbe(pixels.try_into().unwrap(), out.try_into().unwrap());
183                }
184                writer.write_all(&in_scanline)?;
185            } else {
186                writer.write_u8_err(2)?;
187                writer.write_u8_err(2)?;
188                writer.write_u8_err((width >> 8) as u8)?;
189                writer.write_u8_err((width & 255) as u8)?;
190
191                for (pixels, out) in scanline
192                    .chunks_exact(3)
193                    .zip(in_scanline.chunks_exact_mut(4))
194                {
195                    float_to_rgbe(pixels.try_into().unwrap(), out.try_into().unwrap());
196                }
197                for i in 0..4 {
198                    rle(&in_scanline[i..], &mut writer, width)?;
199                }
200            }
201        }
202        // truncate position to where we reached
203        let position = writer.position();
204        Ok(position)
205    }
206}
207
208fn rle(data: &[u8], writer: &mut ZByteWriter, width: usize) -> Result<(), &'static str> {
209    const MIN_RLE: usize = 4;
210    let mut cur = 0;
211
212    while cur < width {
213        let mut run_count = 0;
214        let mut old_run_count = 0;
215        let mut beg_run = cur;
216        let mut buf: [u8; 2] = [0; 2];
217
218        while run_count < MIN_RLE && beg_run < width {
219            beg_run += run_count;
220            old_run_count = run_count;
221            run_count = 1;
222
223            while (beg_run + run_count < width)
224                && (run_count < 127)
225                && (data[beg_run * 4] == data[(beg_run + run_count) * 4])
226            {
227                run_count += 1;
228            }
229        }
230
231        if (old_run_count > 1) && (old_run_count == beg_run - cur) {
232            buf[0] = (128 + old_run_count) as u8;
233            buf[1] = data[cur * 4];
234            writer.write_all(&buf)?;
235            cur = beg_run;
236        }
237
238        while cur < beg_run {
239            let nonrun_count = 128.min(beg_run - cur);
240            buf[0] = nonrun_count as u8;
241            writer.write_u8(buf[0]);
242            for i in 0..nonrun_count {
243                writer.write_u8_err(data[(cur + i) * 4])?;
244            }
245
246            cur += nonrun_count;
247        }
248
249        if run_count >= MIN_RLE {
250            buf[0] = (128 + run_count) as u8;
251            buf[1] = data[beg_run * 4];
252            writer.write_all(&buf)?;
253            cur += run_count;
254        }
255    }
256    Ok(())
257}
258
259fn float_to_rgbe(rgb: &[f32; 3], rgbe: &mut [u8; 4]) {
260    let v = rgb.iter().fold(f32::MIN, |x, y| x.max(*y));
261
262    if v > 1e-32 {
263        let old_v = v;
264        let (mut v, e) = frexp(v);
265        v = v * 256. / old_v;
266        rgbe[0] = (rgb[0] * v).clamp(0.0, 255.0) as u8;
267        rgbe[1] = (rgb[1] * v).clamp(0.0, 255.0) as u8;
268        rgbe[2] = (rgb[2] * v).clamp(0.0, 255.0) as u8;
269
270        rgbe[3] = (e.wrapping_add(128)) as u8;
271    } else {
272        rgbe.fill(0);
273    }
274}
275
276#[rustfmt::skip]
277pub fn abs(num: f32) -> f32
278{
279    // standard compliant
280    // handles NAN and infinity.
281    // pretty cool
282    f32::from_bits(num.to_bits() & (i32::MAX as u32))
283}
284
285/// Implementation of signum that works in no_std
286#[rustfmt::skip]
287pub fn signum(num: f32) -> f32
288{
289    if num.is_nan() { f32::NAN } else if num.is_infinite()
290    {
291        if num.is_sign_positive() { 0.0 } else { -0.0 }
292    } else if num > 0.0 { 1.0 } else { -1.0 }
293}
294
295fn floor(num: f32) -> f32 {
296    if num.is_nan() || num.is_infinite() {
297        /* handle infinities and nan */
298        return num;
299    }
300    let n = num as u64;
301    let d = n as f32;
302
303    if d == num || num >= 0.0 {
304        d
305    } else {
306        d - 1.0
307    }
308}
309
310/// Fast log2 approximation
311/// (we really don't need that accurate)
312#[allow(clippy::cast_precision_loss, clippy::cast_sign_loss)]
313fn fast_log2(x: f32) -> f32 {
314    /*
315     * Fast log approximation from
316     * https://github.com/romeric/fastapprox
317     *
318     * Some pretty good stuff.
319     */
320    let vx = x.to_bits();
321    let mx = (vx & 0x007F_FFFF) | 0x3f00_0000;
322    let mx_f = f32::from_bits(mx);
323
324    let mut y = vx as f32;
325    // 1/(1<<23)
326    y *= 1.192_092_9e-7;
327
328    y - 124.225_52 - 1.498_030_3 * mx_f - 1.725_88 / (0.352_088_72 + mx_f)
329}
330
331/// non standard frexp implementation
332fn frexp(s: f32) -> (f32, i32) {
333    // from https://stackoverflow.com/a/55696477
334    if 0.0 == s {
335        (s, 0)
336    } else {
337        let lg = fast_log2(abs(s));
338        let lg_floor = floor(lg);
339        // Note: This is the only reason we need the standard library
340        // I haven't found a goof exp2 function, fast_exp2 doesn't work
341        // and libm/musl exp2 introduces visible color distortions and is slow, so for
342        // now let's stick to whatever the platform provides
343        let x = (lg - lg_floor - 1.0).exp2();
344        let exp = lg_floor + 1.0;
345        (signum(s) * x, exp as i32)
346    }
347}
zune_hdr/encoder.rs

zune_hdr/
encoder.rs
