Update autoshape .print() and .save() (ultralytics#2022)

fhkiel-mlaip · Jan 23, 2021 · 55f7a18 · 55f7a18
1 parent 20d9b5c
commit 55f7a18
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/models/common.py b/models/common.py
@@ -244,7 +244,7 @@ def __init__(self, imgs, pred, names=None):
     def display(self, pprint=False, show=False, save=False, render=False):
         colors = color_list()
         for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
-            str = f'Image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
+            str = f'image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
             if pred is not None:
                 for c in pred[:, -1].unique():
                     n = (pred[:, -1] == c).sum()  # detections per class
@@ -255,13 +255,13 @@ def display(self, pprint=False, show=False, save=False, render=False):
                         # str += '%s %.2f, ' % (names[int(cls)], conf)  # label
                         ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10])  # plot
             if pprint:
-                print(str)
+                print(str.rstrip(', '))
             if show:
-                img.show(f'Image {i}')  # show
+                img.show(f'image {i}')  # show
             if save:
                 f = f'results{i}.jpg'
-                str += f"saved to '{f}'"
                 img.save(f)  # save
+                print(f"{'Saving' * (i == 0)} {f},", end='' if i < self.n - 1 else ' done.\n')
             if render:
                 self.imgs[i] = np.asarray(img)