Add type hints to monai/visualize/

- use 'if TYPE_CHECKING:' for optional imports
- add type hints
- remove docstring type hints

optional_import does not enable correct type hinting
so 'if TYPE_CHECKING:' is used. The type annotation must
then be a forward reference i.e. it must be in quotes.

Author: Benjamin Gorman

monai/visualize/img2tensorboard.py

@@ -9,7 +9,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Sequence, Union
+from typing import Optional, Sequence, Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.utils.tensorboard
 
 import numpy as np
 import torch
@@ -20,7 +23,6 @@
 PIL, _ = optional_import("PIL")
 GifImage, _ = optional_import("PIL.GifImagePlugin", name="Image")
 summary_pb2, _ = optional_import("tensorboard.compat.proto.summary_pb2")
-SummaryWriter, _ = optional_import("torch.utils.tensorboard", name="SummaryWriter")
 
 
 def _image3_animated_gif(tag: str, image: Union[np.ndarray, torch.Tensor], scale_factor: float = 1.0):
@@ -97,13 +99,13 @@ def make_animated_gif_summary(
 
 
 def add_animated_gif(
-    writer,
+    writer: "torch.utils.tensorboard.SummaryWriter",
     tag: str,
     image_tensor: Union[np.ndarray, torch.Tensor],
     max_out: int,
     scale_factor: float,
     global_step: Optional[int] = None,
-):
+) -> None:
     """Creates an animated gif out of an image tensor in 'CHWD' format and writes it with SummaryWriter.
 
     Args:
@@ -124,13 +126,13 @@ def add_animated_gif(
 
 
 def add_animated_gif_no_channels(
-    writer,
+    writer: "torch.utils.tensorboard.SummaryWriter",
     tag: str,
     image_tensor: Union[np.ndarray, torch.Tensor],
     max_out: int,
     scale_factor: float,
     global_step: Optional[int] = None,
-):
+) -> None:
     """Creates an animated gif out of an image tensor in 'HWD' format that does not have
     a channel dimension and writes it with SummaryWriter. This is similar to the "add_animated_gif"
     after inserting a channel dimension of 1.
@@ -155,28 +157,27 @@ def add_animated_gif_no_channels(
 def plot_2d_or_3d_image(
     data: Union[torch.Tensor, np.ndarray],
     step: int,
-    writer,
+    writer: "torch.utils.tensorboard.SummaryWriter",
     index: int = 0,
     max_channels: int = 1,
     max_frames: int = 64,
     tag: str = "output",
-):
+) -> None:
     """Plot 2D or 3D image on the TensorBoard, 3D image will be converted to GIF image.
 
     Note:
         Plot 3D or 2D image(with more than 3 channels) as separate images.
 
     Args:
-        data (Tensor or np.array): target data to be plotted as image on the TensorBoard.
+        data: target data to be plotted as image on the TensorBoard.
             The data is expected to have 'NCHW[D]' dimensions, and only plot the first in the batch.
         step: current step to plot in a chart.
-        writer (SummaryWriter): specify TensorBoard SummaryWriter to plot the image.
+        writer: specify TensorBoard SummaryWriter to plot the image.
         index: plot which element in the input data batch, default is the first element.
         max_channels: number of channels to plot.
         max_frames: number of frames for 2D-t plot.
         tag: tag of the plotted image on TensorBoard.
     """
-    assert isinstance(writer, SummaryWriter) is True, "must provide a TensorBoard SummaryWriter."
     d = data[index]
     if torch.is_tensor(d):
         d = d.detach().cpu().numpy()