From a0225f55b202902dd3d68129d5274afbd278de47 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 12 Jul 2021 19:32:47 +0200 Subject: [PATCH 1/2] `Detections()` class `print()` overload --- models/common.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/models/common.py b/models/common.py index 418034ddeaac..a0af740db850 100644 --- a/models/common.py +++ b/models/common.py @@ -332,8 +332,7 @@ def display(self, pprint=False, show=False, save=False, crop=False, render=False self.imgs[i] = np.asarray(im) def print(self): - self.display(pprint=True) # print results - print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % self.t) + print(self.__str__()) # print results def show(self): self.display(show=True) # show results @@ -372,6 +371,10 @@ def tolist(self): def __len__(self): return self.n + def __str__(self): + self.display(pprint=True) + return f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % self.t + class Classify(nn.Module): # Classification head, i.e. x(b,c1,20,20) to x(b,c2) From 636fe8ec67955f28a2d330e2b376d40cecb28864 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 12 Jul 2021 19:46:00 +0200 Subject: [PATCH 2/2] Update common.py --- models/common.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/models/common.py b/models/common.py index a0af740db850..05372ae149f5 100644 --- a/models/common.py +++ b/models/common.py @@ -307,7 +307,7 @@ def __init__(self, imgs, pred, files, times=None, names=None, shape=None): def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')): for i, (im, pred) in enumerate(zip(self.imgs, self.pred)): str = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' - if pred is not None: + if pred.shape[0]: for c in pred[:, -1].unique(): n = (pred[:, -1] == c).sum() # detections per class str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string @@ -318,6 +318,8 @@ def display(self, pprint=False, show=False, save=False, crop=False, render=False save_one_box(box, im, file=save_dir / 'crops' / self.names[int(cls)] / self.files[i]) else: # all others plot_one_box(box, im, label=label, color=colors(cls)) + else: + str += '(no detections)' im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np if pprint: @@ -332,7 +334,8 @@ def display(self, pprint=False, show=False, save=False, crop=False, render=False self.imgs[i] = np.asarray(im) def print(self): - print(self.__str__()) # print results + self.display(pprint=True) # print results + print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % self.t) def show(self): self.display(show=True) # show results @@ -371,10 +374,6 @@ def tolist(self): def __len__(self): return self.n - def __str__(self): - self.display(pprint=True) - return f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % self.t - class Classify(nn.Module): # Classification head, i.e. x(b,c1,20,20) to x(b,c2)