Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import numpy as np
-from PIL import Image, ImageDraw
+from PIL import Image
 import torch
 from transformers import AutoProcessor, CLIPSegForImageSegmentation
 
@@ -19,16 +19,16 @@ def segment_everything(image):
 def segment_box(image, box):
     if box is None:
         return image
-    x1, y1, x2, y2 = map(int, box)
-    mask = Image.new('L', (image.shape[1], image.shape[0]), 0)
-    draw = ImageDraw.Draw(mask)
-    draw.rectangle([x1, y1, x2, y2], fill=255)
+    x1, y1, x2, y2 = box
+    cropped_image = image[y1:y2, x1:x2]
     
-    inputs = processor(text=["object in box"], images=[image], mask_pixels=np.array(mask), padding="max_length", return_tensors="pt")
+    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 = (preds.numpy() * 255).astype(np.uint8)
+    
+    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)
 
 def update_image(image, segmentation):
@@ -43,8 +43,7 @@ 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")
-            box_input = gr.Box(label="Select Box")
+            input_image = gr.Image(label="Input Image", tool="select")
             with gr.Row():
                 everything_btn = gr.Button("Everything")
                 box_btn = gr.Button("Box")
@@ -59,7 +58,7 @@ with gr.Blocks() as demo:
     
     box_btn.click(
         fn=segment_box,
-        inputs=[input_image, box_input],
+        inputs=[input_image, input_image],
         outputs=[output_image]
     )