refactor: replace CloudVolume Skeleton with osteoid

.github/workflows/test.yml

@@ -15,7 +15,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
 
     steps:
     - uses: actions/checkout@v2

README.md

@@ -87,7 +87,7 @@ skels = kimimaro.skeletonize(
 #            adjacent or overlapping images.
 
 import kimimaro
-from cloudvolume import Skeleton
+from osteoid import Skeleton
 
 skels = ... # a set of skeletons produced from the same label id
 skel = Skeleton.simple_merge(skels).consolidate()

automated_test.py

@@ -2,7 +2,7 @@
 
 import edt
 import numpy as np
-from cloudvolume import *
+from osteoid import Skeleton
 
 import kimimaro.intake
 import kimimaro.skeletontricks

ext/skeletontricks/skeletontricks.pyx

@@ -45,8 +45,7 @@ import numpy as np
 
 from collections import defaultdict
 
-cdef extern from "math.h":
-  float INFINITY
+cdef float INFINITY = float('inf')
 
 ctypedef fused UINT:
   uint8_t
@@ -191,7 +190,7 @@ def inf2zero(cnp.ndarray[float, cast=True, ndim=3] field):
   for z in range(0, sz):
     for y in range(0, sy):
       for x in range(0, sx):
-        if (field[x,y,z] == INFINITY):
+        if field[x,y,z] == INFINITY:
           field[x,y,z] = 0
 
   return field

kimimaro/intake.py

@@ -27,9 +27,7 @@
 import scipy.spatial
 from tqdm import tqdm
 
-import cloudvolume
-from cloudvolume import CloudVolume, PrecomputedSkeleton, Bbox
-import cloudvolume.sharedmemory as shm
+from osteoid import Skeleton, Bbox
 
 import cc3d # connected components
 import edt # euclidean distance transform
@@ -39,6 +37,7 @@
 import kimimaro.skeletontricks
 import kimimaro.trace
 
+from . import sharedmemory as shm
 from .utility import compute_cc_labels, find_objects
 
 class DimensionError(Exception):
@@ -139,7 +138,7 @@ def skeletonize(
       chunk size is set higher than num tasks // parallel, that number
       is used instead.
 
-  Returns: { $segid: cloudvolume.PrecomputedSkeleton, ... }
+  Returns: { $segid: osteoid.Skeleton, ... }
   """
 
   anisotropy = np.array(anisotropy, dtype=np.float32)
@@ -539,7 +538,7 @@ def compute_border_targets(cc_labels, anisotropy):
 def merge(skeletons):
   merged_skels = {}
   for segid, skels in skeletons.items():
-    skel = PrecomputedSkeleton.simple_merge(skels)
+    skel = Skeleton.simple_merge(skels)
     merged_skels[segid] = skel.consolidate()
 
   return merged_skels

kimimaro/postprocess.py

@@ -34,7 +34,7 @@
 import scipy.sparse.csgraph as csgraph
 import scipy.spatial.distance
 
-from cloudvolume import Skeleton, Bbox
+from osteoid import Skeleton, Bbox
 
 import kimimaro.skeletontricks
 
@@ -90,7 +90,9 @@ def join_close_components(skeletons, radius=None):
   if radius is not None and radius <= 0:
     raise ValueError("radius must be greater than zero: " + str(radius))
 
-  if isinstance(skeletons, Skeleton):
+  try:
+    iter(skeletons)
+  except TypeError:
     skeletons = [ skeletons ]
 
   skels = []

kimimaro/sharedmemory.py

@@ -0,0 +1,211 @@
+from collections import defaultdict
+import errno
+import mmap
+import os
+import sys
+import time
+
+import multiprocessing as mp
+
+import numpy as np
+
+from osteoid import Bbox, Vec
+
+from .utility import mkdir
+
+SHM_DIRECTORY = '/dev/shm/'
+EMULATED_SHM_DIRECTORY = '/tmp/kimimaro-shm'
+
+EMULATE_SHM = not os.path.isdir(SHM_DIRECTORY)
+PLATFORM_SHM_DIRECTORY = SHM_DIRECTORY if not EMULATE_SHM else EMULATED_SHM_DIRECTORY
+
+class SharedMemoryReadError(Exception):
+  pass
+
+class SharedMemoryAllocationError(Exception):
+  pass
+
+def ndarray(shape, dtype, location, order='F', readonly=False, lock=None, **kwargs):
+  """
+  Create a shared memory numpy array. 
+  Lock is only necessary while doing multiprocessing on 
+  platforms without /dev/shm type  shared memory as 
+  filesystem emulation will be used instead.
+
+  Allocating the shared array requires cleanup on your part.
+  A shared memory file will be located at sharedmemory.PLATFORM_SHM_DIRECTORY + location
+  and must be unlinked when you're done. It will outlive the program.
+
+  You should also call .close() on the mmap file handle when done. However,
+  this is less of a problem because the operating system will close the
+  file handle on process termination.
+
+  Parameters:
+  shape: same as numpy.ndarray
+  dtype: same as numpy.ndarray
+  location: the shared memory filename 
+  lock: (optional) multiprocessing.Lock
+
+  Returns: (mmap filehandle, shared ndarray)
+  """
+  if EMULATE_SHM:
+    return ndarray_fs(
+      shape, dtype, location, lock, 
+      readonly, order, emulate_shm=True, **kwargs
+    )
+  return ndarray_shm(shape, dtype, location, readonly, order, **kwargs)
+
+def ndarray_fs(
+    shape, dtype, location, lock, 
+    readonly=False, order='F', emulate_shm=False,
+    **kwargs
+  ):
+  """Emulate shared memory using the filesystem."""
+  dbytes = np.dtype(dtype).itemsize
+  nbytes = Vec(*shape).rectVolume() * dbytes
+
+  if emulate_shm:
+    directory = mkdir(EMULATED_SHM_DIRECTORY)
+    filename = os.path.join(directory, location)
+  else:
+    filename = location
+
+  if lock:
+    lock.acquire()
+
+  try:
+    allocate_shm_file(filename, nbytes, dbytes, readonly)
+  finally:
+    if lock:
+      lock.release()
+
+  with open(filename, 'r+b') as f:
+    array_like = mmap.mmap(f.fileno(), 0) # map entire file
+
+  renderbuffer = np.ndarray(buffer=array_like, dtype=dtype, shape=shape, order=order, **kwargs)
+  renderbuffer.setflags(write=(not readonly))
+  return array_like, renderbuffer
+
+def allocate_shm_file(filename, nbytes, dbytes, readonly):
+  try:
+    size = os.path.getsize(filename)
+    exists = True
+  except FileNotFoundError:
+    size = 0
+    exists = False
+
+  if readonly and not exists:
+    raise SharedMemoryReadError(filename + " has not been allocated. Requested " + str(nbytes) + " bytes.")
+  elif readonly and size != nbytes:
+    raise SharedMemoryReadError("{} exists, but the allocation size ({} bytes) does not match the request ({} bytes).".format(
+      filename, size, nbytes
+    ))
+
+  if exists: 
+    if size > nbytes:
+      with open(filename, 'wb') as f:
+        os.ftruncate(f.fileno(), nbytes)
+    elif size < nbytes:
+      # too small? just remake it below
+      os.unlink(filename) 
+
+  exists = os.path.exists(filename)
+
+  if not exists:
+    # Previously we were writing out real files full of zeros, 
+    # but a) that takes forever and b) modern OSes support sparse
+    # files (i.e. gigabytes of zeros that take up only a few real bytes).
+    #
+    # The following should take advantage of this functionality and be faster.
+    # It should work on Python 2.7 Unix, and Python 3.5+ on Unix and Windows.
+    #
+    # References:
+    #   https://stackoverflow.com/questions/8816059/create-file-of-particular-size-in-python
+    #   https://docs.python.org/3/library/os.html#os.ftruncate
+    #   https://docs.python.org/2/library/os.html#os.ftruncate
+    #
+    with open(filename, 'wb') as f:
+      os.ftruncate(f.fileno(), nbytes)
+
+def ndarray_shm(shape, dtype, location, readonly=False, order='F', **kwargs):
+  """Create a shared memory numpy array. Requires /dev/shm to exist."""
+  import posix_ipc
+  from posix_ipc import O_CREAT
+  import psutil
+
+  nbytes = Vec(*shape).rectVolume() * np.dtype(dtype).itemsize
+  available = psutil.virtual_memory().available
+
+  preexisting = 0
+  # This might only work on Ubuntu
+  shmloc = os.path.join(SHM_DIRECTORY, location)
+  if os.path.exists(shmloc):
+    preexisting = os.path.getsize(shmloc)
+  elif readonly:
+    raise SharedMemoryReadError(shmloc + " has not been allocated. Requested " + str(nbytes) + " bytes.")
+
+  if readonly and preexisting != nbytes:
+    raise SharedMemoryReadError("{} exists, but the allocation size ({} bytes) does not match the request ({} bytes).".format(
+      shmloc, preexisting, nbytes
+    ))
+
+  if (nbytes - preexisting) > available:
+    overallocated = nbytes - preexisting - available
+    overpercent = (100 * overallocated / (preexisting + available))
+    raise SharedMemoryAllocationError("""
+      Requested more memory than is available. 
+
+      Shared Memory Location:  {}
+
+      Shape:                   {}
+      Requested Bytes:         {} 
+
+      Available Bytes:         {} 
+      Preexisting Bytes*:      {} 
+
+      Overallocated Bytes*:    {} (+{:.2f}%)
+
+      * Preexisting is only correct on linux systems that support /dev/shm/""" \
+        .format(location, shape, nbytes, available, preexisting, overallocated, overpercent))
+
+  # This might seem like we're being "extra safe" but consider
+  # a threading condition where the condition of the shared memory
+  # was adjusted between the check above and now. Better to make sure
+  # that we don't accidently change anything if readonly is set.
+  flags = 0 if readonly else O_CREAT 
+  size = 0 if readonly else int(nbytes) 
+
+  try:
+    shared = posix_ipc.SharedMemory(location, flags=flags, size=size)
+    array_like = mmap.mmap(shared.fd, shared.size)
+    os.close(shared.fd)
+    renderbuffer = np.ndarray(buffer=array_like, dtype=dtype, shape=shape, order=order, **kwargs)
+  except OSError as err:
+    if err.errno == errno.ENOMEM: # Out of Memory
+      posix_ipc.unlink_shared_memory(location)      
+    raise
+
+  renderbuffer.setflags(write=(not readonly))
+  return array_like, renderbuffer
+
+def unlink(location):
+  if EMULATE_SHM:
+    return unlink_fs(location)
+  return unlink_shm(location)
+
+def unlink_shm(location):
+  import posix_ipc
+  try:
+    posix_ipc.unlink_shared_memory(location)
+  except posix_ipc.ExistentialError:
+    return False
+  return True
+
+def unlink_fs(location):
+  directory = mkdir(EMULATED_SHM_DIRECTORY)
+  try:
+    filename = os.path.join(directory, location)
+    os.unlink(filename)
+    return True
+  except OSError:
+    return False

kimimaro/trace.py

@@ -3,7 +3,7 @@
 
 Authors: Alex Bae and Will Silversmith
 Affiliation: Seung Lab, Princeton Neuroscience Institue
-Date: June 2018 - August 2021
+Date: June 2018 - Februrary 2025
 
 This file is part of Kimimaro.
 
@@ -31,8 +31,7 @@
 
 import kimimaro.skeletontricks
 
-from cloudvolume import PrecomputedSkeleton
-from cloudvolume.lib import save_images, mkdir
+from osteoid import Skeleton
 
 def trace(
     labels, DBF, 
@@ -131,7 +130,7 @@ def trace(
     root = find_root(labels, anisotropy, voxel_graph)
 
   if root is None:
-    return PrecomputedSkeleton()
+    return Skeleton()
 
   free_space_radius = 0 if not soma_mode else DBF[root]
   # DBF: Distance to Boundary Field
@@ -184,8 +183,8 @@ def trace(
     max_paths, voxel_graph
   )
 
-  skel = PrecomputedSkeleton.simple_merge(
-    [ PrecomputedSkeleton.from_path(path) for path in paths if len(path) > 0 ]
+  skel = Skeleton.simple_merge(
+    [ Skeleton.from_path(path) for path in paths if len(path) > 0 ]
   ).consolidate()
 
   verts = skel.vertices.flatten().astype(np.uint32)
@@ -385,27 +384,8 @@ def point_to_point(
   del DAF
 
   path = dijkstra3d.dijkstra(PDRF, end, start)
-  skel = PrecomputedSkeleton.from_path(path)
+  skel = Skeleton.from_path(path)
 
   verts = skel.vertices.flatten().astype(np.uint32)
   skel.radii = DBF[verts[::3], verts[1::3], verts[2::3]]
   return skel
-
-def xy_path_projection(paths, labels, N=0):
-  """Used for debugging paths."""
-  if type(paths) != list:
-    paths = [ paths ]
-
-  projection = np.zeros( (labels.shape[0], labels.shape[1] ), dtype=np.uint8)
-  outline = labels.any(axis=-1).astype(np.uint8) * 77
-  outline = outline.reshape( (labels.shape[0], labels.shape[1] ) )
-  projection += outline
-  for path in paths:
-    for coord in path:
-      projection[coord[0], coord[1]] = 255
-
-  projection = Image.fromarray(projection.T, 'L')
-  N = str(N).zfill(3)
-  mkdir('./saved_images/projections')
-  projection.save('./saved_images/projections/{}.png'.format(N), 'PNG')
-