This commit is contained in:
Glenn Jocher 2019-07-09 20:56:58 +02:00
parent bf6d96330b
commit d5fd37de26
1 changed files with 23 additions and 3 deletions

View File

@ -10,12 +10,33 @@ import numpy as np
import torch import torch
from torch.utils.data import Dataset from torch.utils.data import Dataset
from tqdm import tqdm from tqdm import tqdm
from PIL import Image, ExifTags
from utils.utils import xyxy2xywh, xywh2xyxy from utils.utils import xyxy2xywh, xywh2xyxy
img_formats = ['.bmp', '.jpg', '.jpeg', '.png', '.tif'] img_formats = ['.bmp', '.jpg', '.jpeg', '.png', '.tif']
vid_formats = ['.mov', '.avi', '.mp4'] vid_formats = ['.mov', '.avi', '.mp4']
# Get orientation exif tag
for orientation in ExifTags.TAGS.keys():
if ExifTags.TAGS[orientation] == 'Orientation':
break
def exif_size(img):
# Returns exif-corrected PIL size
s = img.size # (width, height)
try:
rotation = dict(img._getexif().items())[orientation]
if rotation == 6: # rotation 270
s = (s[1], s[0])
elif rotation == 8: # rotation 90
s = (s[1], s[0])
except:
None
return s
class LoadImages: # for inference class LoadImages: # for inference
def __init__(self, path, img_size=416): def __init__(self, path, img_size=416):
@ -154,12 +175,10 @@ class LoadImagesAndLabels(Dataset): # for training/testing
# Rectangular Training https://github.com/ultralytics/yolov3/issues/232 # Rectangular Training https://github.com/ultralytics/yolov3/issues/232
if self.rect: if self.rect:
from PIL import Image
# Read image shapes # Read image shapes
sp = 'data' + os.sep + path.replace('.txt', '.shapes').split(os.sep)[-1] # shapefile path sp = 'data' + os.sep + path.replace('.txt', '.shapes').split(os.sep)[-1] # shapefile path
if not os.path.exists(sp): # read shapes using PIL and write shapefile for next time (faster) if not os.path.exists(sp): # read shapes using PIL and write shapefile for next time (faster)
s = [Image.open(f).size for f in tqdm(self.img_files, desc='Reading image shapes')] s = [exif_size(Image.open(f)) for f in tqdm(self.img_files, desc='Reading image shapes')]
np.savetxt(sp, s, fmt='%g') np.savetxt(sp, s, fmt='%g')
with open(sp, 'r') as f: # read existing shapefile with open(sp, 'r') as f: # read existing shapefile
@ -352,6 +371,7 @@ def letterbox(img, new_shape=416, color=(127.5, 127.5, 127.5), mode='auto'):
# Resize a rectangular image to a 32 pixel multiple rectangle # Resize a rectangular image to a 32 pixel multiple rectangle
# https://github.com/ultralytics/yolov3/issues/232 # https://github.com/ultralytics/yolov3/issues/232
shape = img.shape[:2] # current shape [height, width] shape = img.shape[:2] # current shape [height, width]
if isinstance(new_shape, int): if isinstance(new_shape, int):
ratio = float(new_shape) / max(shape) ratio = float(new_shape) / max(shape)
else: else: