wim/tiffwrite
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

- add tiffwrite function for python

Browse Source 
- parallel zstd compression
This commit is contained in:
Wim Pomp
2024-10-09 15:07:38 +02:00
parent 52785037b9
commit f62b711692
7 changed files with 97 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
.github/workflows/mypy.yml vendored
View File

@@ -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 .

2
.gitignore vendored
View File

@@ -5,3 +5,5 @@
/.pytest_cache/ /.pytest_cache/
/venv/ /venv/
/target/ /target/
/Cargo.lock
/foo.tif

3
Cargo.toml
View File

@@ -12,12 +12,11 @@ crate-type = ["cdylib", "rlib"]
pyo3 = { version = "0.21.2", features = ["extension-module", "abi3-py310", "generate-import-lib", "anyhow", "multiple-pymethods"] } pyo3 = { version = "0.21.2", features = ["extension-module", "abi3-py310", "generate-import-lib", "anyhow", "multiple-pymethods"] }
anyhow = "1.0.89" anyhow = "1.0.89"
rayon = "1.10.0" rayon = "1.10.0"
fraction = "0.15.3"
num = "0.4.3" num = "0.4.3"
ndarray = "0.15.6" ndarray = "0.15.6"
chrono = "0.4.38" chrono = "0.4.38"
numpy = "0.21.0" numpy = "0.21.0"
futures = "0.3.31" zstd = "0.13.2"
[features] [features]
nopython = [] nopython = []

51
tiffwrite/__init__.py → py/tiffwrite/__init__.py
View File

@@ -1,16 +1,22 @@
from __future__ import annotations from __future__ import annotations
import numpy as np from itertools import product
from typing import Any, Self, Sequence
from pathlib import Path from pathlib import Path
from typing import Any, Sequence
import numpy as np
from numpy.typing import ArrayLike, DTypeLike from numpy.typing import ArrayLike, DTypeLike
from tqdm.auto import tqdm
from . import tiffwrite as rs from . import tiffwrite_rs as rs # noqa
__all__ = ['Tag', 'IJTiffFile', 'tiffwrite'] __all__ = ['Header', 'IJTiffFile', 'IFD', 'FrameInfo', 'Tag', 'Strip', 'tiffwrite']
class Header:
pass
class IFD(dict): class IFD(dict):
pass pass
@@ -19,11 +25,16 @@ class Tag(rs.Tag):
pass pass
Strip = tuple[list[int], list[int]]
CZT = tuple[int, int, int]
FrameInfo = tuple[IFD, Strip, CZT]
class IJTiffFile(rs.IJTiffFile): 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], dtype: DTypeLike = 'uint16',
colors: Sequence[str] = None, colormap: str = None, pxsize: float = None, 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, comment: str = None,
**extratags: Tag.Value | Tag) -> None: **extratags: Tag.Value | Tag) -> IJTiffFile:
new = super().__new__(cls, str(path), shape) new = super().__new__(cls, str(path), shape)
if colors is not None: if colors is not None:
new = new.with_colors(colors) new = new.with_colors(colors)
@@ -41,14 +52,14 @@ class IJTiffFile(rs.IJTiffFile):
new = new.extend_extratags(extratags) new = new.extend_extratags(extratags)
return new return new
def __init__(self, path: str | Path, shape: tuple[int, int, int], dtype: DTypeLike = 'uint16', def __init__(self, path: str | Path, shape: tuple[int, int, int], dtype: DTypeLike = 'uint16', # noqa
colors: Sequence[str] = None, colormap: str = None, pxsize: float = None, colors: Sequence[str] = None, colormap: str = None, pxsize: float = None, # noqa
deltaz: float = None, timeinterval: float = None, comment: str = None, deltaz: float = None, timeinterval: float = None, comment: str = None, # noqa
**extratags: Tag.Value | Tag) -> None: **extratags: Tag.Value | Tag) -> None: # noqa
self.path = Path(path) self.path = Path(path)
self.dtype = np.dtype(dtype) self.dtype = np.dtype(dtype)
def __enter__(self) -> Self: def __enter__(self) -> IJTiffFile:
return self return self
def __exit__(self, exc_type, exc_val, exc_tb): def __exit__(self, exc_type, exc_val, exc_tb):
@@ -81,7 +92,7 @@ class IJTiffFile(rs.IJTiffFile):
raise TypeError(f'Cannot save type {self.dtype}') raise TypeError(f'Cannot save type {self.dtype}')
def tiffwrite(file: str | Path, data: ArrayLike, axes: str = 'TZCXY', dtype: DTypeLike = None, bar: bool = False, def tiffwrite(file: str | Path, data: np.ndarray, axes: str = 'TZCXY', dtype: DTypeLike = None, bar: bool = False,
*args: Any, **kwargs: Any) -> None: *args: Any, **kwargs: Any) -> None:
""" file: string; filename of the new tiff file """ file: string; filename of the new tiff file
data: 2 to 5D numpy array data: 2 to 5D numpy array
@@ -90,3 +101,21 @@ def tiffwrite(file: str | Path, data: ArrayLike, axes: str = 'TZCXY', dtype: DTy
bar: bool; whether to show a progress bar bar: bool; whether to show a progress bar
other args: see IJTiffFile 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), # type: ignore
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

6
pyproject.toml
View File

@@ -13,9 +13,13 @@ classifiers = [
"Programming Language :: Python :: Implementation :: PyPy", "Programming Language :: Python :: Implementation :: PyPy",
] ]
[project.optional-dependencies]
test = ["pytest"]
[tool.maturin] [tool.maturin]
python-source = "py"
features = ["pyo3/extension-module"] features = ["pyo3/extension-module"]
module-name = "tiffwrite" module-name = "tiffwrite.tiffwrite_rs"
[tool.isort] [tool.isort]
line_length = 119 line_length = 119

70
src/lib.rs
View File

@@ -6,19 +6,19 @@ use std::collections::HashMap;
use std::fs::{File, OpenOptions}; use std::fs::{File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write}; use std::io::{Read, Seek, SeekFrom, Write};
use anyhow::Result; use anyhow::Result;
use fraction::Fraction; use num::{Complex, Rational32, Zero};
use num::{Complex, Zero};
use num::complex::ComplexFloat;
use ndarray::{s, Array2}; use ndarray::{s, Array2};
use num::traits::ToBytes; use num::traits::ToBytes;
use std::hash::{DefaultHasher, Hash, Hasher}; use std::hash::{DefaultHasher, Hash, Hasher};
use chrono::Utc; use chrono::Utc;
use zstd::stream::encode_all;
use rayon::prelude::*;
const TAG_SIZE: usize = 20; const TAG_SIZE: usize = 20;
const OFFSET_SIZE: usize = 8; const OFFSET_SIZE: usize = 8;
const OFFSET: u64 = 16; const OFFSET: u64 = 16;
const COMPRESSION: u16 = 1; const COMPRESSION: u16 = 50000;
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
@@ -116,10 +116,10 @@ impl Tag {
Tag::new(code, long.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 4) Tag::new(code, long.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 4)
} }
pub fn rational(code: u16, rational: Vec<Fraction>) -> Self { pub fn rational(code: u16, rational: Vec<Rational32>) -> Self {
Tag::new(code, rational.into_iter().map(|x| Tag::new(code, rational.into_iter().map(|x|
u32::try_from(*x.denom().unwrap()).unwrap().to_le_bytes().into_iter().chain( u32::try_from(*x.denom()).unwrap().to_le_bytes().into_iter().chain(
u32::try_from(*x.numer().unwrap()).unwrap().to_le_bytes()).collect::<Vec<_>>() u32::try_from(*x.numer()).unwrap().to_le_bytes()).collect::<Vec<_>>()
).flatten().collect(), 5) ).flatten().collect(), 5)
} }
@@ -135,10 +135,10 @@ impl Tag {
Tag::new(code, slong.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 9) Tag::new(code, slong.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 9)
} }
pub fn srational(code: u16, srational: Vec<Fraction>) -> Self { pub fn srational(code: u16, srational: Vec<Rational32>) -> Self {
Tag::new(code, srational.into_iter().map(|x| Tag::new(code, srational.into_iter().map(|x|
i32::try_from(*x.denom().unwrap()).unwrap().to_le_bytes().into_iter().chain( i32::try_from(*x.denom()).unwrap().to_le_bytes().into_iter().chain(
i32::try_from(*x.numer().unwrap()).unwrap().to_le_bytes()).collect::<Vec<_>>() i32::try_from(*x.numer()).unwrap().to_le_bytes()).collect::<Vec<_>>()
).flatten().collect(), 10) ).flatten().collect(), 10)
} }
@@ -162,7 +162,7 @@ impl Tag {
pub fn complex(code: u16, complex: Vec<Complex<f32>>) -> Self { pub fn complex(code: u16, complex: Vec<Complex<f32>>) -> Self {
Tag::new(code, complex.into_iter().map(|x| Tag::new(code, complex.into_iter().map(|x|
x.re().to_le_bytes().into_iter().chain(x.im().to_le_bytes()).collect::<Vec<_>>() x.re.to_le_bytes().into_iter().chain(x.im.to_le_bytes()).collect::<Vec<_>>()
).flatten().collect(), 15) ).flatten().collect(), 15)
} }
@@ -237,7 +237,7 @@ impl Tag {
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
struct Frame { struct Frame {
tilebyteoffsets: Vec<u64>, tileoffsets: Vec<u64>,
tilebytecounts: Vec<u64>, tilebytecounts: Vec<u64>,
image_width: u32, image_width: u32,
image_length: u32, image_length: u32,
@@ -250,11 +250,11 @@ struct Frame {
impl Frame { impl Frame {
fn new( fn new(
tilebyteoffsets: Vec<u64>, tilebytecounts: Vec<u64>, image_width: u32, image_length: u32, tileoffsets: Vec<u64>, tilebytecounts: Vec<u64>, image_width: u32, image_length: u32,
bits_per_sample: u16, sample_format: u16, tile_width: u16, tile_length: u16 bits_per_sample: u16, sample_format: u16, tile_width: u16, tile_length: u16
) -> Self { ) -> Self {
Frame { Frame {
tilebyteoffsets, tilebytecounts, image_width, image_length, bits_per_sample, tileoffsets, tilebytecounts, image_width, image_length, bits_per_sample,
sample_format, tile_width, tile_length, extra_tags: Vec::new() sample_format, tile_width, tile_length, extra_tags: Vec::new()
} }
} }
@@ -276,8 +276,7 @@ macro_rules! bytes_impl {
const SAMPLE_FORMAT: u16 = $sample_format; const SAMPLE_FORMAT: u16 = $sample_format;
#[inline] #[inline]
fn bytes(&self) -> Vec<u8> fn bytes(&self) -> Vec<u8> {
{
self.to_le_bytes().to_vec() self.to_le_bytes().to_vec()
} }
} }
@@ -293,7 +292,6 @@ bytes_impl!(u128, 128, 1);
bytes_impl!(usize, 64, 1); bytes_impl!(usize, 64, 1);
#[cfg(target_pointer_width = "32")] #[cfg(target_pointer_width = "32")]
bytes_impl!(usize, 32, 1); bytes_impl!(usize, 32, 1);
bytes_impl!(i8, 8, 2); bytes_impl!(i8, 8, 2);
bytes_impl!(i16, 16, 2); bytes_impl!(i16, 16, 2);
bytes_impl!(i32, 32, 2); bytes_impl!(i32, 32, 2);
@@ -422,30 +420,38 @@ impl IJTiffFile {
pub fn save<T: Bytes + Clone + Zero>(&mut self, frame: Array2<T>, c: usize, z: usize, t: usize, pub fn save<T: Bytes + Clone + Zero>(&mut self, frame: Array2<T>, c: usize, z: usize, t: usize,
extra_tags: Option<Vec<Tag>>) -> Result<()> { extra_tags: Option<Vec<Tag>>) -> Result<()> {
self.compress_frame(frame, c, z, t, extra_tags); self.compress_frame(frame.reversed_axes(), c, z, t, extra_tags)?;
Ok(()) Ok(())
} }
fn compress_frame<T: Bytes + Clone + Zero>(&mut self, frame: Array2<T>, c: usize, z: usize, t: usize, fn compress_frame<T: Bytes + Clone + Zero>(&mut self, frame: Array2<T>,
extra_tags: Option<Vec<Tag>>) { c: usize, z: usize, t: usize,
extra_tags: Option<Vec<Tag>>) -> Result<()> {
let image_width = frame.shape()[0] as u32; let image_width = frame.shape()[0] as u32;
let image_length = frame.shape()[1] as u32; let image_length = frame.shape()[1] as u32;
let mut tilebyteoffsets = Vec::new(); let tile_size = 2usize.pow(((image_width as f64 * image_length as f64 / 64f64).log2() / 2f64).round() as u32).max(16).min(1024);
let mut tileoffsets = Vec::new();
let mut tilebytecounts = Vec::new(); let mut tilebytecounts = Vec::new();
let tiles = IJTiffFile::tile(frame.reversed_axes(), 64); let tiles = IJTiffFile::tile(frame.reversed_axes(), tile_size);
for tile in tiles { let byte_tiles: Vec<Vec<u8>> = tiles.into_iter().map(
let bytes: Vec<u8> = tile.map(|x| x.bytes()).into_iter().flatten().collect(); |tile| tile.map(|x| x.bytes()).into_iter().flatten().collect()
).collect();
for tile in byte_tiles.into_par_iter().map(|x| encode_all(&*x, 3)).collect::<Vec<_>>() {
if let Ok(bytes) = tile {
tilebytecounts.push(bytes.len() as u64); tilebytecounts.push(bytes.len() as u64);
tilebyteoffsets.push(self.write(&bytes).unwrap()); tileoffsets.push(self.write(&bytes)?);
} }
let mut frame = Frame::new(tilebyteoffsets, tilebytecounts, image_width, image_length, }
T::BITS_PER_SAMPLE, T::SAMPLE_FORMAT, 64, 64);
let mut frame = Frame::new(tileoffsets, tilebytecounts, image_width, image_length,
T::BITS_PER_SAMPLE, T::SAMPLE_FORMAT, tile_size as u16, tile_size as u16);
if let Some(tags) = extra_tags { if let Some(tags) = extra_tags {
for tag in tags { for tag in tags {
frame.extra_tags.push(tag); frame.extra_tags.push(tag);
} }
} }
self.frames.insert(self.get_frame_number(c, z, t), frame); self.frames.insert(self.get_frame_number(c, z, t), frame);
Ok(())
} }
fn tile<T: Clone + Zero>(frame: Array2<T>, size: usize) -> Vec<Array2<T>> { fn tile<T: Clone + Zero>(frame: Array2<T>, size: usize) -> Vec<Array2<T>> {
@@ -484,11 +490,11 @@ impl IJTiffFile {
tiles tiles
} }
fn get_colormap(&self, colormap: &Vec<u16>) -> Result<Vec<u16>> { fn get_colormap(&self, _colormap: &Vec<u16>) -> Result<Vec<u16>> {
todo!(); todo!();
} }
fn get_color(&self, colors: (u8, u8, u8)) -> Result<Vec<u16>> { fn get_color(&self, _colors: (u8, u8, u8)) -> Result<Vec<u16>> {
todo!(); todo!();
} }
@@ -497,12 +503,12 @@ impl IJTiffFile {
let mut warn = false; let mut warn = false;
for frame_number in 0..self.n_frames { for frame_number in 0..self.n_frames {
if let Some(frame) = self.frames.get(&(frame_number, 0)) { if let Some(frame) = self.frames.get(&(frame_number, 0)) {
let mut tilebyteoffsets = Vec::new(); let mut tileoffsets = Vec::new();
let mut tilebytecounts = Vec::new(); let mut tilebytecounts = Vec::new();
let mut frame_count = 0; let mut frame_count = 0;
for channel in 0..self.samples_per_pixel { for channel in 0..self.samples_per_pixel {
if let Some(frame_n) = self.frames.get(&(frame_number, channel)) { if let Some(frame_n) = self.frames.get(&(frame_number, channel)) {
tilebyteoffsets.extend(frame_n.tilebyteoffsets.iter()); tileoffsets.extend(frame_n.tileoffsets.iter());
tilebytecounts.extend(frame_n.tilebytecounts.iter()); tilebytecounts.extend(frame_n.tilebytecounts.iter());
frame_count += 1; frame_count += 1;
} else { } else {
@@ -519,7 +525,7 @@ impl IJTiffFile {
ifd.push_tag(Tag::ascii(305, "tiffwrite_rs")); ifd.push_tag(Tag::ascii(305, "tiffwrite_rs"));
ifd.push_tag(Tag::short(322, vec![frame.tile_width])); ifd.push_tag(Tag::short(322, vec![frame.tile_width]));
ifd.push_tag(Tag::short(323, vec![frame.tile_length])); ifd.push_tag(Tag::short(323, vec![frame.tile_length]));
ifd.push_tag(Tag::long8(324, tilebyteoffsets)); ifd.push_tag(Tag::long8(324, tileoffsets));
ifd.push_tag(Tag::long8(325, tilebytecounts)); ifd.push_tag(Tag::long8(325, tilebytecounts));
ifd.push_tag(Tag::short(339, vec![frame.sample_format])); ifd.push_tag(Tag::short(339, vec![frame.sample_format]));
if frame_number == 0 { if frame_number == 0 {

15
src/py.rs
View File

@@ -1,7 +1,6 @@
use pyo3::prelude::*; use pyo3::prelude::*;
use crate::{IJTiffFile, Tag}; use crate::{IJTiffFile, Tag};
use fraction::Fraction; use num::{Complex, Rational32, FromPrimitive};
use num::Complex;
use numpy::{PyReadonlyArray2, PyArrayMethods}; use numpy::{PyReadonlyArray2, PyArrayMethods};
@@ -36,7 +35,7 @@ impl PyTag {
#[staticmethod] #[staticmethod]
fn rational(code: u16, rational: Vec<f64>) -> Self { fn rational(code: u16, rational: Vec<f64>) -> Self {
PyTag { tag: Tag::rational(code, rational.into_iter().map(|x| Fraction::from(x)).collect()) } PyTag { tag: Tag::rational(code, rational.into_iter().map(|x| Rational32::from_f64(x).unwrap()).collect()) }
} }
#[staticmethod] #[staticmethod]
@@ -56,7 +55,7 @@ impl PyTag {
#[staticmethod] #[staticmethod]
fn srational(code: u16, srational: Vec<f64>) -> Self { fn srational(code: u16, srational: Vec<f64>) -> Self {
PyTag { tag: Tag::srational(code, srational.into_iter().map(|x| Fraction::from(x)).collect()) } PyTag { tag: Tag::srational(code, srational.into_iter().map(|x| Rational32::from_f64(x).unwrap()).collect()) }
} }
#[staticmethod] #[staticmethod]
@@ -98,6 +97,10 @@ impl PyTag {
fn ifd8(code: u16, ifd8: Vec<u64>) -> Self { fn ifd8(code: u16, ifd8: Vec<u64>) -> Self {
PyTag { tag: Tag::ifd8(code, ifd8) } PyTag { tag: Tag::ifd8(code, ifd8) }
} }
fn count(&self) -> u64 {
self.tag.count()
}
} }
@@ -197,8 +200,10 @@ impl_save!(i64, save_i64);
impl_save!(f32, save_f32); impl_save!(f32, save_f32);
impl_save!(f64, save_f64); impl_save!(f64, save_f64);
#[pymodule] #[pymodule]
fn tiffwrite(m: &Bound<'_, PyModule>) -> PyResult<()> { #[pyo3(name = "tiffwrite_rs")]
fn tiffwrite_rs(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyTag>()?; m.add_class::<PyTag>()?;
m.add_class::<PyIJTiffFile>()?; m.add_class::<PyIJTiffFile>()?;
Ok(()) Ok(())