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

2
.gitignore vendored Normal file
View file

@ -0,0 +1,2 @@
/__pycache__

5
app.py
View file

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