change return type to model names

This commit is contained in:
Siwat Sirichai 2025-01-09 22:45:49 +07:00
parent 3224295610
commit 93f8974ba0
2 changed files with 6 additions and 1 deletions

5
app.py
View file

@ -13,6 +13,9 @@ model = YOLO("yolov8n.pt")
if torch.cuda.is_available():
model.to('cuda')
# Retrieve class names from the model
class_names = model.names
with open("config.json", "r") as f:
config = json.load(f)
@ -58,7 +61,7 @@ async def detect(websocket: WebSocket):
for r in results:
for box in r.boxes:
boxes.append({
"class": int(box.cls[0]),
"class": class_names[int(box.cls[0])], # Use class name
"confidence": float(box.conf[0]),
})
# Broadcast to all subscribers of this URL