1use std::any::TypeId;
2
3use lerc_band_materialize::{
4 copy_band_values_into_slice, BandLayout as MaterializeLayout, BandMaterializer,
5};
6use lerc_core::{BandLayout, BandSetInfo, BlobInfo, Error, PixelData, Result};
7use ndarray::{ArrayD, IxDyn};
8
9#[derive(Debug, Clone, PartialEq)]
10pub struct Decoded {
11 pub info: BlobInfo,
12 pub pixels: PixelData,
13 pub mask: Option<Vec<u8>>,
14}
15
16impl Decoded {
17 pub fn into_ndarray<T: NdArrayElement>(self) -> Result<ArrayD<T>> {
18 let shape = self.info.ndarray_shape();
19 self.pixels.into_ndarray(&shape)
20 }
21
22 pub fn into_mask_ndarray(self) -> Result<Option<ArrayD<u8>>> {
23 let shape = self.info.mask_ndarray_shape();
24 self.mask
25 .map(|mask| {
26 ArrayD::from_shape_vec(IxDyn(&shape), mask).map_err(|err| {
27 Error::InvalidBlob(format!("failed to build ndarray from decoded mask: {err}"))
28 })
29 })
30 .transpose()
31 }
32}
33
34#[derive(Debug, Clone, PartialEq)]
35pub struct DecodedF64 {
36 pub info: BlobInfo,
37 pub pixels: Vec<f64>,
38 pub mask: Option<Vec<u8>>,
39}
40
41impl DecodedF64 {
42 pub fn into_ndarray(self) -> Result<ArrayD<f64>> {
43 ArrayD::from_shape_vec(IxDyn(&self.info.ndarray_shape()), self.pixels).map_err(|err| {
44 Error::InvalidBlob(format!(
45 "failed to build ndarray from decoded pixels: {err}"
46 ))
47 })
48 }
49
50 pub fn into_mask_ndarray(self) -> Result<Option<ArrayD<u8>>> {
51 let shape = self.info.mask_ndarray_shape();
52 self.mask
53 .map(|mask| {
54 ArrayD::from_shape_vec(IxDyn(&shape), mask).map_err(|err| {
55 Error::InvalidBlob(format!("failed to build ndarray from decoded mask: {err}"))
56 })
57 })
58 .transpose()
59 }
60}
61
62#[derive(Debug, Clone, PartialEq)]
63pub struct DecodedBandSet {
64 pub info: BandSetInfo,
65 pub bands: Vec<PixelData>,
66 pub band_masks: Vec<Option<Vec<u8>>>,
67}
68
69impl DecodedBandSet {
70 pub fn into_ndarray<T: BandElement>(self) -> Result<ArrayD<T>> {
71 self.into_ndarray_with_layout(BandLayout::Interleaved)
72 }
73
74 pub fn into_ndarray_with_layout<T: BandElement>(self, layout: BandLayout) -> Result<ArrayD<T>> {
75 let shape = self.info.ndarray_shape_for_layout(layout);
76 let values = self.into_vec_with_layout(layout)?;
77 ArrayD::from_shape_vec(IxDyn(&shape), values).map_err(|err| {
78 Error::InvalidBlob(format!(
79 "failed to build ndarray from decoded band set: {err}"
80 ))
81 })
82 }
83
84 pub fn into_vec_with_layout<T: BandElement>(self, layout: BandLayout) -> Result<Vec<T>> {
85 if self.bands.len() == 1 {
86 return T::from_pixel_data(self.bands.into_iter().next().unwrap());
87 }
88
89 let mut materializer = BandMaterializer::new(
90 self.info.bands[0].pixel_count()?,
91 self.info.depth() as usize,
92 self.info.band_count(),
93 materialize_layout(layout),
94 )
95 .map_err(materialize_error)?;
96 for (band_index, band) in self.bands.into_iter().enumerate() {
97 copy_pixel_data_into_materializer(&mut materializer, band_index, band)?;
98 }
99 materializer.finish().map_err(materialize_error)
100 }
101
102 pub fn copy_into_slice<T: BandElement>(self, layout: BandLayout, out: &mut [T]) -> Result<()> {
103 let pixel_count = self.info.bands[0].pixel_count()?;
104 let depth = self.info.depth() as usize;
105 let band_count = self.info.band_count();
106 let expected_len = self.info.value_count()?;
107 if out.len() != expected_len {
108 return Err(Error::InvalidBlob(format!(
109 "output slice length {} does not match decoded band set length {}",
110 out.len(),
111 expected_len
112 )));
113 }
114
115 for (band_index, band) in self.bands.into_iter().enumerate() {
116 copy_pixel_data_into_layout_slice(
117 out,
118 band_index,
119 pixel_count,
120 depth,
121 band_count,
122 layout,
123 band,
124 )?;
125 }
126 Ok(())
127 }
128
129 pub fn into_band_mask_ndarray(self) -> Result<Option<ArrayD<u8>>> {
130 into_band_mask_ndarray(self.info, self.band_masks)
131 }
132}
133
134pub fn into_band_mask_ndarray(
135 info: BandSetInfo,
136 band_masks: Vec<Option<Vec<u8>>>,
137) -> Result<Option<ArrayD<u8>>> {
138 if band_masks.iter().all(Option::is_none) {
139 return Ok(None);
140 }
141
142 let pixel_count = info.bands[0].pixel_count()?;
143 let band_count = info.band_count();
144 let shape = info.band_mask_ndarray_shape();
145
146 if band_count == 1 {
147 let mask = band_masks
148 .into_iter()
149 .next()
150 .flatten()
151 .unwrap_or_else(|| vec![1; pixel_count]);
152 return ArrayD::from_shape_vec(IxDyn(&shape), mask)
153 .map(Some)
154 .map_err(|err| {
155 Error::InvalidBlob(format!("failed to build ndarray from decoded mask: {err}"))
156 });
157 }
158
159 let mut merged = Vec::with_capacity(pixel_count * band_count);
160 for pixel in 0..pixel_count {
161 for band_mask in &band_masks {
162 merged.push(band_mask.as_ref().map(|mask| mask[pixel]).unwrap_or(1));
163 }
164 }
165
166 ArrayD::from_shape_vec(IxDyn(&shape), merged)
167 .map(Some)
168 .map_err(|err| {
169 Error::InvalidBlob(format!(
170 "failed to build ndarray from decoded band mask: {err}"
171 ))
172 })
173}
174
175trait SupportedElementValue: Copy + 'static + IntoF64 {
176 const KIND: BandElementKind;
177}
178
179macro_rules! match_pixel_data_values {
180 ($band:expr, |$values:ident| $body:expr) => {
181 match $band {
182 PixelData::I8($values) => $body,
183 PixelData::U8($values) => $body,
184 PixelData::I16($values) => $body,
185 PixelData::U16($values) => $body,
186 PixelData::I32($values) => $body,
187 PixelData::U32($values) => $body,
188 PixelData::F32($values) => $body,
189 PixelData::F64($values) => $body,
190 }
191 };
192}
193
194fn copy_pixel_data_into_materializer<T: BandElement>(
195 materializer: &mut BandMaterializer<T>,
196 band_index: usize,
197 band: PixelData,
198) -> Result<()> {
199 match_pixel_data_values!(band, |values| {
200 copy_typed_values_into_materializer(materializer, band_index, &values)
201 })
202}
203
204fn copy_typed_values_into_materializer<T: BandElement, U: SupportedElementValue>(
205 materializer: &mut BandMaterializer<T>,
206 band_index: usize,
207 values: &[U],
208) -> Result<()> {
209 if T::KIND == U::KIND {
210 let typed = unsafe { cast_slice::<U, T>(values) };
211 return materializer
212 .copy_band(band_index, typed)
213 .map_err(materialize_error);
214 }
215 if T::KIND == BandElementKind::F64 {
216 return materializer
217 .copy_band_with(band_index, |index| {
218 unsafe_cast::<T, f64>(values[index].into_f64())
219 })
220 .map_err(materialize_error);
221 }
222 Err(Error::InvalidBlob(format!(
223 "cannot decode {} pixels into ndarray<{}>",
224 data_type_name::<U>(),
225 std::any::type_name::<T>()
226 .rsplit("::")
227 .next()
228 .unwrap_or("unknown"),
229 )))
230}
231
232fn copy_pixel_data_into_layout_slice<T: BandElement>(
233 out: &mut [T],
234 band_index: usize,
235 pixel_count: usize,
236 depth: usize,
237 band_count: usize,
238 layout: BandLayout,
239 band: PixelData,
240) -> Result<()> {
241 match_pixel_data_values!(band, |values| {
242 copy_typed_values_into_layout_slice(
243 out,
244 band_index,
245 pixel_count,
246 depth,
247 band_count,
248 layout,
249 &values,
250 )
251 })
252}
253
254fn copy_typed_values_into_layout_slice<T: BandElement, U: SupportedElementValue>(
255 out: &mut [T],
256 band_index: usize,
257 pixel_count: usize,
258 depth: usize,
259 band_count: usize,
260 layout: BandLayout,
261 values: &[U],
262) -> Result<()> {
263 if T::KIND == U::KIND {
264 let typed = unsafe { cast_slice::<U, T>(values) };
265 return copy_band_values_into_slice(
266 out,
267 typed,
268 pixel_count,
269 depth,
270 band_index,
271 band_count,
272 materialize_layout(layout),
273 )
274 .map_err(materialize_error);
275 }
276 if T::KIND == BandElementKind::F64 {
277 let band_len = pixel_count
278 .checked_mul(depth.max(1))
279 .ok_or_else(|| Error::InvalidBlob("decoded band length overflows usize".into()))?;
280 if values.len() != band_len {
281 return Err(Error::InvalidBlob(
282 "decoded band length does not match its metadata".into(),
283 ));
284 }
285 for (value_index, value) in values.iter().copied().enumerate() {
286 let out_index = match layout {
287 BandLayout::Interleaved => {
288 if depth <= 1 {
289 value_index * band_count + band_index
290 } else {
291 let pixel = value_index / depth;
292 let sample = value_index % depth;
293 (pixel * band_count + band_index) * depth + sample
294 }
295 }
296 BandLayout::Bsq => band_index * band_len + value_index,
297 };
298 out[out_index] = unsafe_cast::<T, f64>(value.into_f64());
299 }
300 return Ok(());
301 }
302 Err(Error::InvalidBlob(format!(
303 "cannot decode {} pixels into ndarray<{}>",
304 data_type_name::<U>(),
305 std::any::type_name::<T>()
306 .rsplit("::")
307 .next()
308 .unwrap_or("unknown"),
309 )))
310}
311
312unsafe fn cast_slice<U, T>(values: &[U]) -> &[T] {
313 &*(values as *const [U] as *const [T])
314}
315
316fn unsafe_cast<T, U: Copy>(value: U) -> T {
317 unsafe { std::mem::transmute_copy(&value) }
318}
319
320trait IntoF64 {
321 fn into_f64(self) -> f64;
322}
323
324impl IntoF64 for i8 {
325 fn into_f64(self) -> f64 {
326 self as f64
327 }
328}
329impl IntoF64 for u8 {
330 fn into_f64(self) -> f64 {
331 self as f64
332 }
333}
334impl IntoF64 for i16 {
335 fn into_f64(self) -> f64 {
336 self as f64
337 }
338}
339impl IntoF64 for u16 {
340 fn into_f64(self) -> f64 {
341 self as f64
342 }
343}
344impl IntoF64 for i32 {
345 fn into_f64(self) -> f64 {
346 self as f64
347 }
348}
349impl IntoF64 for u32 {
350 fn into_f64(self) -> f64 {
351 self as f64
352 }
353}
354impl IntoF64 for f32 {
355 fn into_f64(self) -> f64 {
356 self as f64
357 }
358}
359impl IntoF64 for f64 {
360 fn into_f64(self) -> f64 {
361 self
362 }
363}
364
365fn data_type_name<T: 'static>() -> &'static str {
366 if TypeId::of::<T>() == TypeId::of::<i8>() {
367 "i8"
368 } else if TypeId::of::<T>() == TypeId::of::<u8>() {
369 "u8"
370 } else if TypeId::of::<T>() == TypeId::of::<i16>() {
371 "i16"
372 } else if TypeId::of::<T>() == TypeId::of::<u16>() {
373 "u16"
374 } else if TypeId::of::<T>() == TypeId::of::<i32>() {
375 "i32"
376 } else if TypeId::of::<T>() == TypeId::of::<u32>() {
377 "u32"
378 } else if TypeId::of::<T>() == TypeId::of::<f32>() {
379 "f32"
380 } else if TypeId::of::<T>() == TypeId::of::<f64>() {
381 "f64"
382 } else {
383 "unknown"
384 }
385}
386
387pub trait NdArrayElement: Sized + Clone {
388 fn from_pixel_data(pixels: PixelData) -> Result<Vec<Self>>;
389}
390
391mod private {
392 pub trait Sealed {}
393
394 impl Sealed for i8 {}
395 impl Sealed for u8 {}
396 impl Sealed for i16 {}
397 impl Sealed for u16 {}
398 impl Sealed for i32 {}
399 impl Sealed for u32 {}
400 impl Sealed for f32 {}
401 impl Sealed for f64 {}
402}
403
404#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
405pub enum BandElementKind {
406 I8,
407 U8,
408 I16,
409 U16,
410 I32,
411 U32,
412 F32,
413 F64,
414}
415
416pub trait BandElement: NdArrayElement + private::Sealed + Copy + Default + 'static {
417 const KIND: BandElementKind;
418}
419
420macro_rules! impl_exact_ndarray_element {
421 ($ty:ty, $variant:ident, $name:literal) => {
422 impl NdArrayElement for $ty {
423 fn from_pixel_data(pixels: PixelData) -> Result<Vec<Self>> {
424 match pixels {
425 PixelData::$variant(values) => Ok(values),
426 other => Err(Error::InvalidBlob(format!(
427 "cannot decode {} pixels into ndarray<{}>",
428 other.data_type().name(),
429 $name
430 ))),
431 }
432 }
433 }
434 };
435}
436
437impl_exact_ndarray_element!(i8, I8, "i8");
438impl_exact_ndarray_element!(u8, U8, "u8");
439impl_exact_ndarray_element!(i16, I16, "i16");
440impl_exact_ndarray_element!(u16, U16, "u16");
441impl_exact_ndarray_element!(i32, I32, "i32");
442impl_exact_ndarray_element!(u32, U32, "u32");
443impl_exact_ndarray_element!(f32, F32, "f32");
444
445impl NdArrayElement for f64 {
446 fn from_pixel_data(pixels: PixelData) -> Result<Vec<Self>> {
447 Ok(pixels.to_f64())
448 }
449}
450
451macro_rules! impl_band_element {
452 ($ty:ty, $kind:ident) => {
453 impl BandElement for $ty {
454 const KIND: BandElementKind = BandElementKind::$kind;
455 }
456
457 impl SupportedElementValue for $ty {
458 const KIND: BandElementKind = BandElementKind::$kind;
459 }
460 };
461}
462
463impl_band_element!(i8, I8);
464impl_band_element!(u8, U8);
465impl_band_element!(i16, I16);
466impl_band_element!(u16, U16);
467impl_band_element!(i32, I32);
468impl_band_element!(u32, U32);
469impl_band_element!(f32, F32);
470impl_band_element!(f64, F64);
471
472fn materialize_layout(layout: BandLayout) -> MaterializeLayout {
473 match layout {
474 BandLayout::Interleaved => MaterializeLayout::Interleaved,
475 BandLayout::Bsq => MaterializeLayout::Bsq,
476 }
477}
478
479fn materialize_error(err: lerc_band_materialize::MaterializeError) -> Error {
480 Error::InvalidBlob(err.to_string())
481}
482
483trait PixelDataExt {
484 fn into_ndarray<T: NdArrayElement>(self, shape: &[usize]) -> Result<ArrayD<T>>;
485}
486
487impl PixelDataExt for PixelData {
488 fn into_ndarray<T: NdArrayElement>(self, shape: &[usize]) -> Result<ArrayD<T>> {
489 ArrayD::from_shape_vec(IxDyn(shape), T::from_pixel_data(self)?).map_err(|err| {
490 Error::InvalidBlob(format!(
491 "failed to build ndarray from decoded pixels: {err}"
492 ))
493 })
494 }
495}