add more type hints to function return values

Author: henrykrumb

ncalab/data.py

@@ -1,5 +1,3 @@
-# TODO implement pattern pool
-# TODO implement damage
 import numpy as np
 import torch  # type: ignore[import-untyped]
 from torch.utils.data import Dataset  # type: ignore[import-untyped]
@@ -11,7 +9,9 @@ def __init__(
         image: np.ndarray,
         num_channels: int,
         batch_size: int = 8,
+        # TODO implement
         use_pattern_pool: bool = False,
+        # TODO implement
         damage: bool = False,
     ):
         """Dedicated dataset for "growing" tasks, like growing emoji.

ncalab/losses.py

@@ -11,17 +11,10 @@ def __init__(self):
         super(DiceScore, self).__init__()
 
     def forward(
-        self, x: torch.Tensor, y: torch.Tensor, smooth: float = 1
+        self, x: torch.Tensor, y: torch.Tensor, smooth: float = 1.0
     ) -> torch.Tensor:
         """
-
-        Args:
-            input (_type_): _description_
-            target (_type_): _description_
-            smooth (int, optional): _description_. Defaults to 1.
-
-        Returns:
-            _type_: _description_
+        :returns: Dice score as Tensor
         """
         x = torch.sigmoid(x)
         x = torch.flatten(x)

ncalab/models/classificationNCA.py

@@ -165,7 +165,7 @@ def loss(self, image: torch.Tensor, label: torch.Tensor) -> Dict[str, torch.Tens
             "classification": loss_classification,
         }
 
-    def metrics(self, pred: torch.Tensor, label: torch.Tensor):
+    def metrics(self, pred: torch.Tensor, label: torch.Tensor) -> Dict[str, float]:
         """
         Return dict of standard evaluation metrics.
 

ncalab/training/earlystopping.py

@@ -16,7 +16,7 @@ def __init__(self, patience: int, min_delta: float = 1e-6):
         self.best_accuracy = 0.0
         self.counter = 0
 
-    def done(self):
+    def done(self) -> bool:
         """
         Checks whether the training can be stopped.
 

ncalab/training/kfold.py

@@ -72,7 +72,7 @@ def __init__(self):
         self.dataloader_test = None
 
     @staticmethod
-    def read(path: PosixPath):
+    def read(path: PosixPath) -> "SplitDefinition":
         """
         Reads json files with split definitions, similar to those created by nnUNet.
 
@@ -105,7 +105,7 @@ def read(path: PosixPath):
         # TODO validate structure
         return sd
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.folds)
 
     def __getitem__(self, idx) -> TrainValRecord:

ncalab/utils.py

@@ -7,6 +7,12 @@
 import torch.nn.functional as F  # type: ignore[import-untyped]
 
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .models import BasicNCAModel
+
+
 def get_compute_device(device: str = "cuda:0") -> torch.device:
     """
     Obtain a pytorch compute device handle based on input string.
@@ -23,7 +29,7 @@ def get_compute_device(device: str = "cuda:0") -> torch.device:
     return d
 
 
-def pad_input(x, nca, noise: bool = True):
+def pad_input(x: torch.Tensor, nca: BasicNCAModel, noise: bool = True) -> torch.Tensor:
     """
     Pads input tensor along channel dimension to match the expected number of
     channels required by the NCA model. Pads with either Gaussian noise or zeros,