1#![allow(clippy::multiple_crate_versions)] #[cfg(test)]
23use ::serde_json as _;
24
25use std::borrow::Cow;
26use std::fmt;
27
28use ndarray::{Array, Array1, ArrayBase, Axis, Data, Dimension, IxDyn, ShapeError};
29use num_traits::identities::Zero;
30use numcodecs::{
31 AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
32 Codec, StaticCodec, StaticCodecConfig,
33};
34use schemars::JsonSchema;
35use serde::{Deserialize, Serialize};
36use thiserror::Error;
37
38mod ffi;
39
40#[derive(Clone, Serialize, Deserialize, JsonSchema)]
41#[schemars(deny_unknown_fields)]
43pub struct Jpeg2000Codec {
50 #[serde(flatten)]
52 pub mode: Jpeg2000CompressionMode,
53}
54
55#[derive(Clone, Serialize, Deserialize, JsonSchema)]
56#[serde(tag = "mode")]
58pub enum Jpeg2000CompressionMode {
59 #[serde(rename = "psnr")]
61 PSNR {
62 psnr: f32,
64 },
65 #[serde(rename = "rate")]
67 Rate {
68 rate: f32,
70 },
71 #[serde(rename = "lossless")]
73 Lossless,
74}
75
76impl Codec for Jpeg2000Codec {
77 type Error = Jpeg2000CodecError;
78
79 fn encode(&self, data: AnyCowArray) -> Result<AnyArray, Self::Error> {
80 match data {
81 AnyCowArray::I8(data) => Ok(AnyArray::U8(
82 Array1::from(compress(data, &self.mode)?).into_dyn(),
83 )),
84 AnyCowArray::U8(data) => Ok(AnyArray::U8(
85 Array1::from(compress(data, &self.mode)?).into_dyn(),
86 )),
87 AnyCowArray::I16(data) => Ok(AnyArray::U8(
88 Array1::from(compress(data, &self.mode)?).into_dyn(),
89 )),
90 AnyCowArray::U16(data) => Ok(AnyArray::U8(
91 Array1::from(compress(data, &self.mode)?).into_dyn(),
92 )),
93 AnyCowArray::I32(data) => Ok(AnyArray::U8(
94 Array1::from(compress(data, &self.mode)?).into_dyn(),
95 )),
96 AnyCowArray::U32(data) => Ok(AnyArray::U8(
97 Array1::from(compress(data, &self.mode)?).into_dyn(),
98 )),
99 AnyCowArray::I64(data) => Ok(AnyArray::U8(
100 Array1::from(compress(data, &self.mode)?).into_dyn(),
101 )),
102 AnyCowArray::U64(data) => Ok(AnyArray::U8(
103 Array1::from(compress(data, &self.mode)?).into_dyn(),
104 )),
105 encoded => Err(Jpeg2000CodecError::UnsupportedDtype(encoded.dtype())),
106 }
107 }
108
109 fn decode(&self, encoded: AnyCowArray) -> Result<AnyArray, Self::Error> {
110 let AnyCowArray::U8(encoded) = encoded else {
111 return Err(Jpeg2000CodecError::EncodedDataNotBytes {
112 dtype: encoded.dtype(),
113 });
114 };
115
116 if !matches!(encoded.shape(), [_]) {
117 return Err(Jpeg2000CodecError::EncodedDataNotOneDimensional {
118 shape: encoded.shape().to_vec(),
119 });
120 }
121
122 decompress(&AnyCowArray::U8(encoded).as_bytes())
123 }
124
125 fn decode_into(
126 &self,
127 encoded: AnyArrayView,
128 mut decoded: AnyArrayViewMut,
129 ) -> Result<(), Self::Error> {
130 let decoded_in = self.decode(encoded.cow())?;
131
132 Ok(decoded.assign(&decoded_in)?)
133 }
134}
135
136impl StaticCodec for Jpeg2000Codec {
137 const CODEC_ID: &'static str = "jpeg2000";
138
139 type Config<'de> = Self;
140
141 fn from_config(config: Self::Config<'_>) -> Self {
142 config
143 }
144
145 fn get_config(&self) -> StaticCodecConfig<Self> {
146 StaticCodecConfig::from(self)
147 }
148}
149
150#[derive(Debug, Error)]
151pub enum Jpeg2000CodecError {
153 #[error("Jpeg2000 does not support the dtype {0}")]
155 UnsupportedDtype(AnyArrayDType),
156 #[error("Jpeg2000 failed to encode the header")]
158 HeaderEncodeFailed {
159 source: Jpeg2000HeaderError,
161 },
162 #[error("Jpeg2000 failed to encode the data")]
164 Jpeg2000EncodeFailed {
165 source: Jpeg2000CodingError,
167 },
168 #[error("Jpeg2000 failed to encode a slice")]
170 SliceEncodeFailed {
171 source: Jpeg2000SliceError,
173 },
174 #[error(
177 "Jpeg2000 can only decode one-dimensional byte arrays but received an array of dtype {dtype}"
178 )]
179 EncodedDataNotBytes {
180 dtype: AnyArrayDType,
182 },
183 #[error("Jpeg2000 can only decode one-dimensional byte arrays but received a byte array of shape {shape:?}")]
186 EncodedDataNotOneDimensional {
187 shape: Vec<usize>,
189 },
190 #[error("Jpeg2000 failed to decode the header")]
192 HeaderDecodeFailed {
193 source: Jpeg2000HeaderError,
195 },
196 #[error("Jpeg2000 failed to decode a slice")]
198 SliceDecodeFailed {
199 source: Jpeg2000SliceError,
201 },
202 #[error("Jpeg2000 failed to decode from an excessive number of slices")]
204 DecodeTooManySlices,
205 #[error("Jpeg2000 failed to decode the data")]
207 Jpeg2000DecodeFailed {
208 source: Jpeg2000CodingError,
210 },
211 #[error("Jpeg2000 decoded into an invalid shape not matching the data size")]
213 DecodeInvalidShape {
214 source: ShapeError,
216 },
217 #[error("Jpeg2000Codec cannot decode into the provided array")]
219 MismatchedDecodeIntoArray {
220 #[from]
222 source: AnyArrayAssignError,
223 },
224}
225
226#[derive(Debug, Error)]
227#[error(transparent)]
228pub struct Jpeg2000HeaderError(postcard::Error);
230
231#[derive(Debug, Error)]
232#[error(transparent)]
233pub struct Jpeg2000SliceError(postcard::Error);
235
236#[derive(Debug, Error)]
237#[error(transparent)]
238pub struct Jpeg2000CodingError(ffi::Jpeg2000Error);
240
241pub fn compress<T: Jpeg2000Element, S: Data<Elem = T>, D: Dimension>(
251 data: ArrayBase<S, D>,
252 mode: &Jpeg2000CompressionMode,
253) -> Result<Vec<u8>, Jpeg2000CodecError> {
254 let mut encoded = postcard::to_extend(
255 &CompressionHeader {
256 dtype: T::DTYPE,
257 shape: Cow::Borrowed(data.shape()),
258 },
259 Vec::new(),
260 )
261 .map_err(|err| Jpeg2000CodecError::HeaderEncodeFailed {
262 source: Jpeg2000HeaderError(err),
263 })?;
264
265 if data.is_empty() {
267 return Ok(encoded);
268 }
269
270 let mut encoded_slice = Vec::new();
271
272 let mut chunk_size = Vec::from(data.shape());
273 let (width, height) = match *chunk_size.as_mut_slice() {
274 [ref mut rest @ .., height, width] => {
275 for r in rest {
276 *r = 1;
277 }
278 (width, height)
279 }
280 [width] => (width, 1),
281 [] => (1, 1),
282 };
283
284 for slice in data.into_dyn().exact_chunks(chunk_size.as_slice()) {
285 encoded_slice.clear();
286
287 ffi::encode_into(
288 slice.iter().copied(),
289 width,
290 height,
291 match mode {
292 Jpeg2000CompressionMode::PSNR { psnr } => ffi::Jpeg2000CompressionMode::PSNR(*psnr),
293 Jpeg2000CompressionMode::Rate { rate } => ffi::Jpeg2000CompressionMode::Rate(*rate),
294 Jpeg2000CompressionMode::Lossless => ffi::Jpeg2000CompressionMode::Lossless,
295 },
296 &mut encoded_slice,
297 )
298 .map_err(|err| Jpeg2000CodecError::Jpeg2000EncodeFailed {
299 source: Jpeg2000CodingError(err),
300 })?;
301
302 encoded = postcard::to_extend(encoded_slice.as_slice(), encoded).map_err(|err| {
303 Jpeg2000CodecError::SliceEncodeFailed {
304 source: Jpeg2000SliceError(err),
305 }
306 })?;
307 }
308
309 Ok(encoded)
310}
311
312pub fn decompress(encoded: &[u8]) -> Result<AnyArray, Jpeg2000CodecError> {
326 fn decompress_typed<T: Jpeg2000Element>(
327 mut encoded: &[u8],
328 shape: &[usize],
329 ) -> Result<Array<T, IxDyn>, Jpeg2000CodecError> {
330 let mut decoded = Array::<T, _>::zeros(shape);
331
332 let mut chunk_size = Vec::from(shape);
333 let (width, height) = match *chunk_size.as_mut_slice() {
334 [ref mut rest @ .., height, width] => {
335 for r in rest {
336 *r = 1;
337 }
338 (width, height)
339 }
340 [width] => (width, 1),
341 [] => (1, 1),
342 };
343
344 for mut slice in decoded.exact_chunks_mut(chunk_size.as_slice()) {
345 let (encoded_slice, rest) =
346 postcard::take_from_bytes::<Cow<[u8]>>(encoded).map_err(|err| {
347 Jpeg2000CodecError::SliceDecodeFailed {
348 source: Jpeg2000SliceError(err),
349 }
350 })?;
351 encoded = rest;
352
353 let (decoded_slice, (_width, _height)) =
354 ffi::decode::<T>(&encoded_slice).map_err(|err| {
355 Jpeg2000CodecError::Jpeg2000DecodeFailed {
356 source: Jpeg2000CodingError(err),
357 }
358 })?;
359 let mut decoded_slice = Array::from_shape_vec((height, width), decoded_slice)
360 .map_err(|source| Jpeg2000CodecError::DecodeInvalidShape { source })?
361 .into_dyn();
362
363 while decoded_slice.ndim() > shape.len() {
364 decoded_slice = decoded_slice.remove_axis(Axis(0));
365 }
366
367 slice.assign(&decoded_slice);
368 }
369
370 if !encoded.is_empty() {
371 return Err(Jpeg2000CodecError::DecodeTooManySlices);
372 }
373
374 Ok(decoded)
375 }
376
377 let (header, encoded) =
378 postcard::take_from_bytes::<CompressionHeader>(encoded).map_err(|err| {
379 Jpeg2000CodecError::HeaderDecodeFailed {
380 source: Jpeg2000HeaderError(err),
381 }
382 })?;
383
384 if header.shape.iter().copied().product::<usize>() == 0 {
386 return match header.dtype {
387 Jpeg2000DType::I8 => Ok(AnyArray::I8(Array::zeros(&*header.shape))),
388 Jpeg2000DType::U8 => Ok(AnyArray::U8(Array::zeros(&*header.shape))),
389 Jpeg2000DType::I16 => Ok(AnyArray::I16(Array::zeros(&*header.shape))),
390 Jpeg2000DType::U16 => Ok(AnyArray::U16(Array::zeros(&*header.shape))),
391 Jpeg2000DType::I32 => Ok(AnyArray::I32(Array::zeros(&*header.shape))),
392 Jpeg2000DType::U32 => Ok(AnyArray::U32(Array::zeros(&*header.shape))),
393 Jpeg2000DType::I64 => Ok(AnyArray::I64(Array::zeros(&*header.shape))),
394 Jpeg2000DType::U64 => Ok(AnyArray::U64(Array::zeros(&*header.shape))),
395 };
396 }
397
398 match header.dtype {
399 Jpeg2000DType::I8 => Ok(AnyArray::I8(decompress_typed(encoded, &header.shape)?)),
400 Jpeg2000DType::U8 => Ok(AnyArray::U8(decompress_typed(encoded, &header.shape)?)),
401 Jpeg2000DType::I16 => Ok(AnyArray::I16(decompress_typed(encoded, &header.shape)?)),
402 Jpeg2000DType::U16 => Ok(AnyArray::U16(decompress_typed(encoded, &header.shape)?)),
403 Jpeg2000DType::I32 => Ok(AnyArray::I32(decompress_typed(encoded, &header.shape)?)),
404 Jpeg2000DType::U32 => Ok(AnyArray::U32(decompress_typed(encoded, &header.shape)?)),
405 Jpeg2000DType::I64 => Ok(AnyArray::I64(decompress_typed(encoded, &header.shape)?)),
406 Jpeg2000DType::U64 => Ok(AnyArray::U64(decompress_typed(encoded, &header.shape)?)),
407 }
408}
409
410pub trait Jpeg2000Element: ffi::Jpeg2000Element + Zero {
412 const DTYPE: Jpeg2000DType;
414}
415
416impl Jpeg2000Element for i8 {
417 const DTYPE: Jpeg2000DType = Jpeg2000DType::I8;
418}
419impl Jpeg2000Element for u8 {
420 const DTYPE: Jpeg2000DType = Jpeg2000DType::U8;
421}
422impl Jpeg2000Element for i16 {
423 const DTYPE: Jpeg2000DType = Jpeg2000DType::I16;
424}
425impl Jpeg2000Element for u16 {
426 const DTYPE: Jpeg2000DType = Jpeg2000DType::U16;
427}
428impl Jpeg2000Element for i32 {
429 const DTYPE: Jpeg2000DType = Jpeg2000DType::I32;
430}
431impl Jpeg2000Element for u32 {
432 const DTYPE: Jpeg2000DType = Jpeg2000DType::U32;
433}
434impl Jpeg2000Element for i64 {
435 const DTYPE: Jpeg2000DType = Jpeg2000DType::I64;
436}
437impl Jpeg2000Element for u64 {
438 const DTYPE: Jpeg2000DType = Jpeg2000DType::U64;
439}
440
441#[derive(Serialize, Deserialize)]
442struct CompressionHeader<'a> {
443 dtype: Jpeg2000DType,
444 #[serde(borrow)]
445 shape: Cow<'a, [usize]>,
446}
447
448#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
450#[expect(missing_docs)]
451pub enum Jpeg2000DType {
452 #[serde(rename = "i8", alias = "int8")]
453 I8,
454 #[serde(rename = "u8", alias = "uint8")]
455 U8,
456 #[serde(rename = "i16", alias = "int16")]
457 I16,
458 #[serde(rename = "u16", alias = "uint16")]
459 U16,
460 #[serde(rename = "i32", alias = "int32")]
461 I32,
462 #[serde(rename = "u32", alias = "uint32")]
463 U32,
464 #[serde(rename = "i64", alias = "int64")]
465 I64,
466 #[serde(rename = "u64", alias = "uint64")]
467 U64,
468}
469
470impl fmt::Display for Jpeg2000DType {
471 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
472 fmt.write_str(match self {
473 Self::I8 => "i8",
474 Self::U8 => "u8",
475 Self::I16 => "i16",
476 Self::U16 => "u16",
477 Self::I32 => "i32",
478 Self::U32 => "u32",
479 Self::I64 => "i64",
480 Self::U64 => "u64",
481 })
482 }
483}
484
485#[cfg(test)]
486#[allow(clippy::unwrap_used)]
487mod tests {
488 use ndarray::{Ix0, Ix1, Ix2, Ix3, Ix4};
489
490 use super::*;
491
492 #[test]
493 fn zero_length() {
494 std::mem::drop(simple_logger::init());
495
496 let encoded = compress(
497 Array::<i16, _>::from_shape_vec([3, 0], vec![]).unwrap(),
498 &Jpeg2000CompressionMode::PSNR { psnr: 42.0 },
499 )
500 .unwrap();
501 let decoded = decompress(&encoded).unwrap();
502
503 assert_eq!(decoded.dtype(), AnyArrayDType::I16);
504 assert!(decoded.is_empty());
505 assert_eq!(decoded.shape(), &[3, 0]);
506 }
507
508 #[test]
509 fn small_2d() {
510 std::mem::drop(simple_logger::init());
511
512 let encoded = compress(
513 Array::<i16, _>::from_shape_vec([1, 1], vec![42]).unwrap(),
514 &Jpeg2000CompressionMode::PSNR { psnr: 42.0 },
515 )
516 .unwrap();
517 let decoded = decompress(&encoded).unwrap();
518
519 assert_eq!(decoded.dtype(), AnyArrayDType::I16);
520 assert_eq!(decoded.len(), 1);
521 assert_eq!(decoded.shape(), &[1, 1]);
522 }
523
524 #[test]
525 fn small_lossless_types() {
526 macro_rules! check {
527 ($T:ident($t:ident)) => {
528 check! { $T($t,$t::MIN,$t::MAX) }
529 };
530 ($T:ident($t:ident,$min:expr,$max:expr)) => {
531 let data = Array::<$t, _>::from_shape_vec([4, 1], vec![$min, 0, 42, $max]).unwrap();
532
533 let encoded = compress(
534 data.view(),
535 &Jpeg2000CompressionMode::Lossless,
536 )
537 .unwrap();
538 let decoded = decompress(&encoded).unwrap();
539
540 assert_eq!(decoded.len(), 4);
541 assert_eq!(decoded.shape(), &[4, 1]);
542 assert_eq!(decoded, AnyArray::$T(data.into_dyn()));
543 };
544 ($($T:ident($($tt:tt),*)),*) => {
545 $(check! { $T($($tt),*) })*
546 };
547 }
548
549 check! {
550 I8(i8), U8(u8), I16(i16), U16(u16),
551 I32(i32,(i32::MIN/(1<<7)),(i32::MAX/(1<<7))),
552 U32(u32,(u32::MIN),(u32::MAX/(1<<7))),
553 I64(i64,(i64::MIN/(1<<(32+7))),(i64::MAX/(1<<(32+7)))),
554 U64(u64,(u64::MIN),(u64::MAX/(1<<(32+7))))
555 }
556 }
557
558 #[test]
559 fn out_of_range() {
560 macro_rules! check {
561 ($T:ident($t:ident,$($v:expr),*)) => {
562 $(
563 let data = Array::<$t, _>::from_shape_vec([1, 1], vec![$v]).unwrap();
564 compress(
565 data.view(),
566 &Jpeg2000CompressionMode::Lossless,
567 )
568 .unwrap_err();
569 )*
570 };
571 ($($T:ident($($tt:tt),*)),*) => {
572 $(check! { $T($($tt),*) })*
573 };
574 }
575
576 check! {
577 I32(i32,(i32::MIN),(i32::MAX)), U32(u32,(u32::MAX)),
578 I64(i64,(i64::MIN),(i64::MAX)), U64(u64,(u64::MAX))
579 }
580 }
581
582 #[test]
583 fn large_2d() {
584 std::mem::drop(simple_logger::init());
585
586 let encoded = compress(
587 Array::<i16, _>::zeros((64, 64)),
588 &Jpeg2000CompressionMode::PSNR { psnr: 42.0 },
589 )
590 .unwrap();
591 let decoded = decompress(&encoded).unwrap();
592
593 assert_eq!(decoded.dtype(), AnyArrayDType::I16);
594 assert_eq!(decoded.len(), 64 * 64);
595 assert_eq!(decoded.shape(), &[64, 64]);
596 }
597
598 #[test]
599 fn all_modes() {
600 std::mem::drop(simple_logger::init());
601
602 for mode in [
603 Jpeg2000CompressionMode::PSNR { psnr: 42.0 },
604 Jpeg2000CompressionMode::Rate { rate: 5.0 },
605 Jpeg2000CompressionMode::Lossless,
606 ] {
607 let encoded = compress(Array::<i16, _>::zeros((64, 64)), &mode).unwrap();
608 let decoded = decompress(&encoded).unwrap();
609
610 assert_eq!(decoded.dtype(), AnyArrayDType::I16);
611 assert_eq!(decoded.len(), 64 * 64);
612 assert_eq!(decoded.shape(), &[64, 64]);
613 }
614 }
615
616 #[test]
617 fn many_dimensions() {
618 std::mem::drop(simple_logger::init());
619
620 for data in [
621 Array::<i16, Ix0>::from_shape_vec([], vec![42])
622 .unwrap()
623 .into_dyn(),
624 Array::<i16, Ix1>::from_shape_vec([2], vec![1, 2])
625 .unwrap()
626 .into_dyn(),
627 Array::<i16, Ix2>::from_shape_vec([2, 2], vec![1, 2, 3, 4])
628 .unwrap()
629 .into_dyn(),
630 Array::<i16, Ix3>::from_shape_vec([2, 2, 2], vec![1, 2, 3, 4, 5, 6, 7, 8])
631 .unwrap()
632 .into_dyn(),
633 Array::<i16, Ix4>::from_shape_vec(
634 [2, 2, 2, 2],
635 vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
636 )
637 .unwrap()
638 .into_dyn(),
639 ] {
640 let encoded = compress(data.view(), &Jpeg2000CompressionMode::Lossless).unwrap();
641 let decoded = decompress(&encoded).unwrap();
642
643 assert_eq!(decoded, AnyArray::I16(data));
644 }
645 }
646}