Spaces:
Sleeping
Sleeping
File size: 868 Bytes
7a1529d 7f4a7eb 7a1529d |
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 |
from ultralytics import YOLO
from PIL import Image
import os
curr = os.getcwd()
detect_model_path = os.path.join(curr, 'models', 'chessDetection3d.pt')
detect_model = YOLO(detect_model_path)
async def detect_pieces(image : Image):
if image is None:
print("No image is there")
return {"error" : "No image detected"}
results = detect_model.predict(image)
if not results or len(results) == 0:
print("No results are there")
return {"error" : "No results found"}
boxes = results[0].boxes.xyxy.tolist()
confidences = results[0].boxes.conf.tolist()
classes = results[0].boxes.cls.tolist()
class_names = []
for idx in classes:
class_names.append(detect_model.names[idx])
return {
"boxes": boxes,
"confidences": confidences,
"classes": class_names
}
|