feat: skeletonize (most) autapses and improve cross section analysis (#…

…179)

* wip(skeletonizing): add support for fixing autapses in graphene volumes

* feat: add support for graphene timestamp

* wip: ensure shape matches atomic chunks

* fix: add back in xyzrange

* fix(SkeletonTask): generalize frag_path for windows

Also interpret bare paths as "file://"

* fix(SkeletonTask): use agglomerate and timestamp correctly

* fix: merge error removed cross_sectional_area params

* fix(SkeletonTask): apply timestamp in correct location

* fix: fixup the edges of the vcg with root vcg

* feat: remove requirements for a queue parameter

* redesign: use vol.info to get skeleton path

* fix: several errors

* feat: add fix_autapses to cli

* fix: compensate for mip levels

* fix: wrong path joining

* fix: ensure mesh bounds are rectified after expansion

* fix: wrong variable name in shard downsample

* fix(shards/image): fix situation where chunk.z == dataset.z

* wip(skeletonizing): add support for fixing autapses in graphene volumes

* feat: add support for graphene timestamp

* wip: ensure shape matches atomic chunks

* fix: add back in xyzrange

* fix(SkeletonTask): generalize frag_path for windows

Also interpret bare paths as "file://"

* fix(SkeletonTask): use agglomerate and timestamp correctly

* fix: merge error removed cross_sectional_area params

* fix(SkeletonTask): apply timestamp in correct location

* fix: fixup the edges of the vcg with root vcg

* feat: remove requirements for a queue parameter

* redesign: use vol.info to get skeleton path

* fix: several errors

* feat: add fix_autapses to cli

* fix: compensate for mip levels

* fix: wrong path joining

* fix: ensure mesh bounds are rectified after expansion

* fix: don't import view from CloudVolume

* fix: don't import view from cloudvolume

* fix: intelligent frag path resolution

* fix: handle binarization issue in repairing

* fix: smart path

* perf: skip downloading the "huge" bbox if no skeletons

* fix: ensure info file uploaded to frag path

* fix: check for key

* feat: support agglomeration for graphene for cross sections

* feat: support static root_ids for a graphene volume

* fix: regression in LuminanceLevelsTask due to CloudVolume upgrade

* refactor: simplify vcg code

* feat: add encoding level support for jpegxl

* fix: don't print error if max_mips is non-zero

---------

Co-authored-by: William Silversmith <[email protected]>

igneous/task_creation/common.py

@@ -226,6 +226,8 @@ def set_encoding(cv, mip, encoding, encoding_level):
 
  if encoding == "jpeg":
  scale["jpeg_quality"] = encoding_level
+ elif encoding == "jpegxl":
+ scale["jpegxl_quality"] = encoding_level
  elif encoding == "png":
  scale["png_level"] = encoding_level
  elif encoding == "fpzip":

igneous/task_creation/image.py

