Spaces:
Sleeping
Sleeping
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) | |
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 {} | |