Dimension mismatch fixed to use MONAI Variance Metric function (#995)

* Dimension mismatch fixed to use MONAI Variance Metric function

Signed-off-by: Vishwesh Nath <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Signed-off-by: Vishwesh Nath <[email protected]>
Co-authored-by: Vishwesh Nath <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: SACHIDANAND ALLE <[email protected]>

Author: 3 people

monailabel/tasks/scoring/epistemic_v2.py

@@ -91,7 +91,7 @@ def variance_volume(self, vol_input, ignore_nans=True):
             variance = np.sum(vari, axis=0)
         else:
             variance_metric = VarianceMetric(threshold=0.0005, spatial_map=True, scalar_reduction="sum")
-            variance = variance_metric(vol_input).cpu().detach().numpy()
+            variance = variance_metric(vol_input)
 
         if self.dimension == 3:
             variance = np.expand_dims(variance, axis=0)
@@ -181,10 +181,13 @@ def run_scoring(self, image_id, simulation_size, model_ts, datastore):
 
         accum = torch.stack(accum_unl_outputs)
         accum = torch.squeeze(accum)
-        if self.dimension == 3:
-            accum = accum[:, 1:, :, :, :] if len(accum.shape) > 4 else accum
-        else:
-            accum = accum[:, 1:, :, :] if len(accum.shape) > 3 else accum
+
+        # Accum Expected shape for 2D images is (N, C, H, W) for 3D (N, C, H, W, D)
+        # To handle cases where only a single class of segmentation is present, an extra dimension is added
+        if self.dimension == 2 and len(accum.shape) == 3:
+            accum = torch.unsqueeze(accum, dim=1)
+        elif self.dimension == 3 and len(accum.shape) == 4:
+            accum = torch.unsqueeze(accum, dim=1)
 
         entropy = self.variance_volume(accum) if self.use_variance else self.entropy_volume(accum)
         entropy = float(np.nanmean(entropy))