|
|
@@ -525,7 +525,7 @@ class AutoShape(nn.Module): |
|
|
|
|
|
|
|
class Detections: |
|
|
|
# YOLOv5 detections class for inference results |
|
|
|
def __init__(self, imgs, pred, files, times=None, names=None, shape=None): |
|
|
|
def __init__(self, imgs, pred, files, times=(0, 0, 0, 0), names=None, shape=None): |
|
|
|
super().__init__() |
|
|
|
d = pred[0].device # device |
|
|
|
gn = [torch.tensor([*(im.shape[i] for i in [1, 0, 1, 0]), 1, 1], device=d) for im in imgs] # normalizations |
|
|
@@ -533,6 +533,7 @@ class Detections: |
|
|
|
self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls) |
|
|
|
self.names = names # class names |
|
|
|
self.files = files # image filenames |
|
|
|
self.times = times # profiling times |
|
|
|
self.xyxy = pred # xyxy pixels |
|
|
|
self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels |
|
|
|
self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized |
|
|
@@ -612,10 +613,11 @@ class Detections: |
|
|
|
|
|
|
|
def tolist(self): |
|
|
|
# return a list of Detections objects, i.e. 'for result in results.tolist():' |
|
|
|
x = [Detections([self.imgs[i]], [self.pred[i]], names=self.names, shape=self.s) for i in range(self.n)] |
|
|
|
for d in x: |
|
|
|
for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']: |
|
|
|
setattr(d, k, getattr(d, k)[0]) # pop out of list |
|
|
|
r = range(self.n) # iterable |
|
|
|
x = [Detections([self.imgs[i]], [self.pred[i]], [self.files[i]], self.times, self.names, self.s) for i in r] |
|
|
|
# for d in x: |
|
|
|
# for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']: |
|
|
|
# setattr(d, k, getattr(d, k)[0]) # pop out of list |
|
|
|
return x |
|
|
|
|
|
|
|
def __len__(self): |