From 31af04b70fcc0e7b46cbe18a52d150eb4e274f0e Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Wed, 17 Feb 2021 20:26:54 +0800 Subject: Type hint fixes --- preprocess.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) (limited to 'preprocess.py') diff --git a/preprocess.py b/preprocess.py index 8af58f2..894c7c1 100644 --- a/preprocess.py +++ b/preprocess.py @@ -1,5 +1,6 @@ import glob import os +from typing import Tuple import torch import torchvision @@ -23,7 +24,7 @@ class CASIABClip(Dataset): video, *_ = torchvision.io.read_video(filename, pts_unit='sec') self.frames = video.permute(0, 3, 1, 2) / 255 - def __getitem__(self, index) -> tuple[int, torch.Tensor]: + def __getitem__(self, index) -> Tuple[int, torch.Tensor]: return index, self.frames[index] def __len__(self) -> int: @@ -35,7 +36,7 @@ model = model.to(DEVICE) model.eval() -def result_handler(frame_: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: +def result_handler(frame_: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: for (box, label, score, mask) in zip(*result.values()): x0, y0, x1, y1 = box height, width = y1 - y0, x1 - x0 -- cgit v1.2.3