Update app.py

app.py

@@ -17,29 +17,34 @@ def segment_everything(image):
     return Image.fromarray(segmentation)
 
 def segment_box(image, box):
+    if box is None:
+        return image
     x1, y1, x2, y2 = map(int, box)
-    mask = Image.new('L', image.size, 0)
+    mask = Image.new('L', (image.shape[1], image.shape[0]), 0)
     draw = ImageDraw.Draw(mask)
     draw.rectangle([x1, y1, x2, y2], fill=255)
     
-    inputs = processor(text=["object in box"], images=[image], mask_pixels=mask, padding="max_length", return_tensors="pt")
+    inputs = processor(text=["object in box"], images=[image], mask_pixels=np.array(mask), padding="max_length", return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
     preds = outputs.logits.squeeze().sigmoid()
     segmentation = (preds.numpy() * 255).astype(np.uint8)
     return Image.fromarray(segmentation)
 
-def update_image(image, segmentation, tool):
+def update_image(image, segmentation):
     if segmentation is None:
         return image
-    blended = Image.blend(image.convert('RGBA'), segmentation.convert('RGBA'), 0.5)
-    return blended
+    image_pil = Image.fromarray((image * 255).astype(np.uint8))
+    seg_pil = Image.fromarray(segmentation).convert('RGBA')
+    blended = Image.blend(image_pil.convert('RGBA'), seg_pil, 0.5)
+    return np.array(blended)
 
 with gr.Blocks() as demo:
     gr.Markdown("# Segment Anything-like Demo")
     with gr.Row():
         with gr.Column(scale=1):
-            input_image = gr.Image(label="Input Image", tool="select")
+            input_image = gr.Image(label="Input Image")
+            box_input = gr.Box(label="Select Box")
             with gr.Row():
                 everything_btn = gr.Button("Everything")
                 box_btn = gr.Button("Box")
@@ -54,13 +59,13 @@ with gr.Blocks() as demo:
     
     box_btn.click(
         fn=segment_box,
-        inputs=[input_image, input_image.sel],
+        inputs=[input_image, box_input],
         outputs=[output_image]
     )
     
     output_image.change(
         fn=update_image,
-        inputs=[input_image, output_image, gr.State("last_tool")],
+        inputs=[input_image, output_image],
         outputs=[output_image]
     )