diff --git a/igneous/task_creation/mesh.py b/igneous/task_creation/mesh.py index 47b110d1..d5af32c7 100644 --- a/igneous/task_creation/mesh.py +++ b/igneous/task_creation/mesh.py @@ -166,7 +166,7 @@ def create_meshing_tasks( ): shape = Vec(*shape) - assert 0 <= fill_holes <= 3, "fill_holes must be between 0 to 3 inclusive." + assert 0 <= fill_holes <= 103, "fill_holes must be between 0 to 103 inclusive." vol = CloudVolume(layer_path, mip) diff --git a/igneous/task_creation/skeleton.py b/igneous/task_creation/skeleton.py index ab368ff3..c7e2e2dd 100644 --- a/igneous/task_creation/skeleton.py +++ b/igneous/task_creation/skeleton.py @@ -150,7 +150,6 @@ def create_skeletonizing_tasks( 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 @@ -202,9 +201,10 @@ def create_skeletonizing_tasks( 0: off 1: simple hole filling 2: also fill borders in 2d on sides of image - 3: also perform a morphological closing using 3x3x3 stencil + 3: also perform a morphological dilation using 3x3x3 stencil + 4+: also decrement merge_threshold by 1% for each point above 3 """ - assert 0 <= fill_holes <= 3, "fill_holes must be between 0 to 3 inclusive." + assert 0 <= fill_holes <= 103, "fill_holes must be between 0 to 103 inclusive." shape = Vec(*shape) vol = CloudVolume(cloudpath, mip=mip, info=info) @@ -378,7 +378,7 @@ def on_finish(self): 'cross_sectional_area_smoothing_window': int(cross_sectional_area_smoothing_window), 'cross_sectional_area_repair_sec_per_label': int(cross_sectional_area_repair_sec_per_label), 'root_ids_cloudpath': root_ids_cloudpath, - 'fill_holes': int(fill_holes) + 'fill_holes': int(fill_holes), }, 'by': operator_contact(), 'date': strftime('%Y-%m-%d %H:%M %Z'), diff --git a/igneous/tasks/mesh/mesh.py b/igneous/tasks/mesh/mesh.py index 0f73fac8..e6508c14 100644 --- a/igneous/tasks/mesh/mesh.py +++ b/igneous/tasks/mesh/mesh.py @@ -87,7 +87,8 @@ def __init__(self, shape, offset, layer_path, **kwargs): 0: off 1: simple hole filling 2: also fill borders in 2d on sides of image - 3: also perform a morphological closing using 3x3x3 stencil + 3: also perform a morphological dilation using 3x3x3 stencil + 4+: also decrement merge_threshold by 1% for each point above 3 """ super(MeshTask, self).__init__(shape, offset, layer_path, **kwargs) self.shape = Vec(*shape) @@ -147,17 +148,14 @@ def execute(self): self._mesher = zmesh.Mesher(self._volume.resolution) - # aggressive morphological hole filling has a 1-2vx - # edge effect that needs to be cropped away fill_level = self.options["fill_holes"] - hole_filling_padding = int(fill_level >= 3) * 2 # Marching cubes loves its 1vx overlaps. # This avoids lines appearing between # adjacent chunks. data_bounds = self._bounds.clone() - data_bounds.minpt -= self.options['low_padding'] + hole_filling_padding - data_bounds.maxpt += self.options['high_padding'] + hole_filling_padding + data_bounds.minpt -= self.options['low_padding'] + data_bounds.maxpt += self.options['high_padding'] self._mesh_dir = self.get_mesh_dir() @@ -203,44 +201,56 @@ def execute(self): data, renumbermap = fastremap.renumber(data, in_place=True) renumbermap = { v:k for k,v in renumbermap.items() } - if fill_level > 0: - filled_data, hole_labels = fastmorph.fill_holes( - data[...,0], - remove_enclosed=True, - return_removed=True, - fix_borders=(fill_level > 1), - morphological_closing=(fill_level > 2), - ) + data = data[...,0] + if fill_level > 0: if fill_level >= 3: - hp = hole_filling_padding - data = np.asfortranarray(data[hp:-hp,hp:-hp,hp:-hp]) - filled_data = np.asfortranarray(filled_data[hp:-hp,hp:-hp,hp:-hp]) - - data = crackle.compress(data) - self._mesher.mesh(filled_data) - meshes, bounding_boxes = self.compute_meshes(renumbermap, left_offset) - del filled_data - data = crackle.decompress(data) - - data *= np.isin(data, list(hole_labels)) - - self._mesher.mesh(data) - hole_meshes, hole_bounding_boxes = self.compute_meshes(renumbermap, left_offset) - meshes.update(hole_meshes) - bounding_boxes.update(hole_bounding_boxes) + data = fastmorph.dilate( + data, + mode=fastmorph.Mode.multilabel, + background_only=True, + parallel=1, + ) + + filled_labels, hole_labels = fastmorph.fill_holes_v2( + data, + return_crackle=True, + fix_borders=(fill_level >= 2), + merge_threshold=( + 1.0 if fill_level <= 3 else (1.0 - 0.01 * (fill_level - 3)) + ), + parallel=1, + ) del data + self._mesher.mesh(filled_labels.numpy()) + + meshes = self.compute_meshes(renumbermap) + del filled_labels + self._mesher.mesh(hole_labels.numpy()) + hole_meshes = self.compute_meshes(renumbermap) + del hole_labels + + for segid in list(hole_meshes.keys()): + if segid in meshes: + meshes[segid] = zmesh.Mesh.concatenate(meshes[segid], hole_meshes[segid], id=segid) + else: + meshes[segid] = hole_meshes[segid] del hole_meshes - del hole_bounding_boxes else: - self._mesher.mesh(data[..., 0]) + self._mesher.mesh(data) del data - - meshes, bounding_boxes = self.compute_meshes(renumbermap, left_offset) + meshes = self.compute_meshes(renumbermap) if self.options['dry_run']: return (meshes, bounding_boxes) + bounding_boxes = {} + + for segid, mesh in meshes.items(): + binary, mesh_bbx = self._create_mesh_binary(mesh, left_offset) + meshes[segid] = binary + bounding_boxes[segid] = mesh_bbx.to_list() + if self.options['sharded']: self._upload_batch(meshes, self._bounds) else: @@ -353,17 +363,19 @@ def _remap(self, data): data = fastremap.mask_except(data, list(remap.keys()), in_place=True) return fastremap.remap(data, remap, in_place=True) - def compute_meshes(self, renumbermap, offset): - bounding_boxes = {} + def compute_meshes(self, renumbermap:dict) -> dict[int,zmesh.Mesh]: meshes = {} - - for obj_id in tqdm(self._mesher.ids(), disable=(not self.progress), desc="Mesh"): + + for obj_id in tqdm(self._mesher.ids(), disable=(not self.progress), desc="Extracting Mesh"): remapped_id = renumbermap[obj_id] - mesh_binary, mesh_bounds = self._create_mesh(obj_id, offset) - bounding_boxes[remapped_id] = mesh_bounds.to_list() - meshes[remapped_id] = mesh_binary + meshes[remapped_id] = self._mesher.get( + obj_id, + reduction_factor=self.options['simplification_factor'], + max_error=self.options['max_simplification_error'], + voxel_centered=True, + ) - return meshes, bounding_boxes + return meshes def _upload_batch(self, meshes, bbox): frag_path = self.options['frag_path'] or self.layer_path @@ -412,16 +424,7 @@ def _upload_individuals(self, mesh_binaries, generate_manifests): cache_control=self.options['cache_control'], ) - def _create_mesh(self, obj_id, left_bound_offset): - mesh = self._mesher.get( - obj_id, - reduction_factor=self.options['simplification_factor'], - max_error=self.options['max_simplification_error'], - voxel_centered=True, - ) - - self._mesher.erase(obj_id) - + def _create_mesh_binary(self, mesh:zmesh.Mesh, left_bound_offset:int) -> tuple[bytes, Bbox]: resolution = self._volume.resolution offset = (self._bounds.minpt - self.options['low_padding']).astype(np.float32) mesh.vertices[:] += (offset - left_bound_offset) * resolution diff --git a/igneous/tasks/skeleton.py b/igneous/tasks/skeleton.py index cd7901ff..ebf62fcb 100644 --- a/igneous/tasks/skeleton.py +++ b/igneous/tasks/skeleton.py @@ -113,10 +113,6 @@ def __init__( self.bounds = Bbox(offset, Vec(*shape) + Vec(*offset)) self.index_bounds = Bbox(offset, Vec(*spatial_grid_shape) + Vec(*offset)) - # aggressive morphological hole filling has a 1-2vx - # edge effect that needs to be cropped away - self.hole_filling_padding = (self.fill_holes >= 3) * 2 - def execute(self): # For graphene volumes, if we've materialized the root IDs # into a static archive, let's use that because it's way more @@ -136,7 +132,7 @@ def execute(self): vol = CloudVolume( cloudpath, mip=self.mip, - bounded=(self.hole_filling_padding == 0), + bounded=True, info=self.info, cdn_cache=False, parallel=self.parallel, @@ -147,9 +143,6 @@ def execute(self): bbox = Bbox.clamp(self.bounds, vol.bounds) index_bbox = Bbox.clamp(self.index_bounds, vol.bounds) - bbox.minpt -= self.hole_filling_padding - bbox.maxpt += self.hole_filling_padding - path = vol.info.get("skeletons", "skeletons") if self.frag_path is None: path = vol.meta.join(self.cloudpath, path) @@ -272,43 +265,40 @@ def should_use_low_memory(self, bbox:Bbox) -> bool: # with delta +250 return bigger_bbx.volume() * 6 > self.cross_sectional_area_low_memory_threshold - def _compute_fill_holes(self, all_labels): - filled_labels, hole_labels_set = fastmorph.fill_holes( - all_labels, - remove_enclosed=True, - return_removed=True, - fix_borders=(self.fill_holes >= 2), - morphological_closing=(self.fill_holes >= 3), - progress=self.progress, - parallel=self.parallel, - ) - - if self.fill_holes >= 3: - hp = self.hole_filling_padding - all_labels = np.asfortranarray(all_labels[hp:-hp,hp:-hp,hp:-hp]) - filled_labels = np.asfortranarray(filled_labels[hp:-hp,hp:-hp,hp:-hp]) - - return (filled_labels, hole_labels_set, all_labels) - - def _do_operation(self, all_labels, fn): + def _do_operation(self, vol, all_labels, fn): if callable(all_labels): all_labels = all_labels() - if self.fill_holes > 0: - filled_labels, hole_labels, all_labels = self._compute_fill_holes(all_labels) - all_labels = crackle.compress(all_labels) - skeletons = fn(filled_labels) - del filled_labels + if self.fill_holes == 0: + return (fn(all_labels), {}) - all_labels = crackle.decompress(all_labels) - hole_labels = fastremap.mask_except(all_labels, list(hole_labels), in_place=True) - del all_labels - hole_skeletons = fn(hole_labels) - skeletons.update(hole_skeletons) - else: - skeletons = fn(all_labels) + if self.fill_holes >= 3: + all_labels = fastmorph.dilate( + all_labels, + mode=fastmorph.Mode.multilabel, + background_only=True, + parallel=self.parallel, + ) - return skeletons + merge_threshold=( + 1.0 if self.fill_holes <= 3 else (1.0 - 0.01 * (self.fill_holes - 3)) + ) + + filled_labels, hole_labels = fastmorph.fill_holes_v2( + all_labels, + fix_borders=(self.fill_holes >= 2), + merge_threshold=merge_threshold, + anisotropy=vol.resolution, + parallel=self.parallel, + return_crackle=True, + ) + del all_labels + + skeletons = fn(filled_labels.numpy()) + del filled_labels + hole_skeletons = fn(hole_labels.numpy()) + + return (skeletons, hole_skeletons) def skeletonize( self, @@ -334,7 +324,15 @@ def do_skeletonize(labels): voxel_graph=voxel_graph, ) - return self._do_operation(all_labels, do_skeletonize) + skeletons, hole_skeletons = self._do_operation(vol, all_labels, do_skeletonize) + + for segid, hole_skel in hole_skeletons.items(): + if segid in skeletons: + skeletons[segid] = Skeleton.simple_merge([ skeletons[segid], hole_skel ]) + else: + skeletons[segid] = hole_skel + + return skeletons def voxel_connectivity_graph( self, @@ -410,7 +408,6 @@ def compute_cross_sectional_area(self, vol, bbox, skeletons): big_bbox = bbox.clone() big_bbox.grow(delta) big_bbox = Bbox.clamp(big_bbox, vol.bounds) - big_bbox.grow(self.hole_filling_padding) true_delta = bbox.minpt - big_bbox.minpt @@ -456,9 +453,10 @@ def do_cross_section(labels): progress=self.progress, in_place=True, fill_holes=False, + multipass=True, ) - skeletons = self._do_operation(download_all_labels, do_cross_section) + skeletons, _ = self._do_operation(vol, download_all_labels, do_cross_section) mapping = { v:k for k,v in mapping.items() } skeletons = { @@ -487,7 +485,6 @@ def compute_cross_sectional_area_low_mem(self, vol, bbox, skeletons): big_bbox = bbox.clone() big_bbox.grow(delta) big_bbox = Bbox.clamp(big_bbox, vol.bounds) - big_bbox.grow(self.hole_filling_padding) true_delta = bbox.minpt - big_bbox.minpt @@ -529,6 +526,8 @@ def __iter__(self): labels=skel_labels, ) + hp = 2 # hole padding for morphological closure + with tqdm( iterator, disable=(not self.progress), @@ -538,16 +537,19 @@ def __iter__(self): pbar.set_postfix(label=str(label)) if self.fill_holes > 0: - binimg = fastmorph.fill_holes( + if self.fill_holes >= 3: + binimg = np.pad(binimg, pad_width=hp, mode='constant') + + binimg = fastmorph.fill_holes_v1( binimg, remove_enclosed=True, fix_borders=(self.fill_holes >= 2), morphological_closing=(self.fill_holes >= 3), progress=False, + parallel=self.parallel, ) if self.fill_holes >= 3: - hp = self.hole_filling_padding binimg = np.asfortranarray(binimg[hp:-hp,hp:-hp,hp:-hp]) bbx = Bbox.from_list(bbxes[label]) @@ -613,9 +615,6 @@ def reprocess_skel(pts, skel): skel_bbx = Bbox.clamp(skel_bbx, vol.bounds) - skel_bbx.minpt -= self.hole_filling_padding - skel_bbx.maxpt += self.hole_filling_padding - binary_image = vol.download( skel_bbx, mip=vol.mip, label=skel.id )[...,0] @@ -628,14 +627,22 @@ def reprocess_skel(pts, skel): segid = skel.id skel.id = 1 + hp = 2 # hole padding + if self.fill_holes > 0: - binary_image = fastmorph.fill_holes( + if self.fill_holes >= 3: + binary_image = np.pad(binary_image, pad_width=hp, mode='constant') + + binary_image = fastmorph.fill_holes_v1( binary_image, + remove_enclosed=True, fix_borders=(self.fill_holes >= 2), morphological_closing=(self.fill_holes >= 3), + progress=False, + parallel=self.parallel, ) + if self.fill_holes >= 3: - hp = self.hole_filling_padding binary_image = np.asfortranarray(binary_image[hp:-hp,hp:-hp,hp:-hp]) skel.vertices -= hp * vol.resolution diff --git a/requirements.txt b/requirements.txt index 24035930..c4ddc0b0 100755 --- a/requirements.txt +++ b/requirements.txt @@ -6,9 +6,9 @@ connected-components-3d>=3.10.1 dbscan DracoPy>=1.0.1,<2.0.0 fastremap>=1.17.4 -fastmorph>=1.6.0 +fastmorph>=1.7.0 google-cloud-logging -kimimaro[accel]>=5.4.0 +kimimaro[accel]>=5.7.0 mapbuffer>=0.7.1 networkx numpy>=1.14.2