Changed image processor

theia_model.py

@@ -296,6 +296,39 @@ class ViTModelReg(ViTModel):
             ).to(module.cls_token.dtype)
 
 
+class TorchImageProcessor:
+    def __init__(self, processor):
+        # converts huggingface image processor to torch processor
+        self.mean = torch.tensor(processor.image_mean, dtype=torch.float32).reshape((1, 3, 1, 1))
+        self.std = torch.tensor(processor.image_std, dtype=torch.float32).reshape((1, 3, 1, 1))
+        self.width = processor.size['width']
+        self.height = processor.size['height']
+
+    def __call__(self, x,
+                 do_resize: bool = True,
+                 do_rescale: bool = True,
+                 do_normalize: bool = True,
+                 device='cuda'):
+        #x = torch.tensor(x, device=device, dtype=torch.float32)
+        if do_resize:
+            #assert x.shape[-1] == self.width
+            #assert x.shape[-2] == self.height
+            x = F.interpolate(
+                x,
+                size=(self.height, self.width),
+                mode='bilinear',
+                align_corners=False
+            )
+            
+            # not implemented. If you really need resize on each forward step, use torch.interpolate
+        if do_rescale:
+            x = x / 255.
+        if do_normalize:
+            x = x - self.mean.to(device)
+            x = x / self.std.to(device)
+        return {'pixel_values': x}
+
+
 class DeiT(nn.Module):
     """DeiT model.
 
@@ -326,7 +359,9 @@ class DeiT(nn.Module):
 
         self.model.pooler = nn.Identity()
 
-        self.processor = AutoProcessor.from_pretrained(model_name)
+        #self.processor = AutoProcessor.from_pretrained(model_name)
+        self.processor = AutoProcessor.from_pretrained(model_name, use_fast=True)
+        self.gpu_processor = TorchImageProcessor(self.processor)
 
     def get_feature_size(
         self,
@@ -378,9 +413,13 @@ class DeiT(nn.Module):
         Returns:
             torch.Tensor: model output.
         """
-        input = self.processor(
-            x, return_tensors="pt", do_resize=do_resize, do_rescale=do_rescale, do_normalize=do_normalize
-        ).to(self.model.device)
+        #input = self.processor(
+        #    x, return_tensors="pt", do_resize=do_resize, do_rescale=do_rescale, do_normalize=do_normalize
+        #).to(self.model.device)
+        if x.shape[-1] == 3:
+            x = x.permute(0, 3, 1, 2)
+        input = self.gpu_processor(x, device=self.model.device)
+
         y = self.model(**input, interpolate_pos_encoding=interpolate_pos_encoding)
         return y.last_hidden_state
 
@@ -1492,4 +1531,4 @@ class TheiaModel(PreTrainedModel):
             "mse_losses_per_model": mse_losses_per_model,
             "cos_losses_per_model": cos_losses_per_model,
             "l1_losses_per_model": l1_losses_per_model,
-        }
+        }