Update app.py

app.py

@@ -1,20 +1,21 @@
-from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
 import gradio as gr
 from PIL import Image
 import torch
 import matplotlib.pyplot as plt
 import cv2
-import torch
 import numpy as np
+import spaces
+from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
 
 processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
-model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
+model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined").cuda()
 
-
-def process_image(image, prompt, threhsold, alpha_value, draw_rectangles):
+@spaces.GPU
+def process_image(image, prompt, threshold, alpha_value, draw_rectangles):
     inputs = processor(
-        text=prompt, images=image, padding="max_length", return_tensors="pt"
+        text=prompt, images=image, return_tensors="pt"
     )
+    inputs = {k: v.cuda() for k, v in inputs.items()}
 
     # predict
     with torch.no_grad():
@@ -22,7 +23,7 @@ def process_image(image, prompt, threhsold, alpha_value, draw_rectangles):
         preds = outputs.logits
 
     pred = torch.sigmoid(preds)
-    mat = pred.cpu().numpy()
+    mat = pred.squeeze().cpu().numpy()
     mask = Image.fromarray(np.uint8(mat * 255), "L")
     mask = mask.convert("RGB")
     mask = mask.resize(image.size)
@@ -34,9 +35,9 @@ def process_image(image, prompt, threhsold, alpha_value, draw_rectangles):
     mask = (mask - mask_min) / (mask_max - mask_min)
 
     # threshold the mask
-    bmask = mask > threhsold
+    bmask = mask > threshold
     # zero out values below the threshold
-    mask[mask < threhsold] = 0
+    mask[mask < threshold] = 0
 
     fig, ax = plt.subplots()
     ax.imshow(image)
@@ -62,12 +63,10 @@ def process_image(image, prompt, threhsold, alpha_value, draw_rectangles):
 
     return fig, mask, output_image
 
-
 title = "Interactive demo: zero-shot image segmentation with CLIPSeg"
 description = "Demo for using CLIPSeg, a CLIP-based model for zero- and one-shot image segmentation. To use it, simply upload an image and add a text to mask (identify in the image), or use one of the examples below and click 'submit'. Results will show up in a few seconds."
 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.10003'>CLIPSeg: Image Segmentation Using Text and Image Prompts</a> | <a href='https://huggingface.co/docs/transformers/main/en/model_doc/clipseg'>HuggingFace docs</a></p>"
 
-
 with gr.Blocks() as demo:
     gr.Markdown("# CLIPSeg: Image Segmentation Using Text and Image Prompts")
     gr.Markdown(article)
@@ -119,4 +118,4 @@ with gr.Blocks() as demo:
         ],
     )
 
-demo.launch()
+demo.launch(share=True)