import cv2
import torch
import numpy as np
from ultralytics import YOLO
from fastapi import FastAPI

app = FastAPI()
model = YOLO('model.pt')
#device = torch.device('cuda')
#model.to(device)
@app.post("/detect/")
def detect(img):
    try:
        objects_dict = {}
        img_arr = cv2.imread(img)
        img_arr = cv2.resize(img_arr, (416, 416))
        img_arr = np.array([img_arr.transpose(2, 0, 1)])
        #img_arr = torch.from_numpy(img_arr).float().to(device)
        results = model(img_arr)

        for r in results:
            n = len(r.boxes.cls)

            for i in range(n):
                cls = int(r.boxes.cls[i].cpu())
                temp_obj = [int(r.boxes.conf[i].cpu()), r.boxes.xyxy[i].cpu().numpy()] #уверенность модели, координаты прямоугольника

                if cls not in objects_dict:
                    objects_dict[cls] = [temp_obj]
                else:
                    objects_dict[cls].append(temp_obj)
        return objects_dict

    except Exception as e:
        print(str(e))
        return {}