fix

inference.py

@@ -202,18 +202,32 @@ class CustomSamPredictor(SamPredictor):
     ) -> None:
         super().__init__(sam_model)
 
-    def encode_image(self, image: np.ndarray, image_format: str = "RGB") -> torch.Tensor:
+    def encode_image(
+        self,
+        image: np.ndarray,
+        image_format: str = "RGB",
+    ) -> None:
         """
-        Encodes the image and returns its embedding.
+        Calculates the image embeddings for the provided image, allowing
+        masks to be predicted with the 'predict' method.
 
         Arguments:
-          image (np.ndarray): The image for which to calculate the embedding.
+          image (np.ndarray): The image for calculating masks. Expects an
+            image in HWC uint8 format, with pixel values in [0, 255].
           image_format (str): The color format of the image, in ['RGB', 'BGR'].
-
-        Returns:
-          torch.Tensor: The image embedding with shape 1xCxHxW.
         """
-        self.set_image(image, image_format)
+        assert image_format in [
+            "RGB",
+            "BGR",
+        ], f"image_format must be in ['RGB', 'BGR'], is {image_format}."
+        if image_format != self.model.image_format:
+            image = image[..., ::-1]
+
+        # Transform the image to the form expected by the model
+        input_image = self.transform.apply_image(image)
+        input_image_torch = torch.as_tensor(input_image, device=self.device)
+        input_image_torch = input_image_torch.permute(2, 0, 1).contiguous()[None, :, :, :]
+        self.set_torch_image(input_image_torch, image.shape[:2])
         return self.get_image_embedding()
 
     def decode_and_predict(