Skip to content

Semantic / Size / Seg constraints + Optimisations - #24

Open
dodamih wants to merge 21 commits into
funkey:masterfrom
ZettaAI:master
Open

Semantic / Size / Seg constraints + Optimisations#24
dodamih wants to merge 21 commits into
funkey:masterfrom
ZettaAI:master

Conversation

@dodamih

@dodamih dodamih commented Mar 31, 2026

Copy link
Copy Markdown

Performance optimizations for waterz agglomeration, plus new features (constraints, semantic taint).

Features

  • Agglomeration constraints: semantic, size heuristic, and seg constraint providers that block merges based on configurable criteria
  • RAG initialization: agglomerate from a precomputed region graph without re-extracting from voxels
  • Region graph metadata: return edge contact sizes alongside the region graph
  • Semantic taint constraint: tainted segments (by label ratio) can only merge with other tainted segments

Performance

  • Pre-compile default agglomeration module in setup.py instead of JIT-compiling via witty at runtime. Single-TU compilation enables cross-file inlining of template-heavy code. (30s → 4.3s)
  • Fix witty build: witty's source_files parameter only hashes files for caching, doesn't compile them. Use build_extra_objects to actually compile frontend_agglomerate.cpp.
  • O(1) findEdge: replace linear scan over incident edge list with per-node unordered_map adjacency map
  • O(1) root lookups: unordered_map instead of std::map for merge-tree root paths
  • O(1) removeIncEdge: swap-and-pop instead of find + erase on incident edge vector
  • Stream affinities directly into statistics provider during RG extraction, eliminating temporary vector<map<ID, vector>> storage
  • Raw pointer arithmetic in RG extraction voxel loop instead of boost multi_array indexing
  • Avoid redundant voxel counting: pass precomputed sizes to SizeHeuristicConstraintProvider
  • Fix DefaultDict double hash lookup in operator[] and erase
  • Constraint providers return false from notifyNodeMerge: constraints don't affect edge scores, only isConstrained() at pop time. -Returning true caused unnecessary stale-edge churn.
  • takeIncEdges + reassignEdge + replaceEdge: avoid vector copy and redundant graph operations during merge

Benchmarks on a 1024x1024x512 dataset:

RG extraction : 52.24s -> 17.73s
Agglomeration : 70.99s -> 28.15s

trivoldus28 and others added 21 commits March 30, 2026 17:40
@trivoldus28 @claude @dodamih
Add getting mean aff metadata
c60123d 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@trivoldus28 @claude @dodamih
Add initialization using rag
38c0d61 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@dodamih @trivoldus28
feat: add semantic/size/seg constraints
5439b20 
Co-Authored-By: trivoldus28 <tri@zetta.ai>
@dodamih @claude
feat: treat semantic label 0 as "no label" in constraint check
6b43d53 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@dodamih @claude
fix: compile frontend_agglomerate.cpp via build_extra_objects
7fe0b43 
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) <noreply@anthropic.com>
@dodamih @claude
perf: Use swap-and-pop in removeIncEdge for O(1) removal
97aeb07 
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) <noreply@anthropic.com>
@dodamih @claude
perf: Use unordered_map for _rootPaths for O(1) root lookups
6f39a5a 
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) <noreply@anthropic.com>
@dodamih @claude
perf: Use hash-based adjacency map for O(1) findEdge lookups
8488f9a 
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<neighbor, edgeId> makes each lookup O(1).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@dodamih @claude
perf: Stream affinities directly into stats provider in get_region_graph
450a3f1 
Previously, all boundary affinities were collected into a temporary
vector<map<ID, vector<F>>> 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) <noreply@anthropic.com>
@dodamih @claude
perf: Use raw pointer arithmetic in get_region_graph voxel loop
ef80a44 
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) <noreply@anthropic.com>
@dodamih @claude
perf: Fix double hash lookup in DefaultDict::operator[] and erase
d7644dd 
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) <noreply@anthropic.com>
@dodamih @claude
perf: Eliminate redundant voxel counting in SizeHeuristicConstraintPr…
ab60c27 
…ovider

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) <noreply@anthropic.com>
@dodamih @claude
feat: add semantic taint constraint for agglomeration
0db59a7 
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) <noreply@anthropic.com>
@dodamih @claude
perf: pre-compile default agglomeration module in setup.py
0431973 
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) <noreply@anthropic.com>
@dodamih @claude
perf: constraint providers return false from notifyNodeMerge
907631c 
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<MeanAffinity>) only uses MeanAffinityProvider (which
already returns false), not RegionSizeProvider.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@dodamih @claude
perf: avoid vector copy and redundant graph ops in mergeRegions
ca7d423 
- 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) <noreply@anthropic.com>
@dodamih
Merge pull request #2 from ZettaAI/dodam/agg_clean
90898fb 
Update to master's main, add constraints, optimise
Nico Kemnitz @claude @dodamih
fix: include <cstdint> where fixed-width int types are used
cfab96f 
Headers that use uint64_t and friends relied on transitive includes
that newer libstdc++ (GCC 13+) no longer provides, causing build
failures such as "'uint64_t' has not been declared" and the cascade of
abstract-class / override mismatch errors that follow from the implicit
int fallback. Add explicit <cstdint> includes to each affected header.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Nico Kemnitz @claude @dodamih
fix: keep maxId 64-bit in initialize_with_rag
a891b34 
edge.u/edge.v are SegID (uint64_t) while maxId was std::size_t. On Linux
these are the same underlying type so std::max deduces fine, but on
macOS uint64_t is unsigned long long and size_t is unsigned long, so
template argument deduction fails ("deduced conflicting types"). Declare
maxId as SegID so both std::max operands share a 64-bit type on every
platform, avoiding both the deduction failure and any narrowing of seg
ids toward a possibly-32-bit size_t.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@dodamih @claude
fix(agglomeration): tolerate float rounding in stale-edge rescore
77eccde 
MeanAffinity edge scores are recomputed from running/area-weighted means
in float32, so a stale edge's rescore can dip ~1 ULP (5.96e-08) below its
queued score and abort on `assert(newScore >= score)`. This surfaces under
heavy over-fragmentation (watershed size_threshold=0): high-affinity
(mean~0.99), tiny-contact edges where the running-mean rounding is most
jittery right where 1-mean is tiny. Observed 23 such 1-ULP violations out
of ~48M edges -- one is enough to abort the process.

Tolerate sub-epsilon rounding noise; a real non-monotonicity exceeds 1e-5
by orders of magnitude (and the discretize_queue bucket width ~1/256).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@dodamih @claude
perf(agglomeration): ship discretize_queue=256 precompiled, no runtim…
a52bafa 
…e JIT

Pre-compile the BinQueue<256> mean-affinity variant in setup.py alongside the
PriorityQueue default, and dispatch discretize_queue in {0, 256} to the shipped
modules. The production agglomeration path (discretize_queue=256) no longer
JIT-compiles at runtime, eliminating the concurrent-compile races that corrupt
the witty cache (ImportError: invalid ELF header) on multi-process workers.
Uncommon combinations still fall back to JIT.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@dodamih @trivoldus28