import uvicorn
from fastapi.staticfiles import StaticFiles
from enum import Enum
from fastapi import FastAPI, UploadFile, File
from PIL import Image
import io
from ultralytics import YOLO
import os
import uuid
from fastapi.responses import Response

app = FastAPI(docs_url='/')
use_gpu = False
output_dir = 'output'
model = YOLO("model.pt", task="detect")

class OutputEnum(str, Enum):
    json = "json"
    image = "image"

@app.post("/detect")
async def detect(
    file: UploadFile = File(...),
    output: OutputEnum = OutputEnum.json
):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents))
    results = model.predict(source=image)
    
    if output == OutputEnum.image:
        filename = f"{uuid.uuid4().hex}.jpg"
        filepath = os.path.join(output_dir, filename)
        results[0].save(filename=filepath)
        with open(filepath, "rb") as f:
            image_bytes = f.read()
        os.remove(filepath)
        return Response(content=image_bytes, media_type="image/jpeg")
    else:    
        detections = [{
            'class': int(box.cls),
            'confidence': float(box.conf),
            'box': [float(x) for x in box.xyxy[0]]
        } for box in results[0].boxes]

        return {'detections': detections}

app.mount("/output", StaticFiles(directory="output", follow_symlink=True, html=True), name="output")

if __name__ == '__main__':
    uvicorn.run(app=app)