Unverified Commit c793b26f authored by amyeroberts's avatar amyeroberts Committed by GitHub
Browse files

load_image - decode b64encode and encodebytes strings (#30192)

* Decode b64encode and encodebytes strings

* Remove conditional encode -- image is always a string
parent e7d52a10
......@@ -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(
......
......@@ -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 @@ class LoadImageTester(unittest.TestCase):
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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment