Update app.py

app.py

@@ -1,23 +1,69 @@
-import gradio as gr
-from rembg import remove
-from PIL import Image
-import io
-
-def remove_background(input_image):
-    if isinstance(input_image, str):
-        input_image = Image.open(input_image)
-    output_image = remove(input_image)
-    return output_image
-
-demo = gr.Interface(
-    fn=remove_background,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Background Remover",
-    description="Upload an image to remove its background",
-    examples=[],
-    theme=gr.themes.Soft()
-)
-
-if __name__ == "__main__":
-    demo.launch()
+from flask import Flask, request, send_file, jsonify
+from gradio_client import Client, handle_file
+import tempfile
+import os
+from werkzeug.utils import secure_filename
+
+app = Flask(__name__)
+
+# Allowed file extensions
+ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
+MAX_FILE_SIZE = 5 * 1024 * 1024  # 5MB
+
+def allowed_file(filename):
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+
+@app.route('/remove-background', methods=['POST'])
+def remove_background():
+    if 'file' not in request.files:
+        return jsonify({"error": "No file provided"}), 400
+
+    file = request.files['file']
+
+    if file.filename == '':
+        return jsonify({"error": "No selected file"}), 400
+
+    if not allowed_file(file.filename):
+        return jsonify({"error": "Invalid file type"}), 400
+
+    # Check file size
+    file.seek(0, os.SEEK_END)
+    file_length = file.tell()
+    file.seek(0)
+
+    if file_length > MAX_FILE_SIZE:
+        return jsonify({"error": "File too large (max 5MB)"}), 400
+
+    try:
+        # Save the uploaded file temporarily
+        temp_input = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+        file.save(temp_input.name)
+        temp_input.close()
+
+        # Call Hugging Face API
+        client = Client("vinaynayak/RemBG_API")
+        result_path = client.predict(
+            input_image=handle_file(temp_input.name),
+            api_name="/predict"
+        )
+
+        # Send the result back
+        return send_file(
+            result_path,
+            mimetype='image/png',
+            as_attachment=True,
+            download_name='background_removed.png'
+        )
+
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+    finally:
+        # Clean up temporary files
+        if 'temp_input' in locals() and os.path.exists(temp_input.name):
+            os.unlink(temp_input.name)
+        if 'result_path' in locals() and os.path.exists(result_path):
+            os.unlink(result_path)
+
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5000, debug=True)