Update app.py

app.py

@@ -9,6 +9,14 @@ processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
 model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
 
 def segment_everything(image):
+    # Check if image is a list and extract the actual image data
+    if isinstance(image, list):
+        image = image[0]
+    
+    # Convert numpy array to PIL Image
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    
     inputs = processor(text=["object"], images=[image], padding="max_length", return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
@@ -17,9 +25,17 @@ def segment_everything(image):
     return Image.fromarray(segmentation)
 
 def segment_box(image, x1, y1, x2, y2):
+    # Check if image is a list and extract the actual image data
+    if isinstance(image, list):
+        image = image[0]
+    
+    # Convert PIL Image to numpy array if necessary
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+    
     x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
     cropped_image = image[y1:y2, x1:x2]
-    inputs = processor(text=["object"], images=[cropped_image], padding="max_length", return_tensors="pt")
+    inputs = processor(text=["object"], images=[Image.fromarray(cropped_image)], padding="max_length", return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
     preds = outputs.logits.squeeze().sigmoid()
@@ -31,9 +47,13 @@ def update_image(image, segmentation):
     if segmentation is None:
         return image
     
+    # Check if image is a list and extract the actual image data
+    if isinstance(image, list):
+        image = image[0]
+    
     # Ensure image is in the correct format (PIL Image)
     if isinstance(image, np.ndarray):
-        image_pil = Image.fromarray((image * 255).astype(np.uint8))
+        image_pil = Image.fromarray(image)
     else:
         image_pil = image