Update app.py

app.py

@@ -5,6 +5,7 @@ import torch
 
 import matplotlib.pyplot as plt
 import numpy as np
+import io
 
 import gradio as gr
 
@@ -14,15 +15,20 @@ processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
 model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
 
 def visualize_segmentation(image, prompts, preds):
-  _, ax = plt.subplots(1, len(prompts) + 1, figsize=(3*(len(prompts) + 1), 4))
-  [a.axis('off') for a in ax.flatten()]
-  ax[0].imshow(image)
-  [ax[i+1].imshow(torch.sigmoid(preds[i][0])) for i in range(len(prompts))];
-  [ax[i+1].text(0, -15, prompt) for i, prompt in enumerate(prompts)];
+    fig, ax = plt.subplots(1, len(prompts) + 1, figsize=(3*(len(prompts) + 1), 4))
+    [a.axis('off') for a in ax.flatten()]
+    ax[0].imshow(image)
+    [ax[i+1].imshow(torch.sigmoid(preds[i][0])) for i in range(len(prompts))];
+    [ax[i+1].text(0, -15, prompt) for i, prompt in enumerate(prompts)];
+
+    buf = io.BytesIO()
+    fig.savefig(buf, format='png')
+    plt.close(fig)
+    
+    return Image.open(buf)
      
 
 def segment(img, clases):
-    print(img)
     image = Image.fromarray(img, 'RGB')
     prompts = clases.split(',')
 
@@ -31,7 +37,7 @@ def segment(img, clases):
     outputs = model(**inputs)
     preds = outputs.logits.unsqueeze(1)
 
-    return "Hello " + prompts[0] + "!!"
+    return visualize_segmentation(image, prompts, preds)
 
-demo = gr.Interface(fn=segment, inputs=["image","text"], outputs="text")
+demo = gr.Interface(fn=segment, inputs=["image","text"], outputs="image")
 demo.launch()