Repository: Project-MONAI/MONAI

Issue:
Dim error in `HoVerNetLoss`
**Describe the bug**
Since in `SobelGradients `, we need input with shape (C, H, W, [D]), but I find in `HoVerNetLoss`, it used shape with (B, H, W).
https://github.com/Project-MONAI/MONAI/blob/27362e482761985daac8d5ef732eb57b50c741f1/monai/transforms/post/array.py#L888-L890
https://github.com/Project-MONAI/MONAI/blob/27362e482761985daac8d5ef732eb57b50c741f1/monai/apps/pathology/losses/hovernet_loss.py#L69-L72


Discussion:
Hi @JHancox and @drbeh, could you please help to confirm it? Thanks!
HoverNetLoss doesn't use the channel dimension, whereas I think SobelGradient can use both channel and batch. Sobel has probably changed since I implemented the HoverNetLoss though, so there may be a better way of calling Sobel now. @drbeh can probably advise.
@KumoLiu @JHancox That's right! The Sobel gradient has changed but this one not. I'll submit a quick fix.

Patch:
diff --git a/monai/apps/pathology/losses/hovernet_loss.py b/monai/apps/pathology/losses/hovernet_loss.py
--- a/monai/apps/pathology/losses/hovernet_loss.py
+++ b/monai/apps/pathology/losses/hovernet_loss.py
@@ -61,15 +61,21 @@ def __init__(
 
         self.dice = DiceLoss(softmax=True, smooth_dr=1e-03, smooth_nr=1e-03, reduction="sum", batch=True)
         self.ce = CrossEntropyLoss(reduction="mean")
-        self.sobel = SobelGradients(kernel_size=5)
+        self.sobel_v = SobelGradients(kernel_size=5, spatial_axes=0)
+        self.sobel_h = SobelGradients(kernel_size=5, spatial_axes=1)
 
     def _compute_sobel(self, image: torch.Tensor) -> torch.Tensor:
+        """Compute the Sobel gradients of the horizontal vertical map (HoVerMap).
+        More specifically, it will compute horizontal gradient of the input horizontal gradient map (channel=0) and
+        vertical gradient of the input vertical gradient map (channel=1).
 
-        batch_size = image.shape[0]
-        result_h = self.sobel(torch.squeeze(image[:, 0], dim=1))[batch_size:]
-        result_v = self.sobel(torch.squeeze(image[:, 1], dim=1))[:batch_size]
+        Args:
+            image: a tensor with the shape of BxCxHxW representing HoVerMap
 
-        return torch.cat([result_h[:, None, ...], result_v[:, None, ...]], dim=1)
+        """
+        result_h = self.sobel_h(image[:, 0])
+        result_v = self.sobel_v(image[:, 1])
+        return torch.stack([result_h, result_v], dim=1)
 
     def _mse_gradient_loss(self, prediction: torch.Tensor, target: torch.Tensor, focus: torch.Tensor) -> torch.Tensor:
         """Compute the MSE loss of the gradients of the horizontal and vertical centroid distance maps"""
diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py
--- a/monai/transforms/post/array.py
+++ b/monai/transforms/post/array.py
@@ -832,7 +832,7 @@ def __call__(self, data):
 
 
 class SobelGradients(Transform):
-    """Calculate Sobel gradients of a grayscale image with the shape of (CxH[xWxDx...]).
+    """Calculate Sobel gradients of a grayscale image with the shape of CxH[xWxDx...] or BxH[xWxDx...].
 
     Args:
         kernel_size: the size of the Sobel kernel. Defaults to 3.
@@ -922,7 +922,7 @@ def __call__(self, image: NdarrayOrTensor) -> torch.Tensor:
         grad_list = []
         for ax in spatial_axes:
             kernels = [kernel_smooth] * n_spatial_dims
-            kernels[ax - 1] = kernel_diff
+            kernels[ax] = kernel_diff
             grad = separable_filtering(image_tensor, kernels, mode=self.padding)
             if self.normalize_gradients:
                 grad_min = grad.min()


Test patch:
diff --git a/tests/test_hovernet_loss.py b/tests/test_hovernet_loss.py
--- a/tests/test_hovernet_loss.py
+++ b/tests/test_hovernet_loss.py
@@ -140,17 +140,17 @@ def test_shape_generator(num_classes=1, num_objects=3, batch_size=1, height=5, w
 
 TEST_CASE_3 = [  # batch size of 2, 3 classes with minor rotation of nuclear prediction
     {"prediction": inputs_test[3].inputs, "target": inputs_test[3].targets},
-    3.6169,
+    3.6348,
 ]
 
 TEST_CASE_4 = [  # batch size of 2, 3 classes with medium rotation of nuclear prediction
     {"prediction": inputs_test[4].inputs, "target": inputs_test[4].targets},
-    4.5079,
+    4.5312,
 ]
 
 TEST_CASE_5 = [  # batch size of 2, 3 classes with medium rotation of nuclear prediction
     {"prediction": inputs_test[5].inputs, "target": inputs_test[5].targets},
-    5.4663,
+    5.4929,
 ]
 
 CASES = [TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5]
diff --git a/tests/test_sobel_gradient.py b/tests/test_sobel_gradient.py
--- a/tests/test_sobel_gradient.py
+++ b/tests/test_sobel_gradient.py
@@ -22,17 +22,17 @@
 
 # Output with reflect padding
 OUTPUT_3x3 = torch.zeros(2, 16, 16, dtype=torch.float32)
-OUTPUT_3x3[1, 7, :] = 0.5
-OUTPUT_3x3[1, 9, :] = -0.5
+OUTPUT_3x3[0, 7, :] = 0.5
+OUTPUT_3x3[0, 9, :] = -0.5
 
 # Output with zero padding
 OUTPUT_3x3_ZERO_PAD = OUTPUT_3x3.clone()
-OUTPUT_3x3_ZERO_PAD[0, 7, 0] = OUTPUT_3x3_ZERO_PAD[0, 9, 0] = 0.125
-OUTPUT_3x3_ZERO_PAD[0, 8, 0] = 0.25
-OUTPUT_3x3_ZERO_PAD[0, 7, -1] = OUTPUT_3x3_ZERO_PAD[0, 9, -1] = -0.125
-OUTPUT_3x3_ZERO_PAD[0, 8, -1] = -0.25
-OUTPUT_3x3_ZERO_PAD[1, 7, 0] = OUTPUT_3x3_ZERO_PAD[1, 7, -1] = 3.0 / 8.0
-OUTPUT_3x3_ZERO_PAD[1, 9, 0] = OUTPUT_3x3_ZERO_PAD[1, 9, -1] = -3.0 / 8.0
+OUTPUT_3x3_ZERO_PAD[1, 7, 0] = OUTPUT_3x3_ZERO_PAD[1, 9, 0] = 0.125
+OUTPUT_3x3_ZERO_PAD[1, 8, 0] = 0.25
+OUTPUT_3x3_ZERO_PAD[1, 7, -1] = OUTPUT_3x3_ZERO_PAD[1, 9, -1] = -0.125
+OUTPUT_3x3_ZERO_PAD[1, 8, -1] = -0.25
+OUTPUT_3x3_ZERO_PAD[0, 7, 0] = OUTPUT_3x3_ZERO_PAD[0, 7, -1] = 3.0 / 8.0
+OUTPUT_3x3_ZERO_PAD[0, 9, 0] = OUTPUT_3x3_ZERO_PAD[0, 9, -1] = -3.0 / 8.0
 
 TEST_CASE_0 = [IMAGE, {"kernel_size": 3, "dtype": torch.float32}, OUTPUT_3x3]
 TEST_CASE_1 = [IMAGE, {"kernel_size": 3, "dtype": torch.float64}, OUTPUT_3x3]
@@ -68,7 +68,7 @@
         "spatial_axes": (0, 1),
         "dtype": torch.float64,
     },
-    torch.cat([OUTPUT_3x3[0:1], OUTPUT_3x3[1:2] + 0.5]),
+    torch.cat([OUTPUT_3x3[0:1] + 0.5, OUTPUT_3x3[1:2]]),
 ]
 TEST_CASE_10 = [  # Normalized gradients but non-normalized kernels
     IMAGE,
@@ -79,7 +79,7 @@
         "spatial_axes": (0, 1),
         "dtype": torch.float64,
     },
