Fixed exif rotation + squashing

handler.py

@@ -6,7 +6,10 @@ from urllib.request import urlopen
 
 import open_clip
 import torch
-from PIL import Image
+from PIL import Image, ImageOps
+from torchvision.transforms import Compose, Normalize, ToTensor
+
+INPUT_SIZE = 224
 
 
 def _is_git_lfs_pointer(path: Path) -> bool:
@@ -25,13 +28,14 @@ class EndpointHandler:
         self._validate_model_files()
 
         model_id = f"local-dir:{self.model_dir}"
-        self.model, self.preprocess = open_clip.create_model_from_pretrained(
+        self.model, preprocess = open_clip.create_model_from_pretrained(
             model_id,
             device=self.device,
             return_transform=True,
         )
         self.tokenizer = open_clip.get_tokenizer(model_id)
         self.model.eval()
+        self.tensor_preprocess = self._build_tensor_preprocess(preprocess)
 
     def _validate_model_files(self) -> None:
         config_path = self.model_dir / "open_clip_config.json"
@@ -61,16 +65,46 @@ class EndpointHandler:
                 "Upload the actual LFS blobs to the Hugging Face model repo before starting the endpoint."
             )
 
+    @staticmethod
+    def _build_tensor_preprocess(original_preprocess) -> Compose:
+        """Extract Normalize from the model's preprocess and build ToTensor + Normalize only.
+
+        The default model preprocess includes Resize + CenterCrop + ToTensor + Normalize.
+        Since we manually squash images to INPUT_SIZE x INPUT_SIZE, we only need
+        ToTensor + Normalize to match the existing embedding pipeline.
+        """
+        normalize = None
+        for t in original_preprocess.transforms:
+            if isinstance(t, Normalize):
+                normalize = t
+                break
+        if normalize is None:
+            normalize = Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
+        return Compose([ToTensor(), normalize])
+
+    @staticmethod
+    def _prepare_image(img: Image.Image) -> Image.Image:
+        """Squash image to INPUT_SIZE x INPUT_SIZE."""
+        return img.resize((INPUT_SIZE, INPUT_SIZE), Image.BICUBIC)
+
     def _load_image(self, image_input: Any) -> Image.Image | None:
         if not isinstance(image_input, str):
             return None
 
         if image_input.startswith(("http://", "https://")):
             with urlopen(image_input, timeout=10) as response:
-                return Image.open(io.BytesIO(response.read())).convert("RGB")
+                img = Image.open(io.BytesIO(response.read()))
+        else:
+            image_bytes = base64.b64decode(image_input.split(",")[-1])
+            img = Image.open(io.BytesIO(image_bytes))
+
+        img = ImageOps.exif_transpose(img)
+        return img.convert("RGB")
 
-        image_bytes = base64.b64decode(image_input.split(",")[-1])
-        return Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    def _preprocess_image(self, image: Image.Image) -> torch.Tensor:
+        """Squash to INPUT_SIZE and apply tensor normalization."""
+        image = self._prepare_image(image)
+        return self.tensor_preprocess(image).unsqueeze(0).to(self.device)
 
     def _tokenize_text(self, text: str | List[str]) -> torch.Tensor:
         texts = text if isinstance(text, list) else [text]
@@ -86,7 +120,7 @@ class EndpointHandler:
 
         with torch.no_grad():
             if image is not None and text is not None:
-                image_tensor = self.preprocess(image).unsqueeze(0).to(self.device)
+                image_tensor = self._preprocess_image(image)
                 text_tensor = self._tokenize_text(text)
 
                 image_features = self.model.encode_image(image_tensor, normalize=True)
@@ -99,7 +133,7 @@ class EndpointHandler:
                     response["text_embedding"] = text_features[0].cpu().tolist()
                 return response
             elif image is not None:
-                image_tensor = self.preprocess(image).unsqueeze(0).to(self.device)
+                image_tensor = self._preprocess_image(image)
                 image_features = self.model.encode_image(image_tensor, normalize=True)
                 return {"image_embedding": image_features[0].cpu().tolist()}
             elif text is not None: