Unverified Commit a8d2227f authored by F-G Fernandez's avatar F-G Fernandez Committed by GitHub
Browse files

Added typing annotations to io/__init__ (#4224)



* style: Added typing annotations

* Specified types for iter and seek.
Co-authored-by: default avatarVasilis Vryniotis <datumbox@users.noreply.github.com>
Co-authored-by: default avatarVasilis Vryniotis <vvryniotis@fb.com>
parent 5e254127
import torch import torch
from typing import Any, Dict, Iterator
from ._video_opt import ( from ._video_opt import (
Timebase, Timebase,
...@@ -33,13 +34,13 @@ from .image import ( ...@@ -33,13 +34,13 @@ from .image import (
if _HAS_VIDEO_OPT: if _HAS_VIDEO_OPT:
def _has_video_opt(): def _has_video_opt() -> bool:
return True return True
else: else:
def _has_video_opt(): def _has_video_opt() -> bool:
return False return False
...@@ -99,7 +100,7 @@ class VideoReader: ...@@ -99,7 +100,7 @@ class VideoReader:
Currently available options include ``['video', 'audio']`` Currently available options include ``['video', 'audio']``
""" """
def __init__(self, path, stream="video"): def __init__(self, path: str, stream: str = "video") -> None:
if not _has_video_opt(): if not _has_video_opt():
raise RuntimeError( raise RuntimeError(
"Not compiled with video_reader support, " "Not compiled with video_reader support, "
...@@ -109,7 +110,7 @@ class VideoReader: ...@@ -109,7 +110,7 @@ class VideoReader:
) )
self._c = torch.classes.torchvision.Video(path, stream) self._c = torch.classes.torchvision.Video(path, stream)
def __next__(self): def __next__(self) -> Dict[str, Any]:
"""Decodes and returns the next frame of the current stream. """Decodes and returns the next frame of the current stream.
Frames are encoded as a dict with mandatory Frames are encoded as a dict with mandatory
data and pts fields, where data is a tensor, and pts is a data and pts fields, where data is a tensor, and pts is a
...@@ -126,10 +127,10 @@ class VideoReader: ...@@ -126,10 +127,10 @@ class VideoReader:
raise StopIteration raise StopIteration
return {"data": frame, "pts": pts} return {"data": frame, "pts": pts}
def __iter__(self): def __iter__(self) -> Iterator['VideoReader']:
return self return self
def seek(self, time_s: float): def seek(self, time_s: float) -> 'VideoReader':
"""Seek within current stream. """Seek within current stream.
Args: Args:
...@@ -144,7 +145,7 @@ class VideoReader: ...@@ -144,7 +145,7 @@ class VideoReader:
self._c.seek(time_s) self._c.seek(time_s)
return self return self
def get_metadata(self): def get_metadata(self) -> Dict[str, Any]:
"""Returns video metadata """Returns video metadata
Returns: Returns:
...@@ -152,7 +153,7 @@ class VideoReader: ...@@ -152,7 +153,7 @@ class VideoReader:
""" """
return self._c.get_metadata() return self._c.get_metadata()
def set_current_stream(self, stream: str): def set_current_stream(self, stream: str) -> bool:
"""Set current stream. """Set current stream.
Explicitly define the stream we are operating on. Explicitly define the stream we are operating on.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment