🐛 Fix HoVer-Net Class Output Information in docs (TissueImageAnal…

…ytics#751)

- [X] Fix bug with HoVer-Net, where currently merging classes 0 and 1. This is resulting in the background class being omitted and merged with the first.
- [x] Add the output nuclear classes for the nuclear instance segmentation models to the documentation. We have also added type dictionaries to the pretrained_yaml file.
- [x] Update notebooks to show this correction.

---------

Co-authored-by: Shan E Ahmed Raza <[email protected]>

docs/pretrained.rst

@@ -185,6 +185,15 @@ The input output configuration is as follows:
 
     - hovernet_fast-pannuke
 
+.. collapse:: Output Nuclear Classes
+
+    - 0: Background
+    - 1: Neoplastic
+    - 2: Inflammatory
+    - 3: Connective
+    - 4: Dead
+    - 5: Non-Neoplastic Epithelial
+
 MoNuSAC Dataset
 ---------------
 
@@ -218,6 +227,14 @@ The input output configuration is as follows:
 
     - hovernet_fast-monusac
 
+.. collapse:: Output Nuclear Classes
+
+    - 0: Background
+    - 1: Epithelial
+    - 2: Lymphocyte
+    - 3: Macrophage
+    - 4: Neutrophil
+
 CoNSeP Dataset
 --------------
 
@@ -251,6 +268,14 @@ The input output configuration is as follows:
 
     - hovernet_original-consep
 
+.. collapse:: Output Nuclear Classes
+
+    - 0: Background
+    - 1: Epithelial
+    - 2: Inflammatory
+    - 3: Spindle-Shaped
+    - 4: Miscellaneous
+
 
 .. collapse:: Input Output Configuration Details
 
@@ -275,11 +300,13 @@ The input output configuration is as follows:
 
     - micronet_hovernet-consep
 
+
 Kumar Dataset
 -------------
 
 We provide the following models trained using the `Kumar dataset <https://monuseg.grand-challenge.org/>`_.
 All model weights trained on Kumar are held under the `Creative Commons Attribution-NonCommercial-ShareAlike Version 4 (CC BY-NC-SA 4.0) License <https://creativecommons.org/licenses/by-nc-sa/4.0/>`_.
+The Kumar dataset does not contain nuclear class information, and so TIAToolbox pretrained models based on Kumar for nuclear segmentation, will only perform segmentation and not classification.
 The input output configuration is as follows:
 
 .. collapse:: Input Output Configuration Details
@@ -428,3 +455,17 @@ The model uses the following input output configuration:
 .. collapse:: Model names
 
     - hovernetplus-oed
+
+.. collapse:: Output Nuclear Classes
+
+    - 0: Background
+    - 1: Other
+    - 2: Epithelial
+
+.. collapse:: Output Region Classes
+
+    - 0: Background
+    - 1: Other Tissue
+    - 2: Basal Epithelium
+    - 3: (Core) Epithelium
+    - 4: Keratin

tiatoolbox/data/pretrained_model.yaml

@@ -643,6 +643,14 @@ hovernet_fast-pannuke:
         kwargs:
             num_types: 6
             mode: "fast"
+            nuc_type_dict: {
+                0: "Background",
+                1: "Neoplastic",
+                2: "Inflammatory",
+                3: "Connective",
+                4: "Dead",
+                5: "Non-Neoplastic Epithelial",
+            }
     ioconfig:
         class: semantic_segmentor.IOSegmentorConfig
         kwargs:
@@ -666,6 +674,13 @@ hovernet_fast-monusac:
         kwargs:
             num_types: 5
             mode: "fast"
+            nuc_type_dict: {
+                0: "Background",
+                1: "Epithelial",
+                2: "Lymphocyte",
+                3: "Macrophage",
+                4: "Neutrophil",
+            }
     ioconfig:
         class: semantic_segmentor.IOSegmentorConfig
         kwargs:
@@ -689,6 +704,13 @@ hovernet_original-consep:
         kwargs:
             num_types: 5
             mode: "original"
+            nuc_type_dict: {
+                0: "Background",
+                1: "Epithelial",
+                2: "Inflammatory",
+                3: "Spindle-Shaped",
+                4: "Miscellaneous",
+            }
     ioconfig:
         class: semantic_segmentor.IOSegmentorConfig
         kwargs:
@@ -734,6 +756,18 @@ hovernetplus-oed:
         kwargs:
             num_types: 3
             num_layers: 5
+            nuc_type_dict: {
+                0: "Background",
+                1: "Other",
+                2: "Epithelial",
+            }
+            layer_type_dict: {
+                0: "Background",
+                1: "Other Tissue",
+                2: "Basal Epithelium",
+                3: "(Core) Epithelium",
+                4: "Keratin",
+            }
     ioconfig:
         class: semantic_segmentor.IOSegmentorConfig
         kwargs:

tiatoolbox/models/architecture/hovernet.py

@@ -329,11 +329,13 @@ def __init__(
         num_input_channels: int = 3,
         num_types: int | None = None,
         mode: str = "original",
+        nuc_type_dict: dict | None = None,
     ) -> None:
         """Initialize :class:`HoVerNet`."""
         super().__init__()
         self.mode = mode
         self.num_types = num_types
+        self.nuc_type_dict = nuc_type_dict
 
         if mode not in ["original", "fast"]:
             msg = (
@@ -771,6 +773,7 @@ def postproc(raw_maps: list[np.ndarray]) -> tuple[np.ndarray, dict]:
         """
         if len(raw_maps) == 3:  # noqa: PLR2004
             np_map, hv_map, tp_map = raw_maps
+            tp_map = np.around(tp_map).astype("uint8")
         else:
             tp_map = None
             np_map, hv_map = raw_maps

tiatoolbox/models/architecture/hovernetplus.py

@@ -85,12 +85,16 @@ def __init__(
         num_input_channels: int = 3,
         num_types: int | None = None,
         num_layers: int | None = None,
+        nuc_type_dict: dict | None = None,
+        layer_type_dict: dict | None = None,
     ) -> None:
         """Initialize :class:`HoVerNetPlus`."""
         super().__init__(mode="fast")
         self.num_input_channels = num_input_channels
         self.num_types = num_types
         self.num_layers = num_layers
+        self.nuc_type_dict = nuc_type_dict
+        self.layer_type_dict = layer_type_dict
         ksize = 3
 
         self.decoder = nn.ModuleDict(