text_sam: Fixed issue when passing PIL Image as input to SamGeo2

opengeos · Nov 18, 2024 · aca2175 · aca2175
1 parent 94cfeab
commit aca2175
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/samgeo/text_sam.py b/samgeo/text_sam.py
@@ -229,7 +229,12 @@ def predict_sam(self, image, boxes):
             )
             return masks.cpu()
         elif self._sam_version == 2:
-            self.sam.set_image(self.source)
+
+            if isinstance(self.source, str):
+                self.sam.set_image(self.source)
+            # If no source is set provide PIL image
+            if self.source is None:
+                self.sam.set_image(image)
             self.sam.boxes = boxes.numpy().tolist()
             masks, _, _ = self.sam.predict(
                 boxes=boxes.numpy().tolist(),