laserborg glenn-jocher commited on
Commit
529fbc1
1 Parent(s): 9f4d113

Use PIL to eliminate chroma subsampling in crops (#7008)

Browse files

* use pillow to save higher-quality jpg (w/o color subsampling)

* Cleanup and doc issue

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

Files changed (1) hide show
  1. utils/plots.py +4 -2
utils/plots.py CHANGED
@@ -458,7 +458,7 @@ def profile_idetection(start=0, stop=0, labels=(), save_dir=''):
458
  plt.savefig(Path(save_dir) / 'idetection_profile.png', dpi=200)
459
 
460
 
461
- def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BGR=False, save=True):
462
  # Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
463
  xyxy = torch.tensor(xyxy).view(-1, 4)
464
  b = xyxy2xywh(xyxy) # boxes
@@ -470,5 +470,7 @@ def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BG
470
  crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
471
  if save:
472
  file.parent.mkdir(parents=True, exist_ok=True) # make directory
473
- cv2.imwrite(str(increment_path(file).with_suffix('.jpg')), crop)
 
 
474
  return crop
 
458
  plt.savefig(Path(save_dir) / 'idetection_profile.png', dpi=200)
459
 
460
 
461
+ def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, BGR=False, save=True):
462
  # Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
463
  xyxy = torch.tensor(xyxy).view(-1, 4)
464
  b = xyxy2xywh(xyxy) # boxes
 
470
  crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
471
  if save:
472
  file.parent.mkdir(parents=True, exist_ok=True) # make directory
473
+ f = str(increment_path(file).with_suffix('.jpg'))
474
+ # cv2.imwrite(f, crop) # https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue
475
+ Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)).save(f, quality=95, subsampling=0)
476
  return crop