Create app.py

app.py

@@ -0,0 +1,141 @@
+from fastapi import FastAPI, File, UploadFile, HTTPException, Header
+from fastapi.responses import Response
+from fastapi.middleware.cors import CORSMiddleware
+import onnxruntime as ort
+import numpy as np
+from PIL import Image
+import cv2
+import io
+from datetime import datetime, timedelta
+from collections import defaultdict
+import os
+
+app = FastAPI(title="Backdrop Studio API", version="2.0.0")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Restrict for production!
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+MODEL_PATH = "models/modnet.onnx"
+MODEL_WIDTH = 512   # Official MODNet ONNX input size (for best speed, use 512; 768 or 1024 for higher res if model supports)
+MODEL_HEIGHT = 512
+
+print("🔄 Loading MODNet ONNX model...")
+onnx_session = ort.InferenceSession(MODEL_PATH, providers=['CPUExecutionProvider'])
+print("✅ MODNet model loaded successfully!")
+
+user_quotas = defaultdict(lambda: {"count": 0, "date": datetime.now().date()})
+MAX_DAILY_IMAGES = 5
+
+def check_and_update_quota(user_id: str) -> bool:
+    today = datetime.now().date()
+    user_data = user_quotas[user_id]
+    if user_data["date"] != today:
+        user_data["count"] = 0
+        user_data["date"] = today
+    if user_data["count"] >= MAX_DAILY_IMAGES:
+        return False
+    user_data["count"] += 1
+    return True
+
+def preprocess_image(image: Image.Image, target_size=(MODEL_WIDTH, MODEL_HEIGHT)):
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    orig_width, orig_height = image.size
+    image_resized = image.resize(target_size, Image.LANCZOS)
+    img_array = np.array(image_resized).astype(np.float32) / 255.0  # shape (512, 512, 3)
+    img_array = np.transpose(img_array, (2, 0, 1))  # (3, 512, 512)
+    img_array = np.expand_dims(img_array, axis=0)   # (1, 3, 512, 512)
+    return img_array, (orig_width, orig_height)
+
+def postprocess_mask(mask: np.ndarray, original_size):
+    # MODNet returns (1,1,H,W) float in [0,1]
+    mask = mask[0, 0]  # (H,W)
+    mask = (mask * 255).round().astype(np.uint8)
+    mask = cv2.resize(mask, original_size, interpolation=cv2.INTER_LINEAR)
+    # Optional: Apply threshold to get crisp mask
+    mask = np.where(mask > 127, 255, 0).astype(np.uint8)
+    return mask
+
+def remove_background(image: Image.Image):
+    input_array, original_size = preprocess_image(image)
+    input_name = onnx_session.get_inputs()[0].name
+    output = onnx_session.run(None, {input_name: input_array})
+    mask = postprocess_mask(output[0], original_size)
+    image_array = np.array(image.convert('RGBA'))
+    image_array[:, :, 3] = mask
+    result_image = Image.fromarray(image_array, 'RGBA')
+    return result_image
+
+@app.get("/")
+async def root():
+    return {"status": "healthy", "service": "Backdrop Studio MODNet API", "version": "2.0.0"}
+
+@app.get("/quota/{user_id}")
+async def get_quota(user_id: str):
+    today = datetime.now().date()
+    user_data = user_quotas[user_id]
+    if user_data["date"] != today:
+        user_data["count"] = 0
+        user_data["date"] = today
+    remaining = MAX_DAILY_IMAGES - user_data["count"]
+    return {
+        "user_id": user_id,
+        "used": user_data["count"],
+        "remaining": max(0, remaining),
+        "limit": MAX_DAILY_IMAGES,
+        "resets_at": str(today + timedelta(days=1))
+    }
+
+@app.post("/remove-background")
+async def remove_background_endpoint(
+    file: UploadFile = File(...),
+    user_id: str = Header(..., alias="X-User-ID")
+):
+    if not user_id or len(user_id) < 10:
+        raise HTTPException(
+            status_code=400,
+            detail="Invalid user ID. Please provide a valid device identifier."
+        )
+    if not check_and_update_quota(user_id):
+        raise HTTPException(
+            status_code=429,
+            detail=f"Daily quota exceeded. You can process {MAX_DAILY_IMAGES} images per day. Try again tomorrow!"
+        )
+    if not file.content_type.startswith('image/'):
+        raise HTTPException(
+            status_code=400,
+            detail="Invalid file type. Please upload an image (JPEG or PNG)."
+        )
+    try:
+        image_bytes = await file.read()
+        image = Image.open(io.BytesIO(image_bytes))
+        result_image = remove_background(image)
+        output_buffer = io.BytesIO()
+        result_image.save(output_buffer, format='PNG')
+        output_buffer.seek(0)
+        return Response(
+            content=output_buffer.getvalue(),
+            media_type="image/png",
+            headers={
+                "X-Quota-Used": str(user_quotas[user_id]["count"]),
+                "X-Quota-Remaining": str(MAX_DAILY_IMAGES - user_quotas[user_id]["count"])
+            }
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error processing image: {str(e)}"
+        )
+    finally:
+        if 'image_bytes' in locals(): del image_bytes
+        if 'image' in locals(): del image
+
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.environ.get("PORT", 8080))
+    uvicorn.run(app, host="0.0.0.0", port=port)