Add Python wrapper for watershed

README.md

@@ -7,3 +7,22 @@ mkdir build && cd build
 cmake ..
 make
 ```
+
+# Python interface
+
+Install from local directory
+```sh
+pip install .
+```
+
+Run watershed
+```python
+# affs is a [width, height, depth, 3] numpy array of float32
+affs = ...
+segmentations = abiss.watershed(
+        affs=affs,
+        aff_threshold_low=0.01,
+        aff_threshold_high=0.09,
+        size_threshold=200,
+    )
+```

pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools", "numpy", "wheel", "cython"]
+build-backend = "setuptools.build_meta"

python/abiss.pyx

@@ -0,0 +1,74 @@
+from libcpp.vector cimport vector
+from libc.stdint cimport uint64_t, uint32_t
+import numpy as np
+cimport numpy as np
+
+
+def watershed(
+        affs,  # assuming XYZC
+        aff_threshold_low=0.01,
+        aff_threshold_high=0.99,
+        size_threshold=50,
+        # agglomeration_threshold=0.0,
+        ):
+
+    if not affs.flags['F_CONTIGUOUS']:
+        print("Creating memory-contiguous affinity array (avoid this by passing F_CONTIGUOUS arrays)")
+        affs = np.asfortranarray(affs)
+
+    volume_shape = (affs.shape[0], affs.shape[1], affs.shape[2])  # XYZ
+    segmentation = np.asfortranarray(np.zeros(volume_shape, dtype=np.uint64))
+
+    # _watershed(affs, segmentation, aff_threshold_low, aff_threshold_high, size_threshold, agglomeration_threshold)
+    _watershed(affs, segmentation, aff_threshold_low, aff_threshold_high, size_threshold)
+    return segmentation  # XYZ
+
+
+def _watershed(
+        np.ndarray[np.float32_t, ndim=4] affs,
+        np.ndarray[uint64_t, ndim=3]     segmentation,
+        aff_threshold_low=0.01,
+        aff_threshold_high=0.99,
+        size_threshold=50,
+        # agglomeration_threshold=0.0,
+        ):
+
+    cdef float*    aff_data
+    cdef uint64_t* segmentation_data
+
+    aff_data = &affs[0,0,0,0]
+    segmentation_data = &segmentation[0,0,0]
+    width, height, depth = affs.shape[0], affs.shape[1], affs.shape[2]
+
+    sv_sizes = run_watershed(
+        width, height, depth,
+        aff_data,
+        segmentation_data,
+        size_threshold,
+        aff_threshold_low,
+        aff_threshold_high)
+
+    # if agglomeration_threshold > 0.0:
+    #     rg = extract_region_graph(
+    #         width, height, depth,
+    #         aff_data,
+    #         segmentation_data)
+    #     remaps = run_agglomeration(rg, sv_sizes, agglomeration_threshold)
+    #     run_update_segmentation(
+    #         width, height, depth,
+    #         segmentation_data,
+    #         remaps)
+
+    return
+
+cdef extern from "abiss_wrapper.h":
+    vector[size_t] run_watershed(
+            size_t          width,
+            size_t          height,
+            size_t          depth,
+            const float*    affinity_data,
+            uint64_t*       segmentation_data,
+            size_t          size_threshold,
+            float           affThresholdLow,
+            float           affThresholdHigh,
+            );

python/abiss_wrapper.cpp

@@ -0,0 +1,47 @@
+
+#include "ws/basic_watershed.hpp"
+#include "ws/region_graph.hpp"
+#include "ws/agglomeration.hpp"
+#include "ws/utils.hpp"
+#include "abiss_wrapper.h"
+
+std::vector<std::size_t> run_watershed(
+        std::size_t     width,
+        std::size_t     height,
+        std::size_t     depth,
+        const aff_t*    affinity_data,
+        seg_t*          seg_out,
+        std::size_t     size_threshold,
+        aff_t           aff_threshold_low,
+        aff_t           aff_threshold_high) {
+
+    std::shared_ptr<affinity_graph<aff_t>> affs_ptr(
+        new affinity_graph<aff_t> (
+            const_cast<aff_t*>(affinity_data),  // remove const for compatibility
+            boost::extents[width][height][depth][3],
+            boost::fortran_storage_order()
+        ),
+        // remove deleter because the obj is actually owned by Python
+        [](affinity_graph<aff_t>*){}
+    );
+
+    // Run basic watershed
+    volume_ptr<seg_t> ws_out;
+    std::vector<std::size_t> counts;
+    std::array<bool, 6> boundary_flags {true, true, true, true, true, true};
+    std::tie(ws_out, counts) = watershed<seg_t>(affs_ptr, aff_threshold_low, aff_threshold_high, boundary_flags);
+
+    // Run size thresholding
+    if (size_threshold > 0) {
+        auto rg = get_region_graph<seg_t>(
+            affs_ptr, ws_out, counts.size()-1, aff_threshold_low, boundary_flags);
+        merge_segments(ws_out, rg, counts, std::make_pair(size_threshold, aff_threshold_low), size_threshold);
+    }
+
+    // Copy output to a numpy allocated array
+    for (uint64_t i = 0; i < depth*height*width; ++i)
+        seg_out[i] = ws_out->data()[i];
+
+    free_container(ws_out);
+    return counts;
+}

python/abiss_wrapper.h

@@ -0,0 +1,17 @@
+#ifndef ABISS_WRAPPER_H
+#define ABISS_WRAPPER_H
+
+using seg_t = uint64_t;
+using aff_t = float;
+
+std::vector<std::size_t> run_watershed(
+        size_t          width,
+        size_t          height,
+        size_t          depth,
+        const float*    affinity_data,
+        uint64_t*       segmentation_data,
+        size_t          size_threshold,
+        float           aff_threshold_low,
+        float           aff_threshold_high);
+
+#endif

setup.py

@@ -0,0 +1,35 @@
+from setuptools import setup, find_packages
+from setuptools.extension import Extension
+import os
+import numpy
+
+# VERSION = '1.0'
+
+source_dir = os.path.dirname(os.path.abspath(__file__))
+include_dirs = [
+    os.path.join(source_dir, 'src'),
+    numpy.get_include(),
+]
+
+extensions = [
+    Extension(
+        'abiss',
+        sources=[
+            'python/abiss.pyx',
+            'python/abiss_wrapper.cpp',
+        ],
+        include_dirs=include_dirs,
+        language='c++',
+        # extra_link_args=['-lboost_iostreams', '-fopenmp', '-ltbb'],  # only necessary for when agg is included
+        extra_compile_args=['-std=c++17', '-w'],
+        # undef_macros=["NDEBUG"],
+    )
+]
+
+
+setup(
+    name='abiss',
+    setup_requires=['numpy'],
+    install_requires=['numpy', 'cython'],
+    ext_modules=extensions,
+)

src/ws/types.hpp

@@ -36,15 +36,27 @@ using id_pair = std::pair<T, T>;
 template < typename T >
 using volume = boost::multi_array<T,3>;
 
+template < typename T >
+using volume_ref = boost::multi_array_ref<T,3>;
+
+template < typename T >
+using const_volume_ref = boost::const_multi_array_ref<T,3>;
+
 template < typename T >
 using affinity_graph = boost::multi_array_ref<T,4>;
 
+template < typename T >
+using const_affinity_graph = boost::const_multi_array_ref<T,4>;
+
 template < typename T >
 using volume_ptr = std::shared_ptr<volume<T>>;
 
 template < typename T >
 using affinity_graph_ptr = std::shared_ptr<affinity_graph<T>>;
 
+template < typename T >
+using const_affinity_graph_ptr = std::shared_ptr<const_affinity_graph<T>>;
+
 template< typename ID, typename F >
 using region_graph = std::vector<std::tuple<F,ID,ID>>;