Update app.py

app.py

@@ -1,95 +1,63 @@
-# Import app from main.py for actual background removal
-try:
-    from main import app as main_app
-    import os
-    import secrets
-    import json
-    import uuid
-    from fastapi import FastAPI, Depends, HTTPException, Header, status, File, UploadFile
-    
-    # Create the FastAPI app
-    app = FastAPI(title="Background Removal API with API Key Support")
-    
-    # API key storage file
-    API_KEYS_FILE = "api_keys.json"
-    
-    # Initialize API keys from file or create default
-    def init_api_keys():
-        if os.path.exists(API_KEYS_FILE):
-            try:
-                with open(API_KEYS_FILE, "r") as f:
-                    return json.load(f)
-            except:
-                # If file is corrupted, create new
-                pass
-                
-        # Generate a default API key if none exists
-        default_key = secrets.token_urlsafe(16)
-        keys = {"default": {"key": default_key, "active": True}}
-        with open(API_KEYS_FILE, "w") as f:
-            json.dump(keys, f)
-        print(f"Created default API key: {default_key}")
-        return keys
-    
-    # API keys dictionary
-    API_KEYS = init_api_keys()
-    
-    # Function to verify API key
-    def verify_api_key(api_key: str = Header(None, alias="X-API-Key")):
-        if api_key is None:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail="API key is missing"
-            )
-            
-        for key_info in API_KEYS.values():
-            if key_info["key"] == api_key and key_info["active"]:
-                return True
-                
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Invalid API key"
-        )
+from flask import Flask, request, jsonify, send_file
+from rembg import remove
+from PIL import Image
+import io
+import os
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+app = Flask(__name__)
+
+# Get API key from environment variable
+API_KEY = os.getenv('API_KEY')
+
+def check_api_key():
+    api_key = request.headers.get('X-API-Key')
+    if not api_key or api_key != API_KEY:
+        return False
+    return True
+
+@app.route('/remove-bg', methods=['POST'])
+def remove_background():
+    # Check API key
+    if not check_api_key():
+        return jsonify({'error': 'Invalid or missing API key'}), 401
+
+    if 'image' not in request.files:
+        return jsonify({'error': 'No image file provided'}), 400
     
-    # Root endpoint
-    @app.get("/")
-    async def root():
-        return {
-            "message": "Background Removal API with API Key Support",
-            "docs": "/docs",
-            "endpoints": ["/remove-bg", "/api-key"]
-        }
+    file = request.files['image']
     
-    # Generate API key endpoint
-    @app.get("/api-key")
-    async def generate_api_key():
-        new_key = secrets.token_urlsafe(16)
-        key_id = str(uuid.uuid4())[:8]
-        API_KEYS[key_id] = {"key": new_key, "active": True}
+    if file.filename == '':
+        return jsonify({'error': 'No selected file'}), 400
+
+    try:
+        # Read the image
+        input_image = Image.open(file.stream)
         
-        # Save to file
-        with open(API_KEYS_FILE, "w") as f:
-            json.dump(API_KEYS, f)
-            
-        return {"api_key": new_key, "message": "Keep this key secure. You'll need it for API calls."}
-    
-    # Import the background removal function from main
-    from main import remove_background
-    
-    # Protect the remove-bg endpoint with API key
-    @app.post("/remove-bg")
-    async def protected_remove_background(api_key_valid: bool = Depends(verify_api_key), file: UploadFile = File(...)):
-        # Forward to the original function
-        return await remove_background(file=file)
-    
-    # Print the default API key for testing
-    if "default" in API_KEYS:
-        print(f"Default API key for testing: {API_KEYS['default']['key']}")
-        print(f"To use: Add 'X-API-Key: {API_KEYS['default']['key']}' to your request headers")
-    
-except ImportError:
-    # Fallback to the simple version if the main version fails
-    from fallback import app
+        # Remove background
+        output = remove(input_image)
+        
+        # Save to bytes
+        img_byte_arr = io.BytesIO()
+        output.save(img_byte_arr, format='PNG')
+        img_byte_arr.seek(0)
+        
+        return send_file(
+            img_byte_arr,
+            mimetype='image/png',
+            as_attachment=True,
+            download_name='removed_bg.png'
+        )
+
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    return jsonify({'status': 'healthy'}), 200
 
-# This ensures the "app" variable is properly exported for ASGI
-# This file ensures compatibility with both app:app and main:app entrypoints 
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=5000)