summaryrefslogtreecommitdiff
path: root/preprocess.py
diff options
context:
space:
mode:
authorJordan Gong <jordan.gong@protonmail.com>2021-02-19 22:43:17 +0800
committerJordan Gong <jordan.gong@protonmail.com>2021-02-19 22:43:17 +0800
commit4049566103a00aa6d5a0b1f73569bdc5435714ca (patch)
treed84604773f05eab030ff2106c43cb2c091b6e8fc /preprocess.py
parentd12dd6b04a4e7c2b1ee43ab6f36f25d0c35ca364 (diff)
parent969030864495e7c2b419400fd81ee0fad83de41e (diff)
Merge branch 'python3.8' into disentangling_only_py3.8
# Conflicts: # models/hpm.py # models/layers.py # models/model.py # models/part_net.py # models/rgb_part_net.py # utils/configuration.py
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 91fa8c2..eef59ba 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