diff --git a/Cargo.toml b/Cargo.toml index 0bc729b..180fbb9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "tiffwrite" -version = "2024.10.1" +version = "2024.10.2" edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html diff --git a/py/tiffwrite/__init__.py b/py/tiffwrite/__init__.py index 85edd53..25cb020 100644 --- a/py/tiffwrite/__init__.py +++ b/py/tiffwrite/__init__.py @@ -33,11 +33,15 @@ FrameInfo = tuple[np.ndarray, None, CZT] class IJTiffFile(rs.IJTiffFile): - def __new__(cls, path: str | Path, shape: tuple[int, int, int], dtype: DTypeLike = 'uint16', + def __new__(cls, path: str | Path, shape: tuple[int, int, int] = None, dtype: DTypeLike = 'uint16', colors: Sequence[str] = None, colormap: str = None, pxsize: float = None, - deltaz: float = None, timeinterval: float = None, comment: str = None, + deltaz: float = None, timeinterval: float = None, compression: int = None, comment: str = None, **extratags: Tag) -> IJTiffFile: - new = super().__new__(cls, str(path), shape) + new = super().__new__(cls, str(path)) + if compression is not None: + if isinstance(compression, Sequence): + compression = compression[-1] + new.set_compression_level(compression) if colors is not None: new.colors = np.array([get_color(color) for color in colors]) if colormap is not None: @@ -54,7 +58,7 @@ class IJTiffFile(rs.IJTiffFile): new.append_extra_tag(extra_tag, None) return new - def __init__(self, path: str | Path, shape: tuple[int, int, int], dtype: DTypeLike = 'uint16', # noqa + def __init__(self, path: str | Path, shape: tuple[int, int, int] = None, dtype: DTypeLike = 'uint16', # noqa colors: Sequence[str] = None, colormap: str = None, pxsize: float = None, # noqa deltaz: float = None, timeinterval: float = None, comment: str = None, # noqa **extratags: Tag.Value | Tag) -> None: # noqa diff --git a/pyproject.toml b/pyproject.toml index 0078504..7ccaf6b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ classifiers = [ dependencies = ["colorcet", "matplotlib", "numpy", "tqdm"] [project.optional-dependencies] -test = ["pytest"] +test = ["pytest", "tifffile"] [tool.maturin] python-source = "py" diff --git a/src/lib.rs b/src/lib.rs index 580fbcf..bf45a63 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,26 +1,43 @@ -// #[cfg(not(feature = "nopython"))] +#[cfg(not(feature = "nopython"))] mod py; use anyhow::Result; use chrono::Utc; use ndarray::{s, Array2}; -use num::traits::ToBytes; -use num::{Complex, FromPrimitive, Rational32, Zero}; +use num::{traits::ToBytes, Complex, FromPrimitive, Rational32, Zero}; use rayon::prelude::*; -use std::cmp::Ordering; -use std::collections::HashMap; +use std::{cmp::Ordering, collections::HashMap}; use std::fs::{File, OpenOptions}; use std::hash::{DefaultHasher, Hash, Hasher}; -use std::io::{Read, Seek, SeekFrom, Write}; -use std::thread; -use std::thread::JoinHandle; -use zstd::stream::encode_all; +use std::io::{copy, Read, Seek, SeekFrom, Write}; +use std::{thread, thread::JoinHandle}; +use zstd::{DEFAULT_COMPRESSION_LEVEL, stream::Encoder}; const TAG_SIZE: usize = 20; const OFFSET_SIZE: usize = 8; const OFFSET: u64 = 16; const COMPRESSION: u16 = 50000; +pub fn encode_all(source: Vec, level: i32) -> Result> { + let mut result = Vec::::new(); + copy_encode(&*source, &mut result, level, source.len() as u64)?; + Ok(result) +} + +/// copy_encode from zstd crate, but let it include the content size in the zstd block header +pub fn copy_encode(mut source: R, destination: W, level: i32, length: u64) -> Result<()> +where + R: Read, + W: Write, +{ + let mut encoder = Encoder::new(destination, level)?; + encoder.include_contentsize(true)?; + encoder.set_pledged_src_size(Some(length))?; + copy(&mut source, &mut encoder)?; + encoder.finish()?; + Ok(()) +} + #[derive(Clone, Debug)] struct IFD { tags: Vec, @@ -366,8 +383,8 @@ struct CompressedFrame { #[derive(Clone, Debug)] struct Frame { - tileoffsets: Vec, - tilebytecounts: Vec, + offsets: Vec, + bytecounts: Vec, image_width: u32, image_length: u32, bits_per_sample: u16, @@ -378,8 +395,8 @@ struct Frame { impl Frame { fn new( - tileoffsets: Vec, - tilebytecounts: Vec, + offsets: Vec, + bytecounts: Vec, image_width: u32, image_length: u32, bits_per_sample: u16, @@ -388,8 +405,8 @@ impl Frame { tile_length: u16, ) -> Self { Frame { - tileoffsets, - tilebytecounts, + offsets, + bytecounts, image_width, image_length, bits_per_sample, @@ -455,7 +472,7 @@ pub struct IJTiffFile { frames: HashMap<(usize, usize, usize), Frame>, hashes: HashMap, threads: HashMap<(usize, usize, usize), JoinHandle>, - pub shape: (usize, usize, usize), + pub compression_level: i32, pub colors: Colors, pub comment: Option, pub px_size: Option, @@ -473,7 +490,7 @@ impl Drop for IJTiffFile { } impl IJTiffFile { - pub fn new(path: &str, shape: (usize, usize, usize)) -> Result { + pub fn new(path: &str) -> Result { let mut file = OpenOptions::new() .create(true) .truncate(true) @@ -490,7 +507,7 @@ impl IJTiffFile { frames: HashMap::new(), hashes: HashMap::new(), threads: HashMap::new(), - shape, + compression_level: DEFAULT_COMPRESSION_LEVEL, colors: Colors::None, comment: None, px_size: None, @@ -500,19 +517,23 @@ impl IJTiffFile { }) } - pub fn description(&self) -> String { + pub fn set_compression_level(&mut self, compression_level: i32) { + self.compression_level = compression_level; + } + + pub fn description(&self, c_size: usize, z_size: usize, t_size: usize) -> String { let mut desc: String = String::from("ImageJ=1.11a"); if let Colors::None = self.colors { - desc += &format!("\nimages={}", self.shape.0); - desc += &format!("\nslices={}", self.shape.1); - desc += &format!("\nframes={}", self.shape.2); + desc += &format!("\nimages={}", c_size); + desc += &format!("\nslices={}", z_size); + desc += &format!("\nframes={}", t_size); } else { - desc += &format!("\nimages={}", self.shape.0 * self.shape.1 * self.shape.2); - desc += &format!("\nchannels={}", self.shape.0); - desc += &format!("\nslices={}", self.shape.1); - desc += &format!("\nframes={}", self.shape.2); + desc += &format!("\nimages={}", c_size * z_size * t_size); + desc += &format!("\nchannels={}", c_size); + desc += &format!("\nslices={}", z_size); + desc += &format!("\nframes={}", t_size); }; - if self.shape.0 == 1 { + if c_size == 1 { desc += "\nmode=grayscale"; } else { desc += "\nmode=composite"; @@ -530,27 +551,33 @@ impl IJTiffFile { desc } - fn get_czt(&self, frame_number: usize, channel: u8) -> (usize, usize, usize) { + fn get_czt( + &self, + frame_number: usize, + channel: u8, + c_size: usize, + z_size: usize, + ) -> (usize, usize, usize) { if let Colors::None = self.colors { ( channel as usize, - frame_number % self.shape.1, - frame_number / self.shape.1, + frame_number % z_size, + frame_number / z_size, ) } else { ( - frame_number % self.shape.0, - frame_number / self.shape.0 % self.shape.1, - frame_number / self.shape.0 / self.shape.1, + frame_number % c_size, + frame_number / c_size % z_size, + frame_number / c_size / z_size, ) } } - fn spp_and_n_frames(&self) -> (u8, usize) { + fn spp_and_n_frames(&self, c_size: usize, z_size: usize, t_size: usize) -> (u8, usize) { if let Colors::None = &self.colors { - (self.shape.0 as u8, self.shape.1 * self.shape.2) + (c_size as u8, z_size * t_size) } else { - (1, self.shape.0 * self.shape.1 * self.shape.2) + (1, c_size * z_size * t_size) } } @@ -599,7 +626,7 @@ impl IJTiffFile { where T: Bytes + Clone + Zero + Send + 'static, { - fn compress(frame: Array2) -> CompressedFrame + fn compress(frame: Array2, compression_level: i32) -> CompressedFrame where T: Bytes + Clone + Zero, { @@ -607,7 +634,7 @@ impl IJTiffFile { let image_length = frame.shape()[1] as u32; let tile_size = 2usize .pow( - ((image_width as f64 * image_length as f64 / 64f64).log2() / 2f64).round() + ((image_width as f64 * image_length as f64 / 2f64).log2() / 2f64).round() as u32, ) .max(16) @@ -619,7 +646,7 @@ impl IJTiffFile { .collect(); let bytes = byte_tiles .into_par_iter() - .map(|x| encode_all(&*x, 3).unwrap()) + .map(|x| encode_all(x, compression_level).unwrap()) .collect::>(); CompressedFrame { bytes, @@ -630,8 +657,11 @@ impl IJTiffFile { sample_format: T::SAMPLE_FORMAT, } } - self.threads - .insert((c, z, t), thread::spawn(move || compress(frame))); + let compression_level = self.compression_level; + self.threads.insert( + (c, z, t), + thread::spawn(move || compress(frame, compression_level)), + ); for key in self .threads .keys() @@ -652,15 +682,15 @@ impl IJTiffFile { } fn write_frame(&mut self, frame: CompressedFrame, c: usize, z: usize, t: usize) -> Result<()> { - let mut tileoffsets = Vec::new(); - let mut tilebytecounts = Vec::new(); + let mut offsets = Vec::new(); + let mut bytecounts = Vec::new(); for tile in frame.bytes { - tilebytecounts.push(tile.len() as u64); - tileoffsets.push(self.write(&tile)?); + bytecounts.push(tile.len() as u64); + offsets.push(self.write(&tile)?); } let frame = Frame::new( - tileoffsets, - tilebytecounts, + offsets, + bytecounts, frame.image_width, frame.image_length, frame.bits_per_sample, @@ -674,8 +704,8 @@ impl IJTiffFile { fn tile(frame: Array2, size: usize) -> Vec> { let shape = frame.shape(); - let mut tiles = Vec::new(); let (n, m) = (shape[0] / size, shape[1] / size); + let mut tiles = Vec::new(); for i in 0..n { for j in 0..m { tiles.push( @@ -726,8 +756,15 @@ impl IJTiffFile { } } - fn get_color(&self, _colors: &Vec, _bits_per_sample: u16) -> Result> { - todo!(); + fn get_color(&self, colors: &Vec, bits_per_sample: u16) -> Result> { + let mut c = Vec::new(); + let lvl = if bits_per_sample == 8 { 255 } else { 65535 }; + for i in 0..=lvl { + c.push(i * (colors[0] as u16) / 255); + c.push(i * (colors[1] as u16) / 255); + c.push(i * (colors[2] as u16) / 255); + } + Ok(c) } fn close(&mut self) -> Result<()> { @@ -736,18 +773,33 @@ impl IJTiffFile { self.write_frame(thread.join().unwrap(), c, z, t)?; } } + let mut c_size = 1; + let mut z_size = 1; + let mut t_size = 1; + for (c, z, t) in self.frames.keys() { + c_size = c_size.max(c + 1); + z_size = z_size.max(z + 1); + t_size = t_size.max(t + 1); + } + let mut where_to_write_next_ifd_offset = OFFSET - OFFSET_SIZE as u64; let mut warn = false; - let (samples_per_pixel, n_frames) = self.spp_and_n_frames(); + let (samples_per_pixel, n_frames) = self.spp_and_n_frames(c_size, t_size, z_size); for frame_number in 0..n_frames { - if let Some(frame) = self.frames.get(&self.get_czt(frame_number, 0)) { - let mut tileoffsets = Vec::new(); - let mut tilebytecounts = Vec::new(); + if let Some(frame) = self + .frames + .get(&self.get_czt(frame_number, 0, c_size, z_size)) + { + let mut offsets = Vec::new(); + let mut bytecounts = Vec::new(); let mut frame_count = 0; for channel in 0..samples_per_pixel { - if let Some(frame_n) = self.frames.get(&self.get_czt(frame_number, channel)) { - tileoffsets.extend(frame_n.tileoffsets.iter()); - tilebytecounts.extend(frame_n.tilebytecounts.iter()); + if let Some(frame_n) = + self.frames + .get(&self.get_czt(frame_number, channel, c_size, z_size)) + { + offsets.extend(frame_n.offsets.iter()); + bytecounts.extend(frame_n.bytecounts.iter()); frame_count += 1; } else { warn = true; @@ -758,30 +810,33 @@ impl IJTiffFile { ifd.push_tag(Tag::long(257, &vec![frame.image_length])); ifd.push_tag(Tag::short(258, &vec![frame.bits_per_sample; frame_count])); ifd.push_tag(Tag::short(259, &vec![COMPRESSION])); - ifd.push_tag(Tag::ascii(270, &self.description())); + ifd.push_tag(Tag::ascii(270, &self.description(c_size, z_size, t_size))); ifd.push_tag(Tag::short(277, &vec![frame_count as u16])); - ifd.push_tag(Tag::ascii(305, "tiffwrite_rs")); + ifd.push_tag(Tag::ascii(305, "tiffwrite_tllab_NKI")); ifd.push_tag(Tag::short(322, &vec![frame.tile_width])); ifd.push_tag(Tag::short(323, &vec![frame.tile_length])); - ifd.push_tag(Tag::long8(324, &tileoffsets)); - ifd.push_tag(Tag::long8(325, &tilebytecounts)); - ifd.push_tag(Tag::short(339, &vec![frame.sample_format])); + ifd.push_tag(Tag::long8(324, &offsets)); + ifd.push_tag(Tag::long8(325, &bytecounts)); + if frame.sample_format > 1 { + ifd.push_tag(Tag::short(339, &vec![frame.sample_format])); + } if let Some(px_size) = self.px_size { let r = vec![Rational32::from_f64(px_size).unwrap()]; ifd.push_tag(Tag::rational(282, &r)); ifd.push_tag(Tag::rational(283, &r)); ifd.push_tag(Tag::short(296, &vec![1])); } - + if let Colors::Colormap(_) = &self.colors { + ifd.push_tag(Tag::short(262, &vec![3])); + } else if let Colors::None = self.colors { + ifd.push_tag(Tag::short(262, &vec![1])); + } if frame_number == 0 { if let Colors::Colormap(colormap) = &self.colors { ifd.push_tag(Tag::short( 320, &self.get_colormap(colormap, frame.bits_per_sample), )); - ifd.push_tag(Tag::short(262, &vec![3])); - } else if let Colors::None = self.colors { - ifd.push_tag(Tag::short(262, &vec![1])); } } if frame_number < samples_per_pixel as usize { @@ -794,12 +849,12 @@ impl IJTiffFile { } } if let Colors::None = &self.colors { - if self.shape.0 > 1 { + if c_size > 1 { ifd.push_tag(Tag::short(284, &vec![2])) } } for channel in 0..samples_per_pixel { - let czt = self.get_czt(frame_number, channel); + let czt = self.get_czt(frame_number, channel, c_size, z_size); if let Some(extra_tags) = self.extra_tags.get(&Some(czt)) { for tag in extra_tags { ifd.push_tag(tag.to_owned()) @@ -811,10 +866,12 @@ impl IJTiffFile { ifd.push_tag(tag.to_owned()) } } - ifd.push_tag(Tag::ascii( - 306, - &format!("{}", Utc::now().format("%Y:%m:%d %H:%M:%S")), - )); + if frame_number == 0 { + ifd.push_tag(Tag::ascii( + 306, + &format!("{}", Utc::now().format("%Y:%m:%d %H:%M:%S")), + )); + } where_to_write_next_ifd_offset = ifd.write(self, where_to_write_next_ifd_offset)?; } else { warn = true; diff --git a/src/main.rs b/src/main.rs index 75a03da..449463d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,7 +4,8 @@ use tiffwrite::IJTiffFile; fn main() -> Result<()> { println!("Hello World!"); - let mut f = IJTiffFile::new("foo.tif", (2, 1, 1))?; + let mut f = IJTiffFile::new("foo.tif")?; + f.set_compression_level(10); let mut arr = Array2::::zeros((100, 100)); for i in 0..arr.shape()[0] { for j in 0..arr.shape()[1] { diff --git a/src/py.rs b/src/py.rs index 693b338..ce567f2 100644 --- a/src/py.rs +++ b/src/py.rs @@ -165,12 +165,18 @@ struct PyIJTiffFile { #[pymethods] impl PyIJTiffFile { #[new] - fn new(path: &str, shape: (usize, usize, usize)) -> PyResult { + fn new(path: &str) -> PyResult { Ok(PyIJTiffFile { - ijtifffile: Some(IJTiffFile::new(path, shape)?), + ijtifffile: Some(IJTiffFile::new(path)?), }) } + fn set_compression_level(&mut self, compression_level: i32) { + if let Some(ref mut ijtifffile) = self.ijtifffile { + ijtifffile.set_compression_level(compression_level); + } + } + #[getter] fn get_colors(&self) -> PyResult>>> { if let Some(ijtifffile) = &self.ijtifffile { diff --git a/tests/test_multiple.py b/tests/test_multiple.py index 37e7abb..936c4bc 100644 --- a/tests/test_multiple.py +++ b/tests/test_multiple.py @@ -12,7 +12,7 @@ def test_mult(tmp_path: Path) -> None: shape = (2, 3, 5) paths = [tmp_path / f'test{i}.tif' for i in range(6)] with ExitStack() as stack: - tifs = [stack.enter_context(IJTiffFile(path, shape)) for path in paths] # noqa + tifs = [stack.enter_context(IJTiffFile(path)) for path in paths] # noqa for c, z, t in tqdm(product(range(shape[0]), range(shape[1]), range(shape[2])), total=np.prod(shape)): # noqa for tif in tifs: tif.save(np.random.randint(0, 255, (64, 64)), c, z, t) diff --git a/tests/test_single.py b/tests/test_single.py index 2d47470..8583b9b 100644 --- a/tests/test_single.py +++ b/tests/test_single.py @@ -1,14 +1,28 @@ -from itertools import product from pathlib import Path import numpy as np +import pytest +from tifffile import imread from tiffwrite import IJTiffFile -def test_single(tmp_path: Path) -> None: - path = tmp_path / 'test.tif' - with IJTiffFile(path, (3, 4, 5)) as tif: - for c, z, t in product(range(3), range(4), range(5)): - tif.save(np.random.randint(0, 255, (64, 64)), c, z, t) - assert path.exists() +@pytest.mark.parametrize('dtype', ('uint8', 'uint16', 'uint32', 'uint64', + 'int8', 'int16', 'int32', 'int64', 'float32', 'float64')) +def test_single(tmp_path: Path, dtype) -> None: + with IJTiffFile(tmp_path / 'test.tif', dtype=dtype) as tif: + a0, b0 = np.meshgrid(range(100), range(100)) + a0[::2, :] = 0 + b0[:, ::2] = 1 + tif.save(a0, 0, 0, 0) + tif.save(b0, 1, 0, 0) + + a1, b1 = np.meshgrid(range(100), range(100)) + a1[:, ::2] = 0 + b1[::2, :] = 1 + tif.save(a1, 0, 0, 1) + tif.save(b1, 1, 0, 1) + + t = imread(tmp_path / 'test.tif') + assert t.dtype == np.dtype(dtype), "data type does not match" + assert np.all(np.stack(((a0, b0), (a1, b1))) == t), "data does not match"