-    torch.cat([OUTPUT_3x3[0:1], OUTPUT_3x3[1:2] + 0.5]),
+    torch.cat([OUTPUT_3x3[0:1] + 0.5, OUTPUT_3x3[1:2]]),
 ]
 
 TEST_CASE_KERNEL_0 = [
diff --git a/tests/test_sobel_gradientd.py b/tests/test_sobel_gradientd.py
--- a/tests/test_sobel_gradientd.py
+++ b/tests/test_sobel_gradientd.py
@@ -22,17 +22,17 @@
 
 # Output with reflect padding
 OUTPUT_3x3 = torch.zeros(2, 16, 16, dtype=torch.float32)
-OUTPUT_3x3[1, 7, :] = 0.5
-OUTPUT_3x3[1, 9, :] = -0.5
+OUTPUT_3x3[0, 7, :] = 0.5
+OUTPUT_3x3[0, 9, :] = -0.5
 
 # Output with zero padding
 OUTPUT_3x3_ZERO_PAD = OUTPUT_3x3.clone()
-OUTPUT_3x3_ZERO_PAD[0, 7, 0] = OUTPUT_3x3_ZERO_PAD[0, 9, 0] = 0.125
-OUTPUT_3x3_ZERO_PAD[0, 8, 0] = 0.25
-OUTPUT_3x3_ZERO_PAD[0, 7, -1] = OUTPUT_3x3_ZERO_PAD[0, 9, -1] = -0.125
-OUTPUT_3x3_ZERO_PAD[0, 8, -1] = -0.25
-OUTPUT_3x3_ZERO_PAD[1, 7, 0] = OUTPUT_3x3_ZERO_PAD[1, 7, -1] = 3.0 / 8.0
-OUTPUT_3x3_ZERO_PAD[1, 9, 0] = OUTPUT_3x3_ZERO_PAD[1, 9, -1] = -3.0 / 8.0
+OUTPUT_3x3_ZERO_PAD[1, 7, 0] = OUTPUT_3x3_ZERO_PAD[1, 9, 0] = 0.125
+OUTPUT_3x3_ZERO_PAD[1, 8, 0] = 0.25
+OUTPUT_3x3_ZERO_PAD[1, 7, -1] = OUTPUT_3x3_ZERO_PAD[1, 9, -1] = -0.125
+OUTPUT_3x3_ZERO_PAD[1, 8, -1] = -0.25
+OUTPUT_3x3_ZERO_PAD[0, 7, 0] = OUTPUT_3x3_ZERO_PAD[0, 7, -1] = 3.0 / 8.0
+OUTPUT_3x3_ZERO_PAD[0, 9, 0] = OUTPUT_3x3_ZERO_PAD[0, 9, -1] = -3.0 / 8.0
 
 TEST_CASE_0 = [{"image": IMAGE}, {"keys": "image", "kernel_size": 3, "dtype": torch.float32}, {"image": OUTPUT_3x3}]
 TEST_CASE_1 = [{"image": IMAGE}, {"keys": "image", "kernel_size": 3, "dtype": torch.float64}, {"image": OUTPUT_3x3}]
@@ -86,7 +86,7 @@
         "normalize_gradients": True,
         "dtype": torch.float32,
     },
-    {"image": torch.cat([OUTPUT_3x3[0:1], OUTPUT_3x3[1:2] + 0.5])},
+    {"image": torch.cat([OUTPUT_3x3[0:1] + 0.5, OUTPUT_3x3[1:2]])},
 ]
 TEST_CASE_11 = [  # Normalized gradients but non-normalized kernels
     {"image": IMAGE},
@@ -98,7 +98,7 @@
         "normalize_gradients": True,
         "dtype": torch.float32,
     },
-    {"image": torch.cat([OUTPUT_3x3[0:1], OUTPUT_3x3[1:2] + 0.5])},
+    {"image": torch.cat([OUTPUT_3x3[0:1] + 0.5, OUTPUT_3x3[1:2]])},
 ]
 
 TEST_CASE_KERNEL_0 = [