|
|
|
@ -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
|
|
|
|
|