fixed normalization

app.py

@@ -54,7 +54,7 @@ else:
     
     image_dict = imageLoader.transform(image)
     
-    image = imageLoader.to_image(image_dict["image"].squeeze(0))
+#     image = imageLoader.to_image(image_dict["image"].squeeze(0))
     
     show = st.image(image)
     show.image(image, "Your Image")

model.py

@@ -11,7 +11,7 @@ import torchvision
 import wordsegment as ws
 
 from virtex.config import Config
-from virtex.factories import TokenizerFactory, PretrainingModelFactory
+from virtex.factories import TokenizerFactory, PretrainingModelFactory, ImageTransformsFactory
 from virtex.utils.checkpointing import CheckpointManager
 
 CONFIG_PATH = "config.yaml"
@@ -21,12 +21,17 @@ SAMPLES_PATH = "./samples/*.jpg"
 
 class ImageLoader():
     def __init__(self):
-        self.transformer = torchvision.transforms.Compose([torchvision.transforms.Resize(256),
-                               torchvision.transforms.CenterCrop(224),
+        self.transformer = torchvision.transforms.Compose([ImageTransformsFactory.create("smallest_resize"),
+                               ImageTransformsFactory.create("center_crop"),
+                               ImageTransformsFactory.create("normalize"),
                                torchvision.transforms.ToTensor()])
     def load(self, im_path):
         im = torch.FloatTensor(self.transformer(Image.open(im_path))).unsqueeze(0)
         return {"image": im}
+    
+    def raw_load(self, im_path):
+        im = torch.FloatTensor(Image.open(im_path)).unsqueeze(0)
+        return {"image": im}
     def transform(self, image):
         im = torch.FloatTensor(self.transformer(image)).unsqueeze(0)
         return {"image": im}