apply exif

stefanklut · Dec 4, 2023 · fbad945 · fbad945
1 parent 5403a76
commit fbad945
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/utils/image_utils.py b/utils/image_utils.py
@@ -9,7 +9,7 @@
 import torch
 import torchvision
 from detectron2.data.detection_utils import convert_PIL_to_numpy
-from PIL import Image
+from PIL import Image, ImageOps
 
 sys.path.append(str(Path(__file__).resolve().parent.joinpath("..")))
 from utils.logging_utils import get_logger_name
@@ -18,6 +18,7 @@
 def load_image_array_from_path(
     image_path: Path | str,
     mode: str = "color",
+    ignore_exif: bool = False,
 ) -> Optional[np.ndarray]:
     """
     Load image from a given path, return None if loading failed due to corruption
@@ -36,7 +37,10 @@ def load_image_array_from_path(
     try:
         # image = cv2.imread(str(image_path), cv2.IMREAD_COLOR if mode == "color" else cv2.IMREAD_GRAYSCALE)
         # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) if mode == "color" else image
-        image = convert_PIL_to_numpy(Image.open(image_path), "RGB" if mode == "color" else "L").copy()
+        image = Image.open(image_path)
+        if not ignore_exif:
+            image = ImageOps.exif_transpose(image)
+        image = convert_PIL_to_numpy(image, "RGB" if mode == "color" else "L").copy()
         if mode == "grayscale":
             image = image.squeeze(axis=2)
         return image
@@ -77,6 +81,7 @@ def load_image_array_from_bytes(
     image_bytes: bytes,
     image_path: Optional[Path] = None,
     mode: str = "color",
+    ignore_exif: bool = False,
 ) -> Optional[np.ndarray]:
     """
     Load image based on given bytes, return None if loading failed due to corruption
@@ -96,7 +101,10 @@ def load_image_array_from_bytes(
         # bytes_array = np.frombuffer(image_bytes, np.uint8)
         # image = cv2.imdecode(bytes_array, cv2.IMREAD_COLOR if mode == "color" else cv2.IMREAD_GRAYSCALE)
         # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) if mode == "color" else image
-        image = convert_PIL_to_numpy(Image.open(BytesIO(image_bytes)), "RGB" if mode == "color" else "L").copy()
+        image = Image.open(BytesIO(image_bytes))
+        if not ignore_exif:
+            image = ImageOps.exif_transpose(image)
+        image = convert_PIL_to_numpy(image, "RGB" if mode == "color" else "L").copy()
         if mode == "grayscale":
             image = image.squeeze(axis=2)
         return image