From 12d5d3355be4b8d1a36e4fedbb8d701c7d3b50d6 Mon Sep 17 00:00:00 2001 From: Wim Pomp Date: Thu, 27 Mar 2025 15:32:31 +0100 Subject: [PATCH] - add deflate compression --- Cargo.toml | 13 +-- py/tiffwrite/__init__.py | 14 ++-- src/lib.rs | 173 +++++++++++++++++++++++++++++---------- src/py.rs | 18 +++- 4 files changed, 161 insertions(+), 57 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 3fa300f..08293d1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "tiffwrite" -version = "2025.2.0" +version = "2025.3.0" edition = "2021" authors = ["Wim Pomp "] license = "MIT" @@ -16,16 +16,17 @@ name = "tiffwrite" crate-type = ["cdylib", "rlib"] [dependencies] -anyhow = "1.0.95" -chrono = "0.4.39" +anyhow = "1.0.97" +chrono = "0.4.40" +flate2 = "1.1.0" ndarray = "0.16.1" num = "0.4.3" rayon = "1.10.0" -zstd = "0.13.2" -numpy = { version = "0.23.0", optional = true } +zstd = "0.13.3" +numpy = { version = "0.24.0", optional = true } [dependencies.pyo3] -version = "0.23.4" +version = "0.24.0" features = ["extension-module", "abi3-py310", "generate-import-lib", "anyhow", "multiple-pymethods"] optional = true diff --git a/py/tiffwrite/__init__.py b/py/tiffwrite/__init__.py index 0becace..b1b51ef 100644 --- a/py/tiffwrite/__init__.py +++ b/py/tiffwrite/__init__.py @@ -48,7 +48,7 @@ class IJTiffFile(rs.IJTiffFile): pxsize: pixel size in um deltaz: z slice interval in um timeinterval: time between frames in seconds - compression: zstd compression level: -7 to 22. + compression: ('zstd', level) for zstd with compression level: -7 to 22, 'deflate' for deflate compresion comment: comment to be saved in tif extratags: other tags to be saved, example: (Tag.ascii(315, 'John Doe'), Tag.bytes(4567, [400, 500]) or (Tag.ascii(33432, 'Made by me'),). @@ -58,14 +58,18 @@ class IJTiffFile(rs.IJTiffFile): def __init__(self, path: str | Path, *, dtype: DTypeLike = 'uint16', colors: Sequence[str] = None, colormap: str = None, pxsize: float = None, - deltaz: float = None, timeinterval: float = None, compression: int = None, comment: str = None, + deltaz: float = None, timeinterval: float = None, + compression: int | str | tuple[int, int] | tuple[str, int] = None, comment: str = None, extratags: Sequence[Tag] = None) -> None: + codecs = {'z': 50000, 'd': 8, 8: 8, 50000: 50000} self.path = Path(path) self.dtype = np.dtype(dtype) if compression is not None: - if isinstance(compression, Sequence): - compression = compression[-1] - self.set_compression_level(compression) + if isinstance(compression, tuple): + compression = codecs.get(compression[0], 50000), (int(compression[1]) if len(compression) == 2 else 22) + else: + compression = codecs.get(compression, 50000), 22 + self.set_compression(*compression) if colors is not None: self.colors = np.array([get_color(color) for color in colors]) if colormap is not None: diff --git a/src/lib.rs b/src/lib.rs index 7abd9c2..e1a955f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ mod py; use anyhow::Result; use chrono::Utc; +use flate2::write::ZlibEncoder; use ndarray::{s, ArcArray2, AsArray, Ix2}; use num::{traits::ToBytes, Complex, FromPrimitive, Rational32}; use rayon::prelude::*; @@ -16,12 +17,28 @@ use std::{ thread, thread::{sleep, JoinHandle}, }; +use zstd::zstd_safe::CompressionLevel; use zstd::{stream::Encoder, DEFAULT_COMPRESSION_LEVEL}; const TAG_SIZE: usize = 20; const OFFSET_SIZE: usize = 8; const OFFSET: u64 = 16; -const COMPRESSION: u16 = 50000; + +/// Compression: deflate or zstd +#[derive(Clone, Debug)] +pub enum Compression { + Deflate, + Zstd(CompressionLevel), +} + +impl Compression { + fn index(&self) -> u16 { + match self { + Compression::Deflate => 8, + Compression::Zstd(_) => 50000, + } + } +} /// Image File Directory #[allow(clippy::upper_case_acronyms)] @@ -345,7 +362,7 @@ struct CompressedFrame { } impl CompressedFrame { - fn new(frame: ArcArray2, compression_level: i32) -> CompressedFrame + fn new(frame: ArcArray2, compression: Compression) -> CompressedFrame where T: Bytes + Send + Sync, { @@ -391,34 +408,68 @@ impl CompressedFrame { } } - let bytes: Vec<_> = if slices.len() > 4 { - slices - .into_par_iter() - .map(|slice| { - CompressedFrame::compress_tile( - frame.clone(), - slice, - tile_size, - tile_size, - compression_level, - ) - .unwrap() - }) - .collect() - } else { - slices - .into_iter() - .map(|slice| { - CompressedFrame::compress_tile( - frame.clone(), - slice, - tile_size, - tile_size, - compression_level, - ) - .unwrap() - }) - .collect() + let bytes: Vec<_> = match compression { + Compression::Deflate => { + if slices.len() > 4 { + slices + .into_par_iter() + .map(|slice| { + CompressedFrame::compress_tile_deflate( + frame.clone(), + slice, + tile_size, + tile_size, + ) + .unwrap() + }) + .collect() + } else { + slices + .into_iter() + .map(|slice| { + CompressedFrame::compress_tile_deflate( + frame.clone(), + slice, + tile_size, + tile_size, + ) + .unwrap() + }) + .collect() + } + } + + Compression::Zstd(level) => { + if slices.len() > 4 { + slices + .into_par_iter() + .map(|slice| { + CompressedFrame::compress_tile_zstd( + frame.clone(), + slice, + tile_size, + tile_size, + level, + ) + .unwrap() + }) + .collect() + } else { + slices + .into_iter() + .map(|slice| { + CompressedFrame::compress_tile_zstd( + frame.clone(), + slice, + tile_size, + tile_size, + level, + ) + .unwrap() + }) + .collect() + } + } }; CompressedFrame { @@ -432,22 +483,18 @@ impl CompressedFrame { } } - fn compress_tile( + fn encode( + mut encoder: W, frame: ArcArray2, slice: (usize, usize, usize, usize), tile_width: usize, tile_length: usize, - compression_level: i32, - ) -> Result> + ) -> Result where + W: Write, T: Bytes, { - let mut dest = Vec::new(); - let mut encoder = Encoder::new(&mut dest, compression_level)?; let bytes_per_sample = (T::BITS_PER_SAMPLE / 8) as usize; - encoder.include_contentsize(true)?; - encoder.set_pledged_src_size(Some((bytes_per_sample * tile_width * tile_length) as u64))?; - encoder.include_checksum(true)?; let shape = (slice.1 - slice.0, slice.3 - slice.2); for i in 0..shape.0 { encoder.write_all( @@ -465,6 +512,40 @@ impl CompressedFrame { 0; bytes_per_sample * tile_width * (tile_length - shape.0) ])?; + Ok(encoder) + } + + fn compress_tile_deflate( + frame: ArcArray2, + slice: (usize, usize, usize, usize), + tile_width: usize, + tile_length: usize, + ) -> Result> + where + T: Bytes, + { + let mut encoder = ZlibEncoder::new(Vec::new(), flate2::Compression::default()); + encoder = CompressedFrame::encode(encoder, frame, slice, tile_width, tile_length)?; + Ok(encoder.finish()?) + } + + fn compress_tile_zstd( + frame: ArcArray2, + slice: (usize, usize, usize, usize), + tile_width: usize, + tile_length: usize, + compression_level: i32, + ) -> Result> + where + T: Bytes, + { + let mut dest = Vec::new(); + let mut encoder = Encoder::new(&mut dest, compression_level)?; + let bytes_per_sample = (T::BITS_PER_SAMPLE / 8) as usize; + encoder.include_contentsize(true)?; + encoder.set_pledged_src_size(Some((bytes_per_sample * tile_width * tile_length) as u64))?; + encoder.include_checksum(true)?; + encoder = CompressedFrame::encode(encoder, frame, slice, tile_width, tile_length)?; encoder.finish()?; Ok(dest) } @@ -569,7 +650,7 @@ pub struct IJTiffFile { hashes: HashMap, threads: HashMap<(usize, usize, usize), JoinHandle>, /// zstd: -7 ..= 22 - pub compression_level: i32, + pub compression: Compression, pub colors: Colors, pub comment: Option, /// um per pixel @@ -610,7 +691,7 @@ impl IJTiffFile { frames: HashMap::new(), hashes: HashMap::new(), threads: HashMap::new(), - compression_level: DEFAULT_COMPRESSION_LEVEL, + compression: Compression::Zstd(DEFAULT_COMPRESSION_LEVEL), colors: Colors::None, comment: None, px_size: None, @@ -620,6 +701,11 @@ impl IJTiffFile { }) } + /// set compression: zstd(level) or deflate + pub fn set_compression(&mut self, compression: Compression) { + self.compression = compression; + } + /// to be saved in description tag (270) pub fn description(&self, c_size: usize, z_size: usize, t_size: usize) -> String { let mut desc: String = String::from("ImageJ=1.11a"); @@ -727,11 +813,11 @@ impl IJTiffFile { } sleep(Duration::from_millis(100)); } - let compression_level = self.compression_level; + let compression = self.compression.clone(); let frame = frame.into().to_shared(); self.threads.insert( (c, z, t), - thread::spawn(move || CompressedFrame::new(frame, compression_level)), + thread::spawn(move || CompressedFrame::new(frame, compression)), ); Ok(()) } @@ -833,7 +919,8 @@ impl IJTiffFile { ifd.tags.insert(Tag::long(257, &[frame.image_length])); ifd.tags .insert(Tag::short(258, &vec![frame.bits_per_sample; frame_count])); - ifd.tags.insert(Tag::short(259, &[COMPRESSION])); + ifd.tags + .insert(Tag::short(259, &[self.compression.index()])); ifd.tags .insert(Tag::ascii(270, &self.description(c_size, z_size, t_size))); ifd.tags.insert(Tag::short(277, &[frame_count as u16])); diff --git a/src/py.rs b/src/py.rs index 5f8c68a..0da86ae 100644 --- a/src/py.rs +++ b/src/py.rs @@ -1,7 +1,8 @@ -use crate::{Colors, IJTiffFile, Tag}; +use crate::{Colors, Compression, IJTiffFile, Tag}; use ndarray::s; use num::{Complex, FromPrimitive, Rational32}; use numpy::{PyArrayMethods, PyReadonlyArray2}; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; #[pyclass(subclass)] @@ -174,10 +175,21 @@ impl PyIJTiffFile { } /// set zstd compression level: -7 ..= 22 - fn set_compression_level(&mut self, compression_level: i32) { + fn set_compression(&mut self, compression: i32, level: i32) -> PyResult<()> { + let c = match compression { + 50000 => Compression::Zstd(level.clamp(-7, 22)), + 8 => Compression::Deflate, + _ => { + return Err(PyValueError::new_err(format!( + "Unknown compression {}", + compression + ))) + } + }; if let Some(ref mut ijtifffile) = self.ijtifffile { - ijtifffile.compression_level = compression_level.clamp(-7, 22); + ijtifffile.set_compression(c) } + Ok(()) } #[getter]