From 6e299a9a6a2a2439d724c600aca200d26b90e192 Mon Sep 17 00:00:00 2001 From: Basar Demir Date: Sat, 28 Mar 2026 04:04:36 -0400 Subject: [PATCH] Fix masked loss mean: divide by non-masked voxel count in LNCC, SquaredLNCC, and MINDSSC --- src/icon_registration/losses.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/src/icon_registration/losses.py b/src/icon_registration/losses.py index 0f3320d..6535dc9 100644 --- a/src/icon_registration/losses.py +++ b/src/icon_registration/losses.py @@ -592,16 +592,17 @@ def __call__(self, image_A, image_B, mask_B=None): I = image_A J = image_B assert I.shape == J.shape, "The shape of image I and J should be the same." - + if mask_B is not None: assert mask_B.shape == J.shape, "The shape of mask and J should be the same." - + lncc = self.calculate_lncc(I, J) - + if mask_B is not None: mask = mask_B.clone() mask[mask != 0] = 1 lncc = lncc * mask + return 1 - torch.sum(lncc) / torch.clamp(torch.sum(mask), min=1) return 1 - torch.mean(lncc) @@ -610,17 +611,18 @@ def __call__(self, image_A, image_B, mask_B=None): I = image_A J = image_B assert I.shape == J.shape, "The shape of image I and J should be the same." - + if mask_B is not None: assert mask_B.shape == J.shape, "The shape of mask and J should be the same." - + lncc = self.calculate_lncc(I, J) - + if mask_B is not None: mask = mask_B.clone() mask[mask != 0] = 1 lncc = lncc * mask - + return 1 - torch.sum(lncc**2) / torch.clamp(torch.sum(mask), min=1) + return 1 - torch.mean(lncc**2) class LNCCOnlyInterpolated(SimilarityBase): @@ -824,7 +826,8 @@ def __call__(self, image_A, image_B, mask_B=None): mask = mask.expand(-1, difference.shape[1], -1, -1, -1) difference = difference * mask - + return torch.sum(difference ** 2) / torch.clamp(torch.sum(mask), min=1) + return torch.mean(difference ** 2) def flips(phi, in_percentage=False):