main.py

import ultralytics.engine.results
from fastapi import FastAPI, UploadFile
from PIL import Image
from yolo_model import YoloModel

app = FastAPI()

model = YoloModel("SHOU-ISD/fire-and-smoke", "yolov8n.pt")


@app.get("/")
async def info():
    return {
        "model": f"Fire and Smoke Detection",
        "version": "1.0",
    }


@app.post("/predict")
async def predict(image: UploadFile):
    if image.content_type not in ["image/jpeg", "image/png"]:
        return {"error": "Invalid file type"}

    im = Image.open(image.file)
    res = model.detect(im)
    return {
        "list": [
            detect_item(box, res_item.names)
            for res_item in res
            for box in res_item.boxes
        ]
    }


CXYWH = {
    "cx": float,
    "cy": float,
    "w": float,
    "h": float,
}

DetectItem = {
    "category": str,
    "bbox": CXYWH,
    "score": float,
}


def detect_item(box: ultralytics.engine.results.Boxes, mapping: dict[int, str]) -> DetectItem:
    cx, cy, w, h = box.xywhn.tolist()[0]
    return {
        "category": mapping[int(box.cls)],
        "bbox": {
            "cx": cx,
            "cy": cy,
            "w": w,
            "h": h,
        },
        "score": float(box.conf),
    }