Repository: Project-MONAI/MONAI

Issue:
HoVerNet Model branch naming
HoVerNet model has hard coded branch names that makes it hard to use these branch names across other components, for instance when using in a loss function. It would be great to have enum keys to define these branch names to be usable elsewhere.


Discussion:

Resolving patch:
diff --git a/monai/networks/nets/hovernet.py b/monai/networks/nets/hovernet.py
--- a/monai/networks/nets/hovernet.py
+++ b/monai/networks/nets/hovernet.py
@@ -25,12 +25,11 @@
 #    year={2019},
 #    publisher={Elsevier}
 # }
-
 # =========================================================================
 
 from collections import OrderedDict
 from enum import Enum
-from typing import Callable, Dict, List, Sequence, Type, Union
+from typing import Callable, Dict, List, Optional, Sequence, Type, Union
 
 import torch
 import torch.nn as nn
@@ -39,8 +38,9 @@
 from monai.networks.layers.factories import Conv, Dropout
 from monai.networks.layers.utils import get_act_layer, get_norm_layer
 from monai.utils import InterpolateMode, UpsampleMode, export
+from monai.utils.enums import StrEnum
 
-__all__ = ["HoverNet", "Hovernet", "HoVernet", "HoVerNet"]
+__all__ = ["HoVerNet", "Hovernet", "HoVernet", "HoVerNet"]
 
 
 class _DenseLayerDecoder(nn.Module):
@@ -371,8 +371,8 @@ def forward(self, xin: torch.Tensor, short_cuts: List[torch.Tensor]) -> torch.Te
 
 
 @export("monai.networks.nets")
-class HoverNet(nn.Module):
-    """HoVerNet
+class HoVerNet(nn.Module):
+    """HoVerNet model
 
     References:
       Graham, Simon et al. Hover-net: Simultaneous segmentation
@@ -391,6 +391,19 @@ class Mode(Enum):
         FAST: int = 0
         ORIGINAL: int = 1
 
+    class Branch(StrEnum):
+        """
+        Three branches of HoVerNet model, which results in three outputs:
+        `HOVER` is horizontal and vertical regressed gradient map of each nucleus,
+        `NUCLEUS` is the segmentation of all nuclei, and
+        `TYPE` is the type of each nucleus.
+
+        """
+
+        HV = "horizontal_vertical"
+        NP = "nucleus_prediction"
+        NC = "type_prediction"
+
     def _mode_to_int(self, mode) -> int:
 
         if mode == self.Mode.FAST:
@@ -484,10 +497,9 @@ def __init__(
         # decode branches
         self.nucleus_prediction = _DecoderBranch(kernel_size=_ksize)
         self.horizontal_vertical = _DecoderBranch(kernel_size=_ksize)
-        self.type_prediction: _DecoderBranch = None  # type: ignore
-
-        if out_classes > 0:
-            self.type_prediction = _DecoderBranch(out_channels=out_classes, kernel_size=_ksize)
+        self.type_prediction: Optional[_DecoderBranch] = (
+            _DecoderBranch(out_channels=out_classes, kernel_size=_ksize) if out_classes > 0 else None
+        )
 
         for m in self.modules():
             if isinstance(m, conv_type):
@@ -518,15 +530,14 @@ def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
         x = self.bottleneck(x)
         x = self.upsample(x)
 
-        x_np = self.nucleus_prediction(x, short_cuts)
-        x_hv = self.horizontal_vertical(x, short_cuts)
-        tp = self.type_prediction
-
-        if tp is not None:
-            x_tp = self.type_prediction(x, short_cuts)
-            return {"nucleus_prediction": x_np, "horizonal_vertical": x_hv, "type_prediction": x_tp}
+        output = {
+            HoVerNet.Branch.NP.value: self.nucleus_prediction(x, short_cuts),
+            HoVerNet.Branch.HV.value: self.horizontal_vertical(x, short_cuts),
+        }
+        if self.type_prediction is not None:
+            output[HoVerNet.Branch.NC.value] = self.type_prediction(x, short_cuts)
 
-        return {"nucleus_prediction": x_np, "horizonal_vertical": x_hv}
+        return output
 
 
-Hovernet = HoVernet = HoVerNet = HoverNet
+Hovernet = HoVernet = HoverNet = HoVerNet


Test patch:
diff --git a/tests/test_hovernet.py b/tests/test_hovernet.py
--- a/tests/test_hovernet.py
+++ b/tests/test_hovernet.py
@@ -15,76 +15,72 @@
 from parameterized import parameterized
 
 from monai.networks import eval_mode
-from monai.networks.nets import HoverNet
+from monai.networks.nets import HoVerNet
 from tests.utils import test_script_save
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 TEST_CASE_0 = [  # fast mode, batch 16
-    {"out_classes": 5, "mode": HoverNet.Mode.FAST},
+    {"out_classes": 5, "mode": HoVerNet.Mode.FAST},
     (1, 3, 256, 256),
-    {
-        "nucleus_prediction": (1, 2, 164, 164),
-        "type_prediction": (1, 5, 164, 164),
-        "horizonal_vertical": (1, 2, 164, 164),
-    },
+    {HoVerNet.Branch.NP: (1, 2, 164, 164), HoVerNet.Branch.NC: (1, 5, 164, 164), HoVerNet.Branch.HV: (1, 2, 164, 164)},
 ]
 
 TEST_CASE_1 = [  # single channel 2D, batch 16
-    {"mode": HoverNet.Mode.FAST},
+    {"mode": HoVerNet.Mode.FAST},
     (1, 3, 256, 256),
-    {"nucleus_prediction": (1, 2, 164, 164), "horizonal_vertical": (1, 2, 164, 164)},
+    {HoVerNet.Branch.NP: (1, 2, 164, 164), HoVerNet.Branch.HV: (1, 2, 164, 164)},
 ]
 
 TEST_CASE_2 = [  # single channel 3D, batch 16
-    {"mode": HoverNet.Mode.ORIGINAL},
+    {"mode": HoVerNet.Mode.ORIGINAL},
     (1, 3, 270, 270),
-    {"nucleus_prediction": (1, 2, 80, 80), "horizonal_vertical": (1, 2, 80, 80)},
+    {HoVerNet.Branch.NP: (1, 2, 80, 80), HoVerNet.Branch.HV: (1, 2, 80, 80)},
 ]
 
 TEST_CASE_3 = [  # 4-channel 3D, batch 16
-    {"out_classes": 6, "mode": HoverNet.Mode.ORIGINAL},
+    {"out_classes": 6, "mode": HoVerNet.Mode.ORIGINAL},
     (1, 3, 270, 270),
-    {"nucleus_prediction": (1, 2, 80, 80), "type_prediction": (1, 6, 80, 80), "horizonal_vertical": (1, 2, 80, 80)},
+    {HoVerNet.Branch.NP: (1, 2, 80, 80), HoVerNet.Branch.NC: (1, 6, 80, 80), HoVerNet.Branch.HV: (1, 2, 80, 80)},
 ]
 
 TEST_CASE_4 = [  # 4-channel 3D, batch 16, batch normalization
-    {"mode": HoverNet.Mode.FAST, "dropout_prob": 0.5},
+    {"mode": HoVerNet.Mode.FAST, "dropout_prob": 0.5},
     (1, 3, 256, 256),
-    {"nucleus_prediction": (1, 2, 164, 164), "horizonal_vertical": (1, 2, 164, 164)},
+    {HoVerNet.Branch.NP: (1, 2, 164, 164), HoVerNet.Branch.HV: (1, 2, 164, 164)},
 ]
 
 CASES = [TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4]
 
 ILL_CASES = [
     [{"out_classes": 6, "mode": 3}],
-    [{"out_classes": 1000, "mode": HoverNet.Mode.ORIGINAL}],
-    [{"out_classes": 1, "mode": HoverNet.Mode.ORIGINAL}],
-    [{"out_classes": 6, "mode": HoverNet.Mode.ORIGINAL, "dropout_prob": 100}],
+    [{"out_classes": 1000, "mode": HoVerNet.Mode.ORIGINAL}],
+    [{"out_classes": 1, "mode": HoVerNet.Mode.ORIGINAL}],
+    [{"out_classes": 6, "mode": HoVerNet.Mode.ORIGINAL, "dropout_prob": 100}],
 ]
 
 
 class TestHoverNet(unittest.TestCase):
     @parameterized.expand(CASES)
     def test_shape(self, input_param, input_shape, expected_shapes):
-        net = HoverNet(**input_param).to(device)
+        net = HoVerNet(**input_param).to(device)
         with eval_mode(net):
             result = net.forward(torch.randn(input_shape).to(device))
             for item in result:
                 self.assertEqual(result[item].shape, expected_shapes[item])
 
     def test_script(self):
-        net = HoverNet(mode=HoverNet.Mode.FAST)
+        net = HoVerNet(mode=HoVerNet.Mode.FAST)
         test_data = torch.randn(1, 3, 256, 256)
         test_script_save(net, test_data)
 
     def test_script_without_running_stats(self):
-        net = HoverNet(mode=HoverNet.Mode.FAST)
+        net = HoVerNet(mode=HoVerNet.Mode.FAST)
         test_data = torch.randn(1, 3, 256, 256)
         test_script_save(net, test_data)
 
     def test_ill_input_shape(self):
-        net = HoverNet(mode=HoverNet.Mode.FAST)
+        net = HoVerNet(mode=HoVerNet.Mode.FAST)
         with eval_mode(net):
             with self.assertRaises(ValueError):
                 net.forward(torch.randn(1, 3, 270, 260))
@@ -92,7 +88,7 @@ def test_ill_input_shape(self):
     @parameterized.expand(ILL_CASES)
     def test_ill_input_hyper_params(self, input_param):
         with self.assertRaises(ValueError):
-            _ = HoverNet(**input_param)
+            _ = HoVerNet(**input_param)
 
 
 if __name__ == "__main__":