summaryrefslogtreecommitdiff
path: root/preprocess.py
diff options
context:
space:
mode:
authorJordan Gong <jordan.gong@protonmail.com>2021-02-17 20:27:11 +0800
committerJordan Gong <jordan.gong@protonmail.com>2021-02-17 20:27:11 +0800
commit3cc4c018c214c4a0e0fbaaa2996390e66fd8ff2b (patch)
treeadb56ca85c469f8be594fe47677910ec9dc9255b /preprocess.py
parent09822ca28894c5b37de06ee24ff47d3b0de83380 (diff)
parent31af04b70fcc0e7b46cbe18a52d150eb4e274f0e (diff)
Merge branch 'python3.8' into python3.7
Diffstat (limited to 'preprocess.py')
-rw-r--r--preprocess.py5
1 files changed, 3 insertions, 2 deletions
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