- make zstd block include content size so fiji can actually read it

- add compression level argument
- remove shape argument
- some more pytest tests

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

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

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"

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<u8>, level: i32) -> Result<Vec<u8>> {
+    let mut result = Vec::<u8>::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<R, W>(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<Tag>,
@@ -366,8 +383,8 @@ struct CompressedFrame {
 
 #[derive(Clone, Debug)]
 struct Frame {
-    tileoffsets: Vec<u64>,
-    tilebytecounts: Vec<u64>,
+    offsets: Vec<u64>,
+    bytecounts: Vec<u64>,
     image_width: u32,
     image_length: u32,
     bits_per_sample: u16,
@@ -378,8 +395,8 @@ struct Frame {
 
 impl Frame {
     fn new(
-        tileoffsets: Vec<u64>,
-        tilebytecounts: Vec<u64>,
+        offsets: Vec<u64>,
+        bytecounts: Vec<u64>,
         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<u64, u64>,
     threads: HashMap<(usize, usize, usize), JoinHandle<CompressedFrame>>,
-    pub shape: (usize, usize, usize),
+    pub compression_level: i32,
     pub colors: Colors,
     pub comment: Option<String>,
     pub px_size: Option<f64>,
@@ -473,7 +490,7 @@ impl Drop for IJTiffFile {
 }
 
 impl IJTiffFile {
-    pub fn new(path: &str, shape: (usize, usize, usize)) -> Result<Self> {
+    pub fn new(path: &str) -> Result<Self> {
         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<T>(frame: Array2<T>) -> CompressedFrame
+        fn compress<T>(frame: Array2<T>, 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::<Vec<_>>();
             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<T: Clone + Zero>(frame: Array2<T>, size: usize) -> Vec<Array2<T>> {
         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<u8>, _bits_per_sample: u16) -> Result<Vec<u16>> {
-        todo!();
+    fn get_color(&self, colors: &Vec<u8>, bits_per_sample: u16) -> Result<Vec<u16>> {
+        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::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())
                     }
                 }
+                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;

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::<u16>::zeros((100, 100));
     for i in 0..arr.shape()[0] {
         for j in 0..arr.shape()[1] {

src/py.rs

@@ -165,12 +165,18 @@ struct PyIJTiffFile {
 #[pymethods]
 impl PyIJTiffFile {
     #[new]
-    fn new(path: &str, shape: (usize, usize, usize)) -> PyResult<Self> {
+    fn new(path: &str) -> PyResult<Self> {
         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<Option<Vec<Vec<u8>>>> {
         if let Some(ijtifffile) = &self.ijtifffile {

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)

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"