tensor conversion

This commit is contained in:
Siwat Sirichai 2025-01-14 23:35:51 +07:00
parent ce53d60702
commit e73bd95677

5
app.py
View file

@ -64,11 +64,14 @@ async def detect(websocket: WebSocket):
boxes = [] boxes = []
for r in model.track(frame, stream=False, persist=True): for r in model.track(frame, stream=False, persist=True):
for box in r.boxes: for box in r.boxes:
track_id = None
if hasattr(box, "id") and box.id is not None:
track_id = box.id.item()
box_cpu = box.cpu() box_cpu = box.cpu()
boxes.append({ boxes.append({
"class": model.names[int(box_cpu.cls[0])], "class": model.names[int(box_cpu.cls[0])],
"confidence": float(box_cpu.conf[0]), "confidence": float(box_cpu.conf[0]),
"id": box_cpu.id, "id": track_id,
}) })
# Broadcast to all subscribers of this URL # Broadcast to all subscribers of this URL
detection_data = { detection_data = {