"""
vessel_classifier
=================
Iterative artery/vein classification on a reduced vascular graph.
The algorithm has four phases:
1. **Score** — each edge gets a continuous artery score and vein score
in [0, 1] based on radius, marker intensity, and binary mask.
2. **Seed** — high-confidence edges are assigned initial labels
(artery, vein, or capillary).
3. **Propagate** — labels spread outward from seeds along the graph.
Artery and vein labels propagate simultaneously and compete for
ambiguous edges. Propagation uses a decaying score threshold.
4. **Cleanup** — remove small connected components (fragments).
All thresholds are collected in :class:`ClassificationConfig`.
See Also
--------
:class:`~ClearMap.pipeline_orchestrators.tube_map.VesselGraphProcessor`
Orchestrator that builds the graph and calls this classifier.
[Kirst2020]_.
"""
from __future__ import annotations
from dataclasses import dataclass
from enum import IntEnum
from typing import Optional
import numpy as np
from ClearMap.Analysis.graphs import graph_processing
from ClearMap.Analysis.graphs.graph_gt import Graph
# ── vessel type enum ─────────────────────────────────────────────────────────
[docs]
class VesselType(IntEnum):
"""
Edge classification labels for vessel type assignment.
IntEnum so labels can be used directly as numpy array values
and in boolean comparisons without casting.
"""
UNLABELLED = 0
ARTERY = 1
VEIN = 2
CAPILLARY = 3
[docs]
@classmethod
def competing_labels(cls) -> tuple['VesselType', ...]:
"""The two vessel types that compete during propagation."""
return cls.ARTERY, cls.VEIN
[docs]
def competitor(self) -> 'VesselType':
"""Return the competing vessel type."""
if self == VesselType.ARTERY:
return VesselType.VEIN
if self == VesselType.VEIN:
return VesselType.ARTERY
raise ValueError(f'{self} has no competitor')
# ── configuration ────────────────────────────────────────────────────────────
[docs]
@dataclass(frozen=True)
class ClassificationConfig:
"""
All thresholds for vessel classification.
Radius thresholds are unit-agnostic: they match whatever units
``VesselSignals.radii`` uses (µm for current graphs, voxels for legacy).
The ``from_config`` factory resolves the correct source.
Attributes
----------
vein_intensity_range : tuple[float, float]
(low, high) intensity range on the artery channel that
characterises veins. Edges *within* this range AND with
large radius are classified as restrictive veins.
restrictive_vein_radius : float
Minimum radius for an edge to qualify as a
*restrictive* (high-confidence) vein.
arteries_min_noise_edges : int
Minimum number of connected artery edges to survive
initial noise removal (step 1).
artery_intensity_min : float
Minimum artery marker intensity for tracing to continue.
artery_trace_radius : float
Minimum edge radius to continue artery tracing.
vein_trace_radius : float
Minimum edge radius to continue vein tracing.
vein_intensity_min : float
Minimum vein marker intensity to continue vein tracing.
Only used when a dedicated vein channel exists.
distance_to_surface_min : float
Minimum distance to brain surface (atlas voxels) for
tracing to continue.
max_artery_trace_iter : int
Maximum hysteresis iterations for artery tracing.
max_vein_trace_iter : int
Maximum hysteresis iterations for vein tracing.
min_artery_component_edges : int
Minimum edges in a connected artery component to survive
final cleanup (step 7).
min_vein_component_edges : int
Minimum edges in a connected vein component to survive
final cleanup (step 7).
"""
# ── pre-filtering ────────────────────────────────────────────────────
vein_intensity_range: tuple[float, float]
restrictive_vein_radius: float
permissive_vein_radius: float
final_vein_radius: float
arteries_min_noise_edges: int
# ── signal thresholds (used in scoring AND as propagation floors) ────
artery_intensity_min: float
artery_trace_radius: float
vein_trace_radius: float
vein_intensity_min: float
distance_to_surface_min: float
max_artery_trace_iter: int
max_vein_trace_iter: int
# ── scoring ──────────────────────────────────────────────────────────
artery_seed_threshold: float = 0.7
vein_seed_threshold: float = 0.6
# ── cleanup ──────────────────────────────────────────────────────────
min_artery_component_edges: int = 30
min_vein_component_edges: int = 30
@property
def max_propagation_iter(self) -> int:
"""Total propagation budget derived from existing tracing iterations."""
return self.max_artery_trace_iter + self.max_vein_trace_iter
@property
def initial_threshold(self) -> float:
"""Starting score threshold for propagation."""
return min(self.artery_seed_threshold, self.vein_seed_threshold) * 0.7
@property
def propagation_floor(self) -> float:
"""Score floor — propagation never accepts below this."""
return self.initial_threshold * 0.3
@property
def threshold_decay(self) -> float:
"""Per-iteration decay. Reaches floor at last iteration."""
return ((self.initial_threshold - self.propagation_floor)
/ max(self.max_propagation_iter, 1))
@property
def seed_thresholds(self) -> dict[VesselType, float]:
"""Per-type seed thresholds."""
return {
VesselType.ARTERY: self.artery_seed_threshold,
VesselType.VEIN: self.vein_seed_threshold,
}
@property
def min_component_edges(self) -> dict[VesselType, int]:
"""Per-type minimum component sizes for cleanup."""
return {
VesselType.ARTERY: self.min_artery_component_edges,
VesselType.VEIN: self.min_vein_component_edges,
}
[docs]
@classmethod
def from_config(cls, cfg: dict, legacy_thresholds: dict = None,
use_legacy: bool = False) -> 'ClassificationConfig':
"""
Build from the ``vessel_type_postprocessing`` config subtree.
Parameters
----------
cfg : dict
The ``vessel_type_postprocessing`` config subtree.
legacy_thresholds : dict or None
``VesselGraphProcessor._LEGACY_THRESHOLDS`` dict.
Required when ``use_legacy=True``.
use_legacy : bool
If True, read voxel-space thresholds from ``legacy_thresholds``
instead of µm values from config.
"""
pre = cfg['pre_filtering']
tr = cfg['tracing']
cap = cfg['capillaries_removal']
sc = cfg.get('scoring', {})
# ── shared kwargs (identical regardless of legacy mode) ──────────
shared = dict(
vein_intensity_range=tuple(pre['vein_intensity_range_on_arteries_ch']),
arteries_min_noise_edges=pre['arteries_min_noise_edges'],
artery_intensity_min=tr['artery_intensity_min'],
vein_intensity_min=tr.get('vein_intensity_min', 200.0),
distance_to_surface_min=tr['distance_to_surface_min'],
max_artery_trace_iter=tr['max_arteries_iterations'],
max_vein_trace_iter=tr['max_veins_iterations'],
artery_seed_threshold=sc.get('artery_seed_threshold', 0.7),
vein_seed_threshold=sc.get('vein_seed_threshold', 0.6),
min_artery_component_edges=cap['min_artery_component_edges'],
min_vein_component_edges=cap['min_vein_component_edges'],
)
# ── unit-dependent kwargs ────────────────────────────────────────
if use_legacy:
if legacy_thresholds is None:
raise ValueError('legacy_thresholds required when use_legacy=True')
t = legacy_thresholds
return cls(
**shared,
restrictive_vein_radius=t['restrictive_vein_radius_vx'],
permissive_vein_radius=t['permissive_vein_radius_vx'],
final_vein_radius=t['final_vein_radius_vx'],
artery_trace_radius=t['artery_trace_radius_vx'],
vein_trace_radius=t['vein_trace_radius_vx'],
)
return cls(
**shared,
restrictive_vein_radius=pre['restrictive_vein_radius_um'],
permissive_vein_radius=pre['permissive_vein_radius_um'],
final_vein_radius=pre['final_vein_radius_um'],
artery_trace_radius=tr['artery_trace_radius_um'],
vein_trace_radius=tr['vein_trace_radius_um'],
)
# ── signal bundle ────────────────────────────────────────────────────────────
[docs]
@dataclass
class VesselSignals:
"""
All per-edge arrays needed by the classifier.
``radii`` and all ``ClassificationConfig`` radius thresholds are
guaranteed to be in the same units. The caller resolves units
before constructing this object.
Fields are ``None`` when the corresponding channel does not exist.
"""
radii: np.ndarray # (n_edges,) — radii in whichever unit was supplied (should match thresholds)
artery_binary: Optional[np.ndarray] = None # (n_edges,) — bool, from binary mask
artery_intensity: Optional[np.ndarray] = None # (n_edges,) — float, raw marker signal
vein_binary: Optional[np.ndarray] = None # (n_edges,) — bool, from binary mask
vein_intensity: Optional[np.ndarray] = None # (n_edges,) — float, raw marker signal
distance_to_surface: Optional[np.ndarray] = None # (n_edges,) — atlas voxels
# ── classification metrics ───────────────────────────────────────────────────
[docs]
@dataclass
class ClassificationMetrics:
"""
Diagnostic metrics produced after classification.
Intended for logging, QA dashboards, and sanity checks.
"""
n_edges: int = 0
counts: dict = None # {VesselType: int}
fractions: dict = None # {VesselType: float}
seed_counts: dict = None # {VesselType: int}
propagation_iterations: int = 0
propagated_counts: dict = None # {VesselType: int}
n_seed_conflicts: int = 0
n_direct_contacts: int = 0
fragments_removed: dict = None # {VesselType: int}
score_stats: dict = None # {VesselType: {'mean': float, 'min': float}}
def __post_init__(self):
if self.counts is None:
self.counts = {t: 0 for t in VesselType}
if self.fractions is None:
self.fractions = {t: 0.0 for t in VesselType}
if self.seed_counts is None:
self.seed_counts = {t: 0 for t in VesselType}
if self.propagated_counts is None:
self.propagated_counts = {t: 0 for t in VesselType.competing_labels()}
if self.fragments_removed is None:
self.fragments_removed = {t: 0 for t in VesselType.competing_labels()}
if self.score_stats is None:
self.score_stats = {t: {'mean': 0.0, 'min': 0.0}
for t in VesselType.competing_labels()}
[docs]
def record_seed_counts(self, labels: np.ndarray) -> None:
"""Snapshot label distribution after seeding phase."""
for t in VesselType:
self.seed_counts[t] = int(np.sum(labels == t))
[docs]
def record_final_counts(self, labels: np.ndarray,
scores: dict[VesselType, np.ndarray]) -> None:
"""Compute final distribution, propagated counts, and score statistics."""
n = max(self.n_edges, 1)
for t in VesselType:
mask = labels == t
self.counts[t] = int(mask.sum())
self.fractions[t] = self.counts[t] / n
for t in VesselType.competing_labels():
self.propagated_counts[t] = self.counts[t] - self.seed_counts[t]
mask = labels == t
if self.counts[t] > 0:
self.score_stats[t] = {
'mean': float(scores[t][mask].mean()),
'min': float(scores[t][mask].min()),
}
@property
def sanity_warnings(self) -> list[str]:
"""Return a list of human-readable warnings if metrics look suspicious."""
warnings = []
n = max(self.n_edges, 1)
if self.n_direct_contacts > 0:
precent = 100 * self.n_direct_contacts / n
warnings.append(f'{self.n_direct_contacts} direct artery-vein contacts '
f'({precent:.2f}%). Biologically these should be separated by capillaries.')
if self.fractions[VesselType.ARTERY] > 0.5:
warnings.append(f'Artery fraction unusually high ({self.fractions[VesselType.ARTERY]:.1%}).')
if self.fractions[VesselType.VEIN] > 0.5:
warnings.append(f'Vein fraction unusually high ({self.fractions[VesselType.VEIN]:.1%}).')
if self.fractions[VesselType.CAPILLARY] < 0.3:
warnings.append(f'Capillary fraction unusually low ({self.fractions[VesselType.CAPILLARY]:.1%}).')
if self.seed_counts[VesselType.UNLABELLED] > 0.5 * self.n_edges:
warnings.append(f'{self.seed_counts[VesselType.UNLABELLED]} edges unlabelled after seeding. '
f'Seed thresholds may be too strict.')
for t in VesselType.competing_labels():
if self.score_stats[t]['min'] < 0.1 and self.counts[t] > 0:
warnings.append(f'Some {t.name.lower()} edges have very low scores '
f'(min={self.score_stats[t]["min"]:.3f}).')
return warnings
[docs]
def summary(self) -> str:
"""Human-readable summary for logging."""
lines = [f'Classification results ({self.n_edges:,} edges):']
for v_type in (VesselType.ARTERY, VesselType.VEIN, VesselType.CAPILLARY):
extra = ''
if v_type in self.propagated_counts:
extra = (f' [seeds: {self.seed_counts[v_type]:,}, '
f'propagated: {self.propagated_counts[v_type]:,}]')
lines.append(f' {v_type.name.title():12s} {self.counts[v_type]:>8,} '
f'({self.fractions[v_type]:6.1%}){extra}')
lines.extend([f' Propagation iterations: {self.propagation_iterations}',
f' Seed conflicts: {self.n_seed_conflicts:,}',
f' Direct artery-vein contacts: {self.n_direct_contacts:,}',
f' Fragments removed: ' + ', '.join(
f'{t.name.lower()}={self.fragments_removed[t]}'
for t in VesselType.competing_labels()),
])
warns = self.sanity_warnings
if warns:
lines.append(' ⚠ Warnings:')
for w in warns:
lines.append(f' - {w}')
return '\n'.join(lines)
# ── scoring helpers ──────────────────────────────────────────────────────────
def _sigmoid(values: np.ndarray, centre: float, width: float = 2.0) -> np.ndarray:
"""
Smooth step function mapping values to [0, 1].
``centre`` is the inflection point (output = 0.5).
``width`` controls steepness — larger = gentler transition.
"""
x = (values - centre) / max(width, 1e-6)
return 1.0 / (1.0 + np.exp(-x))
def _normalise_intensity(intensity: np.ndarray, floor: float,
ceiling: float = None) -> np.ndarray:
"""
Map intensity to [0, 1] with ``floor`` → 0 and ``ceiling`` → 1.
If ``ceiling`` is None, uses the 99th percentile of the data.
Values below floor are clipped to 0.
"""
if ceiling is None:
ceiling = np.percentile(intensity, 99) if len(intensity) > 0 else 1.0
span = max(ceiling - floor, 1e-6)
return np.clip((intensity - floor) / span, 0.0, 1.0)
def _in_range_soft(values: np.ndarray, lo: float, hi: float,
margin: float = 0.1) -> np.ndarray:
"""
Soft membership function for the interval [lo, hi].
**Why not a hard range check?**
A hard ``(values >= lo) & (values <= hi)`` produces a binary mask
with sharp transitions at the boundaries. This is problematic when
used as a score component because:
- An edge with artery intensity 2501 (just above ``hi=2500``) gets
score 0.0, while 2499 gets score 1.0. The biological signal does
not have this kind of precision.
- The sharp boundary creates unstable classification at the edges
of the intensity range — small noise pushes edges in and out.
The soft version uses two opposing sigmoids that multiply together:
::
1.0 ─────────╮ ╭─────────
╲ ╱
0.5 ╲ ╱
╲ ╱
0.0 ──────────────╰───────╯──────────────
lo hi
The ``margin`` parameter (fraction of ``hi - lo``) controls how
gradual the transition is. At ``margin=0.1`` the transition zone
is about 10% of the range width on each side.
Parameters
----------
values : np.ndarray
Input values to evaluate.
lo, hi : float
Boundaries of the target range.
margin : float
Width of the sigmoid transition as a fraction of (hi - lo).
Larger = gentler transition.
Returns
-------
np.ndarray
Soft membership scores in [0, 1]. Values well inside [lo, hi]
score ~1.0, values well outside score ~0.0.
"""
below = _sigmoid(values, centre=lo, width=margin * (hi - lo))
above = 1.0 - _sigmoid(values, centre=hi, width=margin * (hi - lo))
return below * above
def _marker_score(intensity: Optional[np.ndarray], binary: Optional[np.ndarray],
floor: float, n: int) -> np.ndarray:
"""
Combine intensity and binary mask into a single [0, 1] marker score.
The intensity is normalised to [0, 1] using ``floor`` as the zero point.
The binary mask contributes 1.0 where True.
The result is the element-wise maximum of the two contributions.
"""
score = np.zeros(n)
if intensity is not None:
score = _normalise_intensity(intensity, floor=floor)
if binary is not None:
score = np.maximum(score, binary.astype(float))
return score
# ── classifier ───────────────────────────────────────────────────────────────
[docs]
class VesselClassifier:
"""
Score-then-propagate vessel classifier for reduced vascular graphs.
Usage
-----
::
signals = VesselSignals(...)
cfg = ClassificationConfig.from_config(config_dict)
vc = VesselClassifier(graph, cfg, signals)
metrics = vc.classify()
print(metrics.summary())
Parameters
----------
graph : Graph
The annotated, reduced graph. Modified **in place**.
config : ClassificationConfig
All classification thresholds.
signals : VesselSignals
Pre-collected per-edge signal arrays.
"""
def __init__(self, graph: Graph, config: ClassificationConfig, signals: VesselSignals):
self.graph = graph
self.cfg = config
self.signals = signals
self.metrics = ClassificationMetrics(n_edges=graph.n_edges)
# ── public entry point ───────────────────────────────────────────────
[docs]
def classify(self) -> ClassificationMetrics:
"""
Run the full classification pipeline:
score → seed → propagate → cleanup.
After this call the graph has two boolean edge properties:
* ``'artery'`` — True for edges classified as artery
* ``'vein'`` — True for edges classified as vein
Edges that are neither artery nor vein are implicitly capillaries.
Returns
-------
ClassificationMetrics
Diagnostic metrics and sanity check results.
"""
# Phase 1 — score every edge
artery_score, vein_score = self._compute_edge_scores()
scores = {VesselType.ARTERY: artery_score, VesselType.VEIN: vein_score}
labels = self._seed_labels(artery_score, vein_score)
self.metrics.record_seed_counts(labels)
labels = self._propagate_labels(labels, scores)
# Everything still unlabelled becomes capillary
labels[labels == VesselType.UNLABELLED] = VesselType.CAPILLARY
labels = self._cleanup(labels)
# Write results to graph
self._labels_to_properties(labels)
# Compute final metrics
self.metrics.record_final_counts(labels, scores)
self.metrics.n_direct_contacts = self._count_direct_contacts(labels)
return self.metrics
def _relabel_small_components(self, labels: np.ndarray, vessel_type: VesselType,
min_edges: int, new_kind: VesselType) -> int:
"""
Find connected components of ``vessel_type`` smaller than ``min_edges``
and relabel them as ``new_kind``.
Uses the edge-graph trick (see ``_remove_seed_noise``) to find
connected components of edges sharing the same label.
Returns the number of edges relabelled.
"""
mask = labels == vessel_type
if not mask.any():
return 0
sub = self.graph.sub_graph(edge_filter=mask, view=True)
sub_edge, edge_map = sub.edge_graph(return_edge_map=True)
components, sizes = sub_edge.label_components(return_vertex_counts=True)
edge_mask = np.isin(components, np.where(sizes < min_edges)[0]) # indices ?? in G space
too_few = edge_map[edge_mask] # mask in L(G) space
labels[too_few] = new_kind
return len(too_few)
# ── phase 1: scoring ─────────────────────────────────────────────────
def _compute_edge_scores(self) -> tuple[np.ndarray, np.ndarray]:
"""
Compute per-edge artery and vein scores in [0, 1].
Three cases depending on which channels are available:
Case 1 — vessels only (no specific markers)
Cannot distinguish arteries from veins by signal.
Returns zero scores for both — everything stays capillary.
A warning is emitted.
Case 2 — vessels + arteries (most common)
Artery score = radius × artery marker.
Vein score = radius × in_vein_range (anti-correlation proxy:
low artery signal in a large vessel → probably vein).
Case 3 — triple labelling (vessels + arteries + veins)
Artery score = radius × artery marker.
Vein score = radius × max(direct vein marker,
anti-correlation proxy).
The max combines direct evidence (vein channel) with indirect
evidence (low artery channel) — whichever is stronger wins.
"""
s = self.signals
cfg = self.cfg
n_edges = len(s.radii)
has_artery = s.artery_binary is not None
has_vein = s.vein_binary is not None
# Radii contribution
artery_radius_score = _sigmoid(s.radii, centre=cfg.artery_trace_radius)
vein_radius_score = _sigmoid(s.radii, centre=cfg.vein_trace_radius)
if not has_artery:
# Case 1 — no markers, cannot classify vessel types by signal
warnings.warn('No artery channel available. Vessel type classification '
'requires at least an artery marker channel. '
'All edges will be classified as capillary.', RuntimeWarning, stacklevel=3)
return np.zeros(n_edges), np.zeros(n_edges)
# Labeling contribution
# Cases 2 and 3 — artery channel available
artery_marker_score = _marker_score(s.artery_intensity, s.artery_binary, cfg.artery_intensity_min, n_edges)
# combine
artery_score = artery_radius_score * artery_marker_score
# Anti-correlation: low artery signal → positive vein evidence
lo, hi = cfg.vein_intensity_range
in_vein_range = _in_range_soft(s.artery_intensity, lo, hi)
# combine
vein_score = vein_radius_score * in_vein_range
if has_vein:
# Case 3 — direct vein evidence + optional anti-correlation proxy
vein_marker_score = _marker_score(s.vein_intensity, s.vein_binary, cfg.vein_intensity_min, n_edges)
# Direct vein signal OR inferred from low artery signal
# combine with signal
vein_score = vein_radius_score * np.maximum(vein_marker_score, in_vein_range) # Avoids classifiying small vessels with label as veins
return artery_score, vein_score
# ── phase 2: seeding ─────────────────────────────────────────────────
def _seed_labels(self, artery_score: np.ndarray,
vein_score: np.ndarray) -> np.ndarray:
"""
Assign initial labels to high-confidence edges.
Order:
1. Capillaries — too small for either vessel type
2. Score-based artery/vein seeds (non-conflicting)
3. Conflict resolution — higher score wins
4. Large-radius vein fallback — huge + not artery → vein
5. Noise removal — tiny artery components removed
Returns
-------
labels : np.ndarray
dtype int8, values from VesselType enum.
"""
cfg = self.cfg
n = len(artery_score)
labels = np.full(n, VesselType.UNLABELLED, dtype=np.int8)
# 1. Definite capillaries — too small for either vessel type
capillary_radius = min(cfg.artery_trace_radius, cfg.vein_trace_radius)
labels[self.signals.radii < capillary_radius] = VesselType.CAPILLARY
# 2. Score-based seeds (non-conflicting only)
artery_seed = artery_score >= cfg.artery_seed_threshold
vein_seed = vein_score >= cfg.vein_seed_threshold
conflict = artery_seed & vein_seed
self.metrics.n_seed_conflicts = int(conflict.sum())
unlabelled = labels == VesselType.UNLABELLED
labels[artery_seed & ~conflict & unlabelled] = VesselType.ARTERY
labels[vein_seed & ~conflict & unlabelled] = VesselType.VEIN
# 3. Conflicts — higher score wins
labels[conflict & (artery_score >= vein_score)] = VesselType.ARTERY
labels[conflict & (vein_score > artery_score)] = VesselType.VEIN
# 4. Large-radius vein: whatever is still unlabelled, very large and low artery scores → vein
unlabelled = labels == VesselType.UNLABELLED
large_non_artery = ((self.signals.radii >= cfg.restrictive_vein_radius) &
(artery_score < cfg.artery_seed_threshold) &
unlabelled)
labels[large_non_artery] = VesselType.VEIN
# 5. Noise removal on artery seeds
self._relabel_small_components(labels, VesselType.ARTERY, min_edges=cfg.arteries_min_noise_edges,
new_kind=VesselType.UNLABELLED)
return labels
# ── phase 3: propagation ─────────────────────────────────────────────
def _propagate_labels(self, labels: np.ndarray, scores: dict[VesselType, np.ndarray]) -> np.ndarray:
"""
Simultaneously propagate artery and vein labels outward from seeds.
**Why the edge-graph indirection?**
Our labels live on *edges* of the reduced graph. But
``propagate_labels`` (in graph_processing) works by dilating
*vertex* labels on a graph. We need the same "dilate by one
step" semantics but for edges instead of vertices.
The solution is the same edge-graph trick used in
``_remove_seed_noise``:
- Build the edge graph of our reduced graph.
- In the edge graph, each vertex represents an original edge.
- Two edge-graph vertices are connected if their original edges
share a vertex (i.e. are adjacent in the vessel network).
- Propagating vertex labels on the edge graph is equivalent to
propagating edge labels on the original graph.
**What space are we in?**
The reduced graph has ~10M edges. Each edge represents an
entire vessel segment (a chain of degree-2 vertices that was
squeezed during graph reduction). So one "edge" here is not
a single voxel — it is an entire vessel segment with
properties like radius, artery_score, etc. that were
aggregated from the underlying geometry.
we attach all needed arrays (e.g. scores)
as edge properties on the original graph.
When graph-tool builds the edge graph (line graph),
edge properties become vertex properties automatically.
``propagate_labels`` then reads them by name from the edge graph.
::
Reduced graph (our input)
─────────────────────────
Vertices = branch points (few M)
Edges = vessel segments (~10M)
Properties on edges: radii, artery_score, vein_score, labels
↓ build edge graph
Edge graph (propagation space)
──────────────────────────────
Vertices = vessel segments (~10M)
Edges = adjacency between segments
Properties on vertices: same arrays, remapped via edge_map
↓ propagate_labels()
Updated labels in edge-graph space
↓ map back via edge_map
Updated labels on reduced graph edges
Returns
-------
labels : np.ndarray
Updated label array with VesselType.CAPILLARY for remaining
unlabelled edges.
"""
cfg = self.cfg
# Attach all propagation inputs as edge properties
self._set_propagation_properties(labels, scores)
# Build edge graph — all edge properties become vertex properties
edge_graph, edge_map = self.graph.edge_graph(return_edge_map=True)
competing_pairs = {t: t.competitor() for t in VesselType.competing_labels()}
surface_masks = {
VesselType.ARTERY: '_propagation_surface', # property name
VesselType.VEIN: None, # no surface constraint
}
min_radii = {
VesselType.ARTERY: cfg.artery_trace_radius,
VesselType.VEIN: cfg.vein_trace_radius,
}
label_priority = {
VesselType.ARTERY: 2,
VesselType.VEIN: 1,
}
score_properties = {
VesselType.ARTERY: '_propagation_score_artery',
VesselType.VEIN: '_propagation_score_vein',
}
result_labels, n_iter = graph_processing.propagate_labels(
edge_graph,
label_property='_propagation_labels', score_properties=score_properties,
competing_pairs=competing_pairs,
radii_property='_propagation_radii', min_radii=min_radii,
surface_mask_properties=surface_masks,
label_priority=label_priority,
max_iterations=cfg.max_propagation_iter, initial_threshold=cfg.initial_threshold,
threshold_growth=1.15, threshold_ceiling=cfg.propagation_floor,
unlabelled_value=VesselType.UNLABELLED)
self.metrics.propagation_iterations = n_iter
# Map back from edge-graph to original edge space
# caller will decide what to do with unlabelled edges
result = np.full(self.graph.n_edges, VesselType.UNLABELLED, dtype=np.int8)
result[edge_map] = result_labels[:len(edge_map)]
# Cleanup temporary properties
self._remove_propagation_properties()
return result
def _set_propagation_properties(self, labels: np.ndarray,
scores: dict[VesselType, np.ndarray]) -> None:
"""
Attach all arrays needed by propagation as temporary edge
properties on the original graph.
When graph-tool builds the edge graph (line graph), these
become vertex properties automatically — no manual index
mapping needed.
Property names are prefixed with ``_propagation_`` to avoid
collisions with permanent graph properties.
"""
s = self.signals
self.graph.define_edge_property('_propagation_labels',
labels.astype(np.int32))
self.graph.define_edge_property('_propagation_score_artery',
scores[VesselType.ARTERY].astype(np.float64))
self.graph.define_edge_property('_propagation_score_vein',
scores[VesselType.VEIN].astype(np.float64))
self.graph.define_edge_property('_propagation_radii',
s.radii.astype(np.float64))
if s.distance_to_surface is not None:
surface = s.distance_to_surface < self.cfg.distance_to_surface_min
else:
surface = np.zeros(self.graph.n_edges, dtype=bool)
self.graph.define_edge_property('_propagation_surface',
surface)
def _remove_propagation_properties(self) -> None:
"""Remove temporary edge properties added for propagation."""
for name in list(self.graph.edge_properties):
if name.startswith('_propagation_'):
self.graph.remove_edge_property(name)
# ── phase 4: cleanup ─────────────────────────────────────────────────
def _cleanup(self, labels: np.ndarray) -> np.ndarray:
"""Remove small connected components of each competing vessel type."""
for vessel_type, min_size in self.cfg.min_component_edges.items():
n_removed = self._relabel_small_components(labels, vessel_type, min_size, new_kind=VesselType.CAPILLARY)
self.metrics.fragments_removed[vessel_type] = n_removed
return labels
# ── write results ────────────────────────────────────────────────────
def _labels_to_properties(self, labels: np.ndarray) -> None:
"""Write the unified label array back as boolean edge properties."""
self.graph.define_edge_property('artery', labels == VesselType.ARTERY)
self.graph.define_edge_property('vein', labels == VesselType.VEIN)
# ── diagnostics ──────────────────────────────────────────────────────
def _count_direct_contacts(self, labels: np.ndarray) -> int:
"""
Count vertices where an artery edge and a vein edge meet.
Biologically arteries and veins should always be separated
by capillaries. A nonzero count indicates classification
ambiguity at transition zones.
"""
ec = self.graph.edge_connectivity()
n_vertices = self.graph.n_vertices
is_artery = labels == VesselType.ARTERY
is_vein = labels == VesselType.VEIN
touches_artery = np.zeros(n_vertices, dtype=bool)
touches_vein = np.zeros(n_vertices, dtype=bool)
artery_edges = ec[is_artery]
if len(artery_edges):
touches_artery[artery_edges[:, 0]] = True
touches_artery[artery_edges[:, 1]] = True
vein_edges = ec[is_vein]
if len(vein_edges):
touches_vein[vein_edges[:, 0]] = True
touches_vein[vein_edges[:, 1]] = True
return int((touches_artery & touches_vein).sum())
