Update app.py

app.py

@@ -45,10 +45,6 @@ def get_masks(prompts, img, threshold):
 
     return masks
 
-@app.route('/')
-def hello_world():
-    return 'Hello, World!'
-
 # 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)
@@ -62,7 +58,16 @@ def extract_image(pos_prompts, neg_prompts, img, threshold):
     output_image = Image.new("RGBA", img.size, (0, 0, 0, 0))
     output_image.paste(img, mask=final_mask)
 
-    return output_image, final_mask
+    # Convert final mask to base64
+    buffered = io.BytesIO()
+    final_mask.save(buffered, format="PNG")
+    final_mask_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+    return final_mask_base64
+
+@app.route('/')
+def hello_world():
+    return 'Hello, World!'
 
 @app.route('/api', methods=['POST'])
 def process_request():
@@ -79,16 +84,10 @@ def process_request():
     threshold = float(data.get('threshold', 0.4))
 
     # Perform image segmentation
-    output_image, final_mask = extract_image(pos_prompts, neg_prompts, img, threshold)
+    final_mask_base64 = 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})
+    return jsonify({'final_mask_base64': final_mask_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)
-