Semantic / Size / Seg constraints + Optimisations

.gitignore

@@ -13,3 +13,4 @@ src/waterz/**/*.pyc
 src/waterz/**/*.c
 src/waterz/**/*.so
 src/waterz/evaluate.cpp
+/_agglomerate_gen/

setup.py

@@ -1,4 +1,6 @@
 import sys
+import tempfile
+from pathlib import Path
 
 import numpy
 from Cython.Build import cythonize
@@ -24,4 +26,42 @@
     extra_compile_args=["-std=c++11", "-w"],
 )
 
-setup(ext_modules=cythonize([evaluate]))
+# Pre-compile the agglomeration modules for the (scoring, queue) combinations
+# used in production so they never need runtime JIT (witty) compilation on
+# workers -- which avoids the concurrent-compile races that can corrupt the
+# witty cache. Each variant bakes its template parameters in via generated
+# headers and gets a uniquely-named .pyx copy so cythonize emits distinct
+# modules (the same .pyx compiled twice would collide on its PyInit symbol).
+_MEAN_SCORING = "typedef OneMinus<MeanAffinity<RegionGraphType, ScoreValue>> ScoringFunctionType;"
+_PRIORITY_QUEUE = "template<typename T, typename S> using QueueType = PriorityQueue<T, S>;"
+_BIN_QUEUE_256 = "template<typename T, typename S> using QueueType = BinQueue<T, S, 256>;"
+
+_agglomerate_pyx_src = Path("src/waterz/agglomerate.pyx").read_text()
+# Per-variant .pyx copies must live at a relative path (setuptools rejects
+# absolute paths in Extension.sources); generated headers may stay absolute.
+_gen_pyx_dir = Path("_agglomerate_gen")
+_gen_pyx_dir.mkdir(exist_ok=True)
+
+
+def _agglomerate_ext(module_name, queue_decl):
+    headers = Path(tempfile.mkdtemp())
+    (headers / "ScoringFunction.h").write_text(_MEAN_SCORING)
+    (headers / "Queue.h").write_text(queue_decl)
+    pyx_path = _gen_pyx_dir / f"{module_name}.pyx"
+    pyx_path.write_text(_agglomerate_pyx_src)
+    return Extension(
+        name=f"waterz.{module_name}",
+        sources=[str(pyx_path), "src/waterz/frontend_agglomerate.cpp"],
+        include_dirs=include_dirs + [str(headers)],
+        language="c++",
+        extra_link_args=["-std=c++11"],
+        extra_compile_args=["-std=c++11", "-w", "-O3"],
+    )
+
+
+# (MeanAffinity, PriorityQueue)  -> discretize_queue == 0
+agglomerate_default = _agglomerate_ext("_agglomerate_default", _PRIORITY_QUEUE)
+# (MeanAffinity, BinQueue<256>) -> discretize_queue == 256 (production agglomeration)
+agglomerate_bin256 = _agglomerate_ext("_agglomerate_mean_bin256", _BIN_QUEUE_256)
+
+setup(ext_modules=cythonize([evaluate, agglomerate_default, agglomerate_bin256]))

src/waterz/_agglomerate.py

@@ -15,15 +15,28 @@
 
 
 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,
+    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,
     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<MeanAffinity<RegionGraphType, ScoreValue>>",
+    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,
     discretize_queue: int = 0,
     force_rebuild: bool = False,
 ) -> Iterator[tuple | NDArray[np.uint64]]:
@@ -142,48 +155,99 @@ def agglomerate(
             affs, range(100,10000,100), gt, return_merge_history = True):
             # ...
     """
-    import witty
-
-    with TemporaryDirectory() as tmpdir:
-        # supply #include <ScoringFunction.h> in frontend_agglomerate.h
-        tmp_path = Path(tmpdir)
-        scoredef = f"typedef {scoring_function} ScoringFunctionType;"
-        (tmp_path / "ScoringFunction.h").write_text(scoredef)
-
-        # supply #include <Queue.h> in frontend_agglomerate.h
-        queue_src = "template<typename T, typename S> using QueueType = " + (
-            "PriorityQueue<T, S>;"
-            if discretize_queue == 0
-            else f"BinQueue<T, S, {discretize_queue}>;"
-        )
-        (tmp_path / "Queue.h").write_text(queue_src)
-
-        # compile module
-        module = witty.compile_cython(
-            (HERE / "agglomerate.pyx").read_text(),
-            source_files=[str(HERE / "frontend_agglomerate.cpp")],
-            extra_link_args=["-std=c++11"],
-            extra_compile_args=["-std=c++11", "-w"],
-            include_dirs=[
+    _DEFAULT_SCORING = "OneMinus<MeanAffinity<RegionGraphType, ScoreValue>>"
+    # Pre-compiled (shipped) modules keyed by discretize_queue, for the default
+    # mean scoring. Importing a ready .so avoids runtime JIT (witty) compilation
+    # -- and the concurrent-compile cache races that can corrupt it on workers.
+    # Other discretize_queue values / scoring functions fall back to JIT below.
+    _PRECOMPILED_BY_QUEUE = {0: "_agglomerate_default", 256: "_agglomerate_mean_bin256"}
+    _precompiled_name = (
+        _PRECOMPILED_BY_QUEUE.get(discretize_queue)
+        if scoring_function == _DEFAULT_SCORING
+        else None
+    )
+
+    if _precompiled_name is not None:
+        import importlib
+
+        module = importlib.import_module(f"waterz.{_precompiled_name}")
+    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<typename T, typename S> using QueueType = " + (
+                "PriorityQueue<T, S>;"
+                if discretize_queue == 0
+                else f"BinQueue<T, S, {discretize_queue}>;"
+            )
+            (tmp_path / "Queue.h").write_text(queue_src)
+
+            _include_dirs = [
                 str(HERE),
                 tmpdir,
                 str(HERE / "backend"),
                 np.get_include(),
                 "/opt/homebrew/include",
-            ],
-            language="c++",
-            quiet=True,
-            force_rebuild=force_rebuild,
-        )
+            ]
+            _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:
+        return module.agglomerate_rag(
+            rag=input_rag,
+            rag_metadata=input_rag_metadata,
+            thresholds=thresholds,
+            fragments=fragments,
+        )
+
     return module.agglomerate(
-        affs,
-        thresholds,
-        gt,
-        fragments,
-        aff_threshold_low,
-        aff_threshold_high,
-        return_merge_history,
-        return_region_graph,
+        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,
+        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,
+        return_merge_history=return_merge_history,
+        return_region_graph=return_region_graph,
+        return_region_graph_metadata=return_region_graph_metadata,
     )