project Refactor, CNN Classifier Basics

dataset/dataset.py

@@ -11,12 +11,17 @@ from lib.preprocessing.generator import Generator
 
 
 class TrajPairDataset(Dataset):
+    @property
+    def map_shape(self):
+        return self._dataset.map.as_array.shape
+
     def __init__(self, data):
         super(TrajPairDataset, self).__init__()
         self.alternatives = data['alternatives']
         self.trajectory = data['trajectory']
         self.labels = data['labels']
         self.mapname = data['map']['name'][4:] if data['map']['name'].startswith('map_') else data['map']['name']
+        self.map = data['map']['map']
 
     def __len__(self):
         return len(self.alternatives)
@@ -26,7 +31,8 @@ class TrajPairDataset(Dataset):
 
 
 class DatasetMapping(Dataset):
-    def __init__(self, dataset, mapping):
+
+    def __init__(self, dataset: Union[TrajPairDataset, ConcatDataset], mapping):
         self._dataset = dataset
         self._mapping = mapping
 
@@ -38,6 +44,15 @@ class DatasetMapping(Dataset):
 
 
 class TrajPairData(object):
+    @property
+    def map_shapes(self):
+        return [dataset.map_shape for dataset in self._dataset.datasets]
+
+    @property
+    def map_shapes_max(self):
+        shapes = self.map_shapes
+        return map(max, zip(*shapes))
+
     @property
     def name(self):
         return self.__class__.__name__
@@ -52,8 +67,7 @@ class TrajPairData(object):
         self.mapname = mapname
         self.train_split, self.val_split, self.test_split = train_val_test_split
         self.data_root = Path(data_root)
-        self.maps_root = Path(data_root) if data_root else Path() / 'res' / 'maps'
-        self._dataset = None
+        self.maps_root = Path(map_root) if map_root else Path() / 'res' / 'maps'
         self._dataset, self._train_map, self._val_map, self._test_map = self._load_dataset()
 
     def _build_data_on_demand(self):