Update app.py

app.py

@@ -1,8 +1,11 @@
+from flask import Flask, request, jsonify
 from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
-import gradio as gr
 from PIL import Image
 import torch
 import numpy as np
+import io
+import base64
+
 
 # Load CLIPSeg processor and model
 processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
@@ -41,6 +44,10 @@ def get_masks(prompts, img, threshold):
 
     return masks
 
+    @app.route('/')
+def home():
+    return 'Server started. Verify it is running by visiting http://0.0.0.0:7860/'
+
 # Function to extract image using positive and negative prompts
 def extract_image(pos_prompts, neg_prompts, img, threshold):
     positive_masks = get_masks(pos_prompts, img, 0.5)
@@ -56,39 +63,30 @@ def extract_image(pos_prompts, neg_prompts, img, threshold):
 
     return output_image, final_mask
 
-# Gradio UI
-iface_ui = gr.Interface(
-    fn=extract_image,
-    inputs=[
-        gr.Textbox(
-            label="Please describe what you want to identify (comma separated)",
-            key="pos_prompts",
-        ),
-        gr.Textbox(
-            label="Please describe what you want to ignore (comma separated)",
-            key="neg_prompts",
-        ),
-        gr.Image(type="pil", label="Input Image", key="img"),
-        gr.Slider(minimum=0, maximum=1, default=0.4, label="Threshold", key="threshold"),
-    ],
-    outputs=[
-        gr.Image(label="Result", key="output_image"),
-        gr.Image(label="Mask", key="output_mask"),
-    ],
-)
-
-# Launch Gradio UI
-# Launch Gradio UI
-iface_ui.launch(share=True)
-
-# Non-UI Version
-def run_non_ui(image_path, pos_prompts, neg_prompts, threshold):
-    img = Image.open(image_path)
-    output_image, output_mask = extract_image(pos_prompts, neg_prompts, img, threshold)
-
-    # Save or use the output_image and output_mask as needed
-    output_image.show()  # For demonstration purposes, opens the image with the default image viewer
-    output_mask.show()  # For demonstration purposes, opens the mask with the default image viewer
-
-# Example of using non-UI version
-# run_non_ui("path/to/your/image.jpg", "positive prompt", "negative prompt", 0.5)
+@app.route('/api', methods=['POST'])
+def process_request():
+    data = request.json
+
+    # Convert base64 image to PIL Image
+    base64_image = data.get('image')
+    image_data = base64.b64decode(base64_image.split(',')[1])
+    img = Image.open(io.BytesIO(image_data))
+
+    # Get other parameters
+    pos_prompts = data.get('positive_prompts', '')
+    neg_prompts = data.get('negative_prompts', '')
+    threshold = float(data.get('threshold', 0.4))
+
+    # Perform image segmentation
+    output_image, final_mask = extract_image(pos_prompts, neg_prompts, img, threshold)
+
+    # Convert result to base64 for response
+    buffered = io.BytesIO()
+    output_image.save(buffered, format="PNG")
+    result_image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+    return jsonify({'result_image_base64': result_image_base64})
+
+if __name__ == '__main__':
+    print("Server starting. Verify it is running by visiting http://0.0.0.0:7860/")
+    app.run(host='0.0.0.0', port=7860, debug=True)