glenn-jocher commited on
Commit
2fc4760
1 Parent(s): 3b7feea

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

Browse files
Files changed (1) hide show
  1. models/common.py +4 -4
models/common.py CHANGED
@@ -244,7 +244,7 @@ class Detections:
244
  def display(self, pprint=False, show=False, save=False, render=False):
245
  colors = color_list()
246
  for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
247
- str = f'Image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
248
  if pred is not None:
249
  for c in pred[:, -1].unique():
250
  n = (pred[:, -1] == c).sum() # detections per class
@@ -255,13 +255,13 @@ class Detections:
255
  # str += '%s %.2f, ' % (names[int(cls)], conf) # label
256
  ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10]) # plot
257
  if pprint:
258
- print(str)
259
  if show:
260
- img.show(f'Image {i}') # show
261
  if save:
262
  f = f'results{i}.jpg'
263
- str += f"saved to '{f}'"
264
  img.save(f) # save
 
265
  if render:
266
  self.imgs[i] = np.asarray(img)
267
 
 
244
  def display(self, pprint=False, show=False, save=False, render=False):
245
  colors = color_list()
246
  for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
247
+ str = f'image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
248
  if pred is not None:
249
  for c in pred[:, -1].unique():
250
  n = (pred[:, -1] == c).sum() # detections per class
 
255
  # str += '%s %.2f, ' % (names[int(cls)], conf) # label
256
  ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10]) # plot
257
  if pprint:
258
+ print(str.rstrip(', '))
259
  if show:
260
+ img.show(f'image {i}') # show
261
  if save:
262
  f = f'results{i}.jpg'
 
263
  img.save(f) # save
264
+ print(f"{'Saving' * (i == 0)} {f},", end='' if i < self.n - 1 else ' done.\n')
265
  if render:
266
  self.imgs[i] = np.asarray(img)
267