Update app.py

app.py

@@ -1,99 +1,39 @@
 import gradio as gr
-import numpy as np
-from PIL import Image
 import torch
-from transformers import AutoProcessor, CLIPSegForImageSegmentation
-
-# Load the CLIPSeg model and processor
-processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
-model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
-
-# Ensure that the model uses GPU if available
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model.to(device)
-
-def segment_everything(image):
-    if isinstance(image, list):
-        image = image[0]
-    
-    if isinstance(image, np.ndarray):
-        image = Image.fromarray(image)
-    
-    inputs = processor(text=["object"], images=image, padding="max_length", return_tensors="pt").to(device)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    preds = outputs.logits.squeeze().sigmoid().cpu()
-    segmentation = (preds.numpy() * 255).astype(np.uint8)
-    return Image.fromarray(segmentation)
-
-def segment_box(image, x1, y1, x2, y2):
-    if isinstance(image, list):
-        image = image[0]
-    
-    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=Image.fromarray(cropped_image), padding="max_length", return_tensors="pt").to(device)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    preds = outputs.logits.squeeze().sigmoid().cpu()
-    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):
-    if segmentation is None:
-        return image
-    
-    if isinstance(image, list):
-        image = image[0]
-    
-    if isinstance(image, np.ndarray):
-        image_pil = Image.fromarray(image)
-    else:
-        image_pil = image
-    
-    seg_pil = Image.fromarray(segmentation).convert('RGBA')
-    
-    if image_pil.size!= seg_pil.size:
-        seg_pil = seg_pil.resize(image_pil.size, Image.NEAREST)
-    
-    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")
-            with gr.Row():
-                x1_input = gr.Number(label="X1")
-                y1_input = gr.Number(label="Y1")
-                x2_input = gr.Number(label="X2")
-                y2_input = gr.Number(label="Y2")
-            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, x1_input, y1_input, x2_input, y2_input],
-        outputs=[output_image]
-    )
-    output_image.change(
-        fn=update_image,
-        inputs=[input_image, output_image],
-        outputs=[output_image]
-    )
-
+from PIL import Image
+from torchvision import transforms
+
+# Load pre-trained U-Net model
+model = torch.hub.load('nvidia/DeepLearningExamples:torchhub', 'unet', pretrained=True)
+
+# Define a function to segment an image
+def segment_image(image):
+    # Preprocess image
+    image = Image.fromarray(image)
+    image = transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])(image)
+    
+    # Run segmentation model
+    output = model(image.unsqueeze(0))
+    output = torch.argmax(output, dim=1)
+    
+    # Postprocess output
+    output = output.squeeze(0).cpu().numpy()
+    output = Image.fromarray(output.astype('uint8'))
+    
+    return output
+
+# Create Gradio app
+demo = gr.Interface(
+    fn=segment_image,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Segment Anything",
+    description="Segment any image using a pre-trained U-Net model"
+)
+
+# Launch Gradio app
 demo.launch()