Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| from fastapi import FastAPI, UploadFile, File, HTTPException, Security, Depends, Request | |
| from fastapi.security.api_key import APIKeyHeader | |
| from starlette.status import HTTP_403_FORBIDDEN, HTTP_429_TOO_MANY_REQUESTS | |
| from fastapi.responses import StreamingResponse | |
| from io import BytesIO | |
| from PIL import Image | |
| # Import the background removal function | |
| from rembg import remove | |
| # --- Rate Limiting --- | |
| from slowapi import Limiter, _rate_limit_exceeded_handler | |
| from slowapi.util import get_remote_address | |
| from slowapi.errors import RateLimitExceeded | |
| # Initialize the rate limiter | |
| # This allows 5 requests per minute per IP address. Adjust as needed. | |
| limiter = Limiter(key_func=get_remote_address, default_limits=["5/minute"]) | |
| # --- API Initialization --- | |
| app = FastAPI( | |
| title="Background Remover API", | |
| description="A simple API to remove the background from images, hosted on Hugging Face Spaces.", | |
| version="1.0.0" | |
| ) | |
| # Add the rate limiter to the app | |
| app.state.limiter = limiter | |
| app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler) | |
| # --- Security: API Key Authentication --- | |
| # Define the API Key header | |
| api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False) | |
| # Load API keys from Hugging Face Secrets (as a JSON string) | |
| # For local testing, you can create a .env file and use python-dotenv | |
| # Example JSON string: '{"your-internal-key-for-website": "internal", "a-paid-user-key": "paid"}' | |
| API_KEYS_JSON = os.getenv("API_KEYS_JSON") | |
| if not API_KEYS_JSON: | |
| # Fallback for local development if the secret is not set | |
| print("WARNING: API_KEYS_JSON secret not found. Using dummy keys for local development.") | |
| API_KEYS = { | |
| "my-secret-website-key": "internal", | |
| "paid-user-123-key": "paid" | |
| } | |
| else: | |
| API_KEYS = json.loads(API_KEYS_JSON) | |
| async def get_api_key(api_key: str = Security(api_key_header)): | |
| """ | |
| Dependency function to validate the API key from the X-API-Key header. | |
| """ | |
| if not api_key or api_key not in API_KEYS: | |
| raise HTTPException( | |
| status_code=HTTP_403_FORBIDDEN, detail="Could not validate credentials" | |
| ) | |
| return {"key": api_key, "type": API_KEYS[api_key]} | |
| # --- API Endpoints --- | |
| def read_root(): | |
| """ A simple health check endpoint. """ | |
| return {"status": "ok", "message": "Welcome to the Background Remover API!"} | |
| # Apply the rate limit to this specific endpoint | |
| async def remove_background( | |
| request: Request, # Required for the rate limiter | |
| file: UploadFile = File(..., description="Image file to process (JPG, PNG)"), | |
| api_key: dict = Depends(get_api_key) # This enforces API key authentication | |
| ): | |
| """ | |
| This endpoint receives an image, validates it, removes the background, | |
| and returns the resulting image with a transparent background. | |
| - **Authentication**: Requires a valid `X-API-Key` in the header. | |
| - **Rate Limiting**: Limited to 5 requests per minute per IP. | |
| - **File Validation**: Accepts only `image/jpeg` and `image/png` up to 10MB. | |
| """ | |
| # --- Input Validation --- | |
| # 1. Check file type | |
| if file.content_type not in ["image/jpeg", "image/png"]: | |
| raise HTTPException(status_code=400, detail="Invalid file type. Please upload a JPG or PNG.") | |
| # 2. Check file size (e.g., 10 MB limit) | |
| # The `await file.read()` will load the whole file into memory. | |
| # We read it once and then use the bytes. | |
| contents = await file.read() | |
| if len(contents) > 10 * 1024 * 1024: # 10 MB | |
| raise HTTPException(status_code=400, detail="File too large. Maximum size is 10 MB.") | |
| # Log the request (you can expand this for analytics) | |
| print(f"Processing request for API key type: {api_key['type']}") | |
| # --- Core Logic: Background Removal --- | |
| try: | |
| # Use rembg to remove the background | |
| output_bytes = remove(contents) | |
| # Return the result as a streaming response (efficient for binary data) | |
| return StreamingResponse(BytesIO(output_bytes), media_type="image/png") | |
| except Exception as e: | |
| # If anything goes wrong during processing | |
| print(f"Error processing image: {e}") | |
| raise HTTPException(status_code=500, detail="Internal server error while processing the image.") |