Update app.py

app.py

@@ -3,96 +3,18 @@ import gradio as gr
 from PIL import Image
 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__)
-
-
-# Define article as a global variable
-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>"
-
-def process_image(image, prompt):
-    inputs = processor(
-        text=prompt, images=image, padding="max_length", return_tensors="pt"
-    )
-
-    with torch.no_grad():
-        outputs = model(**inputs)
-        preds = outputs.logits
-
-    pred = torch.sigmoid(preds)
-    mat = pred.cpu().numpy()
-    mask = Image.fromarray(np.uint8(mat * 255), "L")
-    mask = mask.convert("RGB")
-    mask = mask.resize(image.size)
-    mask = np.array(mask)[:, :, 0]
-
-    mask_min = mask.min()
-    mask_max = mask.max()
-    mask = (mask - mask_min) / (mask_max - mask_min)
-    return mask
-
-def get_masks(prompts, img, threshold):
-    prompts = prompts.split(",")
-    masks = []
-    for prompt in prompts:
-        mask = process_image(img, prompt)
-        mask = mask > threshold
-        masks.append(mask)
-    return masks
-
-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)
-
-    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
-
-    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
-
-
-@app.route('/api', methods=['POST'])
-def api():
-    data = request.form
-    img_url = data['input_image']
-    positive_prompts = data['positive_prompts']
-    negative_prompts = data['negative_prompts']
-    threshold = float(data['input_slider_T'])
-
-    # Download image from URL
-    response = requests.get(img_url)
-    img = Image.open(BytesIO(response.content))
-
-    # Process image
-    masks = get_masks(positive_prompts, negative_prompts, img, threshold)
-    final_mask = np.any(np.stack(masks), axis=0)
-
-    # Convert mask to image
-    final_mask = Image.fromarray(final_mask.astype(np.uint8) * 255, "L")
-
-    # Convert the final image to bytes
-    img_bytes = BytesIO()
-    final_mask.save(img_bytes, format='PNG')
-    img_bytes.seek(0)
-
-    return send_file(img_bytes, mimetype='image/png')
-    
 # Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# CLIPSeg: Image Segmentation Using Text and Image Prompts")
-    gr.Markdown(article)
-    gr.Markdown(description)
+
+    # Add your article and description here
+    gr.Markdown("Your article goes here")
+    gr.Markdown("Your description goes here")
 
     with gr.Row():
         with gr.Column():
@@ -113,6 +35,49 @@ with gr.Blocks() as demo:
             output_image = gr.Image(label="Result")
             output_mask = gr.Image(label="Mask")
 
+    def process_image(image, prompt):
+        inputs = processor(
+            text=prompt, images=image, padding="max_length", return_tensors="pt"
+        )
+
+        with torch.no_grad():
+            outputs = model(**inputs)
+            preds = outputs.logits
+
+        pred = torch.sigmoid(preds)
+        mat = pred.cpu().numpy()
+        mask = Image.fromarray(np.uint8(mat * 255), "L")
+        mask = mask.convert("RGB")
+        mask = mask.resize(image.size)
+        mask = np.array(mask)[:, :, 0]
+
+        mask_min = mask.min()
+        mask_max = mask.max()
+        mask = (mask - mask_min) / (mask_max - mask_min)
+        return mask
+
+    def get_masks(prompts, img, threshold):
+        prompts = prompts.split(",")
+        masks = []
+        for prompt in prompts:
+            mask = process_image(img, prompt)
+            mask = mask > threshold
+            masks.append(mask)
+        return masks
+
+    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)
+
+        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
+
+        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
+
     btn_process.click(
         extract_image,
         inputs=[
@@ -124,19 +89,5 @@ 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=8080)
-
-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()
-
-    gr_thread.join()
-    flask_thread.join()
+# Launch Gradio API
+demo.launch(share=True)