| import unittest | |
| import torch | |
| from utils.region_smc import blend_region_wise_probabilities | |
| class RegionWiseSMCTest(unittest.TestCase): | |
| def test_returns_spatial_maps_without_gradients(self): | |
| sam_prob = torch.tensor( | |
| [[[[0.8, 0.3]], [[0.2, 0.7]]]], requires_grad=True | |
| ) | |
| unet_teacher_prob = torch.tensor( | |
| [[[[0.7, 0.4]], [[0.3, 0.6]]]], requires_grad=True | |
| ) | |
| unet_student_prob = torch.tensor( | |
| [[[[0.6, 0.5]], [[0.4, 0.5]]]], requires_grad=True | |
| ) | |
| blended, self_conf, mutual_conf, alpha_map = ( | |
| blend_region_wise_probabilities( | |
| sam_prob, | |
| unet_teacher_prob, | |
| unet_student_prob, | |
| ) | |
| ) | |
| self.assertEqual(blended.shape, (1, 2, 1, 2)) | |
| self.assertEqual(self_conf.shape, (1, 1, 1, 2)) | |
| self.assertEqual(mutual_conf.shape, (1, 1, 1, 2)) | |
| self.assertEqual(alpha_map.shape, (1, 1, 1, 2)) | |
| for tensor in (blended, self_conf, mutual_conf, alpha_map): | |
| self.assertFalse(tensor.requires_grad) | |
| for confidence_map in (self_conf, mutual_conf, alpha_map): | |
| self.assertTrue(torch.all(confidence_map >= 0.0)) | |
| self.assertTrue(torch.all(confidence_map <= 1.0)) | |
| torch.testing.assert_close( | |
| blended.sum(dim=1), torch.ones_like(blended[:, 0]) | |
| ) | |
| def test_agreement_moves_alpha_to_half(self): | |
| shared_prob = torch.tensor([[[[0.9]], [[0.1]]]]) | |
| _, _, mutual_conf, alpha_map = blend_region_wise_probabilities( | |
| shared_prob, | |
| shared_prob, | |
| shared_prob, | |
| ) | |
| torch.testing.assert_close(mutual_conf, torch.ones_like(mutual_conf)) | |
| torch.testing.assert_close(alpha_map, torch.full_like(alpha_map, 0.5)) | |
| def test_stable_confident_unet_is_preferred_during_disagreement(self): | |
| sam_prob = torch.tensor([[[[0.55]], [[0.45]]]]) | |
| unet_prob = torch.tensor([[[[0.01]], [[0.99]]]]) | |
| blended, _, _, alpha_map = blend_region_wise_probabilities( | |
| sam_prob, | |
| unet_prob, | |
| unet_prob, | |
| ) | |
| self.assertGreater(alpha_map.item(), 0.5) | |
| self.assertEqual(blended.argmax(dim=1).item(), 1) | |
| def test_unstable_unet_is_deprioritized(self): | |
| sam_prob = torch.tensor([[[[0.99]], [[0.01]]]]) | |
| unet_teacher_prob = torch.tensor([[[[0.45]], [[0.55]]]]) | |
| unet_student_prob = sam_prob.clone() | |
| blended, self_conf, _, alpha_map = blend_region_wise_probabilities( | |
| sam_prob, | |
| unet_teacher_prob, | |
| unet_student_prob, | |
| ) | |
| self.assertLess(self_conf.item(), 0.5) | |
| self.assertLess(alpha_map.item(), 0.5) | |
| self.assertEqual(blended.argmax(dim=1).item(), 0) | |
| def test_supports_multiclass_probability_maps(self): | |
| sam_prob = torch.softmax(torch.randn(2, 3, 4, 5), dim=1) | |
| unet_teacher_prob = torch.softmax(torch.randn(2, 3, 4, 5), dim=1) | |
| unet_student_prob = torch.softmax(torch.randn(2, 3, 4, 5), dim=1) | |
| blended, self_conf, mutual_conf, alpha_map = ( | |
| blend_region_wise_probabilities( | |
| sam_prob, | |
| unet_teacher_prob, | |
| unet_student_prob, | |
| ) | |
| ) | |
| self.assertEqual(blended.shape, (2, 3, 4, 5)) | |
| self.assertEqual(self_conf.shape, (2, 1, 4, 5)) | |
| self.assertEqual(mutual_conf.shape, (2, 1, 4, 5)) | |
| self.assertEqual(alpha_map.shape, (2, 1, 4, 5)) | |
| if __name__ == '__main__': | |
| unittest.main() | |