def get_current_soln(self, edge_weights): p_min = 0.001 p_max = 1. segmentations = [] node_labels = [] n_lbl_offs = [0] for i in range(1, len(self.e_offs)): probs = edge_weights[self.e_offs[i-1]:self.e_offs[i]] costs = (p_max - p_min) * probs + p_min costs = (torch.log((1. - costs) / costs)).detach().cpu().numpy() _node_labels = elf.segmentation.multicut.multicut_decomposition(self.rags[i-1], costs, internal_solver='greedy-additive', n_threads=4) mc_seg = project_node_labels_to_pixels(self.rags[i-1], _node_labels).squeeze() mc_seg = torch.from_numpy(mc_seg.astype(np.long)).to(self.device) segmentations.append(mc_seg) node_labels.append(torch.from_numpy(_node_labels.astype(np.long)).to(self.device) + n_lbl_offs[i-1]) n_lbl_offs.append(_node_labels.max()) node_labels = torch.cat(node_labels) object_node_mask = (node_labels[None] == torch.unique(node_labels)[:, None]).long() object_nodes = object_node_mask * torch.arange(1, len(node_labels) + 1, device=self.device)[None] nodes_per_edge = (object_nodes[:, :, None, None] == (self.edge_ids[None, None] + 1)).long().sum(1).sum(1) object_edge_ind_critic = torch.nonzero(nodes_per_edge == 2)[:, 1] object_edges_actor = (nodes_per_edge >= 1).long() # G = nx.path_graph(4) # G.add_edges_from(self.edge_ids[:, object_edge_ind_critic].T.tolist()) # nx.connected_components(G) # print("############", len(list(nx.connected_components(G))) + (object_node_mask.sum(1)==1).sum() - object_node_mask.shape[0]) return torch.stack(segmentations, dim=0), object_edge_ind_critic, object_node_mask, object_edges_actor
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
def get_multicut_sln(rag, edge_weights): p_min = 1e-5 p_max = 1. probs = edge_weights.copy() costs = (p_max - p_min) * probs + p_min costs = (np.log((1. - costs) / costs.clip(1e-10, 1))) node_labels = multicut_decomposition(rag, costs, internal_solver='kernighan-lin', n_threads=4) # greedy-additive node_labels = vigra.analysis.relabelConsecutive(node_labels, keep_zeros=False, start_label=0)[0] return project_node_labels_to_pixels(rag, node_labels).squeeze()
def segment_volume_lmc_from_seg(boundary_pmaps, nuclei_seg, threshold=0.4, sigma=2.0, sp_min_size=100): watershed = distance_transform_watershed(boundary_pmaps, threshold, sigma, min_size=sp_min_size)[0] # compute the region adjacency graph rag = compute_rag(watershed) # compute the edge costs features = compute_boundary_mean_and_length(rag, boundary_pmaps) costs, sizes = features[:, 0], features[:, 1] # transform the edge costs from [0, 1] to [-inf, inf], which is # necessary for the multicut. This is done by intepreting the values # as probabilities for an edge being 'true' and then taking the negative log-likelihood. # in addition, we weight the costs by the size of the corresponding edge # we choose a boundary bias smaller than 0.5 in order to # decrease the degree of over segmentation boundary_bias = .6 costs = transform_probabilities_to_costs(costs, edge_sizes=sizes, beta=boundary_bias) max_cost = np.abs(np.max(costs)) lifted_uvs, lifted_costs = lifted_problem_from_segmentation( rag, watershed, nuclei_seg, overlap_threshold=0.2, graph_depth=4, same_segment_cost=5 * max_cost, different_segment_cost=-5 * max_cost) # solve the full lifted problem using the kernighan lin approximation introduced in # http://openaccess.thecvf.com/content_iccv_2015/html/Keuper_Efficient_Decomposition_of_ICCV_2015_paper.html node_labels = lmc.lifted_multicut_kernighan_lin(rag, costs, lifted_uvs, lifted_costs) lifted_segmentation = project_node_labels_to_pixels(rag, node_labels) return lifted_segmentation
def get_current_soln(self, edge_weights): p_min = 0.001 p_max = 1. segmentations = [] for i in range(1, len(self.e_offs)): probs = edge_weights[self.e_offs[i - 1]:self.e_offs[i]] probs -= probs.min() probs /= probs.max() costs = (p_max - p_min) * probs + p_min costs = (torch.log((1. - costs) / costs)).detach().cpu().numpy() node_labels = elf.segmentation.multicut.multicut_decomposition( self.rags[i - 1], costs, internal_solver='greedy-additive', n_threads=4) mc_seg = project_node_labels_to_pixels(self.rags[i - 1], node_labels).squeeze() mc_seg = torch.from_numpy(mc_seg.astype(np.long)).to(self.device) # mask = mc_seg[None] == torch.unique(mc_seg)[:, None, None] # mc_seg = (mask * (torch.arange(len(torch.unique(mc_seg)), device=mc_seg.device)[:, None, None] + 1)).sum(0) - 1 segmentations.append(mc_seg) return torch.stack(segmentations, dim=0)
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