def agglomerate_with_waterz(affs, thresholds, fragments, histogram_quantiles=False, discrete_queue=False, merge_function='median_aff', init_with_max=True, return_merge_history=False, return_region_graph=False, **kwargs): print("Agglomerating with %s" % merge_function) if init_with_max: merge_function += '_maxinit' if histogram_quantiles: merge_function += '_histograms' discretize_queue = 0 if discrete_queue: discretize_queue = 256 return waterz.agglomerate( affs, thresholds, fragments=fragments, scoring_function=scoring_functions[merge_function], discretize_queue=discretize_queue, return_merge_history=return_merge_history, return_region_graph=return_region_graph)
def __call__(self, affs: np.ndarray, fragments: np.ndarray = None): """ Parameters: ----------- affs: affinity map with 4 dimensions: channel, z, y, x """ if isinstance(affs, Chunk): # the segmentation is 3d, so we only need the zyx global_offset = affs.global_offset[-3:] else: global_offset = None # our affinity map channel order is x,y,z! # meaning the first channel is x, the second is y, # the third is z. We have to reverse the channel. if self.flip_channel: affs = np.flip(affs, axis=0) # waterz need datatype float32 and # the array memory layout is contiguous affs = np.ascontiguousarray(affs, dtype=np.float32) # the output is a generator, and the computation is delayed seg_generator = agglomerate( affs, [self.threshold], fragments=fragments, aff_threshold_low=self.aff_threshold_low, aff_threshold_high=self.aff_threshold_high, scoring_function=self.scoring_function, force_rebuild=False) # there is only one threshold, so there is also only one result # in generator seg = next(seg_generator) return Chunk(seg, global_offset=global_offset)
def agglomerate_with_waterz(affs, thresholds, histogram_quantiles=False, discrete_queue=False, merge_function='median_aff', init_with_max=True, return_merge_history=False, return_region_graph=False, has_background=None): print("Extracting initial fragments...") fragments, affs_xy, distances, seeds = watershed(affs, 'maxima_distance', has_background) print("Agglomerating with %s", merge_function) if init_with_max: merge_function += '_maxinit' if histogram_quantiles: merge_function += '_histograms' discretize_queue = 0 if discrete_queue: discretize_queue = 256 return (waterz.agglomerate( affs, thresholds, fragments=fragments, scoring_function=scoring_functions[merge_function], discretize_queue=discretize_queue, return_merge_history=return_merge_history, return_region_graph=return_region_graph), fragments, affs_xy, distances, seeds)
def blockwise_segmentation_function_using_contact_area(array_in, roi, thresholds, quantile): #Materialize region of interest predicted_distances = array_in.to_ndarray(roi, fill_value=0) #Normalizing distances normalized_distances = normalize_distances(predicted_distances) predicted_distances = None #Find seeds seeds = find_seeds(normalized_distances) #Watershed fragments fragments = watershed(-normalized_distances, seeds, mask = (normalized_distances.astype(bool))).astype(np.uint64) #Creating affinity arrays normalized_distances/=20 #divide by 20 since that is maximum, https://www.wolframalpha.com/input/?i=%28%28128*tanh%28sqrt%281%2Fpi%29%2F12.5%29%2B127%29-126%29%2F129 affs = np.stack([normalized_distances.data]*3) #print(f"maximum: {np.amax(normalized_distances)}") #Agglolmerate agglomeration = [] for s in waterz.agglomerate( affs=affs, fragments=fragments, thresholds=[thresholds], #max is 20: https://www.wolframalpha.com/input/?i=%28%28128*tanh%28sqrt%281%2Fpi%29%2F12.5%29%2B127%29-126%29%2F129 scoring_function = 1.0 - waterz.QuantileAffinity(quantile, init_with_max=False)/np.tanh( waterz.ContactArea())): agglomeration = s return agglomeration
def blockwise_segmentation_function(array_in, roi, thresholds, quantile): #Materialize region of interest predicted_distances = array_in.to_ndarray(roi, fill_value=0) #Normalizing distances normalized_distances = normalize_distances(predicted_distances) predicted_distances = None #Find seeds seeds = find_seeds(normalized_distances) #Watershed fragments fragments = watershed(-normalized_distances, seeds, mask = (normalized_distances.astype(bool))).astype(np.uint64) #Creating affinity arrays affs = np.stack([normalized_distances.data]*3) #Agglolmerate agglomeration = [] for s in waterz.agglomerate( affs=affs, fragments=fragments, thresholds=[thresholds], scoring_function = 1 - waterz.QuantileAffinity(quantile, init_with_max=False)): agglomeration = s return agglomeration
def evaluate_affs(affs, labels, thresholds): scores = {} segmentations = [] fragments = watershed_from_affinities(affs)[0] a = affs.astype(np.float32) l = labels.astype(np.uint32) f = fragments.astype(np.uint64) i = 0 for segmentation, metrics in waterz.agglomerate( affs=a, thresholds=thresholds, gt=l, fragments=f): segmentations.append(segmentation) scores[f'threshold_{thresholds[i]}'] = { 'voi_split': metrics['V_Info_split'], 'voi_merge': metrics['V_Info_merge'], 'rand_split': metrics['V_Rand_split'], 'rand_merge': metrics['V_Rand_merge'] } i += 1 return segmentations, scores, fragments
def segment_affs(affs, thresholds): #scores = {} segmentations = [] fragments = watershed_from_affinities(affs)[0] a = affs.astype(np.float32) #l = labels.astype(np.uint32) f = fragments.astype(np.uint64) i = 0 for segmentation in waterz.agglomerate( affs=a, thresholds=thresholds, #gt=l, fragments=f): segmentations.append(segmentation) # scores[f'threshold_{thresholds[i]}'] = { # 'voi_split': metrics['V_Info_split'], # 'rand_split': metrics['V_Rand_split'], # 'rand_merge': metrics['V_Rand_merge'] #} i += 1 return segmentations, fragments
def execute(affs: Chunk, fragments: np.ndarray = None): """ Mean/max agglomeration of affinity map including watershed step. Parameters: ----------- affs: affinity map with 4 dimensions: channel, z, y, x """ properties = affs.properties # our affinity map channel order is x,y,z! # meaning the first channel is x, the second is y, # the third is z. We have to reverse the channel for waterz. if flip_channel: affs = np.flip(affs, axis=0) # waterz need datatype float32 and # the array memory layout is contiguous affs = np.ascontiguousarray(affs, dtype=np.float32) # the output is a generator, and the computation is delayed seg_generator = agglomerate( affs, [threshold], fragments=fragments, aff_threshold_low=aff_threshold_low, aff_threshold_high=aff_threshold_high, scoring_function=scoring_function, force_rebuild=False) # there is only one threshold, so there is also only one result # in generator seg = next(seg_generator) seg = Chunk(seg) seg.set_properties(properties) return [seg]
def set_prediction(self, prediction): fragments = watershed_from_affinities(prediction.data)[0] thresholds = self.parameter_range.threshold segmentations = waterz.agglomerate( affs=prediction.data.astype(np.float32), fragments=fragments, thresholds=thresholds, ) self.segmentations = {t: s for t, s in zip(thresholds, segmentations)}
def get_segmentation(affinities, threshold): fragments = watershed_from_affinities(affinities)[0] thresholds = [threshold] segmentations = waterz.agglomerate( affs=affinities.astype(np.float32), fragments=fragments, thresholds=thresholds, ) segmentation = next(segmentations) return segmentation
def __watershed_return_metrics_list(aff_vol, gt_vol, thresh, metric="both", thresh_high=0.9999, thresh_low=0.0001): aff_vol = np.ascontiguousarray(aff_vol, dtype=np.float32) gt_vol = np.ascontiguousarray(gt_vol, dtype=np.uint32) if str(metric) == "choose": metrics = np.zeros((np.shape(thresh)[0], 2)) else: metrics = np.zeros(np.shape(thresh)) segmentations = np.zeros(np.shape(thresh), dtype=np.object) seg = waterz.agglomerate(aff_vol, thresholds=thresh, gt=gt_vol, aff_threshold_high=thresh_high, aff_threshold_low=thresh_low) n = 0 for segmentation in seg: segmentations[n] = segmentation[0] dictionary = segmentation[1] if str(metric) == "both": metrics[ n] = dictionary['V_Rand_split'] + dictionary["V_Rand_merge"] elif str(metric) == "split": metrics[n] = dictionary['V_Rand_split'] elif str(metric) == "merge": metrics[n] = dictionary['V_Rand_merge'] elif str(metric) == "choose": metrics[n, 0] = dictionary['V_Rand_split'] metrics[n, 1] = dictionary['V_Rand_merge'] else: assert 1 == 0, "Bad metric string" n += 1 return segmentations, metrics
def agglomerate(affs, gt, thresholds, custom_fragments=False, histogram_quantiles=False, discrete_queue=False, merge_function=None, init_with_max=True, fragments_mask=None, aff_high=0.9999, aff_low=0.0001): if merge_function is not 'zwatershed': fragments = None if custom_fragments: fragments = watershed(affs, 'maxima_distance') if fragments_mask is not None: fragments[fragments_mask == False] = 0 if init_with_max: merge_function += '_maxinit' if histogram_quantiles: merge_function += '_histograms' discretize_queue = 0 if discrete_queue: discretize_queue = 256 return waterz.agglomerate( affs, thresholds, gt, fragments=fragments, scoring_function=config.scoring_function[merge_function], discretize_queue=discretize_queue, aff_threshold_high=aff_high, aff_threshold_low=aff_low) else: return zwatershed_thresholds(affs, thresholds, gt, aff_high, aff_low)
def agglomerate_fragments( affs, fragments, threshold, scoring_function='OneMinus<HistogramQuantileAffinity<RegionGraphType, 25, ScoreValue, 256, false>>' ): # waterz agglomerate changes fragments --> # Make a copy here fragments = fragments.copy() generator = waterz.agglomerate(affs, thresholds=[threshold], fragments=fragments, scoring_function=scoring_function) # segs = [] for seg in generator: # segs.append(seg) logger.info('Number of segments in volume: {}'.format( len(np.unique(seg)))) return seg
def __watershed_return_metrics_single(aff_vol, gt_vol, thresh, metric="both", no_metric_flag=0, thresh_high=0.9999, thresh_low=0.0001): aff_vol = np.ascontiguousarray(aff_vol, dtype=np.float32) gt_vol = np.ascontiguousarray(gt_vol, dtype=np.uint32) seg = waterz.agglomerate(aff_vol, thresholds=[thresh], gt=gt_vol, aff_threshold_high=thresh_high, aff_threshold_low=thresh_low) for segmentation in seg: seg = segmentation if no_metric_flag == 0: metrics = seg[1] if str(metric) == "both": metric = metrics['V_Rand_split'] + metrics["V_Rand_merge"] elif str(metric) == "split": metric = metrics['V_Rand_split'] elif str(metric) == "merge": metric = metrics['V_Rand_merge'] else: metric = metrics return seg[0], metric else: return seg[0]
def watershed(aff_vol, thresh=0.1, thresh_high=0.9999, thresh_low=0.0001): process = psutil.Process(os.getpid()) mem_b = psutil.virtual_memory().used time_b = time.time() aff_vol = np.ascontiguousarray(aff_vol, dtype=np.float32) seg = waterz.agglomerate(aff_vol, thresholds=[thresh], aff_threshold_high=thresh_high, aff_threshold_low=thresh_low, discretize_queue=256) for segmentation in seg: seg = segmentation time_e = time.time() mem_a = psutil.virtual_memory().used log.log("watershed", np.shape(seg), time_e - time_b, mem_a - mem_b) return seg
import affinities as af import waterz as wz import malis as mal gt = np.load("data/spir_gt.npy") sample = gt[0:100, 0:400, 0:400] nhood = mal.mknhood3d(1) aff = mal.seg_to_affgraph(sample, nhood) num_act = np.shape(np.unique(sample))[0] - 1 aff = np.asarray(aff, dtype=np.float32) seg = wz.agglomerate(aff, thresholds=[1]) for segmentation in seg: seg = segmentation num_calc = np.shape(np.unique(seg))[0] - 1 print("Calculated: %i" % num_calc) print("Actual: %i" % num_act) #print(np.unique(seg)) #print(np.unique(sample)) if (np.equal(seg, sample).all): "PASSED"
def watershed_in_block( affs, block, context, rag_provider, fragments_out, num_voxels_in_block, mask=None, fragments_in_xy=False, epsilon_agglomerate=0.0, filter_fragments=0.0, min_seed_distance=10, replace_sections=None): ''' Args: filter_fragments (float): Filter fragments that have an average affinity lower than this value. min_seed_distance (int): Controls distance between seeds in the initial watershed. Reducing this value improves downsampled segmentation. ''' total_roi = affs.roi logger.debug("reading affs from %s", block.read_roi) affs = affs.intersect(block.read_roi) affs.materialize() if affs.dtype == np.uint8: logger.info("Assuming affinities are in [0,255]") max_affinity_value = 255.0 affs.data = affs.data.astype(np.float32) else: max_affinity_value = 1.0 if mask is not None: logger.debug("reading mask from %s", block.read_roi) mask_data = get_mask_data_in_roi(mask, affs.roi, affs.voxel_size) logger.debug("masking affinities") affs.data *= mask_data # extract fragments fragments_data, _ = watershed_from_affinities( affs.data, max_affinity_value, fragments_in_xy=fragments_in_xy, min_seed_distance=min_seed_distance) if mask is not None: fragments_data *= mask_data.astype(np.uint64) if filter_fragments > 0: if fragments_in_xy: average_affs = np.mean(affs.data[0:2]/max_affinity_value, axis=0) else: average_affs = np.mean(affs.data/max_affinity_value, axis=0) filtered_fragments = [] fragment_ids = np.unique(fragments_data) for fragment, mean in zip( fragment_ids, measurements.mean( average_affs, fragments_data, fragment_ids)): if mean < filter_fragments: filtered_fragments.append(fragment) filtered_fragments = np.array( filtered_fragments, dtype=fragments_data.dtype) replace = np.zeros_like(filtered_fragments) replace_values(fragments_data, filtered_fragments, replace, inplace=True) if epsilon_agglomerate > 0: logger.info( "Performing initial fragment agglomeration until %f", epsilon_agglomerate) generator = waterz.agglomerate( affs=affs.data/max_affinity_value, thresholds=[epsilon_agglomerate], fragments=fragments_data, scoring_function='OneMinus<HistogramQuantileAffinity<RegionGraphType, 25, ScoreValue, 256, false>>', discretize_queue=256, return_merge_history=False, return_region_graph=False) fragments_data[:] = next(generator) # cleanup generator for _ in generator: pass if replace_sections: logger.info("Replacing sections...") block_begin = block.write_roi.get_begin() shape = block.write_roi.get_shape() z_context = context[0]/affs.voxel_size[0] logger.info("Z context: %i",z_context) mapping = {} voxel_offset = block_begin[0]/affs.voxel_size[0] for i,j in zip( range(fragments_data.shape[0]), range(shape[0])): mapping[i] = i mapping[j] = int(voxel_offset + i) \ if block_begin[0] == total_roi.get_begin()[0] \ else int(voxel_offset + (i - z_context)) logging.info('Mapping: %s', mapping) replace = [k for k,v in mapping.items() if v in replace_sections] for r in replace: logger.info("Replacing mapped section %i with zero", r) fragments_data[r] = 0 #todo add key value replacement option fragments = daisy.Array(fragments_data, affs.roi, affs.voxel_size) # crop fragments to write_roi fragments = fragments[block.write_roi] fragments.materialize() max_id = fragments.data.max() # ensure we don't have IDs larger than the number of voxels (that would # break uniqueness of IDs below) if max_id > num_voxels_in_block: logger.warning( "fragments in %s have max ID %d, relabelling...", block.write_roi, max_id) fragments.data, max_id = relabel(fragments.data) assert max_id < num_voxels_in_block # ensure unique IDs id_bump = block.block_id[1]*num_voxels_in_block logger.debug("bumping fragment IDs by %i", id_bump) fragments.data[fragments.data>0] += id_bump fragment_ids = range(id_bump + 1, id_bump + 1 + int(max_id)) # store fragments logger.debug("writing fragments to %s", block.write_roi) fragments_out[block.write_roi] = fragments # following only makes a difference if fragments were found if max_id == 0: return # get fragment centers fragment_centers = { fragment: block.write_roi.get_offset() + affs.voxel_size*daisy.Coordinate(center) for fragment, center in zip( fragment_ids, measurements.center_of_mass(fragments.data, fragments.data, fragment_ids)) if not np.isnan(center[0]) } # store nodes rag = rag_provider[block.write_roi] rag.add_nodes_from([ (node, { 'center_z': c[0], 'center_y': c[1], 'center_x': c[2] } ) for node, c in fragment_centers.items() ]) rag.write_nodes(block.write_roi)
#Image.fromarray(((~mask[20,:,:]).astype(np.uint8).astype(np.float32)*affs[1,20,:,:]*255).astype(np.uint8)).save('temp3.png') # mask affs for d in range(3): affs[d][mask] = 0 # watershed fragments = watershed(affs, 'maxima_distance') # mask fragments fragments[mask] = 0 # agglomerate threshold = [args.thresd] sf = 'OneMinus<EdgeStatisticValue<RegionGraphType, MeanAffinityProvider<RegionGraphType, ScoreValue>>>' #sf = 'OneMinus<QuantileAffinity<RegionGraphType, 85, ScoreValue>>' segmentation = list( waterz.agglomerate(affs, threshold, None, fragments, scoring_function=sf, discretize_queue=256)) # save out_file = os.path.join(output_filename, 'segmentation.hdf') seg = segmentation[0] out = h5py.File(out_file, 'w') out.create_dataset('labels', data=seg, dtype=seg.dtype, compression='gzip') out.close()
def label(**kwargs): logger.info("labelling %s %s", kwargs['sample'], kwargs['gt']) sample = os.path.join(kwargs['pred_folder'], kwargs['sample'] + "." + kwargs['pred_format']) if kwargs['pred_format'] == "zarr": input_file = zarr.open(sample, 'r') elif kwargs['pred_format'] == "hdf": input_file = h5py.File(sample, 'r') else: raise NotImplementedError("invalid pred format") if kwargs.get('waterz'): affinities = np.array(input_file[kwargs['surf_key']]) segmentations = list( waterz.agglomerate( affinities, [kwargs['waterz_threshold']]))[0].astype(np.uint32) print(segmentations.shape, segmentations.dtype, len(np.unique(segmentations)), kwargs['waterz_threshold']) segmentations_dil = np.copy(segmentations) if kwargs['num_dilations'] > 0 and len( np.unique(segmentations)) < 1000: segmentations_dil = np.copy(segmentations) for lbl in np.unique(segmentations_dil): if lbl == 0: continue label_mask = segmentations_dil == lbl dilated_label_mask = scipy.ndimage.binary_dilation( label_mask, iterations=kwargs['num_dilations']) segmentations_dil[dilated_label_mask] = lbl if kwargs['output_format'] == "hdf": out_fn = os.path.join( kwargs['output_folder'], kwargs['sample'] + "." + kwargs['output_format']) with h5py.File(out_fn, 'w') as output_file: # write fgbg prediction to file output_file.create_dataset('volumes/watershed_seg_fg', data=segmentations, compression='gzip') output_file.create_dataset('volumes/watershed_seg_fg_dilated', data=segmentations_dil, compression='gzip') return surf = np.array(input_file[kwargs['surf_key']]) fgbg = np.array(input_file[kwargs['fgbg_key']]) raw = np.array(input_file[kwargs['raw_key']]) if kwargs['pred_format'] == "hdf": input_file.close() # threshold bg/fg fg = 1.0 * (fgbg > kwargs['fg_thresh']) if np.count_nonzero(fg) == 0: logger.warning("{}: no foreground found".format(kwargs['sample'])) # combine surface components surf_scalar = 1.0 - 0.33 * (surf[0] + surf[1] + surf[2]) # load gt if 'gt' in kwargs and kwargs['gt'] is not None: if kwargs['gt_format'] == "hdf": with h5py.File(kwargs['gt'], 'r') as gt: gt_labels = np.array(gt[kwargs['gt_key']]) elif kwargs['gt_format'] == "zarr": gt = zarr.open(kwargs['gt'], 'r') gt_labels = np.array(gt[kwargs['gt_key']]) else: raise NotImplementedError("invalid gt format") logger.debug("%s: gt min %f, max %f", kwargs['sample'], gt_labels.min(), gt_labels.max()) # compute markers for watershed (seeds) seeds = (1 * (surf > kwargs['seed_thresh'])).astype(np.uint8) seeds = (seeds[0] + seeds[1] + seeds[2]) seeds = (seeds > 2).astype(np.uint8) logger.info("%s: seeds min/max %f %f", kwargs['sample'], np.min(seeds), np.max(seeds)) if np.count_nonzero(seeds) == 0: logger.warning("%s: no seed points found for watershed", sample) markers, cnt = scipy.ndimage.label(seeds) logger.debug("%s: markers min %f, max %f, cnt %f", kwargs['sample'], np.min(markers), np.max(markers), cnt) # compute watershed wsUI, wsFGUI, wsFGUIdil = watershed(kwargs['sample'], surf_scalar, markers, fg, its=kwargs['num_dilations']) if kwargs['output_format'] == "hdf": out_fn = os.path.join(kwargs['output_folder'], kwargs['sample'] + "." + kwargs['output_format']) with h5py.File(out_fn, 'w') as output_file: # write fgbg prediction to file output_file.create_dataset('volumes/fgbg', data=fgbg, compression='gzip') output_file.create_dataset('volumes/surf', data=surf_scalar, compression='gzip') output_file.create_dataset('volumes/raw', data=raw, compression='gzip') output_file.create_dataset('volumes/seeds', data=seeds, compression='gzip') output_file.create_dataset('volumes/watershed_seg', data=wsUI, compression='gzip') output_file.create_dataset('volumes/watershed_seg_fg', data=wsFGUI, compression='gzip') output_file.create_dataset('volumes/watershed_seg_fg_dilated', data=wsFGUIdil, compression='gzip') else: raise NotImplementedError("invalid output format")
def evaluate_affs(affs, labels, dims, store_results=None): num_samples = affs.data.shape[0] scores = {} segmentations = [] if dims == 2: # get all fragments at once for all samples # (use samples as z dimension) # (2, s, h, w) a = affs.data.transpose((1, 0, 2, 3)) # (3, s, h, w) a = np.concatenate([np.zeros_like(a[0:1, :]), a]) # (s, h, w) fragments, _ = watershed_from_affinities(a, fragments_in_xy=True) else: fragments = [ watershed_from_affinities(affs[i])[0] for i in range(num_samples) ] for i in tqdm(range(num_samples), desc="evaluate"): a = affs.data[i].astype(np.float32) l = labels.data[i].astype(np.uint32) f = fragments[i].astype(np.uint64) if dims == 2: # convert to 3D a = np.concatenate( [np.zeros((1, 1) + a.shape[1:], dtype=np.float32), a[:,np.newaxis,:,:]]) l = l[np.newaxis,:,:] f = f[np.newaxis,:,:] for segmentation, metrics in waterz.agglomerate( affs=a, thresholds=[0.5], gt=l, fragments=f): segmentations.append(segmentation) scores[f'sample_{i}'] = { 'voi_split': metrics['V_Info_split'], 'voi_merge': metrics['V_Info_merge'], 'rand_split': metrics['V_Rand_split'], 'rand_merge': metrics['V_Rand_merge'] } if store_results: f = zarr.open(store_results) f['fragments'] = np.stack(fragments) f['segmentation'] = np.concatenate(segmentations) f['labels'] = labels.data return scores
def agglomerate_in_block( affs, fragments, rag_provider, block, merge_function, threshold): logger.info( "Agglomerating in block %s with context of %s", block.write_roi, block.read_roi) # get the sub-{affs, fragments, graph} to work on affs = affs.intersect(block.read_roi) fragments = fragments.to_ndarray(affs.roi, fill_value=0) rag = rag_provider[affs.roi] # waterz uses memory proportional to the max label in fragments, therefore # we relabel them here and use those fragments_relabelled, n, fragment_relabel_map = relabel( fragments, return_backwards_map=True) logger.debug("affs shape: %s", affs.shape) logger.debug("fragments shape: %s", fragments.shape) logger.debug("fragments num: %d", n) # convert affs to float32 ndarray with values between 0 and 1 affs = affs.to_ndarray()[0:3] if affs.dtype == np.uint8: affs = affs.astype(np.float32)/255.0 # So far, 'rag' does not contain any edges belonging to write_roi (there # might be a few edges from neighboring blocks, though). Run waterz until # threshold 0 to get the waterz RAG, which tells us which nodes are # neighboring. Use this to populate 'rag' with edges. Then run waterz for # the given threshold. # for efficiency, we create one waterz call with both thresholds generator = waterz.agglomerate( affs=affs, thresholds=[0, threshold], fragments=fragments_relabelled, scoring_function=merge_function, discretize_queue=256, return_merge_history=True, return_region_graph=True) # add edges to RAG _, _, initial_rag = next(generator) for edge in initial_rag: u, v = fragment_relabel_map[edge['u']], fragment_relabel_map[edge['v']] # this might overwrite already existing edges from neighboring blocks, # but that's fine, we only write attributes for edges within write_roi rag.add_edge(u, v, merge_score=None, agglomerated=True) # agglomerate fragments using affs _, merge_history, _ = next(generator) # cleanup generator for _, _, _ in generator: pass # create a merge tree from the merge history merge_tree = MergeTree(fragment_relabel_map) for merge in merge_history: a, b, c, score = merge['a'], merge['b'], merge['c'], merge['score'] merge_tree.merge( fragment_relabel_map[a], fragment_relabel_map[b], fragment_relabel_map[c], score) # mark edges in original RAG with score at time of merging logger.debug("marking merged edges...") num_merged = 0 for u, v, data in rag.edges(data=True): merge_score = merge_tree.find_merge(u, v) data['merge_score'] = merge_score if merge_score is not None: num_merged += 1 logger.info("merged %d edges", num_merged) # write back results (only within write_roi) logger.debug("writing to DB...") rag.write_edges(block.write_roi)
def segment(args): """Run segmentation on contiguous block of affinities from CV Args: args: ArgParse object from main """ bbox_start = Vec(*args.bbox_start) bbox_size = Vec(*args.bbox_size) chunk_size = Vec(*args.chunk_size) bbox = Bbox(bbox_start, bbox_start + bbox_size) src_cv = CloudVolume(args.src_path, fill_missing=True, parallel=args.parallel) info = CloudVolume.create_new_info( num_channels=1, layer_type='segmentation', data_type='uint64', encoding='raw', resolution=src_cv.info['scales'][args.mip]['resolution'], voxel_offset=bbox_start, chunk_size=chunk_size, volume_size=bbox_size, mesh='mesh_mip_{}_err_{}'.format(args.mip, args.max_simplification_error)) dst_cv = CloudVolume(args.dst_path, info=info, parallel=args.parallel) dst_cv.provenance.description = 'ws+agg using waterz' dst_cv.provenance.processing.append({ 'method': { 'task': 'watershed+agglomeration', 'src_path': args.src_path, 'dst_path': args.dst_path, 'mip': args.mip, 'shape': bbox_size.tolist(), 'bounds': [ bbox.minpt.tolist(), bbox.maxpt.tolist(), ], }, 'by': args.owner, 'date': strftime('%Y-%m-%d%H:%M %Z'), }) dst_cv.provenance.owners = [args.owner] dst_cv.commit_info() dst_cv.commit_provenance() if args.segment: print('Downloading affinities') aff = src_cv[bbox.to_slices()] aff = np.transpose(aff, (3, 0, 1, 2)) aff = np.ascontiguousarray(aff, dtype=np.float32) thresholds = [args.threshold] print('Starting ws+agg') seg_gen = waterz.agglomerate(aff, thresholds) seg = next(seg_gen) print('Deleting affinities') del aff print('Uploading segmentation') dst_cv[bbox.to_slices()] = seg if args.mesh: print('Starting meshing') with LocalTaskQueue(parallel=args.parallel) as tq: tasks = tc.create_meshing_tasks( layer_path=args.dst_path, mip=args.mip, shape=args.chunk_size, simplification=True, max_simplification_error=args.max_simplification_error, progress=True) tq.insert_all(tasks) tasks = tc.create_mesh_manifest_tasks(layer_path=args.dst_path, magnitude=args.magnitude) tq.insert_all(tasks) print("Meshing complete")
def watershed_from_affinities(affs, fragments_in_xy=False, return_seeds=False, epsilon_agglomerate=0): '''Extract initial fragments from affinities using a watershed transform. Returns the fragments and the maximal ID in it.''' if affs.dtype == np.uint8: logger.info("Assuming affinities are in [0,255]") max_affinity_value = 255.0 affs = affs.astype(np.float32) else: max_affinity_value = 1.0 if fragments_in_xy: mean_affs = 0.5 * (affs[1] + affs[2]) depth = mean_affs.shape[0] fragments = np.zeros(mean_affs.shape, dtype=np.uint64) if return_seeds: seeds = np.zeros(mean_affs.shape, dtype=np.uint64) id_offset = 0 for z in range(depth): boundary_mask = mean_affs[z] > 0.5 * max_affinity_value boundary_distances = distance_transform_edt(boundary_mask) ret = watershed_from_boundary_distance(boundary_distances, return_seeds=return_seeds, id_offset=id_offset) fragments[z] = ret[0] if return_seeds: seeds[z] = ret[2] id_offset = ret[1] ret = (fragments, id_offset) if return_seeds: ret += (seeds, ) else: boundary_mask = np.mean(affs, axis=0) > 0.5 * max_affinity_value boundary_distances = distance_transform_edt(boundary_mask) ret = watershed_from_boundary_distance(boundary_distances, return_seeds) if epsilon_agglomerate > 0: logger.info("Performing initial fragment agglomeration until %f", epsilon_agglomerate) generator = waterz.agglomerate( affs=affs / max_affinity_value, thresholds=[epsilon_agglomerate], fragments=fragments, scoring_function= 'OneMinus<HistogramQuantileAffinity<RegionGraphType, 25, ScoreValue, 256, false>>', discretize_queue=256, return_merge_history=False, return_region_graph=False) fragments[:] = next(generator) # cleanup generator for _ in generator: pass return ret
def waterz(affs, thresholds, output_prefix='./', merge_function=None, gt=None, gt_border=25 / 4.0, fragments=None, discretize_queue=256, fragments_mask=None, aff_threshold=[0.0001, 0.9999], return_seg=True, save_record=False): # affs shape: 3*z*y*x thresholds = list(thresholds) print("waterz at thresholds " + str(thresholds)) if fragments is None: fragments = watershed(affs, 'maxima_distance') if fragments_mask is not None: fragments[fragments_mask == False] = 0 outs = [] if gt is not None and gt_border != 0: gt = create_border_mask(gt, gt_border, np.uint64(0)) for i, out in enumerate( agglomerate(affs, thresholds, gt=gt, aff_threshold_low=aff_threshold[0], aff_threshold_high=aff_threshold[1], fragments=fragments, scoring_function=getScoreFunc(merge_function), discretize_queue=discretize_queue, force_rebuild=True)): threshold = thresholds[i] output_basename = output_prefix + merge_function + '_%.2f' % threshold if gt is not None: seg = out[0] else: seg = out if return_seg: outs.append(seg.copy()) else: print("Storing segmentation...") writeh5(output_basename + '.hdf', 'main', seg) if gt is not None: metrics = out[1] print("Storing record...") record = { 'threshold': threshold, 'merge_function': merge_function, 'discretize_queue': discretize_queue, 'voi_split': metrics['V_Info_split'], 'voi_merge': metrics['V_Info_merge'], 'rand_split': metrics['V_Rand_split'], 'rand_merge': metrics['V_Rand_merge'], } if save_record == True: with open(output_basename + '.json', 'w') as f: json.dump(record, f) if return_seg: return outs
# raw_slice=random_provider.random_provider_raw((16,128,128),raw) # # raw_slice=np.reshape(raw_slice,(1,16,128,128,1)) # # aff=model.predict(raw_slice) # # aff=np.einsum("bzxyc->bczxy",aff) # # aff=aff[0] # # print(np.shape(aff)) # # ########################## seg = w.agglomerate( aff, thresholds=[thresh], gt=gt ) #, aff_threshold_low=thresh_l,aff_threshold_high=thresh_h,scoring_function="OneMinus<MinAffinity<RegionGraphType, ScoreValue>>") for segmentation in seg: seg = segmentation seg = seg[0] print("FOUND %i UNIQUE NEURONS" % (np.shape(np.unique(seg))[0] - 2)) #seg=filter.top_n(seg,400) seg = add_colors.add_colors(seg) #print(np.shape(seg))
temp0 = pred[0] temp1 = pred[1] temp2 = pred[2] pred[0] = temp2 pred[1] = temp1 pred[2] = temp0 aff = np.asarray(aff, dtype=np.float32) pred = np.asarray(pred, dtype=np.float32) gt = np.asarray(gt, dtype=np.uint32) seg = w.agglomerate(pred, thresholds=[thresh], gt=gt, aff_threshold_low=thresh_l, aff_threshold_high=thresh_h) print("predicted averages: ", np.average(pred[0]), np.average(pred[1]), np.average(pred[2])) print("actual averages: ", np.average(aff[0]), np.average(aff[1]), np.average(aff[2])) for segmentation in seg: seg = segmentation seg = seg[0] print("FOUND %i UNIQUE NEURONS" % (np.shape(np.unique(seg))[0] - 2))