From d232b8bb5846f0e4a13bfd7d59a9e653e97efd2f Mon Sep 17 00:00:00 2001
From: Albert Villanova del Moral
 <8515462+albertvillanova@users.noreply.github.com>
Date: Thu, 18 Apr 2024 18:41:20 +0200
Subject: [PATCH] Fix test transposing image with EXIF Orientation tag (#30319)

* Fix test with exif_transpose image

* Replace datasets with PIL to load image in tests
---
 tests/utils/test_image_utils.py | 36 +++++++++++++++++++++++----------
 1 file changed, 25 insertions(+), 11 deletions(-)

diff --git a/tests/utils/test_image_utils.py b/tests/utils/test_image_utils.py
index ee45300a7e5279..d6bc9a37585899 100644
--- a/tests/utils/test_image_utils.py
+++ b/tests/utils/test_image_utils.py
@@ -16,11 +16,13 @@
 import os
 import tempfile
 import unittest
+from io import BytesIO
+from typing import Optional
 
-import datasets
 import numpy as np
 import pytest
-from huggingface_hub.file_download import http_get
+import requests
+from huggingface_hub.file_download import hf_hub_url, http_get
 from requests import ConnectTimeout, ReadTimeout
 
 from tests.pipelines.test_pipelines_document_question_answering import INVOICE_URL
@@ -39,6 +41,11 @@
     from transformers.image_utils import get_image_size, infer_channel_dimension_format, load_image
 
 
+def get_image_from_hub_dataset(dataset_id: str, filename: str, revision: Optional[str] = None) -> "PIL.Image.Image":
+    url = hf_hub_url(dataset_id, filename, repo_type="dataset", revision=revision)
+    return PIL.Image.open(BytesIO(requests.get(url).content))
+
+
 def get_random_image(height, width):
     random_array = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
     return PIL.Image.fromarray(random_array)
@@ -540,9 +547,11 @@ def test_load_img_base64(self):
     def test_load_img_rgba(self):
         # we use revision="refs/pr/1" until the PR is merged
         # https://hf.co/datasets/hf-internal-testing/fixtures_image_utils/discussions/1
-        dataset = datasets.load_dataset("hf-internal-testing/fixtures_image_utils", split="test", revision="refs/pr/1")
+        img = get_image_from_hub_dataset(
+            "hf-internal-testing/fixtures_image_utils", "0-test-lena.png", revision="refs/pr/1"
+        )
 
-        img = load_image(dataset[0]["image"])  # img with mode RGBA
+        img = load_image(img)  # img with mode RGBA
         img_arr = np.array(img)
 
         self.assertEqual(
@@ -553,9 +562,11 @@ def test_load_img_rgba(self):
     def test_load_img_la(self):
         # we use revision="refs/pr/1" until the PR is merged
         # https://hf.co/datasets/hf-internal-testing/fixtures_image_utils/discussions/1
-        dataset = datasets.load_dataset("hf-internal-testing/fixtures_image_utils", split="test", revision="refs/pr/1")
+        img = get_image_from_hub_dataset(
+            "hf-internal-testing/fixtures_image_utils", "1-test-parrots.png", revision="refs/pr/1"
+        )
 
-        img = load_image(dataset[1]["image"])  # img with mode LA
+        img = load_image(img)  # img with mode LA
         img_arr = np.array(img)
 
         self.assertEqual(
@@ -566,9 +577,11 @@ def test_load_img_la(self):
     def test_load_img_l(self):
         # we use revision="refs/pr/1" until the PR is merged
         # https://hf.co/datasets/hf-internal-testing/fixtures_image_utils/discussions/1
-        dataset = datasets.load_dataset("hf-internal-testing/fixtures_image_utils", split="test", revision="refs/pr/1")
+        img = get_image_from_hub_dataset(
+            "hf-internal-testing/fixtures_image_utils", "2-test-tree.png", revision="refs/pr/1"
+        )
 
-        img = load_image(dataset[2]["image"])  # img with mode L
+        img = load_image(img)  # img with mode L
         img_arr = np.array(img)
 
         self.assertEqual(
@@ -579,9 +592,10 @@ def test_load_img_l(self):
     def test_load_img_exif_transpose(self):
         # we use revision="refs/pr/1" until the PR is merged
         # https://hf.co/datasets/hf-internal-testing/fixtures_image_utils/discussions/1
-        dataset = datasets.load_dataset("hf-internal-testing/fixtures_image_utils", split="test", revision="refs/pr/1")
 
-        img_without_exif_transpose = dataset[3]["image"]
+        img_without_exif_transpose = get_image_from_hub_dataset(
+            "hf-internal-testing/fixtures_image_utils", "3-test-cat-rotated.jpg", revision="refs/pr/1"
+        )
         img_arr_without_exif_transpose = np.array(img_without_exif_transpose)
 
         self.assertEqual(
@@ -589,7 +603,7 @@ def test_load_img_exif_transpose(self):
             (333, 500, 3),
         )
 
-        img_with_exif_transpose = load_image(dataset[3]["image"])
+        img_with_exif_transpose = load_image(img_without_exif_transpose)
         img_arr_with_exif_transpose = np.array(img_with_exif_transpose)
 
         self.assertEqual(