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 <glenn.jocher@ultralytics.com>
single_channel
Simon Lévesque 1 year ago committed by GitHub
parent 29ace4982d
commit a76af55533
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -354,12 +354,14 @@ class Results(SimpleClass):
results = [] results = []
data = self.boxes.data.cpu().tolist() data = self.boxes.data.cpu().tolist()
h, w = self.orig_shape if normalize else (1, 1) 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} box = {'x1': row[0] / w, 'y1': row[1] / h, 'x2': row[2] / w, 'y2': row[3] / h}
conf = row[4] conf = row[-2]
id = int(row[5]) class_id = int(row[-1])
name = self.names[id] name = self.names[class_id]
result = {'name': name, 'class': id, 'confidence': conf, 'box': box} 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: if self.masks:
x, y = self.masks.xy[i][:, 0], self.masks.xy[i][:, 1] # numpy array x, y = self.masks.xy[i][:, 0], self.masks.xy[i][:, 1] # numpy array
result['segments'] = {'x': (x / w).tolist(), 'y': (y / h).tolist()} result['segments'] = {'x': (x / w).tolist(), 'y': (y / h).tolist()}
@ -404,7 +406,7 @@ class Boxes(BaseTensor):
if boxes.ndim == 1: if boxes.ndim == 1:
boxes = boxes[None, :] boxes = boxes[None, :]
n = boxes.shape[-1] 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) super().__init__(boxes, orig_shape)
self.is_track = n == 7 self.is_track = n == 7
self.orig_shape = orig_shape self.orig_shape = orig_shape

Loading…
Cancel
Save