Evaluate specific sets only

Closes #320 (closed)

dan/ocr/manager/ocr.py

@@ -10,6 +10,7 @@ import torch
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
 
+from dan import TEST_NAME, TRAIN_NAME, VAL_NAME
 from dan.ocr.manager.dataset import OCRDataset
 from dan.ocr.transforms import get_augmentation_transforms, get_preprocessing_transforms
 from dan.utils import pad_images, pad_sequences_1D
@@ -66,8 +67,8 @@ class OCRDatasetManager:
         Load training and validation datasets
         """
         self.train_dataset = OCRDataset(
-            set_name="train",
-            paths_and_sets=self.get_paths_and_sets(self.params["train"]["datasets"]),
+            set_name=TRAIN_NAME,
+            paths_and_sets=self.get_paths_and_sets(self.params[TRAIN_NAME]["datasets"]),
             charset=self.charset,
             tokens=self.tokens,
             preprocessing_transforms=self.preprocessing,
@@ -79,10 +80,12 @@ class OCRDatasetManager:
 
         self.mean, self.std = self.train_dataset.compute_std_mean()
 
-        for custom_name in self.params["val"]:
+        for custom_name in self.params[VAL_NAME]:
             self.valid_datasets[custom_name] = OCRDataset(
-                set_name="val",
-                paths_and_sets=self.get_paths_and_sets(self.params["val"][custom_name]),
+                set_name=VAL_NAME,
+                paths_and_sets=self.get_paths_and_sets(
+                    self.params[VAL_NAME][custom_name]
+                ),
                 charset=self.charset,
                 tokens=self.tokens,
                 preprocessing_transforms=self.preprocessing,
@@ -168,7 +171,7 @@ class OCRDatasetManager:
                 {"path": self.params["datasets"][set_info[0]], "set_name": set_info[1]}
             )
         self.test_datasets[custom_name] = OCRDataset(
-            set_name="test",
+            set_name=TEST_NAME,
             paths_and_sets=paths_and_sets,
             charset=self.charset,
             tokens=self.tokens,