util to convert the masks to bool + ruff bump

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

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

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()

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

docs/installation.md

@@ -40,5 +40,6 @@ You can verify the installation by importing the library:
 
 ```python
 import tracksdata as td
+
 print(td.__version__)
 ```

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.

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_prop_dtype, geff_prop_dtype
 
-__all__ = ["compressed_tracks_table", "from_ctc", "to_ctc"]
+__all__ = [
+    "compressed_tracks_table",
+    "convert_geff_prop_dtype",
+    "from_ctc",
+    "geff_prop_dtype",
+    "to_ctc",
+]

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.

src/tracksdata/io/_geff_masks.py

@@ -0,0 +1,221 @@
+"""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_prop_dtype", "geff_prop_dtype"]
+
+
+def geff_prop_dtype(
+    geff_store: StoreLike,
+    key: str = DEFAULT_ATTR_KEYS.MASK,
+    *,
+    node_or_edge: str = "nodes",
+) -> np.dtype | None:
+    """Return the on-disk payload dtype of a property in a geff store.
+
+    Parameters
+    ----------
+    geff_store : StoreLike
+        The store or path to the geff group.
+    key : str
+        Which property to inspect. Defaults to the standard mask key.
+    node_or_edge : str
+        Whether ``key`` is a node (``"nodes"``, default) or edge (``"edges"``)
+        property.
+
+    Returns
+    -------
+    np.dtype | None
+        The dtype of the property's payload buffer (``data`` for variable-length
+        properties such as masks, otherwise ``values``), or ``None`` if there is
+        no such property. Compare against ``np.bool_`` to decide whether
+        :func:`convert_geff_prop_dtype` is worth running.
+    """
+    if node_or_edge not in ("nodes", "edges"):
+        raise ValueError(f"node_or_edge must be 'nodes' or 'edges', got {node_or_edge!r}.")
+    root = zarr.open_group(geff_store, mode="r")
+    try:
+        group = root[f"{node_or_edge}/props/{key}"]
+        payload = "data" if "data" in group else "values"
+        return np.dtype(group[payload].dtype)
+    except KeyError:
+        return None
+
+
+def convert_geff_prop_dtype(
+    geff_store: StoreLike,
+    key: str,
+    dtype: np.typing.DTypeLike,
+    *,
+    node_or_edge: str = "nodes",
+    output_path: StoreLike | None = None,
+) -> bool:
+    """Rewrite one property's payload buffer to ``dtype``.
+
+    Only the payload buffer is rewritten (``data`` for variable-length
+    properties such as masks, otherwise ``values``); for variable-length
+    properties the ``values`` offset/shape array is left untouched. The buffer
+    is read one native zarr chunk at a time so the full 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.
+    key : str
+        The property to convert.
+    dtype : np.typing.DTypeLike
+        The target dtype for the payload buffer.
+    node_or_edge : str
+        Whether ``key`` is a node (``"nodes"``, default) or edge (``"edges"``)
+        property.
+    output_path : StoreLike | None
+        If given, the geff is first copied to this store/path and the copy is
+        converted, leaving the original untouched (safer for shared data). If
+        ``None`` (default), the conversion is done in place.
+
+    Returns
+    -------
+    bool
+        ``True`` if a conversion was performed, ``False`` if the payload already
+        had ``dtype``.
+
+    Raises
+    ------
+    KeyError
+        If ``key`` is not a property of ``node_or_edge`` in the geff.
+    """
+    if node_or_edge not in ("nodes", "edges"):
+        raise ValueError(f"node_or_edge must be 'nodes' or 'edges', got {node_or_edge!r}.")
+    dtype = np.dtype(dtype)
+
+    if output_path is not None:
+        geff_store = _copy_geff(geff_store, output_path)
+
+    root = zarr.open_group(geff_store, mode="r+")
+    try:
+        group = root[f"{node_or_edge}/props/{key}"]
+    except KeyError as e:
+        raise KeyError(f"{key!r} is not a {node_or_edge[:-1]} property in this geff.") from e
+
+    payload = "data" if "data" in group else "values"
+    old = group[payload]
+    if np.dtype(old.dtype) == dtype:
+        return False
+
+    n = int(old.shape[0])
+    LOG.info("Converting geff %s %r %s (%d elements) from %s to %s", node_or_edge, key, payload, n, old.dtype, dtype)
+
+    # Read one native chunk (along the first axis) at a time so the full buffer
+    # is never in memory at once.
+    buf = np.empty(old.shape, dtype=dtype)
+    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(dtype)
+
+    _overwrite_array(group, payload, buf, old.chunks)
+    _set_prop_metadata_dtype(root, node_or_edge, key, dtype)
+    return True
+
+
+def _copy_geff(geff_store: StoreLike, output_path: StoreLike) -> StoreLike:
+    """Copy a geff store to ``output_path`` and return the destination.
+
+    Uses a fast filesystem copy when both ends are local paths; otherwise falls
+    back to a store-agnostic key-by-key copy (zarr v3 has no working
+    ``copy_store``), so in-memory and remote stores work too.
+    """
+    src = _as_path(geff_store)
+    dst = _as_path(output_path)
+    if src is not None and dst is not None:
+        if dst.exists():
+            raise FileExistsError(f"output_path already exists: {dst}")
+        shutil.copytree(src, dst)
+        return dst
+
+    # Fallback to a store-agnostic copy for in-memory or remote stores.
+    _copy_store_contents(geff_store, output_path)
+    return output_path
+
+
+def _as_path(store: StoreLike) -> Path | None:
+    """Return ``store`` as a filesystem ``Path`` if it is one, else ``None``."""
+    try:
+        return Path(os.fspath(store))
+    except TypeError:
+        return None
+
+
+def _copy_store_contents(src: StoreLike, dst: StoreLike) -> None:
+    """Copy every key from one geff store to another, key by key."""
+    from zarr.core.buffer import default_buffer_prototype
+    from zarr.core.sync import sync
+
+    src_store = zarr.open_group(src, mode="r").store
+    dst_store = zarr.open_group(dst, mode="a").store
+    prototype = default_buffer_prototype()
+
+    async def _run() -> None:
+        async for k in src_store.list():
+            value = await src_store.get(k, prototype=prototype)
+            if value is not None:
+                await dst_store.set(k, value)
+
+    sync(_run())
+
+
+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_prop_metadata_dtype(root: zarr.Group, node_or_edge: str, key: str, dtype: np.dtype) -> None:
+    """Update the geff property metadata so ``key``'s dtype reads as ``dtype``."""
+    geff_meta = root.attrs.get("geff")
+    if not geff_meta:
+        return
+    meta_key = "node_props_metadata" if node_or_edge == "nodes" else "edge_props_metadata"
+    props = geff_meta.get(meta_key, {})
+    prop_meta = props.get(key)
+    if prop_meta is not None and prop_meta.get("dtype") != dtype.name:
+        prop_meta["dtype"] = dtype.name
+        # Reassign to persist the nested change back to the store.
+        root.attrs["geff"] = geff_meta