from fastapi import FastAPI, UploadFile
from ultralytics import YOLO
from PIL import Image
import os
from huggingface_hub import hf_hub_download

def init_model(model_id: str):
    # Define models
    MODEL_OPTIONS = {
        "YOLOv11-Nano": "medieval-yolov11n.pt",
        "YOLOv11-Small": "medieval-yolov11s.pt",
        "YOLOv11-Medium": "medieval-yolov11m.pt",
        "YOLOv11-Large": "medieval-yolov11l.pt",
        "YOLOv11-XLarge": "medieval-yolov11x.pt"
    }
    
    if model_id in MODEL_OPTIONS:
        print(MODEL_OPTIONS[model_id])
        path = hf_hub_download(
            repo_id="biglam/medieval-manuscript-yolov11",
            filename=MODEL_OPTIONS[model_id],
        )
        print(path)
        # Initialize and return model
        model = YOLO(path)
        print("Model initialized")
        return model
    else:
        raise ValueError(f"Model {model_id} not found")

app = FastAPI()

@app.get("/")
async def root():
    return {"status": "ok"}

@app.post("/predict")
async def predict(image: UploadFile,
                 model_id: str = "YOLOv11-XLarge",
                 conf: float = 0.25,
                 iou: float = 0.7
                 ):
    print(model_id, conf, iou)
    # Initialize model for each request
    model = init_model(model_id)
    
    # Open image from uploaded file
    image = Image.open(image.file)
    print("Image opened")
    # Run inference with the PIL Image
    results = model.predict(source=image, conf=conf, iou=iou)
    print("Inference done")
    # Extract detection results
    result = results[0]
    detections = []
    
    for box in result.boxes:
        detection = {
            "class": result.names[int(box.cls[0])],
            "confidence": float(box.conf[0]),
            "bbox": box.xyxy[0].tolist()
        }
        detections.append(detection)
    
    return {"detections": detections}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)