Ultralytics Yolo toml and Api Implementation

[umar-anzar]

Readme.md would be added

• `TOML file for torch for cpu and cuda
• Script of one api in readme for inference

import cv2
import torch
import asyncio
import numpy as np
from typing import Optional
from ultralytics import YOLO
from pydantic import BaseModel
from fastapi import FastAPI, File, UploadFile
from concurrent.futures import ThreadPoolExecutor

# Define global loop and executor
LOOP = asyncio.get_event_loop()
MODEL_EXECUTOR = ThreadPoolExecutor(max_workers=1)

# Load a pretrained YOLO model
model = YOLO("yolov8n.pt")

class DetectedObject(BaseModel):
    name: str
    conf: Optional[float]
    bbox: Optional[tuple[int, int, int, int]]
    
async def yolo_detect(model: YOLO, image, null_class: str, conf: float = 0.8) -> list[DetectedObject]:
    
	# Run YOLO detection asynchronously
    results = await LOOP.run_in_executor(MODEL_EXECUTOR, lambda: model.predict(image, conf=conf, verbose=False))
    
	# Move results to CPU if CUDA is used
    if torch.cuda.is_available():
        results = [result.cpu() for result in results]
        torch.cuda.empty_cache()
        
    result_arr: list[DetectedObject] = []

    for result in results:
        boxes = result.boxes

        for cls_index, conf_item, coords_item in zip(boxes.cls, boxes.conf, boxes.xyxy):
            name = result.names[int(cls_index)]
            confidence = conf_item.item()
            bbox = tuple(map(int, coords_item))
			result_arr.append(DetectedObject(name=name, conf=confidence, bbox=bbox))
	
	# Return detected objects or a null-class response if no detections
    return result_arr if result_arr else [DetectedObject(name=null_class, conf=None, bbox=None)]
	
	
app = FastAPI()

class PredictResponse(BaseModel):
    detections: list[DetectedObject]

@app.post("/predict")
async def predict(file: UploadFile = File(...)) -> PredictResponse:
    # Read the uploaded image file
    try:
        image_data = await file.read()
        # Decode the image data to a numpy array
        image = cv2.imdecode(np.frombuffer(image_data, np.uint8), cv2.IMREAD_COLOR)
        if image is None:
            raise ValueError("Image could not be decoded.")
    except Exception as e:
        return {"error": "Invalid image file."}

    # Run detection
    detections = await yolo_detect(model, image, null_class="no_object_detected")
	
    # Format response with detected objects
    return PredictResponse(detections=detections)