From 5b73cf529f127929be624e332c437513b0f6c07f Mon Sep 17 00:00:00 2001 From: Philip Meier Date: Mon, 6 Apr 2020 15:32:16 +0200 Subject: [PATCH] Type annotations for torchvision/utils.py (#2034) * type annotations for torchvision/utils.py * add missing annotation for make_grid * fix annotation for save_image * mirror PIL annotation for fp --- torchvision/utils.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/torchvision/utils.py b/torchvision/utils.py index 1a773b3fd2e..0572e46c7fb 100644 --- a/torchvision/utils.py +++ b/torchvision/utils.py @@ -1,10 +1,15 @@ +from typing import Union, Optional, Sequence, Tuple, Text, BinaryIO +import io +import pathlib import torch import math irange = range -def make_grid(tensor, nrow=8, padding=2, - normalize=False, range=None, scale_each=False, pad_value=0): +def make_grid(tensor: Union[torch.Tensor, Sequence[torch.Tensor]], nrow: int = 8, + padding: int = 2, normalize: bool = False, + range: Optional[Tuple[int, int]] = None, scale_each: bool = False, + pad_value: int = 0) -> torch.Tensor: """Make a grid of images. Args: @@ -88,8 +93,9 @@ def norm_range(t, range): return grid -def save_image(tensor, fp, nrow=8, padding=2, - normalize=False, range=None, scale_each=False, pad_value=0, format=None): +def save_image(tensor: Union[torch.Tensor, Sequence[torch.Tensor]], fp: Union[Text, pathlib.Path, BinaryIO], + nrow: int = 8, padding: int = 2, normalize: bool = False, range: Optional[Tuple[int, int]] = None, + scale_each: bool = False, pad_value: int = 0, format: Optional[str] = None) -> None: """Save a given Tensor into an image file. Args: