From c60123daeb94693e219a3bb73e812519e29e99e2 Mon Sep 17 00:00:00 2001 From: trivoldus28 Date: Fri, 13 Mar 2026 08:39:40 -0700 Subject: [PATCH 01/16] Add getting mean aff metadata Co-Authored-By: Claude Opus 4.6 --- src/waterz/_agglomerate.py | 6 +++-- src/waterz/agglomerate.pyx | 16 +++++++++--- src/waterz/backend/IterativeRegionMerging.hpp | 26 +++++++++++++++++++ src/waterz/backend/MeanAffinityProvider.hpp | 4 +++ src/waterz/backend/StatisticsProvider.hpp | 3 +++ src/waterz/frontend_agglomerate.cpp | 10 +++++++ src/waterz/frontend_agglomerate.h | 1 + 7 files changed, 60 insertions(+), 6 deletions(-) diff --git a/src/waterz/_agglomerate.py b/src/waterz/_agglomerate.py index 63d0f9d..bda950b 100644 --- a/src/waterz/_agglomerate.py +++ b/src/waterz/_agglomerate.py @@ -15,14 +15,15 @@ def agglomerate( - affs: NDArray[np.float32], - thresholds: Sequence[float], + affs: NDArray[np.float32] | None = None, + thresholds: Sequence[float] | None = None, gt: NDArray[np.uint32] | None = None, fragments: NDArray[np.uint64] | None = None, aff_threshold_low: float = 0.0001, aff_threshold_high: float = 0.9999, return_merge_history: bool = False, return_region_graph: bool = False, + return_region_graph_metadata: bool = False, scoring_function: str = "OneMinus>", discretize_queue: int = 0, force_rebuild: bool = False, @@ -186,4 +187,5 @@ def agglomerate( aff_threshold_high, return_merge_history, return_region_graph, + return_region_graph_metadata, ) diff --git a/src/waterz/agglomerate.pyx b/src/waterz/agglomerate.pyx index 07a3583..2d58cb2 100644 --- a/src/waterz/agglomerate.pyx +++ b/src/waterz/agglomerate.pyx @@ -5,18 +5,20 @@ import numpy as np cimport numpy as np def agglomerate( - affs, - thresholds, + affs=None, + thresholds=None, gt=None, fragments=None, aff_threshold_low=0.0001, aff_threshold_high=0.9999, return_merge_history=False, - return_region_graph=False): + return_region_graph=False, + return_region_graph_metadata=False, + ): # the C++ part assumes contiguous memory, make sure we have it (and do # nothing, if we do) - if not affs.flags['C_CONTIGUOUS']: + if affs is not None and not affs.flags['C_CONTIGUOUS']: print("Creating memory-contiguous affinity arrray (avoid this by passing C_CONTIGUOUS arrays)") affs = np.ascontiguousarray(affs) if gt is not None and not gt.flags['C_CONTIGUOUS']: @@ -63,6 +65,10 @@ def agglomerate( result += (getRegionGraph(state),) + if return_region_graph_metadata: + + result += (getRegionGraphMeta(state),) + if len(result) == 1: yield result[0] else: @@ -136,4 +142,6 @@ cdef extern from "frontend_agglomerate.h": vector[ScoredEdge] getRegionGraph(WaterzState& state) + vector[double] getRegionGraphMeta(WaterzState& state) + void free(WaterzState& state) diff --git a/src/waterz/backend/IterativeRegionMerging.hpp b/src/waterz/backend/IterativeRegionMerging.hpp index 9689714..d3f5062 100644 --- a/src/waterz/backend/IterativeRegionMerging.hpp +++ b/src/waterz/backend/IterativeRegionMerging.hpp @@ -163,6 +163,32 @@ class IterativeRegionMerging { return edges; } + template + std::vector extractRegionGraphMeta(EdgeScoringFunction& edgeScoringFunction, StatisticsProviderType& statisticsProvider) { + + std::vector ret; + + for (EdgeIdType e = 0; e < _regionGraph.numEdges(); e++) { + + if (_deleted[e]) + continue; + + ScoreType score; + if (_stale[e]) + score = scoreEdge(e, edgeScoringFunction); + else + score = _edgeScores[e]; + + if (score < _mergedUntil) + continue; + + // ret.push_back(EdgeMetadata(statisticsProvider.getEdgeMetadata(e)) ); + ret.emplace_back(statisticsProvider.getEdgeMetadata(e)); + } + + return ret; + } + private: /** diff --git a/src/waterz/backend/MeanAffinityProvider.hpp b/src/waterz/backend/MeanAffinityProvider.hpp index 5ff45a4..34cef35 100644 --- a/src/waterz/backend/MeanAffinityProvider.hpp +++ b/src/waterz/backend/MeanAffinityProvider.hpp @@ -46,6 +46,10 @@ class MeanAffinityProvider : public StatisticsProvider { return _meanAffinities[e]; } + inline double getEdgeMetadata(EdgeIdType e) const { + return _numValues[e]; + } + private: typename RegionGraphType::template EdgeMap _numValues; diff --git a/src/waterz/backend/StatisticsProvider.hpp b/src/waterz/backend/StatisticsProvider.hpp index fe7ec82..c551518 100644 --- a/src/waterz/backend/StatisticsProvider.hpp +++ b/src/waterz/backend/StatisticsProvider.hpp @@ -37,6 +37,9 @@ class StatisticsProvider { */ template inline bool notifyEdgeMerge(EdgeIdType from, EdgeIdType to) { return false; } + + template + inline double getEdgeMetadata(EdgeIdType e) const { return 0; } }; #endif // WATERZ_STATISTICS_PROVIDER_H__ diff --git a/src/waterz/frontend_agglomerate.cpp b/src/waterz/frontend_agglomerate.cpp index 100bcd6..39831ed 100644 --- a/src/waterz/frontend_agglomerate.cpp +++ b/src/waterz/frontend_agglomerate.cpp @@ -164,6 +164,16 @@ getRegionGraph(WaterzState& state) { return regionMerging->extractRegionGraph(*scoringFunction); } +std::vector +getRegionGraphMeta(WaterzState& state) { + + WaterzContext* context = WaterzContext::get(state.context); + std::shared_ptr regionMerging = context->regionMerging; + std::shared_ptr scoringFunction = context->scoringFunction; + + return regionMerging->extractRegionGraphMeta(*scoringFunction, *(context->statisticsProvider)); +} + void free(WaterzState& state) { diff --git a/src/waterz/frontend_agglomerate.h b/src/waterz/frontend_agglomerate.h index 7724ed7..333465b 100644 --- a/src/waterz/frontend_agglomerate.h +++ b/src/waterz/frontend_agglomerate.h @@ -157,6 +157,7 @@ std::vector mergeUntil( float threshold); std::vector getRegionGraph(WaterzState& state); +std::vector getRegionGraphMeta(WaterzState& state); void free(WaterzState& state); From 38c0d615d15d3168a28a7b2f42843c1cb3d7317e Mon Sep 17 00:00:00 2001 From: trivoldus28 Date: Fri, 13 Mar 2026 08:40:13 -0700 Subject: [PATCH 02/16] Add initialization using rag Co-Authored-By: Claude Opus 4.6 --- src/waterz/_agglomerate.py | 10 +++ src/waterz/agglomerate.pyx | 52 ++++++++++++++- src/waterz/backend/MeanAffinityProvider.hpp | 6 ++ src/waterz/backend/StatisticsProvider.hpp | 3 + src/waterz/backend/region_graph.hpp | 21 ++++++ src/waterz/frontend_agglomerate.cpp | 72 ++++++++++++++++++++- src/waterz/frontend_agglomerate.h | 10 +++ 7 files changed, 171 insertions(+), 3 deletions(-) diff --git a/src/waterz/_agglomerate.py b/src/waterz/_agglomerate.py index bda950b..8017259 100644 --- a/src/waterz/_agglomerate.py +++ b/src/waterz/_agglomerate.py @@ -19,6 +19,8 @@ def agglomerate( thresholds: Sequence[float] | None = None, gt: NDArray[np.uint32] | None = None, fragments: NDArray[np.uint64] | None = None, + input_rag=None, + input_rag_metadata=None, aff_threshold_low: float = 0.0001, aff_threshold_high: float = 0.9999, return_merge_history: bool = False, @@ -178,6 +180,14 @@ def agglomerate( ) # call compiled function + if input_rag is not None or input_rag_metadata is not None: + return module.agglomerate_rag( + rag=input_rag, + rag_metadata=input_rag_metadata, + thresholds=thresholds, + fragments=fragments, + ) + return module.agglomerate( affs, thresholds, diff --git a/src/waterz/agglomerate.pyx b/src/waterz/agglomerate.pyx index 2d58cb2..b2de541 100644 --- a/src/waterz/agglomerate.pyx +++ b/src/waterz/agglomerate.pyx @@ -4,9 +4,30 @@ from libcpp cimport bool import numpy as np cimport numpy as np +def agglomerate_rag( + rag, + rag_metadata, + thresholds, + fragments=None, + ): + + if fragments is not None and not fragments.flags['C_CONTIGUOUS']: + print("Creating memory-contiguous fragments arrray (avoid this by passing C_CONTIGUOUS arrays)") + fragments = np.ascontiguousarray(fragments) + + cdef WaterzState state = __initialize_with_rag(rag, rag_metadata, fragments) + + thresholds.sort() + for threshold in thresholds: + merge_history = mergeUntil(state, threshold) + yield merge_history, fragments + + free(state) + + def agglomerate( - affs=None, - thresholds=None, + affs, + thresholds, gt=None, fragments=None, aff_threshold_low=0.0001, @@ -102,6 +123,25 @@ def __initialize( aff_threshold_high, find_fragments) +def __initialize_with_rag( + rag, + rag_metadata, + np.ndarray[uint64_t, ndim=3] segmentation = None): + + cdef uint64_t* segmentation_data = NULL + shape = (0, 0, 0) + + if segmentation is not None: + segmentation_data = &segmentation[0,0,0] + shape = (segmentation.shape[0], segmentation.shape[1], segmentation.shape[2]) + + + return initialize_with_rag( + rag, + rag_metadata, + segmentation_data, + shape[0], shape[1], shape[2]) + cdef extern from "frontend_agglomerate.h": struct Metrics: @@ -136,6 +176,14 @@ cdef extern from "frontend_agglomerate.h": float affThresholdHigh, bool findFragments); + WaterzState initialize_with_rag( + const vector[ScoredEdge]& rag, + const vector[double]& rag_metadata, + uint64_t* segmentation_data, + size_t width, + size_t height, + size_t depth,); + vector[Merge] mergeUntil( WaterzState& state, float threshold) diff --git a/src/waterz/backend/MeanAffinityProvider.hpp b/src/waterz/backend/MeanAffinityProvider.hpp index 34cef35..a8a3b3a 100644 --- a/src/waterz/backend/MeanAffinityProvider.hpp +++ b/src/waterz/backend/MeanAffinityProvider.hpp @@ -27,6 +27,12 @@ class MeanAffinityProvider : public StatisticsProvider { _numValues[e]++; } + template + inline void addEdge(EdgeIdType e, ScoreType affinity, double value) { + _meanAffinities[e] = affinity; + _numValues[e] = value; + } + inline bool notifyEdgeMerge(EdgeIdType from, EdgeIdType to) { size_t fromN = _numValues[from]; diff --git a/src/waterz/backend/StatisticsProvider.hpp b/src/waterz/backend/StatisticsProvider.hpp index c551518..96a040f 100644 --- a/src/waterz/backend/StatisticsProvider.hpp +++ b/src/waterz/backend/StatisticsProvider.hpp @@ -24,6 +24,9 @@ class StatisticsProvider { template inline void addVoxel(NodeIdType n, std::size_t x, std::size_t y, std::size_t z) {} + template + inline void addAffinity(EdgeIdType e, ScoreType affinity, double metadata) {} + /** * Callback for node merges: 'from' will be merged into 'to'. Return true, * if this changed the statistics of this provider. diff --git a/src/waterz/backend/region_graph.hpp b/src/waterz/backend/region_graph.hpp index 4af7fae..2b2553f 100644 --- a/src/waterz/backend/region_graph.hpp +++ b/src/waterz/backend/region_graph.hpp @@ -82,3 +82,24 @@ get_region_graph( std::cout << "Region graph number of edges: " << rg.edges().size() << std::endl; } + +template +inline +void +initialize_with_region_graph( + StatisticsProviderType& statisticsProvider, + RegionGraph& rg, + const std::vector& edges, + const std::vector& edges_metadata) { + + typedef RegionGraph RegionGraphType; + typedef typename RegionGraphType::EdgeIdType EdgeIdType; + + for (int i = 0; i < edges.size(); ++i) { + const auto& edge = edges[i]; + double size = edges_metadata[i]; + auto aff = 1.0 - edge.score; + EdgeIdType e = rg.addEdge(edge.u, edge.v); + statisticsProvider.addEdge(e, aff, size); + } +} diff --git a/src/waterz/frontend_agglomerate.cpp b/src/waterz/frontend_agglomerate.cpp index 39831ed..6f5836c 100644 --- a/src/waterz/frontend_agglomerate.cpp +++ b/src/waterz/frontend_agglomerate.cpp @@ -13,6 +13,76 @@ std::map WaterzContext::_contexts; int WaterzContext::_nextId = 0; +std::vector getRegionGraph(WaterzState& state); +std::vector getRegionGraphMeta(WaterzState& state); + +WaterzState +initialize_with_rag( + const std::vector& rag, + const std::vector& rag_metadata, + SegID* segmentation_data, + std::size_t width, + std::size_t height, + std::size_t depth) { + + std::size_t maxId = 0; + for (const auto& edge : rag) { + auto max_uv = std::max(edge.u, edge.v); + maxId = std::max(max_uv, maxId); + } + + std::size_t numNodes = maxId + 1; + std::cout << "creating region graph for " << numNodes << " nodes" << std::endl; + + std::shared_ptr regionGraph( + new RegionGraphType(numNodes) + ); + + std::cout << "creating statistics provider" << std::endl; + std::shared_ptr statisticsProvider( + new StatisticsProviderType(*regionGraph) + ); + + std::cout << "initializing region graph..." << std::endl; + + initialize_with_region_graph( + *statisticsProvider, + *regionGraph, + rag, + rag_metadata); + + std::shared_ptr scoringFunction( + new ScoringFunctionType(*regionGraph, *statisticsProvider) + ); + + std::shared_ptr regionMerging( + new RegionMergingType(*regionGraph) + ); + + WaterzContext* context = WaterzContext::createNew(); + context->regionGraph = regionGraph; + context->regionMerging = regionMerging; + context->scoringFunction = scoringFunction; + context->statisticsProvider = statisticsProvider; + // context->segmentation = segmentation_data; + + if (segmentation_data != NULL) { + // wrap data (no copy) + volume_ref_ptr segmentation( + new volume_ref( + segmentation_data, + boost::extents[width][height][depth] + ) + ); + context->segmentation = segmentation; + } + + WaterzState initial_state; + initial_state.context = context->id; + + return initial_state; +} + WaterzState initialize( std::size_t width, @@ -132,7 +202,7 @@ mergeUntil( threshold, mergeHistoryVisitor); - if (merged) { + if (merged && context->segmentation) { std::cout << "extracting segmentation" << std::endl; diff --git a/src/waterz/frontend_agglomerate.h b/src/waterz/frontend_agglomerate.h index 333465b..e9b7618 100644 --- a/src/waterz/frontend_agglomerate.h +++ b/src/waterz/frontend_agglomerate.h @@ -44,6 +44,8 @@ struct Merge { struct ScoredEdge { + ScoredEdge() {}; + ScoredEdge(SegID u_, SegID v_, ScoreValue score_) : u(u_), v(v_), @@ -152,6 +154,14 @@ WaterzState initialize( AffValue affThresholdHigh = 0.9999, bool findFragments = true); +WaterzState initialize_with_rag( + const std::vector& rag, + const std::vector& rag_metadata, + SegID* segmentation_data, + std::size_t width, + std::size_t height, + std::size_t depth); + std::vector mergeUntil( WaterzState& state, float threshold); From 5439b20ca3f19e2ffc27487b15c5fcf62d6df513 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Mon, 30 Mar 2026 17:40:44 -0700 Subject: [PATCH 03/16] feat: add semantic/size/seg constraints Co-Authored-By: trivoldus28 --- src/waterz/_agglomerate.py | 34 +++-- src/waterz/agglomerate.pyx | 131 ++++++++++++++---- src/waterz/backend/ConstraintProvider.hpp | 13 ++ src/waterz/backend/DefaultDict.hpp | 74 ++++++++++ src/waterz/backend/IterativeRegionMerging.hpp | 27 +++- src/waterz/backend/SegConstraintProvider.hpp | 52 +++++++ .../backend/SemanticConstraintProvider.hpp | 72 ++++++++++ .../SizeHeuristicConstraintProvider.hpp | 49 +++++++ src/waterz/frontend_agglomerate.cpp | 62 ++++++++- src/waterz/frontend_agglomerate.h | 21 ++- 10 files changed, 486 insertions(+), 49 deletions(-) create mode 100644 src/waterz/backend/ConstraintProvider.hpp create mode 100644 src/waterz/backend/DefaultDict.hpp create mode 100644 src/waterz/backend/SegConstraintProvider.hpp create mode 100644 src/waterz/backend/SemanticConstraintProvider.hpp create mode 100644 src/waterz/backend/SizeHeuristicConstraintProvider.hpp diff --git a/src/waterz/_agglomerate.py b/src/waterz/_agglomerate.py index 8017259..944ee83 100644 --- a/src/waterz/_agglomerate.py +++ b/src/waterz/_agglomerate.py @@ -19,6 +19,8 @@ def agglomerate( thresholds: Sequence[float] | None = None, gt: NDArray[np.uint32] | None = None, fragments: NDArray[np.uint64] | None = None, + semantic: NDArray[np.uint8] | None = None, + segconstraint: NDArray[np.uint64] | None = None, input_rag=None, input_rag_metadata=None, aff_threshold_low: float = 0.0001, @@ -27,6 +29,12 @@ def agglomerate( return_region_graph: bool = False, return_region_graph_metadata: bool = False, scoring_function: str = "OneMinus>", + semantic_aff_threshold: float = 0.5, + semantic_size_threshold: int = 100_000, + semantic_signal_ratio: float = 0.6, + size_heuristic_aff_threshold: float = 1.0, + size_heuristic_small_threshold: int = 1_000_000, + size_heuristic_large_threshold: int = 10_000_000, discretize_queue: int = 0, force_rebuild: bool = False, ) -> Iterator[tuple | NDArray[np.uint64]]: @@ -189,13 +197,21 @@ def agglomerate( ) return module.agglomerate( - affs, - thresholds, - gt, - fragments, - aff_threshold_low, - aff_threshold_high, - return_merge_history, - return_region_graph, - return_region_graph_metadata, + affs=affs, + thresholds=thresholds, + gt=gt, + fragments=fragments, + semantic=semantic, + segconstraint=segconstraint, + aff_threshold_low=aff_threshold_low, + aff_threshold_high=aff_threshold_high, + semantic_aff_threshold=semantic_aff_threshold, + semantic_size_threshold=semantic_size_threshold, + semantic_signal_ratio=semantic_signal_ratio, + size_heuristic_aff_threshold=size_heuristic_aff_threshold, + size_heuristic_small_threshold=size_heuristic_small_threshold, + size_heuristic_large_threshold=size_heuristic_large_threshold, + return_merge_history=return_merge_history, + return_region_graph=return_region_graph, + return_region_graph_metadata=return_region_graph_metadata, ) diff --git a/src/waterz/agglomerate.pyx b/src/waterz/agglomerate.pyx index b2de541..6083d62 100644 --- a/src/waterz/agglomerate.pyx +++ b/src/waterz/agglomerate.pyx @@ -1,6 +1,7 @@ from libcpp.vector cimport vector -from libc.stdint cimport uint64_t, uint32_t +from libc.stdint cimport uint64_t, uint32_t, uint8_t from libcpp cimport bool + import numpy as np cimport numpy as np @@ -8,7 +9,7 @@ def agglomerate_rag( rag, rag_metadata, thresholds, - fragments=None, + fragments, ): if fragments is not None and not fragments.flags['C_CONTIGUOUS']: @@ -26,16 +27,28 @@ def agglomerate_rag( def agglomerate( - affs, - thresholds, - gt=None, - fragments=None, - aff_threshold_low=0.0001, - aff_threshold_high=0.9999, - return_merge_history=False, - return_region_graph=False, - return_region_graph_metadata=False, - ): + affs: np.array | None, + thresholds: np.array | None, + gt: np.array | None, + fragments: np.array | None, + semantic: np.array | None, + segconstraint: np.array | None, + + aff_threshold_low: float, + aff_threshold_high: float, + + semantic_aff_threshold: float, + semantic_size_threshold: int, + semantic_signal_ratio: float, + + size_heuristic_aff_threshold: float, + size_heuristic_small_threshold: int, + size_heuristic_large_threshold: int, + + return_merge_history: bool, + return_region_graph: bool, + return_region_graph_metadata: bool, + ): # the C++ part assumes contiguous memory, make sure we have it (and do # nothing, if we do) @@ -48,6 +61,12 @@ def agglomerate( if fragments is not None and not fragments.flags['C_CONTIGUOUS']: print("Creating memory-contiguous fragments arrray (avoid this by passing C_CONTIGUOUS arrays)") fragments = np.ascontiguousarray(fragments) + if semantic is not None and not semantic.flags['C_CONTIGUOUS']: + print("Creating memory-contiguous semantic arrray (avoid this by passing C_CONTIGUOUS arrays)") + semantic = np.ascontiguousarray(semantic) + if segconstraint is not None and not segconstraint.flags['C_CONTIGUOUS']: + print("Creating memory-contiguous segconstraint arrray (avoid this by passing C_CONTIGUOUS arrays)") + segconstraint = np.ascontiguousarray(segconstraint) print("Preparing segmentation volume...") @@ -59,7 +78,22 @@ def agglomerate( segmentation = fragments find_fragments = False - cdef WaterzState state = __initialize(affs, segmentation, gt, aff_threshold_low, aff_threshold_high, find_fragments) + cdef WaterzState state = __initialize( + affs=affs, + segmentation=segmentation, + gt=gt, + semantic=semantic, + segconstraint=segconstraint, + aff_threshold_low=aff_threshold_low, + aff_threshold_high=aff_threshold_high, + semantic_aff_threshold=semantic_aff_threshold, + semantic_size_threshold=semantic_size_threshold, + semantic_signal_ratio=semantic_signal_ratio, + size_heuristic_aff_threshold=size_heuristic_aff_threshold, + size_heuristic_small_threshold=size_heuristic_small_threshold, + size_heuristic_large_threshold=size_heuristic_large_threshold, + find_fragments=find_fragments, + ) thresholds.sort() for threshold in thresholds: @@ -99,29 +133,68 @@ def agglomerate( def __initialize( np.ndarray[np.float32_t, ndim=4] affs, - np.ndarray[uint64_t, ndim=3] segmentation, - np.ndarray[uint32_t, ndim=3] gt = None, - aff_threshold_low = 0.0001, - aff_threshold_high = 0.9999, - find_fragments = True): + np.ndarray[uint64_t, ndim=3] segmentation, + np.ndarray[uint32_t, ndim=3] gt, + np.ndarray[uint8_t, ndim=3] semantic, + np.ndarray[uint64_t, ndim=3] segconstraint, + aff_threshold_low: float, + aff_threshold_high: float, + semantic_aff_threshold: float, + semantic_size_threshold: int, + semantic_signal_ratio: float, + size_heuristic_aff_threshold: float, + size_heuristic_small_threshold: int, + size_heuristic_large_threshold: int, + find_fragments: bool, + ): cdef float* aff_data cdef uint64_t* segmentation_data cdef uint32_t* gt_data = NULL + cdef uint8_t* semantic_data = NULL + cdef uint64_t* segconstraint_data = NULL aff_data = &affs[0,0,0,0] segmentation_data = &segmentation[0,0,0] if gt is not None: gt_data = >[0,0,0] + if semantic is not None: + semantic_data = &semantic[0,0,0] + if segconstraint is not None: + segconstraint_data = &segconstraint[0,0,0] + + # return initialize( + # affs.shape[1], affs.shape[2], affs.shape[3], + # aff_data, + # segmentation_data, + # gt_data, + # aff_threshold_low, + # aff_threshold_high, + # find_fragments return initialize( - affs.shape[1], affs.shape[2], affs.shape[3], - aff_data, - segmentation_data, - gt_data, - aff_threshold_low, - aff_threshold_high, - find_fragments) + width=affs.shape[1], + height=affs.shape[2], + depth=affs.shape[3], + affinity_data=aff_data, + segmentation_data=segmentation_data, + groundtruth_data=gt_data, + semantic_data=semantic_data, + segconstraint_data=segconstraint_data, + + affThresholdLow=aff_threshold_low, + affThresholdHigh=aff_threshold_high, + + semantic_aff_threshold=semantic_aff_threshold, + semantic_size_threshold=semantic_size_threshold, + semantic_signal_ratio=semantic_signal_ratio, + + size_heuristic_aff_threshold=size_heuristic_aff_threshold, + size_heuristic_small_threshold=size_heuristic_small_threshold, + size_heuristic_large_threshold=size_heuristic_large_threshold, + + findFragments=find_fragments, + ) def __initialize_with_rag( rag, @@ -172,8 +245,16 @@ cdef extern from "frontend_agglomerate.h": const float* affinity_data, uint64_t* segmentation_data, const uint32_t* groundtruth_data, + const uint8_t* semantic_data, + const uint64_t* segconstraint_data, float affThresholdLow, float affThresholdHigh, + float semantic_aff_threshold, + uint64_t semantic_size_threshold, + float semantic_signal_ratio, + float size_heuristic_aff_threshold, + uint64_t size_heuristic_small_threshold, + uint64_t size_heuristic_large_threshold, bool findFragments); WaterzState initialize_with_rag( diff --git a/src/waterz/backend/ConstraintProvider.hpp b/src/waterz/backend/ConstraintProvider.hpp new file mode 100644 index 0000000..a2cbd09 --- /dev/null +++ b/src/waterz/backend/ConstraintProvider.hpp @@ -0,0 +1,13 @@ +#ifndef WATERZ_CONSTRAINT_PROVIDER_H__ +#define WATERZ_CONSTRAINT_PROVIDER_H__ + +/** + * Base class for statistics providers with fallback implementations. + */ +class ConstraintProvider { +public: + virtual inline bool notifyNodeMerge(uint64_t from, uint64_t to) = 0; + virtual inline bool isConstrained(uint64_t from, uint64_t to, float score) const = 0; +}; + +#endif // WATERZ_CONSTRAINT_PROVIDER_H__ diff --git a/src/waterz/backend/DefaultDict.hpp b/src/waterz/backend/DefaultDict.hpp new file mode 100644 index 0000000..aaa45c2 --- /dev/null +++ b/src/waterz/backend/DefaultDict.hpp @@ -0,0 +1,74 @@ +#ifndef WATERZ_DEFAULTDICT_H__ +#define WATERZ_DEFAULTDICT_H__ + +#include + +template +uint64_t getDictSum(const T& map) { + uint64_t sum = 0; + for (const auto& pair : map) { + sum += pair.second; + } + return sum; +} + +template +auto getDictMaxKey(const T& map) -> decltype(map.begin()->first) { + if (map.empty()) { + throw std::runtime_error("Map is empty"); + } + auto max_key = map.begin()->first; + auto max_value = map.begin()->second; + for (const auto& pair : map) { + if (pair.second > max_value) { + max_key = pair.first; + max_value = pair.second; + } + } + return max_key; +} + +template +class DefaultDict { +private: + std::unordered_map container; + V default_value; + +public: + DefaultDict() : default_value() {} + DefaultDict(const V& default_val) : default_value(default_val) {} + + V& operator[](const K& key) { + if (container.find(key) == container.end()) { + container[key] = default_value; + } + return container[key]; + } + + V operator[](const K& key) const { + auto it = container.find(key); + if (it != container.end()) { + return it->second; + } + return default_value; + } + + // V at(const K& key) const { + // auto it = container.find(key); + // if (it != container.end()) { + // return it->second; + // } + // return default_value; + // } + + void erase(const K& key) { + if (container.find(key) != container.end()) { + container.erase(key); + } + } + + const std::unordered_map& getContainer() const { return container; } +}; + + +#endif // WATERZ_DEFAULTDICT_H__ diff --git a/src/waterz/backend/IterativeRegionMerging.hpp b/src/waterz/backend/IterativeRegionMerging.hpp index d3f5062..5665f39 100644 --- a/src/waterz/backend/IterativeRegionMerging.hpp +++ b/src/waterz/backend/IterativeRegionMerging.hpp @@ -10,6 +10,7 @@ #include "RegionGraph.hpp" #include "PriorityQueue.hpp" +#include "ConstraintProvider.hpp" template class QueueType = PriorityQueue> class IterativeRegionMerging { @@ -37,6 +38,7 @@ class IterativeRegionMerging { std::size_t mergeUntil( EdgeScoringFunction& edgeScoringFunction, StatisticsProviderType& statisticsProvider, + std::vector& constraints, ScoreType threshold, Visitor& visitor) { @@ -100,12 +102,23 @@ class IterativeRegionMerging { continue; } - NodeIdType newRegion = mergeRegions(next, statisticsProvider); + NodeIdType a = _regionGraph.edge(next).u; + NodeIdType b = _regionGraph.edge(next).v; + + bool is_constrained = false; + for (auto constraint : constraints) { + is_constrained |= constraint->isConstrained(a, b, score); + } + + if (is_constrained) { + continue; // skip merging + } + + NodeIdType newRegion = mergeRegions(next, statisticsProvider, constraints); merged++; visitor.onMerge( - _regionGraph.edge(next).u, - _regionGraph.edge(next).v, + a, b, newRegion, score); } @@ -182,7 +195,6 @@ class IterativeRegionMerging { if (score < _mergedUntil) continue; - // ret.push_back(EdgeMetadata(statisticsProvider.getEdgeMetadata(e)) ); ret.emplace_back(statisticsProvider.getEdgeMetadata(e)); } @@ -197,7 +209,8 @@ class IterativeRegionMerging { template NodeIdType mergeRegions( EdgeIdType e, - StatisticsProviderType& statisticsProvider) { + StatisticsProviderType& statisticsProvider, + std::vector& constraints) { NodeIdType a = _regionGraph.edge(e).u; NodeIdType b = _regionGraph.edge(e).v; @@ -205,6 +218,10 @@ class IterativeRegionMerging { // assign new node a = a + b bool nodeStatisticsChanged = statisticsProvider.notifyNodeMerge(b, a); + for (auto constraint : constraints) { + nodeStatisticsChanged |= constraint->notifyNodeMerge(b, a); + } + // set path _rootPaths[b] = a; diff --git a/src/waterz/backend/SegConstraintProvider.hpp b/src/waterz/backend/SegConstraintProvider.hpp new file mode 100644 index 0000000..d28a724 --- /dev/null +++ b/src/waterz/backend/SegConstraintProvider.hpp @@ -0,0 +1,52 @@ +#include "ConstraintProvider.hpp" +#include "DefaultDict.hpp" + +#include + +template +class SegConstraintProvider : public ConstraintProvider { + + typedef typename RegionGraphType::NodeIdType NodeIdType; + +private: + std::unordered_map _constraint; + +public: + SegConstraintProvider( + const SegType *segconstraint_data, + const SegType *seg_data, + size_t num_voxels) { + + // Collect all mapped constraint, accounting for when an object + // is straddled across multiple constraint + DefaultDict> constraint{ + DefaultDict{0} + }; + for (std::size_t i = 0; i < num_voxels; i++) { + constraint[seg_data[i]][segconstraint_data[i]] += 1; + } + + // Now we get a single dominant constraint per object + for (const auto& pair : constraint.getContainer()) { + const auto& segid = pair.first; + const auto& mapped_constraint_ids = pair.second; + _constraint[segid] = getDictMaxKey(mapped_constraint_ids.getContainer()); + } + } + + inline bool notifyNodeMerge(NodeIdType from, NodeIdType to) { + if (_constraint.at(to) == 0) { + _constraint[to] = _constraint.at(from); + } + _constraint.erase(from); + return true; + } + + inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const { + if (_constraint.at(from) == 0 || _constraint.at(to) == 0) + return false; + if (_constraint.at(from) == _constraint.at(to)) + return false; + return true; + } +}; diff --git a/src/waterz/backend/SemanticConstraintProvider.hpp b/src/waterz/backend/SemanticConstraintProvider.hpp new file mode 100644 index 0000000..c8c8495 --- /dev/null +++ b/src/waterz/backend/SemanticConstraintProvider.hpp @@ -0,0 +1,72 @@ +#include "ConstraintProvider.hpp" +#include "DefaultDict.hpp" + +#include + +template +class SemanticConstraintProvider : public ConstraintProvider { + + typedef SemValue ValueType; + typedef typename RegionGraphType::EdgeIdType EdgeIdType; + typedef typename RegionGraphType::NodeIdType NodeIdType; + +private: + DefaultDict> _semantic{ + DefaultDict{0} + }; + float semantic_aff_threshold; + uint64_t semantic_size_threshold; + float semantic_signal_ratio; + +public: + SemanticConstraintProvider( + const SemValue *semantic_data, + const SegType *seg_data, + size_t num_voxels, + float semantic_aff_threshold, + uint64_t semantic_size_threshold, + float semantic_signal_ratio + ) : + semantic_aff_threshold(semantic_aff_threshold), + semantic_size_threshold(semantic_size_threshold), + semantic_signal_ratio(semantic_signal_ratio) + { + for (std::size_t i = 0; i < num_voxels; i++) { + _semantic[seg_data[i]][semantic_data[i]] += 1; + } + } + + inline bool notifyNodeMerge(NodeIdType from, NodeIdType to) { + for (auto k : _semantic[from].getContainer()) { + // _semantic[to][k] += _semantic[from][k]; + _semantic[to][k.first] += k.second; + } + _semantic.erase(from); + return true; + } + + inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const { + + if (score < semantic_aff_threshold) + return false; + + auto max_sem1_label = getDictMaxKey(_semantic[from].getContainer()); + auto max_sem1 = _semantic[from][max_sem1_label]; + auto total_sem1 = getDictSum(_semantic[from].getContainer()); + auto max_sem2_label = getDictMaxKey(_semantic[to].getContainer()); + auto max_sem2 = _semantic[to][max_sem2_label]; + auto total_sem2 = getDictSum(_semantic[to].getContainer()); + + if (total_sem1 < semantic_size_threshold) + return false; + if (total_sem2 < semantic_size_threshold) + return false; + if (max_sem1 < semantic_signal_ratio * total_sem1) + return false; + if (max_sem2 < semantic_signal_ratio * total_sem2) + return false; + if (max_sem1_label == max_sem2_label) + return false; + return true; + } +}; diff --git a/src/waterz/backend/SizeHeuristicConstraintProvider.hpp b/src/waterz/backend/SizeHeuristicConstraintProvider.hpp new file mode 100644 index 0000000..ee79fe3 --- /dev/null +++ b/src/waterz/backend/SizeHeuristicConstraintProvider.hpp @@ -0,0 +1,49 @@ +#include "StatisticsProvider.hpp" +#include "DefaultDict.hpp" + +template +class SizeHeuristicConstraintProvider : public ConstraintProvider { + + typedef typename RegionGraphType::NodeIdType NodeIdType; + +private: + DefaultDict _size{0}; + float size_heuristic_aff_threshold; + size_t size_heuristic_small_threshold; + size_t size_heuristic_large_threshold; + +public: + SizeHeuristicConstraintProvider( + const SegType *seg_data, + size_t num_voxels, + float size_heuristic_aff_threshold, + size_t size_heuristic_small_threshold, + size_t size_heuristic_large_threshold + ): + size_heuristic_aff_threshold(size_heuristic_aff_threshold), + size_heuristic_small_threshold(size_heuristic_small_threshold), + size_heuristic_large_threshold(size_heuristic_large_threshold) + { + for (std::size_t i = 0; i < num_voxels; i++) { + _size[seg_data[i]] += 1; + } + } + + inline bool notifyNodeMerge(NodeIdType from, NodeIdType to) override { + _size[to] += _size[from]; + _size.erase(from); + return true; // statistics changed + } + + inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const override { + if (score < size_heuristic_aff_threshold) + return false; + if (_size[from] < size_heuristic_small_threshold) + return false; + if (_size[to] < size_heuristic_small_threshold) + return false; + if ((_size[from] + _size[to]) < size_heuristic_large_threshold) + return false; + return true; + } +}; diff --git a/src/waterz/frontend_agglomerate.cpp b/src/waterz/frontend_agglomerate.cpp index 6f5836c..e0a236e 100644 --- a/src/waterz/frontend_agglomerate.cpp +++ b/src/waterz/frontend_agglomerate.cpp @@ -10,6 +10,10 @@ #include "backend/basic_watershed.hpp" #include "backend/region_graph.hpp" +#include "backend/SemanticConstraintProvider.hpp" +#include "backend/SegConstraintProvider.hpp" +#include "backend/SizeHeuristicConstraintProvider.hpp" + std::map WaterzContext::_contexts; int WaterzContext::_nextId = 0; @@ -59,12 +63,18 @@ initialize_with_rag( new RegionMergingType(*regionGraph) ); + + std::shared_ptr> dummy_constraints( + new vector() + ); + WaterzContext* context = WaterzContext::createNew(); context->regionGraph = regionGraph; context->regionMerging = regionMerging; context->scoringFunction = scoringFunction; context->statisticsProvider = statisticsProvider; // context->segmentation = segmentation_data; + context->constraints = dummy_constraints; if (segmentation_data != NULL) { // wrap data (no copy) @@ -91,8 +101,16 @@ initialize( const AffValue* affinity_data, SegID* segmentation_data, const GtID* ground_truth_data, + const SemValue* semantic_data, + const SegID* segconstraint_data, AffValue affThresholdLow, AffValue affThresholdHigh, + AffValue semantic_aff_threshold, + size_t semantic_size_threshold, + AffValue semantic_signal_ratio, + AffValue size_heuristic_aff_threshold, + size_t size_heuristic_small_threshold, + size_t size_heuristic_large_threshold, bool findFragments) { std::size_t num_voxels = width*height*depth; @@ -114,15 +132,10 @@ initialize( counts_t sizes; if (findFragments) { - std::cout << "performing initial watershed segmentation..." << std::endl; - watershed(affinities, affThresholdLow, affThresholdHigh, *segmentation, sizes); - } else { - std::cout << "counting regions and sizes..." << std::endl; - std::size_t maxId = *std::max_element(segmentation_data, segmentation_data + num_voxels); sizes.resize(maxId + 1); for (std::size_t i = 0; i < num_voxels; i++) @@ -158,12 +171,47 @@ initialize( new RegionMergingType(*regionGraph) ); + std::shared_ptr> constraints( + new vector() + ); + + if (semantic_data != NULL) { + std::cout << "getting semantic information..." << std::endl; + constraints->push_back( + new SemanticConstraintProvider( + semantic_data, segmentation_data, num_voxels, + semantic_aff_threshold, + semantic_size_threshold, + semantic_signal_ratio + ) + ); + } + + if (segconstraint_data != NULL) { + std::cout << "getting seg constraint information..." << std::endl; + constraints->push_back( + new SegConstraintProvider( + segconstraint_data, segmentation_data, num_voxels + ) + ); + } + + constraints->push_back( + new SizeHeuristicConstraintProvider( + segmentation_data, num_voxels, + size_heuristic_aff_threshold, + size_heuristic_small_threshold, + size_heuristic_large_threshold + ) + ); + WaterzContext* context = WaterzContext::createNew(); context->regionGraph = regionGraph; context->regionMerging = regionMerging; context->scoringFunction = scoringFunction; context->statisticsProvider = statisticsProvider; context->segmentation = segmentation; + context->constraints = constraints; WaterzState initial_state; initial_state.context = context->id; @@ -199,8 +247,10 @@ mergeUntil( std::size_t merged = context->regionMerging->mergeUntil( *context->scoringFunction, *context->statisticsProvider, + *context->constraints, threshold, - mergeHistoryVisitor); + mergeHistoryVisitor + ); if (merged && context->segmentation) { diff --git a/src/waterz/frontend_agglomerate.h b/src/waterz/frontend_agglomerate.h index e9b7618..04a2417 100644 --- a/src/waterz/frontend_agglomerate.h +++ b/src/waterz/frontend_agglomerate.h @@ -11,10 +11,13 @@ #include "backend/PriorityQueue.hpp" #include "backend/HistogramQuantileProvider.hpp" #include "backend/VectorQuantileProvider.hpp" +#include "backend/ConstraintProvider.hpp" #include "evaluate.hpp" + typedef uint64_t SegID; typedef uint32_t GtID; +typedef uint8_t SemValue; typedef float AffValue; typedef float ScoreValue; typedef RegionGraph RegionGraphType; @@ -101,6 +104,7 @@ class WaterzContext { std::shared_ptr regionMerging; std::shared_ptr scoringFunction; std::shared_ptr statisticsProvider; + std::shared_ptr> constraints; volume_ref_ptr segmentation; volume_const_ref_ptr groundtruth; @@ -149,10 +153,19 @@ WaterzState initialize( size_t depth, const AffValue* affinity_data, SegID* segmentation_data, - const GtID* groundtruth_data = NULL, - AffValue affThresholdLow = 0.0001, - AffValue affThresholdHigh = 0.9999, - bool findFragments = true); + const GtID* groundtruth_data, + const SemValue* semantic_data, + const SegID* segconstraint_data, + AffValue affThresholdLow, + AffValue affThresholdHigh, + AffValue semantic_aff_threshold, + size_t semantic_size_threshold, + AffValue semantic_signal_ratio, + AffValue size_heuristic_aff_threshold, + size_t size_heuristic_small_threshold, + size_t size_heuristic_large_threshold, + bool findFragments + ); WaterzState initialize_with_rag( const std::vector& rag, From 6b43d53e572759c3089f473075a2663510b5a5f9 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 13 Mar 2026 08:42:21 -0700 Subject: [PATCH 04/16] feat: treat semantic label 0 as "no label" in constraint check Co-Authored-By: Claude Opus 4.6 --- src/waterz/backend/SemanticConstraintProvider.hpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/waterz/backend/SemanticConstraintProvider.hpp b/src/waterz/backend/SemanticConstraintProvider.hpp index c8c8495..885e69e 100644 --- a/src/waterz/backend/SemanticConstraintProvider.hpp +++ b/src/waterz/backend/SemanticConstraintProvider.hpp @@ -65,6 +65,8 @@ class SemanticConstraintProvider : public ConstraintProvider { return false; if (max_sem2 < semantic_signal_ratio * total_sem2) return false; + if (max_sem1_label == 0 || max_sem2_label == 0) + return false; if (max_sem1_label == max_sem2_label) return false; return true; From 7fe0b4325afcc773686fc0ea1db3ef4791c965a8 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Mon, 30 Mar 2026 17:25:30 -0700 Subject: [PATCH 05/16] fix: compile frontend_agglomerate.cpp via build_extra_objects witty's source_files parameter only uses files for cache hashing, not compilation. Use build_extra_objects to pre-compile the cpp into an object file that gets linked into the final module. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/waterz/_agglomerate.py | 34 ++++++++++++++++++++++++++-------- 1 file changed, 26 insertions(+), 8 deletions(-) diff --git a/src/waterz/_agglomerate.py b/src/waterz/_agglomerate.py index 944ee83..59d8bc0 100644 --- a/src/waterz/_agglomerate.py +++ b/src/waterz/_agglomerate.py @@ -170,18 +170,36 @@ def agglomerate( (tmp_path / "Queue.h").write_text(queue_src) # compile module + _include_dirs = [ + str(HERE), + tmpdir, + str(HERE / "backend"), + np.get_include(), + "/opt/homebrew/include", + ] + _compile_args = ["-std=c++11", "-w"] + + def _build_frontend(cache_dir: Path) -> list[str]: + import subprocess + + obj_path = cache_dir / "frontend_agglomerate.o" + cpp_path = HERE / "frontend_agglomerate.cpp" + if not obj_path.exists() or obj_path.stat().st_mtime < cpp_path.stat().st_mtime: + cmd = [ + "c++", *_compile_args, + *[f"-I{d}" for d in _include_dirs], + "-fPIC", "-c", str(cpp_path), "-o", str(obj_path), + ] + subprocess.check_call(cmd) + return [str(obj_path)] + module = witty.compile_cython( (HERE / "agglomerate.pyx").read_text(), source_files=[str(HERE / "frontend_agglomerate.cpp")], + build_extra_objects=_build_frontend, extra_link_args=["-std=c++11"], - extra_compile_args=["-std=c++11", "-w"], - include_dirs=[ - str(HERE), - tmpdir, - str(HERE / "backend"), - np.get_include(), - "/opt/homebrew/include", - ], + extra_compile_args=_compile_args, + include_dirs=_include_dirs, language="c++", quiet=True, force_rebuild=force_rebuild, From 97aeb07421336fdb4037effd7f3909fbbb298e35 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 13 Mar 2026 19:55:11 -0700 Subject: [PATCH 06/16] perf: Use swap-and-pop in removeIncEdge for O(1) removal The previous std::vector::erase from the middle was O(n) per call. Since incidence list order doesn't matter, swap the target with the last element and pop_back instead. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/waterz/backend/RegionGraph.hpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/waterz/backend/RegionGraph.hpp b/src/waterz/backend/RegionGraph.hpp index 4dda0ae..c8de5e5 100644 --- a/src/waterz/backend/RegionGraph.hpp +++ b/src/waterz/backend/RegionGraph.hpp @@ -339,7 +339,8 @@ class RegionGraph { auto it = std::find(_incEdges[n].begin(), _incEdges[n].end(), e); assert(it != _incEdges[n].end()); - _incEdges[n].erase(it); + std::swap(*it, _incEdges[n].back()); + _incEdges[n].pop_back(); assert(std::find(_incEdges[n].begin(), _incEdges[n].end(), e) == _incEdges[n].end()); } From 6f39a5a48ba25049397fc0301016a9c47f842251 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 13 Mar 2026 19:55:43 -0700 Subject: [PATCH 07/16] perf: Use unordered_map for _rootPaths for O(1) root lookups extractSegmentation calls getRoot() for every voxel (e.g. 16M for 256^3). Switching from std::map (O(log n)) to std::unordered_map (O(1) amortized) speeds up both merging and segmentation extraction. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/waterz/backend/IterativeRegionMerging.hpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/waterz/backend/IterativeRegionMerging.hpp b/src/waterz/backend/IterativeRegionMerging.hpp index 5665f39..3e40d33 100644 --- a/src/waterz/backend/IterativeRegionMerging.hpp +++ b/src/waterz/backend/IterativeRegionMerging.hpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -367,7 +368,7 @@ class IterativeRegionMerging { // root nodes are not in the map // // paths will be compressed when read - std::map _rootPaths; + std::unordered_map _rootPaths; // current state of merging ScoreType _mergedUntil; From 8488f9a51df9667c5a763509caa49c6427960d7d Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 13 Mar 2026 19:57:01 -0700 Subject: [PATCH 08/16] perf: Use hash-based adjacency map for O(1) findEdge lookups findEdge previously did a linear scan over a node's incident edge list. During mergeRegions, this is called for every neighbor of the absorbed node, leading to O(degree^2) per merge. A per-node unordered_map makes each lookup O(1). Co-Authored-By: Claude Opus 4.6 (1M context) --- src/waterz/backend/RegionGraph.hpp | 53 ++++++++++++++++++------------ 1 file changed, 32 insertions(+), 21 deletions(-) diff --git a/src/waterz/backend/RegionGraph.hpp b/src/waterz/backend/RegionGraph.hpp index c8de5e5..cb2a572 100644 --- a/src/waterz/backend/RegionGraph.hpp +++ b/src/waterz/backend/RegionGraph.hpp @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -164,7 +165,8 @@ class RegionGraph { RegionGraph(ID numNodes = 0) : _numNodes(numNodes), - _incEdges(numNodes) {} + _incEdges(numNodes), + _adjMap(numNodes) {} ID numNodes() const { return _numNodes; } @@ -175,6 +177,7 @@ class RegionGraph { NodeIdType id = _numNodes; _numNodes++; _incEdges.emplace_back(); + _adjMap.emplace_back(); for (RegionGraphNodeMapBase* map : _nodeMaps) map->onNewNode(id); @@ -189,6 +192,8 @@ class RegionGraph { _incEdges[u].push_back(id); _incEdges[v].push_back(id); + _adjMap[u][v] = id; + _adjMap[v][u] = id; for (RegionGraphEdgeMapBase* map : _edgeMaps) map->onNewEdge(id); @@ -198,8 +203,12 @@ class RegionGraph { void removeEdge(EdgeIdType e) { - removeIncEdge(_edges[e].u, e); - removeIncEdge(_edges[e].v, e); + NodeIdType u = _edges[e].u; + NodeIdType v = _edges[e].v; + removeIncEdge(u, e); + removeIncEdge(v, e); + _adjMap[u].erase(v); + _adjMap[v].erase(u); } void moveEdge(EdgeIdType e, NodeIdType u, NodeIdType v) { @@ -273,22 +282,9 @@ class RegionGraph { */ inline EdgeIdType findEdge(NodeIdType u, NodeIdType v) { - return findEdge(u, v, (_incEdges[u].size() < _incEdges[v].size() ? _incEdges[u] : _incEdges[v])); - } - - /** - * Same as findEdge(u, v), but restricted to edges in pool. - */ - inline EdgeIdType findEdge(NodeIdType u, NodeIdType v, const std::vector& pool) { - - NodeIdType min = std::min(u, v); - NodeIdType max = std::max(u, v); - - for (EdgeIdType e : pool) - if (std::min(_edges[e].u, _edges[e].v) == min && - std::max(_edges[e].u, _edges[e].v) == max) - return e; - + auto it = _adjMap[u].find(v); + if (it != _adjMap[u].end()) + return it->second; return NoEdge; } @@ -323,15 +319,27 @@ class RegionGraph { inline void moveEdgeNodeV(EdgeIdType e, NodeIdType v) { - removeIncEdge(_edges[e].v, e); + NodeIdType oldV = _edges[e].v; + NodeIdType otherNode = _edges[e].u; + removeIncEdge(oldV, e); + _adjMap[oldV].erase(otherNode); + _adjMap[otherNode].erase(oldV); _incEdges[v].push_back(e); + _adjMap[v][otherNode] = e; + _adjMap[otherNode][v] = e; _edges[e].v = v; } inline void moveEdgeNodeU(EdgeIdType e, NodeIdType u) { - removeIncEdge(_edges[e].u, e); + NodeIdType oldU = _edges[e].u; + NodeIdType otherNode = _edges[e].v; + removeIncEdge(oldU, e); + _adjMap[oldU].erase(otherNode); + _adjMap[otherNode].erase(oldU); _incEdges[u].push_back(e); + _adjMap[u][otherNode] = e; + _adjMap[otherNode][u] = e; _edges[e].u = u; } @@ -350,6 +358,9 @@ class RegionGraph { std::vector> _incEdges; + // per-node adjacency map: neighbor -> edge id (O(1) findEdge) + std::vector> _adjMap; + std::vector*> _nodeMaps; std::vector*> _edgeMaps; }; From 450a3f1571daebeccf3ef4584c2ba1457f21034d Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 13 Mar 2026 19:58:07 -0700 Subject: [PATCH 09/16] perf: Stream affinities directly into stats provider in get_region_graph Previously, all boundary affinities were collected into a temporary vector>> of size max_segid+1, then iterated again to add edges and affinities. This was a large memory overhead. Now uses the O(1) findEdge (from adjacency map) to create edges on first encounter and stream each affinity directly into the statistics provider, eliminating the temporary storage entirely. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/waterz/backend/region_graph.hpp | 26 ++++++-------------------- 1 file changed, 6 insertions(+), 20 deletions(-) diff --git a/src/waterz/backend/region_graph.hpp b/src/waterz/backend/region_graph.hpp index 2b2553f..2c84b1a 100644 --- a/src/waterz/backend/region_graph.hpp +++ b/src/waterz/backend/region_graph.hpp @@ -4,7 +4,6 @@ #include #include -#include /** * Extract the region graph from a segmentation. Edges are annotated with the @@ -40,10 +39,6 @@ get_region_graph( std::ptrdiff_t ydim = aff.shape()[2]; std::ptrdiff_t xdim = aff.shape()[3]; - // list of affinities between pairs of regions - std::vector>> affinities(max_segid+1); - - EdgeIdType e; std::size_t p[3]; for (p[0] = 0; p[0] < zdim; ++p[0]) for (p[1] = 0; p[1] < ydim; ++p[1]) @@ -61,25 +56,16 @@ get_region_graph( if (id1 != id2) { - auto mm = std::minmax(id1, id2); - affinities[mm.first][mm.second].push_back(aff[d][p[0]][p[1]][p[2]]); + EdgeIdType e = rg.findEdge(id1, id2); + if (e == RegionGraphType::NoEdge) { + e = rg.addEdge(id1, id2); + statisticsProvider.notifyNewEdge(e); + } + statisticsProvider.addAffinity(e, aff[d][p[0]][p[1]][p[2]]); } } } - for (ID id1 = 1; id1 <= max_segid; ++id1) { - for (const auto& p: affinities[id1]) { - - // p.first is ID - // p.second is list of affiliated edges - EdgeIdType e = rg.addEdge(id1, p.first); - statisticsProvider.notifyNewEdge(e); - - for (F affinity : p.second) - statisticsProvider.addAffinity(e, affinity); - } - } - std::cout << "Region graph number of edges: " << rg.edges().size() << std::endl; } From ef80a4409bfe1af541076d121ebc1a586b26d5f3 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 13 Mar 2026 20:12:29 -0700 Subject: [PATCH 10/16] perf: Use raw pointer arithmetic in get_region_graph voxel loop Replace boost multi_array operator[][][] indexing with direct pointer arithmetic for seg and aff data access. Avoids per-access overhead from boost's bounds checking and multi-level indirection across 134M+ voxel iterations. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/waterz/backend/region_graph.hpp | 57 +++++++++++++++++++++-------- 1 file changed, 42 insertions(+), 15 deletions(-) diff --git a/src/waterz/backend/region_graph.hpp b/src/waterz/backend/region_graph.hpp index 2c84b1a..656db6e 100644 --- a/src/waterz/backend/region_graph.hpp +++ b/src/waterz/backend/region_graph.hpp @@ -39,29 +39,56 @@ get_region_graph( std::ptrdiff_t ydim = aff.shape()[2]; std::ptrdiff_t xdim = aff.shape()[3]; - std::size_t p[3]; - for (p[0] = 0; p[0] < zdim; ++p[0]) - for (p[1] = 0; p[1] < ydim; ++p[1]) - for (p[2] = 0; p[2] < xdim; ++p[2]) { - - ID id1 = seg[p[0]][p[1]][p[2]]; - statisticsProvider.addVoxel(id1, p[2], p[1], p[0]); - - for (int d = 0; d < 3; d++) { - - if (p[d] == 0) - continue; - - ID id2 = seg[p[0]-(d==0)][p[1]-(d==1)][p[2]-(d==2)]; + // Use raw pointers for direct memory access instead of boost operator[][] + const ID* seg_data = seg.data(); + const F* aff_data = aff.data(); + const std::size_t slice_size = ydim * xdim; + const std::size_t aff_channel_size = zdim * slice_size; + + for (std::ptrdiff_t z = 0; z < zdim; ++z) + for (std::ptrdiff_t y = 0; y < ydim; ++y) + for (std::ptrdiff_t x = 0; x < xdim; ++x) { + + std::size_t idx = z * slice_size + y * xdim + x; + ID id1 = seg_data[idx]; + statisticsProvider.addVoxel(id1, x, y, z); + + // d=0: z-affinity, neighbor at z-1 + if (z > 0) { + ID id2 = seg_data[idx - slice_size]; + if (id1 != id2) { + EdgeIdType e = rg.findEdge(id1, id2); + if (e == RegionGraphType::NoEdge) { + e = rg.addEdge(id1, id2); + statisticsProvider.notifyNewEdge(e); + } + statisticsProvider.addAffinity(e, aff_data[idx]); + } + } + // d=1: y-affinity, neighbor at y-1 + if (y > 0) { + ID id2 = seg_data[idx - xdim]; if (id1 != id2) { + EdgeIdType e = rg.findEdge(id1, id2); + if (e == RegionGraphType::NoEdge) { + e = rg.addEdge(id1, id2); + statisticsProvider.notifyNewEdge(e); + } + statisticsProvider.addAffinity(e, aff_data[aff_channel_size + idx]); + } + } + // d=2: x-affinity, neighbor at x-1 + if (x > 0) { + ID id2 = seg_data[idx - 1]; + if (id1 != id2) { EdgeIdType e = rg.findEdge(id1, id2); if (e == RegionGraphType::NoEdge) { e = rg.addEdge(id1, id2); statisticsProvider.notifyNewEdge(e); } - statisticsProvider.addAffinity(e, aff[d][p[0]][p[1]][p[2]]); + statisticsProvider.addAffinity(e, aff_data[2 * aff_channel_size + idx]); } } } From d7644ddfb4cc857b759e8cae6e825b250ca2e618 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 13 Mar 2026 20:12:37 -0700 Subject: [PATCH 11/16] perf: Fix double hash lookup in DefaultDict::operator[] and erase Mutable operator[] did find() then operator[] (two lookups). Use emplace() for a single lookup. erase() also did redundant find() before erasing; unordered_map::erase(key) handles missing keys. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/waterz/backend/DefaultDict.hpp | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/waterz/backend/DefaultDict.hpp b/src/waterz/backend/DefaultDict.hpp index aaa45c2..4faa08e 100644 --- a/src/waterz/backend/DefaultDict.hpp +++ b/src/waterz/backend/DefaultDict.hpp @@ -39,10 +39,8 @@ class DefaultDict { DefaultDict(const V& default_val) : default_value(default_val) {} V& operator[](const K& key) { - if (container.find(key) == container.end()) { - container[key] = default_value; - } - return container[key]; + auto result = container.emplace(key, default_value); + return result.first->second; } V operator[](const K& key) const { @@ -62,9 +60,7 @@ class DefaultDict { // } void erase(const K& key) { - if (container.find(key) != container.end()) { - container.erase(key); - } + container.erase(key); } const std::unordered_map& getContainer() const { return container; } From ab60c27679033ddb859e22cc825ed259b135a439 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Fri, 13 Mar 2026 20:12:44 -0700 Subject: [PATCH 12/16] perf: Eliminate redundant voxel counting in SizeHeuristicConstraintProvider initialize() already counts region sizes in a vector. Pass those precomputed sizes to SizeHeuristicConstraintProvider instead of having it re-scan all voxels (134M+ for 512^3 volumes). Co-Authored-By: Claude Opus 4.6 (1M context) --- src/waterz/backend/SizeHeuristicConstraintProvider.hpp | 8 ++++---- src/waterz/frontend_agglomerate.cpp | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/waterz/backend/SizeHeuristicConstraintProvider.hpp b/src/waterz/backend/SizeHeuristicConstraintProvider.hpp index ee79fe3..2fd6011 100644 --- a/src/waterz/backend/SizeHeuristicConstraintProvider.hpp +++ b/src/waterz/backend/SizeHeuristicConstraintProvider.hpp @@ -14,8 +14,7 @@ class SizeHeuristicConstraintProvider : public ConstraintProvider { public: SizeHeuristicConstraintProvider( - const SegType *seg_data, - size_t num_voxels, + const std::vector& precomputed_sizes, float size_heuristic_aff_threshold, size_t size_heuristic_small_threshold, size_t size_heuristic_large_threshold @@ -24,8 +23,9 @@ class SizeHeuristicConstraintProvider : public ConstraintProvider { size_heuristic_small_threshold(size_heuristic_small_threshold), size_heuristic_large_threshold(size_heuristic_large_threshold) { - for (std::size_t i = 0; i < num_voxels; i++) { - _size[seg_data[i]] += 1; + for (std::size_t i = 0; i < precomputed_sizes.size(); i++) { + if (precomputed_sizes[i] > 0) + _size[static_cast(i)] = precomputed_sizes[i]; } } diff --git a/src/waterz/frontend_agglomerate.cpp b/src/waterz/frontend_agglomerate.cpp index e0a236e..46c708a 100644 --- a/src/waterz/frontend_agglomerate.cpp +++ b/src/waterz/frontend_agglomerate.cpp @@ -198,7 +198,7 @@ initialize( constraints->push_back( new SizeHeuristicConstraintProvider( - segmentation_data, num_voxels, + sizes, size_heuristic_aff_threshold, size_heuristic_small_threshold, size_heuristic_large_threshold From 0db59a79d8c75a785b4af9a06227354af538ee9d Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Mon, 30 Mar 2026 17:26:28 -0700 Subject: [PATCH 13/16] feat: add semantic taint constraint for agglomeration Tainted segments (where taint_voxels/total_voxels > threshold for any label in semantic_taint_labels) can only merge with other tainted segments. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/waterz/_agglomerate.py | 4 ++ src/waterz/agglomerate.pyx | 16 ++++++ .../SemanticTaintConstraintProvider.hpp | 57 +++++++++++++++++++ src/waterz/frontend_agglomerate.cpp | 14 +++++ src/waterz/frontend_agglomerate.h | 2 + 5 files changed, 93 insertions(+) create mode 100644 src/waterz/backend/SemanticTaintConstraintProvider.hpp diff --git a/src/waterz/_agglomerate.py b/src/waterz/_agglomerate.py index 59d8bc0..03c98f2 100644 --- a/src/waterz/_agglomerate.py +++ b/src/waterz/_agglomerate.py @@ -32,6 +32,8 @@ def agglomerate( semantic_aff_threshold: float = 0.5, semantic_size_threshold: int = 100_000, semantic_signal_ratio: float = 0.6, + semantic_taint_labels: list[int] = [], + semantic_taint_threshold: float = 0.0, size_heuristic_aff_threshold: float = 1.0, size_heuristic_small_threshold: int = 1_000_000, size_heuristic_large_threshold: int = 10_000_000, @@ -226,6 +228,8 @@ def _build_frontend(cache_dir: Path) -> list[str]: semantic_aff_threshold=semantic_aff_threshold, semantic_size_threshold=semantic_size_threshold, semantic_signal_ratio=semantic_signal_ratio, + semantic_taint_labels=semantic_taint_labels, + semantic_taint_threshold=semantic_taint_threshold, size_heuristic_aff_threshold=size_heuristic_aff_threshold, size_heuristic_small_threshold=size_heuristic_small_threshold, size_heuristic_large_threshold=size_heuristic_large_threshold, diff --git a/src/waterz/agglomerate.pyx b/src/waterz/agglomerate.pyx index 6083d62..218a8a3 100644 --- a/src/waterz/agglomerate.pyx +++ b/src/waterz/agglomerate.pyx @@ -1,6 +1,7 @@ from libcpp.vector cimport vector from libc.stdint cimport uint64_t, uint32_t, uint8_t from libcpp cimport bool +from libcpp.string cimport string import numpy as np cimport numpy as np @@ -41,6 +42,9 @@ def agglomerate( semantic_size_threshold: int, semantic_signal_ratio: float, + semantic_taint_labels: list, + semantic_taint_threshold: float, + size_heuristic_aff_threshold: float, size_heuristic_small_threshold: int, size_heuristic_large_threshold: int, @@ -89,6 +93,8 @@ def agglomerate( semantic_aff_threshold=semantic_aff_threshold, semantic_size_threshold=semantic_size_threshold, semantic_signal_ratio=semantic_signal_ratio, + semantic_taint_labels=semantic_taint_labels, + semantic_taint_threshold=semantic_taint_threshold, size_heuristic_aff_threshold=size_heuristic_aff_threshold, size_heuristic_small_threshold=size_heuristic_small_threshold, size_heuristic_large_threshold=size_heuristic_large_threshold, @@ -142,6 +148,8 @@ def __initialize( semantic_aff_threshold: float, semantic_size_threshold: int, semantic_signal_ratio: float, + semantic_taint_labels: list, + semantic_taint_threshold: float, size_heuristic_aff_threshold: float, size_heuristic_small_threshold: int, size_heuristic_large_threshold: int, @@ -153,6 +161,9 @@ def __initialize( cdef uint32_t* gt_data = NULL cdef uint8_t* semantic_data = NULL cdef uint64_t* segconstraint_data = NULL + cdef vector[uint8_t] taint_labels_vec + for label in semantic_taint_labels: + taint_labels_vec.push_back(label) aff_data = &affs[0,0,0,0] segmentation_data = &segmentation[0,0,0] @@ -189,6 +200,9 @@ def __initialize( semantic_size_threshold=semantic_size_threshold, semantic_signal_ratio=semantic_signal_ratio, + semantic_taint_labels=taint_labels_vec, + semantic_taint_threshold=semantic_taint_threshold, + size_heuristic_aff_threshold=size_heuristic_aff_threshold, size_heuristic_small_threshold=size_heuristic_small_threshold, size_heuristic_large_threshold=size_heuristic_large_threshold, @@ -252,6 +266,8 @@ cdef extern from "frontend_agglomerate.h": float semantic_aff_threshold, uint64_t semantic_size_threshold, float semantic_signal_ratio, + const vector[uint8_t]& semantic_taint_labels, + float semantic_taint_threshold, float size_heuristic_aff_threshold, uint64_t size_heuristic_small_threshold, uint64_t size_heuristic_large_threshold, diff --git a/src/waterz/backend/SemanticTaintConstraintProvider.hpp b/src/waterz/backend/SemanticTaintConstraintProvider.hpp new file mode 100644 index 0000000..2bcabb7 --- /dev/null +++ b/src/waterz/backend/SemanticTaintConstraintProvider.hpp @@ -0,0 +1,57 @@ +#ifndef WATERZ_SEMANTIC_TAINT_CONSTRAINT_PROVIDER_H__ +#define WATERZ_SEMANTIC_TAINT_CONSTRAINT_PROVIDER_H__ + +#include "ConstraintProvider.hpp" +#include "DefaultDict.hpp" + +#include +#include + +template +class SemanticTaintConstraintProvider : public ConstraintProvider { + + typedef typename RegionGraphType::NodeIdType NodeIdType; + +private: + DefaultDict _taint_counts{0}; + DefaultDict _total_counts{0}; + float _threshold; + +public: + SemanticTaintConstraintProvider( + const SemValue* semantic_data, + const SegType* seg_data, + size_t num_voxels, + const std::vector& taint_labels, + float threshold + ) : _threshold(threshold) { + std::unordered_set taint_set(taint_labels.begin(), taint_labels.end()); + for (size_t i = 0; i < num_voxels; i++) { + _total_counts[seg_data[i]] += 1; + if (taint_set.count(semantic_data[i])) { + _taint_counts[seg_data[i]] += 1; + } + } + } + + inline bool notifyNodeMerge(NodeIdType from, NodeIdType to) override { + _taint_counts[to] += _taint_counts[from]; + _taint_counts.erase(from); + _total_counts[to] += _total_counts[from]; + _total_counts.erase(from); + return true; + } + + inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const override { + size_t total_from = _total_counts[from]; + size_t total_to = _total_counts[to]; + bool from_tainted = total_from > 0 && (float)_taint_counts[from] / total_from > _threshold; + bool to_tainted = total_to > 0 && (float)_taint_counts[to] / total_to > _threshold; + // tainted segments can only merge with other tainted segments + if (from_tainted != to_tainted) + return true; + return false; + } +}; + +#endif // WATERZ_SEMANTIC_TAINT_CONSTRAINT_PROVIDER_H__ diff --git a/src/waterz/frontend_agglomerate.cpp b/src/waterz/frontend_agglomerate.cpp index 46c708a..ebb3205 100644 --- a/src/waterz/frontend_agglomerate.cpp +++ b/src/waterz/frontend_agglomerate.cpp @@ -11,6 +11,7 @@ #include "backend/region_graph.hpp" #include "backend/SemanticConstraintProvider.hpp" +#include "backend/SemanticTaintConstraintProvider.hpp" #include "backend/SegConstraintProvider.hpp" #include "backend/SizeHeuristicConstraintProvider.hpp" @@ -108,6 +109,8 @@ initialize( AffValue semantic_aff_threshold, size_t semantic_size_threshold, AffValue semantic_signal_ratio, + const std::vector& semantic_taint_labels, + AffValue semantic_taint_threshold, AffValue size_heuristic_aff_threshold, size_t size_heuristic_small_threshold, size_t size_heuristic_large_threshold, @@ -187,6 +190,17 @@ initialize( ); } + if (semantic_data != NULL && !semantic_taint_labels.empty()) { + std::cout << "getting semantic taint constraint information..." << std::endl; + constraints->push_back( + new SemanticTaintConstraintProvider( + semantic_data, segmentation_data, num_voxels, + semantic_taint_labels, + semantic_taint_threshold + ) + ); + } + if (segconstraint_data != NULL) { std::cout << "getting seg constraint information..." << std::endl; constraints->push_back( diff --git a/src/waterz/frontend_agglomerate.h b/src/waterz/frontend_agglomerate.h index 04a2417..9551a30 100644 --- a/src/waterz/frontend_agglomerate.h +++ b/src/waterz/frontend_agglomerate.h @@ -161,6 +161,8 @@ WaterzState initialize( AffValue semantic_aff_threshold, size_t semantic_size_threshold, AffValue semantic_signal_ratio, + const std::vector& semantic_taint_labels, + AffValue semantic_taint_threshold, AffValue size_heuristic_aff_threshold, size_t size_heuristic_small_threshold, size_t size_heuristic_large_threshold, From 0431973c6bad996c4edff3963753306196fbea3d Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Mon, 30 Mar 2026 17:59:54 -0700 Subject: [PATCH 14/16] perf: pre-compile default agglomeration module in setup.py Compile agglomerate.pyx + frontend_agglomerate.cpp together as a single Extension at install time for the default scoring function (OneMinus MeanAffinity, PriorityQueue). Falls back to witty JIT for custom scoring functions or discretized queues. Single-TU compilation enables cross-file inlining of template-heavy code, reducing agglomeration time from ~30s to ~4s on 512^3 volumes. Co-Authored-By: Claude Opus 4.6 (1M context) --- setup.py | 26 ++++++++- src/waterz/_agglomerate.py | 105 +++++++++++++++++++------------------ 2 files changed, 79 insertions(+), 52 deletions(-) diff --git a/setup.py b/setup.py index a90e566..61a24c1 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,6 @@ import sys +import tempfile +from pathlib import Path import numpy from Cython.Build import cythonize @@ -24,4 +26,26 @@ extra_compile_args=["-std=c++11", "-w"], ) -setup(ext_modules=cythonize([evaluate])) +# Pre-compile default agglomeration module (OneMinus, PriorityQueue) +# so it doesn't need JIT compilation at runtime. +_generated_dir = Path(tempfile.mkdtemp()) +(_generated_dir / "ScoringFunction.h").write_text( + "typedef OneMinus> ScoringFunctionType;" +) +(_generated_dir / "Queue.h").write_text( + "template using QueueType = PriorityQueue;" +) + +agglomerate = Extension( + name="waterz._agglomerate_default", + sources=[ + "src/waterz/agglomerate.pyx", + "src/waterz/frontend_agglomerate.cpp", + ], + include_dirs=include_dirs + [str(_generated_dir)], + language="c++", + extra_link_args=["-std=c++11"], + extra_compile_args=["-std=c++11", "-w", "-O3"], +) + +setup(ext_modules=cythonize([evaluate, agglomerate])) diff --git a/src/waterz/_agglomerate.py b/src/waterz/_agglomerate.py index 03c98f2..9b0f7bd 100644 --- a/src/waterz/_agglomerate.py +++ b/src/waterz/_agglomerate.py @@ -155,57 +155,60 @@ def agglomerate( affs, range(100,10000,100), gt, return_merge_history = True): # ... """ - import witty - - with TemporaryDirectory() as tmpdir: - # supply #include in frontend_agglomerate.h - tmp_path = Path(tmpdir) - scoredef = f"typedef {scoring_function} ScoringFunctionType;" - (tmp_path / "ScoringFunction.h").write_text(scoredef) - - # supply #include in frontend_agglomerate.h - queue_src = "template using QueueType = " + ( - "PriorityQueue;" - if discretize_queue == 0 - else f"BinQueue;" - ) - (tmp_path / "Queue.h").write_text(queue_src) - - # compile module - _include_dirs = [ - str(HERE), - tmpdir, - str(HERE / "backend"), - np.get_include(), - "/opt/homebrew/include", - ] - _compile_args = ["-std=c++11", "-w"] - - def _build_frontend(cache_dir: Path) -> list[str]: - import subprocess - - obj_path = cache_dir / "frontend_agglomerate.o" - cpp_path = HERE / "frontend_agglomerate.cpp" - if not obj_path.exists() or obj_path.stat().st_mtime < cpp_path.stat().st_mtime: - cmd = [ - "c++", *_compile_args, - *[f"-I{d}" for d in _include_dirs], - "-fPIC", "-c", str(cpp_path), "-o", str(obj_path), - ] - subprocess.check_call(cmd) - return [str(obj_path)] - - module = witty.compile_cython( - (HERE / "agglomerate.pyx").read_text(), - source_files=[str(HERE / "frontend_agglomerate.cpp")], - build_extra_objects=_build_frontend, - extra_link_args=["-std=c++11"], - extra_compile_args=_compile_args, - include_dirs=_include_dirs, - language="c++", - quiet=True, - force_rebuild=force_rebuild, - ) + _DEFAULT_SCORING = "OneMinus>" + _use_precompiled = (scoring_function == _DEFAULT_SCORING and discretize_queue == 0) + + if _use_precompiled: + from waterz import _agglomerate_default as module + else: + import witty + + with TemporaryDirectory() as tmpdir: + tmp_path = Path(tmpdir) + scoredef = f"typedef {scoring_function} ScoringFunctionType;" + (tmp_path / "ScoringFunction.h").write_text(scoredef) + + queue_src = "template using QueueType = " + ( + "PriorityQueue;" + if discretize_queue == 0 + else f"BinQueue;" + ) + (tmp_path / "Queue.h").write_text(queue_src) + + _include_dirs = [ + str(HERE), + tmpdir, + str(HERE / "backend"), + np.get_include(), + "/opt/homebrew/include", + ] + _compile_args = ["-std=c++11", "-w"] + + def _build_frontend(cache_dir: Path) -> list[str]: + import subprocess + + obj_path = cache_dir / "frontend_agglomerate.o" + cpp_path = HERE / "frontend_agglomerate.cpp" + if not obj_path.exists() or obj_path.stat().st_mtime < cpp_path.stat().st_mtime: + cmd = [ + "c++", *_compile_args, + *[f"-I{d}" for d in _include_dirs], + "-fPIC", "-c", str(cpp_path), "-o", str(obj_path), + ] + subprocess.check_call(cmd) + return [str(obj_path)] + + module = witty.compile_cython( + (HERE / "agglomerate.pyx").read_text(), + source_files=[str(HERE / "frontend_agglomerate.cpp")], + build_extra_objects=_build_frontend, + extra_link_args=["-std=c++11"], + extra_compile_args=_compile_args, + include_dirs=_include_dirs, + language="c++", + quiet=True, + force_rebuild=force_rebuild, + ) # call compiled function if input_rag is not None or input_rag_metadata is not None: From 907631c34c8386604501597dd5193fe3dca56d1b Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Mon, 30 Mar 2026 18:50:06 -0700 Subject: [PATCH 15/16] perf: constraint providers return false from notifyNodeMerge Constraint providers update node-level state for isConstrained() checks at pop time, not for edge score computation. Returning true caused all incident edges to be marked stale and re-scored unnecessarily. Safe because the default scoring function (OneMinus) only uses MeanAffinityProvider (which already returns false), not RegionSizeProvider. Co-Authored-By: Claude Opus 4.6 (1M context) --- src/waterz/backend/SegConstraintProvider.hpp | 2 +- src/waterz/backend/SemanticConstraintProvider.hpp | 2 +- src/waterz/backend/SemanticTaintConstraintProvider.hpp | 2 +- src/waterz/backend/SizeHeuristicConstraintProvider.hpp | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/waterz/backend/SegConstraintProvider.hpp b/src/waterz/backend/SegConstraintProvider.hpp index d28a724..f9902f1 100644 --- a/src/waterz/backend/SegConstraintProvider.hpp +++ b/src/waterz/backend/SegConstraintProvider.hpp @@ -39,7 +39,7 @@ class SegConstraintProvider : public ConstraintProvider { _constraint[to] = _constraint.at(from); } _constraint.erase(from); - return true; + return false; } inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const { diff --git a/src/waterz/backend/SemanticConstraintProvider.hpp b/src/waterz/backend/SemanticConstraintProvider.hpp index 885e69e..0c32d30 100644 --- a/src/waterz/backend/SemanticConstraintProvider.hpp +++ b/src/waterz/backend/SemanticConstraintProvider.hpp @@ -42,7 +42,7 @@ class SemanticConstraintProvider : public ConstraintProvider { _semantic[to][k.first] += k.second; } _semantic.erase(from); - return true; + return false; } inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const { diff --git a/src/waterz/backend/SemanticTaintConstraintProvider.hpp b/src/waterz/backend/SemanticTaintConstraintProvider.hpp index 2bcabb7..2fbe38c 100644 --- a/src/waterz/backend/SemanticTaintConstraintProvider.hpp +++ b/src/waterz/backend/SemanticTaintConstraintProvider.hpp @@ -39,7 +39,7 @@ class SemanticTaintConstraintProvider : public ConstraintProvider { _taint_counts.erase(from); _total_counts[to] += _total_counts[from]; _total_counts.erase(from); - return true; + return false; } inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const override { diff --git a/src/waterz/backend/SizeHeuristicConstraintProvider.hpp b/src/waterz/backend/SizeHeuristicConstraintProvider.hpp index 2fd6011..9092f2f 100644 --- a/src/waterz/backend/SizeHeuristicConstraintProvider.hpp +++ b/src/waterz/backend/SizeHeuristicConstraintProvider.hpp @@ -32,7 +32,7 @@ class SizeHeuristicConstraintProvider : public ConstraintProvider { inline bool notifyNodeMerge(NodeIdType from, NodeIdType to) override { _size[to] += _size[from]; _size.erase(from); - return true; // statistics changed + return false; } inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const override { From ca7d423129aa2c8ec17be917dfbfc346e0ee5e58 Mon Sep 17 00:00:00 2001 From: Dodam Ih Date: Mon, 30 Mar 2026 19:05:09 -0700 Subject: [PATCH 16/16] perf: avoid vector copy and redundant graph ops in mergeRegions - takeIncEdges() swaps b's incident list instead of copying it - reassignEdge() directly moves an edge from oldNode to newNode without the generic 5-case moveEdge branching - removeEdgeSkipNode() skips touching the absorbed node's incEdges since the caller already took ownership via takeIncEdges Co-Authored-By: Claude Opus 4.6 (1M context) --- src/waterz/backend/IterativeRegionMerging.hpp | 22 ++--- src/waterz/backend/RegionGraph.hpp | 86 +++++++++++++++++++ src/waterz/backend/SegConstraintProvider.hpp | 2 +- .../backend/SemanticConstraintProvider.hpp | 2 +- .../SemanticTaintConstraintProvider.hpp | 2 +- .../SizeHeuristicConstraintProvider.hpp | 2 +- 6 files changed, 101 insertions(+), 15 deletions(-) diff --git a/src/waterz/backend/IterativeRegionMerging.hpp b/src/waterz/backend/IterativeRegionMerging.hpp index 3e40d33..513a0e3 100644 --- a/src/waterz/backend/IterativeRegionMerging.hpp +++ b/src/waterz/backend/IterativeRegionMerging.hpp @@ -234,7 +234,10 @@ class IterativeRegionMerging { } // ...and update incident edges of b - std::vector neighborEdges = _regionGraph.incEdges(b); + // Take b's incident list via swap to avoid copying the vector. + // b's list is cleared, so use reassignEdge/removeEdgeSkipNode + // which don't touch b's incEdges. + std::vector neighborEdges = _regionGraph.takeIncEdges(b); for (EdgeIdType neighborEdge : neighborEdges) { if (neighborEdge == e) @@ -253,8 +256,7 @@ class IterativeRegionMerging { // We encountered an exclusive neighbor of b. - _regionGraph.moveEdge(neighborEdge, a, neighbor); - assert(_regionGraph.findEdge(a, neighbor) == neighborEdge); + _regionGraph.reassignEdge(neighborEdge, b, a); if (nodeStatisticsChanged) _stale[neighborEdge] = true; @@ -267,32 +269,30 @@ class IterativeRegionMerging { // * mark the cheaper one as stale (if it isn't already) // * delete the more expensive one // - // This ensures that the stale edge bubbles up early enough - // to consider it's real score (which is assumed to be + // This ensures that the stale edge bubbles up early enough + // to consider it's real score (which is assumed to be // larger than the minium of the two original scores). if (_edgeScores[neighborEdge] > _edgeScores[aNeighborEdge]) { - // We got lucky, we can reuse the edge that is attached to a + // We got lucky, we can reuse the edge that is attached to a // already bool edgeStatisticChanged = statisticsProvider.notifyEdgeMerge(neighborEdge, aNeighborEdge); - _regionGraph.removeEdge(neighborEdge); + _regionGraph.removeEdgeSkipNode(neighborEdge, b); _deleted[neighborEdge] = true; if (edgeStatisticChanged) _stale[aNeighborEdge] = true; } else { - // Bummer. The new edge should be the one pointing from + // Bummer. The new edge should be the one pointing from // a to neighbor. bool edgeStatisticChanged = statisticsProvider.notifyEdgeMerge(aNeighborEdge, neighborEdge); - _regionGraph.removeEdge(aNeighborEdge); - _regionGraph.moveEdge(neighborEdge, a, neighbor); - assert(_regionGraph.findEdge(a, neighbor) == neighborEdge); + _regionGraph.replaceEdge(aNeighborEdge, neighborEdge, a, neighbor, b); if (edgeStatisticChanged) _stale[neighborEdge] = true; diff --git a/src/waterz/backend/RegionGraph.hpp b/src/waterz/backend/RegionGraph.hpp index cb2a572..ac7a540 100644 --- a/src/waterz/backend/RegionGraph.hpp +++ b/src/waterz/backend/RegionGraph.hpp @@ -272,11 +272,97 @@ class RegionGraph { inline const std::vector& incEdges(ID node) const { return _incEdges[node]; } + inline std::vector takeIncEdges(ID node) { + std::vector result; + result.swap(_incEdges[node]); + return result; + } + inline NodeIdType getOpposite(NodeIdType n, EdgeIdType e) const { return (_edges[e].u == n ? _edges[e].v : _edges[e].u); } + /** + * Fast edge reassignment: move edge e from oldNode to newNode. + * Caller must ensure one endpoint of e is oldNode. + * Does NOT touch oldNode's incEdges (caller handles that). + */ + void reassignEdge(EdgeIdType e, NodeIdType oldNode, NodeIdType newNode) { + + NodeIdType other = getOpposite(oldNode, e); + + // Update adjacency maps + _adjMap[oldNode].erase(other); + _adjMap[other].erase(oldNode); + _adjMap[newNode][other] = e; + _adjMap[other][newNode] = e; + + // Update edge endpoints + if (_edges[e].u == oldNode) + _edges[e].u = newNode; + else + _edges[e].v = oldNode == _edges[e].v ? newNode : _edges[e].v; + + // Ensure u < v + if (_edges[e].u > _edges[e].v) + std::swap(_edges[e].u, _edges[e].v); + + // Add to newNode's incident list (don't remove from oldNode - caller owns that) + _incEdges[newNode].push_back(e); + } + + /** + * Remove edge from graph but don't touch incEdges of the given skipNode. + * Used when the caller already owns/cleared that node's incident list. + */ + void removeEdgeSkipNode(EdgeIdType e, NodeIdType skipNode) { + + NodeIdType u = _edges[e].u; + NodeIdType v = _edges[e].v; + NodeIdType other = (u == skipNode) ? v : u; + removeIncEdge(other, e); + _adjMap[u].erase(v); + _adjMap[v].erase(u); + } + + /** + * Replace oldEdge (between survivor and neighbor) with newEdge + * (being reassigned from oldNode to survivor). Fuses removeEdge + + * reassignEdge to avoid redundant incEdge and adjMap operations. + * Caller must have already taken oldNode's incEdges. + */ + void replaceEdge(EdgeIdType oldEdge, EdgeIdType newEdge, + NodeIdType survivor, NodeIdType neighbor, NodeIdType oldNode) { + + // Remove oldEdge from survivor's and neighbor's incEdges + removeIncEdge(survivor, oldEdge); + removeIncEdge(neighbor, oldEdge); + + // Remove newEdge from neighbor's incEdges (oldNode's already taken) + removeIncEdge(neighbor, newEdge); + + // Update newEdge endpoints: oldNode -> survivor + if (_edges[newEdge].u == oldNode) + _edges[newEdge].u = survivor; + else + _edges[newEdge].v = survivor; + if (_edges[newEdge].u > _edges[newEdge].v) + std::swap(_edges[newEdge].u, _edges[newEdge].v); + + // Add newEdge to survivor's and neighbor's incEdges + _incEdges[survivor].push_back(newEdge); + _incEdges[neighbor].push_back(newEdge); + + // Update adjMaps: survivor↔neighbor now points to newEdge + _adjMap[survivor][neighbor] = newEdge; + _adjMap[neighbor][survivor] = newEdge; + + // Clean up oldNode's adjMap entries + _adjMap[oldNode].erase(neighbor); + _adjMap[neighbor].erase(oldNode); + } + /** * Find the edge connecting u and v. Returns NoEdge, if there is none. */ diff --git a/src/waterz/backend/SegConstraintProvider.hpp b/src/waterz/backend/SegConstraintProvider.hpp index f9902f1..d28a724 100644 --- a/src/waterz/backend/SegConstraintProvider.hpp +++ b/src/waterz/backend/SegConstraintProvider.hpp @@ -39,7 +39,7 @@ class SegConstraintProvider : public ConstraintProvider { _constraint[to] = _constraint.at(from); } _constraint.erase(from); - return false; + return true; } inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const { diff --git a/src/waterz/backend/SemanticConstraintProvider.hpp b/src/waterz/backend/SemanticConstraintProvider.hpp index 0c32d30..885e69e 100644 --- a/src/waterz/backend/SemanticConstraintProvider.hpp +++ b/src/waterz/backend/SemanticConstraintProvider.hpp @@ -42,7 +42,7 @@ class SemanticConstraintProvider : public ConstraintProvider { _semantic[to][k.first] += k.second; } _semantic.erase(from); - return false; + return true; } inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const { diff --git a/src/waterz/backend/SemanticTaintConstraintProvider.hpp b/src/waterz/backend/SemanticTaintConstraintProvider.hpp index 2fbe38c..2bcabb7 100644 --- a/src/waterz/backend/SemanticTaintConstraintProvider.hpp +++ b/src/waterz/backend/SemanticTaintConstraintProvider.hpp @@ -39,7 +39,7 @@ class SemanticTaintConstraintProvider : public ConstraintProvider { _taint_counts.erase(from); _total_counts[to] += _total_counts[from]; _total_counts.erase(from); - return false; + return true; } inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const override { diff --git a/src/waterz/backend/SizeHeuristicConstraintProvider.hpp b/src/waterz/backend/SizeHeuristicConstraintProvider.hpp index 9092f2f..2fd6011 100644 --- a/src/waterz/backend/SizeHeuristicConstraintProvider.hpp +++ b/src/waterz/backend/SizeHeuristicConstraintProvider.hpp @@ -32,7 +32,7 @@ class SizeHeuristicConstraintProvider : public ConstraintProvider { inline bool notifyNodeMerge(NodeIdType from, NodeIdType to) override { _size[to] += _size[from]; _size.erase(from); - return false; + return true; // statistics changed } inline bool isConstrained(NodeIdType from, NodeIdType to, float score) const override {