Skip to content

Commit

Permalink
Fixes Results.tojson when tracking (ultralytics#4270)
Browse files Browse the repository at this point in the history
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <[email protected]>
  • Loading branch information
3 people authored Aug 9, 2023
1 parent 29ace49 commit a76af55
Showing 1 changed file with 8 additions and 6 deletions.
14 changes: 8 additions & 6 deletions ultralytics/engine/results.py
Original file line number Diff line number Diff line change
Expand Up @@ -354,12 +354,14 @@ def tojson(self, normalize=False):
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()}
Expand Down Expand Up @@ -404,7 +406,7 @@ def __init__(self, boxes, orig_shape) -> None:
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
Expand Down

0 comments on commit a76af55

Please sign in to comment.