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]>
- models/common.py +5 -4
- 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 =
|
|
|
256 |
elif isinstance(im, Image.Image): # PIL Image
|
257 |
-
im, f = np.asarray(im), getattr(im, 'filename'
|
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
|