From 0d25454509f634cea88af7c48f521c2fe7f0b29a Mon Sep 17 00:00:00 2001 From: Norman Rzepka Date: Fri, 16 Aug 2019 17:22:55 +0200 Subject: [PATCH] Map futures (#120) * convert to map_to_futures * tests * tests * check black for all + remove testdata * black all --- .gitignore | 1 + setup.py | 12 ++- tests/scripts/check_formatting.sh | 2 +- tests/scripts/downsampling.sh | 3 +- tests/test_downsampling.py | 40 ++++--- tests/test_export_wkw_as_tiff.py | 171 ++++++++++++++++++++++-------- tests/test_utils.py | 49 +++++---- wkcuber/compress.py | 15 ++- wkcuber/convert_knossos.py | 14 ++- wkcuber/cubing.py | 30 +++--- wkcuber/downsampling.py | 28 ++--- wkcuber/tile_cubing.py | 12 +-- 12 files changed, 239 insertions(+), 138 deletions(-) diff --git a/.gitignore b/.gitignore index 9542db530..e820af145 100644 --- a/.gitignore +++ b/.gitignore @@ -96,3 +96,4 @@ ENV/ testoutput testdata/WT1_wkw +testdata/tiff_mag_2_reference diff --git a/setup.py b/setup.py index 00cb02f59..cdbf3e1fe 100644 --- a/setup.py +++ b/setup.py @@ -4,8 +4,16 @@ name="wkcuber", packages=find_packages(exclude=("tests",)), use_scm_version=True, - setup_requires=['setuptools_scm'], - install_requires=["scipy", "numpy", "pillow", "pyyaml", "wkw", "cluster_tools==1.36", "natsort"], + setup_requires=["setuptools_scm"], + install_requires=[ + "scipy", + "numpy", + "pillow", + "pyyaml", + "wkw", + "cluster_tools==1.36", + "natsort", + ], description="A cubing tool for webKnossos", author="Norman Rzepka", author_email="norman.rzepka@scalableminds.com", diff --git a/tests/scripts/check_formatting.sh b/tests/scripts/check_formatting.sh index 9cd77670f..7c3e451da 100755 --- a/tests/scripts/check_formatting.sh +++ b/tests/scripts/check_formatting.sh @@ -4,4 +4,4 @@ docker run \ --entrypoint "/bin/bash" \ -w "/" \ scalableminds/webknossos-cuber:${CIRCLE_BUILD_NUM} \ - -c "black --check /app/wkcuber" \ No newline at end of file + -c "black --check /app" \ No newline at end of file diff --git a/tests/scripts/downsampling.sh b/tests/scripts/downsampling.sh index ae1e474b2..57f580bb3 100755 --- a/tests/scripts/downsampling.sh +++ b/tests/scripts/downsampling.sh @@ -17,4 +17,5 @@ docker run \ [ $(find testoutput/tiff/color/2 -mindepth 3 -name "*.wkw" | wc -l) -eq 1 ] [ $(find testoutput/tiff/color/4 -mindepth 3 -name "*.wkw" | wc -l) -eq 1 ] [ $(find testoutput/tiff/color/8 -mindepth 3 -name "*.wkw" | wc -l) -eq 1 ] -cmp --silent testoutput/tiff/color/2/z0/y0/x0.wkw testdata/tiff_mag_2_reference/color/2/z0/y0/x0.wkw \ No newline at end of file +cmp --silent testoutput/tiff/color/2/z0/y0/x0.wkw testdata/tiff_mag_2_reference/color/2/z0/y0/x0.wkw +cmp --silent testoutput/tiff/color/2/header.wkw testdata/tiff_mag_2_reference/color/2/header.wkw \ No newline at end of file diff --git a/tests/test_downsampling.py b/tests/test_downsampling.py index 1de24a5c0..395883105 100644 --- a/tests/test_downsampling.py +++ b/tests/test_downsampling.py @@ -95,7 +95,7 @@ def downsample_test_helper(use_compress): )[0] assert np.any(source_buffer != 0) - downsample_cube_job( + downsample_args = ( source_info, target_info, (2, 2, 2), @@ -104,6 +104,7 @@ def downsample_test_helper(use_compress): offset, use_compress, ) + downsample_cube_job(downsample_args) assert np.any(source_buffer != 0) block_type = ( @@ -155,7 +156,7 @@ def test_downsample_multi_channel(): wkw_dataset.write(offset, source_data) assert np.any(source_data != 0) - downsample_cube_job( + downsample_args = ( source_info, target_info, (2, 2, 2), @@ -164,6 +165,7 @@ def test_downsample_multi_channel(): tuple(a * WKW_CUBE_SIZE for a in offset), False, ) + downsample_cube_job(downsample_args) channels = [] for channel_index in range(num_channels): @@ -184,21 +186,27 @@ def test_downsample_multi_channel(): assert np.all(target_buffer == joined_buffer) + def test_anisotropic_mag_calculation(): mag_tests = [ - [(10.5, 10.5, 24), Mag(1), Mag((2, 2, 1))], - [(10.5, 10.5, 21), Mag(1), Mag((2, 2, 1))], - [(10.5, 24, 10.5), Mag(1), Mag((2, 1, 2))], - [(24, 10.5, 10.5), Mag(1), Mag((1, 2, 2))], - [(10.5, 10.5, 10.5), Mag(1), Mag((2, 2, 2))], - [(10.5, 10.5, 24), Mag((2, 2, 1)), Mag((4, 4, 1))], - [(10.5, 10.5, 21), Mag((2, 2, 1)), Mag((4, 4, 2))], - [(10.5, 24, 10.5), Mag((2, 1, 2)), Mag((4, 1, 4))], - [(24, 10.5, 10.5), Mag((1, 2, 2)), Mag((1, 4, 4))], - [(10.5, 10.5, 10.5), Mag(2), Mag(4)], + [(10.5, 10.5, 24), Mag(1), Mag((2, 2, 1))], + [(10.5, 10.5, 21), Mag(1), Mag((2, 2, 1))], + [(10.5, 24, 10.5), Mag(1), Mag((2, 1, 2))], + [(24, 10.5, 10.5), Mag(1), Mag((1, 2, 2))], + [(10.5, 10.5, 10.5), Mag(1), Mag((2, 2, 2))], + [(10.5, 10.5, 24), Mag((2, 2, 1)), Mag((4, 4, 1))], + [(10.5, 10.5, 21), Mag((2, 2, 1)), Mag((4, 4, 2))], + [(10.5, 24, 10.5), Mag((2, 1, 2)), Mag((4, 1, 4))], + [(24, 10.5, 10.5), Mag((1, 2, 2)), Mag((1, 4, 4))], + [(10.5, 10.5, 10.5), Mag(2), Mag(4)], ] for i in range(len(mag_tests)): - assert mag_tests[i][2] == get_next_anisotropic_mag(mag_tests[i][1], mag_tests[i][0]), "The next anisotropic" \ - " Magnification of {} with " \ - "the size {} should be {}"\ - .format(mag_tests[i][1], mag_tests[i][0], mag_tests[i][2]) + assert mag_tests[i][2] == get_next_anisotropic_mag( + mag_tests[i][1], mag_tests[i][0] + ), ( + "The next anisotropic" + " Magnification of {} with " + "the size {} should be {}".format( + mag_tests[i][1], mag_tests[i][0], mag_tests[i][2] + ) + ) diff --git a/tests/test_export_wkw_as_tiff.py b/tests/test_export_wkw_as_tiff.py index a9ca9dd4b..abc17efb1 100644 --- a/tests/test_export_wkw_as_tiff.py +++ b/tests/test_export_wkw_as_tiff.py @@ -9,53 +9,87 @@ def test_export_tiff_stack(): destination_path = os.path.join("testoutput", "WT1_wkw") - args_list = ["--source_path", os.path.join("testdata", "WT1_wkw"), - "--destination_path", destination_path, - "--layer_name", "color", - "--name", "test_export", - "--bbox", "0,0,0,100,100,5", - "--mag", "1"] + args_list = [ + "--source_path", + os.path.join("testdata", "WT1_wkw"), + "--destination_path", + destination_path, + "--layer_name", + "color", + "--name", + "test_export", + "--bbox", + "0,0,0,100,100,5", + "--mag", + "1", + ] bbox = {"topleft": [0, 0, 0], "size": [100, 100, 5]} run(args_list) - test_wkw_file_path = os.path.join("testdata", "WT1_wkw", "color", Mag(1).to_layer_name()) + test_wkw_file_path = os.path.join( + "testdata", "WT1_wkw", "color", Mag(1).to_layer_name() + ) with wkw.Dataset.open(test_wkw_file_path) as dataset: slice_bbox = bbox slice_bbox["size"] = [slice_bbox["size"][0], slice_bbox["size"][1], 1] for data_slice_index in range(1, bbox["size"][2] + 1): - slice_bbox["topleft"] = [slice_bbox["topleft"][0], slice_bbox["topleft"][1], bbox["topleft"][2] + data_slice_index] - tiff_path = os.path.join(destination_path, wkw_name_and_bbox_to_tiff_name("test_export", data_slice_index)) - - assert os.path.isfile(tiff_path), f"Expected a tiff to be written at: {tiff_path}." + slice_bbox["topleft"] = [ + slice_bbox["topleft"][0], + slice_bbox["topleft"][1], + bbox["topleft"][2] + data_slice_index, + ] + tiff_path = os.path.join( + destination_path, + wkw_name_and_bbox_to_tiff_name("test_export", data_slice_index), + ) + + assert os.path.isfile( + tiff_path + ), f"Expected a tiff to be written at: {tiff_path}." test_image = np.array(Image.open(tiff_path)) test_image.transpose((1, 0)) - correct_image = dataset.read(off=slice_bbox["topleft"], shape=slice_bbox["size"]) + correct_image = dataset.read( + off=slice_bbox["topleft"], shape=slice_bbox["size"] + ) correct_image = np.squeeze(correct_image) - assert np.array_equal(correct_image, test_image), f"The tiff file {tiff_path} that was written is not " \ - f"equal to the original wkw_file." + assert np.array_equal(correct_image, test_image), ( + f"The tiff file {tiff_path} that was written is not " + f"equal to the original wkw_file." + ) def test_export_tiff_stack_tile_size(): destination_path = os.path.join("testoutput", "WT1_wkw_tile_size") - args_list = ["--source_path", os.path.join("testdata", "WT1_wkw"), - "--destination_path", destination_path, - "--layer_name", "color", - "--name", "test_export", - "--bbox", "0,0,0,100,100,5", - "--mag", "1", - "--tile_size", "30,30"] + args_list = [ + "--source_path", + os.path.join("testdata", "WT1_wkw"), + "--destination_path", + destination_path, + "--layer_name", + "color", + "--name", + "test_export", + "--bbox", + "0,0,0,100,100,5", + "--mag", + "1", + "--tile_size", + "30,30", + ] bbox = {"topleft": [0, 0, 0], "size": [100, 100, 5]} run(args_list) tile_bbox = {"topleft": bbox["topleft"], "size": [30, 30, 1]} - test_wkw_file_path = os.path.join("testdata", "WT1_wkw", "color", Mag(1).to_layer_name()) + test_wkw_file_path = os.path.join( + "testdata", "WT1_wkw", "color", Mag(1).to_layer_name() + ) with wkw.Dataset.open(test_wkw_file_path) as dataset: slice_bbox = {"topleft": bbox["topleft"], "size": bbox["size"]} slice_bbox["size"] = [slice_bbox["size"][0], slice_bbox["size"][1], 1] @@ -63,64 +97,107 @@ def test_export_tiff_stack_tile_size(): for y_tile_index in range(ceil(bbox["size"][1] / tile_bbox["size"][1])): for x_tile_index in range(ceil(bbox["size"][0] / tile_bbox["size"][0])): - tiff_path = os.path.join(destination_path, f"{data_slice_index + 1}", f"{y_tile_index + 1}", f"{x_tile_index + 1}.tiff") + tiff_path = os.path.join( + destination_path, + f"{data_slice_index + 1}", + f"{y_tile_index + 1}", + f"{x_tile_index + 1}.tiff", + ) - assert os.path.isfile(tiff_path), f"Expected a tiff to be written at: {tiff_path}." + assert os.path.isfile( + tiff_path + ), f"Expected a tiff to be written at: {tiff_path}." test_image = np.array(Image.open(tiff_path)) test_image.transpose((1, 0)) - correct_image = dataset.read(off=[tile_bbox["topleft"][0] + tile_bbox["size"][0] * x_tile_index, - tile_bbox["topleft"][1] + tile_bbox["size"][1] * y_tile_index, - tile_bbox["topleft"][2] + data_slice_index], - shape=tile_bbox["size"]) + correct_image = dataset.read( + off=[ + tile_bbox["topleft"][0] + + tile_bbox["size"][0] * x_tile_index, + tile_bbox["topleft"][1] + + tile_bbox["size"][1] * y_tile_index, + tile_bbox["topleft"][2] + data_slice_index, + ], + shape=tile_bbox["size"], + ) correct_image = np.squeeze(correct_image) - assert np.array_equal(correct_image, test_image), f"The tiff file {tiff_path} that was written " \ + assert np.array_equal(correct_image, test_image), ( + f"The tiff file {tiff_path} that was written " f"is not equal to the original wkw_file." + ) def test_export_tiff_stack_tiles_per_dimension(): destination_path = os.path.join("testoutput", "WT1_wkw_tiles_per_dimension") - args_list = ["--source_path", os.path.join("testdata", "WT1_wkw"), - "--destination_path", destination_path, - "--layer_name", "color", - "--name", "test_export", - "--bbox", "0,0,0,100,100,5", - "--mag", "1", - "--tiles_per_dimension", "6,6"] + args_list = [ + "--source_path", + os.path.join("testdata", "WT1_wkw"), + "--destination_path", + destination_path, + "--layer_name", + "color", + "--name", + "test_export", + "--bbox", + "0,0,0,100,100,5", + "--mag", + "1", + "--tiles_per_dimension", + "6,6", + ] bbox = {"topleft": [0, 0, 0], "size": [100, 100, 5]} run(args_list) tile_bbox = {"topleft": bbox["topleft"], "size": [17, 17, 1]} - test_wkw_file_path = os.path.join("testdata", "WT1_wkw", "color", Mag(1).to_layer_name()) + test_wkw_file_path = os.path.join( + "testdata", "WT1_wkw", "color", Mag(1).to_layer_name() + ) with wkw.Dataset.open(test_wkw_file_path) as dataset: slice_bbox = bbox slice_bbox["size"] = [slice_bbox["size"][0], slice_bbox["size"][1], 1] for data_slice_index in range(bbox["size"][2]): for y_tile_index in range(ceil(bbox["size"][1] / tile_bbox["size"][1])): - for x_tile_index in range(ceil(tile_bbox["size"][0] / tile_bbox["size"][0])): - tiff_path = os.path.join(destination_path, f"{data_slice_index + 1}", f"{y_tile_index + 1}", f"{x_tile_index + 1}.tiff") - - assert os.path.isfile(tiff_path), f"Expected a tiff to be written at: {tiff_path}." + for x_tile_index in range( + ceil(tile_bbox["size"][0] / tile_bbox["size"][0]) + ): + tiff_path = os.path.join( + destination_path, + f"{data_slice_index + 1}", + f"{y_tile_index + 1}", + f"{x_tile_index + 1}.tiff", + ) + + assert os.path.isfile( + tiff_path + ), f"Expected a tiff to be written at: {tiff_path}." test_image = np.array(Image.open(tiff_path)) test_image.transpose((1, 0)) - correct_image = dataset.read(off=[tile_bbox["topleft"][0] + tile_bbox["size"][0] * x_tile_index, - tile_bbox["topleft"][1] + tile_bbox["size"][1] * y_tile_index, - tile_bbox["topleft"][2] + data_slice_index], - shape=tile_bbox["size"]) + correct_image = dataset.read( + off=[ + tile_bbox["topleft"][0] + + tile_bbox["size"][0] * x_tile_index, + tile_bbox["topleft"][1] + + tile_bbox["size"][1] * y_tile_index, + tile_bbox["topleft"][2] + data_slice_index, + ], + shape=tile_bbox["size"], + ) correct_image = np.squeeze(correct_image) - assert np.array_equal(correct_image, test_image), f"The tiff file {tiff_path} that was written " \ + assert np.array_equal(correct_image, test_image), ( + f"The tiff file {tiff_path} that was written " f"is not equal to the original wkw_file." + ) if __name__ == "__main__": - test_export_tiff_stack_tile_size() \ No newline at end of file + test_export_tiff_stack_tile_size() diff --git a/tests/test_utils.py b/tests/test_utils.py index 00290f241..e54a64d76 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -35,22 +35,27 @@ def test_get_regular_chunks_max_inclusive(): def test_buffered_slice_writer(): test_img = np.arange(24 * 24).reshape(24, 24).astype(np.uint16) + 1 dtype = test_img.dtype - bbox = {'topleft': (0, 0, 0), 'size': (24, 24, 35)} + bbox = {"topleft": (0, 0, 0), "size": (24, 24, 35)} origin = [0, 0, 0] - dataset_dir = 'testoutput/buffered_slice_writer' - layer_name = 'color' + dataset_dir = "testoutput/buffered_slice_writer" + layer_name = "color" mag = Mag(1) dataset_path = os.path.join(dataset_dir, layer_name, mag.to_layer_name()) - with BufferedSliceWriter(dataset_dir, layer_name, dtype, bbox, origin, mag=mag) as writer: + with BufferedSliceWriter( + dataset_dir, layer_name, dtype, bbox, origin, mag=mag + ) as writer: for i in range(13): writer.write_slice(i, test_img) with wkw.Dataset.open(dataset_path, wkw.Header(dtype)) as data: try: read_data = data.read(origin, (24, 24, 13)) if read_data[read_data.nonzero()].size != 0: - raise AssertionError('Nothing should be written on the disk. But found data with shape: {}' - .format(read_data.shape)) + raise AssertionError( + "Nothing should be written on the disk. But found data with shape: {}".format( + read_data.shape + ) + ) except wkw.wkw.WKWException: pass @@ -58,11 +63,13 @@ def test_buffered_slice_writer(): writer.write_slice(i, test_img) with wkw.Dataset.open(dataset_path, wkw.Header(dtype)) as data: read_data = data.read(origin, (24, 24, 32)) - assert np.squeeze(read_data).shape == (24, 24, 32), "The read data should have the shape: (24, 24, 32) " \ - "but has a shape of: {}"\ - .format(np.squeeze(read_data).shape) - assert read_data.size == read_data[read_data.nonzero()].size, "The read data contains zeros while the " \ - "written image has no zeros" + assert np.squeeze(read_data).shape == (24, 24, 32), ( + "The read data should have the shape: (24, 24, 32) " + "but has a shape of: {}".format(np.squeeze(read_data).shape) + ) + assert read_data.size == read_data[read_data.nonzero()].size, ( + "The read data contains zeros while the " "written image has no zeros" + ) for i in range(32, 35): writer.write_slice(i, test_img) @@ -70,18 +77,20 @@ def test_buffered_slice_writer(): with wkw.Dataset.open(dataset_path, wkw.Header(dtype)) as data: read_data = data.read(origin, (24, 24, 35)) read_data = np.squeeze(read_data) - assert read_data.shape == (24, 24, 35), "The read data should have the shape: (24, 24, 35) " \ - "but has a shape of: {}"\ - .format(np.squeeze(read_data).shape) - assert read_data.size == read_data[read_data.nonzero()].size, "The read data contains zeros while the " \ - "written image has no zeros" + assert read_data.shape == (24, 24, 35), ( + "The read data should have the shape: (24, 24, 35) " + "but has a shape of: {}".format(np.squeeze(read_data).shape) + ) + assert read_data.size == read_data[read_data.nonzero()].size, ( + "The read data contains zeros while the " "written image has no zeros" + ) test_img_3d = np.zeros((test_img.shape[0], test_img.shape[1], 35)) for i in np.arange(35): test_img_3d[:, :, i] = test_img # transpose because the slice writer takes [y, x] data and transposes it to [x, y] before writing test_img_3d = np.transpose(test_img_3d, (1, 0, 2)) # check if the data are correct - assert np.array_equal(test_img_3d, read_data), "The data from the disk is not the same " \ - "as the data that should be written." - - + assert np.array_equal(test_img_3d, read_data), ( + "The data from the disk is not the same " + "as the data that should be written." + ) diff --git a/wkcuber/compress.py b/wkcuber/compress.py index c6c131733..b17937271 100644 --- a/wkcuber/compress.py +++ b/wkcuber/compress.py @@ -52,7 +52,8 @@ def create_parser(): return parser -def compress_file_job(source_path, target_path): +def compress_file_job(args): + source_path, target_path = args try: logging.debug("Compressing '{}' to '{}'".format(source_path, target_path)) ref_time = time.time() @@ -89,16 +90,14 @@ def compress_mag(source_path, layer_name, target_path, mag: Mag, args=None): with open_wkw(source_wkw_info) as source_wkw: source_wkw.compress(target_mag_path) with get_executor_for_args(args) as executor: - futures = [] + job_args = [] for file in source_wkw.list_files(): rel_file = path.relpath(file, source_wkw.root) - futures.append( - executor.submit( - compress_file_job, file, path.join(target_mag_path, rel_file) - ) - ) + job_args.append((file, path.join(target_mag_path, rel_file))) - wait_and_ensure_success(futures) + wait_and_ensure_success( + executor.map_to_futures(compress_file_job, job_args) + ) logging.info("Mag {0} successfully compressed".format(str(mag))) diff --git a/wkcuber/convert_knossos.py b/wkcuber/convert_knossos.py index 8934db5d3..41ca6e54c 100644 --- a/wkcuber/convert_knossos.py +++ b/wkcuber/convert_knossos.py @@ -53,7 +53,8 @@ def create_parser(): return parser -def convert_cube_job(cube_xyz, source_knossos_info, target_wkw_info): +def convert_cube_job(args): + cube_xyz, source_knossos_info, target_wkw_info = args logging.info("Converting {},{},{}".format(cube_xyz[0], cube_xyz[1], cube_xyz[2])) ref_time = time.time() offset = tuple(x * CUBE_EDGE_LEN for x in cube_xyz) @@ -87,14 +88,11 @@ def convert_knossos( exit(1) knossos_cubes.sort() - futures = [] + job_args = [] for cube_xyz in knossos_cubes: - futures.append( - executor.submit( - convert_cube_job, cube_xyz, source_knossos_info, target_wkw_info - ) - ) - wait_and_ensure_success(futures) + job_args.append((cube_xyz, source_knossos_info, target_wkw_info)) + + wait_and_ensure_success(executor.map_to_futures(convert_cube_job, job_args)) if __name__ == "__main__": diff --git a/wkcuber/cubing.py b/wkcuber/cubing.py index d19f561dc..db65f3b8b 100644 --- a/wkcuber/cubing.py +++ b/wkcuber/cubing.py @@ -100,15 +100,16 @@ def prepare_slices_for_wkw(slices, num_channels=None): return buffer -def cubing_job( - target_wkw_info, - z_batches, - source_file_batches, - batch_size, - image_size, - num_channels, - pad=False, -): +def cubing_job(args): + ( + target_wkw_info, + z_batches, + source_file_batches, + batch_size, + image_size, + num_channels, + pad, + ) = args if len(z_batches) == 0: return @@ -187,16 +188,15 @@ def cubing(source_path, target_path, layer_name, dtype, batch_size, args=None) - ensure_wkw(target_wkw_info, num_channels=num_channels) with get_executor_for_args(args) as executor: - futures = [] + job_args = [] # We iterate over all z sections for z in range(0, num_z, BLOCK_LEN): # Prepare z batches max_z = min(num_z, z + BLOCK_LEN) z_batch = list(range(z, max_z)) - # Execute - futures.append( - executor.submit( - cubing_job, + # Prepare job + job_args.append( + ( target_wkw_info, z_batch, source_files[z:max_z], @@ -207,7 +207,7 @@ def cubing(source_path, target_path, layer_name, dtype, batch_size, args=None) - ) ) - wait_and_ensure_success(futures) + wait_and_ensure_success(executor.map_to_futures(cubing_job, job_args)) # Return Bounding Box return {"topLeft": [0, 0, 0], "width": num_x, "height": num_y, "depth": num_z} diff --git a/wkcuber/downsampling.py b/wkcuber/downsampling.py index 5e4fddf4e..5cf25095c 100644 --- a/wkcuber/downsampling.py +++ b/wkcuber/downsampling.py @@ -177,11 +177,10 @@ def downsample( ) with get_executor_for_args(args) as executor: - futures = [] + job_args = [] for target_cube_xyz in target_cube_addresses: - futures.append( - executor.submit( - downsample_cube_job, + job_args.append( + ( source_wkw_info, target_wkw_info, mag_factors, @@ -191,20 +190,21 @@ def downsample( compress, ) ) - wait_and_ensure_success(futures) + wait_and_ensure_success(executor.map_to_futures(downsample_cube_job, job_args)) logging.info("Mag {0} successfully cubed".format(target_mag)) -def downsample_cube_job( - source_wkw_info, - target_wkw_info, - mag_factors, - interpolation_mode, - cube_edge_len, - target_cube_xyz, - compress, -): +def downsample_cube_job(args): + ( + source_wkw_info, + target_wkw_info, + mag_factors, + interpolation_mode, + cube_edge_len, + target_cube_xyz, + compress, + ) = args logging.info("Downsampling of {}".format(target_cube_xyz)) try: diff --git a/wkcuber/tile_cubing.py b/wkcuber/tile_cubing.py index 20eae458a..ef9f69410 100644 --- a/wkcuber/tile_cubing.py +++ b/wkcuber/tile_cubing.py @@ -52,7 +52,8 @@ def find_source_files(source_section_path): return natsorted(all_source_files) -def tile_cubing_job(target_wkw_info, z_batches, source_path, batch_size, tile_size): +def tile_cubing_job(args): + target_wkw_info, z_batches, source_path, batch_size, tile_size = args if len(z_batches) == 0: return @@ -154,12 +155,11 @@ def tile_cubing( target_wkw_info = WkwDatasetInfo(target_path, layer_name, dtype, 1) ensure_wkw(target_wkw_info) with get_executor_for_args(args) as executor: - futures = [] + job_args = [] # Iterate over all z batches for z_batch in get_regular_chunks(min_z, max_z, BLOCK_LEN): - futures.append( - executor.submit( - tile_cubing_job, + job_args.append( + ( target_wkw_info, list(z_batch), source_path, @@ -167,7 +167,7 @@ def tile_cubing( tile_size, ) ) - wait_and_ensure_success(futures) + wait_and_ensure_success(executor.map_to_futures(tile_cubing_job, job_args)) if __name__ == "__main__":