main.py

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import StreamingResponse
import io
import cv2
import numpy as np
from PIL import Image
from rembg import remove, new_session
import mediapipe as mp

app = FastAPI()

# Global variables
session = None
face_detection = None

@app.on_event("startup")
async def startup_event():
    global session, face_detection
    print("Initializing models...")
    # Use standard U2Net for stability
    session = new_session("u2net")
    # Use standard solutions API which is built-in to all mediapipe Linux wheels
    face_detection = mp.solutions.face_detection.FaceDetection(
        model_selection=1, 
        min_detection_confidence=0.5
    )
    print("Models initialized.")

@app.post("/generate")
async def generate_passport(file: UploadFile = File(...)):
    global face_detection, session
    
    # 1. Load image
    contents = await file.read()
    pil_img = Image.open(io.BytesIO(contents)).convert("RGBA")
    np_img = np.array(pil_img.convert("RGB"))
    h, w, _ = np_img.shape

    # 2. Detection (Standard Solutions API)
    results = face_detection.process(cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR))
    if not results.detections:
        raise HTTPException(status_code=400, detail="No face detected")
    
    # 3. Smart Crop
    bbox = results.detections[0].location_data.relative_bounding_box
    y_start = max(0, int((bbox.ymin - 0.2) * h))
    y_end = min(h, int((bbox.ymin + bbox.height + 0.5) * h))
    crop = pil_img.crop((0, y_start, w, y_end))
    
    # 4. Background Removal
    no_bg = remove(crop, session=session)
    
    # 5. Composite White BG
    final = Image.new("RGBA", no_bg.size, (255, 255, 255, 255))
    final.paste(no_bg, (0, 0), no_bg)
    final = final.resize((600, 600), Image.Resampling.LANCZOS)
    
    # 6. Export
    buf = io.BytesIO()
    final.convert("RGB").save(buf, format="JPEG", quality=95)
    buf.seek(0)
    return StreamingResponse(buf, media_type="image/jpeg")