Speed up preloading by not loading pixel data

PIL lazy loads pixel data, so image size can be accessed without loading the full image.
This commit makes it so only image size and EXIF metadata are fetched from disk, speeding up the preload stage while still supporting transposing images.
This commit is contained in:
MFAlex 2023-05-31 23:16:01 +12:00 committed by Victor Hall
parent 5c98cdee70
commit 7dcfa7acbf
1 changed files with 28 additions and 2 deletions

View File

@ -162,6 +162,29 @@ class ImageTrainItem:
pass
return image
def _needs_transpose(self, image, print_error=False):
try:
exif = image.getexif()
orientation = exif.get(0x0112)
"""
https://pillow.readthedocs.io/en/stable/_modules/PIL/ImageOps.html#exif_transpose
method = {
2: Image.Transpose.FLIP_LEFT_RIGHT,
3: Image.Transpose.ROTATE_180,
4: Image.Transpose.FLIP_TOP_BOTTOM,
5: Image.Transpose.TRANSPOSE,
6: Image.Transpose.ROTATE_270,
7: Image.Transpose.TRANSVERSE,
8: Image.Transpose.ROTATE_90,
}.get(orientation)
"""
return orientation in [5, 6, 7, 8]
except Exception as e:
logging.warning(F"Error rotating image: {e} on {self.pathname}, image will be loaded as is, EXIF may be corrupt") if print_error else None
pass
return False
def _percent_random_crop(self, image, crop_jitter=0.02):
"""
randomly crops the image by a percentage of the image size on each of the four sides
@ -270,7 +293,10 @@ class ImageTrainItem:
self.target_wh = None
try:
with PIL.Image.open(self.pathname) as image:
image = self._try_transpose(image, print_error=True).convert('RGB')
needs_transpose = self._needs_transpose(image)
if needs_transpose:
height, width = image.size
else:
width, height = image.size
image_aspect = width / height
target_wh = min(self.aspects, key=lambda aspects:abs(aspects[0]/aspects[1] - image_aspect))