Update app.py

app.py

@@ -2,20 +2,21 @@ from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
 import gradio as gr
 from PIL import Image
 import torch
-import matplotlib.pyplot as plt
-import torch
 import numpy as np
+from flask import Flask, request, jsonify, send_file
+from io import BytesIO
+import threading
 
 processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
 model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
 
+app = Flask(__name__)
 
 def process_image(image, prompt):
     inputs = processor(
         text=prompt, images=image, padding="max_length", return_tensors="pt"
     )
 
-    # predict
     with torch.no_grad():
         outputs = model(**inputs)
         preds = outputs.logits
@@ -27,44 +28,34 @@ def process_image(image, prompt):
     mask = mask.resize(image.size)
     mask = np.array(mask)[:, :, 0]
 
-    # normalize the mask
     mask_min = mask.min()
     mask_max = mask.max()
     mask = (mask - mask_min) / (mask_max - mask_min)
     return mask
 
-
-def get_masks(prompts, img, threhsold):
+def get_masks(prompts, img, threshold):
     prompts = prompts.split(",")
     masks = []
     for prompt in prompts:
         mask = process_image(img, prompt)
-        mask = mask > threhsold
+        mask = mask > threshold
         masks.append(mask)
     return masks
 
-
-def extract_image(pos_prompts, neg_prompts, img, threhsold):
+def extract_image(pos_prompts, neg_prompts, img, threshold):
     positive_masks = get_masks(pos_prompts, img, 0.5)
     negative_masks = get_masks(neg_prompts, img, 0.5)
 
-    # combine masks into one masks, logic OR
     pos_mask = np.any(np.stack(positive_masks), axis=0)
     neg_mask = np.any(np.stack(negative_masks), axis=0)
     final_mask = pos_mask & ~neg_mask
 
-    # extract the final image
     final_mask = Image.fromarray(final_mask.astype(np.uint8) * 255, "L")
     output_image = Image.new("RGBA", img.size, (0, 0, 0, 0))
     output_image.paste(img, mask=final_mask)
     return output_image, final_mask
 
-
-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>"
-
-
+# Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# CLIPSeg: Image Segmentation Using Text and Image Prompts")
     gr.Markdown(article)
@@ -99,6 +90,19 @@ with gr.Blocks() as demo:
         ],
         outputs=[output_image, output_mask],
     )
+def run_demo():
+    demo.launch()
+
+def run_flask():
+    app.run(host='127.0.0.1', port=7860)
+
+if __name__ == '__main__':
+    # Run Gradio UI and Flask in separate threads
+    gr_thread = threading.Thread(target=run_demo)
+    flask_thread = threading.Thread(target=run_flask)
 
+    gr_thread.start()
+    flask_thread.start()
 
-demo.launch()
+    gr_thread.join()
+    flask_thread.join()