diff --git a/.github/workflows/mypy.yml b/.github/workflows/mypy.yml deleted file mode 100644 index 9146014..0000000 --- a/.github/workflows/mypy.yml +++ /dev/null @@ -1,21 +0,0 @@ -name: MyPy - -on: [push, pull_request] - -jobs: - mypy: - runs-on: ubuntu-latest - strategy: - matrix: - python-version: ["3.10", "3.12"] - - steps: - - uses: actions/checkout@v4 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - name: Install dependencies - run: pip install .[test] - - name: Test with mypy - run: mypy . \ No newline at end of file diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..901cc38 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,24 @@ +name: Publish + +on: workflow_dispatch + +jobs: + publish_wheels: + uses: ./.github/workflows/wheels.yml + publish: + name: publish + needs: publish_wheels + runs-on: ubuntu-latest + environment: + name: pypi + url: https://pypi.org/p/tiffwrite + permissions: + id-token: write + steps: + - name: Download all the dists + uses: actions/download-artifact@v4 + with: + name: python-package-distributions + path: dist/ + - name: Publish distribution 📦 to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 \ No newline at end of file diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index c61decd..dafe279 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -1,6 +1,6 @@ name: PyTest -on: [push, pull_request] +on: [workflow_call, push, pull_request] jobs: pytest: diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml new file mode 100644 index 0000000..a93bec7 --- /dev/null +++ b/.github/workflows/wheels.yml @@ -0,0 +1,42 @@ +name: Wheels + +on: workflow_call + +jobs: + wheels_pytest: + uses: ./.github/workflows/pytest.yml + build_wheels: + name: Build wheels on ${{ matrix.os }} + needs: [ wheels_pytest ] + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ ubuntu-latest, windows-latest, macos-13, macos-latest ] + + steps: + - uses: actions/checkout@v4 + + - name: Build wheels + uses: pypa/cibuildwheel@v2.21.2 + + - uses: actions/upload-artifact@v4 + with: + name: tiffwrite-wheels-${{ matrix.os }}-${{ strategy.job-index }} + path: ./wheelhouse/*.whl + + build_sdist: + name: Build source distribution + needs: [ wheels_pytest ] + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Build sdist + run: | + pip install build + python -m build --sdist + + - uses: actions/upload-artifact@v4 + with: + name: tiffwrite-sdist + path: dist/*.tar.gz \ No newline at end of file diff --git a/.gitignore b/.gitignore index 5a5a156..cb58d10 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,7 @@ /tiffwrite.egg-info/ /.pytest_cache/ /venv/ +/target/ +/Cargo.lock +/foo.tif +*.so diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..a14a93e --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "tiffwrite" +version = "2024.10.4" +edition = "2021" + +[lib] +name = "tiffwrite" +crate-type = ["cdylib", "rlib"] + +[dependencies] +anyhow = "1.0.89" +chrono = "0.4.38" +ndarray = "0.16.1" +num = "0.4.3" +rayon = "1.10.0" +zstd = "0.13.2" +numpy = { version = "0.22.0", optional = true } + +[dependencies.pyo3] +version = "0.22.5" +features = ["extension-module", "abi3-py310", "generate-import-lib", "anyhow", "multiple-pymethods"] +optional = true + +[features] +python = ["dep:pyo3", "dep:numpy"] diff --git a/README.md b/README.md index bb9af0a..308162a 100644 --- a/README.md +++ b/README.md @@ -1,23 +1,26 @@ -[![mypy](https://github.com/wimpomp/tiffwrite/actions/workflows/mypy.yml/badge.svg)](https://github.com/wimpomp/tiffwrite/actions/workflows/mypy.yml) [![pytest](https://github.com/wimpomp/tiffwrite/actions/workflows/pytest.yml/badge.svg)](https://github.com/wimpomp/tiffwrite/actions/workflows/pytest.yml) # Tiffwrite -Exploiting [tifffile](https://pypi.org/project/tifffile/) in parallel to write BioFormats/ImageJ compatible tiffs with -good compression. +Write BioFormats/ImageJ compatible tiffs with zstd compression in parallel using Rust. ## Features - Writes bigtiff files that open in ImageJ as hyperstack with correct dimensions. - Parallel compression. - Write individual frames in random order. - Compresses even more by referencing tag or image data which otherwise would have been saved several times. -For example empty frames, or a long string tag on every frame. +For example empty frames, or a long string tag on every frame. Editing tiffs becomes mostly impossible, but compression +makes that very hard anyway. - Enables memory efficient scripts by saving frames whenever they're ready to be saved, not waiting for the whole stack. -- Colormaps, extra tags, globally or frame dependent. +- Colormaps +- Extra tags, globally or frame dependent. ## Installation pip install tiffwrite or +- install [rust](https://rustup.rs/) + + pip install tiffwrite@git+https://github.com/wimpomp/tiffwrite # Usage @@ -67,11 +70,10 @@ or from tiffwrite import IJTiffFile import numpy as np - shape = (3, 5, 10) # channels, z, time - with IJTiffFile('file.tif', shape, pxsize=0.09707) as tif: - for c in range(shape[0]): - for z in range(shape[1]): - for t in range(shape[2]): + with IJTiffFile('file.tif', pxsize=0.09707) as tif: + for c in range(3): + for z in range(5): + for t in range(10): tif.save(np.random.randint(0, 10, (32, 32)), c, z, t) ## Saving multiple tiffs simultaneously @@ -79,7 +81,7 @@ or import numpy as np shape = (3, 5, 10) # channels, z, time - with IJTiffFile('fileA.tif', shape) as tif_a, IJTiffFile('fileB.tif', shape) as tif_b: + with IJTiffFile('fileA.tif') as tif_a, IJTiffFile('fileB.tif') as tif_b: for c in range(shape[0]): for z in range(shape[1]): for t in range(shape[2]): diff --git a/py/tiffwrite/__init__.py b/py/tiffwrite/__init__.py new file mode 100644 index 0000000..322236e --- /dev/null +++ b/py/tiffwrite/__init__.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +from itertools import product +from pathlib import Path +from typing import Any, Sequence +from warnings import warn + +import colorcet +import matplotlib +import numpy as np +from matplotlib import colors as mpl_colors +from numpy.typing import ArrayLike, DTypeLike +from tqdm.auto import tqdm + +from . import tiffwrite_rs as rs # noqa + +__all__ = ['IJTiffFile', 'IJTiffParallel', 'FrameInfo', 'Tag', 'tiffwrite'] + + +Tag = rs.Tag +FrameInfo = tuple[np.ndarray, int, int, int] + + +class Header: + """ deprecated """ + + +class IFD(dict): + """ deprecated """ + + +class TiffWriteWarning(UserWarning): + pass + + +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 + colors: a tuple with a color per channel, chosen from matplotlib.colors, html colors are also possible + colormap: name of a colormap from colorcet + pxsize: pixel size in um + deltaz: z slice interval in um + timeinterval: time between frames in seconds + extratags: other tags to be saved, example: (Tag.ascii(315, 'John Doe'), Tag.bytes(4567, [400, 500]) + or (Tag.ascii(33432, 'Made by me'),). + """ + def __new__(cls, path: str | Path, *args, **kwargs) -> IJTiffFile: + return super().__new__(cls, str(path)) + + def __init__(self, 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, compression: int = None, comment: str = None, + extratags: Sequence[Tag] = None) -> None: + self.path = Path(path) + self.shape = shape + self.dtype = np.dtype(dtype) + if compression is not None: + if isinstance(compression, Sequence): + compression = compression[-1] + self.set_compression_level(compression) + if colors is not None: + self.colors = np.array([get_color(color) for color in colors]) + if colormap is not None: + self.colormap = get_colormap(colormap) + if pxsize is not None: + self.px_size = float(pxsize) + if deltaz is not None: + self.delta_z = float(deltaz) + if timeinterval is not None: + self.time_interval = float(timeinterval) + if comment is not None: + self.comment = comment + if extratags is not None: + for extra_tag in extratags: + self.append_extra_tag(extra_tag, None) + if self.dtype.itemsize == 1 and colors is not None: + warn('Fiji will not interpret colors saved in an (u)int8 tif, save as (u)int16 instead.', + TiffWriteWarning, stacklevel=2) + if shape is not None: + warn('Providing shape is not needed anymore, the argument will be removed in the future.', + DeprecationWarning, stacklevel=2) + if colors is not None and colormap is not None: + warn('Cannot have colors and colormap simultaneously.', TiffWriteWarning, stacklevel=2) + + def __enter__(self) -> IJTiffFile: + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + + 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 """ + frame = np.asarray(frame).astype(self.dtype) + match self.dtype: + case np.uint8: + self.save_u8(frame, c, z, t) + case np.uint16: + self.save_u16(frame, c, z, t) + case np.uint32: + self.save_u32(frame, c, z, t) + case np.uint64: + self.save_u64(frame, c, z, t) + case np.int8: + self.save_i8(frame, c, z, t) + case np.int16: + self.save_i16(frame, c, z, t) + case np.int32: + self.save_i32(frame, c, z, t) + case np.int64: + self.save_i64(frame, c, z, t) + case np.float32: + self.save_f32(frame, c, z, t) + case np.float64: + self.save_f64(frame, c, z, t) + case _: + raise TypeError(f'Cannot save type {self.dtype}') + if extratags is not None: + for extra_tag in extratags: + self.append_extra_tag(extra_tag, (c, z, t)) + + +def get_colormap(colormap: str) -> np.ndarray: + if hasattr(colorcet, colormap.rstrip('_r')): + cm = np.array([[int(''.join(i), 16) for i in zip(*[iter(s[1:])] * 2)] + for s in getattr(colorcet, colormap.rstrip('_r'))]).astype('uint8') + if colormap.endswith('_r'): + cm = cm[::-1] + if colormap.startswith('glasbey') or colormap.endswith('glasbey'): + cm[0] = 255, 255, 255 + cm[-1] = 0, 0, 0 + else: + cmap = matplotlib.colormaps.get_cmap(colormap) + if cmap.N < 256: + cm = (255 * np.vstack(((1, 1, 1), + matplotlib.cm.ScalarMappable(matplotlib.colors.Normalize(1, 254), + cmap).to_rgba(np.arange(1, 254))[:, :3], + (0, 0, 0)))).astype('uint8') + else: + cm = (255 * matplotlib.cm.ScalarMappable(matplotlib.colors.Normalize(0, 255), cmap) + .to_rgba(np.arange(256))[:, :3]).astype('uint8') + return cm + + +def get_color(color: str) -> np.ndarray: + return np.array([int(''.join(i), 16) for i in zip(*[iter(mpl_colors.to_hex(color)[1:])] * 2)]).astype('uint8') + + +def tiffwrite(file: str | Path, data: np.ndarray, axes: str = 'TZCXY', dtype: DTypeLike = None, bar: bool = False, + *args: Any, **kwargs: Any) -> None: + """ file: string; filename of the new tiff file + data: 2 to 5D numpy array + axes: string; order of dimensions in data, default: TZCXY for 5D, ZCXY for 4D, CXY for 3D, XY for 2D data + dtype: string; datatype to use when saving to tiff + bar: bool; whether to show a progress bar + other args: see IJTiffFile + """ + + axes = axes[-np.ndim(data):].upper() + if not axes == 'CZTXY': + axes_shuffle = [axes.find(i) for i in 'CZTXY'] + axes_add = [i for i, j in enumerate(axes_shuffle) if j < 0] + axes_shuffle = [i for i in axes_shuffle if i >= 0] + data = np.transpose(data, axes_shuffle) + for axis in axes_add: + data = np.expand_dims(data, axis) + + shape = data.shape[:3] + with IJTiffFile(file, dtype=data.dtype if dtype is None else dtype, *args, **kwargs) as f: + 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 diff --git a/pyproject.toml b/pyproject.toml index 2db2f95..7bddf15 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,39 +1,26 @@ -[tool.poetry] +[build-system] +requires = ["maturin>=1.5,<2.0"] +build-backend = "maturin" + +[project] name = "tiffwrite" -version = "2024.10.1" -description = "Parallel tiff writer compatible with ImageJ." -authors = ["Wim Pomp, Lenstra lab NKI "] -license = "GPL-3.0-or-later" -readme = "README.md" -packages = [{include = "tiffwrite"}] -repository = "https://github.com/wimpomp/tiffwrite" +dynamic = ["version"] +authors = [{ name = "Wim Pomp", email = "w.pomp@nki.nl" }] +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", + "Programming Language :: Python :: Implementation :: PyPy", +] +dependencies = ["colorcet", "matplotlib", "numpy", "tqdm"] -[tool.poetry.dependencies] -python = "^3.10" -tifffile = "*" -imagecodecs = "*" -numpy = "*" -tqdm = "*" -colorcet = "*" -matplotlib = "*" -parfor = ">=2024.9.2" -pytest = { version = "*", optional = true } -mypy = { version = "*", optional = true } +[project.optional-dependencies] +test = ["pytest", "tifffile", "imagecodecs"] -[tool.poetry.extras] -test = ["pytest", "mypy"] - -[tool.pytest.ini_options] -filterwarnings = ["ignore:::(?!tiffwrite)"] +[tool.maturin] +python-source = "py" +features = ["pyo3/extension-module", "python"] +module-name = "tiffwrite.tiffwrite_rs" [tool.isort] -line_length = 119 - -[tool.mypy] -disable_error_code = ["import-untyped", "return"] -implicit_optional = true -exclude = ["build"] - -[build-system] -requires = ["poetry-core"] -build-backend = "poetry.core.masonry.api" +line_length = 119 \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..2df7557 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,915 @@ +#[cfg(feature = "python")] +mod py; + +use anyhow::Result; +use chrono::Utc; +use ndarray::{s, Array2}; +use num::{traits::ToBytes, Complex, FromPrimitive, Rational32, Zero}; +use rayon::prelude::*; +use std::collections::HashSet; +use std::fs::{File, OpenOptions}; +use std::hash::{DefaultHasher, Hash, Hasher}; +use std::io::{copy, Read, Seek, SeekFrom, Write}; +use std::time::Duration; +use std::{cmp::Ordering, collections::HashMap}; +use std::{ + thread, + thread::{sleep, JoinHandle}, +}; +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; + +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 +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: HashSet, +} + +impl IFD { + pub fn new() -> Self { + IFD { + tags: HashSet::new(), + } + } + + fn write(&mut self, ijtifffile: &mut IJTiffFile, where_to_write_offset: u64) -> Result { + let mut tags = self.tags.drain().collect::>(); + tags.sort(); + ijtifffile.file.seek(SeekFrom::End(0))?; + if ijtifffile.file.stream_position()? % 2 == 1 { + ijtifffile.file.write(&[0])?; + } + let offset = ijtifffile.file.stream_position()?; + ijtifffile.file.write(&(tags.len() as u64).to_le_bytes())?; + + for tag in tags.iter_mut() { + tag.write_tag(ijtifffile)?; + } + let where_to_write_next_ifd_offset = ijtifffile.file.stream_position()?; + ijtifffile.file.write(&vec![0u8; OFFSET_SIZE])?; + for tag in tags.iter() { + tag.write_data(ijtifffile)?; + } + ijtifffile + .file + .seek(SeekFrom::Start(where_to_write_offset))?; + ijtifffile.file.write(&offset.to_le_bytes())?; + Ok(where_to_write_next_ifd_offset) + } +} + +#[derive(Clone, Debug, Eq)] +pub struct Tag { + code: u16, + bytes: Vec, + ttype: u16, + offset: u64, +} + +impl PartialOrd for Tag { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for Tag { + fn cmp(&self, other: &Self) -> Ordering { + self.code.cmp(&other.code) + } +} + +impl PartialEq for Tag { + fn eq(&self, other: &Self) -> bool { + self.code == other.code + } +} + +impl Hash for Tag { + fn hash(&self, state: &mut H) { + self.code.hash(state); + } +} + +impl Tag { + pub fn new(code: u16, bytes: Vec, ttype: u16) -> Self { + Tag { + code, + bytes, + ttype, + offset: 0, + } + } + + pub fn byte(code: u16, value: &Vec) -> Self { + Tag::new(code, value.to_owned(), 1) + } + + pub fn ascii(code: u16, value: &str) -> Self { + let mut bytes = value.as_bytes().to_vec(); + bytes.push(0); + Tag::new(code, bytes, 2) + } + + pub fn short(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 3, + ) + } + + pub fn long(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value.into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 4, + ) + } + + pub fn rational(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value + .into_iter() + .map(|x| { + u32::try_from(*x.denom()) + .unwrap() + .to_le_bytes() + .into_iter() + .chain(u32::try_from(*x.numer()).unwrap().to_le_bytes()) + .collect::>() + }) + .flatten() + .collect(), + 5, + ) + } + + pub fn sbyte(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value.iter().map(|x| x.to_le_bytes()).flatten().collect(), + 6, + ) + } + + pub fn sshort(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 8, + ) + } + + pub fn slong(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 9, + ) + } + + pub fn srational(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value + .into_iter() + .map(|x| { + i32::try_from(*x.denom()) + .unwrap() + .to_le_bytes() + .into_iter() + .chain(i32::try_from(*x.numer()).unwrap().to_le_bytes()) + .collect::>() + }) + .flatten() + .collect(), + 10, + ) + } + + pub fn float(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 11, + ) + } + + pub fn double(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 12, + ) + } + + pub fn ifd(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), + 13, + ) + } + + pub fn unicode(code: u16, value: &str) -> Self { + let mut bytes: Vec = value + .encode_utf16() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(); + bytes.push(0); + Tag::new(code, bytes, 14) + } + + pub fn complex(code: u16, value: &Vec>) -> Self { + Tag::new( + code, + value + .into_iter() + .map(|x| { + x.re.to_le_bytes() + .into_iter() + .chain(x.im.to_le_bytes()) + .collect::>() + }) + .flatten() + .collect(), + 15, + ) + } + + pub fn long8(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 16, + ) + } + + pub fn slong8(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 17, + ) + } + + pub fn ifd8(code: u16, value: &Vec) -> Self { + Tag::new( + code, + value + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 18, + ) + } + + pub fn short_long_or_long8(code: u16, value: &Vec) -> Self { + let m = *value.iter().max().unwrap(); + if m < 65536 { + Tag::short(code, &value.into_iter().map(|x| *x as u16).collect()) + } else if m < 4294967296 { + Tag::long(code, &value.into_iter().map(|x| *x as u32).collect()) + } else { + Tag::long8(code, value) + } + } + + pub fn count(&self) -> u64 { + let c = match self.ttype { + 1 => self.bytes.len(), // BYTE + 2 => self.bytes.len(), // ASCII + 3 => self.bytes.len() / 2, // SHORT + 4 => self.bytes.len() / 4, // LONG + 5 => self.bytes.len() / 8, // RATIONAL + 6 => self.bytes.len(), // SBYTE + 7 => self.bytes.len(), // UNDEFINED + 8 => self.bytes.len() / 2, // SSHORT + 9 => self.bytes.len() / 4, // SLONG + 10 => self.bytes.len() / 8, // SRATIONAL + 11 => self.bytes.len() / 4, // FLOAT + 12 => self.bytes.len() / 8, // DOUBLE + 13 => self.bytes.len() / 4, // IFD + 14 => self.bytes.len() / 2, // UNICODE + 15 => self.bytes.len() / 8, // COMPLEX + 16 => self.bytes.len() / 8, // LONG8 + 17 => self.bytes.len() / 8, // SLONG8 + 18 => self.bytes.len() / 8, // IFD8 + _ => self.bytes.len(), + }; + c as u64 + } + + fn write_tag(&mut self, ijtifffile: &mut IJTiffFile) -> Result<()> { + self.offset = ijtifffile.file.stream_position()?; + ijtifffile.file.write(&self.code.to_le_bytes())?; + ijtifffile.file.write(&self.ttype.to_le_bytes())?; + ijtifffile.file.write(&self.count().to_le_bytes())?; + if self.bytes.len() <= OFFSET_SIZE { + ijtifffile.file.write(&self.bytes)?; + for _ in self.bytes.len()..OFFSET_SIZE { + ijtifffile.file.write(&[0])?; + } + } else { + ijtifffile.file.write(&vec![0u8; OFFSET_SIZE])?; + } + Ok(()) + } + + fn write_data(&self, ijtifffile: &mut IJTiffFile) -> Result<()> { + if self.bytes.len() > OFFSET_SIZE { + ijtifffile.file.seek(SeekFrom::End(0))?; + let offset = ijtifffile.write(&self.bytes)?; + ijtifffile.file.seek(SeekFrom::Start( + self.offset + (TAG_SIZE - OFFSET_SIZE) as u64, + ))?; + ijtifffile.file.write(&offset.to_le_bytes())?; + if ijtifffile.file.stream_position()? % 2 == 1 { + ijtifffile.file.write(&[0u8])?; + } + } + Ok(()) + } +} + +#[derive(Debug)] +struct CompressedFrame { + bytes: Vec>, + image_width: u32, + image_length: u32, + tile_size: usize, + bits_per_sample: u16, + sample_format: u16, +} + +#[derive(Clone, Debug)] +struct Frame { + offsets: Vec, + bytecounts: Vec, + image_width: u32, + image_length: u32, + bits_per_sample: u16, + sample_format: u16, + tile_width: u16, + tile_length: u16, +} + +impl Frame { + fn new( + offsets: Vec, + bytecounts: Vec, + image_width: u32, + image_length: u32, + bits_per_sample: u16, + sample_format: u16, + tile_width: u16, + tile_length: u16, + ) -> Self { + Frame { + offsets, + bytecounts, + image_width, + image_length, + bits_per_sample, + sample_format, + tile_width, + tile_length, + } + } +} + +pub trait Bytes { + const BITS_PER_SAMPLE: u16; + const SAMPLE_FORMAT: u16; + + fn bytes(&self) -> Vec; +} + +macro_rules! bytes_impl { + ($T:ty, $bits_per_sample:expr, $sample_format:expr) => { + impl Bytes for $T { + const BITS_PER_SAMPLE: u16 = $bits_per_sample; + const SAMPLE_FORMAT: u16 = $sample_format; + + #[inline] + fn bytes(&self) -> Vec { + self.to_le_bytes().to_vec() + } + } + }; +} + +bytes_impl!(u8, 8, 1); +bytes_impl!(u16, 16, 1); +bytes_impl!(u32, 32, 1); +bytes_impl!(u64, 64, 1); +bytes_impl!(u128, 128, 1); +#[cfg(target_pointer_width = "64")] +bytes_impl!(usize, 64, 1); +#[cfg(target_pointer_width = "32")] +bytes_impl!(usize, 32, 1); +bytes_impl!(i8, 8, 2); +bytes_impl!(i16, 16, 2); +bytes_impl!(i32, 32, 2); +bytes_impl!(i64, 64, 2); +bytes_impl!(i128, 128, 2); +#[cfg(target_pointer_width = "64")] +bytes_impl!(isize, 64, 2); +#[cfg(target_pointer_width = "32")] +bytes_impl!(isize, 32, 2); +bytes_impl!(f32, 32, 3); +bytes_impl!(f64, 64, 3); + +#[derive(Clone, Debug)] +pub enum Colors { + None, + Colors(Vec>), + Colormap(Vec>), +} + +#[derive(Debug)] +pub struct IJTiffFile { + file: File, + frames: HashMap<(usize, usize, usize), Frame>, + hashes: HashMap, + threads: HashMap<(usize, usize, usize), JoinHandle>, + pub compression_level: i32, + pub colors: Colors, + pub comment: Option, + pub px_size: Option, + pub delta_z: Option, + pub time_interval: Option, + pub extra_tags: HashMap, Vec>, +} + +impl Drop for IJTiffFile { + fn drop(&mut self) { + if let Err(e) = self.close() { + println!("Error closing IJTiffFile: {:?}", e); + } + } +} + +impl IJTiffFile { + pub fn new(path: &str) -> Result { + let mut file = OpenOptions::new() + .create(true) + .truncate(true) + .write(true) + .read(true) + .open(path)?; + file.write(b"II")?; + file.write(&43u16.to_le_bytes())?; + file.write(&8u16.to_le_bytes())?; + file.write(&0u16.to_le_bytes())?; + file.write(&OFFSET.to_le_bytes())?; + Ok(IJTiffFile { + file, + frames: HashMap::new(), + hashes: HashMap::new(), + threads: HashMap::new(), + compression_level: DEFAULT_COMPRESSION_LEVEL, + colors: Colors::None, + comment: None, + px_size: None, + delta_z: None, + time_interval: None, + extra_tags: HashMap::new(), + }) + } + + pub fn set_compression_level(&mut self, compression_level: i32) { + self.compression_level = compression_level.max(-7).min(22); + } + + 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={}", c_size); + desc += &format!("\nslices={}", z_size); + desc += &format!("\nframes={}", t_size); + } else { + 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 c_size == 1 { + desc += "\nmode=grayscale"; + } else { + desc += "\nmode=composite"; + } + desc += "\nhyperstack=true\nloop=false\nunit=micron"; + if let Some(delta_z) = self.delta_z { + desc += &format!("\nspacing={}", delta_z); + } + if let Some(timeinterval) = self.time_interval { + desc += &format!("\ninterval={}", timeinterval); + } + if let Some(comment) = &self.comment { + desc += &format!("\ncomment={}", comment); + } + desc + } + + 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 % z_size, + frame_number / z_size, + ) + } else { + ( + frame_number % c_size, + frame_number / c_size % z_size, + frame_number / c_size / z_size, + ) + } + } + + fn spp_and_n_frames(&self, c_size: usize, z_size: usize, t_size: usize) -> (u8, usize) { + if let Colors::None = &self.colors { + (c_size as u8, z_size * t_size) + } else { + (1, c_size * z_size * t_size) + } + } + + fn hash(value: &T) -> u64 { + let mut hasher = DefaultHasher::new(); + value.hash(&mut hasher); + hasher.finish() + } + + fn hash_check(&mut self, bytes: &Vec, offset: u64) -> Result { + let current_offset = self.file.stream_position()?; + self.file.seek(SeekFrom::Start(offset))?; + let mut buffer = vec![0u8; bytes.len()]; + self.file.read_exact(&mut buffer)?; + let same = bytes == &buffer; + self.file.seek(SeekFrom::Start(current_offset))?; + Ok(same) + } + + fn write(&mut self, bytes: &Vec) -> Result { + let hash = IJTiffFile::hash(&bytes); + if self.hashes.contains_key(&hash) + && self.hash_check(&bytes, *self.hashes.get(&hash).unwrap())? + { + Ok(*self.hashes.get(&hash).unwrap()) + } else { + if self.file.stream_position()? % 2 == 1 { + self.file.write(&[0])?; + } + let offset = self.file.stream_position()?; + self.hashes.insert(hash, offset); + self.file.write(&bytes)?; + Ok(offset) + } + } + + pub fn save(&mut self, frame: Array2, c: usize, z: usize, t: usize) -> Result<()> + where + T: Bytes + Clone + Send + Sync + Zero + 'static, + { + self.compress_frame(frame.reversed_axes(), c, z, t)?; + Ok(()) + } + + fn compress_frame(&mut self, frame: Array2, c: usize, z: usize, t: usize) -> Result<()> + where + T: Bytes + Clone + Zero + Send + 'static, + { + fn compress(frame: Array2, compression_level: i32) -> CompressedFrame + where + T: Bytes + Clone + Zero, + { + let image_width = frame.shape()[0] as u32; + let image_length = frame.shape()[1] as u32; + let tile_size = 2usize + .pow( + ((image_width as f64 * image_length as f64 / 2f64).log2() / 2f64).round() + as u32, + ) + .max(16) + .min(1024); + let tiles = IJTiffFile::tile(frame.reversed_axes(), tile_size); + let byte_tiles: Vec> = tiles + .into_iter() + .map(|tile| tile.map(|x| x.bytes()).into_iter().flatten().collect()) + .collect(); + let bytes = if byte_tiles.len() > 4 { + byte_tiles + .into_par_iter() + .map(|x| encode_all(x, compression_level).unwrap()) + .collect::>() + } else { + byte_tiles + .into_iter() + .map(|x| encode_all(x, compression_level).unwrap()) + .collect::>() + }; + CompressedFrame { + bytes, + image_width, + image_length, + tile_size, + bits_per_sample: T::BITS_PER_SAMPLE, + sample_format: T::SAMPLE_FORMAT, + } + } + loop { + self.collect_threads(false)?; + if self.threads.len() < 48 { + break; + } + sleep(Duration::from_millis(100)); + } + let compression_level = self.compression_level; + self.threads.insert( + (c, z, t), + thread::spawn(move || compress(frame, compression_level)), + ); + Ok(()) + } + + fn collect_threads(&mut self, block: bool) -> Result<()> { + for (c, z, t) in self.threads.keys().cloned().collect::>() { + if block || self.threads[&(c, z, t)].is_finished() { + if let Some(thread) = self.threads.remove(&(c, z, t)) { + self.write_frame(thread.join().unwrap(), c, z, t)?; + } + } + } + Ok(()) + } + + fn write_frame(&mut self, frame: CompressedFrame, c: usize, z: usize, t: usize) -> Result<()> { + let mut offsets = Vec::new(); + let mut bytecounts = Vec::new(); + for tile in frame.bytes { + bytecounts.push(tile.len() as u64); + offsets.push(self.write(&tile)?); + } + let frame = Frame::new( + offsets, + bytecounts, + frame.image_width, + frame.image_length, + frame.bits_per_sample, + frame.sample_format, + frame.tile_size as u16, + frame.tile_size as u16, + ); + self.frames.insert((c, z, t), frame); + Ok(()) + } + + fn tile(frame: Array2, size: usize) -> Vec> { + let shape = frame.shape(); + 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( + frame + .slice(s![i * size..(i + 1) * size, j * size..(j + 1) * size]) + .to_owned(), + ); + } + if shape[1] % size != 0 { + let mut tile = Array2::::zeros((size, size)); + tile.slice_mut(s![.., ..shape[1] - m * size]) + .assign(&frame.slice(s![i * size..(i + 1) * size, m * size..])); + tiles.push(tile); + } + } + if shape[0] % size != 0 { + for j in 0..m { + let mut tile = Array2::::zeros((size, size)); + tile.slice_mut(s![..shape[0] - n * size, ..]) + .assign(&frame.slice(s![n * size.., j * size..(j + 1) * size])); + tiles.push(tile); + } + if shape[1] % size != 0 { + let mut tile = Array2::::zeros((size, size)); + tile.slice_mut(s![..shape[0] - n * size, ..shape[1] - m * size]) + .assign(&frame.slice(s![n * size.., m * size..])); + tiles.push(tile); + } + } + tiles + } + + fn get_colormap(&self, colormap: &Vec>, bits_per_sample: u16) -> Vec { + let mut r = Vec::new(); + let mut g = Vec::new(); + let mut b = Vec::new(); + let n = 2usize.pow(bits_per_sample as u32 - 8); + for color in colormap { + r.extend(vec![(color[0] as u16) * 257; n]); + g.extend(vec![(color[1] as u16) * 257; n]); + b.extend(vec![(color[2] as u16) * 257; n]); + } + r.extend(g); + r.extend(b); + r + } + + fn get_color(&self, colors: &Vec, bits_per_sample: u16) -> Vec { + let mut c = Vec::new(); + let n = 2usize.pow(bits_per_sample as u32 - 8); + for color in colors { + for i in 0..256 { + c.extend(vec![i * (*color as u16) / 255 * 257; n]) + } + } + c + } + + fn close(&mut self) -> Result<()> { + self.collect_threads(true)?; + 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 = 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 + .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, c_size, z_size)) + { + offsets.extend(frame_n.offsets.iter()); + bytecounts.extend(frame_n.bytecounts.iter()); + frame_count += 1; + } else { + warn.push((frame_number, channel)); + } + } + let mut ifd = IFD::new(); + ifd.tags.insert(Tag::long(256, &vec![frame.image_width])); + ifd.tags.insert(Tag::long(257, &vec![frame.image_length])); + ifd.tags + .insert(Tag::short(258, &vec![frame.bits_per_sample; frame_count])); + ifd.tags.insert(Tag::short(259, &vec![COMPRESSION])); + ifd.tags + .insert(Tag::ascii(270, &self.description(c_size, z_size, t_size))); + ifd.tags.insert(Tag::short(277, &vec![frame_count as u16])); + ifd.tags.insert(Tag::ascii(305, "tiffwrite_rs")); + ifd.tags.insert(Tag::short(322, &vec![frame.tile_width])); + ifd.tags.insert(Tag::short(323, &vec![frame.tile_length])); + ifd.tags.insert(Tag::short_long_or_long8(324, &offsets)); + ifd.tags.insert(Tag::short_long_or_long8(325, &bytecounts)); + if frame.sample_format > 1 { + ifd.tags.insert(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.tags.insert(Tag::rational(282, &r)); + ifd.tags.insert(Tag::rational(283, &r)); + } + if let Colors::Colormap(_) = &self.colors { + ifd.tags.insert(Tag::short(262, &vec![3])); + } else if let Colors::None = self.colors { + ifd.tags.insert(Tag::short(262, &vec![1])); + } + if frame_number == 0 { + if let Colors::Colormap(colormap) = &self.colors { + ifd.tags.insert(Tag::short( + 320, + &self.get_colormap(colormap, frame.bits_per_sample), + )); + } + } + if frame_number < c_size { + if let Colors::Colors(colors) = &self.colors { + ifd.tags.insert(Tag::short( + 320, + &self.get_color(&colors[frame_number], frame.bits_per_sample), + )); + ifd.tags.insert(Tag::short(262, &vec![3])); + } + } + if let Colors::None = &self.colors { + if c_size > 1 { + ifd.tags.insert(Tag::short(284, &vec![2])); + } + } + for channel in 0..samples_per_pixel { + 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.tags.insert(tag.to_owned()); + } + } + } + if let Some(extra_tags) = self.extra_tags.get(&None) { + for tag in extra_tags { + ifd.tags.insert(tag.to_owned()); + } + } + if frame_number == 0 { + ifd.tags.insert(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.push((frame_number, 0)); + } + if warn.len() > 0 { + println!("The following frames were not added to the tif file"); + for (frame_number, channel) in &warn { + 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 + .seek(SeekFrom::Start(where_to_write_next_ifd_offset))?; + self.file.write(&0u64.to_le_bytes())?; + Ok(()) + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..449463d --- /dev/null +++ b/src/main.rs @@ -0,0 +1,23 @@ +use anyhow::Result; +use ndarray::{s, Array2}; +use tiffwrite::IJTiffFile; + +fn main() -> Result<()> { + println!("Hello World!"); + 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] { + arr[[i, j]] = i as u16; + } + } + f.save(arr.to_owned(), 0, 0, 0)?; + + let mut arr = Array2::::zeros((100, 100)); + arr.slice_mut(s![64.., ..64]).fill(1); + arr.slice_mut(s![..64, 64..]).fill(2); + arr.slice_mut(s![64.., 64..]).fill(3); + f.save(arr.to_owned(), 1, 0, 0)?; + Ok(()) +} diff --git a/src/py.rs b/src/py.rs new file mode 100644 index 0000000..ce567f2 --- /dev/null +++ b/src/py.rs @@ -0,0 +1,360 @@ +use crate::{Colors, IJTiffFile, Tag}; +use ndarray::s; +use num::{Complex, FromPrimitive, Rational32}; +use numpy::{PyArrayMethods, PyReadonlyArray2}; +use pyo3::prelude::*; + +#[pyclass(subclass)] +#[pyo3(name = "Tag")] +#[derive(Clone, Debug)] +struct PyTag { + tag: Tag, +} + +#[pymethods] +impl PyTag { + #[staticmethod] + fn byte(code: u16, byte: Vec) -> Self { + PyTag { + tag: Tag::byte(code, &byte), + } + } + + #[staticmethod] + fn ascii(code: u16, ascii: &str) -> Self { + PyTag { + tag: Tag::ascii(code, ascii), + } + } + + #[staticmethod] + fn short(code: u16, short: Vec) -> Self { + PyTag { + tag: Tag::short(code, &short), + } + } + + #[staticmethod] + fn long(code: u16, long: Vec) -> Self { + PyTag { + tag: Tag::long(code, &long), + } + } + + #[staticmethod] + fn rational(code: u16, rational: Vec) -> Self { + PyTag { + tag: Tag::rational( + code, + &rational + .into_iter() + .map(|x| Rational32::from_f64(x).unwrap()) + .collect(), + ), + } + } + + #[staticmethod] + fn sbyte(code: u16, sbyte: Vec) -> Self { + PyTag { + tag: Tag::sbyte(code, &sbyte), + } + } + + #[staticmethod] + fn sshort(code: u16, sshort: Vec) -> Self { + PyTag { + tag: Tag::sshort(code, &sshort), + } + } + + #[staticmethod] + fn slong(code: u16, slong: Vec) -> Self { + PyTag { + tag: Tag::slong(code, &slong), + } + } + + #[staticmethod] + fn srational(code: u16, srational: Vec) -> Self { + PyTag { + tag: Tag::srational( + code, + &srational + .into_iter() + .map(|x| Rational32::from_f64(x).unwrap()) + .collect(), + ), + } + } + + #[staticmethod] + fn float(code: u16, float: Vec) -> Self { + PyTag { + tag: Tag::float(code, &float), + } + } + + #[staticmethod] + fn double(code: u16, double: Vec) -> Self { + PyTag { + tag: Tag::double(code, &double), + } + } + + #[staticmethod] + fn ifd(code: u16, ifd: Vec) -> Self { + PyTag { + tag: Tag::ifd(code, &ifd), + } + } + + #[staticmethod] + fn unicode(code: u16, unicode: &str) -> Self { + PyTag { + tag: Tag::unicode(code, unicode), + } + } + + #[staticmethod] + fn complex(code: u16, complex: Vec<(f32, f32)>) -> Self { + PyTag { + tag: Tag::complex( + code, + &complex + .into_iter() + .map(|(x, y)| Complex { re: x, im: y }) + .collect(), + ), + } + } + + #[staticmethod] + fn long8(code: u16, long8: Vec) -> Self { + PyTag { + tag: Tag::long8(code, &long8), + } + } + + #[staticmethod] + fn slong8(code: u16, slong8: Vec) -> Self { + PyTag { + tag: Tag::slong8(code, &slong8), + } + } + + #[staticmethod] + fn ifd8(code: u16, ifd8: Vec) -> Self { + PyTag { + tag: Tag::ifd8(code, &ifd8), + } + } + + fn count(&self) -> u64 { + self.tag.count() + } +} + +#[pyclass(subclass)] +#[pyo3(name = "IJTiffFile")] +#[derive(Debug)] +struct PyIJTiffFile { + ijtifffile: Option, +} + +#[pymethods] +impl PyIJTiffFile { + #[new] + fn new(path: &str) -> PyResult { + Ok(PyIJTiffFile { + 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 { + if let Colors::Colors(colors) = &ijtifffile.colors { + return Ok(Some(colors.to_owned())); + } + } + Ok(None) + } + + #[setter] + fn set_colors(&mut self, colors: PyReadonlyArray2) -> PyResult<()> { + if let Some(ijtifffile) = &mut self.ijtifffile { + let a = colors.to_owned_array(); + ijtifffile.colors = Colors::Colors( + (0..a.shape()[0]) + .map(|i| Vec::from(a.slice(s![i, ..]).as_slice().unwrap())) + .collect(), + ); + } + Ok(()) + } + + #[getter] + fn get_colormap(&mut self) -> PyResult>>> { + if let Some(ijtifffile) = &self.ijtifffile { + if let Colors::Colormap(colormap) = &ijtifffile.colors { + return Ok(Some(colormap.to_owned())); + } + } + Ok(None) + } + + #[setter] + fn set_colormap(&mut self, colormap: PyReadonlyArray2) -> PyResult<()> { + if let Some(ijtifffile) = &mut self.ijtifffile { + let a = colormap.to_owned_array(); + ijtifffile.colors = Colors::Colormap( + (0..a.shape()[0]) + .map(|i| Vec::from(a.slice(s![i, ..]).as_slice().unwrap())) + .collect(), + ); + } + Ok(()) + } + + #[getter] + fn get_px_size(&self) -> PyResult> { + if let Some(ijtifffile) = &self.ijtifffile { + Ok(ijtifffile.px_size) + } else { + Ok(None) + } + } + + #[setter] + fn set_px_size(&mut self, px_size: f64) -> PyResult<()> { + if let Some(ijtifffile) = &mut self.ijtifffile { + ijtifffile.px_size = Some(px_size); + } + Ok(()) + } + + #[getter] + fn get_delta_z(&self) -> PyResult> { + if let Some(ijtifffile) = &self.ijtifffile { + Ok(ijtifffile.delta_z) + } else { + Ok(None) + } + } + + #[setter] + fn set_delta_z(&mut self, delta_z: f64) -> PyResult<()> { + if let Some(ijtifffile) = &mut self.ijtifffile { + ijtifffile.delta_z = Some(delta_z); + } + Ok(()) + } + + #[getter] + fn get_time_interval(&self) -> PyResult> { + if let Some(ijtifffile) = &self.ijtifffile { + Ok(ijtifffile.time_interval) + } else { + Ok(None) + } + } + + #[setter] + fn set_time_interval(&mut self, time_interval: f64) -> PyResult<()> { + if let Some(ijtifffile) = &mut self.ijtifffile { + ijtifffile.time_interval = Some(time_interval); + } + Ok(()) + } + + #[getter] + fn get_comment(&self) -> PyResult> { + if let Some(ijtifffile) = &self.ijtifffile { + Ok(ijtifffile.comment.clone()) + } else { + Ok(None) + } + } + + #[setter] + fn set_comment(&mut self, comment: &str) -> PyResult<()> { + if let Some(ijtifffile) = &mut self.ijtifffile { + ijtifffile.comment = Some(String::from(comment)); + } + Ok(()) + } + + fn append_extra_tag(&mut self, tag: PyTag, czt: Option<(usize, usize, usize)>) { + if let Some(ijtifffile) = self.ijtifffile.as_mut() { + if let Some(extra_tags) = ijtifffile.extra_tags.get_mut(&czt) { + extra_tags.push(tag.tag) + } + } + } + + fn get_tags(&self, czt: Option<(usize, usize, usize)>) -> PyResult> { + if let Some(ijtifffile) = &self.ijtifffile { + if let Some(extra_tags) = ijtifffile.extra_tags.get(&czt) { + let v = extra_tags + .iter() + .map(|tag| PyTag { + tag: tag.to_owned(), + }) + .collect(); + return Ok(v); + } + } + Ok(Vec::new()) + } + + fn close(&mut self) -> PyResult<()> { + self.ijtifffile.take(); + Ok(()) + } +} + +macro_rules! impl_save { + ($T:ty, $t:ident) => { + #[pymethods] + impl PyIJTiffFile { + fn $t( + &mut self, + frame: PyReadonlyArray2<$T>, + c: usize, + t: usize, + z: usize, + ) -> PyResult<()> { + if let Some(ijtifffile) = self.ijtifffile.as_mut() { + ijtifffile.save(frame.to_owned_array(), c, t, z)?; + } + Ok(()) + } + } + }; +} + +impl_save!(u8, save_u8); +impl_save!(u16, save_u16); +impl_save!(u32, save_u32); +impl_save!(u64, save_u64); +impl_save!(i8, save_i8); +impl_save!(i16, save_i16); +impl_save!(i32, save_i32); +impl_save!(i64, save_i64); +impl_save!(f32, save_f32); +impl_save!(f64, save_f64); + +#[pymodule] +#[pyo3(name = "tiffwrite_rs")] +fn tiffwrite_rs(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + Ok(()) +} 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..0d88e69 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, pxsize=0.1, deltaz=0.5, timeinterval=6.5) 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" diff --git a/tiffwrite/__init__.py b/tiffwrite/__init__.py deleted file mode 100644 index e4c0d1e..0000000 --- a/tiffwrite/__init__.py +++ /dev/null @@ -1,603 +0,0 @@ -from __future__ import annotations - -import struct -import warnings -from collections.abc import Iterable -from contextlib import contextmanager -from datetime import datetime -from fractions import Fraction -from functools import cached_property -from hashlib import sha1 -from importlib.metadata import version -from io import BytesIO -from itertools import product -from pathlib import Path -from typing import Any, BinaryIO, Callable, Generator, Literal, Optional, Sequence - -import colorcet -import numpy as np -import tifffile -from matplotlib import colors as mpl_colors -from numpy.typing import DTypeLike -from parfor import ParPool, PoolSingleton -from tqdm.auto import tqdm - -__all__ = ["IJTiffFile", "Tag", "tiffwrite"] - - -try: - __version__ = version("tiffwrite") -except Exception: # noqa - __version__ = "unknown" - - -Strip = tuple[list[int], list[int]] -CZT = tuple[int, int, int] - - -def tiffwrite(file: str | Path, data: np.ndarray, axes: str = 'TZCXY', dtype: DTypeLike = None, bar: bool = False, - *args: Any, **kwargs: Any) -> None: - """ file: string; filename of the new tiff file - data: 2 to 5D numpy array - axes: string; order of dimensions in data, default: TZCXY for 5D, ZCXY for 4D, CXY for 3D, XY for 2D data - dtype: string; datatype to use when saving to tiff - bar: bool; whether to show a progress bar - other args: see IJTiffFile - """ - - axes = axes[-np.ndim(data):].upper() - if not axes == 'CZTXY': - axes_shuffle = [axes.find(i) for i in 'CZTXY'] - axes_add = [i for i, j in enumerate(axes_shuffle) if j < 0] - axes_shuffle = [i for i in axes_shuffle if i >= 0] - data = np.transpose(data, axes_shuffle) - for axis in axes_add: - data = np.expand_dims(data, axis) - - shape = data.shape[:3] - with IJTiffFile(file, shape, data.dtype if dtype is None else dtype, *args, **kwargs) as f: # type: ignore - at_least_one = False - for n in tqdm(product(*[range(i) for i in shape]), total=np.prod(shape), desc='Saving tiff', disable=not bar): - if np.any(data[n]) or not at_least_one: - f.save(data[n], *n) - at_least_one = True - - -class Header: - def __init__(self, filehandle_or_byteorder: BinaryIO | Literal['>', '<'] | None = None, - bigtiff: bool = True) -> None: - if filehandle_or_byteorder is None or isinstance(filehandle_or_byteorder, str): - self.byteorder = filehandle_or_byteorder or '<' - self.bigtiff = bigtiff - if self.bigtiff: - self.tagsize = 20 - self.tagnoformat = 'Q' - self.offsetsize = 8 - self.offsetformat = 'Q' - self.offset = 16 - else: - self.tagsize = 12 - self.tagnoformat = 'H' - self.offsetsize = 4 - self.offsetformat = 'I' - self.offset = 8 - else: - fh = filehandle_or_byteorder - fh.seek(0) - self.byteorder = '>' if fh.read(2) == b'MM' else '<' - self.bigtiff = {42: False, 43: True}[struct.unpack(self.byteorder + 'H', fh.read(2))[0]] - if self.bigtiff: - self.tagsize = 20 - self.tagnoformat = 'Q' - self.offsetsize = struct.unpack(self.byteorder + 'H', fh.read(2))[0] - self.offsetformat = {8: 'Q', 16: '2Q'}[self.offsetsize] - assert struct.unpack(self.byteorder + 'H', fh.read(2))[0] == 0, 'Not a TIFF-file' - self.offset = struct.unpack(self.byteorder + self.offsetformat, fh.read(self.offsetsize))[0] - else: - self.tagsize = 12 - self.tagnoformat = 'H' - self.offsetformat = 'I' - self.offsetsize = 4 - self.offset = struct.unpack(self.byteorder + self.offsetformat, fh.read(self.offsetsize))[0] - - def write(self, fh: BinaryIO) -> None: - fh.write({'<': b'II', '>': b'MM'}[self.byteorder]) - if self.bigtiff: - fh.write(struct.pack(self.byteorder + 'H', 43)) - fh.write(struct.pack(self.byteorder + 'H', 8)) - fh.write(struct.pack(self.byteorder + 'H', 0)) - fh.write(struct.pack(self.byteorder + 'Q', self.offset)) - else: - fh.write(struct.pack(self.byteorder + 'H', 42)) - fh.write(struct.pack(self.byteorder + 'I', self.offset)) - - -class Tag: - Value = bytes | str | float | Fraction | Sequence[bytes | str | float | Fraction] - tiff_tag_registry = tifffile.TiffTagRegistry({key: value.lower() for key, value in tifffile.TIFF.TAGS.items()}) - - @staticmethod - def from_dict(tags: dict[str | int, Value | Tag]) -> dict[int, Tag]: - return {(key if isinstance(key, int) - else (int(key[3:]) if key.lower().startswith('tag') - else Tag.tiff_tag_registry[key.lower()])): tag if isinstance(tag, Tag) else Tag(tag) - for key, tag in tags.items()} - - @staticmethod - def fraction(numerator: float = 0, denominator: int = None) -> Fraction: - return Fraction(numerator, denominator).limit_denominator( # type: ignore - 2 ** (31 if numerator < 0 or (denominator is not None and denominator < 0) else 32) - 1) - - def __init__(self, ttype_or_value: str | Value, value: Value = None, - offset: int = None) -> None: - self._value: bytes | str | Sequence[bytes | str | float | Fraction] - self.fh: Optional[BinaryIO] = None - self.header: Optional[Header] = None - self.bytes_data: Optional[bytes] = None - if value is None: - self.value = ttype_or_value # type: ignore - if isinstance(self.value, (str, bytes)) or all([isinstance(value, (str, bytes)) for value in self.value]): - ttype = 'ascii' - elif all([isinstance(value, int) for value in self.value]): - min_value: int = np.min(self.value) # type: ignore - max_value: int = np.max(self.value) # type: ignore - type_map = {'uint8': 'byte', 'int8': 'sbyte', 'uint16': 'short', 'int16': 'sshort', - 'uint32': 'long', 'int32': 'slong', 'uint64': 'long8', 'int64': 'slong8'} - for dtype, ttype in type_map.items(): - if np.iinfo(dtype).min <= min_value and max_value <= np.iinfo(dtype).max: - break - else: - ttype = 'undefined' - elif all([isinstance(value, Fraction) for value in self.value]): - if all([value.numerator < 0 or value.denominator < 0 for value in self.value]): # type: ignore - ttype = 'srational' - else: - ttype = 'rational' - elif all([isinstance(value, (float, int)) for value in self.value]): - min_value = np.min(np.asarray(self.value)[np.isfinite(self.value)]) # type: ignore - max_value = np.max(np.asarray(self.value)[np.isfinite(self.value)]) # type: ignore - type_map = {'float32': 'float', 'float64': 'double'} - for dtype, ttype in type_map.items(): - if np.finfo(dtype).min <= min_value and max_value <= np.finfo(dtype).max: - break - else: - ttype = 'undefined' - elif all([isinstance(value, complex) for value in self.value]): - ttype = 'complex' - else: - ttype = 'undefined' - self.ttype = tifffile.TIFF.DATATYPES[ttype.upper()] # noqa - else: - self.value = value # type: ignore - self.ttype = tifffile.TIFF.DATATYPES[ttype_or_value.upper()] if isinstance(ttype_or_value, str) \ - else ttype_or_value # type: ignore - self.dtype = tifffile.TIFF.DATA_FORMATS[self.ttype] - self.offset = offset - self.type_check() - - @property - def value(self) -> bytes | str | Sequence[bytes | str | float | Fraction]: - return self._value - - @value.setter - def value(self, value: Value) -> None: - self._value = value if isinstance(value, Iterable) else (value,) - - def __repr__(self) -> str: - if self.offset is None: - return f'{tifffile.TIFF.DATATYPES(self.ttype).name}: {self.value!r}' - else: - return f'{tifffile.TIFF.DATATYPES(self.ttype).name} @ {self.offset}: {self.value!r}' - - def type_check(self) -> None: - try: - self.bytes_and_count(Header()) - except Exception: - raise ValueError(f"tif tag type '{tifffile.TIFF.DATATYPES(self.ttype).name}' and " - f"data type '{type(self.value[0]).__name__}' do not correspond") - - def bytes_and_count(self, header: Header) -> tuple[bytes, int]: - if isinstance(self.value, bytes): - return self.value, len(self.value) // struct.calcsize(self.dtype) - elif self.ttype in (2, 14): - if isinstance(self.value, str): - bytes_value = self.value.encode('ascii') + b'\x00' - else: - bytes_value = b'\x00'.join([value.encode('ascii') for value in self.value]) + b'\x00' # type: ignore - return bytes_value, len(bytes_value) - elif self.ttype in (5, 10): - return b''.join([struct.pack(header.byteorder + self.dtype, # type: ignore - *((value.denominator, value.numerator) if isinstance(value, Fraction) - else value)) for value in self.value]), len(self.value) - else: - return b''.join([struct.pack(header.byteorder + self.dtype, value) for value in self.value]), \ - len(self.value) - - def write_tag(self, fh: BinaryIO, key: int, header: Header, offset: int = None) -> None: - self.fh = fh - self.header = header - if offset is None: - self.offset = fh.tell() - else: - fh.seek(offset) - self.offset = offset - fh.write(struct.pack(header.byteorder + 'HH', key, self.ttype)) - bytes_tag, count = self.bytes_and_count(header) - fh.write(struct.pack(header.byteorder + header.offsetformat, count)) - len_bytes = len(bytes_tag) - if len_bytes <= header.offsetsize: - fh.write(bytes_tag) - self.bytes_data = None - empty_bytes = header.offsetsize - len_bytes - else: - self.bytes_data = bytes_tag - empty_bytes = header.offsetsize - if empty_bytes: - fh.write(empty_bytes * b'\x00') - - def write_data(self, write: Callable[[BinaryIO, bytes], None] = None) -> None: - if self.bytes_data and self.fh is not None and self.header is not None and self.offset is not None: - self.fh.seek(0, 2) - if write is None: - offset = self.write(self.bytes_data) - else: - offset = write(self.fh, self.bytes_data) - self.fh.seek(self.offset + self.header.tagsize - self.header.offsetsize) - self.fh.write(struct.pack(self.header.byteorder + self.header.offsetformat, offset)) - - def write(self, bytes_value: bytes) -> Optional[int]: - if self.fh is not None: - if self.fh.tell() % 2: - self.fh.write(b'\x00') - offset = self.fh.tell() - self.fh.write(bytes_value) - return offset - - def copy(self) -> Tag: - return self.__class__(self.ttype, self.value[:], self.offset) - - -class IFD(dict): - def __init__(self, fh: BinaryIO = None) -> None: - super().__init__() - self.fh = fh - self.header: Optional[Header] = None - self.offset: Optional[int] = None - self.where_to_write_next_ifd_offset: Optional[int] = None - if fh is not None: - header = Header(fh) - fh.seek(header.offset) - n_tags = struct.unpack(header.byteorder + header.tagnoformat, - fh.read(struct.calcsize(header.tagnoformat)))[0] - assert n_tags < 4096, 'Too many tags' - addr = [] - addroffset = [] - - length = 8 if header.bigtiff else 2 - length += n_tags * header.tagsize + header.offsetsize - - for i in range(n_tags): - pos = header.offset + struct.calcsize(header.tagnoformat) + header.tagsize * i - fh.seek(pos) - - code, ttype = struct.unpack(header.byteorder + 'HH', fh.read(4)) - count = struct.unpack(header.byteorder + header.offsetformat, fh.read(header.offsetsize))[0] - - dtype = tifffile.TIFF.DATA_FORMATS[ttype] - dtypelen = struct.calcsize(dtype) - - toolong = struct.calcsize(dtype) * count > header.offsetsize - if toolong: - addr.append(fh.tell() - header.offset) - caddr = struct.unpack(header.byteorder + header.offsetformat, fh.read(header.offsetsize))[0] - addroffset.append(caddr - header.offset) - cp = fh.tell() - fh.seek(caddr) - - if ttype == 1: - value: Tag.Value = fh.read(count) - elif ttype == 2: - value = fh.read(count).decode('ascii').rstrip('\x00') - elif ttype in (5, 10): - value = [struct.unpack(header.byteorder + dtype, fh.read(dtypelen)) # type: ignore - for _ in range(count)] - else: - value = [struct.unpack(header.byteorder + dtype, fh.read(dtypelen))[0] for _ in range(count)] - - if toolong: - fh.seek(cp) # noqa - - self[code] = Tag(ttype, value, pos) - fh.seek(header.offset) - - def __setitem__(self, key: str | int, tag: str | float | Fraction | Tag) -> None: - super().__setitem__(Tag.tiff_tag_registry[key.lower()] if isinstance(key, str) else key, - tag if isinstance(tag, Tag) else Tag(tag)) - - def items(self) -> Generator[tuple[int, Tag], None, None]: # type: ignore[override] - return ((key, self[key]) for key in sorted(self)) - - def keys(self) -> Generator[int, None, None]: # type: ignore[override] - return (key for key in sorted(self)) - - def values(self) -> Generator[Tag, None, None]: # type: ignore[override] - return (self[key] for key in sorted(self)) - - def write(self, fh: BinaryIO, header: Header, write: Callable[[BinaryIO, bytes], None] = None) -> BinaryIO: - self.fh = fh - self.header = header - if fh.seek(0, 2) % 2: - fh.write(b'\x00') - self.offset = fh.tell() - fh.write(struct.pack(header.byteorder + header.tagnoformat, len(self))) - for key, tag in self.items(): - tag.write_tag(fh, key, header) - self.where_to_write_next_ifd_offset = fh.tell() - fh.write(b'\x00' * header.offsetsize) - for tag in self.values(): - tag.write_data(write) - return fh - - def write_offset(self, where_to_write_offset: int) -> None: - if self.fh is not None and self.header is not None: - self.fh.seek(where_to_write_offset) - self.fh.write(struct.pack(self.header.byteorder + self.header.offsetformat, self.offset)) - - def copy(self) -> IFD: - new = self.__class__() - new.update({key: tag.copy() for key, tag in self.items()}) - return new - - -FrameInfo = tuple[IFD, Strip, CZT] - - -class IJTiffFile: - """ Writes a tiff file in a format that the BioFormats reader in Fiji understands. - file: filename of the new tiff file - shape: shape (CZT) of the data to be written - dtype: datatype to use when saving to tiff - colors: a tuple with a color per channel, chosen from matplotlib.colors, html colors are also possible - colormap: name of a colormap from colorcet - pxsize: pixel size in um - deltaz: z slice interval in um - timeinterval: time between frames in seconds - extratags: other tags to be saved, example: Artist='John Doe', Tag4567=[400, 500] - or Copyright=Tag('ascii', 'Made by me'). See tiff_tag_registry.items(). - wp@tl20200214 - """ - def __init__(self, path: str | Path, shape: tuple[int, int, int], dtype: DTypeLike = 'uint16', - colors: Sequence[str] = None, colormap: str = None, pxsize: float = None, - deltaz: float = None, timeinterval: float = None, - compression: tuple[int, int] = (50000, 22), comment: str = None, - **extratags: Tag.Value | Tag) -> None: - assert len(shape) >= 3, 'please specify all c, z, t for the shape' - assert len(shape) <= 3, 'please specify only c, z, t for the shape' - assert np.dtype(dtype).char in 'BbHhf', 'datatype not supported' - assert colors is None or colormap is None, 'cannot have colors and colormap simultaneously' - - self.path = Path(path) - self.shape = shape - self.dtype = np.dtype(dtype) - self.colors = colors - self.colormap = colormap - self.pxsize = pxsize - self.deltaz = deltaz - self.timeinterval = timeinterval - self.compression = compression - self.comment = comment - self.extratags = {} if extratags is None else Tag.from_dict(extratags) # type: ignore - if pxsize is not None: - pxsize_fraction = Tag.fraction(pxsize) - self.extratags.update({282: Tag(pxsize_fraction), 283: Tag(pxsize_fraction)}) - - self.header = Header() - self.spp = self.shape[0] if self.colormap is None and self.colors is None else 1 # samples/pixel - self.nframes = np.prod(self.shape[1:]) if self.colormap is None and self.colors is None else np.prod(self.shape) - self.frame_extra_tags: dict[tuple[int, int, int], dict[int, Tag]] = {} - self.fh = FileHandle(self.path) - self.pool = ParPool(self.compress_frame) # type: ignore - self.hashes = PoolSingleton().manager.dict() - self.main_process = True - - with self.fh.lock() as fh: # noqa - self.header.write(fh) - - def __setstate__(self, state: dict[str, Any]) -> None: - self.__dict__.update(state) - self.main_process = False - - def __hash__(self) -> int: - return hash(self.path) - - def get_frame_number(self, n: tuple[int, int, int]) -> tuple[int, int]: - if self.colormap is None and self.colors is None: - return n[1] + n[2] * self.shape[1], n[0] - else: - return n[0] + n[1] * self.shape[0] + n[2] * self.shape[0] * self.shape[1], 0 - - def ij_tiff_frame(self, frame: np.ndarray) -> bytes: - with BytesIO() as frame_bytes: - with tifffile.TiffWriter(frame_bytes, bigtiff=self.header.bigtiff, - byteorder=self.header.byteorder) as t: # type: ignore - # predictor=True might save a few bytes, but requires the package imagecodes to save floats - t.write(frame, compression=self.compression, contiguous=True, predictor=False) # type: ignore - return frame_bytes.getvalue() - - def save(self, frame: np.ndarray | Any, c: int, z: int, t: int, - **extratags: Tag.Value | Tag) -> None: - """ save a 2d numpy array to the tiff at channel=c, slice=z, time=t, with optional extra tif tags - """ - assert (c, z, t) not in self.pool.tasks, f'frame {c} {z} {t} is added already' - assert all([0 <= i < s for i, s in zip((c, z, t), self.shape)]), \ - 'frame {} {} {} is outside shape {} {} {}'.format(c, z, t, *self.shape) - self.pool(frame.astype(self.dtype) if hasattr(frame, 'astype') else frame, handle=(c, z, t)) - if extratags: - self.frame_extra_tags[(c, z, t)] = Tag.from_dict(extratags) # type: ignore - - @property - def description(self) -> bytes: - desc = ['ImageJ=1.11a'] - if self.colormap is None and self.colors is None: - desc.extend((f'images={np.prod(self.shape[:1])}', f'slices={self.shape[1]}', f'frames={self.shape[2]}')) - else: - desc.extend((f'images={np.prod(self.shape)}', f'channels={self.shape[0]}', f'slices={self.shape[1]}', - f'frames={self.shape[2]}')) - if self.shape[0] == 1: - desc.append('mode=grayscale') - else: - desc.append('mode=composite') - desc.extend(('hyperstack=true', 'loop=false', 'unit=micron')) - if self.deltaz is not None: - desc.append(f'spacing={self.deltaz}') - if self.timeinterval is not None: - desc.append(f'interval={self.timeinterval}') - desc_bytes = [bytes(d, 'ascii') for d in desc] - if self.comment is not None: - desc_bytes.append(b'') - if isinstance(self.comment, bytes): - desc_bytes.append(self.comment) - else: - desc_bytes.append(bytes(self.comment, 'ascii')) - return b'\n'.join(desc_bytes) + b'\0' - - @cached_property - def colormap_bytes(self) -> Optional[bytes]: - if self.colormap: - colormap = getattr(colorcet, self.colormap) - colormap[0] = '#ffffff' - colormap[-1] = '#000000' - colormap = 65535 * np.array( - [[int(''.join(i), 16) for i in zip(*[iter(s[1:])] * 2)] for s in colormap]) // 255 - if np.dtype(self.dtype).itemsize == 2: - colormap = np.tile(colormap, 256).reshape((-1, 3)) - return b''.join([struct.pack(self.header.byteorder + 'H', c) for c in colormap.T.flatten()]) - - @cached_property - def colors_bytes(self) -> list[bytes]: - return [b''.join([struct.pack(self.header.byteorder + 'H', c) - for c in np.linspace(0, 65535 * np.array(mpl_colors.to_rgb(color)), - 65536 if np.dtype(self.dtype).itemsize == 2 else 256, - dtype=int).T.flatten()]) for color in self.colors] if self.colors else [] - - def close(self) -> None: - if self.main_process: - ifds, strips = {}, {} - for n in list(self.pool.tasks): - for ifd, strip, delta in self.pool[n]: - framenr, channel = self.get_frame_number(tuple(i + j for i, j in zip(n, delta))) # type: ignore - ifds[framenr], strips[(framenr, channel)] = ifd, strip - - self.pool.close() - with self.fh.lock() as fh: # noqa - for n, tags in self.frame_extra_tags.items(): - framenr, _ = self.get_frame_number(n) - ifds[framenr].update(tags) - if 0 in ifds and self.colormap is not None: - ifds[0][320] = Tag('SHORT', self.colormap_bytes) - ifds[0][262] = Tag('SHORT', 3) - if self.colors is not None: - for c, color in enumerate(self.colors_bytes): - if c in ifds: - ifds[c][320] = Tag('SHORT', color) - ifds[c][262] = Tag('SHORT', 3) - if 0 in ifds and 306 not in ifds[0]: - ifds[0][306] = Tag('ASCII', datetime.now().strftime('%Y:%m:%d %H:%M:%S')) - wrn = False - for framenr in range(self.nframes): - if framenr in ifds and all([(framenr, channel) in strips for channel in range(self.spp)]): - stripbyteoffsets, stripbytecounts = zip(*[strips[(framenr, channel)] - for channel in range(self.spp)]) - ifds[framenr][258].value = self.spp * ifds[framenr][258].value - ifds[framenr][270] = Tag('ASCII', self.description) - ifds[framenr][273] = Tag('LONG8', sum(stripbyteoffsets, [])) - ifds[framenr][277] = Tag('SHORT', self.spp) - ifds[framenr][279] = Tag('LONG8', sum(stripbytecounts, [])) - ifds[framenr][305] = Tag('ASCII', 'tiffwrite_tllab_NKI') - if self.extratags is not None: - ifds[framenr].update(self.extratags) - if self.colormap is None and self.colors is None and self.shape[0] > 1: - ifds[framenr][284] = Tag('SHORT', 2) - ifds[framenr].write(fh, self.header, self.write) - if framenr: - ifds[framenr].write_offset(ifds[framenr - 1].where_to_write_next_ifd_offset) - else: - ifds[framenr].write_offset(self.header.offset - self.header.offsetsize) - else: - wrn = True - if wrn: - warnings.warn('Some frames were not added to the tif file, either you forgot them, ' - 'or an error occurred and the tif file was closed prematurely.') - - def __enter__(self) -> IJTiffFile: - return self - - def __exit__(self, *args: Any, **kwargs: Any) -> None: - self.close() - - @staticmethod - def hash_check(fh: BinaryIO, bvalue: bytes, offset: int) -> bool: - addr = fh.tell() - fh.seek(offset) - same = bvalue == fh.read(len(bvalue)) - fh.seek(addr) - return same - - def write(self, fh: BinaryIO, bvalue: bytes) -> int: - hash_value = sha1(bvalue).hexdigest() # hash uses a random seed making hashes different in different processes - if hash_value in self.hashes and self.hash_check(fh, bvalue, self.hashes[hash_value]): - return self.hashes[hash_value] # reuse previously saved data - else: - if fh.tell() % 2: - fh.write(b'\x00') - offset = fh.tell() - self.hashes[hash_value] = offset - fh.write(bvalue) - return offset - - def compress_frame(self, frame: np.ndarray) -> Sequence[FrameInfo]: - """ This is run in a different process. Turns an image into bytes, writes them and returns the ifd, strip info - and czt delta. When subclassing IJTiffWrite this can be overridden to write one or more (using czt delta) - frames. - """ - stripbytecounts, ifd, chunks = self.get_chunks(self.ij_tiff_frame(frame)) - stripbyteoffsets = [] - with self.fh.lock() as fh: # noqa - for chunk in chunks: - stripbyteoffsets.append(self.write(fh, chunk)) - return (ifd, (stripbyteoffsets, stripbytecounts), (0, 0, 0)), - - @staticmethod - def get_chunks(frame: bytes) -> tuple[list[int], IFD, list[bytes]]: - with BytesIO(frame) as fh: - ifd = IFD(fh) - stripoffsets = ifd[273].value - stripbytecounts = ifd[279].value - chunks = [] - for stripoffset, stripbytecount in zip(stripoffsets, stripbytecounts): - fh.seek(stripoffset) - chunks.append(fh.read(stripbytecount)) - return stripbytecounts, ifd, chunks - - -class FileHandle: - """ Process safe file handle """ - def __init__(self, path: Path) -> None: - manager = PoolSingleton().manager - if path.exists(): - path.unlink() - with open(path, 'xb'): - pass - self.path = path - self._lock = manager.RLock() - self._pos = manager.Value('i', 0) - - @contextmanager - def lock(self) -> Generator[BinaryIO, None, None]: - with self._lock: - with open(self.path, 'rb+') as f: - try: - f.seek(self._pos.value) - yield f - finally: - self._pos.value = f.tell()