feat: add oversegment function (#85)

* feat: add oversegmentation using skeletons

* refactor: put binary image iterator code in a separate function

* install: bump dijkstra3d to 1.15.0

This has the new EDF feature map

* feat: return copies of skeletons with segments attribute

* merge: add visualize_section_planes feature

* docs: show how to use oversegment

README.md

@@ -127,7 +127,7 @@ skel = kimimaro.join_close_components([skel1, skel2], radius=None) # no threshol
 extra_targets = kimimaro.synapses_to_targets(labels, synapses)
 
 
-# LISTING 3: Drawing a centerline between
+# LISTING 4: Drawing a centerline between
 #   preselected points on a binary image.
 #   This is a much simpler option for when
 #   you know exactly what you want, but may
@@ -139,6 +139,19 @@ skel = kimimaro.connect_points(
   end=(121, 426, 227),
   anisotropy=(32,32,40),
 )
+
+# LISTING 5: Using skeletons to oversegment existing
+#  segmentations for integration into proofreading systems 
+#  that on merging atomic labels. oversegmented_labels 
+#  is returned numbered from 1. skels is a copy returned 
+#  with the property skel.segments that associates a label
+#  to each vertex (labels will not be unique if downsampling 
+#  is used)
+oversegmented_labels, skels = kimimaro.oversegment(
+  labels, skels, 
+  anisotropy=(32,32,40), 
+  downsample=10,
+)
 ```
 
 `connectomics.npy` is multilabel connectomics data derived from pinky40, a 2018 experimental automated segmentation of ~1.5 million cubic micrometers of mouse visual cortex. It is an early predecessor to the now public pinky100_v185 segmentation that can be found at https://microns-explorer.org/phase1 You will need to run `lzma -d connectomics.npy.lzma` to obtain the 512x512x512 uint32 volume at 32x32x40 nm<sup>3</sup> resolution.  

kimimaro/__init__.py

@@ -17,4 +17,4 @@
 
 from .intake import skeletonize, DimensionError, synapses_to_targets, connect_points
 from .postprocess import postprocess, join_close_components
-from .utility import extract_skeleton_from_binary_image, cross_sectional_area
+from .utility import extract_skeleton_from_binary_image, cross_sectional_area, oversegment

kimimaro/utility.py

@@ -1,4 +1,6 @@
-from typing import Dict, Union, List
+from typing import Dict, Union, List, Tuple
+
+import copy
 
 import numpy as np
 import scipy.ndimage
@@ -8,6 +10,7 @@
 import kimimaro.skeletontricks
 
 import cc3d
+import dijkstra3d
 import fastremap
 import fill_voids
 import xs3d
@@ -50,6 +53,54 @@ def add_property(skel, prop):
   if needs_prop:
     skel.extra_attributes.append(prop)
 
+def shape_iterator(all_labels, skeletons, fill_holes, in_place, progress, fn):
+  iterator = skeletons
+  if type(skeletons) == dict:
+    iterator = skeletons.values()
+    total = len(skeletons)
+  elif type(skeletons) == Skeleton:
+    iterator = [ skeletons ]
+    total = 1
+  else:
+    total = len(skeletons)
+
+  if all_labels.dtype == bool:
+    remapping = { True: 1, False: 0, 1:1, 0:0 }
+  else:
+    all_labels, remapping = fastremap.renumber(all_labels, in_place=in_place)
+
+  all_slices = find_objects(all_labels)
+
+  for skel in tqdm(iterator, desc="Labels", disable=(not progress), total=total):
+    if all_labels.dtype == bool:
+      label = 1
+    else:
+      label = skel.id
+
+    if label == 0:
+      continue
+
+    label = remapping[label]
+    slices = all_slices[label - 1]
+    if slices is None:
+      continue
+
+    roi = Bbox.from_slices(slices)
+    if roi.volume() <= 1:
+      continue
+
+    roi.grow(1)
+    roi.minpt = Vec.clamp(roi.minpt, Vec(0,0,0), roi.maxpt)
+    slices = roi.to_slices()
+
+    binimg = np.asfortranarray(all_labels[slices] == label)
+    if fill_holes:
+      binimg = fill_voids.fill(binimg, in_place=True)
+
+    fn(skel, binimg, roi)
+
+  return iterator
+
 def cross_sectional_area(
   all_labels:np.ndarray, 
   skeletons:Union[Dict[int,Skeleton],List[Skeleton],Skeleton],
@@ -100,54 +151,11 @@ def cross_sectional_area(
     "num_components": 1,
   }
 
-  iterator = skeletons
-  if type(skeletons) == dict:
-    iterator = skeletons.values()
-    total = len(skeletons)
-  elif type(skeletons) == Skeleton:
-    iterator = [ skeletons ]
-    total = 1
-  else:
-    total = len(skeletons)
-
-  if all_labels.dtype == bool:
-    remapping = { True: 1, False: 0, 1:1, 0:0 }
-  else:
-    all_labels, remapping = fastremap.renumber(all_labels, in_place=in_place)
-
-  all_slices = find_objects(all_labels)
-
-  for skel in tqdm(iterator, desc="Labels", disable=(not progress), total=total):
-    if all_labels.dtype == bool:
-      label = 1
-    else:
-      label = skel.id
-
-    if label == 0:
-      continue
-
-    label = remapping[label]
-    slices = all_slices[label - 1]
-    if slices is None:
-      continue
-
-    roi = Bbox.from_slices(slices)
-    if roi.volume() <= 1:
-      continue
-
-    roi.grow(1)
-    roi.minpt = Vec.clamp(roi.minpt, Vec(0,0,0), roi.maxpt)
-    slices = roi.to_slices()
-
-    binimg = np.asfortranarray(all_labels[slices] == label)
-
+  def cross_sectional_area_helper(skel, binimg, roi):
     cross_sections = None
     if visualize_section_planes:
       cross_sections = np.zeros(binimg.shape, dtype=np.uint32, order="F")
 
-    if fill_holes:
-      binimg = fill_voids.fill(binimg, in_place=True)
-
     all_verts = (skel.vertices / anisotropy).round().astype(int)
     all_verts -= roi.minpt
 
@@ -207,11 +215,90 @@ def cross_sectional_area(
       microviewer.view(cross_sections, seg=True)
 
     add_property(skel, prop)
+
     skel.cross_sectional_area = areas
     skel.cross_sectional_area_contacts = contacts
 
+  shape_iterator(
+    all_labels, skeletons, 
+    fill_holes, in_place, progress, 
+    cross_sectional_area_helper
+  )
+
   return skeletons
 
+def oversegment(
+  all_labels:np.ndarray, 
+  skeletons:Union[Dict[int,Skeleton],List[Skeleton],Skeleton],
+  anisotropy:np.ndarray = np.array([1,1,1], dtype=np.float32),
+  progress:bool = False,
+  fill_holes:bool = False,
+  in_place:bool = False,
+  downsample:int = 0,
+) -> Tuple[np.ndarray, Union[Dict[int,Skeleton],List[Skeleton],Skeleton]]:
+  """
+  Use skeletons to create an oversegmentation of a pre-existing set
+  of labels. This is useful for proofreading systems that work by merging
+  labels.
+
+  For each skeleton, get the feature map from its euclidean distance
+  field. The final image is the composite of all these feature maps
+  numbered from 1.
+
+  Each skeleton will have a new property skel.segments that associates
+  a label to each vertex.
+  """
+  prop = {
+    "id": "segments",
+    "data_type": "uint64",
+    "num_components": 1,
+  }
+
+  skeletons = copy.deepcopy(skeletons)
+
+  all_features = np.zeros(all_labels.shape, dtype=np.uint64, order="F")
+  next_label = 0
+
+  def oversegment_helper(skel, binimg, roi):
+    nonlocal next_label
+    nonlocal all_features
+
+    segment_skel = skel
+    if downsample > 0:
+      segment_skel = skel.downsample(downsample)
+
+    vertices = (segment_skel.vertices / anisotropy).round().astype(int)
+    vertices -= roi.minpt
+
+    field, feature_map = dijkstra3d.euclidean_distance_field(
+      binimg, vertices, 
+      anisotropy=anisotropy, 
+      return_feature_map=True
+    )
+    del field
+
+    add_property(skel, prop)
+
+    vertices = (skel.vertices / anisotropy).round().astype(int)
+    vertices -= roi.minpt
+
+    feature_map[binimg] += next_label
+    skel.segments = feature_map[vertices[:,0], vertices[:,1], vertices[:,2]]
+    next_label += vertices.shape[0]
+    all_features[roi.to_slices()] += feature_map
+
+  # iterator is an iterable list of skeletons, not the shape iterator
+  iterator = shape_iterator(
+    all_labels, skeletons, fill_holes, in_place, progress, 
+    oversegment_helper
+  )
+
+  all_features, mapping = fastremap.renumber(all_features)
+  for skel in iterator:
+    skel.segments = fastremap.remap(skel.segments, mapping, in_place=True)
+
+  return all_features, skeletons
+
 # From SO: https://stackoverflow.com/questions/14313510/how-to-calculate-rolling-moving-average-using-python-numpy-scipy
 def moving_average(a:np.ndarray, n:int, mode:str = "symmetric") -> np.ndarray:
   if n <= 0:

requirements.txt

@@ -1,7 +1,7 @@
 click
 connected-components-3d>=1.5.0
 cloud-volume>=0.57.6
-dijkstra3d>=1.12.0
+dijkstra3d>=1.15.0
 fill-voids>=2.0.0
 edt>=2.1.0
 fastremap>=1.10.2