change return type to model names
This commit is contained in:
parent
3224295610
commit
93f8974ba0
2 changed files with 6 additions and 1 deletions
2
.gitignore
vendored
Normal file
2
.gitignore
vendored
Normal file
|
@ -0,0 +1,2 @@
|
||||||
|
|
||||||
|
/__pycache__
|
5
app.py
5
app.py
|
@ -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
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue