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))
