vaaliferov glenn-jocher commited on
Commit
831773f
·
unverified ·
1 Parent(s): 4717a3b

Add EXIF rotation to YOLOv5 Hub inference (#3852)

Browse files

* rotating an image according to its exif tag

* Update common.py

* Update datasets.py

* Update datasets.py

faster

* delete extraneous gpg file

* Update common.py

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

Files changed (2) hide show
  1. models/common.py +5 -4
  2. utils/datasets.py +26 -0
models/common.py CHANGED
@@ -1,9 +1,9 @@
1
  # YOLOv5 common modules
2
 
3
- import math
4
  from copy import copy
5
  from pathlib import Path
6
 
 
7
  import numpy as np
8
  import pandas as pd
9
  import requests
@@ -12,7 +12,7 @@ import torch.nn as nn
12
  from PIL import Image
13
  from torch.cuda import amp
14
 
15
- from utils.datasets import letterbox
16
  from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh, save_one_box
17
  from utils.plots import colors, plot_one_box
18
  from utils.torch_utils import time_synchronized
@@ -252,9 +252,10 @@ class AutoShape(nn.Module):
252
  for i, im in enumerate(imgs):
253
  f = f'image{i}' # filename
254
  if isinstance(im, str): # filename or uri
255
- im, f = np.asarray(Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im)), im
 
256
  elif isinstance(im, Image.Image): # PIL Image
257
- im, f = np.asarray(im), getattr(im, 'filename', f) or f
258
  files.append(Path(f).with_suffix('.jpg').name)
259
  if im.shape[0] < 5: # image in CHW
260
  im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
 
1
  # YOLOv5 common modules
2
 
 
3
  from copy import copy
4
  from pathlib import Path
5
 
6
+ import math
7
  import numpy as np
8
  import pandas as pd
9
  import requests
 
12
  from PIL import Image
13
  from torch.cuda import amp
14
 
15
+ from utils.datasets import exif_transpose, letterbox
16
  from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh, save_one_box
17
  from utils.plots import colors, plot_one_box
18
  from utils.torch_utils import time_synchronized
 
252
  for i, im in enumerate(imgs):
253
  f = f'image{i}' # filename
254
  if isinstance(im, str): # filename or uri
255
+ im, f = Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im), im
256
+ im = np.asarray(exif_transpose(im))
257
  elif isinstance(im, Image.Image): # PIL Image
258
+ im, f = np.asarray(exif_transpose(im)), getattr(im, 'filename') or f
259
  files.append(Path(f).with_suffix('.jpg').name)
260
  if im.shape[0] < 5: # image in CHW
261
  im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
utils/datasets.py CHANGED
@@ -64,6 +64,32 @@ def exif_size(img):
64
  return s
65
 
66
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
67
  def create_dataloader(path, imgsz, batch_size, stride, single_cls=False, hyp=None, augment=False, cache=False, pad=0.0,
68
  rect=False, rank=-1, workers=8, image_weights=False, quad=False, prefix=''):
69
  # Make sure only the first process in DDP process the dataset first, and the following others can use the cache
 
64
  return s
65
 
66
 
67
+ def exif_transpose(image):
68
+ """
69
+ Transpose a PIL image accordingly if it has an EXIF Orientation tag.
70
+ From https://github.com/python-pillow/Pillow/blob/master/src/PIL/ImageOps.py
71
+
72
+ :param image: The image to transpose.
73
+ :return: An image.
74
+ """
75
+ exif = image.getexif()
76
+ orientation = exif.get(0x0112, 1) # default 1
77
+ if orientation > 1:
78
+ method = {2: Image.FLIP_LEFT_RIGHT,
79
+ 3: Image.ROTATE_180,
80
+ 4: Image.FLIP_TOP_BOTTOM,
81
+ 5: Image.TRANSPOSE,
82
+ 6: Image.ROTATE_270,
83
+ 7: Image.TRANSVERSE,
84
+ 8: Image.ROTATE_90,
85
+ }.get(orientation)
86
+ if method is not None:
87
+ image = image.transpose(method)
88
+ del exif[0x0112]
89
+ image.info["exif"] = exif.tobytes()
90
+ return image
91
+
92
+
93
  def create_dataloader(path, imgsz, batch_size, stride, single_cls=False, hyp=None, augment=False, cache=False, pad=0.0,
94
  rect=False, rank=-1, workers=8, image_weights=False, quad=False, prefix=''):
95
  # Make sure only the first process in DDP process the dataset first, and the following others can use the cache