Merge 2c2ac30 into 62ce90b

Dockerfile

@@ -1,5 +1,5 @@
 # backend: base | pytorch | pznet | pytorch-cuda9
-ARG BACKEND=pytorch
+ARG BACKEND=base
 
 FROM seunglab/chunkflow:${BACKEND}
 

chunkflow/flow/flow.py

@@ -3,6 +3,7 @@
 import sys
 from functools import update_wrapper, wraps
 from time import time
+
 import numpy as np
 import click
 
@@ -988,7 +989,7 @@ def copy_var(tasks, name, from_name, to_name):
 @click.option('--batch-size', '-b',
               type=int, default=1, help='mini batch size of input patch.')
 @click.option('--bump', type=click.Choice(['wu', 'zung']), default='wu',
-              help='bump function type. only works with pytorch-multitask backend.')
+              help='bump function type (only support wu now!).')
 @click.option('--mask-output-chunk/--no-mask-output-chunk', default=False,
               help='mask output chunk will make the whole chunk like one output patch. '
               + 'This will also work with non-aligned chunk size.')
@@ -1150,15 +1151,15 @@ def crop_margin(tasks, name, margin_size,
     type=int, default=None, help='do not mesh segments with voxel number less than threshold.')
 @click.option('--ids', type=str, default=None, 
               help='a list of segment ids to mesh. This is for sparse meshing. ' + 
-              'The ids should be separated by comma without space, such as "34,56,78,90"')
+              'The ids should be separated by comma without space, such as "34,56,78,90"' +
+              ' it can also be a json file contains a list of ids. The json file should be '+
+              'put in the output_path.')
 @click.option('--manifest/--no-manifest', default=False, help='create manifest file or not.')
 @operator
 def mesh(tasks, name, input_chunk_name, mip, voxel_size, output_path, output_format,
          simplification_factor, max_simplification_error, dust_threshold, 
          ids, manifest):
     """Perform meshing for segmentation chunk."""
-    if ids:
-        ids = set([int(id) for id in ids.split(',')])
     if mip is None:
         mip = state['mip']
 

chunkflow/flow/mesh.py

@@ -24,16 +24,16 @@ def __init__(self,
                  simplification_factor: int = 100,
                  max_simplification_error: int = 8,
                  dust_threshold: int = None,
-                 ids: set = None,
+                 ids: str = None,
                  manifest: bool = False,
                  name: str = 'meshing',
                  verbose: bool = True):
         """
         Parameters
         ------------
-        output_path: 
+        output_path:
             path to store mesh files
-        output_format: 
+        output_format:
             format of output {'ply', 'obj', 'precomputed'}
         voxel_size:
             size of voxels
@@ -62,8 +62,8 @@ def __init__(self,
         self.output_path = output_path
         self.output_format = output_format
         self.dust_threshold = dust_threshold
-        self.ids = ids
         self.manifest = manifest
+        self.ids = ids
 
         if manifest:
             assert output_format == 'precomputed'
@@ -86,6 +86,20 @@ def __init__(self,
             self.mesher = Mesher(voxel_size[::-1])
 
         self.storage = Storage(mesh_path)
+
+        if ids:
+            if ids.endswith('.json'):
+                # assume that ids is a json file in the storage path
+                json_storage = Storage(output_path)
+                ids_str = json_storage.get_file(ids)
+                self.ids = set(json.loads(ids_str))
+                assert len(self.ids) > 0
+                if self.verbose:
+                    print(f'number of selected objects: {len(self.ids)}')
+            else:
+                # a simple string, like "34,45,56,23"
+                # this is used for small object numbers
+                self.ids = set([int(id) for id in ids.split(',')])
 
     def _get_mesh_data(self, obj_id, offset):
         mesh = self.mesher.get_mesh(
@@ -116,8 +130,8 @@ def _remove_dust(self, seg: np.ndarray):
 
         if self.dust_threshold or self.ids:
             segids, voxel_nums = np.unique(seg, return_counts=True)
-            dust_segids = [sid for sid, ct in 
-                            zip(segids, voxel_nums) if ct < self.dust_threshold]
+            dust_segids = [sid for sid, ct in
+                           zip(segids, voxel_nums) if ct < self.dust_threshold]
 
             seg = fastremap.mask(seg, dust_segids, in_place=True)
         return seg