Skip to content

Commit

Permalink
PyTorch Hub crops = results.crop() return values (ultralytics#3282)
Browse files Browse the repository at this point in the history
* Changing save_one_box 

Made to work with other changes to common.py

* PEP8 and single line BGR

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
  • Loading branch information
yeric1789 and glenn-jocher committed May 23, 2021
1 parent 0e2f2cb commit a039b77
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions utils/general.py
Original file line number Diff line number Diff line change
Expand Up @@ -657,17 +657,19 @@ def apply_classifier(x, model, img, im0):
return x


def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BGR=False):
# Save an image crop as {file} with crop size multiplied by {gain} and padded by {pad} pixels
def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BGR=False, save=True):
# Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
xyxy = torch.tensor(xyxy).view(-1, 4)
b = xyxy2xywh(xyxy) # boxes
if square:
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square
b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
xyxy = xywh2xyxy(b).long()
clip_coords(xyxy, im.shape)
crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2])]
cv2.imwrite(str(increment_path(file, mkdir=True).with_suffix('.jpg')), crop if BGR else crop[..., ::-1])
crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
if save:
cv2.imwrite(str(increment_path(file, mkdir=True).with_suffix('.jpg')), crop)
return crop


def increment_path(path, exist_ok=False, sep='', mkdir=False):
Expand Down

0 comments on commit a039b77

Please sign in to comment.