def compute_brick_index(cls, brick, wall_box, wall_grid):
"""
For a given brick in a BrickWall corresponding to wall_box and wall_grid,
compute a unique index based on the brick's logical_box.
(We just compute the scan-order index.)
Note that in the case of a sparsely populated brick wall,
brick indexes will not be consecutive.
"""
wall_box = round_box(wall_box, wall_grid.block_shape, 'out')
wall_shape = wall_box[1] - wall_box[0]
wall_index_box = np.array(
([0, 0, 0], wall_shape)) // wall_grid.block_shape
wall_index_shape = wall_index_box[1] - wall_index_box[0]
location_id = (brick.logical_box[0] -
wall_box[0]) // wall_grid.block_shape
def _scan_order_index(coord, shape):
# Example:
# coord = (z,y,x,c)
# shape = (Z,Y,X,C)
# strides = (Y*X*C, X*C, C, 1)
# index = z*Y*X*C + y*X*C + x*C + c*1
# Accumulate products in reverse
reverse_shape = tuple(shape[::-1]) # (C,X,Y,Z)
reverse_strides = np.multiply.accumulate(
(1, ) + reverse_shape[:-1]) # (1, C, C*X, C*X*Y)
strides = reverse_strides[::-1]
index = (coord * strides).sum()
return np.int32(index)
return _scan_order_index(location_id, wall_index_shape)
def main():
parser = argparse.ArgumentParser(
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument(
'--expand-for-alignment',
'-e',
action='store_true',
help='Auto-expand the given subvolume bounds to make it aligned.')
parser.add_argument('--grayscale',
'-g',
action='store_true',
help='Copy grayscale only')
parser.add_argument('--segmentation',
'-s',
action='store_true',
help='Copy segmentation only')
parser.add_argument(
'offset_xyz',
help="Starting offset for the subvolume, e.g. '25280, 42048, 48448'")
parser.add_argument('shape_xyz',
help="Shape of the subvolume, e.g. '512, 512, 512'")
args = parser.parse_args()
if not args.grayscale and not args.segmentation:
args.grayscale = True
args.segmentation = True
from neuclease import configure_default_logging
configure_default_logging()
import re
import numpy as np
from neuclease.util import round_box
args.offset_xyz = re.sub(r'\D', ' ', args.offset_xyz)
args.shape_xyz = re.sub(r'\D', ' ', args.shape_xyz)
offset_xyz = np.array([*map(int, args.offset_xyz.split())])
shape_xyz = np.array([*map(int, args.shape_xyz.split())])
box_xyz = np.array([offset_xyz, offset_xyz + shape_xyz])
box_zyx = box_xyz[:, ::-1]
del box_xyz
if args.expand_for_alignment:
box_zyx = round_box(box_zyx, 64, 'out')
shape_zyx = box_zyx[1] - box_zyx[0]
logger.info(
f"Expanded box to {box_zyx[:, ::-1].tolist()} (shape = {shape_zyx[::-1].tolist()})"
)
elif (box_zyx % 64).any():
raise RuntimeError(
"Only 64px block-aligned volumes can be copied.\n"
"Adjust your offset/shape or try the --expand-for-alignment option."
)
copy_vnc_subvolume(box_zyx, args.grayscale, args.segmentation)
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 _create_grayscale_instances(self, volume_config):
"""
Create the grayscale instance(s) in DVID for the given volume configuration.
In DVID, grayscale data is stored in instances of type 'uint8blk',
which has no concept of scale.
Instead, multi-scale volumes are represented by creating multiple instances,
with the scale indicated by a suffix (except for scale 0).
For example:
- grayscale # scale 0
- grayscale_1
- grayscale_2
- grayscale_3
- ...
"""
settings = volume_config["dvid"]["creation-settings"]
block_width = volume_config["geometry"]["block-width"]
pyramid_depth = settings["max-scale"]
if pyramid_depth == -1:
pyramid_depth = choose_pyramid_depth(self.bounding_box_zyx, 512)
repo_instances = fetch_repo_instances(self.server, self.uuid)
# Bottom level of pyramid is listed as neuroglancer-compatible
extend_list_value(self.server, self.uuid, '.meta', 'neuroglancer',
[self.instance_name])
for scale in range(pyramid_depth + 1):
scaled_output_box_zyx = round_box(self.bounding_box_zyx, 2**scale,
'out') // 2**scale
if scale == 0:
scaled_instance_name = self.instance_name
else:
scaled_instance_name = f"{self.instance_name}_{scale}"
if scaled_instance_name in repo_instances:
logger.info(
f"'{scaled_instance_name}' already exists, skipping creation"
)
else:
create_voxel_instance(
self.server, self.uuid, scaled_instance_name, 'uint8blk',
settings["versioned"], settings["compression"],
settings["tags"], block_width, settings["voxel-size"],
settings["voxel-units"], settings["background"])
update_extents(self.server, self.uuid, scaled_instance_name,
scaled_output_box_zyx)
# Higher-levels of the pyramid should not appear in the DVID console.
extend_list_value(self.server, self.uuid, '.meta', 'restrictions',
[scaled_instance_name])
def init_brickwall(self, volume_service, subset_labels, roi):
sbm = None
if roi["name"]:
base_service = volume_service.base_service
if not roi["server"] or not roi["uuid"]:
assert isinstance(base_service, DvidVolumeService), \
"Since you aren't using a DVID input source, you must specify the ROI server and uuid."
roi["server"] = (roi["server"] or volume_service.server)
roi["uuid"] = (roi["uuid"] or volume_service.uuid)
if roi["scale"] is not None:
scale = roi["scale"]
elif isinstance(volume_service, ScaledVolumeService):
scale = volume_service.scale_delta
assert scale <= 5, \
"The 'roi' option doesn't support volumes downscaled beyond level 5"
else:
scale = 0
brick_shape = volume_service.preferred_message_shape
assert not (brick_shape % 2**(5-scale)).any(), \
"If using an ROI, select a brick shape that is divisible by 32"
seg_box = volume_service.bounding_box_zyx
seg_box = round_box(seg_box, 2**(5-scale))
seg_box_s0 = seg_box * 2**scale
seg_box_s5 = seg_box // 2**(5-scale)
with Timer(f"Fetching mask for ROI '{roi['name']}' ({seg_box_s0[:, ::-1].tolist()})", logger):
roi_mask_s5, _ = fetch_roi(roi["server"], roi["uuid"], roi["name"], format='mask', mask_box=seg_box_s5)
# SBM 'full-res' corresponds to the input service voxels, not necessarily scale-0.
sbm = SparseBlockMask(roi_mask_s5, seg_box, 2**(5-scale))
elif subset_labels:
try:
sbm = volume_service.sparse_block_mask_for_labels([*subset_labels])
if ((sbm.box[1] - sbm.box[0]) == 0).any():
raise RuntimeError("Could not find sparse masks for any of the subset-labels")
except NotImplementedError:
sbm = None
with Timer("Initializing BrickWall", logger):
# Aim for 2 GB RDD partitions when loading segmentation
GB = 2**30
target_partition_size_voxels = 2 * GB // np.uint64().nbytes
# Apply halo WHILE downloading the data.
# TODO: Allow the user to configure whether or not the halo should
# be fetched from the outset, or added after the blocks are loaded.
halo = self.config["connectedcomponents"]["halo"]
brickwall = BrickWall.from_volume_service(volume_service, 0, None, self.client, target_partition_size_voxels, halo, sbm, compression='lz4_2x')
return brickwall
def init_boxes(self, volume_service, subset_labels, roi):
sbm = None
if roi:
base_service = volume_service.base_service
assert isinstance(base_service, DvidVolumeService), \
"Can't specify an ROI unless you're using a dvid input"
assert isinstance(volume_service, (ScaledVolumeService, DvidVolumeService)), \
"The 'roi' option doesn't support adapters other than 'rescale-level'"
scale = 0
if isinstance(volume_service, ScaledVolumeService):
scale = volume_service.scale_delta
assert scale <= 5, \
"The 'roi' option doesn't support volumes downscaled beyond level 5"
server, uuid, _seg_instance = base_service.instance_triple
brick_shape = volume_service.preferred_message_shape
assert not (brick_shape % 2**(5-scale)).any(), \
"If using an ROI, select a brick shape that is divisible by 32"
seg_box = volume_service.bounding_box_zyx
seg_box = round_box(seg_box, brick_shape)
seg_box_s0 = seg_box * 2**scale
seg_box_s5 = seg_box // 2**(5 - scale)
with Timer(
f"Fetching mask for ROI '{roi}' ({seg_box_s0[:, ::-1].tolist()})",
logger):
roi_mask_s5, _ = fetch_roi(server,
uuid,
roi,
format='mask',
mask_box=seg_box_s5)
# SBM 'full-res' corresponds to the input service voxels, not necessarily scale-0.
sbm = SparseBlockMask.create_from_highres_mask(
roi_mask_s5, 2**(5 - scale), seg_box, brick_shape)
elif subset_labels:
try:
sbm = volume_service.sparse_block_mask_for_labels(
[*subset_labels])
if ((sbm.box[1] - sbm.box[0]) == 0).any():
raise RuntimeError(
"Could not find sparse masks for any of the subset-labels"
)
except NotImplementedError:
sbm = None
if sbm is None:
boxes = boxes_from_grid(volume_service.bounding_box_zyx,
volume_service.preferred_message_shape,
clipped=True)
return np.array([*boxes])
else:
return sbm.sparse_boxes(brick_shape)
def init_boxes(self, volume_service, roi):
if not roi["name"]:
boxes = boxes_from_grid(volume_service.bounding_box_zyx,
volume_service.preferred_message_shape,
clipped=True)
return np.array([*boxes])
base_service = volume_service.base_service
if not roi["server"] or not roi["uuid"]:
assert isinstance(base_service, DvidVolumeService), \
"Since you aren't using a DVID input source, you must specify the ROI server and uuid."
roi["server"] = (roi["server"] or volume_service.server)
roi["uuid"] = (roi["uuid"] or volume_service.uuid)
if roi["scale"] is not None:
scale = roi["scale"]
elif isinstance(volume_service, ScaledVolumeService):
scale = volume_service.scale_delta
assert scale <= 5, \
"The 'roi' option doesn't support volumes downscaled beyond level 5"
else:
scale = 0
brick_shape = volume_service.preferred_message_shape
assert not (brick_shape % 2**(5-scale)).any(), \
"If using an ROI, select a brick shape that is divisible by 32"
seg_box = volume_service.bounding_box_zyx
seg_box = round_box(seg_box, 2**(5 - scale))
seg_box_s0 = seg_box * 2**scale
seg_box_s5 = seg_box // 2**(5 - scale)
with Timer(
f"Fetching mask for ROI '{roi['name']}' ({seg_box_s0[:, ::-1].tolist()})",
logger):
roi_mask_s5, _ = fetch_roi(roi["server"],
roi["uuid"],
roi["name"],
format='mask',
mask_box=seg_box_s5)
# SBM 'full-res' corresponds to the input service voxels, not necessarily scale-0.
sbm = SparseBlockMask(roi_mask_s5, seg_box, 2**(5 - scale))
boxes = sbm.sparse_boxes(brick_shape)
# Clip boxes to the true (not rounded) bounding box
boxes[:, 0] = np.maximum(boxes[:, 0],
volume_service.bounding_box_zyx[0])
boxes[:, 1] = np.minimum(boxes[:, 1],
volume_service.bounding_box_zyx[1])
return boxes
def init_boxes(self, volume_service, roi):
if not roi:
boxes = boxes_from_grid(volume_service.bounding_box_zyx,
volume_service.preferred_message_shape,
clipped=True)
return np.array([*boxes])
base_service = volume_service.base_service
assert isinstance(base_service, DvidVolumeService), \
"Can't specify an ROI unless you're using a dvid input"
assert isinstance(volume_service, (ScaledVolumeService, DvidVolumeService)), \
"The 'roi' option doesn't support adapters other than 'rescale-level'"
scale = 0
if isinstance(volume_service, ScaledVolumeService):
scale = volume_service.scale_delta
assert scale <= 5, \
"The 'roi' option doesn't support volumes downscaled beyond level 5"
server, uuid, _seg_instance = base_service.instance_triple
brick_shape = volume_service.preferred_message_shape
assert not (brick_shape % 2**(5-scale)).any(), \
"If using an ROI, select a brick shape that is divisible by 32"
seg_box = volume_service.bounding_box_zyx
seg_box = round_box(seg_box, 2**(5 - scale))
seg_box_s0 = seg_box * 2**scale
seg_box_s5 = seg_box // 2**(5 - scale)
with Timer(
f"Fetching mask for ROI '{roi}' ({seg_box_s0[:, ::-1].tolist()})",
logger):
roi_mask_s5, _ = fetch_roi(server,
uuid,
roi,
format='mask',
mask_box=seg_box_s5)
# SBM 'full-res' corresponds to the input service voxels, not necessarily scale-0.
sbm = SparseBlockMask(roi_mask_s5, seg_box, 2**(5 - scale))
boxes = sbm.sparse_boxes(brick_shape)
# Clip boxes to the true (not rounded) bounding box
boxes[:, 0] = np.maximum(boxes[:, 0],
volume_service.bounding_box_zyx[0])
boxes[:, 1] = np.minimum(boxes[:, 1],
volume_service.bounding_box_zyx[1])
return boxes
def init_boxes(self, volume_service, roi, chunk_shape_s0):
"""
Return a set of bounding boxes to tile the given ROI.
Scale 0 of the volume service should correspond to full-res data,
which is 32x higher-res than ROI resolution.
"""
if not roi["name"]:
boxes = boxes_from_grid(volume_service.bounding_box_zyx,
chunk_shape_s0,
clipped=True)
return np.array([*boxes])
base_service = volume_service.base_service
if not roi["server"] or not roi["uuid"]:
assert isinstance(base_service, DvidVolumeService), \
"Since you aren't using a DVID input source, you must specify the ROI server and uuid."
roi["server"] = (roi["server"] or volume_service.server)
roi["uuid"] = (roi["uuid"] or volume_service.uuid)
assert not (chunk_shape_s0 % 2**5).any(), \
"If using an ROI, select a chunk shape that is divisible by 32"
seg_box_s0 = volume_service.bounding_box_zyx
seg_box_s0 = round_box(seg_box_s0, 2**5)
seg_box_s5 = seg_box_s0 // 2**5
with Timer(
f"Fetching mask for ROI '{roi['name']}' ({seg_box_s0[:, ::-1].tolist()})",
logger):
roi_mask_s5, _ = fetch_roi(roi["server"],
roi["uuid"],
roi["name"],
format='mask',
mask_box=seg_box_s5)
# SBM 'full-res' corresponds to the input service voxels, not necessarily scale-0.
sbm = SparseBlockMask(roi_mask_s5, seg_box_s0, 2**5)
boxes = sbm.sparse_boxes(chunk_shape_s0)
# Clip boxes to the true (not rounded) bounding box
boxes[:, 0] = np.maximum(boxes[:, 0],
volume_service.bounding_box_zyx[0])
boxes[:, 1] = np.minimum(boxes[:, 1],
volume_service.bounding_box_zyx[1])
return boxes
def _create_segmentation_instance(self, volume_config):
"""
Create a segmentation volume in DVID according to the given configuration.
In DVID, segmentation instances are stored via a special instance type, 'labelmap',
which has several features, including built-in multiscale support, supervoxel-to-body
mapping, and sparse retrieval of body locations.
"""
if self.instance_name in fetch_repo_instances(self.server, self.uuid):
logger.info(
f"'{self.instance_name}' already exists, skipping creation")
return
settings = volume_config["dvid"]["creation-settings"]
block_width = volume_config["geometry"]["block-width"]
pyramid_depth = settings["max-scale"]
if pyramid_depth == -1:
pyramid_depth = choose_pyramid_depth(self.bounding_box_zyx, 512)
if settings["compression"] != DvidInstanceCreationSettingsSchema[
"properties"]["compression"]["default"]:
raise RuntimeError(
"Alternative compression methods are not permitted on labelmap instances. "
"Please remove the 'compression' setting from your config.")
if settings["background"] != 0:
raise RuntimeError(
"Labelmap instances do not support custom background values. "
"Please remove 'background' from your config.")
create_labelmap_instance(self.server, self.uuid, self.instance_name,
settings["versioned"], settings["tags"],
block_width, settings["voxel-size"],
settings["voxel-units"],
settings["enable-index"], pyramid_depth)
# Workaround for https://github.com/janelia-flyem/dvid/issues/344
# Write an empty block to the first and last blocks to set the
# bounding-box in DVID now, without any concurrency issues.
empty_block = np.zeros((64, 64, 64), np.uint64)
first_block_start, last_block_stop = round_box(self.bounding_box_zyx,
64, 'out')
last_block_start = last_block_stop - 64
self.write_subvolume(empty_block, first_block_start)
self.write_subvolume(empty_block, last_block_start)
def compute_brick_indexes(cls, logical_corners, wall_box, wall_grid):
"""
For an array of brick locations (logical_box[0]) in a BrickWall
corresponding to wall_box and wall_grid,
compute unique indexes for each brick based on the brick's logical_box.
(We just compute the scan-order index.)
Note that in the case of a sparsely populated brick wall,
brick indexes will not be consecutive.
"""
logical_corners = np.asarray(logical_corners)
assert logical_corners.ndim == 2
assert logical_corners.shape[1] == 3
wall_box = round_box(wall_box, wall_grid.block_shape, 'out')
wall_shape = wall_box[1] - wall_box[0]
wall_index_box = np.array(
([0, 0, 0], wall_shape)) // wall_grid.block_shape
wall_index_shape = wall_index_box[1] - wall_index_box[0]
location_ids = (logical_corners - wall_box[0]) // wall_grid.block_shape
def _scan_order_indexes(coords, shape):
# Example:
# coords = [(z,y,x,c),
# (z,y,x,c),
# ...]
#
# shape = (Z,Y,X,C)
# strides = (Y*X*C, X*C, C, 1)
# indexes = [z*Y*X*C + y*X*C + x*C + c*1,
# z*Y*X*C + y*X*C + x*C + c*1,
# ...]
# Accumulate products in reverse
reverse_shape = tuple(shape[::-1]) # (C,X,Y,Z)
reverse_strides = np.multiply.accumulate(
(1, ) + reverse_shape[:-1]) # (1, C, C*X, C*X*Y)
strides = reverse_strides[::-1]
indexes = (coords * strides).sum(axis=1)
return indexes.astype(np.int32)
return _scan_order_indexes(location_ids, wall_index_shape)
def export_sparsevol(server, uuid, instance, neurons_df, scale=5, format='tiff', output_dir='.'):
import os
import vigra
import numpy as np
from neuclease.util import round_box, tqdm_proxy
from neuclease.dvid import fetch_sparsevol, resolve_ref, fetch_volume_box, box_to_slicing
uuid = resolve_ref(server, uuid)
# Determine the segmentation bounding box at the given scale,
# which is used as the mask shape.
seg = (server, uuid, instance)
box = round_box(fetch_volume_box(*seg), 64, 'out')
box[0] = (0,0,0)
box_scaled = box // 2**scale
# How many digits will we need in each slice file name?
digits = int(np.ceil(np.log10(box_scaled[1, 0])))
# Export a mask stack for each group.
groups = neurons_df.groupby('group', sort=False)
num_groups = neurons_df['group'].nunique()
group_prog = tqdm_proxy(groups, total=num_groups)
for group, df in group_prog:
group_prog.write(f'Group "{group}": Assembling mask')
group_mask = np.zeros(box_scaled[1], dtype=bool)
group_mask = vigra.taggedView(group_mask, 'zyx')
# Overlay each body mask in the current group
for body in tqdm_proxy(df['body'], leave=False):
body_mask, mask_box = fetch_sparsevol(*seg, body, scale=scale, format='mask')
group_mask[box_to_slicing(*mask_box)] |= body_mask
# Write out the slice files
group_prog.write(f'Group "{group}": Writing slices')
d = f'{output_dir}/{group}.stack'
os.makedirs(d, exist_ok=True)
for z in tqdm_proxy(range(group_mask.shape[0]), leave=False):
p = ('{d}/{z:' + f'0{digits}' + 'd}.{f}').format(d=d, z=z, f=format)
vigra.impex.writeImage(group_mask[z].astype(np.uint8), p)
def process_point(seg_src, seg_dst, point, radius, src_body, dst_body):
"""
Generate a neighborhood segment around a particular point.
Upload the voxels for the segment and the corresponding mesh.
"""
r = radius
src_box = np.asarray((point - r, point + r + 1))
src_vol = fetch_labelmap_voxels(*seg_src, src_box)
if src_body is None:
src_body = src_vol[r, r, r]
if dst_body is None:
# Generate a neighborhood segment ID from the coordinate.
# Divide by 4 to ensure the coordinates fit within 2^53.
# (The segment ID will not retain the full resolution of
# the coordinate, but that's usually OK for our purposes.)
dst_body = encode_point_to_uint64(point // 4, 17)
mask = (src_vol == src_body) & sphere_mask(r)
dst_box = round_box(src_box, 64, 'out')
dst_vol = fetch_labelmap_voxels(*seg_dst, dst_box)
dst_view = dst_vol[b2s(*(src_box - dst_box[0]))]
dst_view[mask] = dst_body
post_labelmap_voxels(*seg_dst, dst_box[0], dst_vol, downres=True)
# Mesh needs to be written in nm, hence 8x
mesh = Mesh.from_binary_vol(mask, 8 * src_box, smoothing_rounds=2)
mesh.simplify(0.05, in_memory=True)
post_key(*seg_dst[:2], f'{seg_dst[2]}_meshes', f'{dst_body}.ngmesh',
mesh.serialize(fmt='ngmesh'))
centroid = src_box[0] + mask_centroid(mask, True)
top_z = mask.sum(axis=(1, 2)).nonzero()[0][0]
top_coords = np.transpose(mask[top_z].nonzero())
top_point = src_box[0] + (top_z, *top_coords[len(top_coords) // 2])
return point, centroid, top_point, src_body, dst_body, mask.sum()
def downsample_brick(brick):
assert (brick.logical_box % factor == 0).all()
# If the factor doesn't perfectly divide into
# the brick's physical dimensions,
# then chop off the edges until it does.
if (brick.physical_box % factor != 0).any():
clipped_box = round_box(brick.physical_box, factor, 'in')
volume = extract_subvol(brick.volume,
clipped_box - brick.physical_box[0])
else:
clipped_box = brick.physical_box
volume = brick.volume
downsampled_volume = downsample(volume, factor, method)
downsampled_logical_box = brick.logical_box // factor
downsampled_physical_box = clipped_box // factor
return Brick(downsampled_logical_box,
downsampled_physical_box,
downsampled_volume,
compression=brick.compression)
def init_boxes(self, volume_service, roi):
if not roi["name"]:
boxes = boxes_from_grid(volume_service.bounding_box_zyx,
volume_service.preferred_message_shape,
clipped=True)
return np.array([*boxes])
server, uuid, roi_name = roi["server"], roi["uuid"], roi["name"]
roi_scale = roi["relative-scale"]
brick_shape = volume_service.preferred_message_shape
assert not (brick_shape % 2**roi_scale).any(), \
"If using an ROI, select a brick shape that is divisible by 32"
seg_box = volume_service.bounding_box_zyx
seg_box = round_box(seg_box, 2**roi_scale)
seg_box_s5 = seg_box // 2**roi_scale
with Timer(
f"Fetching mask for ROI '{roi_name}' ({seg_box[:, ::-1].tolist()})",
logger):
roi_mask_s5, _ = fetch_roi(server,
uuid,
roi_name,
format='mask',
mask_box=seg_box_s5)
# SBM 'full-res' corresponds to the input service voxels, not necessarily scale-0.
sbm = SparseBlockMask(roi_mask_s5, seg_box, 2**roi_scale)
boxes = sbm.sparse_boxes(brick_shape)
# Clip boxes to the true (not rounded) bounding box
boxes[:, 0] = np.maximum(boxes[:, 0],
volume_service.bounding_box_zyx[0])
boxes[:, 1] = np.minimum(boxes[:, 1],
volume_service.bounding_box_zyx[1])
return boxes
def mitos_in_neighborhood(mito_roi_source, neighborhood_origin_xyz,
neighborhood_id, mito_res_scale_diff):
"""
Determine how many non-trivial mito objects overlap with the given "neighborhood object",
and return a table of their IDs and sizes.
1. Download the neighborhood mask for the given neighborhood_id.
2. Erode the neighborhood mask by 1 px (see note in the comment above).
3. Fetch the mito segmentation for the voxels within the neighborhood.
4. Fetch (from dvid) the sizes of each mito object.
5. Filter out the mitos that are smaller than the minimum size that is
actually used in our published mito analyses.
6. Just for additional info, determine how many connected components
are formed by the mito objects.
7. Return the mito IDs, sizses, and CC info as a DataFrame.
"""
# The neighborhood segmentation source
protocol, url = mito_roi_source.split('://')[-2:]
server, uuid, instance = url.split('/')
server = f'{protocol}://{server}'
origin_zyx = np.array(neighborhood_origin_xyz[::-1])
box = [origin_zyx - RADIUS, 1 + origin_zyx + RADIUS]
# Align box to the analysis scale before scaling it.
box = round_box(box, (2**ANALYSIS_SCALE))
# Scale box
box //= (2**ANALYSIS_SCALE)
neighborhood_seg = fetch_labelmap_voxels(server,
uuid,
instance,
box,
scale=ANALYSIS_SCALE)
neighborhood_mask = (neighborhood_seg == neighborhood_id)
# This is equivalent to a 1-px erosion
# See note above for why we do this.
neighborhood_mask ^= binary_edge_mask(neighborhood_mask, 'inner')
mito_seg = fetch_labelmap_voxels(*MITO_SEG,
box,
supervoxels=True,
scale=ANALYSIS_SCALE -
mito_res_scale_diff)
assert neighborhood_mask.shape == mito_seg.shape
mito_seg = np.where(neighborhood_mask, mito_seg, 0)
# The mito segmentation includes little scraps and slivers
# that were filtered out of the "real" mito set.
# Filter those scraps out of our results here.
mito_ids = set(pd.unique(mito_seg.ravel())) - {0}
mito_sizes = fetch_sizes(*MITO_SEG, [*mito_ids], supervoxels=True)
mito_sizes = mito_sizes.rename_axis('mito')
mito_sizes *= (2**mito_res_scale_diff)**3
# This is our main result: mito IDs (and their sizes)
mito_sizes = mito_sizes.loc[mito_sizes >= MIN_MITO_SIZE]
# Just for extra info, group the mitos we found into connected components.
mito_mask = mask_for_labels(mito_seg, mito_sizes.index)
mito_box = compute_nonzero_box(mito_mask)
mito_mask = extract_subvol(mito_mask, mito_box)
mito_seg = extract_subvol(mito_seg, mito_box)
mito_cc = label(mito_mask, connectivity=1)
ct = contingency_table(mito_seg, mito_cc).reset_index()
ct = ct.rename(columns={
'left': 'mito',
'right': 'cc',
'voxel_count': 'cc_size'
})
ct = ct.set_index('mito')
mito_sizes = pd.DataFrame(mito_sizes).merge(ct,
'left',
left_index=True,
right_index=True)
return mito_sizes
def load_roi_label_volume(server,
uuid,
rois_or_neuprint,
box_s5=[None, None],
export_path=None,
export_labelmap=None):
"""
Fetch several ROIs from DVID and combine them into a single label volume or mask.
The label values in the returned volume correspond to the order in which the ROI
names were passed in, starting at label 1.
This function is essentially a convenience function around fetch_combined_roi_volume(),
but in this case it will optionally auto-fetch the ROI list, and auto-export the volume.
Args:
server:
DVID server
uuid:
DVID uuid
rois_or_neuprint:
Either a list of ROIs or a neuprint server from which to obtain the roi list.
box_s5:
If you want to restrict the ROIs to a particular subregion,
you may pass your own bounding box (at scale 5).
Alternatively, you may pass the name of a segmentation
instance from DVID whose bounding box will be used.
export_path:
If you want the ROI volume to be exported to disk,
provide a path name ending with .npy or .h5.
export_labelmap:
If you want the ROI volume to be exported to a DVID labelmap instance,
Provide the instance name, or a tuple of (server, uuid, instance).
Returns:
(roi_vol, roi_box), containing the fetched label volume and the
bounding box it corresponds to, in DVID scale-5 coordinates.
Note:
If you have a list of (full-res) points to extract from the returned volume,
pass a DataFrame with columns ['z','y','x'] to the following function.
If you already downloaded the roi_vol (above), provide it.
Otherwise, leave out those args and it will be fetched first.
Adds columns to the input DF (in-place) for 'roi' (str) and 'roi_label' (int).
>>> from neuclease.dvid import determine_point_rois
>>> determine_point_rois(*master, rois, point_df, roi_vol, roi_box)
"""
if isinstance(box_s5, str):
# Assume that this is a segmentation instance whose dimensions should be used
# Fetch the maximum extents of the segmentation,
# and rescale it for scale-5.
seg_box = fetch_volume_box(server, uuid, box_s5)
box_s5 = round_box(seg_box, (2**5), 'out') // 2**5
box_s5[0] = (0, 0, 0)
if export_labelmap:
assert isinstance(box_s5, np.ndarray)
assert not (box_s5 % 64).any(), \
("If exporting to a labelmap instance, please supply "
"an explicit box and make sure it is block-aligned.")
if isinstance(rois_or_neuprint, (str, neuprint.Client)):
if isinstance(rois_or_neuprint, str):
npclient = neuprint.Client(rois_or_neuprint)
else:
npclient = rois_or_neuprint
# Fetch ROI names from neuprint
q = "MATCH (m: Meta) RETURN m.superLevelRois as rois"
rois = npclient.fetch_custom(q)['rois'].iloc[0]
rois = sorted(rois)
# # Remove '.*ACA' ROIs. Apparently there is some
# # problem with them. (They overlap with other ROIs.)
# rois = [*filter(lambda r: 'ACA' not in r, rois)]
else:
assert isinstance(rois_or_neuprint, collections.abc.Iterable)
rois = rois_or_neuprint
# Fetch each ROI and write it into a volume
with Timer(f"Fetching combined ROI volume for {len(rois)} ROIs", logger):
roi_vol, roi_box, overlap_stats = fetch_combined_roi_volume(
server, uuid, rois, box_zyx=box_s5)
if len(overlap_stats) > 0:
logger.warn(
f"Some ROIs overlap! Here's an incomplete list of overlapping pairs:\n{overlap_stats}"
)
# Export to npy/h5py for external use
if export_path:
with Timer(f"Exporting to {export_path}", logger):
if export_path.endswith('.npy'):
np.save(export_path, roi_vol)
elif export_path.endswith('.h5'):
with h5py.File(export_path, 'w') as f:
f.create_dataset('rois_scale_5', data=roi_vol, chunks=True)
if export_labelmap:
if isinstance(export_labelmap, str):
export_labelmap = (server, uuid, export_labelmap)
assert len(export_labelmap) == 3
with Timer(f"Exporting to {export_labelmap[2]}", logger):
if export_labelmap[2] not in fetch_repo_instances(
server, uuid, 'labelmap'):
create_labelmap_instance(
*export_labelmap, voxel_size=8 * (2**5),
max_scale=6) # FIXME: hard-coded voxel size
# It's really important to use this block shape.
# See https://github.com/janelia-flyem/dvid/issues/342
boxes = boxes_from_grid(roi_box, (256, 256, 256), clipped=True)
for box in tqdm_proxy(boxes):
block = extract_subvol(roi_vol, box - roi_box[0])
post_labelmap_voxels(*export_labelmap,
box[0],
block,
scale=0,
downres=True)
return roi_vol, roi_box, rois
def _fetch_body_mask(seg_src, primary_point_s0, body, search_cfg, tbar_points_s0, logger):
"""
Fetch a mask for the given body around the given point, with the given radius.
The mask will be downloaded using download_scale and then rescaled to
the analysis_scale using continuity-preserving downsampling.
The returned mask is NOT a binary (boolean) volume. Instead, a uint8 volume is returned,
with labels 1 and 2, indicating which portion of the mask belongs to the body (2) and
which portion was added due to dilation (1).
Providing a non-binary result is convenient for debugging. It is also used to
restrict the voxels that are preserved when filtering mitos, if no valid_mitos
are provided to that function.
"""
(radius_s0, download_scale, analysis_scale,
dilation_radius_s0, dilation_exclusion_buffer_s0, _is_final) = search_cfg
scale_diff = analysis_scale - download_scale
with Timer("Fetching body segmentation", logger):
assert _have_flyemflows and isinstance(seg_src, VolumeService)
p = np.asarray(primary_point_s0) // (2**download_scale)
R = radius_s0 // (2**download_scale)
seg_box = [p-R, p+R+1]
# Align to 64-px for consistency with the dvid case,
# and compatibility with the code below.
seg_box = round_box(seg_box, 64 * (2**scale_diff), 'out')
p_local = (p - seg_box[0])
seg = seg_src.get_subvolume(seg_box, download_scale)
# Extract mask
raw_mask = (seg == body).view(np.uint8)
del seg
# Downsample mask conservatively, i.e. keeping 'on' pixels no matter what
seg_box //= (2**scale_diff)
p //= (2**scale_diff)
p_local //= (2**scale_diff)
raw_mask = downsample_mask(raw_mask, 2**scale_diff, 'or')
# Due to downsampling effects in the original data, it's possible
# that the main tbar fell off its body in the downsampled image.
# Make sure it's part of the mask
raw_mask[(*p_local,)] = True
# The same is true for all of the other tbars that fall within the mask box.
# Fix them all.
tbar_points = tbar_points_s0 // (2**analysis_scale)
in_box = (tbar_points >= seg_box[0]).all(axis=1) & (tbar_points < seg_box[1]).all(axis=1)
tbar_points = tbar_points[in_box]
local_tbar_points = tbar_points - seg_box[0]
raw_mask[(*local_tbar_points.transpose(),)] = True
assert raw_mask.dtype == bool
raw_mask = vigra.taggedView(raw_mask.view(np.uint8), 'zyx')
dilation_buffer = dilation_exclusion_buffer_s0 // (2**analysis_scale)
dilation_box = np.array((seg_box[0] + dilation_buffer, seg_box[1] - dilation_buffer))
# Shrink to fit the data.
local_box = compute_nonzero_box(raw_mask)
raw_mask = raw_mask[box_to_slicing(*local_box)]
mask_box = local_box + seg_box[0]
p_local -= local_box[0]
# Fill gaps
repaired_mask = _fill_gaps(raw_mask, mask_box, analysis_scale, dilation_radius_s0, dilation_box)
# Label the voxels:
# 1: filled gaps
# 2: filled gaps AND raw
assert raw_mask.dtype == repaired_mask.dtype == np.uint8
body_mask = np.where(repaired_mask, raw_mask + repaired_mask, 0)
assert body_mask[(*p - mask_box[0],)]
return body_mask, mask_box
def setup_dvid_segmentation_input(setup_dvid_repo, random_segmentation):
dvid_address, repo_uuid = setup_dvid_repo
input_segmentation_name = 'labelmapcopy-segmentation-input'
output_segmentation_name = 'labelmapcopy-segmentation-output'
partial_output_segmentation_name = 'labelmapcopy-segmentation-partial-output'
max_scale = 3
already_exists = False
try:
create_labelmap_instance(dvid_address,
repo_uuid,
input_segmentation_name,
max_scale=max_scale)
create_labelmap_instance(dvid_address,
repo_uuid,
partial_output_segmentation_name,
max_scale=max_scale)
except HTTPError as ex:
if ex.response is not None and 'already exists' in ex.response.content.decode(
'utf-8'):
already_exists = True
expected_vols = {}
for scale in range(1 + max_scale):
if scale == 0:
scaled_vol = random_segmentation
else:
scaled_vol = downsample(scaled_vol, 2, 'labels-numba')
expected_vols[scale] = scaled_vol
if not already_exists:
scaled_box = round_box([(0, 0, 0), scaled_vol.shape], 64, 'out')
aligned_vol = np.zeros(scaled_box[1], np.uint64)
overwrite_subvol(aligned_vol, [(0, 0, 0), scaled_vol.shape],
scaled_vol)
post_labelmap_voxels(dvid_address,
repo_uuid,
input_segmentation_name, (0, 0, 0),
aligned_vol,
scale=scale)
if not already_exists:
# Create a 'partial' output volume that is the same (bitwise) as the input except for some blocks.
scaled_box = np.array([(0, 0, 0), random_segmentation.shape])
scaled_box[1, -1] = 192
for scale in range(1 + max_scale):
scaled_box = round_box(scaled_box // (2**scale), 64, 'out')
raw_blocks = fetch_labelmap_voxels(dvid_address,
repo_uuid,
input_segmentation_name,
scaled_box,
scale,
supervoxels=True,
format='raw-response')
post_labelmap_blocks(dvid_address,
repo_uuid,
partial_output_segmentation_name, [(0, 0, 0)],
raw_blocks,
scale,
is_raw=True)
block = np.random.randint(1_000_000,
1_000_010,
size=(64, 64, 64),
dtype=np.uint64)
post_labelmap_voxels(dvid_address,
repo_uuid,
partial_output_segmentation_name, (0, 128, 64),
block,
0,
downres=True)
partial_vol = fetch_labelmap_voxels(dvid_address,
repo_uuid,
partial_output_segmentation_name,
[(0, 0, 0), random_segmentation.shape],
supervoxels=True)
template_dir = tempfile.mkdtemp(suffix="labelmapcopy-template")
config_text = textwrap.dedent(f"""\
workflow-name: labelmapcopy
cluster-type: {CLUSTER_TYPE}
input:
dvid:
server: {dvid_address}
uuid: {repo_uuid}
segmentation-name: {input_segmentation_name}
supervoxels: true
geometry:
message-block-shape: [512,64,64]
available-scales: [0,1,2,3]
output:
dvid:
server: {dvid_address}
uuid: {repo_uuid}
segmentation-name: {output_segmentation_name}
supervoxels: true
disable-indexing: true
create-if-necessary: true
labelmapcopy:
slab-shape: [512,128,64]
dont-overwrite-identical-blocks: true
""")
with open(f"{template_dir}/workflow.yaml", 'w') as f:
f.write(config_text)
yaml = YAML()
with StringIO(config_text) as f:
config = yaml.load(f)
return template_dir, config, expected_vols, partial_vol, dvid_address, repo_uuid, output_segmentation_name, partial_output_segmentation_name
def init_boxes(self, volume_service, subset_labels, roi):
sbm = None
if roi:
base_service = volume_service.base_service
assert isinstance(base_service, DvidVolumeService), \
"Can't specify an ROI unless you're using a dvid input"
assert isinstance(volume_service, (ScaledVolumeService, DvidVolumeService)), \
"The 'roi' option doesn't support adapters other than 'rescale-level'"
scale = 0
if isinstance(volume_service, ScaledVolumeService):
scale = volume_service.scale_delta
assert scale <= 5, \
"The 'roi' option doesn't support volumes downscaled beyond level 5"
server, uuid, _seg_instance = base_service.instance_triple
brick_shape = volume_service.preferred_message_shape
assert not (brick_shape % 2**(5-scale)).any(), \
"If using an ROI, select a brick shape that is divisible by 32"
seg_box = volume_service.bounding_box_zyx
seg_box = round_box(seg_box, brick_shape)
seg_box_s5 = seg_box // 2**(5 - scale)
with Timer(f"Fetching mask for ROI '{roi}'", logger):
roi_mask_s5, roi_box_s5 = fetch_roi(server,
uuid,
roi,
format='mask')
# Restrict to input bounding box
clipped_roi_box_s5 = box_intersection(seg_box_s5, roi_box_s5)
clipped_roi_mask_s5 = extract_subvol(
roi_mask_s5, clipped_roi_box_s5 - roi_box_s5[0])
# Align to brick grid
aligned_roi_box_s5 = round_box(clipped_roi_box_s5,
brick_shape // 2**5, 'out')
padding = (aligned_roi_box_s5 - clipped_roi_box_s5)
padding[0] *= -1
aligned_roi_mask_s5 = np.pad(clipped_roi_mask_s5,
padding.transpose())
# At the service native scale
aligned_roi_box = (2**(5 - scale) * aligned_roi_box_s5)
logger.info(
f"Brick-aligned ROI '{roi}' has bounding-box {aligned_roi_box[:, ::-1].tolist()}"
)
# SBM 'full-res' corresponds to the input service voxels, not necessarily scale-0.
sbm = SparseBlockMask.create_from_highres_mask(
aligned_roi_mask_s5, 2**(5 - scale), aligned_roi_box,
brick_shape)
elif subset_labels:
try:
sbm = volume_service.sparse_block_mask_for_labels(
[*subset_labels])
if ((sbm.box[1] - sbm.box[0]) == 0).any():
raise RuntimeError(
"Could not find sparse masks for any of the subset-labels"
)
except NotImplementedError:
sbm = None
if sbm is None:
boxes = boxes_from_grid(volume_service.bounding_box_zyx,
volume_service.preferred_message_shape,
clipped=True)
return np.array([*boxes])
else:
boxes = sbm.sparse_boxes(brick_shape)
boxes = np.array(boxes)
# Clip
boxes[:, 0, :] = np.maximum(volume_service.bounding_box_zyx[0],
boxes[:, 0, :])
boxes[:, 1, :] = np.minimum(volume_service.bounding_box_zyx[1],
boxes[:, 1, :])
assert (boxes[:,0,:] < boxes[:,1,:]).all(), \
"After cropping to input volume, some bricks disappeared."
return boxes
def _init_mask(self):
"""
- read the mask ROI as a volume
- dilate/erode it if necessary
- invert it if necessary
- save to .h5 (just for offline debug)
- return the scale-5 mask and its scale-5 bounding-box
"""
options = self.config["masksegmentation"]
roi = options["mask-roi"]
invert_mask = options["invert-mask"]
max_scale = options["max-pyramid-scale"]
roi_dilation = options["dilate-roi"]
roi_erosion = options["erode-roi"]
seg_dilation = options["dilate-segmentation"]
block_width = self.output_service.block_width
# Select a mask_box that's large enough to divide evenly into the
# block width even when reduced to the highest scale we'll be processing.
seg_box = round_box(self.input_service.bounding_box_zyx,
block_width * 2**max_scale)
seg_box_s5 = round_box(seg_box, 2**5) // (2**5)
with Timer(f"Loading ROI '{roi}'", logger):
roi_mask, _ = fetch_roi(self.input_service.server,
self.input_service.uuid,
roi,
format='mask',
mask_box=seg_box_s5)
with h5py.File('roi-mask.h5', 'w') as f:
f.create_dataset('mask',
data=roi_mask.view(np.uint8),
chunks=(128, 128, 128))
assert not (roi_dilation and roi_erosion)
if roi_dilation > 0:
with Timer(f"Dilating ROI by {roi_dilation}", logger):
roi_mask = vigra.filters.multiBinaryDilation(
roi_mask, roi_dilation)
with h5py.File('dilated-roi-mask.h5', 'w') as f:
f.create_dataset('mask',
data=roi_mask.view(np.uint8),
chunks=(128, 128, 128))
if roi_erosion > 0:
with Timer(f"Eroding ROI by {roi_erosion}", logger):
roi_mask = vigra.filters.multiBinaryErosion(
roi_mask, roi_erosion)
with h5py.File('eroded-roi-mask.h5', 'w') as f:
f.create_dataset('mask',
data=roi_mask.view(np.uint8),
chunks=(128, 128, 128))
assert not seg_dilation or invert_mask, \
"Can't use 'dilate-segmentation'. The segmentation isn't downloaded unless 'invert-mask' is used."
if invert_mask:
with Timer("Inverting mask", logger):
# Initialize the mask with entire segmentation at scale 5,
# then subtract the roi from it.
boxes = [
*boxes_from_grid(seg_box_s5, (64, 64, 2048), clipped=True)
]
input_service = self.input_service
def fetch_seg_mask_s5(box_s5):
seg_s5 = input_service.get_subvolume(box_s5, scale=5)
return box_s5, (seg_s5 != 0)
boxes_and_mask = dask.bag.from_sequence(
boxes, 1).map(fetch_seg_mask_s5).compute()
seg_mask = np.zeros(box_shape(seg_box_s5), bool)
for box_s5, box_mask in boxes_and_mask:
overwrite_subvol(seg_mask, box_s5, box_mask)
if seg_dilation == 0:
with h5py.File('segmentation-mask.h5', 'w') as f:
f.create_dataset('mask',
data=seg_mask.view(np.uint8),
chunks=(128, 128, 128))
else:
with Timer(f"Dilating segmentation by {seg_dilation}",
logger):
seg_mask = vigra.filters.multiBinaryDilation(
seg_mask, seg_dilation)
with h5py.File('dilated-segmentation-mask.h5', 'w') as f:
f.create_dataset('mask',
data=seg_mask.view(np.uint8),
chunks=(128, 128, 128))
seg_mask[roi_mask] = False
roi_mask = seg_mask
with h5py.File('final-mask.h5', 'w') as f:
f.create_dataset('mask',
data=roi_mask.view(np.uint8),
chunks=(128, 128, 128))
# Downsample the roi_mask to dvid-block resolution, just to see how many blocks it touches.
block_mask = view_as_blocks(roi_mask, (2, 2, 2)).any(axis=(3, 4, 5))
blocks_touched = block_mask.sum()
voxel_total = blocks_touched * (block_width**3)
logger.info(
f"Mask touches {blocks_touched} blocks ({voxel_total / 1e9:.1f} Gigavoxels)"
)
return roi_mask, seg_box_s5
def overwrite_box(box, lowres_mask):
assert lowres_mask.dtype == np.bool
assert not (box[0] % block_width).any()
assert lowres_mask.any(), \
"This function is supposed to be called on bricks that actually need masking"
# Crop box and mask to only include the extent of the masked voxels
nonzero_mask_box = compute_nonzero_box(lowres_mask)
nonzero_mask_box = round_box(nonzero_mask_box,
(block_width * 2**scale) // 2**5)
lowres_mask = extract_subvol(lowres_mask, nonzero_mask_box)
box = box[0] + (nonzero_mask_box * 2**(5 - scale))
box = box.astype(np.int32)
if scale <= 5:
mask = upsample(lowres_mask, 2**(5 - scale))
else:
# Downsample, but favor UNmasked voxels
mask = ~view_as_blocks(~lowres_mask, 3 *
(2**(scale - 5), )).any(axis=(3, 4, 5))
old_seg = input_service.get_subvolume(box, scale)
assert mask.dtype == np.bool
new_seg = old_seg.copy()
new_seg[mask] = 0
if (new_seg == old_seg).all():
# It's possible that there are no changed voxels, but only
# at high scales where the masked voxels were downsampled away.
#
# So if the original downscale pyramids are perfect,
# then the following assumption ought to hold.
#
# But I'm commenting it out in case the DVID pyramid at scale 5
# isn't pixel-perfect in some places.
#
# assert scale > 5
return None
def post_changed_blocks(old_seg, new_seg):
# If we post the whole volume, we'll be overwriting blocks that haven't changed,
# wasting space in DVID (for duplicate blocks stored in the child uuid).
# Instead, we need to only post the blocks that have changed.
# So, can't just do this:
# output_service.write_subvolume(new_seg, box[0], scale)
seg_diff = (old_seg != new_seg)
block_diff = view_as_blocks(seg_diff, 3 * (block_width, ))
changed_block_map = block_diff.any(axis=(3, 4, 5)).nonzero()
changed_block_corners = box[0] + np.transpose(
changed_block_map) * block_width
changed_blocks = view_as_blocks(
new_seg, 3 * (block_width, ))[changed_block_map]
encoded_blocks = encode_labelarray_blocks(
changed_block_corners, changed_blocks)
mgr = output_service.resource_manager_client
with mgr.access_context(output_service.server, True, 1,
changed_blocks.nbytes):
post_labelmap_blocks(*output_service.instance_triple,
None,
encoded_blocks,
scale,
downres=False,
noindexing=True,
throttle=False,
is_raw=True)
assert not (box % block_width).any(), \
"Should not write partial blocks"
post_changed_blocks(old_seg, new_seg)
del new_seg
if scale != 0:
# Don't collect statistics for higher scales
return None
erased_seg = old_seg.copy()
erased_seg[~mask] = 0
block_shape = 3 * (input_service.block_width, )
erased_stats_df = block_stats_for_volume(block_shape, erased_seg,
box)
return erased_stats_df
def scale_box(box, scale):
# Scale down, then round up to the nearest multiple of the block width
box = np.ceil(box / 2**scale).astype(np.int32)
return round_box(box, block_width)
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
def copy_box(box, scale):
assert not record_only or scale == 0
box = round_box(box, 64, 'out')
box_shape = (box[1] - box[0])
# Read input blocks
with mgr_client.access_context(input_service.server, True, 1,
np.prod(box_shape)):
input_raw_blocks = fetch_labelmap_voxels(
*input_service.instance_triple,
box,
scale,
False,
input_service.supervoxels,
format='raw-response')
# If we're just recording, parse and return
if scale == 0 and record_only:
_input_spans, input_labels = parse_labelarray_data(
input_raw_blocks, extract_labels=True)
return list(set(chain(*input_labels.values())))
# If not checking the output, just copy input to output
if not check_existing:
with mgr_client.access_context(output_service.server, False, 1,
np.prod(box_shape)):
post_labelmap_blocks(*output_service.instance_triple,
None,
input_raw_blocks,
scale,
output_service.enable_downres,
output_service.disable_indexing,
False,
is_raw=True)
if scale == 0 and record_labels:
_input_spans, input_labels = parse_labelarray_data(
input_raw_blocks, extract_labels=True)
return list(set(chain(*input_labels.values())))
return []
# Read from output
with mgr_client.access_context(output_service.server, True, 1,
np.prod(box_shape)):
output_raw_blocks = fetch_labelmap_voxels(
*output_service.instance_triple,
box,
scale,
False,
output_service.supervoxels,
format='raw-response')
# If no differences, no need to parse
if (input_raw_blocks == output_raw_blocks):
return []
input_spans = parse_labelarray_data(input_raw_blocks,
extract_labels=False)
output_spans = parse_labelarray_data(output_raw_blocks,
extract_labels=False)
# Compare block IDs
input_ids = set(input_spans.keys())
output_ids = set(output_spans.keys())
missing_from_output = input_ids - output_ids
missing_from_input = output_ids - input_ids
common_ids = input_ids & output_ids
for block_id in missing_from_input:
# FIXME: We should pass this in the result so it can be logged in the client, not the worker.
logger.error(
f"Not overwriting block-id: {block_id}. It doesn't exist in the input."
)
# Filter the input blocks so only the new/different ones remain
filtered_input_list = []
for block_id in missing_from_output:
start, stop = input_spans[block_id]
filtered_input_list.append(
(block_id, input_raw_blocks[start:stop]))
filtered_output_list = []
for block_id in common_ids:
in_start, in_stop = input_spans[block_id]
out_start, out_stop = output_spans[block_id]
in_buf = input_raw_blocks[in_start:in_stop]
out_buf = output_raw_blocks[out_start:out_stop]
if in_buf != out_buf:
filtered_input_list.append((block_id, in_buf))
filtered_output_list.append((block_id, out_buf))
# Sort filtered blocks so they appear in the same order in which we received them.
filtered_input_list = sorted(
filtered_input_list, key=lambda k_v: input_spans[k_v[0]][0])
# Post them
filtered_input_buf = b''.join(
[buf for (_, buf) in filtered_input_list])
with mgr_client.access_context(output_service.server, False, 1,
np.prod(box_shape)):
post_labelmap_blocks(*output_service.instance_triple,
None,
filtered_input_buf,
scale,
output_service.enable_downres,
output_service.disable_indexing,
False,
is_raw=True)
if scale == 0 and record_labels:
filtered_output_buf = b''.join(
[buf for (_, buf) in filtered_output_list])
_, filtered_input_labels = parse_labelarray_data(
filtered_input_buf, extract_labels=True)
_, filtered_output_labels = parse_labelarray_data(
filtered_output_buf, extract_labels=True)
input_set = set(chain(*filtered_input_labels.values()))
output_set = set(chain(*filtered_output_labels.values()))
return list(input_set - output_set)
return []
def _fetch_body_mask(seg_src, p, radius, scale, body, tbar_points):
"""
Fetch a mask for the given body around the given point, with the given radius.
Only the connected component that covers the given point will be returned.
If the component doesn't extend out to the given radius in all dimensions,
then the returned subvolume may be smaller than the requested radius would seem to imply.
"""
with Timer("Fetching body segmentation", logger):
if _have_flyemflows and isinstance(seg_src, VolumeService):
p = np.asarray(p)
R = radius
seg_box = [p - R, p + R + 1]
# Align to 64-px for consistency with the dvid case,
# and compatibility with the code below.
seg_box = round_box(seg_box, 64, 'out')
p_local = (p - seg_box[0])
seg = seg_src.get_subvolume(seg_box, scale)
else:
assert len(seg_src) == 3 and all(
isinstance(s, str) for s in seg_src)
seg, seg_box, p_local, _ = fetch_seg_around_point(
*seg_src,
p,
radius,
scale,
body,
sparse_component_only=True,
map_on_client=True,
threads=16)
# Due to downsampling effects, it's possible that the
# main tbar fell off its body in the downsampled image.
# Force it to have the correct label.
seg[(*p_local, )] = body
# The same is true for all of the other tbars that fall within the mask box.
# Fix them all.
in_box = (tbar_points >= seg_box[0]).all(
axis=1) & (tbar_points < seg_box[1]).all(axis=1)
tbar_points = tbar_points[in_box]
local_tbar_points = tbar_points - seg_box[0]
seg[(*local_tbar_points.transpose(), )] = body
# Compute mito CC within body mask
body_mask = (seg == body).view(np.uint8)
del seg
# Find the connected component that contains our point of interest
# amd limit our analysis to those voxels.
body_mask = vigra.taggedView(body_mask, 'zyx')
body_cc = labelMultiArrayWithBackground(body_mask)
# Update mask. Limit to extents of the main cc, but align box to nearest 64px
body_mask = (body_cc == body_cc[(*p_local, )])
body_block_mask = view_as_blocks(body_mask,
(64, 64, 64)).any(axis=(3, 4, 5))
body_block_corners = seg_box[0] + (64 * np.argwhere(body_block_mask))
mask_box = (body_block_corners.min(axis=0),
body_block_corners.max(axis=0) + 64)
local_box = mask_box - seg_box[0]
body_mask = body_mask[box_to_slicing(*local_box)]
assert body_mask[(*p - mask_box[0], )]
return body_mask, mask_box, body_block_corners
def autogen_points(input_seg,
count,
roi,
body,
tbars,
use_skeleton,
random_seed=None,
minimum_distance=0):
"""
Generate a list of points within the input segmentation, based on the given criteria.
See the main help text below for details.
"""
if tbars and not body:
sys.exit(
"If you want to auto-generate tbar points, please specify a body.")
if not tbars and not count:
sys.exit(
"You must supply a --count unless you are generating all tbars of a body."
)
if use_skeleton:
if not body:
sys.exit(
"You must supply a body ID if you want to use a skeleton.")
if tbars:
sys.exit(
"You can't select both tbar points and skeleton points. Pick one or the other."
)
if not count and minimum_distance > 0:
sys.exit(
"You must supply a --count if you want skeleton point samples to respect the minimum distance."
)
if not count and not roi and minimum_distance == 0:
logger.warning(
"You are using all nodes of a skeleton without any ROI filter! Is that what you meant?"
)
rng = default_rng(random_seed)
if tbars:
logger.info(f"Fetching synapses for body {body}")
syn_df = fetch_annotation_label(*input_seg[:2],
'synapses',
body,
format='pandas')
tbars = syn_df.query('kind == "PreSyn"')[[*'zyx']]
if roi:
logger.info(f"Filtering tbars for roi {roi}")
determine_point_rois(*input_seg[:2], [roi], tbars)
tbars = tbars.query('roi == @roi')[[*'zyx']]
if minimum_distance:
logger.info(
f"Pruning close points from {len(tbars)} total tbar points")
tbars = prune_close_pairs(tbars, minimum_distance, rng)
logger.info(f"After pruning, {len(tbars)} tbars remain.")
if count:
count = min(count, len(tbars))
logger.info(f"Sampling {count} tbars")
choices = rng.choice(tbars.index, size=count, replace=False)
tbars = tbars.loc[choices]
logger.info(f"Returning {len(tbars)} tbar points")
return tbars
elif use_skeleton:
assert body
logger.info(f"Fetching skeleton for body {body}")
skeleton_instance = f'{input_seg[2]}_skeletons'
swc = fetch_key(*input_seg[:2], skeleton_instance, f'{body}_swc')
skeleton_df = swc_to_dataframe(swc)
skeleton_df['x'] = skeleton_df['x'].astype(int)
skeleton_df['y'] = skeleton_df['y'].astype(int)
skeleton_df['z'] = skeleton_df['z'].astype(int)
if roi:
logger.info(f"Filtering skeleton for roi {roi}")
determine_point_rois(*input_seg[:2], [roi], skeleton_df)
skeleton_df = skeleton_df.query('roi == @roi')[[*'zyx']]
if minimum_distance:
assert count
# Distance-pruning is very expensive on a huge number of close points.
# If skeleton is large, first reduce the workload by pre-selecting a
# random sample of skeleton points, and prune more from there.
if len(skeleton_df) > 10_000:
# FIXME: random_state can't use rng until I upgrade to pandas 1.0
skeleton_df = skeleton_df.sample(min(4 * count,
len(skeleton_df)),
random_state=None)
logger.info(
f"Pruning close points from {len(skeleton_df)} skeleton points"
)
prune_close_pairs(skeleton_df, minimum_distance, rng)
logger.info(
f"After pruning, {len(skeleton_df)} skeleton points remain.")
if count:
count = min(count, len(skeleton_df))
logger.info(f"Sampling {count} skeleton points")
choices = rng.choice(skeleton_df.index, size=count, replace=False)
skeleton_df = skeleton_df.loc[choices]
logger.info(f"Returning {len(skeleton_df)} skeleton points")
return skeleton_df
elif body:
assert count
if roi:
# TODO: intersect the ranges with the ROI.
raise NotImplementedError(
"Sorry, I haven't yet implemented support for "
"body+roi filtering. Pick one or the other, "
"or ask Stuart to fix this.")
logger.info(f"Fetching sparsevol for body {body}")
ranges = fetch_sparsevol(*input_seg, body, format='ranges')
logger.info("Sampling from sparsevol")
if minimum_distance > 0:
# Sample 4x extra so we still have enough after pruning.
points = sample_points_from_ranges(ranges, 4 * count, rng)
else:
points = sample_points_from_ranges(ranges, count, rng)
points = pd.DataFrame(points, columns=[*'zyx'])
if minimum_distance > 0:
logger.info(f"Pruning close points from {len(points)} body points")
prune_close_pairs(points, minimum_distance, rng)
logger.info(f"After pruning, {len(points)} body points remain")
points = points.iloc[:count]
logger.info(f"Returning {len(points)} body points")
return points
elif roi:
assert count
logger.info(f"Fetching roi {roi}")
roi_ranges = fetch_roi_roi(*input_seg[:2], roi, format='ranges')
logger.info("Sampling from ranges")
if minimum_distance > 0:
# Sample 4x extra so we can prune some out if necessary.
points_s5 = sample_points_from_ranges(roi_ranges, 4 * count, rng)
else:
points_s5 = sample_points_from_ranges(roi_ranges, count, rng)
corners_s0 = points_s5 * (2**5)
points_s0 = rng.integers(corners_s0, corners_s0 + (2**5))
points = pd.DataFrame(points_s0, columns=[*'zyx'])
if minimum_distance > 0:
logger.info(f"Pruning close points from {len(points)} roi points")
prune_close_pairs(points, minimum_distance, rng)
logger.info(
f"After pruning, points from {len(points)} roi points remain")
points = points.iloc[:count]
logger.info(f"Returning {len(points)} roi points")
return points
else:
# No body or roi specified, just sample from the whole non-zero segmentation area
assert count
logger.info("Sampling random points from entire input segmentation")
logger.info("Fetching low-res input volume")
box_s6 = round_box(fetch_volume_box(*input_seg), 2**6, 'out') // 2**6
seg_s6 = fetch_labelmap_voxels(*input_seg, box_s6, scale=6)
mask_s6 = seg_s6.astype(bool)
logger.info("Encoding segmentation as ranges")
seg_ranges = runlength_encode_mask_to_ranges(mask_s6, box_s6)
logger.info("Sampling from ranges")
if minimum_distance > 0:
# Sample 4x extra so we can prune some out if necessary.
points_s6 = sample_points_from_ranges(seg_ranges, 4 * count, rng)
else:
points_s6 = sample_points_from_ranges(seg_ranges, count, rng)
corners_s0 = points_s6 * (2**6)
points_s0 = rng.integers(corners_s0, corners_s0 + (2**6))
points = pd.DataFrame(points_s0, columns=[*'zyx'])
if minimum_distance > 0:
logger.info(
f"Pruning close points from {len(points)} segmentation points")
prune_close_pairs(points, minimum_distance, rng)
logger.info(
f"After pruning, points from {len(points)} segmentation points remain"
)
points = points.iloc[:count]
logger.info(f"Returning {len(points)} segmentation points")
return points
