perf: faster foreground offset

Author: william-silversmith

cc3d/__init__.py

@@ -410,7 +410,8 @@ def connected_components_stack(
       return_N=True, out_dtype=np.uint64,
       binary_image=bool(binary_image),
     )
-    np.add(cc_labels, offset, out=cc_labels, where=(cc_labels != 0))
+    # np.add(cc_labels, offset, out=cc_labels, where=(cc_labels != 0))
+    fastcc3d.offset_foreground(cc_labels, offset)
     offset += N
     binary = crackle.compress(cc_labels)
 

cc3d/fastcc3d.pyx

@@ -1370,3 +1370,20 @@ def each(
   if in_place:
     return InPlaceImageIterator()
   return ImageIterator()
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def offset_foreground(
+  cnp.ndarray[cnp.uint64_t, ndim=3] labels,
+  cnp.uint64_t offset,
+):
+  cdef:
+    Py_ssize_t i
+    Py_ssize_t n = labels.size
+    cnp.uint64_t* data = <cnp.uint64_t*>labels.data
+
+  with nogil:
+    for i in range(n):
+      if data[i] != 0:
+        data[i] += offset
+