j2k_ml/cpu/
batch.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use core::marker::PhantomData;
4
5use burn_core::tensor::{backend::Backend, DType, Tensor, TensorData};
6use j2k::{
7    BatchDecodeOptions, CpuBatchDecodeResult, CpuBatchSamples, EncodedImage, NativeSampleType,
8    PreparedBatch, PreparedImage,
9};
10
11use crate::batch_contract::{dtype, tensor_shape};
12use crate::{BurnBatchDecode, BurnBatchGroup, BurnBatchTensor, BurnDecodeError};
13
14/// Persistent CPU codec session followed by one Burn materialization per output group.
15#[derive(Debug)]
16pub struct CpuBurnDecoder<B: Backend> {
17    codec: j2k::CpuBatchDecoder,
18    device: B::Device,
19    backend: PhantomData<B>,
20}
21
22impl<B: Backend> CpuBurnDecoder<B> {
23    /// Construct a CPU codec session for one Burn device.
24    #[must_use]
25    pub fn new(device: B::Device, options: BatchDecodeOptions) -> Self {
26        Self {
27            codec: j2k::CpuBatchDecoder::new(options),
28            device,
29            backend: PhantomData,
30        }
31    }
32
33    /// Borrow the retained codec session.
34    #[must_use]
35    pub const fn codec(&self) -> &j2k::CpuBatchDecoder {
36        &self.codec
37    }
38
39    /// Borrow the Burn device receiving decoded tensors.
40    #[must_use]
41    pub const fn device(&self) -> &B::Device {
42        &self.device
43    }
44
45    /// Parse and group owned inputs once for repeated training epochs.
46    pub fn prepare(&self, inputs: Vec<EncodedImage>) -> Result<PreparedBatch, BurnDecodeError> {
47        Ok(self.codec.prepare(inputs)?)
48    }
49
50    /// Regroup caller-supplied prepared images without reparsing encoded bytes.
51    ///
52    /// Source indices in the returned batch are positions in `images`; each
53    /// [`PreparedImage::source_index`] remains its original preparation index.
54    pub fn prepare_prepared_images(
55        &self,
56        images: Vec<PreparedImage>,
57    ) -> Result<PreparedBatch, BurnDecodeError> {
58        Ok(self.codec.prepare_prepared_images(images)?)
59    }
60
61    /// Prepare, decode, and materialize one owned batch.
62    pub fn decode(
63        &mut self,
64        inputs: Vec<EncodedImage>,
65    ) -> Result<BurnBatchDecode<B>, BurnDecodeError> {
66        let prepared = self.prepare(inputs)?;
67        self.decode_prepared(&prepared)
68    }
69
70    /// Regroup, decode, and materialize caller-supplied prepared images.
71    pub fn decode_prepared_images(
72        &mut self,
73        images: Vec<PreparedImage>,
74    ) -> Result<BurnBatchDecode<B>, BurnDecodeError> {
75        let prepared = self.prepare_prepared_images(images)?;
76        self.decode_prepared(&prepared)
77    }
78
79    /// Decode a reusable prepared batch into ordinary rank-4 Burn integer tensors.
80    pub fn decode_prepared(
81        &mut self,
82        prepared: &PreparedBatch,
83    ) -> Result<BurnBatchDecode<B>, BurnDecodeError> {
84        self.ensure_prepared_dtypes(prepared)?;
85        let decoded = self.codec.decode_prepared(prepared)?;
86        materialize::<B>(decoded, &self.device)
87    }
88
89    fn ensure_prepared_dtypes(&self, prepared: &PreparedBatch) -> Result<(), BurnDecodeError> {
90        for group in prepared.groups() {
91            let dtype = dtype(group.info().sample_type)?;
92            if !B::supports_dtype(&self.device, dtype) {
93                return Err(BurnDecodeError::UnsupportedDType { dtype });
94            }
95        }
96        Ok(())
97    }
98}
99
100fn materialize<B: Backend>(
101    decoded: CpuBatchDecodeResult,
102    device: &B::Device,
103) -> Result<BurnBatchDecode<B>, BurnDecodeError> {
104    let (codec_groups, errors) = decoded.into_parts();
105    let mut groups = Vec::new();
106    groups
107        .try_reserve_exact(codec_groups.len())
108        .map_err(|_| BurnDecodeError::SizeOverflow)?;
109
110    for group in codec_groups {
111        let (info, source_indices, decoded_rects, warnings, samples) = group.into_parts();
112        let shape = tensor_shape(source_indices.len(), &info)?;
113        let tensor = match (info.sample_type, samples) {
114            (NativeSampleType::U8, CpuBatchSamples::U8(samples)) => BurnBatchTensor::U8(
115                Tensor::from_data(TensorData::new(samples, shape), (device, DType::U8)),
116            ),
117            (NativeSampleType::U16, CpuBatchSamples::U16(samples)) => BurnBatchTensor::U16(
118                Tensor::from_data(TensorData::new(samples, shape), (device, DType::U16)),
119            ),
120            (NativeSampleType::I16, CpuBatchSamples::I16(samples)) => BurnBatchTensor::I16(
121                Tensor::from_data(TensorData::new(samples, shape), (device, DType::I16)),
122            ),
123            _ => return Err(BurnDecodeError::SampleTypeMismatch),
124        };
125        groups.push(BurnBatchGroup {
126            tensor,
127            info,
128            source_indices,
129            decoded_rects,
130            warnings,
131        });
132    }
133
134    Ok(BurnBatchDecode {
135        groups,
136        errors,
137        group_errors: Vec::new(),
138    })
139}
j2k_ml/cpu/batch.rs

j2k_ml/cpu/
batch.rs
