From 3a8c48c722f687a47180221f807bd0a59c9f616e Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Sun, 27 Dec 2020 14:43:27 +0800 Subject: Add dataset selector to config type hint, change ClipLabels typo to ClipViews --- utils/configuration.py | 5 ++++- utils/dataset.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) (limited to 'utils') diff --git a/utils/configuration.py b/utils/configuration.py index 84bd064..445badd 100644 --- a/utils/configuration.py +++ b/utils/configuration.py @@ -1,7 +1,9 @@ -from typing import TypedDict +from typing import TypedDict, Optional, Union import torch +from utils.dataset import ClipLabels, ClipConditions, ClipViews + class SystemConfiguration(TypedDict): device: torch.device @@ -15,6 +17,7 @@ class DatasetConfiguration(TypedDict): train_size: int num_sampled_frames: int discard_threshold: int + selector: Optional[dict[str, Union[ClipLabels, ClipConditions, ClipViews]]] num_input_channels: int frame_size: tuple[int, int] cache_on: bool diff --git a/utils/dataset.py b/utils/dataset.py index 9f9229a..1a9d595 100644 --- a/utils/dataset.py +++ b/utils/dataset.py @@ -26,7 +26,7 @@ class CASIAB(data.Dataset): num_sampled_frames: int = 30, discard_threshold: int = 15, selector: Optional[dict[ - str, Union[ClipLabels, ClipConditions, ClipLabels] + str, Union[ClipLabels, ClipConditions, ClipViews] ]] = None, num_input_channels: int = 3, frame_size: tuple[int, int] = (64, 32), -- cgit v1.2.3