yeric1789 glenn-jocher commited on
Commit
a039b77
·
unverified ·
1 Parent(s): 0e2f2cb

PyTorch Hub `crops = results.crop()` return values (#3282)

Browse files

* Changing save_one_box

Made to work with other changes to common.py

* PEP8 and single line BGR

Co-authored-by: Glenn Jocher <[email protected]>

Files changed (1) hide show
  1. utils/general.py +6 -4
utils/general.py CHANGED
@@ -657,8 +657,8 @@ def apply_classifier(x, model, img, im0):
657
  return x
658
 
659
 
660
- def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BGR=False):
661
- # Save an image crop as {file} with crop size multiplied by {gain} and padded by {pad} pixels
662
  xyxy = torch.tensor(xyxy).view(-1, 4)
663
  b = xyxy2xywh(xyxy) # boxes
664
  if square:
@@ -666,8 +666,10 @@ def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BG
666
  b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
667
  xyxy = xywh2xyxy(b).long()
668
  clip_coords(xyxy, im.shape)
669
- crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2])]
670
- cv2.imwrite(str(increment_path(file, mkdir=True).with_suffix('.jpg')), crop if BGR else crop[..., ::-1])
 
 
671
 
672
 
673
  def increment_path(path, exist_ok=False, sep='', mkdir=False):
 
657
  return x
658
 
659
 
660
+ def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BGR=False, save=True):
661
+ # Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
662
  xyxy = torch.tensor(xyxy).view(-1, 4)
663
  b = xyxy2xywh(xyxy) # boxes
664
  if square:
 
666
  b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
667
  xyxy = xywh2xyxy(b).long()
668
  clip_coords(xyxy, im.shape)
669
+ crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
670
+ if save:
671
+ cv2.imwrite(str(increment_path(file, mkdir=True).with_suffix('.jpg')), crop)
672
+ return crop
673
 
674
 
675
  def increment_path(path, exist_ok=False, sep='', mkdir=False):