from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from PIL import Image
from ultralytics import YOLO
import io
import base64
import torch

app = FastAPI()

# Carica il modello all'avvio (singleton per efficienza)
model = YOLO("YoloV8-FaceDetection.pt")  # Sostituisci con "yolov8n-face.pt" se disponibile

# Modello Pydantic per validare l'input Base64
class ImageRequest(BaseModel):
    image_base64: str

def preprocess_image(image: Image.Image, size=(320, 320)):
    """Ridimensiona l'immagine per velocizzare l'inferenza"""
    img = image.convert("RGB")
    img = img.resize(size, Image.LANCZOS)  # 320x320 è veloce su CPU
    return img

@app.post("/detect_single_face")
async def detect_single_face(request: ImageRequest):
    try:
        # Decodifica la stringa Base64
        image_data = base64.b64decode(request.image_base64)
        image = Image.open(io.BytesIO(image_data))
        image = preprocess_image(image)

        # Esegui l'inferenza con YOLOv8
        results = model(image)

        # Conta i volti (classe "person" per yolov8n, o "face" se usi modello specifico)
        boxes = results[0].boxes
        face_count = 0
        for box in boxes:
            # Se usi yolov8n, filtra per classe "person" (class_id 0 in COCO)
            # Se usi yolov8n-face, conta direttamente tutte le detection
            if int(box.cls) == 0:  # "person" in yolov8n
                face_count += 1

        # Logica di risposta
        if face_count == 1:
            return {"has_single_face": True}
        elif face_count > 1:
            return {"has_single_face": False, "face_count": face_count}
        else:
            return {"has_single_face": False}

    except Exception as e:
        # In caso di errore (es. Base64 non valido), restituisci False
        return {"has_single_face": False, "error": str(e), "debug": request.image_base64}

# Avvio del server (Hugging Face lo gestisce automaticamente)
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)