Repository: huggingface/transformers (Python)

Issue:
Image classification pipeline fails with base64.encodebytes() input
### System Info

transformers - 4.39.1
python 3.10.11
platform - ubuntu 22.04

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

``` python
from PIL import Image
import base64
from io import BytesIO
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-18")
tokenizer = AutoFeatureExtractor.from_pretrained("microsoft/resnet-18")
tr_model = {"model": model, "image_processor": tokenizer}
vision_model = transformers.pipeline(
    task="image-classification", model=model, image_processor=tokenizer
)

buffered = BytesIO()
im = torch.rand((256, 256,3))
image = Image.fromarray(im.numpy().astype('uint8'), 'RGB')
image.save(buffered, format="JPEG")
img_str1 = base64.b64encode(buffered.getvalue()).decode("utf-8")
pred1 = vision_model.predict(img_str1)

img_str2 = base64.encodebytes(buffered.getvalue()).decode("utf-8")
pred2 = vision_model.predict(img_str2)`
```

results in this error:
Traceback (most recent call last):
  File "/home/azureuser/workspace/AzureMlCli/test_oss.py", line 22, in <module>
    pred2 = vision_model.predict(img_str2)
  File "/anaconda/envs/vision_finetune/lib/python3.10/site-packages/transformers/pipelines/base.py", line 921, in predict
    return self(X)
  File "/anaconda/envs/vision_finetune/lib/python3.10/site-packages/transformers/pipelines/image_classification.py", line 158, in __call__
    return super().__call__(images, **kwargs)
  File "/anaconda/envs/vision_finetune/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1162, in __call__
    return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)
  File "/anaconda/envs/vision_finetune/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1168, in run_single
    model_inputs = self.preprocess(inputs, **preprocess_params)
  File "/anaconda/envs/vision_finetune/lib/python3.10/site-packages/transformers/pipelines/image_classification.py", line 161, in preprocess
    image = load_image(image, timeout=timeout)
  File "/anaconda/envs/vision_finetune/lib/python3.10/site-packages/transformers/image_utils.py", line 326, in load_image
    raise ValueError(
ValueError: Incorrect image source. Must be a valid URL starting with `http://` or `https://`, a valid path to an image file, or a base64 encoded string. Got /9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0a
HBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/2wBDAQkJCQwLDBgNDRgyIRwhMjIyMjIy
MjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjIyMjL/wAARCAEAAQADASIA
AhEBAxEB/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQA
AAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3
ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWm
p6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/8QAHwEA
AwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoL/8QAtREAAgECBAQDBAcFBAQAAQJ3AAECAxEEBSEx
BhJBUQdhcRMiMoEIFEKRobHBCSMzUvAVYnLRChYkNOEl8RcYGRomJygpKjU2Nzg5OkNERUZHSElK
U1RVVldYWVpjZGVmZ2hpanN0dXZ3eHl6goOEhYaHiImKkpOUlZaXmJmaoqOkpaanqKmqsrO0tba3
uLm6wsPExcbHyMnK0tPU1dbX2Nna4uPk5ebn6Onq8vP09fb3+Pn6/9oADAMBAAIRAxEAPwD5/ooo
oAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiig
AooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKAC
iiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKK
KKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooo
oAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiig
AooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKAC
iiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKK
KKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooo
oAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiig
AooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKAC
iiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKK
KKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooo
oAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiig
AooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKAC
iiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKK
KKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooo
oAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiig
AooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooAKKKKACiiigAooooA//2Q==
. Failed with Non-base64 digit found


Image encoded with b64encode is successful, but encodebytes results in the above error.

### Expected behavior

Image classification pipeline works well when input is given in base64.base64encode() format but fails with base64.encodebytes() format. Shouldn't the pipeline work on both the input formats?


Discussion:
Hi @MadhuM02 - thanks for raising this issue and providing a clear reproducer! I've opened a PR to enable this
> Hi @MadhuM02 - thanks for raising this issue and providing a clear reproducer! I've opened a PR to enable this

thanks for raising a PR to address the issue.



Patch:
diff --git a/src/transformers/image_utils.py b/src/transformers/image_utils.py
index e4a55b3455a3..7d71fc982b43 100644
--- a/src/transformers/image_utils.py
+++ b/src/transformers/image_utils.py
@@ -320,7 +320,7 @@ def load_image(image: Union[str, "PIL.Image.Image"], timeout: Optional[float] =
 
             # Try to load as base64
             try:
-                b64 = base64.b64decode(image, validate=True)
+                b64 = base64.decodebytes(image.encode())
                 image = PIL.Image.open(BytesIO(b64))
             except Exception as e:
                 raise ValueError(


Test patch:
diff --git a/tests/utils/test_image_utils.py b/tests/utils/test_image_utils.py
index d6bc9a375858..f360c4bb8253 100644
--- a/tests/utils/test_image_utils.py
+++ b/tests/utils/test_image_utils.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import codecs
 import os
 import tempfile
 import unittest
@@ -544,6 +545,23 @@ def test_load_img_base64(self):
 
         self.assertEqual(img_arr.shape, (64, 32, 3))
 
+    def test_load_img_base64_encoded_bytes(self):
+        try:
+            tmp_file = tempfile.mktemp()
+            with open(tmp_file, "wb") as f:
+                http_get(
+                    "https://huggingface.co/datasets/hf-internal-testing/dummy-base64-images/raw/main/image_2.txt", f
+                )
+
+            with codecs.open(tmp_file, encoding="unicode_escape") as b64:
+                img = load_image(b64.read())
+                img_arr = np.array(img)
+
+        finally:
+            os.remove(tmp_file)
+
+        self.assertEqual(img_arr.shape, (256, 256, 3))
+
     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