def test_malis_2d_gradient_descent(self):
from affogato.learning import compute_malis_2d
from affogato.segmentation import connected_components
shape = (100, 100)
labels = np.zeros(shape)
for i in range(10):
for j in range(10):
labels[10 * i:10 * (i + 1),
10 * j:10 * (j + 1)] = 10 * i + j + 1
affs = 0.5 * np.ones((2, 100, 100))
offsets = [[-1, 0], [0, -1]]
for epoch in range(40):
loss, grads = compute_malis_2d(affs, labels, offsets)
affs -= 10000 * grads
affs = np.clip(affs, 0, 1)
labels1, _ = connected_components(affs, 0.5)
self.assertEqual(grads.shape, affs.shape)
self.assertTrue(np.allclose(grads, 0))
self.assertEqual(loss, 0)
edges1 = self.seg2edges(labels1)
edges2 = self.seg2edges(labels)
self.assertTrue(np.allclose(edges1, edges2))
def test_cc_2d(self):
from affogato.segmentation import connected_components
shape = (2, 100, 100)
affs = np.random.rand(*shape)
ccs, max_label = connected_components(affs, 0.5)
self.assertEqual(ccs.shape, shape[1:])
self.assertGreater(max_label, 10)
self.assertEqual(max_label, ccs.max())
def cc_segmenter(prediction,
thresholds=[.9, .925, .95, .975, .99],
invert=True):
from affogato.segmentation import connected_components
if invert:
prediction = 1. - prediction
return [
connected_components(prediction, thresh)[0] for thresh in thresholds
]
def input_to_segmentation(self, input_batch, thresh):
dim = input_batch.ndim - 2
# if specified, invert and / or normalize the affinities
if self.invert_affinities:
input_batch = 1. - input_batch
if self.normalize_affinities:
input_batch = input_batch / input_batch.max()
# Compute the segmentation via connected components on the affinities
ccs = np.array([
connected_components(batch[:dim], thresh)[0]
for batch in input_batch
])
# NOTE: we add a singleton channel axis here, which is expected by the arand metrics
return torch.from_numpy(ccs[:, None].astype('int32'))