From 93f8974ba0265c1fd18866f39a103cb2c75f3327 Mon Sep 17 00:00:00 2001 From: Siwat Sirichai Date: Thu, 9 Jan 2025 22:45:49 +0700 Subject: [PATCH] change return type to model names --- .gitignore | 2 ++ app.py | 5 ++++- 2 files changed, 6 insertions(+), 1 deletion(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0ce25c6 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ + +/__pycache__ diff --git a/app.py b/app.py index f8894fa..51f4e75 100644 --- a/app.py +++ b/app.py @@ -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