From 1197806a6fd0d070d66eaf73d7e385c34a34d76e Mon Sep 17 00:00:00 2001 From: Wim Pomp Date: Thu, 10 Oct 2024 15:28:14 +0200 Subject: [PATCH] - start some color(map) support in python - compress_frame function in python for backwards compatibility - save all extra tags in a single hashmap - construct tags from references - store frames by c, z, t - save px_size in tiff - some getters and setters in py.rs --- py/tiffwrite/__init__.py | 83 +++--- pyproject.toml | 2 +- src/lib.rs | 616 ++++++++++++++++++++++++++------------- src/main.rs | 7 +- src/py.rs | 266 +++++++++++++---- 5 files changed, 681 insertions(+), 293 deletions(-) diff --git a/py/tiffwrite/__init__.py b/py/tiffwrite/__init__.py index bef208d..85edd53 100644 --- a/py/tiffwrite/__init__.py +++ b/py/tiffwrite/__init__.py @@ -4,7 +4,9 @@ from itertools import product from pathlib import Path from typing import Any, Sequence +import colorcet import numpy as np +from matplotlib import colors as mpl_colors from numpy.typing import ArrayLike, DTypeLike from tqdm.auto import tqdm @@ -27,29 +29,29 @@ class Tag(rs.Tag): Strip = tuple[list[int], list[int]] CZT = tuple[int, int, int] -FrameInfo = tuple[IFD, Strip, CZT] +FrameInfo = tuple[np.ndarray, None, CZT] class IJTiffFile(rs.IJTiffFile): def __new__(cls, 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, comment: str = None, - **extratags: Tag.Value | Tag) -> IJTiffFile: + **extratags: Tag) -> IJTiffFile: new = super().__new__(cls, str(path), shape) if colors is not None: - new = new.with_colors(colors) + new.colors = np.array([get_color(color) for color in colors]) if colormap is not None: - new = new.with_colormap(colormap) + new.colormap = get_colormap(colormap) if pxsize is not None: - new = new.with_pxsize(pxsize) + new.px_size = float(pxsize) if deltaz is not None: - new = new.with_deltaz(deltaz) + new.delta_z = float(deltaz) if timeinterval is not None: - new = new.with_timeinterval(timeinterval) + new.time_interval = float(timeinterval) if comment is not None: - new = new.with_comment(comment) - if extratags: - new = new.extend_extratags(extratags) + new.comment = comment + for extra_tag in extratags: + new.append_extra_tag(extra_tag, None) return new def __init__(self, path: str | Path, shape: tuple[int, int, int], dtype: DTypeLike = 'uint16', # noqa @@ -66,30 +68,43 @@ class IJTiffFile(rs.IJTiffFile): self.close() def save(self, frame: ArrayLike, c: int, z: int, t: int) -> None: - 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}') + for frame, _, (cn, zn, tn) in self.compress_frame(frame): + frame = np.asarray(frame).astype(self.dtype) + match self.dtype: + case np.uint8: + self.save_u8(frame, c + cn, z + zn, t + tn) + case np.uint16: + self.save_u16(frame, c + cn, z + zn, t + tn) + case np.uint32: + self.save_u32(frame, c + cn, z + zn, t + tn) + case np.uint64: + self.save_u64(frame, c + cn, z + zn, t + tn) + case np.int8: + self.save_i8(frame, c + cn, z + zn, t + tn) + case np.int16: + self.save_i16(frame, c + cn, z + zn, t + tn) + case np.int32: + self.save_i32(frame, c + cn, z + zn, t + tn) + case np.int64: + self.save_i64(frame, c + cn, z + zn, t + tn) + case np.float32: + self.save_f32(frame, c + cn, z + zn, t + tn) + case np.float64: + self.save_f64(frame, c + cn, z + zn, t + tn) + case _: + raise TypeError(f'Cannot save type {self.dtype}') + + def compress_frame(self, frame: ArrayLike) -> tuple[FrameInfo]: # noqa + return (frame, None, (0, 0, 0)), + +def get_colormap(colormap: str) -> np.ndarray: + colormap = getattr(colorcet, colormap) + colormap[0] = '#ffffff' + colormap[-1] = '#000000' + return np.array([[int(''.join(i), 16) for i in zip(*[iter(s[1:])] * 2)] for s in colormap]).astype('uint8') + +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, diff --git a/pyproject.toml b/pyproject.toml index a658afd..0078504 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] -dependencies = ["numpy", "tqdm"] +dependencies = ["colorcet", "matplotlib", "numpy", "tqdm"] [project.optional-dependencies] test = ["pytest"] diff --git a/src/lib.rs b/src/lib.rs index 758632b..580fbcf 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,31 +1,29 @@ // #[cfg(not(feature = "nopython"))] mod py; +use anyhow::Result; +use chrono::Utc; +use ndarray::{s, Array2}; +use num::traits::ToBytes; +use num::{Complex, FromPrimitive, Rational32, Zero}; +use rayon::prelude::*; use std::cmp::Ordering; use std::collections::HashMap; use std::fs::{File, OpenOptions}; -use std::io::{Read, Seek, SeekFrom, Write}; -use anyhow::Result; -use num::{Complex, Rational32, Zero}; -use ndarray::{s, Array2}; -use num::traits::ToBytes; use std::hash::{DefaultHasher, Hash, Hasher}; +use std::io::{Read, Seek, SeekFrom, Write}; use std::thread; use std::thread::JoinHandle; -use chrono::Utc; use zstd::stream::encode_all; -use rayon::prelude::*; - const TAG_SIZE: usize = 20; const OFFSET_SIZE: usize = 8; const OFFSET: u64 = 16; const COMPRESSION: u16 = 50000; - #[derive(Clone, Debug)] struct IFD { - tags: Vec + tags: Vec, } impl IFD { @@ -39,12 +37,6 @@ impl IFD { } } - fn extend_tags(&mut self, tags: Vec) { - for tag in tags { - self.push_tag(tag); - } - } - fn write(&mut self, ijtifffile: &mut IJTiffFile, where_to_write_offset: u64) -> Result { self.tags.sort(); ijtifffile.file.seek(SeekFrom::End(0))?; @@ -52,7 +44,9 @@ impl IFD { ijtifffile.file.write(&[0])?; } let offset = ijtifffile.file.stream_position()?; - ijtifffile.file.write(&(self.tags.len() as u64).to_le_bytes())?; + ijtifffile + .file + .write(&(self.tags.len() as u64).to_le_bytes())?; for tag in self.tags.iter_mut() { tag.write_tag(ijtifffile)?; @@ -62,19 +56,20 @@ impl IFD { for tag in self.tags.iter() { tag.write_data(ijtifffile)?; } - ijtifffile.file.seek(SeekFrom::Start(where_to_write_offset))?; + 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 + offset: u64, } impl PartialOrd for Tag { @@ -97,11 +92,16 @@ impl PartialEq for Tag { impl Tag { pub fn new(code: u16, bytes: Vec, ttype: u16) -> Self { - Tag { code, bytes, ttype, offset: 0 } + Tag { + code, + bytes, + ttype, + offset: 0, + } } - pub fn byte(code: u16, byte: Vec) -> Self { - Tag::new(code, byte, 1) + pub fn byte(code: u16, byte: &Vec) -> Self { + Tag::new(code, byte.to_owned(), 1) } pub fn ascii(code: u16, ascii: &str) -> Self { @@ -110,96 +110,213 @@ impl Tag { Tag::new(code, bytes, 2) } - pub fn short(code: u16, short: Vec) -> Self { - Tag::new(code, short.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 3) + pub fn short(code: u16, short: &Vec) -> Self { + Tag::new( + code, + short + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 3, + ) } - pub fn long(code: u16, long: Vec) -> Self { - Tag::new(code, long.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 4) + pub fn long(code: u16, long: &Vec) -> Self { + Tag::new( + code, + long.into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 4, + ) } - pub fn rational(code: u16, rational: Vec) -> Self { - Tag::new(code, rational.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 rational(code: u16, rational: &Vec) -> Self { + Tag::new( + code, + rational + .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, sbyte: Vec) -> Self { - Tag::new(code, sbyte.iter().map(|x| x.to_le_bytes()).flatten().collect(), 6) + pub fn sbyte(code: u16, sbyte: &Vec) -> Self { + Tag::new( + code, + sbyte.iter().map(|x| x.to_le_bytes()).flatten().collect(), + 6, + ) } - pub fn sshort(code: u16, sshort: Vec) -> Self { - Tag::new(code, sshort.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 8) + pub fn sshort(code: u16, sshort: &Vec) -> Self { + Tag::new( + code, + sshort + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 8, + ) } - pub fn slong(code: u16, slong: Vec) -> Self { - Tag::new(code, slong.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 9) + pub fn slong(code: u16, slong: &Vec) -> Self { + Tag::new( + code, + slong + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 9, + ) } - pub fn srational(code: u16, srational: Vec) -> Self { - Tag::new(code, srational.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 srational(code: u16, srational: &Vec) -> Self { + Tag::new( + code, + srational + .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, float: Vec) -> Self { - Tag::new(code, float.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 11) + pub fn float(code: u16, float: &Vec) -> Self { + Tag::new( + code, + float + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 11, + ) } - pub fn double(code: u16, double: Vec) -> Self { - Tag::new(code, double.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 12) + pub fn double(code: u16, double: &Vec) -> Self { + Tag::new( + code, + double + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 12, + ) } - pub fn ifd(code: u16, ifd: Vec) -> Self { - Tag::new(code, ifd.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 13) + pub fn ifd(code: u16, ifd: &Vec) -> Self { + Tag::new( + code, + ifd.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), + 13, + ) } pub fn unicode(code: u16, unicode: &str) -> Self { - let mut bytes: Vec = unicode.encode_utf16().map(|x| x.to_le_bytes()).flatten().collect(); + let mut bytes: Vec = unicode + .encode_utf16() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(); bytes.push(0); Tag::new(code, bytes, 14) } - pub fn complex(code: u16, complex: Vec>) -> Self { - Tag::new(code, complex.into_iter().map(|x| - x.re.to_le_bytes().into_iter().chain(x.im.to_le_bytes()).collect::>() - ).flatten().collect(), 15) + pub fn complex(code: u16, complex: &Vec>) -> Self { + Tag::new( + code, + complex + .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, long8: Vec) -> Self { - Tag::new(code, long8.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 16) + pub fn long8(code: u16, long8: &Vec) -> Self { + Tag::new( + code, + long8 + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 16, + ) } - pub fn slong8(code: u16, slong8: Vec) -> Self { - Tag::new(code, slong8.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 17) + pub fn slong8(code: u16, slong8: &Vec) -> Self { + Tag::new( + code, + slong8 + .into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 17, + ) } - pub fn ifd8(code: u16, ifd8: Vec) -> Self { - Tag::new(code, ifd8.into_iter().map(|x| x.to_le_bytes()).flatten().collect(), 18) + pub fn ifd8(code: u16, ifd8: &Vec) -> Self { + Tag::new( + code, + ifd8.into_iter() + .map(|x| x.to_le_bytes()) + .flatten() + .collect(), + 18, + ) } pub fn count(&self) -> u64 { let c = match self.ttype { - 1 => self.bytes.len(), // BYTE - 2 => self.bytes.len(), // ASCII + 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 + 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 + 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 @@ -226,7 +343,8 @@ impl Tag { 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))?; + 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])?; @@ -236,17 +354,16 @@ impl Tag { } } - +#[derive(Debug)] struct CompressedFrame { bytes: Vec>, image_width: u32, image_length: u32, tile_size: usize, bits_per_sample: u16, - sample_format: u16 + sample_format: u16, } - #[derive(Clone, Debug)] struct Frame { tileoffsets: Vec, @@ -261,17 +378,28 @@ struct Frame { impl Frame { fn new( - tileoffsets: Vec, tilebytecounts: Vec, image_width: u32, image_length: u32, - bits_per_sample: u16, sample_format: u16, tile_width: u16, tile_length: u16 + tileoffsets: Vec, + tilebytecounts: Vec, + image_width: u32, + image_length: u32, + bits_per_sample: u16, + sample_format: u16, + tile_width: u16, + tile_length: u16, ) -> Self { Frame { - tileoffsets, tilebytecounts, image_width, image_length, bits_per_sample, - sample_format, tile_width, tile_length + tileoffsets, + tilebytecounts, + 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; @@ -279,7 +407,6 @@ pub trait Bytes { fn bytes(&self) -> Vec; } - macro_rules! bytes_impl { ($T:ty, $bits_per_sample:expr, $sample_format:expr) => { impl Bytes for $T { @@ -315,23 +442,26 @@ 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, u8), Frame>, + frames: HashMap<(usize, usize, usize), Frame>, hashes: HashMap, - threads: HashMap<(usize, u8), JoinHandle>, + threads: HashMap<(usize, usize, usize), JoinHandle>, pub shape: (usize, usize, usize), - pub n_frames: usize, - pub samples_per_pixel: u8, - pub colormap: Option>, - pub colors: Option>, + pub colors: Colors, pub comment: Option, + pub px_size: Option, pub delta_z: Option, - pub timeinterval: Option, - pub extra_tags: Vec, - pub extra_tags_frame: HashMap> + pub time_interval: Option, + pub extra_tags: HashMap, Vec>, } impl Drop for IJTiffFile { @@ -344,28 +474,35 @@ impl Drop for IJTiffFile { impl IJTiffFile { pub fn new(path: &str, shape: (usize, usize, usize)) -> Result { - let mut file = OpenOptions::new().create(true).truncate(true) - .write(true).read(true).open(path)?; + 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())?; - let colormap: Option> = None; - let (spp, n_frames) = if let None = &colormap { - (shape.0 as u8, shape.1 * shape.2) - } else { - (1, shape.0 * shape.1 * shape.2) - }; - Ok(IJTiffFile { file, frames: HashMap::new(), hashes: HashMap::new(), - threads: HashMap::new(), shape, n_frames, - samples_per_pixel: spp, colormap: None, colors: None, comment: None, delta_z: None, - timeinterval: None, extra_tags: Vec::new(), extra_tags_frame: HashMap::new() } ) + Ok(IJTiffFile { + file, + frames: HashMap::new(), + hashes: HashMap::new(), + threads: HashMap::new(), + shape, + colors: Colors::None, + comment: None, + px_size: None, + delta_z: None, + time_interval: None, + extra_tags: HashMap::new(), + }) } pub fn description(&self) -> String { let mut desc: String = String::from("ImageJ=1.11a"); - if let (None, None) = (self.colormap.as_ref(), self.colors.as_ref()) { + if let Colors::None = self.colors { desc += &format!("\nimages={}", self.shape.0); desc += &format!("\nslices={}", self.shape.1); desc += &format!("\nframes={}", self.shape.2); @@ -384,7 +521,7 @@ impl IJTiffFile { if let Some(delta_z) = self.delta_z { desc += &format!("\nspacing={}", delta_z); } - if let Some(timeinterval) = self.timeinterval { + if let Some(timeinterval) = self.time_interval { desc += &format!("\ninterval={}", timeinterval); } if let Some(comment) = &self.comment { @@ -393,11 +530,27 @@ impl IJTiffFile { desc } - fn get_frame_number(&self, c: usize, z: usize, t: usize) -> (usize, u8) { - if let (None, None) = (self.colormap.as_ref(), self.colors.as_ref()) { - (z + t * self.shape.1, c as u8) + fn get_czt(&self, frame_number: usize, channel: u8) -> (usize, usize, usize) { + if let Colors::None = self.colors { + ( + channel as usize, + frame_number % self.shape.1, + frame_number / self.shape.1, + ) } else { - (c + z * self.shape.0 + t * self.shape.0 * self.shape.1, 0) + ( + frame_number % self.shape.0, + frame_number / self.shape.0 % self.shape.1, + frame_number / self.shape.0 / self.shape.1, + ) + } + } + + fn spp_and_n_frames(&self) -> (u8, usize) { + if let Colors::None = &self.colors { + (self.shape.0 as u8, self.shape.1 * self.shape.2) + } else { + (1, self.shape.0 * self.shape.1 * self.shape.2) } } @@ -419,7 +572,9 @@ impl IJTiffFile { 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())? { + 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 { @@ -432,64 +587,88 @@ impl IJTiffFile { } } - pub fn save(&mut self, frame: Array2, c: usize, z: usize, t: usize, - extra_tags: Option>) -> Result<()> - where T: Bytes + Clone + Send + Sync + Zero + 'static { - let key = self.get_frame_number(c, z, t); - if let Some(extra_tags) = extra_tags { - if let Some(extra_tags_frame) = self.extra_tags_frame.get_mut(&key.0) { - extra_tags_frame.extend(extra_tags); - } else { - self.extra_tags_frame.insert(key.0, extra_tags); - } - } - self.compress_frame(frame.reversed_axes(), key)?; + 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, key: (usize, u8)) -> Result<()> - where T: Bytes + Clone + Zero + Send + 'static { + 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) -> CompressedFrame - where T: Bytes + Clone + Zero { + 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 / 64f64 - ).log2() / 2f64).round() as u32).max(16).min(1024); + let tile_size = 2usize + .pow( + ((image_width as f64 * image_length as f64 / 64f64).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 = byte_tiles.into_par_iter().map(|x| encode_all(&*x, 3).unwrap()).collect::>(); - CompressedFrame { bytes, image_width, image_length, tile_size, - bits_per_sample: T::BITS_PER_SAMPLE, sample_format: T::SAMPLE_FORMAT } - } - self.threads.insert(key, thread::spawn(move || compress(frame))); - for key in self.threads.keys().cloned().collect::>() { - if self.threads[&key].is_finished() { - + let byte_tiles: Vec> = tiles + .into_iter() + .map(|tile| tile.map(|x| x.bytes()).into_iter().flatten().collect()) + .collect(); + let bytes = byte_tiles + .into_par_iter() + .map(|x| encode_all(&*x, 3).unwrap()) + .collect::>(); + CompressedFrame { + bytes, + image_width, + image_length, + tile_size, + bits_per_sample: T::BITS_PER_SAMPLE, + sample_format: T::SAMPLE_FORMAT, } } + self.threads + .insert((c, z, t), thread::spawn(move || compress(frame))); + for key in self + .threads + .keys() + .cloned() + .collect::>() + { + if self.threads[&key].is_finished() {} + } - for key in self.threads.keys().cloned().collect::>() { - if self.threads[&key].is_finished() { - if let Some(thread) = self.threads.remove(&key) { - self.write_frame(thread.join().unwrap(), key)?; + for (c, z, t) in self.threads.keys().cloned().collect::>() { + if 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(()) + Ok(()) } - fn write_frame(&mut self, frame: CompressedFrame, key: (usize, u8)) -> Result<()> { + fn write_frame(&mut self, frame: CompressedFrame, c: usize, z: usize, t: usize) -> Result<()> { let mut tileoffsets = Vec::new(); let mut tilebytecounts = Vec::new(); for tile in frame.bytes { tilebytecounts.push(tile.len() as u64); tileoffsets.push(self.write(&tile)?); } - let frame = Frame::new(tileoffsets, tilebytecounts, 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(key, frame); + let frame = Frame::new( + tileoffsets, + tilebytecounts, + 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(()) } @@ -499,59 +678,74 @@ impl IJTiffFile { let (n, m) = (shape[0] / size, shape[1] / size); 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()); + 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..])); + 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])); + 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..])); + 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) -> Result> { - todo!(); + fn get_colormap(&self, colormap: &Vec>, bits_per_sample: u16) -> Vec { + if bits_per_sample == 8 { + colormap + .iter() + .flatten() + .map(|x| (*x as u16) * 256) + .collect() + } else { + colormap + .iter() + .map(|x| vec![x; 256]) + .flatten() + .flatten() + .map(|x| (*x as u16) * 256) + .collect() + } } - fn get_color(&self, _colors: (u8, u8, u8)) -> Result> { + fn get_color(&self, _colors: &Vec, _bits_per_sample: u16) -> Result> { todo!(); } fn close(&mut self) -> Result<()> { - for key in self.threads.keys().cloned().collect::>() { - if let Some(thread) = self.threads.remove(&key) { - self.write_frame(thread.join().unwrap(), key)?; + for (c, z, t) in self.threads.keys().cloned().collect::>() { + if let Some(thread) = self.threads.remove(&(c, z, t)) { + self.write_frame(thread.join().unwrap(), c, z, t)?; } } let mut where_to_write_next_ifd_offset = OFFSET - OFFSET_SIZE as u64; let mut warn = false; - for frame_number in 0..self.n_frames { - if let Some(frame) = self.frames.get(&(frame_number, 0)) { + let (samples_per_pixel, n_frames) = self.spp_and_n_frames(); + for frame_number in 0..n_frames { + if let Some(frame) = self.frames.get(&self.get_czt(frame_number, 0)) { let mut tileoffsets = Vec::new(); let mut tilebytecounts = Vec::new(); let mut frame_count = 0; - for channel in 0..self.samples_per_pixel { - if let Some(frame_n) = self.frames.get(&(frame_number, channel)) { + for channel in 0..samples_per_pixel { + if let Some(frame_n) = self.frames.get(&self.get_czt(frame_number, channel)) { tileoffsets.extend(frame_n.tileoffsets.iter()); tilebytecounts.extend(frame_n.tilebytecounts.iter()); frame_count += 1; @@ -560,52 +754,80 @@ impl IJTiffFile { } } let mut ifd = IFD::new(); - ifd.push_tag(Tag::long(256, vec![frame.image_width])); - ifd.push_tag(Tag::long(257, vec![frame.image_length])); - ifd.push_tag(Tag::short(258, vec![frame.bits_per_sample; frame_count])); - ifd.push_tag(Tag::short(259, vec![COMPRESSION])); + ifd.push_tag(Tag::long(256, &vec![frame.image_width])); + ifd.push_tag(Tag::long(257, &vec![frame.image_length])); + ifd.push_tag(Tag::short(258, &vec![frame.bits_per_sample; frame_count])); + ifd.push_tag(Tag::short(259, &vec![COMPRESSION])); ifd.push_tag(Tag::ascii(270, &self.description())); - ifd.push_tag(Tag::short(277, vec![frame_count as u16])); + ifd.push_tag(Tag::short(277, &vec![frame_count as u16])); ifd.push_tag(Tag::ascii(305, "tiffwrite_rs")); - ifd.push_tag(Tag::short(322, vec![frame.tile_width])); - ifd.push_tag(Tag::short(323, vec![frame.tile_length])); - ifd.push_tag(Tag::long8(324, tileoffsets)); - ifd.push_tag(Tag::long8(325, tilebytecounts)); - ifd.push_tag(Tag::short(339, vec![frame.sample_format])); + ifd.push_tag(Tag::short(322, &vec![frame.tile_width])); + ifd.push_tag(Tag::short(323, &vec![frame.tile_length])); + ifd.push_tag(Tag::long8(324, &tileoffsets)); + ifd.push_tag(Tag::long8(325, &tilebytecounts)); + ifd.push_tag(Tag::short(339, &vec![frame.sample_format])); + if let Some(px_size) = self.px_size { + let r = vec![Rational32::from_f64(px_size).unwrap()]; + ifd.push_tag(Tag::rational(282, &r)); + ifd.push_tag(Tag::rational(283, &r)); + ifd.push_tag(Tag::short(296, &vec![1])); + } + if frame_number == 0 { - if let Some(colormap) = &self.colormap { - ifd.push_tag(Tag::short(320, self.get_colormap(colormap)?)); - ifd.push_tag(Tag::short(262, vec![3])); // PhotometricInterpretation PHOTOMETRIC_PALETTE - } else if let None = self.colors { - ifd.push_tag(Tag::short(262, vec![1])); // PhotometricInterpretation PHOTOMETRIC_PALETTE + if let Colors::Colormap(colormap) = &self.colors { + ifd.push_tag(Tag::short( + 320, + &self.get_colormap(colormap, frame.bits_per_sample), + )); + ifd.push_tag(Tag::short(262, &vec![3])); + } else if let Colors::None = self.colors { + ifd.push_tag(Tag::short(262, &vec![1])); } } - if frame_number < self.samples_per_pixel as usize { - if let Some(color) = &self.colors { - ifd.push_tag(Tag::short(320, self.get_color(color[frame_number])?)); - ifd.push_tag(Tag::short(262, vec![3])); // PhotometricInterpretation PHOTOMETRIC_PALETTE + if frame_number < samples_per_pixel as usize { + if let Colors::Colors(colors) = &self.colors { + ifd.push_tag(Tag::short( + 320, + &self.get_color(&colors[frame_number], frame.bits_per_sample)?, + )); + ifd.push_tag(Tag::short(262, &vec![3])); } } - if let (None, None) = (&self.colormap, &self.colors) { + if let Colors::None = &self.colors { if self.shape.0 > 1 { - ifd.push_tag(Tag::short(284, vec![2])) + ifd.push_tag(Tag::short(284, &vec![2])) } } - if self.extra_tags_frame.contains_key(&frame_number) { - ifd.extend_tags(self.extra_tags_frame[&frame_number].to_owned()); + for channel in 0..samples_per_pixel { + let czt = self.get_czt(frame_number, channel); + if let Some(extra_tags) = self.extra_tags.get(&Some(czt)) { + for tag in extra_tags { + ifd.push_tag(tag.to_owned()) + } + } } - ifd.extend_tags(self.extra_tags.to_owned()); - ifd.push_tag(Tag::ascii(306, &format!("{}", Utc::now().format("%Y:%m:%d %H:%M:%S")))); + if let Some(extra_tags) = self.extra_tags.get(&None) { + for tag in extra_tags { + ifd.push_tag(tag.to_owned()) + } + } + ifd.push_tag(Tag::ascii( + 306, + &format!("{}", Utc::now().format("%Y:%m:%d %H:%M:%S")), + )); where_to_write_next_ifd_offset = ifd.write(self, where_to_write_next_ifd_offset)?; } else { warn = true; } if warn { - println!("Some frames were not added to the tif file, either you forgot them, \ - or an error occurred and the tif file was closed prematurely.") + println!( + "Some frames were not added to the tif file, 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 + .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 index 658eac3..75a03da 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,7 +2,6 @@ use anyhow::Result; use ndarray::{s, Array2}; use tiffwrite::IJTiffFile; - fn main() -> Result<()> { println!("Hello World!"); let mut f = IJTiffFile::new("foo.tif", (2, 1, 1))?; @@ -12,12 +11,12 @@ fn main() -> Result<()> { arr[[i, j]] = i as u16; } } - f.save(arr.to_owned(), 0, 0, 0, None)?; + 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, None)?; + f.save(arr.to_owned(), 1, 0, 0)?; Ok(()) -} \ No newline at end of file +} diff --git a/src/py.rs b/src/py.rs index 3c91f4b..693b338 100644 --- a/src/py.rs +++ b/src/py.rs @@ -1,101 +1,153 @@ +use crate::{Colors, IJTiffFile, Tag}; +use ndarray::s; +use num::{Complex, FromPrimitive, Rational32}; +use numpy::{PyArrayMethods, PyReadonlyArray2}; use pyo3::prelude::*; -use crate::{IJTiffFile, Tag}; -use num::{Complex, Rational32, FromPrimitive}; -use numpy::{PyReadonlyArray2, PyArrayMethods}; - #[pyclass(subclass)] #[pyo3(name = "Tag")] #[derive(Clone, Debug)] struct PyTag { - tag: Tag + tag: Tag, } #[pymethods] impl PyTag { #[staticmethod] fn byte(code: u16, byte: Vec) -> Self { - PyTag { tag: Tag::byte(code, byte) } + PyTag { + tag: Tag::byte(code, &byte), + } } #[staticmethod] fn ascii(code: u16, ascii: &str) -> Self { - PyTag { tag: Tag::ascii(code, ascii) } + PyTag { + tag: Tag::ascii(code, ascii), + } } #[staticmethod] fn short(code: u16, short: Vec) -> Self { - PyTag { tag: Tag::short(code, short) } + PyTag { + tag: Tag::short(code, &short), + } } #[staticmethod] fn long(code: u16, long: Vec) -> Self { - PyTag { tag: Tag::long(code, long) } + 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()) } + 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) } + PyTag { + tag: Tag::sbyte(code, &sbyte), + } } #[staticmethod] fn sshort(code: u16, sshort: Vec) -> Self { - PyTag { tag: Tag::sshort(code, sshort) } + PyTag { + tag: Tag::sshort(code, &sshort), + } } #[staticmethod] fn slong(code: u16, slong: Vec) -> Self { - PyTag { tag: Tag::slong(code, slong) } + 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()) } + 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) } + PyTag { + tag: Tag::float(code, &float), + } } #[staticmethod] fn double(code: u16, double: Vec) -> Self { - PyTag { tag: Tag::double(code, double) } + PyTag { + tag: Tag::double(code, &double), + } } #[staticmethod] fn ifd(code: u16, ifd: Vec) -> Self { - PyTag { tag: Tag::ifd(code, ifd) } + PyTag { + tag: Tag::ifd(code, &ifd), + } } #[staticmethod] fn unicode(code: u16, unicode: &str) -> Self { - PyTag { tag: Tag::unicode(code, unicode) } + 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()) } + 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) } + PyTag { + tag: Tag::long8(code, &long8), + } } #[staticmethod] fn slong8(code: u16, slong8: Vec) -> Self { - PyTag { tag: Tag::slong8(code, slong8) } + PyTag { + tag: Tag::slong8(code, &slong8), + } } #[staticmethod] fn ifd8(code: u16, ifd8: Vec) -> Self { - PyTag { tag: Tag::ifd8(code, ifd8) } + PyTag { + tag: Tag::ifd8(code, &ifd8), + } } fn count(&self) -> u64 { @@ -103,55 +155,157 @@ impl PyTag { } } - #[pyclass(subclass)] #[pyo3(name = "IJTiffFile")] #[derive(Debug)] struct PyIJTiffFile { - ijtifffile: Option + ijtifffile: Option, } #[pymethods] impl PyIJTiffFile { #[new] fn new(path: &str, shape: (usize, usize, usize)) -> PyResult { - Ok(PyIJTiffFile { ijtifffile: Some(IJTiffFile::new(path, shape)?) } ) + Ok(PyIJTiffFile { + ijtifffile: Some(IJTiffFile::new(path, shape)?), + }) } - fn with_colors(&mut self, colors: (u8, u8, u8)) -> Self { - todo!() + #[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) } - fn with_colormap(&mut self, colormap: Vec<(u8, u8, u8)>) -> Self { - todo!() + #[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(()) } - fn with_px_size(&mut self, pxsize: f64) -> Self { - todo!() + #[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) } - fn with_delta_z(&mut self, delta_z: f64) -> Self { - todo!() + #[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(()) } - fn with_time_interval(&mut self, time_interval: f64) -> Self { - todo!() - } - - fn with_comments(&mut self, comments: String) -> Self { - todo!() - } - - fn append_extra_tag(&mut self, tag: PyTag) { - if let Some(ijtifffile) = self.ijtifffile.as_mut() { - ijtifffile.extra_tags.push(tag.tag); + #[getter] + fn get_px_size(&self) -> PyResult> { + if let Some(ijtifffile) = &self.ijtifffile { + Ok(ijtifffile.px_size) + } else { + Ok(None) } } - fn extend_extra_tags(&mut self, tags: Vec) { - if let Some(ijtifffile) = self.ijtifffile.as_mut() { - ijtifffile.extra_tags.extend(tags.into_iter().map(|x| x.tag)); + #[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<()> { @@ -160,20 +314,19 @@ impl PyIJTiffFile { } } - macro_rules! impl_save { ($T:ty, $t:ident) => { #[pymethods] impl PyIJTiffFile { - fn $t(&mut self, frame: PyReadonlyArray2<$T>, c: usize, t: usize, z: usize, - extra_tags: Option>) -> PyResult<()> { - let extra_tags = if let Some(extra_tags) = extra_tags { - Some(extra_tags.into_iter().map(|x| x.tag).collect()) - } else { - None - }; + 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, extra_tags)?; + ijtifffile.save(frame.to_owned_array(), c, t, z)?; } Ok(()) } @@ -192,7 +345,6 @@ 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<()> {