diff --git a/models/common.py b/models/common.py index 418034ddeaac..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: