ultrahdr_rs/
decode.rs

1//! Ultra HDR decoder.
2
3use ultrahdr_core::gainmap::apply::{apply_gainmap, HdrOutputFormat};
4use ultrahdr_core::metadata::{
5    mpf::{find_jpeg_boundaries, parse_mpf},
6    xmp::parse_xmp,
7};
8use ultrahdr_core::{
9    ColorGamut, ColorTransfer, Error, GainMap, GainMapMetadata, PixelFormat, RawImage, Result,
10    Unstoppable,
11};
12
13use crate::jpeg::{extract_icc_profile, find_xmp_data};
14
15/// Ultra HDR decoder.
16///
17/// Decodes Ultra HDR JPEGs, extracting the SDR base image, gain map,
18/// and metadata. Can reconstruct HDR content at various display
19/// brightness levels.
20pub struct Decoder {
21    data: Vec<u8>,
22    metadata: Option<GainMapMetadata>,
23    primary_jpeg: Option<(usize, usize)>,
24    gainmap_jpeg: Option<(usize, usize)>,
25    is_ultrahdr: bool,
26}
27
28impl Decoder {
29    /// Create a new decoder from JPEG data.
30    pub fn new(data: &[u8]) -> Result<Self> {
31        let mut decoder = Self {
32            data: data.to_vec(),
33            metadata: None,
34            primary_jpeg: None,
35            gainmap_jpeg: None,
36            is_ultrahdr: false,
37        };
38
39        decoder.parse()?;
40        Ok(decoder)
41    }
42
43    /// Check if this is a valid Ultra HDR image.
44    pub fn is_ultrahdr(&self) -> bool {
45        self.is_ultrahdr
46    }
47
48    /// Get the gain map metadata.
49    pub fn metadata(&self) -> Option<&GainMapMetadata> {
50        self.metadata.as_ref()
51    }
52
53    /// Get the raw gain map JPEG data.
54    pub fn gainmap_jpeg(&self) -> Option<&[u8]> {
55        self.gainmap_jpeg.map(|(start, end)| &self.data[start..end])
56    }
57
58    /// Decode the SDR base image.
59    pub fn decode_sdr(&self) -> Result<RawImage> {
60        let (start, end) = self
61            .primary_jpeg
62            .ok_or_else(|| Error::DecodeError("No primary image found".into()))?;
63
64        let primary_data = &self.data[start..end];
65        decode_jpeg_to_rgb(primary_data)
66    }
67
68    /// Decode the gain map.
69    pub fn decode_gainmap(&self) -> Result<GainMap> {
70        let (start, end) = self
71            .gainmap_jpeg
72            .ok_or_else(|| Error::DecodeError("No gain map found".into()))?;
73
74        let gainmap_data = &self.data[start..end];
75        let decoded = decode_jpeg_to_grayscale(gainmap_data)?;
76
77        Ok(GainMap {
78            width: decoded.width,
79            height: decoded.height,
80            channels: 1,
81            data: decoded.data,
82        })
83    }
84
85    /// Decode to HDR at the specified display boost level.
86    ///
87    /// `display_boost` is the ratio of display peak brightness to SDR white.
88    /// For example:
89    /// - 1.0 = SDR display (no HDR enhancement)
90    /// - 4.0 = Display capable of 4x SDR brightness
91    /// - ~49.0 = Full HDR10 (10000 nits / 203 SDR nits)
92    pub fn decode_hdr(&self, display_boost: f32) -> Result<RawImage> {
93        self.decode_hdr_with_format(display_boost, HdrOutputFormat::LinearFloat)
94    }
95
96    /// Decode to HDR with a specific output format.
97    pub fn decode_hdr_with_format(
98        &self,
99        display_boost: f32,
100        format: HdrOutputFormat,
101    ) -> Result<RawImage> {
102        if !self.is_ultrahdr {
103            return Err(Error::DecodeError("Not an Ultra HDR image".into()));
104        }
105
106        let metadata = self
107            .metadata
108            .as_ref()
109            .ok_or_else(|| Error::DecodeError("No gain map metadata".into()))?;
110
111        let sdr = self.decode_sdr()?;
112        let gainmap = self.decode_gainmap()?;
113
114        apply_gainmap(&sdr, &gainmap, metadata, display_boost, format, Unstoppable)
115    }
116
117    /// Parse the Ultra HDR structure.
118    fn parse(&mut self) -> Result<()> {
119        // Check for valid JPEG
120        if self.data.len() < 4 || self.data[0] != 0xFF || self.data[1] != 0xD8 {
121            return Err(Error::DecodeError("Not a valid JPEG".into()));
122        }
123
124        // Try to find XMP metadata with hdrgm namespace
125        if let Some(xmp) = find_xmp_data(&self.data) {
126            if xmp.contains("hdrgm:") || xmp.contains("http://ns.adobe.com/hdr-gain-map/") {
127                if let Ok((metadata, _gainmap_len)) = parse_xmp(&xmp) {
128                    self.metadata = Some(metadata);
129                    self.is_ultrahdr = true;
130                }
131            }
132        }
133
134        // Try to parse MPF to find gain map
135        if let Ok(images) = parse_mpf(&self.data) {
136            if images.len() >= 2 {
137                // First image is primary, second is gain map
138                self.primary_jpeg = Some(images[0]);
139                self.gainmap_jpeg = Some(images[1]);
140                self.is_ultrahdr = true;
141            }
142        }
143
144        // Fallback: look for multiple JPEGs in the file
145        if self.gainmap_jpeg.is_none() {
146            let boundaries = find_jpeg_boundaries(&self.data);
147            if boundaries.len() >= 2 {
148                self.primary_jpeg = Some(boundaries[0]);
149                self.gainmap_jpeg = Some(boundaries[1]);
150            }
151        }
152
153        // Set primary to full data if not found via MPF
154        if self.primary_jpeg.is_none() {
155            self.primary_jpeg = Some((0, self.data.len()));
156        }
157
158        Ok(())
159    }
160
161    /// Get the ICC profile from the primary image if present.
162    pub fn icc_profile(&self) -> Option<Vec<u8>> {
163        extract_icc_profile(&self.data)
164    }
165
166    /// Get information about the decoded image dimensions.
167    pub fn dimensions(&self) -> Result<(u32, u32)> {
168        let sdr = self.decode_sdr()?;
169        Ok((sdr.width, sdr.height))
170    }
171}
172
173/// Decode JPEG to RGB.
174fn decode_jpeg_to_rgb(jpeg_data: &[u8]) -> Result<RawImage> {
175    use jpegli::decoder::{Decoder as JpegDecoder, PixelFormat as JpegPixelFormat};
176    let decoded = JpegDecoder::new()
177        .output_format(JpegPixelFormat::Rgb)
178        .decode(jpeg_data)
179        .map_err(|e| Error::DecodeError(format!("JPEG decode failed: {}", e)))?;
180
181    let width = decoded.width;
182    let height = decoded.height;
183    let pixels = &decoded.data;
184    let bpp = decoded.bytes_per_pixel();
185
186    // Convert to RGBA if needed
187    let data = if bpp == 3 {
188        // RGB -> RGBA
189        let mut rgba = Vec::with_capacity((width * height * 4) as usize);
190        for chunk in pixels.chunks(3) {
191            rgba.push(chunk[0]);
192            rgba.push(chunk[1]);
193            rgba.push(chunk[2]);
194            rgba.push(255);
195        }
196        rgba
197    } else if bpp == 4 {
198        pixels.to_vec()
199    } else if bpp == 1 {
200        // Grayscale -> RGBA
201        let mut rgba = Vec::with_capacity((width * height * 4) as usize);
202        for &g in pixels {
203            rgba.push(g);
204            rgba.push(g);
205            rgba.push(g);
206            rgba.push(255);
207        }
208        rgba
209    } else {
210        return Err(Error::DecodeError(format!(
211            "Unsupported bytes per pixel: {}",
212            bpp
213        )));
214    };
215
216    Ok(RawImage {
217        width,
218        height,
219        stride: width * 4,
220        data,
221        format: PixelFormat::Rgba8,
222        gamut: ColorGamut::Bt709, // Assume sRGB for SDR
223        transfer: ColorTransfer::Srgb,
224    })
225}
226
227/// Decode JPEG to grayscale.
228fn decode_jpeg_to_grayscale(jpeg_data: &[u8]) -> Result<RawImage> {
229    use jpegli::decoder::{Decoder as JpegDecoder, PixelFormat as JpegPixelFormat};
230    let decoded = JpegDecoder::new()
231        .output_format(JpegPixelFormat::Gray)
232        .decode(jpeg_data)
233        .map_err(|e| Error::DecodeError(format!("JPEG decode failed: {}", e)))?;
234
235    let width = decoded.width;
236    let height = decoded.height;
237    let pixels = &decoded.data;
238    let bpp = decoded.bytes_per_pixel();
239
240    // Convert to grayscale if needed
241    let data = if bpp == 1 {
242        pixels.to_vec()
243    } else if bpp == 3 {
244        // RGB -> Grayscale (using luminance)
245        pixels
246            .chunks(3)
247            .map(|rgb| {
248                let r = rgb[0] as f32;
249                let g = rgb[1] as f32;
250                let b = rgb[2] as f32;
251                // BT.709 luminance
252                (0.2126_f32 * r + 0.7152 * g + 0.0722 * b).clamp(0.0, 255.0) as u8
253            })
254            .collect()
255    } else {
256        return Err(Error::DecodeError(format!(
257            "Unsupported bytes per pixel for grayscale: {}",
258            bpp
259        )));
260    };
261
262    Ok(RawImage {
263        width,
264        height,
265        stride: width,
266        data,
267        format: PixelFormat::Gray8,
268        gamut: ColorGamut::Bt709,
269        transfer: ColorTransfer::Srgb,
270    })
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    #[test]
278    fn test_decoder_invalid_data() {
279        let result = Decoder::new(&[0, 1, 2, 3]);
280        assert!(result.is_err());
281    }
282
283    #[test]
284    fn test_decoder_minimal_jpeg() {
285        // Minimal JPEG (just SOI + EOI)
286        let data = vec![0xFF, 0xD8, 0xFF, 0xD9];
287        let decoder = Decoder::new(&data);
288        // This will fail to parse as a valid image but won't error on construction
289        assert!(decoder.is_ok());
290        assert!(!decoder.unwrap().is_ultrahdr());
291    }
292}
ultrahdr_rs/decode.rs

ultrahdr_rs/
decode.rs
