feat: add additional output option for voxel counting

igneous/task_creation/image.py

@@ -1762,7 +1762,8 @@ def on_finish(self):
 
 def accumulate_voxel_counts(
   cloudpath, mip, 
-  progress=True, compress=None
+  progress=True, compress=None,
+  additional_output:Optional[str] = None,
 ):
   """
   Accumulate counts from each task.
@@ -1801,10 +1802,17 @@ def task(self, shape, offset):
   cf = CloudFiles(final_path)
   im = IntMap(final_counts)
 
+  binary = im.tobytes()
+  del im
   del final_counts
+
+  if additional_output:
+    with open(additional_output, "wb") as f:
+      f.write(binary)
+
   cf.put(
     'voxel_counts.im', 
-    im.tobytes(),
+    binary,
     content_type="application/x-intmap",
     compress=compress,
   )

igneous_cli/cli.py

@@ -509,15 +509,19 @@ def count_voxels(ctx, path, mip, queue):
 @click.argument("path", type=CloudPath())
 @click.option('--mip', default=0, help="Count this mip level of the image pyramid.", show_default=True)
 @click.option('--compress', default="zstd", help="What compression algorithm to apply. These files can be pretty big and must be downloaded by workers.", show_default=True)
+@click.option('-o', '--output', default=None, help="Also output the result as an IntMap file locally at this path. This is an additional output to avoid needing to re-download the result.", show_default=True)
 @click.pass_context
-def sum_voxel_counts(ctx, path, mip, compress):
+def sum_voxel_counts(ctx, path, mip, compress, output):
   """Accumulate counts from each task.
 
   Results are saved in a mapbuffer IntMap file:
   $cloudpath/$KEY/stats/voxel_counts.im
   """
   tc.accumulate_voxel_counts(
-    path, mip=mip, compress=compress
+    path, 
+    mip=mip, 
+    compress=compress, 
+    additional_output=output,
   )
 
 @imagegroup.group("contrast")

test/test_tasks.py

@@ -555,7 +555,13 @@ def test_voxel_counting_task():
     tq = MockTaskQueue()
     tasks = tc.create_voxel_counting_tasks(layer_path, mip=0)
     tq.insert_all(tasks)
-    tc.accumulate_voxel_counts(layer_path, mip=0)
+
+    additional_path = os.path.join(directory, "voxel_counts_additional.im")
+
+    tc.accumulate_voxel_counts(
+        layer_path, mip=0,
+        additional_output=additional_path,
+    )
 
     from mapbuffer import IntMap
     im = CloudFiles(layer_path).get(f"{cv.key}/stats/voxel_counts.im")
@@ -567,6 +573,13 @@ def test_voxel_counting_task():
 
     assert cts_dict_task == cts_dict_gt
 
+    with open(additional_path, "rb") as f:
+        im = IntMap(f)
+
+    cts_dict_task = im.todict()
+    assert cts_dict_task == cts_dict_gt
+
+    os.remove(additional_path)
 
 def test_num_mips_from_memory_target():
     from igneous.task_creation.image import num_mips_from_memory_target