def test_copygrayscale_from_hdf5_to_dvid_multiscale(setup_hdf5_grayscale_input, disable_auto_retry): _template_dir, config, _volume, dvid_address, repo_uuid, output_grayscale_name = setup_hdf5_grayscale_input # Modify the config from above to compute pyramid scales, # and choose a bounding box that is aligned with the bricks even at scale 2 # (just for easier testing). config["input"]["geometry"]["bounding-box"] = [[0, 0, 0], [128, 256, 128]] config["copygrayscale"]["min-pyramid-scale"] = 1 config["copygrayscale"]["max-pyramid-scale"] = 2 config["copygrayscale"][ "downsample-method"] = "label" # This test is easier to write if we use this downsampling method box_zyx, scale_0_vol = _run_to_dvid(setup_hdf5_grayscale_input, check_scale_0=False) scale_1_vol = downsample_labels(scale_0_vol, 2, True) scale_2_vol = downsample_labels(scale_1_vol, 2, True) # Check the other scales -- be careful to extract exactly one brick. output_vol = fetch_raw(dvid_address, repo_uuid, output_grayscale_name + '_1', box_zyx // 2) assert (output_vol == scale_1_vol).all(), \ "Scale 1: Written vol does not match expected" # Check the other scales -- be careful to extract exactly one brick. output_vol = fetch_raw(dvid_address, repo_uuid, output_grayscale_name + '_2', box_zyx // 4) assert (output_vol == scale_2_vol).all(), \ "Scale 2: Written vol does not match expected"
def _run_to_dvid(setup, check_scale_0=True): template_dir, config, volume, dvid_address, repo_uuid, output_segmentation_name = setup yaml = YAML() yaml.default_flow_style = False # re-dump config in case it's been changed by a specific test with open(f"{template_dir}/workflow.yaml", 'w') as f: yaml.dump(config, f) _execution_dir, workflow = launch_flow(template_dir, 1) final_config = workflow.config input_box_xyz = np.array(final_config['input']['geometry']['bounding-box']) input_box_zyx = input_box_xyz[:, ::-1] expected_vol = extract_subvol(volume, input_box_zyx) output_box_xyz = np.array( final_config['output']['geometry']['bounding-box']) output_box_zyx = output_box_xyz[:, ::-1] output_vol = fetch_raw(dvid_address, repo_uuid, output_segmentation_name, output_box_zyx, dtype=np.uint64) np.save('/tmp/output_vol.npy', output_vol) np.save('/tmp/expected_vol.npy', expected_vol) if check_scale_0: assert (output_vol == expected_vol).all(), \ "Written vol does not match expected" return input_box_zyx, expected_vol, output_vol
def copy_vnc_subvolume(box_zyx, copy_grayscale=True, copy_segmentation=True, chunk_shape=(64, 64, 2048)): assert not (box_zyx % 64).any(), \ "Only 64px block-aligned volumes can be copied." import numpy as np from neuclease.util import boxes_from_grid, tqdm_proxy, round_box from neuclease.dvid import find_master, fetch_raw, post_raw, fetch_subvol, post_labelmap_voxels vnc_master = ('emdata4:8200', find_master('emdata4:8200')) NUM_SCALES = 8 num_voxels = np.prod(box_zyx[1] - box_zyx[0]) if copy_grayscale: logger.info( f"Copying grayscale from box {box_zyx[:,::-1].tolist()} ({num_voxels/1e6:.1f} Mvox) for {NUM_SCALES} scales" ) for scale in tqdm_proxy(range(NUM_SCALES)): if scale == 0: input_name = 'grayscalejpeg' output_name = 'local-grayscalejpeg' else: input_name = f'grayscalejpeg_{scale}' output_name = f'local-grayscalejpeg_{scale}' scaled_box_zyx = np.maximum(box_zyx // 2**scale, 1) scaled_box_zyx = round_box(scaled_box_zyx, 64, 'out') for chunk_box in tqdm_proxy(boxes_from_grid(scaled_box_zyx, chunk_shape, clipped=True), leave=False): chunk = fetch_subvol(*vnc_master, input_name, chunk_box, progress=False) post_raw(*vnc_master, output_name, chunk_box[0], chunk) if copy_segmentation: logger.info( f"Copying segmentation from box {box_zyx[:,::-1].tolist()} ({num_voxels/1e6:.2f} Mvox)" ) for chunk_box in tqdm_proxy( boxes_from_grid(box_zyx, chunk_shape, clipped=True)): chunk = fetch_raw(*vnc_master, 'segmentation', chunk_box, dtype=np.uint64) post_labelmap_voxels(*vnc_master, 'local-segmentation', chunk_box[0], chunk, downres=True) # TODO: Update label indexes? logger.info("DONE")
def _run_to_dvid(setup, check_scale_0=True): template_dir, config, volume, dvid_address, repo_uuid, output_grayscale_name = setup yaml = YAML() yaml.default_flow_style = False # re-dump config in case it's been changed by a specific test with open(f"{template_dir}/workflow.yaml", 'w') as f: yaml.dump(config, f) _execution_dir, workflow = launch_flow(template_dir, 1) final_config = workflow.config box_xyz = np.array(final_config['input']['geometry']['bounding-box']) box_zyx = box_xyz[:, ::-1] output_vol = fetch_raw(dvid_address, repo_uuid, output_grayscale_name, box_zyx) expected_vol = volume[box_to_slicing(*box_zyx)] if check_scale_0: assert (output_vol == expected_vol).all(), \ "Written vol does not match expected" return box_zyx, expected_vol
def test_fetch_raw(labelmap_setup): dvid_server, dvid_repo, _merge_table_path, _mapping_path, supervoxel_vol = labelmap_setup instance_info = (dvid_server, dvid_repo, 'segmentation') # fetch_raw() returns mapped labels voxels = fetch_raw(*instance_info, [(0, 0, 0), supervoxel_vol.shape], dtype=np.uint64) assert (voxels == 1).all()
def download(bounding_box_zyx, output_path): shape = bounding_box_zyx[1] - bounding_box_zyx[0] with h5py.File(output_path, 'w') as f: gray_dset = f.create_dataset('grayscale', shape=shape, dtype=np.uint8, chunks=True) seg_dset = f.create_dataset('segmentation', shape=shape, dtype=np.uint64, chunks=True, compression='gzip') print("Downloading grayscale...") block_shape = (256, 256, 256) block_boxes = boxes_from_grid(bounding_box_zyx, block_shape, clipped=True) for block_box in tqdm(block_boxes): relative_box = block_box - bounding_box_zyx[0] block_gray = fetch_raw(*GRAYSCALE, block_box) overwrite_subvol(gray_dset, relative_box, block_gray) print("") print("Downloading segmentation...") block_boxes = boxes_from_grid(bounding_box_zyx, block_shape, clipped=True) for block_box in tqdm(block_boxes): relative_box = block_box - bounding_box_zyx[0] block_seg = fetch_labelmap_voxels(*SEGMENTATION, block_box) overwrite_subvol(seg_dset, relative_box, block_seg) print("") print("DONE")
def get_subvolume(self, box_zyx, scale=0): req_bytes = self._dtype_nbytes * np.prod(box_zyx[1] - box_zyx[0]) instance_name = self._instance_name if self._instance_type.endswith('blk') and scale > 0: # Grayscale multi-scale is achieved via multiple instances instance_name = f"{instance_name}_{scale}" scale = 0 try: if self._instance_type in ('labelarray', 'labelmap'): # Obtain permission from the resource manager while fetching the compressed data, # but release the resource token before inflating the data. with self._resource_manager_client.access_context( self._server, True, 1, req_bytes): aligned_box = round_box(box_zyx, 64, 'out') if 8 * np.prod(aligned_box[1] - aligned_box[0]) < 2**31: vol_proxy = fetch_labelmap_voxels( self._server, self._uuid, instance_name, box_zyx, scale, self._throttle, supervoxels=self.supervoxels, format='lazy-array') else: # Requested subvolume is too large to download in one request. # Download it in chunks, with a somewhat arbitrary chunkshape chunk_shape = (64, 128, 10240) vol_proxy = fetch_labelmap_voxels_chunkwise( self._server, self._uuid, instance_name, box_zyx, scale, self._throttle, supervoxels=self.supervoxels, format='lazy-array', chunk_shape=chunk_shape) # Inflate after releasing resource token return vol_proxy() else: with self._resource_manager_client.access_context( self._server, True, 1, req_bytes): return fetch_raw(self._server, self._uuid, instance_name, box_zyx, self._throttle) except Exception as ex: # In cluster scenarios, a chained 'raise ... from ex' traceback # doesn't get fully transmitted to the driver, # so we simply append this extra info to the current exception # rather than using exception chaining. # Also log it now so it at least appears in the worker log. # See: https://github.com/dask/dask/issues/4384 import traceback, io sio = io.StringIO() traceback.print_exc(file=sio) logger.log(logging.ERROR, sio.getvalue()) host = socket.gethostname() msg = f"Host {host}: Failed to fetch subvolume: box_zyx = {box_zyx.tolist()}" ex.args += (msg, ) raise