Skip to content

Commit

Permalink
Alert (no detections) (ultralytics#3984)
Browse files Browse the repository at this point in the history
* `Detections()` class `print()` overload

* Update common.py
  • Loading branch information
glenn-jocher committed Jul 12, 2021
1 parent b8737bc commit e6874b5
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion models/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down

0 comments on commit e6874b5

Please sign in to comment.