Spaces:
Sleeping
Sleeping
File size: 1,123 Bytes
4aee5ae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
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 {}
|