Split transform and evaluate method

1 files changed, 22 insertions, 5 deletions

diff --git a/models/model.py b/models/model.py
index f79b832..0c3e5eb 100644
--- a/models/model.py
+++ b/models/model.py
@@ -269,6 +269,24 @@ class Model:
             ],
             dataloader_config: DataloaderConfiguration,
     ) -> dict[str, torch.Tensor]:
+        # Transform data to features
+        gallery_samples, probe_samples = self.transform(
+            iters, dataset_config, dataset_selectors, dataloader_config
+        )
+        # Evaluate features
+        accuracy = self.evaluate(gallery_samples, probe_samples)
+
+        return accuracy
+
+    def transform(
+            self,
+            iters: tuple[int],
+            dataset_config: DatasetConfiguration,
+            dataset_selectors: dict[
+                str, dict[str, Union[ClipClasses, ClipConditions, ClipViews]]
+            ],
+            dataloader_config: DataloaderConfiguration
+    ):
         self.is_train = False
         # Split gallery and probe dataset
         gallery_dataloader, probe_dataloaders = self._split_gallery_probe(
@@ -278,15 +296,15 @@ class Model:
         checkpoints = self._load_pretrained(
             iters, dataset_config, dataset_selectors
         )
+
         # Init models
         model_hp = self.hp.get('model', {})
         self.rgb_pn = RGBPartNet(ae_in_channels=self.in_channels, **model_hp)
         # Try to accelerate computation using CUDA or others
         self.rgb_pn = self.rgb_pn.to(self.device)
-
         self.rgb_pn.eval()
-        gallery_samples, probe_samples = [], {}
 
+        gallery_samples, probe_samples = [], {}
         # Gallery
         checkpoint = torch.load(list(checkpoints.values())[0])
         self.rgb_pn.load_state_dict(checkpoint['model_state_dict'])
@@ -294,7 +312,6 @@ class Model:
                            desc='Transforming gallery', unit='clips'):
             gallery_samples.append(self._get_eval_sample(sample))
         gallery_samples = default_collate(gallery_samples)
-
         # Probe
         for (condition, dataloader) in probe_dataloaders.items():
             checkpoint = torch.load(checkpoints[condition])
@@ -306,7 +323,7 @@ class Model:
                 probe_samples_c.append(self._get_eval_sample(sample))
             probe_samples[condition] = default_collate(probe_samples_c)
 
-        return self._evaluate(gallery_samples, probe_samples)
+        return gallery_samples, probe_samples
 
     def _get_eval_sample(self, sample: dict[str, Union[list, torch.Tensor]]):
         label = sample.pop('label').item()
@@ -318,7 +335,7 @@ class Model:
             **{'feature': feature}
         }
 
-    def _evaluate(
+    def evaluate(
             self,
             gallery_samples: dict[str, Union[list[str], torch.Tensor]],
             probe_samples: dict[str, dict[str, Union[list[str], torch.Tensor]]],