Update app.py

app.py

@@ -19,12 +19,10 @@ def segment_everything(image):
 def segment_box(image, x1, y1, x2, y2):
     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")
     with torch.no_grad():
         outputs = model(**inputs)
     preds = outputs.logits.squeeze().sigmoid()
-    
     segmentation = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
     segmentation[y1:y2, x1:x2] = (preds.numpy() * 255).astype(np.uint8)
     return Image.fromarray(segmentation)
@@ -32,9 +30,23 @@ def segment_box(image, x1, y1, x2, y2):
 def update_image(image, segmentation):
     if segmentation is None:
         return image
-    image_pil = Image.fromarray((image * 255).astype(np.uint8))
+    
+    # Ensure image is in the correct format (PIL Image)
+    if isinstance(image, np.ndarray):
+        image_pil = Image.fromarray((image * 255).astype(np.uint8))
+    else:
+        image_pil = image
+    
+    # Convert segmentation to RGBA
     seg_pil = Image.fromarray(segmentation).convert('RGBA')
+    
+    # Resize segmentation to match input image if necessary
+    if image_pil.size != seg_pil.size:
+        seg_pil = seg_pil.resize(image_pil.size, Image.NEAREST)
+    
+    # Blend images
     blended = Image.blend(image_pil.convert('RGBA'), seg_pil, 0.5)
+    
     return np.array(blended)
 
 with gr.Blocks() as demo:
@@ -52,19 +64,17 @@ with gr.Blocks() as demo:
                 box_btn = gr.Button("Box")
         with gr.Column(scale=1):
             output_image = gr.Image(label="Segmentation Result")
-    
+
     everything_btn.click(
         fn=segment_everything,
         inputs=[input_image],
         outputs=[output_image]
     )
-    
     box_btn.click(
         fn=segment_box,
         inputs=[input_image, x1_input, y1_input, x2_input, y2_input],
         outputs=[output_image]
     )
-    
     output_image.change(
         fn=update_image,
         inputs=[input_image, output_image],