def test_simple(self):
embedding = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]],
dtype=np.float32).reshape((1, 1, 3, 3))
segmentation = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]],
dtype=np.int64).reshape((1, 3, 3))
# 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*1 = 9
# loss on negative pairs: 27*1 = 27
# total loss = 36
# total loss per edge = 1
embedding = tf.constant(embedding, dtype=tf.float32)
segmentation = tf.constant(segmentation)
loss = ultrametric_loss_op(embedding,
segmentation,
alpha=2,
add_coordinates=False,
balance=False,
name='um_test_simple_unbalanced')
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(loss)
self.assertEqual(loss, 1.0)
self.assertAlmostEqual(np.sum(distances), 8, places=4)
loss = ultrametric_loss_op(embedding,
segmentation,
alpha=2,
add_coordinates=False,
name='um_test_simple_balanced')
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(loss)
self.assertEqual(loss, 2.0)
self.assertAlmostEqual(np.sum(distances), 8, places=4)
def test_quadrupel_loss(self):
embedding = np.array([[0, 1, 2], [4, 5, 6], [8, 9, 10]],
dtype=np.float32).reshape((1, 1, 3, 3))
segmentation = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]],
dtype=np.int64).reshape((1, 3, 3))
# 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 = tf.constant(embedding, dtype=tf.float32)
segmentation = tf.constant(segmentation)
loss = ultrametric_loss_op(embedding,
segmentation,
alpha=3,
add_coordinates=False,
quadrupel_loss=True,
name='um_test_quadrupel_loss')
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(loss)
self.assertEqual(loss, 4.0)
self.assertAlmostEqual(np.sum(distances), 10, places=4)
def test_ambiguous_ambiguous(self):
embedding = np.array([[0, 1]], dtype=np.float32).reshape((1, 1, 1, 2))
segmentation = np.array([[-1, -1]], dtype=np.int64).reshape((1, 1, 2))
# number of positive pairs: 0
# number of negative pairs: 0
# total number of pairs: 0
# loss on positive pairs: 0
# loss on negative pairs: 0
# total loss: 0
embedding = tf.constant(embedding, dtype=tf.float32)
segmentation = tf.constant(segmentation)
loss = ultrametric_loss_op(embedding,
segmentation,
alpha=2,
balance=True,
add_coordinates=False,
constrained_emst=True,
name='um_test_constrained')
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(loss)
self.assertAlmostEqual(loss, 0, places=3)
self.assertAlmostEqual(np.sum(distances), 1, places=4)
def test_constrained(self):
embedding = np.array([[0, 1, 101], [2, 3, 4], [5, 6, 7]],
dtype=np.float32).reshape((1, 1, 3, 3))
segmentation = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]],
dtype=np.int64).reshape((1, 3, 3))
# 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 = tf.constant(embedding, dtype=tf.float32)
segmentation = tf.constant(segmentation)
loss = ultrametric_loss_op(embedding,
segmentation,
alpha=2,
add_coordinates=False,
constrained_emst=True,
name='um_test_constrained')
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(loss)
self.assertAlmostEqual(loss, 2224.0, places=1)
self.assertAlmostEqual(np.sum(distances), 107, places=4)
def test_zero(self):
embedding = np.zeros((3, 10, 10, 10), dtype=np.float32)
segmentation = np.ones((10, 10, 10), dtype=np.int64)
embedding = tf.constant(embedding, dtype=tf.float32)
segmentation = tf.constant(segmentation)
loss = ultrametric_loss_op(embedding,
segmentation,
add_coordinates=False,
name='um_test_zero')
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(loss)
self.assertEqual(loss, 0)
self.assertEqual(np.sum(distances), 0)
def test_background(self):
embedding = np.array([[0, 1, 2], [4, 5, 6], [8, 9, 10]],
dtype=np.float32).reshape((1, 1, 3, 3))
segmentation = np.array([[1, 1, 1], [0, 0, 0], [3, 3, 3]],
dtype=np.int64).reshape((1, 3, 3))
# 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 = tf.constant(embedding, dtype=tf.float32)
segmentation = tf.constant(segmentation)
loss = ultrametric_loss_op(embedding,
segmentation,
alpha=4,
add_coordinates=False,
balance=False,
name='um_test_background')
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(loss)
self.assertAlmostEqual(loss, 3.4545, places=4)
self.assertAlmostEqual(np.sum(distances), 10, places=4)
def test_constrained_mask(self):
embedding = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8]],
dtype=np.float32).reshape((1, 1, 3, 3))
segmentation = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]],
dtype=np.int64).reshape((1, 3, 3))
embedding = tf.constant(embedding, dtype=tf.float32)
segmentation = tf.constant(segmentation)
# empty mask
mask = np.zeros((1, 3, 3), dtype=np.bool)
mask = tf.constant(mask)
loss = ultrametric_loss_op(embedding,
segmentation,
mask=mask,
constrained_emst=True,
alpha=2,
add_coordinates=False,
name='um_test_constrained_mask')
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(loss)
self.assertEqual(loss, 0.0)
self.assertAlmostEqual(np.sum(distances), 0, places=4)
# mask with only one point
mask = np.zeros((1, 3, 3), dtype=np.bool)
mask[0, 1, 1] = True
mask = tf.constant(mask)
loss = ultrametric_loss_op(embedding,
segmentation,
mask=mask,
constrained_emst=True,
alpha=2,
add_coordinates=False,
name='um_test_constrained_mask')
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(loss)
self.assertEqual(loss, 0.0)
self.assertAlmostEqual(np.sum(distances), 0, places=4)
# mask with two points
mask = np.zeros((1, 3, 3), dtype=np.bool)
mask[0, 1, 1] = True
mask[0, 0, 0] = True
mask = tf.constant(mask)
loss = ultrametric_loss_op(embedding,
segmentation,
mask=mask,
constrained_emst=True,
alpha=5,
add_coordinates=False,
name='um_test_constrained_mask')
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(loss)
self.assertEqual(loss, 1.0)
self.assertAlmostEqual(np.sum(distances), 4.0, places=4)
def test_large_example(self):
"""
nodes:
[[a, b, c],
[d, e, f],
[g, h, i],
[j, k, l]]
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 1 0
ac, 1 1 0
gh, 1 2 0
hi, 1 0 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
"""
embedding = np.array(
[[0.5, 1, 1.5], [3.5, 3.5, 3.5], [8, 9, 10], [5, 5, 6]],
dtype=np.float32).reshape((1, 1, 4, 3))
embedding = tf.constant(embedding, dtype=tf.float32)
segmentation = np.array([[1, 0, 1], [2, 2, -1], [3, 3, 3], [0, 0, 0]],
dtype=np.int64).reshape((1, 4, 3))
segmentation = tf.constant(segmentation)
loss_unbalanced_unconstrained = ultrametric_loss_op(
embedding,
segmentation,
constrained_emst=False,
balance=False,
alpha=2,
add_coordinates=False,
name='um_test_unbalanced_unconstrained')
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(
loss_unbalanced_unconstrained)
self.assertAlmostEqual(loss, 10 / 53, places=4)
self.assertAlmostEqual(np.sum(distances), 9.5, places=4)
loss_unbalanced_constrained = ultrametric_loss_op(
embedding,
segmentation,
constrained_emst=True,
balance=False,
alpha=2,
add_coordinates=False,
name="um_test_unbalanced_constrained")
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(
loss_unbalanced_constrained)
self.assertAlmostEqual(loss, 26 / 53, places=4)
self.assertAlmostEqual(np.sum(distances), 12, places=4)
loss_balanced_unconstrained = ultrametric_loss_op(
embedding,
segmentation,
constrained_emst=False,
balance=True,
alpha=2,
add_coordinates=False,
name="um_test_balanced_unconstrained")
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(
loss_balanced_unconstrained)
self.assertAlmostEqual(loss, 0.790625, places=4)
self.assertAlmostEqual(np.sum(distances), 9.5, places=4)
loss_balanced_constrained = ultrametric_loss_op(
embedding,
segmentation,
constrained_emst=True,
balance=True,
alpha=2,
add_coordinates=False,
name="um_test_balanced_constrained")
with tf.Session() as s:
s.run(tf.global_variables_initializer())
loss, emst, edges_u, edges_v, distances = s.run(
loss_balanced_constrained)
self.assertAlmostEqual(loss, 1.258333, places=4)
self.assertAlmostEqual(np.sum(distances), 12, places=4)
def add_loss(graph):
output_shape = np.array(config["OUTPUT_SHAPE"])
voxel_size = np.array(config["VOXEL_SIZE"], dtype=int)
max_filter_size = (np.array(config["MAX_FILTER_SIZE"], dtype=int) //
voxel_size).tolist()
maxima_threshold = np.array(config["MAXIMA_THRESHOLD"], dtype=int)
voxel_size = np.array(config["VOXEL_SIZE"], dtype=int)
maxima_threshold = (maxima_threshold / voxel_size[0]).tolist()
coordinate_scale = config["COORDINATE_SCALE"]
alpha = config["ALPHA"]
constrained = config["CONSTRAINED"]
# k, h, w
embedding = graph.get_tensor_by_name(mknet_tensor_names["embedding"])
# h, w
fg_pred = graph.get_tensor_by_name(mknet_tensor_names["fg_pred"])
# h, w
gt_labels = graph.get_tensor_by_name(mknet_tensor_names["gt_labels"])
# h, w
_, maxima = max_detection(
tf.reshape(fg_pred, (1, *output_shape, 1)),
window_size=(1, *max_filter_size),
threshold=maxima_threshold,
)
# maxima = tf.Print(maxima, [maxima.name])
embedding_shape = tuple(embedding.get_shape().as_list())
embedding_dims = len(embedding_shape)
if embedding_dims == 2:
# 1, k, h, w
embedding = tf.reshape(embedding, (1, ) +
tuple(embedding.get_shape().as_list()))
# k, 1, h, w
embedding = tf.transpose(embedding, perm=[1, 0, 2, 3])
elif embedding_dims == 3:
pass
um_loss, emst, edges_u, edges_v, _ = ultrametric_loss_op(
embedding,
gt_labels,
mask=maxima,
coordinate_scale=coordinate_scale,
alpha=alpha,
constrained_emst=constrained,
)
ratio_pos = graph.get_tensor_by_name(ratio_pos_name)
ratio_neg = graph.get_tensor_by_name(ratio_neg_name)
dist = graph.get_tensor_by_name(dist_name)
# make sure ratio_pos, ratio_neg and dist are not empty
ratio_pos = tf.concat([ratio_pos, [0]], axis=0)
ratio_neg = tf.concat([ratio_neg, [0]], axis=0)
dist = tf.concat([dist, [0]], axis=0)
# Reorder edges to process mst in the proper order
# In the case of a constrained um_loss, edges won't be in ascending order
order = tf.argsort(dist, axis=-1, direction="ASCENDING")
ratio_pos = tf.gather(ratio_pos, order)
ratio_neg = tf.gather(ratio_neg, order)
dist = tf.gather(dist, order)
# Calculate a score:
false_pos = tf.math.cumsum(ratio_neg)
false_neg = 1 - tf.math.cumsum(ratio_pos)
scores = false_pos + false_neg
best_score = tf.math.reduce_min(scores)
best_score_index = tf.math.argmin(scores)
# best score is a range of values
best_alpha_min = tf.gather(dist, best_score_index)
best_alpha_max = tf.gather(dist, best_score_index + 1)
assert emst.name == emst_name, f"{emst.name} != {emst_name}"
assert edges_u.name == edges_u_name, f"{edges_u.name} != {edges_u_name}"
assert edges_v.name == edges_v_name, f"{edges_v.name} != {edges_v_name}"
# loss = um_loss + fg_loss
loss = um_loss
tf.summary.scalar("um_loss", um_loss)
tf.summary.scalar("best_alpha_min", best_alpha_min)
tf.summary.scalar("best_alpha_max", best_alpha_max)
tf.summary.scalar("best_score", best_score)
summaries = tf.summary.merge_all()
assert summaries.name == summaries_name, f"{summaries.name} != {summaries_name}"
opt = tf.train.AdamOptimizer(learning_rate=0.5e-5,
beta1=0.95,
beta2=0.999,
epsilon=1e-8)
optimizer = opt.minimize(loss)
return (loss, optimizer)
