Update app.py

app.py

@@ -1,69 +1,22 @@
-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)
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import StreamingResponse
+from rembg import remove
+from PIL import Image
+import io
+
+app = FastAPI()
+
+@app.post("/remove-bg")
+async def remove_bg(file: UploadFile = File(...)):
+    input_bytes = await file.read()
+    input_image = Image.open(io.BytesIO(input_bytes)).convert("RGBA")
+    
+    # Remove background
+    output_image = remove(input_image)
+
+    # Convert output to byte stream
+    img_byte_arr = io.BytesIO()
+    output_image.save(img_byte_arr, format='PNG')
+    img_byte_arr.seek(0)
+
+    return StreamingResponse(img_byte_arr, media_type="image/png")