diff --git a/src/polystore/config.py b/src/polystore/config.py index af36d0e..87cef68 100644 --- a/src/polystore/config.py +++ b/src/polystore/config.py @@ -133,44 +133,106 @@ def __new__(cls, value: str, tifffile_codec: str | None): return member +class TiffPhotometric(Enum): + """Declared TIFF pixel interpretation, independent of array dimensions.""" + + MINISBLACK = "minisblack" + RGB = "rgb" + + +class TiffPlanarConfig(Enum): + """Declared storage layout for TIFF sample dimensions.""" + + CONTIG = "contig" + SEPARATE = "separate" + + @dataclass(frozen=True) class TiffConfig: - """TIFF writer settings; defaults preserve the uncompressed file contract. + """TIFF codec and declared pixel axes; defaults preserve legacy writes. Args: compression: Lossless disk-TIFF codec. ``NONE`` preserves the existing uncompressed output; ``DEFLATE`` compresses TIFF pixels exactly. compression_level: Deflate level from 1 to 9. Ignored when compression is ``NONE``. + photometric: Explicit grayscale or RGB pixel interpretation. + axes: Dimension identities such as ``ZYX`` or ``YXS``. Requires + photometric; ``S`` declares samples rather than a spatial axis. + planarconfig: Sample layout. Declared axes select the layout when absent. """ compression: TiffCompression = TiffCompression.NONE compression_level: int = 1 + photometric: TiffPhotometric | None = None + axes: str | None = None + planarconfig: TiffPlanarConfig | None = None def __post_init__(self) -> None: if not isinstance(self.compression, TiffCompression): raise TypeError("compression must be a TiffCompression value") - if isinstance(self.compression_level, bool) or not isinstance( - self.compression_level, int - ) or not 1 <= self.compression_level <= 9: + if ( + isinstance(self.compression_level, bool) + or not isinstance(self.compression_level, int) + or not 1 <= self.compression_level <= 9 + ): raise ValueError("compression_level must be an integer from 1 to 9") + if self.photometric is not None and not isinstance(self.photometric, TiffPhotometric): + raise TypeError("photometric must be a TiffPhotometric value") + if self.planarconfig is not None and not isinstance(self.planarconfig, TiffPlanarConfig): + raise TypeError("planarconfig must be a TiffPlanarConfig value") + if self.axes is not None: + if ( + not isinstance(self.axes, str) + or not self.axes + or not self.axes.isascii() + or not self.axes.isalpha() + or not self.axes.isupper() + or len(set(self.axes)) != len(self.axes) + ): + raise ValueError("axes must contain distinct uppercase dimension identities") + if self.photometric is None: + raise ValueError("axes require an explicit photometric interpretation") + sample_axis = self.axes.find("S") + if sample_axis >= 0: + if sample_axis == len(self.axes) - 1: + layout = TiffPlanarConfig.CONTIG + elif sample_axis == len(self.axes) - 3: + layout = TiffPlanarConfig.SEPARATE + else: + raise ValueError("TIFF sample axis must be last or immediately before YX") + if self.planarconfig is not None and self.planarconfig is not layout: + raise ValueError("planarconfig disagrees with the declared sample axis") + object.__setattr__(self, "planarconfig", layout) + elif self.photometric is TiffPhotometric.RGB or self.planarconfig is not None: + raise ValueError("RGB and planar TIFF layouts require a declared S sample axis") + + def validate_shape(self, shape: Sequence[int]) -> None: + """Reject a declared axis roster that cannot describe the actual pixels.""" + if self.axes is not None and len(self.axes) != len(shape): + raise ValueError( + f"Declared TIFF axes {self.axes!r} do not match pixel shape {tuple(shape)!r}" + ) def tifffile_write_kwargs(self) -> dict[str, object]: - """Return only codec options accepted by tifffile.imwrite.""" - + """Encode codec options and declared pixel semantics for tifffile.imwrite.""" + options: dict[str, object] = {} codec = self.compression.tifffile_codec - if codec is None: - return {} - return { - "compression": codec, - "compressionargs": {"level": self.compression_level}, - } + if codec is not None: + options.update(compression=codec, compressionargs={"level": self.compression_level}) + if self.photometric is not None: + options["photometric"] = self.photometric.value + if self.axes is not None: + options["metadata"] = {"axes": self.axes} + if self.planarconfig is not None: + options["planarconfig"] = self.planarconfig.value + return options def applies_to_path(self, path: str | Path) -> bool: - """Whether this configured codec applies to this disk output path.""" + """Whether these configured TIFF options apply to this disk path.""" return ( - self.compression is not TiffCompression.NONE + bool(self.tifffile_write_kwargs()) and Path(path).suffix.lower() in FileFormat.TIFF.extensions ) @@ -179,11 +241,11 @@ def tiff_write_batches( paths: Sequence[str | Path], config: TiffConfig | None, ) -> tuple[tuple[tuple[int, ...], TiffConfig | None], ...]: - """Partition consecutive writes by their exact TIFF codec requirement.""" + """Partition consecutive writes by their exact TIFF configuration.""" if not paths: return () - if config is None or config.compression is TiffCompression.NONE: + if config is None or not config.tifffile_write_kwargs(): return ((tuple(range(len(paths))), None),) batches: list[tuple[tuple[int, ...], TiffConfig | None]] = [] start = 0 diff --git a/src/polystore/disk.py b/src/polystore/disk.py index e31daa2..5165451 100644 --- a/src/polystore/disk.py +++ b/src/polystore/disk.py @@ -201,17 +201,15 @@ def _tensorflow_reader(self, path, **kwargs): return tf.io.parse_tensor(tf.io.read_file(path.as_posix()), out_type=tf.dtypes.float32) def _tiff_writer(self, path, data, *, tiff_config: TiffConfig | None = None, **kwargs): - """Write TIFF pixels with optional lossless, per-call codec settings.""" + """Write pixels with their declared TIFF semantics and codec options.""" tifffile = FileFormat.TIFF.load_dependency() config = TiffConfig() if tiff_config is None else tiff_config if not isinstance(config, TiffConfig): raise TypeError("tiff_config must be a TiffConfig value") - tifffile.imwrite( - path, - storage_numpy_array(data), - **config.tifffile_write_kwargs(), - ) + pixels = storage_numpy_array(data) + config.validate_shape(pixels.shape) + tifffile.imwrite(path, pixels, **config.tifffile_write_kwargs()) def _tiff_reader(self, path): tifffile = FileFormat.TIFF.load_dependency() diff --git a/tests/test_disk_backend.py b/tests/test_disk_backend.py index a452b08..618c357 100644 --- a/tests/test_disk_backend.py +++ b/tests/test_disk_backend.py @@ -13,7 +13,13 @@ from polystore import FileManager from polystore import disk as disk_module -from polystore.config import TiffCompression, TiffConfig, tiff_write_batches +from polystore.config import ( + TiffCompression, + TiffConfig, + TiffPhotometric, + TiffPlanarConfig, + tiff_write_batches, +) from polystore.disk import DiskBackend from polystore.exceptions import StorageResolutionError @@ -137,6 +143,69 @@ def test_tiff_deflate_is_lossless_and_opt_in(tmp_path: Path) -> None: np.testing.assert_array_equal(tifffile.imread(compressed_path), pixels) +@pytest.mark.parametrize("depth", [3, 4]) +def test_declared_grayscale_volume_does_not_become_rgb(tmp_path, depth) -> None: + pixels = np.arange(depth * 4 * 5, dtype=np.uint16).reshape(depth, 4, 5) + path = tmp_path / "volume.tif" + DiskBackend().save( + pixels, + path, + tiff_config=TiffConfig(photometric=TiffPhotometric.MINISBLACK, axes="ZYX"), + ) + + with tifffile.TiffFile(path) as image: + assert image.series[0].axes == "ZYX" + assert len(image.pages) == depth + assert image.pages[0].photometric.name == "MINISBLACK" + assert image.pages[0].samplesperpixel == 1 + assert image.pages[0].compression.value == 1 + np.testing.assert_array_equal(image.asarray(), pixels) + + +@pytest.mark.parametrize( + ("axes", "shape", "layout"), + [ + ("YXS", (4, 5, 3), TiffPlanarConfig.CONTIG), + ("SYX", (3, 4, 5), TiffPlanarConfig.SEPARATE), + ("ZSYX", (2, 3, 4, 5), TiffPlanarConfig.SEPARATE), + ], +) +def test_declared_rgb_preserves_samples_and_spatial_axes(tmp_path, axes, shape, layout) -> None: + pixels = np.arange(np.prod(shape), dtype=np.uint16).reshape(shape) + path = tmp_path / "rgb.tif" + DiskBackend().save( + pixels, + path, + tiff_config=TiffConfig(photometric=TiffPhotometric.RGB, axes=axes), + ) + + with tifffile.TiffFile(path) as image: + assert image.series[0].axes == axes + assert image.pages[0].photometric.name == "RGB" + assert image.pages[0].planarconfig.name == layout.name + assert image.pages[0].samplesperpixel == 3 + np.testing.assert_array_equal(image.asarray(), pixels) + + +def test_declared_axes_must_describe_actual_storage_layout(tmp_path) -> None: + path = tmp_path / "invalid.tif" + with pytest.raises(ValueError, match="do not match pixel shape"): + DiskBackend().save( + np.zeros((3, 4, 5)), + path, + tiff_config=TiffConfig(photometric=TiffPhotometric.MINISBLACK, axes="YX"), + ) + assert not path.exists() + with pytest.raises(ValueError, match="disagrees"): + TiffConfig( + photometric=TiffPhotometric.RGB, + axes="SYX", + planarconfig=TiffPlanarConfig.CONTIG, + ) + with pytest.raises(ValueError, match="sample axis"): + TiffConfig(photometric=TiffPhotometric.RGB, axes="SZYX") + + def test_tiff_write_batches_preserve_order_and_exclude_other_formats() -> None: config = TiffConfig(compression=TiffCompression.DEFLATE) assert tiff_write_batches(("a.csv", "b.tif", "c.tiff", "d.npy"), config) == ( @@ -144,8 +213,12 @@ def test_tiff_write_batches_preserve_order_and_exclude_other_formats() -> None: ((1, 2), config), ((3,), None), ) - assert tiff_write_batches(("a.csv", "b.tif"), TiffConfig()) == ( - ((0, 1), None), + assert tiff_write_batches(("a.csv", "b.tif"), TiffConfig()) == (((0, 1), None),) + declared = TiffConfig(photometric=TiffPhotometric.MINISBLACK, axes="ZYX") + assert tiff_write_batches(("a.csv", "b.tif", "c.tiff", "d.npy"), declared) == ( + ((0,), None), + ((1, 2), declared), + ((3,), None), ) with pytest.raises(ValueError, match="compression_level"): TiffConfig(compression_level=10)