- special IJTiffParallel class to help generate frames in parallel

- warning now shows which frames are missing

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tiffwrite"
-version = "2024.10.3"
+version = "2024.10.4"
 edition = "2021"
 
 [lib]
@@ -17,7 +17,7 @@ zstd = "0.13.2"
 numpy = { version = "0.22.0", optional = true }
 
 [dependencies.pyo3]
-version = "0.22.2"
+version = "0.22.5"
 features = ["extension-module", "abi3-py310", "generate-import-lib", "anyhow", "multiple-pymethods"]
 optional = true
 

py/tiffwrite/__init__.py

@@ -14,24 +14,19 @@ from tqdm.auto import tqdm
 
 from . import tiffwrite_rs as rs  # noqa
 
-__all__ = ['Header', 'IJTiffFile', 'IFD', 'FrameInfo', 'Tag', 'Strip', 'tiffwrite']
+__all__ = ['IJTiffFile', 'IJTiffParallel', 'FrameInfo', 'Tag', 'tiffwrite']
+
+
+Tag = rs.Tag
+FrameInfo = tuple[np.ndarray, int, int, int]
 
 
 class Header:
-    pass
+    """ deprecated """
 
 
 class IFD(dict):
-    pass
+    """ deprecated """
-
-
-class Tag(rs.Tag):
-    pass
-
-
-Strip = tuple[list[int], list[int]]
-CZT = tuple[int, int, int]
-FrameInfo = tuple[np.ndarray, None, CZT]
 
 
 class TiffWriteWarning(UserWarning):
