diff options
Diffstat (limited to 'models/model.py')
| -rw-r--r-- | models/model.py | 54 |
1 files changed, 27 insertions, 27 deletions
diff --git a/models/model.py b/models/model.py index 0829f33..46987ca 100644 --- a/models/model.py +++ b/models/model.py @@ -1,7 +1,7 @@ import copy import os import random -from typing import Union, Optional +from typing import Union, Optional, Tuple, List, Dict, Set import numpy as np import torch @@ -55,11 +55,11 @@ class Model: self.is_train: bool = True self.in_channels: int = 3 - self.in_size: tuple[int, int] = (64, 48) + self.in_size: Tuple[int, int] = (64, 48) self.batch_size: Optional[int] = None - self._gallery_dataset_meta: Optional[dict[str, list]] = None - self._probe_datasets_meta: Optional[dict[str, dict[str, list]]] = None + self._gallery_dataset_meta: Optional[Dict[str, List]] = None + self._probe_datasets_meta: Optional[Dict[str, Dict[str, List]]] = None self._model_name: str = self.meta.get('name', 'RGB-GaitPart') self._hp_sig: str = self._make_signature(self.hp) @@ -108,8 +108,8 @@ class Model: def fit_all( self, dataset_config: DatasetConfiguration, - dataset_selectors: dict[ - str, dict[str, Union[ClipClasses, ClipConditions, ClipViews]] + dataset_selectors: Dict[ + str, Dict[str, Union[ClipClasses, ClipConditions, ClipViews]] ], dataloader_config: DataloaderConfiguration, ): @@ -157,7 +157,7 @@ class Model: )) # Prepare for model, optimizer and scheduler model_hp: dict = self.hp.get('model', {}).copy() - optim_hp: dict = self.hp.get('optimizer', {}).copy() + optim_hp: Dict = self.hp.get('optimizer', {}).copy() sched_hp = self.hp.get('scheduler', {}) self.rgb_pn = RGBPartNet(self.in_channels, self.in_size, **model_hp, @@ -283,10 +283,10 @@ class Model: def transform( self, - iters: tuple[int], + iters: Tuple[int], dataset_config: DatasetConfiguration, - dataset_selectors: dict[ - str, dict[str, Union[ClipClasses, ClipConditions, ClipViews]] + dataset_selectors: Dict[ + str, Dict[str, Union[ClipClasses, ClipConditions, ClipViews]] ], dataloader_config: DataloaderConfiguration, is_train: bool = False @@ -331,7 +331,7 @@ class Model: return gallery_samples, probe_samples - def _get_eval_sample(self, sample: dict[str, Union[list, torch.Tensor]]): + def _get_eval_sample(self, sample: Dict[str, Union[List, torch.Tensor]]): label, condition, view, clip = sample.values() with torch.no_grad(): feature_c, feature_p = self.rgb_pn(clip.to(self.device)) @@ -344,12 +344,12 @@ class Model: def _load_pretrained( self, - iters: tuple[int], + iters: Tuple[int], dataset_config: DatasetConfiguration, - dataset_selectors: dict[ - str, dict[str, Union[ClipClasses, ClipConditions, ClipViews]] + dataset_selectors: Dict[ + str, Dict[str, Union[ClipClasses, ClipConditions, ClipViews]] ] - ) -> dict[str, str]: + ) -> Dict[str, str]: checkpoints = {} for (iter_, total_iter, (condition, selector)) in zip( iters, self.total_iters, dataset_selectors.items() @@ -368,7 +368,7 @@ class Model: dataset_config: DatasetConfiguration, dataloader_config: DataloaderConfiguration, is_train: bool = False - ) -> tuple[DataLoader, dict[str, DataLoader]]: + ) -> Tuple[DataLoader, Dict[str, DataLoader]]: dataset_name = dataset_config.get('name', 'CASIA-B') if dataset_name == 'CASIA-B': self.is_train = is_train @@ -425,7 +425,7 @@ class Model: dataset_config, popped_keys=['root_dir', 'cache_on'] ) - config: dict = dataset_config.copy() + config: Dict = dataset_config.copy() name = config.pop('name', 'CASIA-B') if name == 'CASIA-B': return CASIAB(**config, is_train=self.is_train) @@ -439,7 +439,7 @@ class Model: dataset: Union[CASIAB], dataloader_config: DataloaderConfiguration ) -> DataLoader: - config: dict = dataloader_config.copy() + config: Dict = dataloader_config.copy() self.batch_size = config.pop('batch_size', 16) if self.is_train: dis_sampler = DisentanglingSampler(dataset, self.batch_size) @@ -452,9 +452,9 @@ class Model: def _batch_splitter( self, - batch: list[dict[str, Union[np.int64, str, torch.Tensor]]] - ) -> tuple[dict[str, Union[list[str], torch.Tensor]], - dict[str, Union[list[str], torch.Tensor]]]: + batch: List[Dict[str, Union[np.int64, str, torch.Tensor]]] + ) -> Tuple[Dict[str, Union[List[str], torch.Tensor]], + Dict[str, Union[List[str], torch.Tensor]]]: """ Disentanglement need two random conditions, this function will split batch_size * 2 samples to 2 dicts each containing @@ -467,8 +467,8 @@ class Model: return default_collate(batch_0), default_collate(batch_1) def _make_signature(self, - config: dict, - popped_keys: Optional[list] = None) -> str: + config: Dict, + popped_keys: Optional[List] = None) -> str: _config = config.copy() if popped_keys: for key in popped_keys: @@ -476,16 +476,16 @@ class Model: return self._gen_sig(list(_config.values())) - def _gen_sig(self, values: Union[tuple, list, set, str, int, float]) -> str: + def _gen_sig(self, values: Union[Tuple, List, Set, str, int, float]) -> str: strings = [] for v in values: if isinstance(v, str): strings.append(v) - elif isinstance(v, (tuple, list)): + elif isinstance(v, (Tuple, List)): strings.append(self._gen_sig(v)) - elif isinstance(v, set): + elif isinstance(v, Set): strings.append(self._gen_sig(sorted(list(v)))) - elif isinstance(v, dict): + elif isinstance(v, Dict): strings.append(self._gen_sig(list(v.values()))) else: strings.append(str(v)) |