# Copyright (c) 2024-2026 Li Chao, Dong Mengshi and HABIT Contributors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Connected-region node extraction for habitat graph features."""
from __future__ import annotations
from dataclasses import replace
from typing import Dict, List, Optional, Tuple
import numpy as np
from scipy import ndimage as ndi
from habit.kernels.habitat_graph.models import (
HabitatGraphNode,
HabitatNodeExtractionResult,
NodeMethod,
)
from habit.kernels.habitat_graph.union_find import label_painted_components
__all__ = ["extract_habitat_nodes"]
def _nonzero_bbox_slices(
label_array: np.ndarray,
pad: int = 1,
) -> Optional[Tuple[Tuple[slice, ...], Tuple[int, ...]]]:
"""
Bounding-box slices of non-background labels, plus one voxel of pad.
Pad is clipped to the input array so we never invent voxels outside
the original lattice. The extra zero layer keeps erosion / adjacency
at the tumour rim identical to running on the full CT (tumour-adjacent
background is still 0).
Args:
label_array: Integer habitat map; ``0`` is background.
pad: Voxels of background to keep around the nonempty bbox.
Returns:
``(slices, offset)`` where ``offset`` is the inclusive origin of
the crop in the original index space, or ``None`` when empty.
"""
hits = np.argwhere(np.asarray(label_array) != 0)
if hits.size == 0:
return None
lo = np.maximum(hits.min(axis=0) - int(pad), 0)
hi = np.minimum(hits.max(axis=0) + 1 + int(pad), label_array.shape)
slices = tuple(slice(int(start), int(stop)) for start, stop in zip(lo, hi))
offset = tuple(int(v) for v in lo)
return slices, offset
def _shift_node(node: HabitatGraphNode, offset: np.ndarray) -> HabitatGraphNode:
"""Translate one node's centroid and bbox into the original index space."""
ndim = int(offset.size)
bbox = list(node.bbox)
shifted = [int(bbox[axis] + offset[axis]) for axis in range(ndim)]
shifted.extend(
int(bbox[ndim + axis] + offset[axis]) for axis in range(ndim)
)
return HabitatGraphNode(
node_id=node.node_id,
habitat_label=node.habitat_label,
component_id=node.component_id,
centroid=np.asarray(node.centroid, dtype=float) + offset,
voxel_count=int(node.voxel_count),
bbox=tuple(shifted),
)
def _apply_crop_offset(
result: HabitatNodeExtractionResult,
offset: Tuple[int, ...],
) -> HabitatNodeExtractionResult:
"""
Keep cropped maps, report nodes / lattice origin in original indices.
``component_maps`` and ``label_array`` stay on the cropped grid so
later adjacency / min-distance walks skip empty CT slices.
"""
shift = np.asarray(offset, dtype=float)
nodes = {
int(habitat_id): [_shift_node(node, shift) for node in group]
for habitat_id, group in result.nodes_by_habitat.items()
}
grid_origin = result.grid_origin
if grid_origin is not None:
grid_origin = tuple(
int(axis + extra) for axis, extra in zip(grid_origin, offset)
)
return replace(
result,
nodes_by_habitat=nodes,
grid_origin=grid_origin,
crop_offset=offset,
)
def _connectivity_structure(ndim: int, connectivity: str) -> np.ndarray:
"""
Build an n-dimensional binary structure for connected-component labeling.
Args:
ndim: Number of array dimensions.
connectivity: ``"face"`` for 4-neighborhood in 2D / 6-neighborhood in
3D, or ``"full"`` for diagonal-inclusive connectivity.
Returns:
np.ndarray: Binary structure accepted by ``scipy.ndimage.label``.
"""
if connectivity == "face":
return ndi.generate_binary_structure(rank=ndim, connectivity=1)
if connectivity == "full":
return ndi.generate_binary_structure(rank=ndim, connectivity=ndim)
raise ValueError("connectivity must be 'face' or 'full'.")
def _component_bbox(coords: np.ndarray) -> Tuple[int, ...]:
"""
Return a half-open bounding box for component coordinates.
Args:
coords: Array of component voxel coordinates with shape ``(n, ndim)``.
Returns:
Tuple[int, ...]: ``(min_dim0, ..., min_dimN, max_dim0, ..., max_dimN)``.
"""
mins = coords.min(axis=0)
maxs = coords.max(axis=0) + 1
return tuple(int(v) for v in np.concatenate([mins, maxs]))
def _subdivide_component(
coords: np.ndarray,
block_size: int,
ndim: int,
min_coverage: float,
origin: Optional[np.ndarray] = None,
) -> List[np.ndarray]:
"""
Split a connected component into fixed-size grid blocks.
Each voxel is assigned to an n-dimensional block of edge length
``block_size``. A block is kept only when the fraction of its volume covered
by the component exceeds ``min_coverage``, which mirrors the source PathPrism
behavior of dropping sparsely covered boundary blocks.
Args:
coords: Component voxel coordinates with shape ``(n, ndim)``.
block_size: Edge length of each grid block in voxels.
ndim: Number of array dimensions.
min_coverage: Minimum covered fraction of a block volume to keep it.
origin: Lattice origin in voxel indices. ``None`` uses this
component's own bounding-box minimum (legacy per-component
grid). Pass the tumour-VOI minimum for a global lattice.
Returns:
List[np.ndarray]: One coordinate array per kept block. Empty when no
block reaches the coverage threshold.
"""
mins = coords.min(axis=0) if origin is None else np.asarray(origin)
# Integer block index per voxel along every dimension.
block_indices = (coords - mins) // block_size
unique_blocks, inverse = np.unique(block_indices, axis=0, return_inverse=True)
block_volume = float(block_size**ndim)
kept_blocks: List[np.ndarray] = []
for block_id in range(unique_blocks.shape[0]):
block_coords = coords[inverse == block_id]
coverage = block_coords.shape[0] / block_volume
if coverage > min_coverage:
kept_blocks.append(block_coords)
return kept_blocks
def _voi_grid_origin(label_array: np.ndarray) -> Optional[np.ndarray]:
"""
Return the inclusive voxel-index origin of the non-background VOI.
Args:
label_array: Integer habitat label map (background encoded as 0).
Returns:
np.ndarray | None: One integer per axis, or ``None`` when empty.
"""
coords = np.argwhere(label_array > 0)
if coords.size == 0:
return None
return coords.min(axis=0).astype(int, copy=False)
def _eroded_label_array(
label_array: np.ndarray,
labels: List[int],
structure: np.ndarray,
erosion_radius: int,
) -> np.ndarray:
"""
Optionally erode each habitat and rebuild the integer label map.
Args:
label_array: Integer habitat label map.
labels: Positive habitat ids to process.
structure: Binary structure for ``scipy.ndimage.binary_erosion``.
erosion_radius: Erosion iterations; ``0`` returns ``label_array``.
Returns:
np.ndarray: Label map after per-habitat erosion (copy if eroded).
"""
if erosion_radius <= 0:
return label_array
eroded = np.zeros_like(label_array)
for habitat_label in labels:
mask = ndi.binary_erosion(
label_array == habitat_label,
structure=structure,
iterations=erosion_radius,
border_value=0,
)
eroded[mask] = habitat_label
return eroded
def _node_from_coords(
habitat_label: int,
component_id: int,
coords: np.ndarray,
) -> HabitatGraphNode:
"""Build one graph node from a voxel-coordinate set."""
return HabitatGraphNode(
node_id=f"h{habitat_label}_c{component_id}",
habitat_label=habitat_label,
component_id=component_id,
centroid=coords.mean(axis=0).astype(float),
voxel_count=int(coords.shape[0]),
bbox=_component_bbox(coords),
)
def _extract_uniform_grid_nodes(
label_array: np.ndarray,
labels: List[int],
structure: np.ndarray,
min_region_voxels: int,
erosion_radius: int,
block_size: int,
block_min_coverage: float,
connectivity: str = "full",
grid_origin: Optional[np.ndarray] = None,
) -> HabitatNodeExtractionResult:
"""
Tessellate the tumour VOI and emit one node per cell subregion.
Every non-background voxel is assigned to an axis-aligned cube of edge
``block_size`` whose origin is the VOI bounding-box minimum. A cube is
kept when its occupied fraction is strictly greater than
``block_min_coverage`` (cell-level filter for nearly-empty cubes).
Inside each kept cube, every connected component of every habitat
becomes its own node at that subregion's voxel-index centroid, so a
mixed cube can contribute several nodes. Subregions smaller than
``min_region_voxels`` are dropped (fragment filter). Habitats that
occupy no kept subregion still emit one residual node so a present
label is never silently dropped.
Args:
label_array: Integer habitat label map (background encoded as 0).
labels: Positive habitat ids present before erosion.
structure: Neighbourhood for optional erosion (same as
``connectivity``).
connectivity: In-cell CCL rule. ``'full'`` is 8-connected in
2-D / 26-connected in 3-D; ``'face'`` is 4 / 6. Must match
the neighbourhood encoded by ``structure``.
min_region_voxels: Drop in-cell subregions (and residual nodes)
smaller than this voxel count.
erosion_radius: Optional per-habitat erosion iterations.
block_size: Cube edge length in voxels.
block_min_coverage: Minimum occupied fraction of a cube to keep
the cell (strictly greater than this value).
grid_origin: Optional lattice origin in voxel indices. ``None``
uses the non-background bounding-box minimum of ``working``.
Pass the tumour-VOI origin when extracting the background
shell so habitat cubes stay on the same lattice.
Returns:
HabitatNodeExtractionResult: Nodes, component maps, and lattice.
"""
working = _eroded_label_array(label_array, labels, structure, erosion_radius)
origin = (
np.asarray(grid_origin, dtype=int)
if grid_origin is not None
else _voi_grid_origin(working)
)
nodes_by_habitat: Dict[int, List[HabitatGraphNode]] = {}
component_maps: Dict[int, np.ndarray] = {
int(label): np.zeros(working.shape, dtype=np.int32) for label in labels
}
if origin is None:
return HabitatNodeExtractionResult(
label_array=label_array,
nodes_by_habitat=nodes_by_habitat,
component_maps=component_maps,
grid_origin=None,
grid_block_size=int(block_size),
)
coords = np.argwhere(working != 0)
voxel_labels = working[tuple(coords.T)]
block_indices = (coords - origin) // block_size
unique_blocks, inverse = np.unique(block_indices, axis=0, return_inverse=True)
block_volume = float(block_size ** working.ndim)
cube_shape = tuple(int(block_size) for _ in range(working.ndim))
next_component_id = 1
kept_by_habitat: Dict[int, List[HabitatGraphNode]] = {
int(label): [] for label in labels
}
# Sort voxels by cube id so each cell is a contiguous slice (O(N log N)
# once) instead of a full-volume mask per cube (O(N * n_cubes)).
order = np.argsort(inverse, kind="stable")
inverse_sorted = inverse[order]
coords_sorted = coords[order]
labels_sorted = voxel_labels[order]
breaks = np.flatnonzero(inverse_sorted[1:] != inverse_sorted[:-1]) + 1
starts = np.concatenate((np.asarray([0], dtype=np.int64), breaks))
stops = np.concatenate((breaks, np.asarray([inverse_sorted.size], dtype=np.int64)))
for start, stop in zip(starts.tolist(), stops.tolist()):
block_id = int(inverse_sorted[start])
block_coords = coords_sorted[start:stop]
block_lab = labels_sorted[start:stop]
coverage = block_coords.shape[0] / block_volume
if coverage <= block_min_coverage:
continue
# Local coordinates of painted voxels (no dense 8^3 cube).
local = block_coords - origin - unique_blocks[block_id] * block_size
local_i = np.ascontiguousarray(local, dtype=np.int32)
# Split the cube by habitat id so mixed cells become separate nodes.
for habitat_label in (
int(v) for v in np.unique(block_lab) if int(v) != 0
):
habitat_mask = block_lab == habitat_label
hab_local = local_i[habitat_mask]
hab_global = block_coords[habitat_mask]
cc_labels, n_cc = label_painted_components(
hab_local, cube_shape, connectivity
)
for cc_id in range(1, int(n_cc) + 1):
member = cc_labels == cc_id
if int(member.sum()) < min_region_voxels:
continue
global_coords = hab_global[member]
component_id = next_component_id
next_component_id += 1
component_maps[habitat_label][tuple(global_coords.T)] = component_id
kept_by_habitat[habitat_label].append(
_node_from_coords(habitat_label, component_id, global_coords)
)
present_after = {int(v) for v in np.unique(working) if int(v) != 0}
for habitat_label in labels:
habitat_nodes = kept_by_habitat.get(habitat_label, [])
if habitat_nodes:
nodes_by_habitat[habitat_label] = habitat_nodes
continue
if habitat_label not in present_after:
continue
leftover = np.argwhere(working == habitat_label)
if leftover.shape[0] < min_region_voxels:
continue
# Residual node: the habitat is present but no subregion was kept.
component_id = next_component_id
next_component_id += 1
component_maps[habitat_label][tuple(leftover.T)] = component_id
nodes_by_habitat[habitat_label] = [
_node_from_coords(habitat_label, component_id, leftover)
]
return HabitatNodeExtractionResult(
label_array=label_array,
nodes_by_habitat=nodes_by_habitat,
component_maps=component_maps,
grid_origin=tuple(int(v) for v in origin),
grid_block_size=int(block_size),
)
def _extract_component_nodes(
label_array: np.ndarray,
labels: List[int],
structure: np.ndarray,
min_region_voxels: int,
erosion_radius: int,
subdivide_region_voxels: int,
block_size: int,
block_min_coverage: float,
) -> HabitatNodeExtractionResult:
"""Connected-component nodes (optional size split). Habitat labels only."""
nodes_by_habitat: Dict[int, List[HabitatGraphNode]] = {}
component_maps: Dict[int, np.ndarray] = {}
for habitat_label in labels:
mask = label_array == habitat_label
if erosion_radius > 0:
mask = ndi.binary_erosion(
mask,
structure=structure,
iterations=erosion_radius,
border_value=0,
)
labeled_components, component_count = ndi.label(mask, structure=structure)
kept_component_map = np.zeros_like(labeled_components, dtype=np.int32)
habitat_nodes: List[HabitatGraphNode] = []
# Block nodes need unique component ids that never collide with the
# original connected-component ids painted into the component map.
next_block_component_id = int(component_count) + 1
for component_id in range(1, int(component_count) + 1):
coords = np.argwhere(labeled_components == component_id)
voxel_count = int(coords.shape[0])
if voxel_count < min_region_voxels:
continue
should_subdivide = (
subdivide_region_voxels > 0
and voxel_count > subdivide_region_voxels
)
block_groups: List[np.ndarray] = []
if should_subdivide:
block_groups = _subdivide_component(
coords=coords,
block_size=block_size,
ndim=label_array.ndim,
min_coverage=block_min_coverage,
)
if block_groups:
for block_coords in block_groups:
block_component_id = next_block_component_id
next_block_component_id += 1
kept_component_map[tuple(block_coords.T)] = block_component_id
habitat_nodes.append(
HabitatGraphNode(
node_id=f"h{habitat_label}_c{block_component_id}",
habitat_label=habitat_label,
component_id=block_component_id,
centroid=block_coords.mean(axis=0).astype(float),
voxel_count=int(block_coords.shape[0]),
bbox=_component_bbox(block_coords),
)
)
else:
kept_component_map[tuple(coords.T)] = component_id
habitat_nodes.append(
HabitatGraphNode(
node_id=f"h{habitat_label}_c{component_id}",
habitat_label=habitat_label,
component_id=component_id,
centroid=coords.mean(axis=0).astype(float),
voxel_count=voxel_count,
bbox=_component_bbox(coords),
)
)
nodes_by_habitat[habitat_label] = habitat_nodes
component_maps[habitat_label] = kept_component_map
return HabitatNodeExtractionResult(
label_array=label_array,
nodes_by_habitat=nodes_by_habitat,
component_maps=component_maps,
)