edgefirst_codec/
decoder.rs

1// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! [`ImageDecoder`] — reusable decoder state for zero-allocation hot loops.
5
6use crate::error::CodecError;
7use crate::options::ImageInfo;
8use crate::pixel::ImagePixel;
9use edgefirst_tensor::{Tensor, TensorDyn};
10use std::io::Read;
11
12/// Reusable image decoder with internal scratch buffers.
13///
14/// Create one `ImageDecoder` at program initialisation and pass it to
15/// [`ImageLoad::load_image`](crate::ImageLoad::load_image) in the hot loop.
16/// The scratch buffers grow to the high-water mark and are reused across calls
17/// — no per-frame allocations after the first few frames.
18///
19/// The decoder always produces the source's native format and configures the
20/// destination tensor's dimensions + pixel format accordingly (JPEG →
21/// `Nv12`/`Grey`, PNG → `Rgb`/`Rgba`/`Grey`). It never colour-converts or
22/// rotates; use `ImageProcessor::convert()` for that.
23///
24/// # Example
25///
26/// ```rust,no_run
27/// use edgefirst_codec::{ImageDecoder, ImageLoad};
28/// use edgefirst_tensor::{CpuAccess, Tensor, TensorTrait, TensorMemory, PixelFormat};
29///
30/// let mut decoder = ImageDecoder::new();
31/// // Allocate a buffer large enough for the biggest expected image. The
32/// // decoder CPU-writes the pixels, so declare `CpuAccess::Write`.
33/// let mut tensor = Tensor::<u8>::image(1920, 1080, PixelFormat::Nv12, Some(TensorMemory::Mem),
34///                                       CpuAccess::Write)
35///     .expect("alloc");
36///
37/// for _ in 0..100 {
38///     let frame = std::fs::read("frame.jpg").unwrap();
39///     let _info = tensor.load_image(&mut decoder, &frame);
40/// }
41/// ```
42pub struct ImageDecoder {
43    /// Reusable JPEG decoder state (Huffman tables, MCU scratch buffers).
44    pub(crate) jpeg_state: crate::jpeg::JpegDecoderState,
45    /// Scratch buffer for PNG decode output.
46    pub(crate) scratch: Vec<u8>,
47    /// Buffer for `Read` → `&[u8]` conversion.
48    pub(crate) input_buffer: Vec<u8>,
49}
50
51impl ImageDecoder {
52    /// Create a new decoder with empty scratch buffers.
53    pub fn new() -> Self {
54        Self {
55            jpeg_state: crate::jpeg::JpegDecoderState::new(),
56            scratch: Vec::new(),
57            input_buffer: Vec::new(),
58        }
59    }
60
61    /// Decode image data into a typed tensor, configuring its dimensions and
62    /// pixel format to the decoded native format.
63    ///
64    /// Detects the image format (JPEG or PNG) from magic bytes.
65    ///
66    /// # Errors
67    ///
68    /// - [`CodecError::InsufficientCapacity`] if the image is larger than the
69    ///   tensor's allocation
70    /// - [`CodecError::UnsupportedDtype`] if `T` is not valid for the native
71    ///   format (JPEG NV12/GREY require `u8`)
72    /// - [`CodecError::InvalidData`] if the data is not a valid JPEG or PNG
73    pub fn decode_into<T: ImagePixel>(
74        &mut self,
75        data: &[u8],
76        dst: &mut Tensor<T>,
77    ) -> crate::Result<ImageInfo> {
78        if is_jpeg(data) {
79            crate::jpeg::decode_jpeg_into(data, dst, &mut self.jpeg_state)
80        } else if is_png(data) {
81            crate::png::decode_png_into(data, dst, &mut self.scratch)
82        } else {
83            Err(CodecError::InvalidData(
84                "unrecognized image format (expected JPEG or PNG magic bytes)".into(),
85            ))
86        }
87    }
88
89    /// Decode image data into a type-erased tensor.
90    pub fn decode_into_dyn(
91        &mut self,
92        data: &[u8],
93        dst: &mut TensorDyn,
94    ) -> crate::Result<ImageInfo> {
95        match dst {
96            TensorDyn::U8(t) => self.decode_into(data, t),
97            TensorDyn::I8(t) => self.decode_into(data, t),
98            TensorDyn::U16(t) => self.decode_into(data, t),
99            TensorDyn::I16(t) => self.decode_into(data, t),
100            TensorDyn::F32(t) => self.decode_into(data, t),
101            other => Err(CodecError::UnsupportedDtype(other.dtype())),
102        }
103    }
104
105    /// Read all bytes from a `Read` source into the internal input buffer, then
106    /// decode. The input buffer is reused across calls.
107    pub fn decode_from_reader<T: ImagePixel, R: Read>(
108        &mut self,
109        mut reader: R,
110        dst: &mut Tensor<T>,
111    ) -> crate::Result<ImageInfo> {
112        self.input_buffer.clear();
113        reader.read_to_end(&mut self.input_buffer)?;
114        decode_into_inner(
115            &mut self.jpeg_state,
116            &mut self.scratch,
117            &self.input_buffer,
118            dst,
119        )
120    }
121
122    /// Read all bytes from a `Read` source, then decode into a type-erased
123    /// tensor.
124    pub fn decode_from_reader_dyn<R: Read>(
125        &mut self,
126        mut reader: R,
127        dst: &mut TensorDyn,
128    ) -> crate::Result<ImageInfo> {
129        self.input_buffer.clear();
130        reader.read_to_end(&mut self.input_buffer)?;
131        decode_into_inner_dyn(
132            &mut self.jpeg_state,
133            &mut self.scratch,
134            &self.input_buffer,
135            dst,
136        )
137    }
138}
139
140/// Internal decode entry point parameterised over disjoint `&mut` borrows so
141/// callers can read from a `&[u8]` borrowed from one field of [`ImageDecoder`]
142/// while mutably borrowing the others.
143pub(crate) fn decode_into_inner<T: ImagePixel>(
144    jpeg_state: &mut crate::jpeg::JpegDecoderState,
145    scratch: &mut Vec<u8>,
146    data: &[u8],
147    dst: &mut Tensor<T>,
148) -> crate::Result<ImageInfo> {
149    if is_jpeg(data) {
150        crate::jpeg::decode_jpeg_into(data, dst, jpeg_state)
151    } else if is_png(data) {
152        crate::png::decode_png_into(data, dst, scratch)
153    } else {
154        Err(CodecError::InvalidData(
155            "unrecognized image format (expected JPEG or PNG magic bytes)".into(),
156        ))
157    }
158}
159
160pub(crate) fn decode_into_inner_dyn(
161    jpeg_state: &mut crate::jpeg::JpegDecoderState,
162    scratch: &mut Vec<u8>,
163    data: &[u8],
164    dst: &mut TensorDyn,
165) -> crate::Result<ImageInfo> {
166    match dst {
167        TensorDyn::U8(t) => decode_into_inner(jpeg_state, scratch, data, t),
168        TensorDyn::I8(t) => decode_into_inner(jpeg_state, scratch, data, t),
169        TensorDyn::U16(t) => decode_into_inner(jpeg_state, scratch, data, t),
170        TensorDyn::I16(t) => decode_into_inner(jpeg_state, scratch, data, t),
171        TensorDyn::F32(t) => decode_into_inner(jpeg_state, scratch, data, t),
172        other => Err(CodecError::UnsupportedDtype(other.dtype())),
173    }
174}
175
176impl Default for ImageDecoder {
177    fn default() -> Self {
178        Self::new()
179    }
180}
181
182/// Parse image headers and return the native dimensions, format, and EXIF
183/// orientation without decoding pixels.
184///
185/// Recommended for one-shot flows that allocate a tensor sized to the image:
186///
187/// ```rust,no_run
188/// use edgefirst_codec::{peek_info, ImageDecoder, ImageLoad};
189/// use edgefirst_tensor::{CpuAccess, Tensor, TensorMemory};
190///
191/// let data = std::fs::read("image.jpg").unwrap();
192/// let info = peek_info(&data).unwrap();
193/// let mut tensor = Tensor::<u8>::image(info.width, info.height, info.format,
194///                                       Some(TensorMemory::Mem),
195///                                       CpuAccess::Write).unwrap();
196/// let mut decoder = ImageDecoder::new();
197/// tensor.load_image(&mut decoder, &data).unwrap();
198/// ```
199pub fn peek_info(data: &[u8]) -> crate::Result<ImageInfo> {
200    if is_jpeg(data) {
201        crate::jpeg::peek_jpeg_info(data)
202    } else if is_png(data) {
203        crate::png::peek_png_info(data)
204    } else {
205        Err(CodecError::InvalidData(
206            "unrecognized image format (expected JPEG or PNG magic bytes)".into(),
207        ))
208    }
209}
210
211/// Check for JPEG magic bytes (FFD8FF).
212fn is_jpeg(data: &[u8]) -> bool {
213    data.len() >= 3 && data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF
214}
215
216/// Check for PNG magic bytes (89504E47).
217fn is_png(data: &[u8]) -> bool {
218    data.len() >= 4 && data[0] == 0x89 && data[1] == 0x50 && data[2] == 0x4E && data[3] == 0x47
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224
225    #[test]
226    fn magic_bytes_jpeg() {
227        assert!(is_jpeg(&[0xFF, 0xD8, 0xFF, 0xE0]));
228        assert!(!is_jpeg(&[0x89, 0x50, 0x4E, 0x47]));
229        assert!(!is_jpeg(&[]));
230        assert!(!is_jpeg(&[0xFF]));
231    }
232
233    #[test]
234    fn magic_bytes_png() {
235        assert!(is_png(&[0x89, 0x50, 0x4E, 0x47]));
236        assert!(!is_png(&[0xFF, 0xD8, 0xFF, 0xE0]));
237        assert!(!is_png(&[]));
238    }
239
240    #[test]
241    fn invalid_data() {
242        let mut decoder = ImageDecoder::new();
243        let mut tensor = Tensor::<u8>::image(
244            100,
245            100,
246            edgefirst_tensor::PixelFormat::Nv12,
247            Some(edgefirst_tensor::TensorMemory::Mem),
248            edgefirst_tensor::CpuAccess::ReadWrite,
249        )
250        .unwrap();
251        let result = decoder.decode_into(b"not an image", &mut tensor);
252        assert!(matches!(result, Err(CodecError::InvalidData(_))));
253    }
254}
edgefirst_codec/decoder.rs

edgefirst_codec/
decoder.rs
