From a76af55533ce63cb5717ace6e4648b5f3b52a001 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Simon=20L=C3=A9vesque?= Date: Wed, 9 Aug 2023 18:19:28 -0400 Subject: [PATCH] Fixes Results.tojson when tracking (#4270) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher --- ultralytics/engine/results.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/ultralytics/engine/results.py b/ultralytics/engine/results.py index 60b57bf..fd0d48b 100644 --- a/ultralytics/engine/results.py +++ b/ultralytics/engine/results.py @@ -354,12 +354,14 @@ class Results(SimpleClass): results = [] data = self.boxes.data.cpu().tolist() h, w = self.orig_shape if normalize else (1, 1) - for i, row in enumerate(data): + for i, row in enumerate(data): # xyxy, track_id if tracking, conf, class_id box = {'x1': row[0] / w, 'y1': row[1] / h, 'x2': row[2] / w, 'y2': row[3] / h} - conf = row[4] - id = int(row[5]) - name = self.names[id] - result = {'name': name, 'class': id, 'confidence': conf, 'box': box} + conf = row[-2] + class_id = int(row[-1]) + name = self.names[class_id] + result = {'name': name, 'class': class_id, 'confidence': conf, 'box': box} + if self.boxes.is_track: + result['track_id'] = int(row[-3]) # track ID if self.masks: x, y = self.masks.xy[i][:, 0], self.masks.xy[i][:, 1] # numpy array result['segments'] = {'x': (x / w).tolist(), 'y': (y / h).tolist()} @@ -404,7 +406,7 @@ class Boxes(BaseTensor): if boxes.ndim == 1: boxes = boxes[None, :] n = boxes.shape[-1] - assert n in (6, 7), f'expected `n` in [6, 7], but got {n}' # xyxy, (track_id), conf, cls + assert n in (6, 7), f'expected `n` in [6, 7], but got {n}' # xyxy, track_id, conf, cls super().__init__(boxes, orig_shape) self.is_track = n == 7 self.orig_shape = orig_shape