feat(mesh): exclude certain segids from meshing (#221)

* feat: add object exclusion logic to meshing

This can be used to avoid meshing a very large object that
needs to be specially treated.

* feat(cli): add --exclude-labels to mesh forging

Author: william-silversmith

igneous/task_creation/mesh.py

@@ -156,14 +156,28 @@ def __iter__(self):
 
 
 def create_meshing_tasks(
-    layer_path, mip, shape=(448, 448, 448), 
-    simplification=True, max_simplification_error=40,
-    mesh_dir=None, cdn_cache=False, dust_threshold=None,
-    object_ids=None, progress=False, fill_missing=False,
-    encoding='precomputed', spatial_index=True, frag_path=None, sharded=False,
-    compress='gzip', closed_dataset_edges=True, dust_global=False,
-    fill_holes=0, dry_run=False,
-  ):
+  layer_path:str, 
+  mip:int,
+  shape:tuple[int,int,int] = (448, 448, 448), 
+  simplification:bool = True,
+  max_simplification_error:float = 40.0,
+  mesh_dir:Optional[str] = None,
+  cdn_cache:bool = False,
+  dust_threshold:Optional[int] = None,
+  object_ids:Optional[list[int]] = None,
+  progress:bool = False,
+  fill_missing:bool = False,
+  encoding:str = 'precomputed',
+  spatial_index:bool = True,
+  frag_path:Optional[str] = None,
+  sharded:bool = False,
+  compress:Optional[str] = 'gzip',
+  closed_dataset_edges:bool = True,
+  dust_global:bool = False,
+  fill_holes:int = 0,
+  dry_run:bool = False,
+  exclude_object_ids:list[int] = [],
+):
   shape = Vec(*shape)
 
   assert 0 <= fill_holes <= 103, "fill_holes must be between 0 to 103 inclusive."
@@ -208,6 +222,7 @@ def task(self, shape, offset):
         dust_global=bool(dust_global),
         progress=progress,
         object_ids=object_ids,
+        exclude_object_ids=exclude_object_ids,
         fill_missing=fill_missing,
         encoding=encoding,
         spatial_index=spatial_index,
@@ -234,6 +249,7 @@ def on_finish(self):
           'dust_threshold': dust_threshold,
           'encoding': encoding,
           'object_ids': object_ids,
+          'exclude_object_ids': exclude_object_ids,
           'spatial_index': spatial_index,
           'frag_path': frag_path,
           'sharded': sharded,

igneous/tasks/mesh/mesh.py

@@ -76,6 +76,7 @@ def __init__(self, shape, offset, layer_path, **kwargs):
       draco_create_metadata: (bool: False) only applies to draco encoding
       progress: (bool: False) show progress bars for meshing 
       object_ids: (list of ints) if specified, only mesh these ids
+      exclude_object_ids: (list of ints) exclude any label present in this list from meshing
       fill_missing: (bool: False) replace missing segmentation files with zeros instead of erroring
       spatial_index: (bool: False) generate a JSON spatial index of which meshes are available in
         a given bounding box. 
@@ -112,6 +113,7 @@ def __init__(self, shape, offset, layer_path, **kwargs):
       'frag_path': kwargs.get('frag_path', None),
       'mip': kwargs.get('mip', 0),
       'object_ids': kwargs.get('object_ids', None),
+      'exclude_object_ids': kwargs.get('exclude_object_ids', []),
       'parallel_download': kwargs.get('parallel_download', 1),
       'progress': kwargs.get('progress', False),
       'remap_table': kwargs.get('remap_table', None),
@@ -198,6 +200,9 @@ def execute(self):
     if self.options['object_ids']:
       data = fastremap.mask_except(data, self.options['object_ids'], in_place=True)
 
+    if self.options['exclude_object_ids']:
+      data = fastremap.mask(data, self.options['exclude_object_ids'], in_place=True)
+
     data, renumbermap = fastremap.renumber(data, in_place=True)
     renumbermap = { v:k for k,v in renumbermap.items() }
 

igneous_cli/cli.py

@@ -1046,13 +1046,15 @@ def mesh_xfer(
 @click.option('--sharded', is_flag=True, default=False, help="Generate shard fragments instead of outputing mesh fragments.", show_default=True)
 @click.option('--closed-edge/--open-edge', is_flag=True, default=True, help="Whether meshes are closed on the side that contacts the dataset boundary.", show_default=True)
 @click.option('--labels', type=ListN(), default=None, help="Mesh only this comma separated list of labels.", show_default=True)
+@click.option('--exclude-labels', type=ListN(), default=[], help="Do not mesh labels in this comma separated list.", show_default=True)
 @click.pass_context
 def mesh_forge(
   ctx, path, queue, mip, shape, 
   simplify, fill_missing, max_error, 
   dust_threshold, dir, compress, 
   spatial_index, sharded, closed_edge,
   dust_global, labels, fill_holes,
+  exclude_labels,
 ):
   """
   (1) Synthesize meshes from segmentation cutouts.
@@ -1079,6 +1081,7 @@ def mesh_forge(
     encoding='precomputed', spatial_index=spatial_index, 
     sharded=sharded, compress=compress, closed_dataset_edges=closed_edge,
     dust_global=dust_global, fill_holes=fill_holes,
+    exclude_object_ids=exclude_labels,
   )
 
   enqueue_tasks(ctx, queue, tasks)

test/test_tasks.py

@@ -430,6 +430,36 @@ def test_mesh(compress):
     assert cf.get('mesh/10:0:0-64_0-64_0-64') is not None 
     assert list(cf.list('mesh/')) == ['mesh/10:0:0-64_0-64_0-64']
 
+def test_mesh_object_ids():
+    delete_layer()
+    cf, _ = create_layer(size=(64,64,64,1), offset=(0,0,0), layer_type="segmentation")
+    cv = CloudVolume(cf.cloudpath)
+    # create a box of ones surrounded by zeroes
+    data = np.zeros(shape=(64,64,64,1), dtype=np.uint32)
+    data[1:-1,1:-1,1:-1,:] = 1
+    data[1:-1,1:-1,32:63,:] = 2
+    cv[:] = data
+    cv.info['mesh'] = 'mesh'
+    cv.commit_info()
+
+    MeshTask(
+        shape=(64,64,64),
+        offset=(0,0,0),
+        layer_path=cf.cloudpath,
+        mip=0,
+        exclude_object_ids=[2],
+    ).execute()
+    assert cf.get('mesh/1:0:0-64_0-64_0-64') is not None 
+    assert list(cf.list('mesh/')) == ['mesh/1:0:0-64_0-64_0-64']
+
+    MeshTask(
+        shape=(64,64,64),
+        offset=(0,0,0),
+        layer_path=cf.cloudpath,
+        mip=0,
+        object_ids=[2],
+    ).execute()
+    assert cf.get('mesh/2:0:0-64_0-64_0-64') is not None 
 
 def test_quantize():
     qpath = 'file:///tmp/removeme/quantized/'