Update app.py

app.py

@@ -1,70 +1,67 @@
 import gradio as gr
-import torch
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageDraw
+import torch
 from transformers import AutoProcessor, CLIPSegForImageSegmentation
-import traceback
 
 # Load the CLIPSeg model and processor
 processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
 model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
 
-def segment_image(input_image, text_prompt):
-    try:
-        # Ensure input_image is a PIL Image
-        if not isinstance(input_image, Image.Image):
-            input_image = Image.fromarray(input_image)
-
-        # Resize image if it's too large
-        max_size = 1024
-        if max(input_image.size) > max_size:
-            input_image.thumbnail((max_size, max_size))
-
-        # Preprocess the image
-        inputs = processor(text=[text_prompt], images=[input_image], padding="max_length", return_tensors="pt")
-
-        # Perform segmentation
-        with torch.no_grad():
-            outputs = model(**inputs)
-
-        # Get the predicted segmentation
-        preds = outputs.logits.squeeze().sigmoid()
-
-        # Convert the prediction to a numpy array and scale to 0-255
-        segmentation = (preds.numpy() * 255).astype(np.uint8)
-
-        # Resize segmentation to match input image size
-        segmentation = Image.fromarray(segmentation).resize(input_image.size)
-        segmentation = np.array(segmentation)
-
-        # Create a colored heatmap
-        heatmap = np.zeros((segmentation.shape[0], segmentation.shape[1], 3), dtype=np.uint8)
-        heatmap[:, :, 0] = segmentation  # Red channel
-        heatmap[:, :, 2] = 255 - segmentation  # Blue channel
+def segment_everything(image):
+    inputs = processor(text=["object"], images=[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)
+    return Image.fromarray(segmentation)
 
-        # Blend the heatmap with the original image
-        original_image = np.array(input_image)
-        blended = (0.7 * original_image + 0.3 * heatmap).astype(np.uint8)
+def segment_box(image, box):
+    x1, y1, x2, y2 = map(int, box)
+    mask = Image.new('L', image.size, 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")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    preds = outputs.logits.squeeze().sigmoid()
+    segmentation = (preds.numpy() * 255).astype(np.uint8)
+    return Image.fromarray(segmentation)
 
-        return Image.fromarray(blended)
-    except Exception as e:
-        error_msg = f"An error occurred: {str(e)}\n\nStacktrace:\n{traceback.format_exc()}"
-        return Image.new('RGB', (400, 200), color = (255, 0, 0))  # Red image to indicate error
+def update_image(image, segmentation, tool):
+    if segmentation is None:
+        return image
+    blended = Image.blend(image.convert('RGBA'), segmentation.convert('RGBA'), 0.5)
+    return blended
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=segment_image,
-    inputs=[
-        gr.Image(type="pil", label="Input Image"),
-        gr.Textbox(label="Text Prompt", placeholder="Enter a description of what to segment...")
-    ],
-    outputs=[
-        gr.Image(type="pil", label="Segmentation Result"),
-        gr.Textbox(label="Error Message", visible=False)
-    ],
-    title="CLIPSeg Image Segmentation",
-    description="Upload an image and provide a text prompt to segment objects."
-)
+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")
+            with gr.Row():
+                everything_btn = gr.Button("Everything")
+                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, input_image.sel],
+        outputs=[output_image]
+    )
+    
+    output_image.change(
+        fn=update_image,
+        inputs=[input_image, output_image, gr.State("last_tool")],
+        outputs=[output_image]
+    )
 
-# Launch the interface
-iface.launch()
+demo.launch()