From 525959249fb420fcd269d439fe362b9fbab81070 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 <glenn.jocher@ultralytics.com>
(cherry picked from commit a039b7796dc495a5e5ec5ac5ba60cbecfb8b3c60)
---
 utils/general.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/utils/general.py b/utils/general.py
index 05715197c392..fc59ad1e5dd1 100644
--- a/utils/general.py
+++ b/utils/general.py
@@ -660,8 +660,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:
@@ -669,8 +669,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):