make decorator annotations more concrete

Author: pmeier

torchvision/prototype/features/_image.py

@@ -3,6 +3,7 @@
 import warnings
 from typing import Any, Optional, Union, Tuple, cast
 
+import PIL.Image
 import torch
 from torchvision.prototype.utils._internal import StrEnum
 from torchvision.transforms.functional import to_pil_image
@@ -78,7 +79,7 @@ def guess_color_space(data: torch.Tensor) -> ColorSpace:
     def show(self) -> None:
         # TODO: this is useful for developing and debugging but we should remove or at least revisit this before we
         #  promote this out of the prototype state
-        to_pil_image(make_grid(self.view(-1, *self.shape[-3:]))).show()
+        cast(PIL.Image.Image, to_pil_image(make_grid(self.view(-1, *self.shape[-3:])))).show()
 
     def draw_bounding_box(self, bounding_box: BoundingBox, **kwargs: Any) -> Image:
         # TODO: this is useful for developing and debugging but we should remove or at least revisit this before we

torchvision/prototype/transforms/_presets.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple, cast
+from typing import Dict, Optional, Tuple
 
 import torch
 from torch import Tensor, nn
@@ -41,7 +41,7 @@ def forward(self, img: Tensor) -> Tensor:
             img = F.pil_to_tensor(img)
         img = F.convert_image_dtype(img, torch.float)
         img = F.normalize(img, mean=self._mean, std=self._std)
-        return cast(Tensor, img)
+        return img
 
 
 class Kinect400Eval(nn.Module):
@@ -65,7 +65,7 @@ def forward(self, vid: Tensor) -> Tensor:
         vid = F.resize(vid, self._size, interpolation=self._interpolation)
         vid = F.center_crop(vid, self._crop_size)
         vid = F.convert_image_dtype(vid, torch.float)
-        vid = cast(Tensor, F.normalize(vid, mean=self._mean, std=self._std))
+        vid = F.normalize(vid, mean=self._mean, std=self._std)
         return vid.permute(1, 0, 2, 3)  # (T, C, H, W) => (C, T, H, W)
 
 

torchvision/prototype/transforms/kernels/_geometry.py

@@ -39,7 +39,7 @@ def resize_image(
     new_height, new_width = size
     num_channels, old_height, old_width = image.shape[-3:]
     batch_shape = image.shape[:-3]
-    return _F.resize(  # type: ignore[no-any-return]
+    return _F.resize(
         image.reshape((-1, num_channels, old_height, old_width)),
         size=size,
         interpolation=interpolation,

torchvision/transforms/functional.py

@@ -21,9 +21,9 @@
 from . import functional_tensor as F_t
 
 
-def log_api_usage_once(fn: Callable) -> Callable:
+def log_api_usage_once(fn: Callable[..., Tensor]) -> Callable[..., Tensor]:
     @functools.wraps(fn)
-    def wrapper(*args, **kwargs):
+    def wrapper(*args: Any, **kwargs: Any) -> Tensor:
         _log_api_usage_once(fn)
         return fn(*args, **kwargs)