From f0a82f6a02acfc851fc011d2babcc07787246919 Mon Sep 17 00:00:00 2001 From: yeric1789 <76454253+yeric1789@users.noreply.github.com> Date: Sun, 23 May 2021 10:00:55 -0400 Subject: [PATCH] PyTorch Hub `crops = results.crop()` return values (#3282) * Changing save_one_box Made to work with other changes to common.py * PEP8 and single line BGR Co-authored-by: Glenn Jocher --- utils/general.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/utils/general.py b/utils/general.py index 4b3d4ab3b189..47eea78eda30 100755 --- a/utils/general.py +++ b/utils/general.py @@ -657,8 +657,8 @@ 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: @@ -666,8 +666,10 @@ def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BG 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):