def test_ambiguous_unkown(): embedding = np.array([[[[0, 1]]]], dtype=np.float32) segmentation = np.array([[[[-1, 0]]]], dtype=np.int64) # number of positive pairs: 0 # number of negative pairs: 1 # total number of pairs: 0 # loss on positive pairs: 0 # loss on negative pairs: 1 # total loss: 1 embedding = torch.from_numpy(embedding).float() segmentation = torch.from_numpy(segmentation) loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, alpha=2, balance=True, add_coordinates=False, constrained_emst=True, ) assert pytest.approx(float(loss)) == 1 assert pytest.approx(np.sum(distances.numpy())) == 1 embedding = np.array([[[[0, 1]]]], dtype=np.float32) segmentation = np.array([[[[0, -1]]]], dtype=np.int64) # number of positive pairs: 0 # number of negative pairs: 1 # total number of pairs: 0 # loss on positive pairs: 0 # loss on negative pairs: 1 # total loss: 1 embedding = torch.from_numpy(embedding).float() segmentation = torch.from_numpy(segmentation) loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, alpha=2, balance=True, add_coordinates=False, constrained_emst=True, ) assert pytest.approx(float(loss)) == 1 assert pytest.approx(np.sum(distances.numpy())) == 1
def test_background(): embedding = np.array([[[[0, 1, 2], [4, 5, 6], [8, 9, 10]]]], dtype=np.float32) segmentation = np.array([[[[1, 1, 1], [0, 0, 0], [3, 3, 3]]]], dtype=np.int64) # number of positive pairs: 2*3 = 6 # number of negative pairs: 3*3*3 = 27 # number of background pairs: 3 # total number of pairs (without background pairs): 33 # loss on positive pairs: 6*1 = 6 # loss on negative pairs: 27*2^2 = 108 # total loss = 114 # total loss per pair = 3.455 # total loss per pos pair = 1 # total loss per neg pair = 4 embedding = torch.from_numpy(embedding).float() segmentation = torch.from_numpy(segmentation) loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, alpha=4, add_coordinates=False, balance=False, ) assert pytest.approx(float(loss), 1e-3) == 3.4545 assert pytest.approx(np.sum(distances.numpy()), ) == 10
def test_quadrupel_loss(): embedding = np.array([[[[0, 1, 2], [4, 5, 6], [8, 9, 10]]]], dtype=np.float32) segmentation = np.array([[[[1, 1, 1], [2, 2, 2], [3, 3, 3]]]], dtype=np.int64) # number of positive pairs: 3*3 = 9 # number of negative pairs: 3*3*3 = 27 # number of quadrupels: 9*27 = 243 # loss per quadrupel: max(0, d(p) - d(n) + alpha)^2 = (1 - 2 + 3)^2 = 4 embedding = torch.from_numpy(embedding).float() segmentation = torch.from_numpy(segmentation) loss, emst, edges_u, edges_v, distances = ultrametric_loss( embedding, segmentation, alpha=3, add_coordinates=False, quadrupel_loss=True, ) assert pytest.approx(float(loss)) == 4.0 assert pytest.approx(np.sum(distances.numpy())) == 10
def test_constrained(): embedding = np.array([[[[0, 1, 101], [2, 3, 4], [5, 6, 7]]]], dtype=np.float32) segmentation = np.array([[[[1, 1, 1], [2, 2, 2], [3, 3, 3]]]], dtype=np.int64) # number of positive pairs: 3*3 = 9 # number of negative pairs: 3*3*3 = 27 # total number of pairs: 9*8/2 = 36 # loss on positive pairs: 6*1 + 1 + 2*100^2 = 20007 # loss on negative pairs: 27*1 = 27 # total loss: 1/9*20007 + 1/27*27 = 2224 embedding = torch.from_numpy(embedding).float() segmentation = torch.from_numpy(segmentation) loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, alpha=2, add_coordinates=False, constrained_emst=True, ) assert pytest.approx(float(loss)) == 2224 assert pytest.approx(np.sum(distances.numpy())) == 107
def test_add_coordinates(): embedding = np.array([[[[0, 1, 2], [3, 4, 5], [6, 7, 8]]]], dtype=np.float32) segmentation = np.array([[[[1, 1, 1], [2, 2, 2], [3, 3, 3]]]], dtype=np.int64) # number of positive pairs: 3*3 = 9 # number of negative pairs: 3*3*3 = 27 # total number of pairs: 9*8/2 = 36 # loss on positive pairs: 9*2 = 18 # loss on negative pairs: 27*0 = 27 # total loss = 18 # total loss per edge = 1 embedding = torch.from_numpy(embedding).float() segmentation = torch.from_numpy(segmentation) loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, alpha=2, add_coordinates=True, balance=False, ) assert pytest.approx(float(loss)) == 18 / 36 assert (pytest.approx(np.sum(distances.numpy())) == 6 * (1**2 + 1**2 + 0**2)**0.5 + 2 * (2**2 + 1**2 + 1**2)**0.5) loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, alpha=2, add_coordinates=False, ) assert pytest.approx(float(loss)) == 2.0 assert pytest.approx(np.sum(distances.numpy())) == 8
def test_zero(): embedding = np.zeros((1, 3, 10, 10, 10), dtype=np.float32) segmentation = np.ones((1, 1, 10, 10, 10), dtype=np.int64) embedding = torch.from_numpy(embedding).float() segmentation = torch.from_numpy(segmentation) loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, add_coordinates=False, ) assert pytest.approx(float(loss)) == 0 assert pytest.approx(np.sum(distances.numpy())) == 0
def forward( self, input, target, mask, neighborhood=None, neighborhood_mask=None, neighborhood_target=None, ): # ultrametric loss is computed on cpu, and no gpu accelerated operations # use the outputs. input = input.cpu() target = target.cpu() mask = mask.cpu() loss, emst, edges_u, edges_v, dist, ratio_pos, ratio_neg = ultrametric_loss( embedding=input, gt_seg=target, mask=mask, alpha=self.alpha, add_coordinates=self.add_coordinates, coordinate_scale=self.coordinate_scale, balance=self.balance, quadrupel_loss=self.quadrupel_loss, constrained_emst=self.constrained_emst, ) if self.aux_task: neighborhood_loss = self.mse_loss(neighborhood * neighborhood_mask, neighborhood_target).cpu() loss = neighborhood_loss * loss return ( loss, emst, edges_u, edges_v, dist, ratio_pos, ratio_neg, neighborhood_loss, ) return loss, emst, edges_u, edges_v, dist, ratio_pos, ratio_neg
def test_gradients(): """ Naieve test that trains a simple network on a simple example and shows that the loss decreases. Better test would be to calculate the actual expected gradients resulting from a specific loss, and show that the back propogation is working as expected. """ n = 50 embedder = torch.nn.Sequential( torch.nn.Conv2d(2, 3, [3, 3], padding=1), torch.nn.ReLU(), torch.nn.Conv2d(3, 1, [3, 3], padding=1), torch.nn.Tanh(), ) raw = torch.Tensor([[ [ [0, 0, 0, 0, 0, 0, 0, 0], [0, 2, 0, 0, 0, 0, 2, 0], [0, 1, 0, 0, 0, 0, 1, 0], [0, 1, 0, 0, 0, 0, 1, 0], [0, 1, 0, 0, 0, 0, 1, 0], [0, 1, 0, 0, 0, 0, 1, 0], [0, 2, 0, 0, 0, 0, 2, 0], [0, 0, 0, 0, 0, 0, 0, 0], ], [ [0, 0, 0, 0, 0, 0, 0, 0], [0, 2, 0, 0, 0, 0, 2, 0], [0, 1, 0, 0, 0, 0, 1, 0], [0, 1, 0, 0, 0, 0, 1, 0], [0, 1, 0, 0, 0, 0, 1, 0], [0, 1, 0, 0, 0, 0, 1, 0], [0, 2, 0, 0, 0, 0, 2, 0], [0, 0, 0, 0, 0, 0, 0, 0], ], ]]) mask = torch.Tensor([[[ [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1], ]]]).bool() gt_labels = torch.Tensor([[[ [0, 0, 0, 0, 0, 0, 0, 0], [0, 3, 0, 0, 0, 0, 4, 0], [0, 1, 0, 0, 0, 0, 2, 0], [0, 1, 0, 0, 0, 0, 2, 0], [0, 1, 0, 0, 0, 0, 2, 0], [0, 1, 0, 0, 0, 0, 2, 0], [0, 5, 0, 0, 0, 0, 6, 0], [0, 0, 0, 0, 0, 0, 0, 0], ]]]).to(torch.int64) embeddings = embedder(raw) optimizer = torch.optim.Adam(embedder.parameters()) assert embeddings.shape == (1, 1, 8, 8) last_loss = float("inf") embedding_vol = np.zeros([n, 8, 8]) for i in range(n): optimizer.zero_grad() embeddings = embedder(raw) embedding_vol[i, :, :] = embeddings.detach().numpy()[0, 0, :, :] loss, *_ = ultrametric_loss( embeddings, gt_labels, mask, alpha=0.5, coordinate_scale=0.2, constrained_emst=False, ) assert float(loss) < float(last_loss) assert raw.sum() == 32 last_loss = loss loss.backward() optimizer.step() np.save("embeddings.npy", embedding_vol)
def test_large_example(): """ nodes: labels: [[a, b, c], [[1, 0, 1], [d, e, f], [2, 2, -1], [g, h, i], [3, 3, 3], [j, k, l]] [0, 0, 0]] mst edges (unconstrained), dist, ratio_pos, ratio_neg de, 0 0 0 ef, 0 1 0 jk, 0 0 0 ab, 0.5 0 1 bc, 0.5 1 1 gh, 1 0 0 hi, 1 1 0 kl, 1 2 0 dj, 1.5 0 9 fj, 2 0 13 gl, 2 0 24 mst edges (constrained), dist, ratio_pos, ratio_neg de, 0 1 0 jk, 0 0 0 ac, 1 1 0 gh, 1 1 0 hi, 1 2 0 kl, 1 0 0 bk, 4 0 0 ef, 0 0 0 ab, 0.5 0 8 dj, 1.5 0 16 gl, 2 0 24 """ num_pos = 5 num_neg = 48 num_edges = num_pos + num_neg embedding = np.array( [[[[0.5, 1, 1.5], [3.5, 3.5, 3.5], [8, 9, 10], [5, 5, 6]]]], dtype=np.float32) segmentation = np.array([[[[1, 0, 1], [2, 2, -1], [3, 3, 3], [0, 0, 0]]]], dtype=np.int64) embedding = torch.from_numpy(embedding).float() segmentation = torch.from_numpy(segmentation) loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, constrained_emst=False, balance=False, alpha=2, add_coordinates=False, ) # numerator = dist**2 * ratio_pos + (alpha-dist)**2 * ratio_neg assert pytest.approx(float(loss)) == (3.25 + 6.75) / num_edges assert pytest.approx(np.sum(distances.numpy())) == 9.5 loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, constrained_emst=True, balance=False, alpha=2, add_coordinates=False, ) assert pytest.approx(float(loss)) == (4 + 22) / num_edges assert pytest.approx(np.sum(distances.numpy())) == 12 loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, constrained_emst=False, balance=True, alpha=2, add_coordinates=False, ) assert pytest.approx(float(loss)) == (3.25 / num_pos + 6.75 / num_neg) assert pytest.approx(np.sum(distances.numpy())) == 9.5 loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, constrained_emst=True, balance=True, alpha=2, add_coordinates=False, ) assert pytest.approx(float(loss)) == (4 / num_pos + 22 / num_neg) assert pytest.approx(np.sum(distances.numpy())) == 12
def test_constrained_mask(): embedding = np.array([[[[0, 1, 2], [3, 4, 5], [6, 7, 8]]]], dtype=np.float32) segmentation = np.array([[[[1, 1, 1], [2, 2, 2], [3, 3, 3]]]], dtype=np.int64) embedding = torch.from_numpy(embedding).float() segmentation = torch.from_numpy(segmentation) # empty mask mask = np.zeros((1, 1, 3, 3), dtype=np.bool) mask = torch.from_numpy(mask) loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, mask=mask, constrained_emst=True, alpha=2, add_coordinates=False, ) assert pytest.approx(float(loss)) == 0.0 assert pytest.approx(np.sum(distances.numpy())) == 0 # mask with only one point mask = np.zeros((1, 1, 3, 3), dtype=np.bool) mask[0, 0, 1, 1] = True mask = torch.from_numpy(mask) loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, mask=mask, constrained_emst=True, alpha=2, add_coordinates=False, ) assert pytest.approx(float(loss)) == 0.0 assert pytest.approx(np.sum(distances.numpy())) == 0 # mask with two points mask = np.zeros((1, 1, 3, 3), dtype=np.bool) mask[0, 0, 1, 1] = True mask[0, 0, 0, 0] = True mask = torch.from_numpy(mask) loss, emst, edges_u, edges_v, distances, *_ = ultrametric_loss( embedding, segmentation, mask=mask, constrained_emst=True, alpha=5, add_coordinates=False, ) assert pytest.approx(float(loss)) == 1.0 assert pytest.approx(np.sum(distances.numpy())) == 4.0