def get_edge_features_1d(sp_seg, offsets, affinities):
offsets_3d = []
for off in offsets:
offsets_3d.append([0] + off)
rag = feats.compute_rag(np.expand_dims(sp_seg, axis=0))
edge_feat = feats.compute_affinity_features(
rag, np.expand_dims(affinities, axis=1), offsets_3d)[:, :]
return edge_feat
pin_memory=True)
neighbors, nodes, seg, gt_seg, affs, gt_affs = next(iter(dloader_disc))
offsets = [[0, 0, -1], [0, -1, 0], [0, -3, 0], [0, 0, -3]]
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(