Skip to content

Commit 7a1968f

Browse files
authored
add 'save_one_box' method for '--save-crop' argument
method from [https://github.com/WongKinYiu/yolov7/blob/u5/utils/plots.py#L474]
1 parent 36266e1 commit 7a1968f

File tree

1 file changed

+19
-1
lines changed

1 file changed

+19
-1
lines changed

track.py

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,13 +32,31 @@
3232
from yolov7.models.experimental import attempt_load
3333
from yolov7.utils.datasets import LoadImages, LoadStreams
3434
from yolov7.utils.general import (check_img_size, non_max_suppression, scale_coords, check_requirements, cv2,
35-
check_imshow, xyxy2xywh, increment_path, strip_optimizer, colorstr, check_file)
35+
check_imshow, xyxy2xywh, xywh2xyxy, clip_coords, increment_path, strip_optimizer, colorstr, check_file)
3636
from yolov7.utils.torch_utils import select_device, time_synchronized
3737
from yolov7.utils.plots import plot_one_box
3838
from strong_sort.utils.parser import get_config
3939
from strong_sort.strong_sort import StrongSORT
4040

4141

42+
# method from [https://github.com/WongKinYiu/yolov7/blob/u5/utils/plots.py#L474] for availabling '--save-crop' argument
43+
def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, BGR=False, save=True):
44+
# Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
45+
xyxy = torch.tensor(xyxy).view(-1, 4)
46+
b = xyxy2xywh(xyxy) # boxes
47+
if square:
48+
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square
49+
b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
50+
xyxy = xywh2xyxy(b).long()
51+
clip_coords(xyxy, im.shape)
52+
crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
53+
if save:
54+
file.parent.mkdir(parents=True, exist_ok=True) # make directory
55+
f = str(Path(increment_path(file)).with_suffix('.jpg'))
56+
# cv2.imwrite(f, crop) # https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue
57+
Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)).save(f, quality=95, subsampling=0)
58+
return crop
59+
4260
VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv' # include video suffixes
4361

4462

0 commit comments

Comments
 (0)