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