diff --git a/igneous/task_creation/image.py b/igneous/task_creation/image.py index 937f672..14bb036 100644 --- a/igneous/task_creation/image.py +++ b/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, ) diff --git a/igneous_cli/cli.py b/igneous_cli/cli.py index 9a2bda8..9fb9145 100644 --- a/igneous_cli/cli.py +++ b/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") diff --git a/test/test_tasks.py b/test/test_tasks.py index 30b8318..0f085ca 100755 --- a/test/test_tasks.py +++ b/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