app.py

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:
        os.makedirs("models", exist_ok=True)
        path = hf_hub_download(
            repo_id="biglam/medieval-manuscript-yolov11",
            filename=MODEL_OPTIONS[model_id]
        )
        local_path = os.path.join("models", path)
        # Initialize and return model
        model = YOLO(path)
        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
                 ):
    # Initialize model for each request
    model = init_model(model_id)
    
    # Open image from uploaded file
    image = Image.open(image.file)
    
    # Run inference with the PIL Image
    results = model.predict(source=image, conf=conf, iou=iou)
    
    # 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)