Fix missing type annotations (#5655)

Fixes #5605.

### Description

Fix missing type annotations so `./runtests.sh --mypy` now finds 0
warnings.

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [x] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

Signed-off-by: Felix Schnabel <f.schnabel@tum.de>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: f-schnabel

monai/apps/auto3dseg/ensemble_builder.py

@@ -221,7 +221,7 @@ def __init__(self, n_fold: int = 5):
         super().__init__()
         self.n_fold = n_fold
 
-    def collect_algos(self):
+    def collect_algos(self) -> None:
         """
         Rank the algos by finding the best model in each cross-validation fold
         """

monai/apps/auto3dseg/hpo_gen.py

@@ -12,6 +12,7 @@
 import os
 from abc import abstractmethod
 from copy import deepcopy
+from typing import Optional
 from warnings import warn
 
 from monai.apps.auto3dseg.bundle_gen import BundleAlgo
@@ -96,7 +97,7 @@ class NNIGen(HPOGen):
         NNI command manually.
     """
 
-    def __init__(self, algo=None, params=None):
+    def __init__(self, algo: Optional[Algo] = None, params=None):
         self.algo: Algo
         self.hint = ""
         self.obj_filename = ""
@@ -115,7 +116,7 @@ def __init__(self, algo=None, params=None):
             else:
                 self.algo = algo
 
-            if isinstance(algo, BundleAlgo):
+            if isinstance(self.algo, BundleAlgo):
                 self.obj_filename = algo_to_pickle(self.algo, template_path=self.algo.template_path)
                 self.print_bundle_algo_instruction()
             else:
@@ -271,7 +272,7 @@ class OptunaGen(HPOGen):
 
     """
 
-    def __init__(self, algo=None, params=None):
+    def __init__(self, algo: Optional[Algo] = None, params=None) -> None:
         self.algo: Algo
         self.obj_filename = ""
 
@@ -289,7 +290,7 @@ def __init__(self, algo=None, params=None):
             else:
                 self.algo = algo
 
-            if isinstance(algo, BundleAlgo):
+            if isinstance(self.algo, BundleAlgo):
                 self.obj_filename = algo_to_pickle(self.algo, template_path=self.algo.template_path)
             else:
                 self.obj_filename = algo_to_pickle(self.algo)

monai/apps/deepgrow/transforms.py

@@ -14,7 +14,7 @@
 import numpy as np
 import torch
 
-from monai.config import IndexSelection, KeysCollection
+from monai.config import IndexSelection, KeysCollection, NdarrayOrTensor
 from monai.networks.layers import GaussianFilter
 from monai.transforms import Resize, SpatialCrop
 from monai.transforms.transform import MapTransform, Randomizable, Transform
@@ -50,7 +50,7 @@ def _apply(self, label):
                 sids.append(sid)
         return np.asarray(sids)
 
-    def __call__(self, data):
+    def __call__(self, data) -> Dict:
         d: Dict = dict(data)
         label = d[self.label].numpy() if isinstance(data[self.label], torch.Tensor) else data[self.label]
         if label.shape[0] != 1:
@@ -654,7 +654,7 @@ def bounding_box(self, points, img_shape):
             box_start[di], box_end[di] = min_d, max_d
         return box_start, box_end
 
-    def __call__(self, data):
+    def __call__(self, data) -> Dict:
         d: Dict = dict(data)
         first_key: Hashable = self.first_key(d)
         if first_key == ():
@@ -741,7 +741,7 @@ def __init__(
         self.meta_key_postfix = meta_key_postfix
         self.cropped_shape_key = cropped_shape_key
 
-    def __call__(self, data):
+    def __call__(self, data) -> Dict:
         d = dict(data)
         guidance = d[self.guidance]
         meta_dict: Dict = d[self.meta_keys or f"{self.ref_image}_{self.meta_key_postfix}"]
@@ -836,7 +836,7 @@ def __init__(
         self.original_shape_key = original_shape_key
         self.cropped_shape_key = cropped_shape_key
 
-    def __call__(self, data):
+    def __call__(self, data) -> Dict:
         d = dict(data)
         meta_dict: Dict = d[f"{self.ref_image}_{self.meta_key_postfix}"]
 
@@ -857,8 +857,9 @@ def __call__(self, data):
             box_end = meta_dict[self.end_coord_key]
 
             spatial_dims = min(len(box_start), len(image.shape[1:]))
-            slices = [slice(None)] + [slice(s, e) for s, e in zip(box_start[:spatial_dims], box_end[:spatial_dims])]
-            slices = tuple(slices)
+            slices = tuple(
+                [slice(None)] + [slice(s, e) for s, e in zip(box_start[:spatial_dims], box_end[:spatial_dims])]
+            )
             result[slices] = image
 
             # Undo Spacing
@@ -869,10 +870,11 @@ def __call__(self, data):
 
             if np.any(np.not_equal(current_size, spatial_size)):
                 resizer = Resize(spatial_size=spatial_size, mode=mode)
-                result = resizer(result, mode=mode, align_corners=align_corners)
+                result = resizer(result, mode=mode, align_corners=align_corners)  # type: ignore
 
             # Undo Slicing
             slice_idx = meta_dict.get("slice_idx")
+            final_result: NdarrayOrTensor
             if slice_idx is None or self.slice_only:
                 final_result = result if len(result.shape) <= 3 else result[0]
             else:

monai/auto3dseg/analyzer.py

@@ -12,7 +12,7 @@
 import time
 from abc import ABC, abstractmethod
 from copy import deepcopy
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Hashable, List, Mapping, Optional
 
 import numpy as np
 import torch
@@ -377,7 +377,7 @@ def __init__(self, image_key: str, label_key: str, stats_name: str = "label_stat
         self.label_key = label_key
         self.do_ccp = do_ccp
 
-        report_format: Dict[str, Any] = {
+        report_format: Dict[LabelStatsKeys, Any] = {
             LabelStatsKeys.LABEL_UID: None,
             LabelStatsKeys.IMAGE_INTST: None,
             LabelStatsKeys.LABEL: [{LabelStatsKeys.PIXEL_PCT: None, LabelStatsKeys.IMAGE_INTST: None}],
@@ -394,7 +394,7 @@ def __init__(self, image_key: str, label_key: str, stats_name: str = "label_stat
         id_seq = ID_SEP_KEY.join([LabelStatsKeys.LABEL, "0", LabelStatsKeys.IMAGE_INTST])
         self.update_ops_nested_label(id_seq, SampleOperations())
 
-    def __call__(self, data):
+    def __call__(self, data: Mapping[Hashable, MetaTensor]) -> Dict[Hashable, MetaTensor]:
         """
         Callable to execute the pre-defined functions.
 
@@ -442,7 +442,7 @@ def __call__(self, data):
             The stats operation uses numpy and torch to compute max, min, and other
             functions. If the input has nan/inf, the stats results will be nan/inf.
         """
-        d = dict(data)
+        d: Dict[Hashable, MetaTensor] = dict(data)
         start = time.time()
         if isinstance(d[self.image_key], (torch.Tensor, MetaTensor)) and d[self.image_key].device.type == "cuda":
             using_cuda = True
@@ -451,13 +451,13 @@ def __call__(self, data):
         restore_grad_state = torch.is_grad_enabled()
         torch.set_grad_enabled(False)
 
-        ndas = [d[self.image_key][i] for i in range(d[self.image_key].shape[0])]
-        ndas_label = d[self.label_key]  # (H,W,D)
+        ndas: List[MetaTensor] = [d[self.image_key][i] for i in range(d[self.image_key].shape[0])]  # type: ignore
+        ndas_label: MetaTensor = d[self.label_key]  # (H,W,D)
 
         if ndas_label.shape != ndas[0].shape:
             raise ValueError(f"Label shape {ndas_label.shape} is different from image shape {ndas[0].shape}")
 
-        nda_foregrounds = [get_foreground_label(nda, ndas_label) for nda in ndas]
+        nda_foregrounds: List[torch.Tensor] = [get_foreground_label(nda, ndas_label) for nda in ndas]
         nda_foregrounds = [nda if nda.numel() > 0 else torch.Tensor([0]) for nda in nda_foregrounds]
 
         unique_label = unique(ndas_label)

monai/data/dataset.py

@@ -1075,7 +1075,7 @@ def randomize(self, data: Sequence) -> None:
         except TypeError as e:
             warnings.warn(f"input data can't be shuffled in SmartCacheDataset with numpy.random.shuffle(): {e}.")
 
-    def _compute_data_idx(self):
+    def _compute_data_idx(self) -> None:
         """
         Update the replacement data position in the total data.
 
@@ -1209,7 +1209,7 @@ def _try_manage_replacement(self, check_round):
                 self._compute_replacements()
             return False, self._round
 
-    def manage_replacement(self):
+    def manage_replacement(self) -> None:
         """
         Background thread for replacement.
 

monai/data/image_reader.py

@@ -277,7 +277,7 @@ def read(self, data: Union[Sequence[PathLike], PathLike], **kwargs):
                 img_.append(itk.imread(name, **kwargs_))
         return img_ if len(filenames) > 1 else img_[0]
 
-    def get_data(self, img):
+    def get_data(self, img) -> Tuple[np.ndarray, Dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function returns two objects, first is numpy array of image data, second is dict of metadata.
@@ -564,7 +564,7 @@ def _combine_dicom_series(self, data):
 
         return stack_array, stack_metadata
 
-    def get_data(self, data):
+    def get_data(self, data) -> Tuple[np.ndarray, Dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function returns two objects, first is numpy array of image data, second is dict of metadata.
@@ -929,7 +929,7 @@ def read(self, data: Union[Sequence[PathLike], PathLike], **kwargs):
             img_.append(img)
         return img_ if len(filenames) > 1 else img_[0]
 
-    def get_data(self, img):
+    def get_data(self, img) -> Tuple[np.ndarray, Dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function returns two objects, first is numpy array of image data, second is dict of metadata.
@@ -1095,7 +1095,7 @@ def read(self, data: Union[Sequence[PathLike], PathLike], **kwargs):
 
         return img_ if len(img_) > 1 else img_[0]
 
-    def get_data(self, img):
+    def get_data(self, img) -> Tuple[np.ndarray, Dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function returns two objects, first is numpy array of image data, second is dict of metadata.
@@ -1113,7 +1113,7 @@ def get_data(self, img):
             img = (img,)
 
         for i in ensure_tuple(img):
-            header = {}
+            header: Dict[MetaKeys, Any] = {}
             if isinstance(i, np.ndarray):
                 # if `channel_dim` is None, can not detect the channel dim, use all the dims as spatial_shape
                 spatial_shape = np.asarray(i.shape)
@@ -1184,7 +1184,7 @@ def read(self, data: Union[Sequence[PathLike], PathLike, np.ndarray], **kwargs):
 
         return img_ if len(filenames) > 1 else img_[0]
 
-    def get_data(self, img):
+    def get_data(self, img) -> Tuple[np.ndarray, Dict]:
         """
         Extract data array and metadata from loaded image and return them.
         This function returns two objects, first is numpy array of image data, second is dict of metadata.

monai/data/meta_obj.py

@@ -79,7 +79,7 @@ class MetaObj:
 
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._meta: dict = MetaObj.get_default_meta()
         self._applied_operations: list = MetaObj.get_default_applied_operations()
         self._pending_operations: list = MetaObj.get_default_applied_operations()  # the same default as applied_ops

monai/data/meta_tensor.py

@@ -354,7 +354,7 @@ def get_array(self, output_type=np.ndarray, dtype=None, device=None, *_args, **_
         """
         return convert_data_type(self, output_type=output_type, dtype=dtype, device=device, wrap_sequence=True)[0]
 
-    def set_array(self, src, non_blocking=False, *_args, **_kwargs):
+    def set_array(self, src, non_blocking: bool = False, *_args, **_kwargs):
         """
         Copies the elements from src into self tensor and returns self.
         The src tensor must be broadcastable with the self tensor.
@@ -369,11 +369,11 @@ def set_array(self, src, non_blocking=False, *_args, **_kwargs):
             _args: currently unused parameters.
             _kwargs:  currently unused parameters.
         """
-        src: torch.Tensor = convert_to_tensor(src, track_meta=False, wrap_sequence=True)
+        converted: torch.Tensor = convert_to_tensor(src, track_meta=False, wrap_sequence=True)
         try:
-            return self.copy_(src, non_blocking=non_blocking)
+            return self.copy_(converted, non_blocking=non_blocking)
         except RuntimeError:  # skip the shape checking
-            self.data = src
+            self.data = converted
             return self
 
     @property

monai/metrics/metric.py

@@ -166,7 +166,7 @@ class Cumulative:
 
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """
         Initialize the internal buffers.
         `self._buffers` are local buffers, they are not usually used directly.

monai/networks/blocks/warp.py

@@ -151,7 +151,7 @@ def __init__(
         self.num_steps = num_steps
         self.warp_layer = Warp(mode=mode, padding_mode=padding_mode)
 
-    def forward(self, dvf):
+    def forward(self, dvf: torch.Tensor) -> torch.Tensor:
         """
         Args:
             dvf: dvf to be transformed, in shape (batch, ``spatial_dims``, H, W[,D])