apps/detection/controller.py

from .provider import model
from io import BytesIO
from PIL import Image


def detect_object(image: bytes):
    # Convert bytes to PIL image
    image = Image.open(BytesIO(image))

    detections = model.predict(image)

    # Format detections
    bboxs = detections[0].boxes.data
    labels = detections[0].names

    # Get boxs result
    boxs_result = []
    for box in bboxs:
        result_conf = box[4].item()
        if result_conf > model.conf:
            boxs_result.append({
                "class_name": labels[int(box[5].item())],
                "confidence": result_conf,
                "xmin": box[0].item(),
                "ymin": box[1].item(),
                "xmax": box[2].item(),
                "ymax": box[3].item()
            })

    print(boxs_result)

    return {
        "detections": boxs_result
    }