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