@@ -406,14 +406,6 @@ def create_sharded_image_info(
  # maximum amount of information in the morton codes
  grid_size = np.ceil(Vec(*dataset_size) / Vec(*chunk_size)).astype(np.int64)
  max_bits = sum([ math.ceil(math.log2(size)) for size in grid_size ])
- if max_bits > 64:
- raise ValueError(
- f"{max_bits}, more than a 64-bit integer, "
- "would be required to describe the chunk positions "
- "in this dataset. Try increasing the chunk size or "
- "increasing dataset bounds."
- f"Dataset Size: {dataset_size} Chunk Size: {chunk_size}"
- )
 
  chunks_per_shard = math.ceil(uncompressed_shard_bytesize / (chunk_voxels * byte_width))
  chunks_per_shard = 2 ** int(math.log2(chunks_per_shard))
@@ -423,7 +415,7 @@ def create_sharded_image_info(
 
  # approximate, would need to account for rounding effects to be exact
  # rounding is corrected for via max_bits - pre - mini below.
- num_shards = num_chunks / chunks_per_shard 
+ num_shards = num_chunks / chunks_per_shard
 
  def update_bits():
  shard_bits = int(math.ceil(math.log2(num_shards)))
@@ -465,7 +457,25 @@ def update_bits():
  # in the morton codes, so if there's any slack from rounding, the
  # remainder goes into shard bits.
  preshift_bits = preshift_bits - minishard_bits
- shard_bits = max_bits - preshift_bits - minishard_bits
+ if dataset_size[2] == chunk_size[2]:
+ additional_bits = (preshift_bits // 3)
+ i = 0
+ while i < additional_bits:
+ max_bits += 1
+ preshift_bits += 1
+ if preshift_bits % 3 != 0:
+ i += 1
+
+ shard_bits = max(max_bits - preshift_bits - minishard_bits, 0)
+
+ if max_bits > 64:
+ raise ValueError(
+ f"{max_bits}, more than a 64-bit integer, "
+ "would be required to describe the chunk positions "
+ "in this dataset. Try increasing the chunk size or "
+ "increasing dataset bounds."
+ f"Dataset Size: {dataset_size} Chunk Size: {chunk_size}"
+ )
 
  if preshift_bits < 0:
  raise ValueError(f"Preshift bits cannot be negative. ({shard_bits}, {minishard_bits}, {preshift_bits}), total info: {max_bits} bits")

igneous/task_creation/skeleton.py

@@ -15,7 +15,7 @@
 from cloudvolume import CloudVolume
 from cloudvolume.lib import Vec, Bbox, max2, min2, xyzrange, find_closest_divisor, yellow, jsonify
 from cloudvolume.datasource.precomputed.sharding import ShardingSpecification
-from cloudfiles import CloudFiles
+from cloudfiles import CloudFiles, CloudFile
 
 from igneous.tasks import ( 
  SkeletonTask, UnshardedSkeletonMergeTask, 
@@ -58,6 +58,7 @@ def bounds_from_mesh(
  bbxes.append(bounds)
 
  bounds = Bbox.expand(*bbxes)
+ bounds = bounds.expand_to_chunk_size(shape, offset=vol.voxel_offset)
  return Bbox.clamp(bounds, vol.bounds)
 
 def create_skeletonizing_tasks(
@@ -71,9 +72,11 @@ def create_skeletonizing_tasks(
  parallel=1, fill_missing=False, 
  sharded=False, frag_path=None, spatial_index=True,
  synapses=None, num_synapses=None,
- dust_global=False, 
+ dust_global=False, fix_autapses=False,
  cross_sectional_area=False,
  cross_sectional_area_smoothing_window=5,
+ timestamp=None,
+ root_ids_cloudpath=None,
  ):
  """
  Assign tasks with one voxel overlap in a regular grid 
@@ -121,6 +124,17 @@ def create_skeletonizing_tasks(
  fix_borders: Allows trivial merging of single overlap tasks. You'll only
  want to set this to false if you're working on single or non-overlapping
  volumes.
+ fix_autapses: Only possible for graphene volumes. Uses PyChunkGraph (PCG) information
+ to fix autapses (when a neuron synapses onto itself). This requires splitting
+ contacts between the edges of two touching voxels. The algorithm for doing this
+ requires much more memory.
+
+ This works by comparing the PYC L2 and root layers. L1 is watershed. L2 is the
+ connections only within an atomic chunk. The root layer provides the global
+ connectivity. Autapses can be distinguished at the L2 level, above that, they
+ may not be (and certainly not at the root level). We extract the voxel connectivity
+ graph from L2 and perform the overall trace at root connectivity.
+
  dust_threshold: don't skeletonize labels smaller than this number of voxels
  as seen by a single task.
  dust_global: Use global voxel counts for the dust threshold instead of from
@@ -155,10 +169,24 @@ def create_skeletonizing_tasks(
  to the total computation.)
  cross_sectional_area_smoothing_window: Perform a rolling average of the 
  normal vectors across these many vectors.
+ timestamp: for graphene volumes only, you can specify the timepoint to use
+ root_ids_cloudpath: for graphene volumes, if you have a materialized archive
+ if your desired timepoint, you can use this path for fetching root ID 
+ segmentation as it is far more efficient.
  """
  shape = Vec(*shape)
  vol = CloudVolume(cloudpath, mip=mip, info=info)
 
+ if fix_autapses:
+ if vol.meta.path.format != "graphene":
+ raise ValueError("fix_autapses can only be performed on graphene volumes.")
+
+ if not np.all(shape % vol.meta.graph_chunk_size == 0):
+ raise ValueError(
+ f"shape must be a multiple of the graph chunk size. Got: {shape}, "
+ f"{vol.meta.graph_chunk_size}"
+ )
+
  if dust_threshold > 0 and dust_global:
  cf = CloudFiles(cloudpath)
  vxctfile = cf.join(vol.key, 'stats', 'voxel_counts.mb')
@@ -201,6 +229,15 @@ def create_skeletonizing_tasks(
 
  vol.skeleton.meta.commit_info()
 
+ if frag_path:
+ frag_info_path = CloudFiles(frag_path).join(frag_path, "info")
+ frag_info = CloudFile(frag_info_path).get_json()
+ if not frag_info:
+ CloudFile(frag_info_path).put_json(vol.skeleton.meta.info)
+ elif 'scales' in frag_info:
+ frag_info_path = CloudFiles(frag_path).join(frag_path, vol.info["skeletons"], "info")
+ CloudFile(frag_info_path).put_json(vol.skeleton.meta.info)
+
  will_postprocess = bool(np.any(vol.bounds.size3() > shape))
  bounds = vol.bounds.clone()
 
@@ -247,8 +284,11 @@ def task(self, shape, offset):
  spatial_grid_shape=shape.clone(), # used for writing index filenames
  synapses=bbox_synapses,
  dust_global=dust_global,
+ fix_autapses=bool(fix_autapses),
+ timestamp=timestamp,
  cross_sectional_area=bool(cross_sectional_area),
  cross_sectional_area_smoothing_window=int(cross_sectional_area_smoothing_window),
+ root_ids_cloudpath=root_ids_cloudpath,
  )
 
  def synapses_for_bbox(self, shape, offset):
@@ -292,8 +332,11 @@ def on_finish(self):
  'spatial_index': bool(spatial_index),
  'synapses': bool(synapses),
  'dust_global': bool(dust_global),
+ 'fix_autapses': bool(fix_autapses),
+ 'timestamp': timestamp,
  'cross_sectional_area': bool(cross_sectional_area),
  'cross_sectional_area_smoothing_window': int(cross_sectional_area_smoothing_window),
+ 'root_ids_cloudpath': root_ids_cloudpath,
  },
  'by': operator_contact(),
  'date': strftime('%Y-%m-%d %H:%M %Z'),

igneous/tasks/image/image.py

@@ -70,7 +70,7 @@ def downsample_and_upload(
  if max_mips is not None:
  factors = factors[:max_mips]
 
- if len(factors) == 0:
+ if len(factors) == 0 and max_mips:
  print("No factors generated. Image Shape: {}, Downsample Shape: {}, Volume Shape: {}, Bounds: {}".format(
  image.shape, ds_shape, vol.volume_size, bounds)
  )
@@ -327,6 +327,9 @@ def execute(self):
  cts = np.bincount(img2d)
  levels[0:len(cts)] += cts.astype(np.uint64)
 
+ if len(bboxes) == 0:
+ return
+
  covered_area = sum([bbx.volume() for bbx in bboxes])
 
  bboxes = [(bbox.volume(), bbox.size3()) for bbox in bboxes]
@@ -376,7 +379,8 @@ def select_bounding_boxes(self, dataset_bounds):
  patch_start += self.offset
  bbox = Bbox(patch_start, patch_start + sample_shape.size3())
  bbox = Bbox.clamp(bbox, dataset_bounds)
- bboxes.append(bbox)
+ if not bbox.subvoxel():
+ bboxes.append(bbox)
  return bboxes
 
 @queueable