from fastapi import FastAPI, File, UploadFile
from fastapi.responses import StreamingResponse
from retinaface import RetinaFace
import cv2
import numpy as np
from typing import Any
from rembg import remove
from PIL import Image
import io

app = FastAPI()

def convert_to_python_types(obj: Any) -> Any:
    if isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, dict):
        return {key: convert_to_python_types(value) for key, value in obj.items()}
    elif isinstance(obj, list):
        return [convert_to_python_types(item) for item in obj]
    else:
        return obj

@app.post("/detect_faces")
async def detect_faces(file: UploadFile = File(...)):
    contents = await file.read()
    nparr = np.frombuffer(contents, np.uint8)
    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    
    # Detect faces
    faces = RetinaFace.detect_faces(image)
    
    return convert_to_python_types(faces)

@app.post("/remove_background")
async def remove_background(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    
    image_rembg = remove(image)
    image_np = np.array(image_rembg)
    image = Image.fromarray(image_np)
    
    # Convert the image to bytes
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')
    img_byte_arr.seek(0)
    
    return StreamingResponse(io.BytesIO(img_byte_arr.read()), media_type="image/png")



@app.get("/")
def greet_json():
    return {"Hello": "World!"}