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