Manage image resizing

handler.py

@@ -5,7 +5,7 @@ import base64
 import numpy as np
 from PIL import Image
 import torch
-from transformers import SamModel, SamProcessor
+from transformers import SamModel, SamImageProcessor
 from typing import Dict, List, Any
 import torch.nn.functional as F
 
@@ -21,13 +21,13 @@ class EndpointHandler():
         try:
             # Load the model and processor from the local path
             self.model = SamModel.from_pretrained(path).to(device)
-            self.processor = SamProcessor.from_pretrained(path)
+            self.processor = SamImageProcessor.from_pretrained(path, do_resize=False)
         except Exception as e:
             # Fallback to loading from a known SAM model if local loading fails
             print("Failed to load from local path: {}".format(e))
             print("Attempting to load from facebook/sam-vit-base")
             self.model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
-            self.processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
+            self.processor = SamImageProcessor.from_pretrained("facebook/sam-vit-base", do_resize=False)
 
     def __call__(self, data):
         """
@@ -52,11 +52,11 @@ class EndpointHandler():
             raise TypeError("Unsupported input type. Expected a PIL Image or a base64 encoded string.")
         
         # 2. Prepare prompts and process the image
-        height, width = img.size[1], img.size[0]
-        input_points = [[[width // 2, height // 2]]]
-        input_labels = [[1]]
+        # height, width = img.size[1], img.size[0]
+        # input_points = [[[width // 2, height // 2]]]
+        # input_labels = [[1]]
         
-        inputs = self.processor(img, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(device)
+        inputs = self.processor(img, return_tensors="pt").to(device)
         
         # 3. Generate masks
         with torch.no_grad():