ELHoussineT glenn-jocher commited on
Commit
0d8a184
1 Parent(s): 8e94bf6

Add `crops = results.crops()` dictionary (#4676)

Browse files

* adding get cropped functionality

* Add target logic in existing functions

* Crops cleanup

* Add dictionary keys: conf, cls, box

* Bug fixes - avoid return after first image

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

Files changed (1) hide show
  1. models/common.py +11 -5
models/common.py CHANGED
@@ -365,6 +365,7 @@ class Detections:
365
  self.s = shape # inference BCHW shape
366
 
367
  def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')):
 
368
  for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
369
  str = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} '
370
  if pred.shape[0]:
@@ -376,7 +377,9 @@ class Detections:
376
  for *box, conf, cls in reversed(pred): # xyxy, confidence, class
377
  label = f'{self.names[int(cls)]} {conf:.2f}'
378
  if crop:
379
- save_one_box(box, im, file=save_dir / 'crops' / self.names[int(cls)] / self.files[i])
 
 
380
  else: # all others
381
  annotator.box_label(box, label, color=colors(cls))
382
  im = annotator.im
@@ -395,6 +398,10 @@ class Detections:
395
  LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
396
  if render:
397
  self.imgs[i] = np.asarray(im)
 
 
 
 
398
 
399
  def print(self):
400
  self.display(pprint=True) # print results
@@ -408,10 +415,9 @@ class Detections:
408
  save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir
409
  self.display(save=True, save_dir=save_dir) # save results
410
 
411
- def crop(self, save_dir='runs/detect/exp'):
412
- save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir
413
- self.display(crop=True, save_dir=save_dir) # crop results
414
- LOGGER.info(f'Saved results to {save_dir}\n')
415
 
416
  def render(self):
417
  self.display(render=True) # render results
 
365
  self.s = shape # inference BCHW shape
366
 
367
  def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')):
368
+ crops = []
369
  for i, (im, pred) in enumerate(zip(self.imgs, self.pred)):
370
  str = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} '
371
  if pred.shape[0]:
 
377
  for *box, conf, cls in reversed(pred): # xyxy, confidence, class
378
  label = f'{self.names[int(cls)]} {conf:.2f}'
379
  if crop:
380
+ file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None
381
+ crops.append({'box': box, 'conf': conf, 'cls': cls, 'label': label,
382
+ 'im': save_one_box(box, im, file=file, save=save)})
383
  else: # all others
384
  annotator.box_label(box, label, color=colors(cls))
385
  im = annotator.im
 
398
  LOGGER.info(f"Saved {self.n} image{'s' * (self.n > 1)} to {colorstr('bold', save_dir)}")
399
  if render:
400
  self.imgs[i] = np.asarray(im)
401
+ if crop:
402
+ if save:
403
+ LOGGER.info(f'Saved results to {save_dir}\n')
404
+ return crops
405
 
406
  def print(self):
407
  self.display(pprint=True) # print results
 
415
  save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir
416
  self.display(save=True, save_dir=save_dir) # save results
417
 
418
+ def crop(self, save=True, save_dir='runs/detect/exp'):
419
+ save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) if save else None
420
+ return self.display(crop=True, save=save, save_dir=save_dir) # crop results
 
421
 
422
  def render(self):
423
  self.display(render=True) # render results