From 8d7490f060a5320010670a387d8c6947257a2091 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 16 Aug 2023 19:31:35 +0200
Subject: [PATCH] `ultralytics 8.0.156` fix PIL `*.png` EXIF bug (#4406)

---
 tests/test_cli.py         |  2 +-
 ultralytics/__init__.py   |  2 +-
 ultralytics/data/utils.py | 20 +++++++++-----------
 3 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/tests/test_cli.py b/tests/test_cli.py
index 88888a1..7c1083c 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -115,4 +115,4 @@ def test_mobilesam():
 @pytest.mark.parametrize('task,model,data', TASK_ARGS)
 def test_train_gpu(task, model, data):
     run(f'yolo train {task} model={model}.yaml data={data} imgsz=32 epochs=1 device="0"')  # single GPU
-    run(f'yolo train {task} model={model}.pt data={data} imgsz=32 epochs=1 device="0,1"')  # Multi GPU
+    run(f'yolo train {task} model={model}.pt data={data} imgsz=32 epochs=1 device="0,1"')  # multi GPU
diff --git a/ultralytics/__init__.py b/ultralytics/__init__.py
index 5586d04..7e1a551 100644
--- a/ultralytics/__init__.py
+++ b/ultralytics/__init__.py
@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = '8.0.155'
+__version__ = '8.0.156'
 
 from ultralytics.hub import start
 from ultralytics.models import RTDETR, SAM, YOLO
diff --git a/ultralytics/data/utils.py b/ultralytics/data/utils.py
index 68b423b..0c07547 100644
--- a/ultralytics/data/utils.py
+++ b/ultralytics/data/utils.py
@@ -1,5 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
+import contextlib
 import hashlib
 import json
 import os
@@ -13,7 +14,7 @@ from tarfile import is_tarfile
 
 import cv2
 import numpy as np
-from PIL import ExifTags, Image, ImageOps
+from PIL import Image, ImageOps
 from tqdm import tqdm
 
 from ultralytics.nn.autobackend import check_class_names
@@ -28,11 +29,6 @@ IMG_FORMATS = 'bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp',
 VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv', 'webm'  # video suffixes
 PIN_MEMORY = str(os.getenv('PIN_MEMORY', True)).lower() == 'true'  # global pin_memory for dataloaders
 
-# Get orientation exif tag
-for orientation in ExifTags.TAGS.keys():
-    if ExifTags.TAGS[orientation] == 'Orientation':
-        break
-
 
 def img2label_paths(img_paths):
     """Define label paths as a function of image paths."""
@@ -51,11 +47,13 @@ def get_hash(paths):
 def exif_size(img: Image.Image):
     """Returns exif-corrected PIL size."""
     s = img.size  # (width, height)
-    exif = img.getexif()
-    if exif:
-        rotation = exif.get(274, None)  # the key for the orientation tag in the EXIF data is 274 (in decimal)
-        if rotation in [6, 8]:  # rotation 270 or 90
-            s = s[1], s[0]
+    if img.format == 'JPEG':  # only support JPEG images
+        with contextlib.suppress(Exception):
+            exif = img.getexif()
+            if exif:
+                rotation = exif.get(274, None)  # the EXIF key for the orientation tag is 274
+                if rotation in [6, 8]:  # rotation 270 or 90
+                    s = s[1], s[0]
     return s