Remove TypeDict for python 3.7

1 files changed, 14 insertions, 17 deletions

diff --git a/models/model.py b/models/model.py
index 1154d7f..7cf6ed0 100644
--- a/models/model.py
+++ b/models/model.py
@@ -13,9 +13,6 @@ from torch.utils.tensorboard import SummaryWriter
 from tqdm import tqdm
 
 from models.rgb_part_net import RGBPartNet
-from utils.configuration import DataloaderConfiguration, \
-    HyperparameterConfiguration, DatasetConfiguration, ModelConfiguration, \
-    SystemConfiguration
 from utils.dataset import CASIAB, ClipConditions, ClipViews, ClipClasses
 from utils.sampler import TripletSampler
 
@@ -23,9 +20,9 @@ from utils.sampler import TripletSampler
 class Model:
     def __init__(
             self,
-            system_config: SystemConfiguration,
-            model_config: ModelConfiguration,
-            hyperparameter_config: HyperparameterConfiguration
+            system_config: Dict,
+            model_config: Dict,
+            hyperparameter_config: Dict
     ):
         self.disable_acc = system_config['disable_acc']
         if self.disable_acc:
@@ -89,11 +86,11 @@ class Model:
 
     def fit_all(
             self,
-            dataset_config: DatasetConfiguration,
+            dataset_config: Dict,
             dataset_selectors: Dict[
                 str, Dict[str, Union[ClipClasses, ClipConditions, ClipViews]]
             ],
-            dataloader_config: DataloaderConfiguration,
+            dataloader_config: Dict,
     ):
         for (condition, selector) in dataset_selectors.items():
             print(f'Training model {condition} ...')
@@ -104,8 +101,8 @@ class Model:
 
     def fit(
             self,
-            dataset_config: DatasetConfiguration,
-            dataloader_config: DataloaderConfiguration,
+            dataset_config: Dict,
+            dataloader_config: Dict,
     ):
         self.is_train = True
         dataset = self._parse_dataset_config(dataset_config)
@@ -184,11 +181,11 @@ class Model:
     def predict_all(
             self,
             iter_: int,
-            dataset_config: DatasetConfiguration,
+            dataset_config: dict,
             dataset_selectors: Dict[
                 str, Dict[str, Union[ClipClasses, ClipConditions, ClipViews]]
             ],
-            dataloader_config: DataloaderConfiguration,
+            dataloader_config: dict,
     ) -> Dict[str, torch.Tensor]:
         self.is_train = False
         # Split gallery and probe dataset
@@ -296,7 +293,7 @@ class Model:
     def _load_pretrained(
             self,
             iter_: int,
-            dataset_config: DatasetConfiguration,
+            dataset_config: Dict,
             dataset_selectors: Dict[
                 str, Dict[str, Union[ClipClasses, ClipConditions, ClipViews]]
             ]
@@ -313,8 +310,8 @@ class Model:
 
     def _split_gallery_probe(
             self,
-            dataset_config: DatasetConfiguration,
-            dataloader_config: DataloaderConfiguration,
+            dataset_config: Dict,
+            dataloader_config: Dict,
     ) -> Tuple[DataLoader, Dict[str: DataLoader]]:
         dataset_name = dataset_config.get('name', 'CASIA-B')
         if dataset_name == 'CASIA-B':
@@ -364,7 +361,7 @@ class Model:
 
     def _parse_dataset_config(
             self,
-            dataset_config: DatasetConfiguration
+            dataset_config: Dict
     ) -> Union[CASIAB]:
         self.train_size = dataset_config.get('train_size', 74)
         self.in_channels = dataset_config.get('num_input_channels', 3)
@@ -385,7 +382,7 @@ class Model:
     def _parse_dataloader_config(
             self,
             dataset: Union[CASIAB],
-            dataloader_config: DataloaderConfiguration
+            dataloader_config: Dict
     ) -> DataLoader:
         config: Dict = dataloader_config.copy()
         (self.pr, self.k) = config.pop('batch_size')