[feat] Add ViLT image text processors

Add ViT image processor and bert text processor.

ghstack-source-id: 34098c0c17b6920a4230d28535bc09cddb8f99b2
Pull Request resolved: #1097

mmf/datasets/processors/bert_processors.py

@@ -353,3 +353,14 @@ def __init__(self, config, *args, **kwargs):
         self.fusion_strategy = config.get("fusion", "concat")
         self.tokenizer = RobertaTokenizer(config, *args, **kwargs)
         self._probability = config.get("mask_probability", 0)
+
+
+@registry.register_processor("vilt_text_processor")
+class VILTTextProcessor(BertTokenizer):
+    from omegaconf import OmegaConf
+
+    def __init__(self, config, *args, **kwargs):
+        config.tokenizer_config = self.OmegaConf.create(
+            {"type": "bert-base-uncased", "params": {"do_lower_case": True}}
+        )
+        super().__init__(config, *args, **kwargs)

mmf/datasets/processors/image_processors.py

@@ -162,3 +162,28 @@ def __call__(self, image):
 
             return padded_image
         return image
+
+
+@registry.register_processor("vilt_image_processor")
+class VILTImageProcessor(BaseProcessor):
+    from torchvision.transforms import Resize, ToTensor, Normalize, Compose
+
+    def __init__(self, config, *args, **kwargs):
+        image_size = getattr(config, "size", [224, 224])
+        transforms_list = []
+        transforms_list.append(self.Resize(image_size))
+        transforms_list.append(self.ToTensor())
+        transforms_list.append(GrayScaleTo3Channels())
+
+        mean = [0.485, 0.456, 0.406]
+        std = [0.229, 0.224, 0.225]
+        transforms_list.append(self.Normalize(mean, std))
+        self.transform = self.Compose(transforms_list)
+
+    def __call__(self, x):
+        # Support both dict and normal mode
+        if isinstance(x, collections.abc.Mapping):
+            x = x["image"]
+            return {"image": self.transform(x)}
+        else:
+            return self.transform(x)