Load image using torch + use training pre-processing function during prediction

Closes #122 (closed)

dan/manager/dataset.py

@@ -3,12 +3,10 @@ import json
 import os
 
 import numpy as np
-import torch
 from torch.utils.data import Dataset
-from torchvision.io import ImageReadMode, read_image
 
 from dan.datasets.utils import natural_sort
-from dan.utils import token_to_ind
+from dan.utils import read_image, token_to_ind
 
 
 class OCRDataset(Dataset):
@@ -82,14 +80,6 @@ class OCRDataset(Dataset):
         )
         return sample
 
-    @staticmethod
-    def load_image(path):
-        """
-        Load an image as a torch.Tensor and scale the values between 0 and 1.
-        """
-        img = read_image(path, mode=ImageReadMode.RGB)
-        return img.to(dtype=torch.get_default_dtype()).div(255)
-
     def load_samples(self, paths_and_sets):
         """
         Load images and labels
@@ -116,7 +106,7 @@ class OCRDataset(Dataset):
                 )
                 if self.load_in_memory:
                     samples[-1]["img"] = self.preprocessing_transforms(
-                        self.load_image(filename)
+                        read_image(filename)
                     )
         return samples
 
@@ -126,10 +116,8 @@ class OCRDataset(Dataset):
         """
         if self.load_in_memory:
             return self.samples[i]["img"]
-        else:
-            return self.preprocessing_transforms(
-                self.load_image(self.samples[i]["path"])
-            )
+
+        return self.preprocessing_transforms(read_image(self.samples[i]["path"]))
 
     def compute_final_size(self, img):
         """