@@ -40,6 +35,7 @@ class TiffWriteWarning(UserWarning):
 
 class IJTiffFile(rs.IJTiffFile):
     """ Writes a tiff file in a format that the BioFormats reader in Fiji understands.
+        Zstd compression is done in parallel using Rust.
         file:           filename of the new tiff file
         shape:          not used anymore
         dtype:          datatype to use when saving to tiff
@@ -97,38 +93,33 @@ class IJTiffFile(rs.IJTiffFile):
 
     def save(self, frame: ArrayLike, c: int, z: int, t: int, extratags: Sequence[Tag] = None) -> None:
         """ save a 2d numpy array to the tiff at channel=c, slice=z, time=t, with optional extra tif tags """
-        for frame, _, (cn, zn, tn) in self.compress_frame(frame):
+        frame = np.asarray(frame).astype(self.dtype)
-            frame = np.asarray(frame).astype(self.dtype)
+        match self.dtype:
-            match self.dtype:
+            case np.uint8:
-                case np.uint8:
+                self.save_u8(frame, c, z, t)
-                    self.save_u8(frame, c + cn, z + zn, t + tn)
+            case np.uint16:
-                case np.uint16:
+                self.save_u16(frame, c, z, t)
-                    self.save_u16(frame, c + cn, z + zn, t + tn)
+            case np.uint32:
-                case np.uint32:
+                self.save_u32(frame, c, z, t)
-                    self.save_u32(frame, c + cn, z + zn, t + tn)
+            case np.uint64:
-                case np.uint64:
+                self.save_u64(frame, c, z, t)
-                    self.save_u64(frame, c + cn, z + zn, t + tn)
+            case np.int8:
-                case np.int8:
+                self.save_i8(frame, c, z, t)
-                    self.save_i8(frame, c + cn, z + zn, t + tn)
+            case np.int16:
-                case np.int16:
+                self.save_i16(frame, c, z, t)
-                    self.save_i16(frame, c + cn, z + zn, t + tn)
+            case np.int32:
-                case np.int32:
+                self.save_i32(frame, c, z, t)
-                    self.save_i32(frame, c + cn, z + zn, t + tn)
+            case np.int64:
-                case np.int64:
+                self.save_i64(frame, c, z, t)
-                    self.save_i64(frame, c + cn, z + zn, t + tn)
+            case np.float32:
-                case np.float32:
+                self.save_f32(frame, c, z, t)
-                    self.save_f32(frame, c + cn, z + zn, t + tn)
+            case np.float64:
-                case np.float64:
+                self.save_f64(frame, c, z, t)
-                    self.save_f64(frame, c + cn, z + zn, t + tn)
+            case _:
-                case _:
+                raise TypeError(f'Cannot save type {self.dtype}')
-                    raise TypeError(f'Cannot save type {self.dtype}')
+        if extratags is not None:
-            if extratags is not None:
+            for extra_tag in extratags:
-                for extra_tag in extratags:
+                self.append_extra_tag(extra_tag, (c, z, t))
-                    self.append_extra_tag(extra_tag, (c, z, t))
-
-    def compress_frame(self, frame: ArrayLike) -> tuple[FrameInfo]:  # noqa
-        """ backwards compatibility """
-        return (frame, None, (0, 0, 0)),
 
 
 def get_colormap(colormap: str) -> np.ndarray:
@@ -181,3 +172,45 @@ def tiffwrite(file: str | Path, data: np.ndarray, axes: str = 'TZCXY', dtype: DT
         for n in tqdm(product(*[range(i) for i in shape]), total=np.prod(shape),  # type: ignore
                       desc='Saving tiff', disable=not bar):
             f.save(data[n], *n)
+
+
+try:
+    from parfor import ParPool, Task
+    from abc import abstractmethod, ABCMeta
+    from functools import wraps
+
+    class IJTiffParallel(ParPool, metaclass=ABCMeta):
+        """ wraps IJTiffFile.save in a parallel pool, the method 'parallel' needs to be overloaded """
+
+        @abstractmethod
+        def parallel(self, frame: Any) -> Sequence[tuple[ArrayLike, int, int, int]]:
+            """ does something with frame in a parallel process,
+                and returns a sequence of frames and offsets to c, z and t to save in the tif """
+
+        @wraps(IJTiffFile.__init__)
+        def __init__(self, *args: Any, **kwargs: Any) -> None:
+            self.ijtifffile = IJTiffFile(*args, **kwargs)
+            super().__init__(self.parallel)  # noqa
+
+        def done(self, task: Task) -> None:
+            c, z, t = task.handle
+            super().done(task)
+            for frame, cn, zn, tn in self[c, z, t]:
+                self.ijtifffile.save(frame, c + cn, z + zn, t + tn)
+
+        @wraps(IJTiffFile.close)
+        def close(self) -> None:
+            while len(self.tasks):
+                self.get_newest()
+            super().close()
+            self.ijtifffile.close()
+
+        @wraps(IJTiffFile.save)
+        def save(self, frame: Any, c: int, z: int, t: int, extratags: Sequence[Tag] = None) -> None:
+            self[c, z, t] = frame
+            if extratags is not None:
+                for extra_tag in extratags:
+                    self.ijtifffile.append_extra_tag(extra_tag, (c, z, t))
+
+except ImportError:
+    IJTiffPool = None

src/lib.rs

@@ -803,7 +803,7 @@ impl IJTiffFile {
         }
 
         let mut where_to_write_next_ifd_offset = OFFSET - OFFSET_SIZE as u64;
-        let mut warn = false;
+        let mut warn = Vec::new();
         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
@@ -822,7 +822,7 @@ impl IJTiffFile {
                         bytecounts.extend(frame_n.bytecounts.iter());
                         frame_count += 1;
                     } else {
-                        warn = true;
+                        warn.push((frame_number, channel));
                     }
                 }
                 let mut ifd = IFD::new();
@@ -895,13 +895,16 @@ impl IJTiffFile {
                 }
                 where_to_write_next_ifd_offset = ifd.write(self, where_to_write_next_ifd_offset)?;
             } else {
-                warn = true;
+                warn.push((frame_number, 0));
             }
-            if warn {
+            if warn.len() > 0 {
-                println!(
+                println!("The following frames were not added to the tif file");
-                    "Some frames were not added to the tif file, either you forgot them, \
+                for (frame_number, channel) in &warn {
-                    or an error occurred and the tif file was closed prematurely."
+                    let (c, z, t) = self.get_czt(*frame_number, *channel, c_size, z_size);
-                )
+                    println!("{c}, {z}, {t}")
+                }
+                println!("Either you forgot them, \
+                         or an error occurred and the tif file was closed prematurely.")
             }
         }
         self.file