def segment_volume(self, pmaps):
if self.ws_2D:
# WS in 2D
ws = self.ws_dt_2D(pmaps)
else:
# WS in 3D
ws, _ = distance_transform_watershed(pmaps,
self.ws_threshold,
self.ws_sigma,
sigma_weights=self.ws_w_sigma,
min_size=self.ws_minsize)
rag = compute_rag(ws, 1)
# Computing edge features
features = nrag.accumulateEdgeMeanAndLength(
rag, pmaps, numberOfThreads=1) # DO NOT CHANGE numberOfThreads
probs = features[:, 0] # mean edge prob
edge_sizes = features[:, 1]
# Prob -> edge costs
costs = transform_probabilities_to_costs(probs,
edge_sizes=edge_sizes,
beta=self.beta)
# Creating graph
graph = nifty.graph.undirectedGraph(rag.numberOfNodes)
graph.insertEdges(rag.uvIds())
# Solving Multicut
node_labels = multicut_kernighan_lin(graph, costs)
return nifty.tools.take(node_labels, ws)
def segment_volume_mc(pmaps,
threshold=0.4,
sigma=2.0,
beta=0.6,
ws=None,
sp_min_size=100):
if ws is None:
ws = distance_transform_watershed(pmaps,
threshold,
sigma,
min_size=sp_min_size)[0]
rag = compute_rag(ws, 1)
features = nrag.accumulateEdgeMeanAndLength(rag, pmaps, numberOfThreads=1)
probs = features[:, 0] # mean edge prob
edge_sizes = features[:, 1]
costs = transform_probabilities_to_costs(probs,
edge_sizes=edge_sizes,
beta=beta)
graph = nifty.graph.undirectedGraph(rag.numberOfNodes)
graph.insertEdges(rag.uvIds())
node_labels = multicut_kernighan_lin(graph, costs)
return nifty.tools.take(node_labels, ws)
def segment_mc(pred, seg, delta):
rag = feats.compute_rag(seg)
edge_probs = embed.edge_probabilities_from_embeddings(
pred, seg, rag, delta)
edge_sizes = feats.compute_boundary_mean_and_length(rag, pred[0])[:, 1]
costs = mc.transform_probabilities_to_costs(edge_probs,
edge_sizes=edge_sizes)
mc_seg = mc.multicut_kernighan_lin(rag, costs)
mc_seg = feats.project_node_labels_to_pixels(rag, mc_seg)
return mc_seg
def multicut_from_probas(segmentation, edges, edge_weights):
rag = compute_rag(segmentation)
edge_dict = dict(zip(list(map(tuple, edges)), edge_weights))
costs = np.empty(len(edge_weights))
for i, neighbors in enumerate(rag.uvIds()):
if tuple(neighbors) in edge_dict:
costs[i] = edge_dict[tuple(neighbors)]
else:
costs[i] = edge_dict[(neighbors[1], neighbors[0])]
costs = transform_probabilities_to_costs(costs)
node_labels = multicut_kernighan_lin(rag, costs)
return project_node_labels_to_pixels(rag, node_labels).squeeze()
def supervoxel_merging(mem, sv, beta=0.5, verbose=False):
rag = feats.compute_rag(sv)
costs = feats.compute_boundary_features(rag, mem)[:, 0]
edge_sizes = feats.compute_boundary_mean_and_length(rag, mem)[:, 1]
costs = mc.transform_probabilities_to_costs(costs,
edge_sizes=edge_sizes,
beta=beta)
node_labels = mc.multicut_kernighan_lin(rag, costs)
segmentation = feats.project_node_labels_to_pixels(rag, node_labels)
return segmentation
affs = np.transpose(affs.cpu().numpy(), (1, 0, 2, 3))
gt_affs = np.transpose(gt_affs.cpu().numpy(), (1, 0, 2, 3))
seg = seg.cpu().numpy()
gt_seg = gt_seg.cpu().numpy()
boundary_input = np.mean(affs, axis=0)
gt_boundary_input = np.mean(gt_affs, axis=0)
rag = feats.compute_rag(seg)
# edges rag.uvIds() [[1, 2], ...]
costs = feats.compute_affinity_features(rag, affs, offsets)[:, 0]
gt_costs = calculate_gt_edge_costs(rag.uvIds(), seg.squeeze(),
gt_seg.squeeze())
edge_sizes = feats.compute_boundary_mean_and_length(rag, boundary_input)[:, 1]
gt_edge_sizes = feats.compute_boundary_mean_and_length(rag,
gt_boundary_input)[:, 1]
costs = mc.transform_probabilities_to_costs(costs, edge_sizes=edge_sizes)
gt_costs = mc.transform_probabilities_to_costs(gt_costs, edge_sizes=edge_sizes)
node_labels = mc.multicut_kernighan_lin(rag, costs)
gt_node_labels = mc.multicut_kernighan_lin(rag, gt_costs)
segmentation = feats.project_node_labels_to_pixels(rag, node_labels)
gt_segmentation = feats.project_node_labels_to_pixels(rag, gt_node_labels)
plt.imshow(
np.concatenate(
(gt_segmentation.squeeze(), segmentation.squeeze(), seg.squeeze()),
axis=1))
plt.show()
def refine_seg(raw,
seeds,
restrict_to_seeds=True,
restrict_to_bb=False,
return_intermediates=False):
pred = get_prediction(raw, cache=False)
n_threads = 1
# make watershed
ws, _ = stacked_watershed(pred,
threshold=.5,
sigma_seeds=1.,
n_threads=n_threads)
rag = compute_rag(ws, n_threads=n_threads)
edge_feats = compute_boundary_mean_and_length(rag,
pred,
n_threads=n_threads)
edge_feats, edge_sizes = edge_feats[:, 0], edge_feats[:, 1]
z_edges = compute_z_edge_mask(rag, ws)
edge_costs = compute_edge_costs(edge_feats,
beta=.4,
weighting_scheme='xyz',
edge_sizes=edge_sizes,
z_edge_mask=z_edges)
# make seeds and map them to edges
bb = tuple(
slice(sh // 2 - ha // 2, sh // 2 + ha // 2)
for sh, ha in zip(pred.shape, seeds.shape))
seeds[seeds < 0] = 0
seeds = vigra.analysis.labelVolumeWithBackground(seeds.astype('uint32'))
seed_ids = np.unique(seeds)
seed_mask = binary_erosion(seeds, iterations=2)
seeds_new = seeds.copy()
seeds_new[~seed_mask] = 0
seed_ids_new = np.unique(seeds_new)
for seed_id in seed_ids:
if seed_id in seed_ids_new:
continue
seeds_new[seeds == seed_id] = seed_id
seeds_full = np.zeros(pred.shape, dtype=seeds.dtype)
seeds_full[bb] = seeds
seeds = seeds_full
seed_labels = compute_maximum_label_overlap(ws, seeds, ignore_zeros=True)
edge_ids = rag.uvIds()
labels_u = seed_labels[edge_ids[:, 0]]
labels_v = seed_labels[edge_ids[:, 1]]
seed_mask = np.logical_and(labels_u != 0, labels_v != 0)
same_seed = np.logical_and(seed_mask, labels_u == labels_v)
diff_seed = np.logical_and(seed_mask, labels_u != labels_v)
max_att = edge_costs.max() + .1
max_rep = edge_costs.min() - .1
edge_costs[same_seed] = max_att
edge_costs[diff_seed] = max_rep
# run multicut
node_labels = multicut_kernighan_lin(rag, edge_costs)
if restrict_to_seeds:
seed_nodes = np.unique(node_labels[seed_labels > 0])
node_labels[~np.isin(node_labels, seed_nodes)] = 0
vigra.analysis.relabelConsecutive(node_labels, out=node_labels)
seg = project_node_labels_to_pixels(rag, node_labels, n_threads=n_threads)
if restrict_to_bb:
bb_mask = np.zeros(seg.shape, dtype='bool')
bb_mask[bb] = 1
seg[~bb_mask] = 0
if return_intermediates:
return pred, ws, seeds, seg
else:
return seg
