From ff08c0a9ad67da54797a33d9f18f93dc7d24aefe Mon Sep 17 00:00:00 2001 From: Yohsuke Fukai Date: Tue, 28 Jul 2026 11:04:58 +0900 Subject: [PATCH] feat(nodes): add RegionPropsNodes.add_node_attrs to recompute region properties Allows (re-)computing scikit-image region properties for nodes that already exist in a graph, evaluating regionprops on each node's stored `Mask` instead of on a labels image. Useful to compute properties that were not requested when the nodes were created (e.g. intensity features of an additional channel) or to refresh properties after masks were modified, without rebuilding the graph. Missing output attribute keys are registered in the graph with dtypes inferred from the first computed values; existing keys are overwritten. The attribute holding the masks is configurable through the new `mask_key` argument. Property computation and attribute-key registration are factored out into module-level `_region_property_attrs` / `_add_missing_node_attr_keys` helpers shared with `add_nodes`. Co-Authored-By: Claude Opus 5 --- src/tracksdata/nodes/_regionprops.py | 249 ++++++++++++++++-- .../nodes/_test/test_regionprops.py | 144 ++++++++++ 2 files changed, 365 insertions(+), 28 deletions(-) diff --git a/src/tracksdata/nodes/_regionprops.py b/src/tracksdata/nodes/_regionprops.py index 11ad7261..94910a24 100644 --- a/src/tracksdata/nodes/_regionprops.py +++ b/src/tracksdata/nodes/_regionprops.py @@ -9,6 +9,7 @@ from skimage.measure._regionprops import RegionProperties, regionprops from typing_extensions import override +from tracksdata.attrs import NodeAttr from tracksdata.constants import DEFAULT_ATTR_KEYS from tracksdata.graph._base_graph import BaseGraph from tracksdata.nodes._base_nodes import BaseNodesOperator @@ -17,14 +18,99 @@ from tracksdata.utils._multiprocessing import multiprocessing_apply +def _validate_properties(properties: list[str | Callable[[RegionProperties], Any]]) -> None: + """ + Reject properties that are already added by default by `RegionPropsNodes`. + + Parameters + ---------- + properties : list[str | Callable[[RegionProperties], Any]] + The requested region properties. + """ + if "centroid" in properties: + raise ValueError( + "`centroid` is not supported as an extra property. It's already included by default as (z), y, x." + ) + if "bbox" in properties: + raise ValueError("`bbox` is not supported as an extra property. It's already included by default.") + + +def _add_missing_node_attr_keys(graph: BaseGraph, node_attrs: dict[str, Any]) -> None: + """ + Register node attribute keys in the graph, inferring dtypes from sample values. + + Keys already present in the graph are left untouched. + + Parameters + ---------- + graph : BaseGraph + The graph to register the attribute keys in. + node_attrs : dict[str, Any] + A sample of node attributes, mapping each key to a single value + used to infer the dtype. + """ + node_attr_keys = graph.node_attr_keys(return_ids=True) + for key, value in node_attrs.items(): + if key not in node_attr_keys: + if isinstance(value, np.ndarray): + default_value = np.zeros_like(value) + graph.add_node_attr_key( + key, pl.Array(numpy_char_code_to_dtype(value.dtype), value.shape), default_value + ) + elif np.isscalar(value): + dtype = numpy_char_code_to_dtype(value.dtype) if hasattr(value, "dtype") else type(value) + graph.add_node_attr_key(key, dtype) + elif type(value).__module__ != "builtins": + graph.add_node_attr_key(key, pl.Object) + else: + graph.add_node_attr_key(key, type(value)) + + +def _region_property_attrs( + obj: RegionProperties, + properties: list[str | Callable[[RegionProperties], Any]], +) -> dict[str, Any]: + """ + Compute the requested region properties for a single region. + + Parameters + ---------- + obj : RegionProperties + The scikit-image region to compute the properties for. + properties : list[str | Callable[[RegionProperties], Any]] + The properties to compute. Strings are looked up on ``obj``, callables + are called with ``obj`` and named after their ``__name__``. + + Returns + ------- + dict[str, Any] + Mapping from attribute key to value for this region. + """ + attrs: dict[str, Any] = {} + for prop in properties: + if callable(prop): + attrs[prop.__name__] = prop(obj) + else: + attrs[prop] = getattr(obj, prop) + + return attrs + + class RegionPropsNodes(BaseNodesOperator): """ - Operator that adds nodes to a graph using scikit-image's regionprops. + Operator that adds nodes and (re-)computes their region properties using scikit-image's regionprops. Extracts region properties from labeled images to create graph nodes using scikit-image's regionprops function to compute geometric and intensity-based features. Automatically adds centroid coordinates and mask information, with - additional properties computed based on the extra_properties parameter. + additional properties computed based on the ``extra_properties`` parameter. + + The same operator can also (re-)compute properties for nodes that already + exist in a graph, evaluating regionprops on each node's stored + [Mask][tracksdata.nodes.Mask] via [add_node_attrs][tracksdata.nodes.RegionPropsNodes.add_node_attrs]. + This is useful to compute properties that were not requested when the nodes + were created (e.g. intensity features of an additional channel) or to refresh + properties after masks were modified, without rebuilding the graph. Parameters ---------- @@ -37,6 +123,9 @@ class RegionPropsNodes(BaseNodesOperator): Physical spacing between pixels. If provided, affects distance-based measurements. Should be (row_spacing, col_spacing) for 2D or (depth_spacing, row_spacing, col_spacing) for 3D. + mask_key : str, optional + The key of the node attribute holding the [Mask][tracksdata.nodes.Mask] + objects, used by [add_node_attrs][tracksdata.nodes.RegionPropsNodes.add_node_attrs]. Attributes ---------- @@ -86,22 +175,26 @@ def custom_property(region): labels_series = np.random.randint(0, 10, (10, 100, 100)) node_op.add_nodes(graph, labels=labels_series) ``` + + Recompute properties of an additional channel on an existing graph: + + ```python + node_op = RegionPropsNodes(extra_properties=["intensity_mean", "intensity_max"]) + node_op.add_node_attrs(graph, intensity_image=second_channel) + ``` """ def __init__( self, extra_properties: list[str | Callable[[RegionProperties], Any]] | None = None, spacing: tuple[float, float] | None = None, + mask_key: str = DEFAULT_ATTR_KEYS.MASK, ): super().__init__() self._extra_properties = extra_properties or [] - if "centroid" in self._extra_properties: - raise ValueError( - "`centroid` is not supported as an extra property. It's already included by default as (z), y, x." - ) - if "bbox" in self._extra_properties: - raise ValueError("`bbox` is not supported as an extra property. It's already included by default.") + _validate_properties(self._extra_properties) self._spacing = spacing + self._mask_key = mask_key def _axis_names(self, labels: NDArray[np.integer]) -> list[str]: """ @@ -128,21 +221,7 @@ def _init_node_attrs(self, graph: BaseGraph, node_attrs: dict[str, Any]) -> None """ Initialize the node attributes for the graph. """ - node_attr_keys = graph.node_attr_keys(return_ids=True) - for key, value in node_attrs.items(): - if key not in node_attr_keys: - if isinstance(value, np.ndarray): - default_value = np.zeros_like(value) - graph.add_node_attr_key( - key, pl.Array(numpy_char_code_to_dtype(value.dtype), value.shape), default_value - ) - elif np.isscalar(value): - dtype = numpy_char_code_to_dtype(value.dtype) if hasattr(value, "dtype") else type(value) - graph.add_node_attr_key(key, dtype) - elif type(value).__module__ != "builtins": - graph.add_node_attr_key(key, pl.Object) - else: - graph.add_node_attr_key(key, type(value)) + _add_missing_node_attr_keys(graph, node_attrs) def attr_keys(self) -> list[str]: """ @@ -300,11 +379,7 @@ def _nodes_per_time( ): attrs = dict(zip(axis_names, obj.centroid, strict=False)) - for prop in self._extra_properties: - if callable(prop): - attrs[prop.__name__] = prop(obj) - else: - attrs[prop] = getattr(obj, prop) + attrs.update(_region_property_attrs(obj, self._extra_properties)) attrs[DEFAULT_ATTR_KEYS.MASK] = Mask(obj.image, obj.bbox) attrs[DEFAULT_ATTR_KEYS.BBOX] = np.asarray(obj.bbox, dtype=int) @@ -317,3 +392,121 @@ def _nodes_per_time( LOG.warning("No valid nodes found for time point %d", t) return nodes_data + + def add_node_attrs( + self, + graph: BaseGraph, + *, + t: int | None = None, + intensity_image: NDArray | None = None, + ) -> None: + """ + (Re-)compute region properties from the node masks and store them as node attributes. + + For each node, scikit-image's regionprops is evaluated on the node's + [Mask][tracksdata.nodes.Mask] attribute (``mask_key``), optionally combined + with a given intensity image cropped to the mask bounding box. Missing + output attribute keys are registered in the graph with dtypes inferred from + the first computed values; existing keys are overwritten. + + Parameters + ---------- + graph : BaseGraph + The graph to add attributes to. + t : int | None, optional + The time point to compute attributes for. + If None, attributes are computed for all time points of the graph. + intensity_image : NDArray | None, optional + Intensity image used for computing intensity-based properties, + indexed by time point such that `intensity_image[t]` is the frame + matching the masks at time point `t`. + + Examples + -------- + Compute intensity features from an additional channel on an existing graph: + + ```python + node_op = RegionPropsNodes(extra_properties=["intensity_mean", "intensity_max"]) + node_op.add_node_attrs(graph, intensity_image=second_channel) + ``` + """ + if not self._extra_properties: + raise ValueError("`extra_properties` must contain at least one region property to compute node attributes.") + + if self._mask_key not in graph.node_attr_keys(): + raise ValueError(f"Mask key '{self._mask_key}' not found in graph. Expected '{graph.node_attr_keys()}'") + + if t is None: + time_points = graph.time_points() + else: + time_points = [t] + + initialized = False + for node_ids, node_attrs in multiprocessing_apply( + func=partial(self._node_attrs_per_time, graph=graph, intensity_image=intensity_image), + sequence=time_points, + desc="Computing region properties attributes", + ): + if len(node_ids) == 0: + continue + if not initialized: + sample_attrs = {key: values[0] for key, values in node_attrs.items()} + _add_missing_node_attr_keys(graph, sample_attrs) + initialized = True + graph.update_node_attrs(node_ids=node_ids, attrs=node_attrs) + + def _node_attrs_per_time( + self, + t: int, + *, + graph: BaseGraph, + intensity_image: NDArray | None = None, + ) -> tuple[list[int], dict[str, list[Any]]]: + """ + Compute region properties for the nodes of a single time point. + + Parameters + ---------- + t : int + The time point to compute attributes for. + graph : BaseGraph + The graph to add attributes to. + intensity_image : NDArray | None, optional + Intensity image indexed by time point, see `add_node_attrs`. + + Returns + ------- + tuple[list[int], dict[str, list[Any]]] + The node ids and the attributes to add to the graph. + """ + graph_filter = graph.filter(NodeAttr(DEFAULT_ATTR_KEYS.T) == t) + node_ids = graph_filter.node_ids() + + if len(node_ids) == 0: + LOG.warning("No nodes found for time point %d", t) + return [], {} + + masks = graph_filter.node_attrs(attr_keys=[self._mask_key])[self._mask_key].to_list() + + frame = np.asarray(intensity_image[t]) if intensity_image is not None else None + + results: dict[str, list[Any]] = {} + for mask in masks: + if not isinstance(mask, Mask): + raise TypeError( + f"Expected `Mask` object in '{self._mask_key}' attribute, got '{type(mask)}'. " + "Use `mask_key` to select the attribute holding the masks." + ) + + regionprops_kwargs: dict[str, Any] = {"spacing": self._spacing} + if frame is not None: + regionprops_kwargs["intensity_image"] = mask.crop(frame) + + obj = mask.regionprops(**regionprops_kwargs) + + for key, value in _region_property_attrs(obj, self._extra_properties).items(): + results.setdefault(key, []).append(value) + + obj._cache.clear() # clearing to reduce memory footprint + + return node_ids, results diff --git a/src/tracksdata/nodes/_test/test_regionprops.py b/src/tracksdata/nodes/_test/test_regionprops.py index 567c62e0..318615db 100644 --- a/src/tracksdata/nodes/_test/test_regionprops.py +++ b/src/tracksdata/nodes/_test/test_regionprops.py @@ -1,4 +1,5 @@ import numpy as np +import polars as pl import pytest from skimage.measure._regionprops import RegionProperties @@ -334,3 +335,146 @@ def test_regionprops_multiprocessing_isolation() -> None: """Test that multiprocessing options don't affect subsequent tests.""" # Verify default n_workers is 1 assert get_options().n_workers == 1 + + +TIMELAPSE_LABELS = np.array( + [ + [[1, 1, 0], [1, 0, 2], [0, 2, 2]], # t=0 + [[0, 3, 3], [0, 3, 0], [4, 0, 0]], # t=1 + ], + dtype=np.int32, +) + +TIMELAPSE_INTENSITY = np.array( + [ + [[10, 20, 0], [30, 0, 40], [0, 50, 60]], # t=0 + [[0, 70, 80], [0, 90, 0], [100, 0, 0]], # t=1 + ], + dtype=np.float32, +) + + +def test_regionprops_attrs_init_validation() -> None: + """Test RegionPropsNodes property validation for add_node_attrs.""" + operator = RegionPropsNodes(extra_properties=["area", "intensity_mean"], spacing=(1.0, 2.0)) + assert operator.attr_keys() == ["area", "intensity_mean"] + assert operator._spacing == (1.0, 2.0) + + # empty properties are only rejected when (re-)computing attributes on existing nodes + graph = RustWorkXGraph() + RegionPropsNodes().add_nodes(graph, labels=TIMELAPSE_LABELS) + with pytest.raises(ValueError, match="at least one region property"): + RegionPropsNodes().add_node_attrs(graph) + + with pytest.raises(ValueError, match="`centroid` is not supported"): + RegionPropsNodes(extra_properties=["centroid"]) + + with pytest.raises(ValueError, match="`bbox` is not supported"): + RegionPropsNodes(extra_properties=["bbox"]) + + +@pytest.mark.parametrize("n_workers", [1, 2]) +def test_regionprops_attrs_matches_nodes_operator(n_workers: int) -> None: + """Test that recomputed properties match those computed at node creation.""" + properties = ["area", "intensity_mean", "intensity_max"] + + # ground truth: properties computed directly from the labels + expected_graph = RustWorkXGraph() + RegionPropsNodes(extra_properties=properties).add_nodes( + expected_graph, labels=TIMELAPSE_LABELS, intensity_image=TIMELAPSE_INTENSITY + ) + expected_df = expected_graph.node_attrs(attr_keys=properties) + + # recompute: nodes created without properties, then add_node_attrs + graph = RustWorkXGraph() + RegionPropsNodes().add_nodes(graph, labels=TIMELAPSE_LABELS) + + operator = RegionPropsNodes(extra_properties=properties) + with options_context(n_workers=n_workers): + operator.add_node_attrs(graph, intensity_image=TIMELAPSE_INTENSITY) + + result_df = graph.node_attrs(attr_keys=properties) + + for prop in properties: + np.testing.assert_allclose(result_df[prop].to_numpy(), expected_df[prop].to_numpy()) + + +def test_regionprops_attrs_callable_property() -> None: + """Test recomputing properties with a custom callable.""" + + def double_area(region: RegionProperties) -> float: + return region.area * 2 + + graph = RustWorkXGraph() + RegionPropsNodes(extra_properties=["area"]).add_nodes(graph, labels=TIMELAPSE_LABELS) + + RegionPropsNodes(extra_properties=[double_area]).add_node_attrs(graph) + + nodes_df = graph.node_attrs(attr_keys=["area", "double_area"]) + np.testing.assert_array_equal( + nodes_df["double_area"].to_numpy(), + nodes_df["area"].to_numpy() * 2, + ) + + +def test_regionprops_attrs_single_time_point() -> None: + """Test recomputing properties for a single time point only.""" + graph = RustWorkXGraph() + RegionPropsNodes().add_nodes(graph, labels=TIMELAPSE_LABELS, intensity_image=TIMELAPSE_INTENSITY) + + RegionPropsNodes(extra_properties=["intensity_mean"]).add_node_attrs( + graph, t=1, intensity_image=TIMELAPSE_INTENSITY + ) + + nodes_df = graph.node_attrs(attr_keys=[DEFAULT_ATTR_KEYS.T, "intensity_mean"]) + at_t1 = nodes_df.filter(nodes_df[DEFAULT_ATTR_KEYS.T] == 1) + + # region 3: pixels (70, 80, 90) -> mean = 80; region 4: pixel (100,) -> mean = 100 + np.testing.assert_allclose(sorted(at_t1["intensity_mean"]), [80.0, 100.0]) + + +def test_regionprops_attrs_overwrites_existing_values() -> None: + """Test that existing attribute values are overwritten by the recomputation.""" + graph = RustWorkXGraph() + RegionPropsNodes(extra_properties=["area"]).add_nodes(graph, labels=TIMELAPSE_LABELS) + + expected_areas = graph.node_attrs(attr_keys=["area"])["area"].to_numpy().copy() + + # corrupt the stored values + graph.update_node_attrs(node_ids=graph.node_ids(), attrs={"area": [-1] * graph.num_nodes()}) + + RegionPropsNodes(extra_properties=["area"]).add_node_attrs(graph) + + np.testing.assert_array_equal( + graph.node_attrs(attr_keys=["area"])["area"].to_numpy(), + expected_areas, + ) + + +def test_regionprops_attrs_spacing() -> None: + """Test that spacing affects the recomputed measurements.""" + graph = RustWorkXGraph() + RegionPropsNodes(extra_properties=["area"]).add_nodes(graph, labels=TIMELAPSE_LABELS) + + pixel_areas = graph.node_attrs(attr_keys=["area"])["area"].to_numpy().copy() + + RegionPropsNodes(extra_properties=["area"], spacing=(2.0, 3.0)).add_node_attrs(graph) + + np.testing.assert_allclose( + graph.node_attrs(attr_keys=["area"])["area"].to_numpy(), + pixel_areas * 6.0, + ) + + +def test_regionprops_attrs_missing_mask_key() -> None: + """Test error handling when the mask key is missing or invalid.""" + graph = RustWorkXGraph() + + with pytest.raises(ValueError, match="Mask key 'mask' not found"): + RegionPropsNodes(extra_properties=["area"]).add_node_attrs(graph) + + graph.add_node_attr_key(DEFAULT_ATTR_KEYS.MASK, pl.Object) + graph.add_node({"t": 0, "mask": "not a mask"}) + + with pytest.raises(TypeError, match="Expected `Mask` object"): + RegionPropsNodes(extra_properties=["area"]).add_node_attrs(graph)