Update app.py

app.py

@@ -14,20 +14,22 @@ def process_image(image, prompt):
   # predict
   with torch.no_grad():
     outputs = model(**inputs)
-    preds = outputs.logits.unsqueeze(1)
+    preds = outputs.logits
   
   filename = f"mask.png"
-  plt.imsave(filename,torch.sigmoid(preds[0][0]))
+  plt.imsave(filename, torch.sigmoid(preds))
   
-  img2 = cv2.imread(filename)
-  gray_image = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
+  # # img2 = cv2.imread(filename)
+  # # gray_image = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
 
-  (thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)
+  # # (thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)
   
-  # fix color format
-  cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)
+  # # # fix color format
+  # # cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)
   
-  return Image.fromarray(bw_image)
+  # # return Image.fromarray(bw_image)
+
+  # return Image.open("mask.png").convert("RGB")
   
 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."
@@ -36,7 +38,7 @@ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2112.100
 examples = [["example_image.png", "wood"]]
    
 interface = gr.Interface(fn=process_image, 
-                     inputs=[gr.Image(type="pil"), gr.Textbox(label="What do you want to identify (separated by comma)?")],
+                     inputs=[gr.Image(type="pil"), gr.Textbox(label="Please describe what you want to identify")],
                      outputs=gr.Image(type="pil"),
                      title=title,
                      description=description,