diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6f51714b..64f4c90e 100755 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -36,12 +36,12 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: v0.11.12 + rev: v0.16.0 hooks: # Run the linter. - id: ruff - types_or: [python, pyi, jupyter] + types_or: [python, pyi, jupyter, markdown] args: [--fix] # Run the formatter. - id: ruff-format - types_or: [python, pyi, jupyter] + types_or: [python, pyi, jupyter, markdown] diff --git a/docs/examples.md b/docs/examples.md index 496ff285..4b929e8a 100644 --- a/docs/examples.md +++ b/docs/examples.md @@ -7,7 +7,7 @@ This section provides practical examples of using TracksData for multi-object tr Here's a complete basic example that demonstrates the core workflow of TracksData. This example is available as an executable Python file at [`docs/examples/basic.py`](examples/basic.py). ```python ---8<-- "docs/examples/basic.py" +--8 < --"docs/examples/basic.py" ``` ## Key Components Explained diff --git a/docs/faq.md b/docs/faq.md index e54dac97..ad243543 100644 --- a/docs/faq.md +++ b/docs/faq.md @@ -47,6 +47,7 @@ Inherit from :class:`tracksdata.edges.BaseEdgesOperator` or :class:`tracksdata.n ```python import tracksdata as td + class CustomNodes(td.nodes.BaseNodesOperator): def add_nodes( self, @@ -69,7 +70,9 @@ import tracksdata as td labels = ... tracks_df, track_graph, track_labels = td.functional.to_napari_format( - solution_graph, shape=labels.shape, mask_key="mask", + solution_graph, + shape=labels.shape, + mask_key="mask", ) viewer = napari.Viewer() diff --git a/docs/getting_started.md b/docs/getting_started.md index a025bd4f..055b0af1 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -20,18 +20,11 @@ import tracksdata as td graph = td.graph.InMemoryGraph() # Generate random nodes for testing -node_generator = td.nodes.RandomNodes( - n_time_points=5, - n_nodes_per_tp=(10, 15), - n_dim=2 -) +node_generator = td.nodes.RandomNodes(n_time_points=5, n_nodes_per_tp=(10, 15), n_dim=2) node_generator.add_nodes(graph) # Connect nearby nodes across time -edge_generator = td.edges.DistanceEdges( - distance_threshold=0.3, - n_neighbors=3 -) +edge_generator = td.edges.DistanceEdges(distance_threshold=0.3, n_neighbors=3) edge_generator.add_edges(graph) # Solve the tracking problem diff --git a/docs/installation.md b/docs/installation.md index b55d8683..148803b3 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -40,5 +40,6 @@ You can verify the installation by importing the library: ```python import tracksdata as td + print(td.__version__) ``` diff --git a/src/tracksdata/graph/_base_graph.py b/src/tracksdata/graph/_base_graph.py index d6317b78..3a214e35 100644 --- a/src/tracksdata/graph/_base_graph.py +++ b/src/tracksdata/graph/_base_graph.py @@ -1055,7 +1055,7 @@ def to_ctc( shape: tuple[int, ...] | None = None, tracklet_id_key: str = DEFAULT_ATTR_KEYS.TRACKLET_ID, overwrite: bool = False, - dtype: None | DTypeLike = None, + dtype: DTypeLike | None = None, ) -> None: """ Save the graph to a CTC ground truth directory. diff --git a/src/tracksdata/io/__init__.py b/src/tracksdata/io/__init__.py index 1165917b..e5dc3df9 100644 --- a/src/tracksdata/io/__init__.py +++ b/src/tracksdata/io/__init__.py @@ -1,5 +1,12 @@ """Input/output utilities for loading and saving tracking data in various formats.""" from tracksdata.io._ctc import compressed_tracks_table, from_ctc, to_ctc +from tracksdata.io._geff_masks import convert_geff_mask_to_bool, geff_mask_dtype -__all__ = ["compressed_tracks_table", "from_ctc", "to_ctc"] +__all__ = [ + "compressed_tracks_table", + "convert_geff_mask_to_bool", + "from_ctc", + "geff_mask_dtype", + "to_ctc", +] diff --git a/src/tracksdata/io/_ctc.py b/src/tracksdata/io/_ctc.py index baf68281..c5dc6ae3 100644 --- a/src/tracksdata/io/_ctc.py +++ b/src/tracksdata/io/_ctc.py @@ -236,7 +236,7 @@ def to_ctc( shape: tuple[int, ...] | None = None, tracklet_id_key: str = DEFAULT_ATTR_KEYS.TRACKLET_ID, overwrite: bool = False, - dtype: None | DTypeLike = None, + dtype: DTypeLike | None = None, ) -> None: """ Save a graph to a CTC data directory. diff --git a/src/tracksdata/io/_geff_masks.py b/src/tracksdata/io/_geff_masks.py new file mode 100644 index 00000000..a8392f42 --- /dev/null +++ b/src/tracksdata/io/_geff_masks.py @@ -0,0 +1,168 @@ +"""Utilities for working with segmentation masks stored in geff files. + +Segmentation masks are binary, but geff files written by older versions of +tracksdata stored the mask ``data`` buffer as ``uint64`` (see +https://github.com/royerlab/tracksdata/pull/318). That is 8x larger than a +boolean buffer both on disk and, more importantly, when read into memory, +which can cause out-of-memory errors when loading large datasets. + +New files store masks as ``bool`` at write time. This module provides a +one-time conversion for legacy files so they can be read cheaply. + +The caller names the mask attribute to convert (defaulting to the standard +``DEFAULT_ATTR_KEYS.MASK`` key). Nothing on disk distinguishes a mask from any +other variable-length attribute, so these functions never guess which +attributes are masks — call once per mask key. +""" + +from __future__ import annotations + +import os +import shutil +from pathlib import Path + +import numpy as np +import zarr +from zarr.storage import StoreLike + +from tracksdata.constants import DEFAULT_ATTR_KEYS +from tracksdata.utils._logging import LOG + +__all__ = ["convert_geff_mask_to_bool", "geff_mask_dtype"] + + +def geff_mask_dtype(geff_store: StoreLike, mask_key: str = DEFAULT_ATTR_KEYS.MASK) -> np.dtype | None: + """Return the on-disk dtype of a mask ``data`` buffer in a geff store. + + Parameters + ---------- + geff_store : StoreLike + The store or path to the geff group. + mask_key : str + Which mask attribute to inspect. Defaults to the standard mask key. + + Returns + ------- + np.dtype | None + The dtype of the mask ``data`` array, or ``None`` if there is no such + variable-length attribute. Compare against ``np.bool_`` to decide + whether :func:`convert_geff_mask_to_bool` is worth running. + """ + root = zarr.open_group(geff_store, mode="r") + try: + return np.dtype(root[f"nodes/props/{mask_key}/data"].dtype) + except KeyError: + return None + + +def convert_geff_mask_to_bool( + geff_store: StoreLike, + mask_key: str = DEFAULT_ATTR_KEYS.MASK, + *, + output_path: StoreLike | None = None, +) -> bool: + """Rewrite one mask attribute's ``data`` buffer to ``bool``. + + Only the mask ``data`` buffer is rewritten; the ``values`` (offset/shape) + array is left untouched. The buffer is read one native zarr chunk at a time + so the full non-boolean buffer is never materialized at once. The geff + metadata dtype is updated to match. + + Parameters + ---------- + geff_store : StoreLike + The store or path to the geff group. + mask_key : str + The mask attribute to convert. Defaults to the standard mask key. Call + once per mask attribute if a graph carries more than one. + output_path : StoreLike | None + If given, the geff is first copied to this path and the copy is + converted, leaving the original untouched (safer for shared data). Both + ``geff_store`` and ``output_path`` must then be filesystem paths. + If ``None`` (default), the conversion is done in place. + + Returns + ------- + bool + ``True`` if a conversion was performed, ``False`` if the mask was + already boolean. + + Raises + ------ + KeyError + If ``mask_key`` is not a variable-length attribute in the geff. + ValueError + If the attribute's buffer is neither integer nor boolean (i.e. not a + mask), to avoid silently corrupting it. + """ + if output_path is not None: + geff_store = _copy_geff(geff_store, output_path) + + root = zarr.open_group(geff_store, mode="r+") + try: + old = root[f"nodes/props/{mask_key}/data"] + except KeyError as e: + raise KeyError(f"{mask_key!r} is not a variable-length (mask) attribute in this geff.") from e + + dtype = np.dtype(old.dtype) + if dtype == np.bool_: + return False + if not np.issubdtype(dtype, np.integer): + raise ValueError(f"Refusing to boolify mask {mask_key!r}: buffer dtype {dtype} is not integer.") + + n = int(old.shape[0]) + LOG.info("Converting geff mask %r data (%d elements) from %s to bool", mask_key, n, dtype) + + # Read one native chunk at a time so the full non-boolean buffer is never + # in memory at once; the resulting boolean buffer is 1/8th its size. + buf = np.empty(n, dtype=bool) + step = old.chunks[0] + for i in range(0, n, step): + j = min(i + step, n) + buf[i:j] = np.asarray(old[i:j]).astype(bool) + + _overwrite_array(root[f"nodes/props/{mask_key}"], "data", buf, old.chunks) + _set_mask_metadata_bool(root, mask_key) + return True + + +def _copy_geff(geff_store: StoreLike, output_path: StoreLike) -> Path: + """Copy a geff directory to ``output_path`` and return the new path.""" + try: + src = Path(os.fspath(geff_store)) + dst = Path(os.fspath(output_path)) + except TypeError as e: + raise TypeError("output_path is only supported for filesystem-path geff stores.") from e + if dst.exists(): + raise FileExistsError(f"output_path already exists: {dst}") + shutil.copytree(src, dst) + return dst + + +def _overwrite_array(parent: zarr.Group, name: str, data: np.ndarray, chunks) -> None: + """Create (overwriting any existing) a zarr array and write ``data`` to it. + + Works with both zarr v2 and v3 array-creation APIs. + """ + try: + # zarr v3 + arr = parent.create_array(name=name, shape=data.shape, chunks=chunks, dtype=data.dtype, overwrite=True) + except (TypeError, AttributeError): + # zarr v2 + if name in parent: + del parent[name] + arr = parent.create_dataset(name, shape=data.shape, chunks=chunks, dtype=data.dtype, overwrite=True) + arr[:] = data + + +def _set_mask_metadata_bool(root: zarr.Group, mask_key: str) -> None: + """Update the geff node-property metadata so the mask dtype reads as bool.""" + geff_meta = root.attrs.get("geff") + if not geff_meta: + return + node_props = geff_meta.get("node_props_metadata", {}) + mask_meta = node_props.get(mask_key) + if mask_meta is not None and mask_meta.get("dtype") != "bool": + mask_meta["dtype"] = "bool" + # Reassign to persist the nested change back to the store. + root.attrs["geff"] = geff_meta diff --git a/src/tracksdata/io/_test/test_geff_masks.py b/src/tracksdata/io/_test/test_geff_masks.py new file mode 100644 index 00000000..de089631 --- /dev/null +++ b/src/tracksdata/io/_test/test_geff_masks.py @@ -0,0 +1,153 @@ +from pathlib import Path + +import numpy as np +import polars as pl +import pytest +import zarr + +from tracksdata.constants import DEFAULT_ATTR_KEYS +from tracksdata.graph import IndexedRXGraph, RustWorkXGraph +from tracksdata.io import convert_geff_mask_to_bool, geff_mask_dtype +from tracksdata.io._geff_masks import _overwrite_array, _set_mask_metadata_bool +from tracksdata.nodes._mask import Mask + +MASK_KEY = DEFAULT_ATTR_KEYS.MASK + + +def _make_masked_geff(tmp_path: Path, name: str = "masks.geff") -> tuple[Path, dict[int, np.ndarray]]: + """Write a small geff with a few masked nodes; return path and node->mask map.""" + graph = RustWorkXGraph() + graph.add_node_attr_key("x", pl.Float64) + graph.add_node_attr_key(DEFAULT_ATTR_KEYS.MASK, pl.Object) + graph.add_node_attr_key(DEFAULT_ATTR_KEYS.BBOX, pl.Array(pl.Int64, 4)) + + masks = [ + np.array([[True, True], [True, False]], dtype=bool), + np.array([[True, False, True]], dtype=bool), + np.array([[True], [True], [True]], dtype=bool), + ] + bboxes = [ + np.array([6, 6, 8, 8]), + np.array([0, 0, 1, 3]), + np.array([2, 5, 5, 6]), + ] + node_masks = {} + for t, (mask, bbox) in enumerate(zip(masks, bboxes, strict=True)): + node_id = graph.add_node({"t": t, "x": float(t), MASK_KEY: Mask(mask, bbox=bbox), "bbox": bbox}) + node_masks[node_id] = mask + + geff_path = tmp_path / name + graph.to_geff(geff_store=geff_path) + return geff_path, node_masks + + +def _downgrade_masks_to_uint64(geff_path: Path, mask_key: str = MASK_KEY) -> None: + """Simulate a legacy (pre-#318) geff by rewriting a mask data buffer as uint64.""" + root = zarr.open_group(geff_path, mode="r+") + data = root[f"nodes/props/{mask_key}/data"] + as_uint64 = np.asarray(data[:]).astype(np.uint64) + _overwrite_array(root[f"nodes/props/{mask_key}"], "data", as_uint64, data.chunks) + geff_meta = root.attrs.get("geff") + if mask_key in geff_meta["node_props_metadata"]: + geff_meta["node_props_metadata"][mask_key]["dtype"] = "uint64" + root.attrs["geff"] = geff_meta + + +def test_convert_geff_mask_to_bool_roundtrip(tmp_path: Path) -> None: + geff_path, node_masks = _make_masked_geff(tmp_path) + _downgrade_masks_to_uint64(geff_path) + + assert geff_mask_dtype(geff_path) == np.uint64 + + assert convert_geff_mask_to_bool(geff_path) is True + assert geff_mask_dtype(geff_path) == np.bool_ + # metadata dtype updated too + root = zarr.open_group(geff_path, mode="r") + assert root.attrs["geff"]["node_props_metadata"][MASK_KEY]["dtype"] == "bool" + + # masks still load correctly and are boolean + graph, _ = IndexedRXGraph.from_geff(geff_path) + for node_id in graph.node_ids(): + loaded = graph.nodes[node_id][MASK_KEY] + assert loaded.mask.dtype == np.bool_ + np.testing.assert_array_equal(loaded.mask, node_masks[node_id]) + + +def test_convert_is_noop_when_already_bool(tmp_path: Path) -> None: + geff_path, _ = _make_masked_geff(tmp_path) + # to_geff already writes bool + assert geff_mask_dtype(geff_path) == np.bool_ + assert convert_geff_mask_to_bool(geff_path) is False + + +def test_convert_is_idempotent(tmp_path: Path) -> None: + geff_path, _ = _make_masked_geff(tmp_path) + _downgrade_masks_to_uint64(geff_path) + assert convert_geff_mask_to_bool(geff_path) is True + # second call is a no-op + assert convert_geff_mask_to_bool(geff_path) is False + + +def test_missing_mask_key_raises(tmp_path: Path) -> None: + geff_path, _ = _make_masked_geff(tmp_path) + assert geff_mask_dtype(geff_path, mask_key="nope") is None + with pytest.raises(KeyError): + convert_geff_mask_to_bool(geff_path, mask_key="nope") + + +def test_output_path_leaves_original_untouched(tmp_path: Path) -> None: + geff_path, node_masks = _make_masked_geff(tmp_path) + _downgrade_masks_to_uint64(geff_path) + + out_path = tmp_path / "converted.geff" + assert convert_geff_mask_to_bool(geff_path, output_path=out_path) is True + + # original is still uint64, copy is bool + assert geff_mask_dtype(geff_path) == np.uint64 + assert geff_mask_dtype(out_path) == np.bool_ + + graph, _ = IndexedRXGraph.from_geff(out_path) + for node_id in graph.node_ids(): + np.testing.assert_array_equal(graph.nodes[node_id][MASK_KEY].mask, node_masks[node_id]) + + +def test_explicit_named_mask_key(tmp_path: Path) -> None: + """A caller with a second mask attribute converts it by naming it explicitly.""" + geff_path, _ = _make_masked_geff(tmp_path) + # add a second variable-length mask attribute (uint64) by copying the first + root = zarr.open_group(geff_path, mode="r+") + props = root["nodes/props"] + src = props[MASK_KEY] + dst = props.create_group("nucleus_mask") + for sub in ("values", "data"): + arr = src[sub] + cast = np.uint64 if sub == "data" else arr.dtype + _overwrite_array(dst, sub, np.asarray(arr[:]).astype(cast), arr.chunks) + + # default call only touches 'mask' + _downgrade_masks_to_uint64(geff_path) + assert convert_geff_mask_to_bool(geff_path) is True + assert geff_mask_dtype(geff_path, "nucleus_mask") == np.uint64 # untouched by default + # explicit key converts the second mask + assert convert_geff_mask_to_bool(geff_path, "nucleus_mask") is True + assert geff_mask_dtype(geff_path, "nucleus_mask") == np.bool_ + + +def test_non_integer_buffer_refused(tmp_path: Path) -> None: + """Naming a non-integer variable-length attribute is refused, not corrupted.""" + geff_path, _ = _make_masked_geff(tmp_path) + root = zarr.open_group(geff_path, mode="r+") + grp = root["nodes/props"].create_group("floatvar") + _overwrite_array(grp, "values", np.zeros((3, 2), dtype=np.uint64), (3, 2)) + _overwrite_array(grp, "data", np.array([1.5, 2.5, 3.5], dtype=np.float64), (3,)) + + with pytest.raises(ValueError, match="not integer"): + convert_geff_mask_to_bool(geff_path, "floatvar") + # left untouched + assert geff_mask_dtype(geff_path, "floatvar") == np.float64 + + +def test_set_mask_metadata_bool_without_geff_attrs(tmp_path: Path) -> None: + """Guard: metadata helper is a no-op when the group has no geff attrs.""" + root = zarr.open_group(tmp_path / "empty.zarr", mode="w") + _set_mask_metadata_bool(root, MASK_KEY) # should not raise diff --git a/src/tracksdata/solvers/_ilp_solver.py b/src/tracksdata/solvers/_ilp_solver.py index 3f6676d5..f4a12aca 100644 --- a/src/tracksdata/solvers/_ilp_solver.py +++ b/src/tracksdata/solvers/_ilp_solver.py @@ -139,7 +139,7 @@ def __init__( appearance_weight: str | ExprInput = 0.0, disappearance_weight: str | ExprInput = 0.0, division_weight: str | ExprInput = 0.0, - merge_weight: None | str | ExprInput = None, + merge_weight: str | ExprInput | None = None, output_key: str = DEFAULT_ATTR_KEYS.SOLUTION, num_threads: int = 1, reset: bool = True,