Handle recoverable truncated image payloads

tests/test_smp_vlm.py

@@ -0,0 +1,33 @@
+import base64
+import importlib.util
+import io
+from pathlib import Path
+
+from PIL import Image, ImageFile
+
+
+def _load_vlm_module():
+    module_path = Path(__file__).resolve().parents[1] / "vlmeval" / "smp" / "vlm.py"
+    spec = importlib.util.spec_from_file_location("vlmeval_smp_vlm_under_test", module_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _truncated_png_base64():
+    image = Image.new("RGB", (20, 20), "white")
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()[:-20]).decode("utf-8")
+
+
+def test_decode_base64_to_image_file_allows_recoverable_truncated_png(tmp_path):
+    ImageFile.LOAD_TRUNCATED_IMAGES = False
+    vlm = _load_vlm_module()
+    output_path = tmp_path / "recovered.png"
+
+    vlm.decode_base64_to_image_file(_truncated_png_base64(), output_path)
+
+    with Image.open(output_path) as recovered:
+        recovered.load()
+        assert recovered.size == (20, 20)

vlmeval/smp/vlm.py

@@ -4,9 +4,11 @@
 import os.path as osp
 
 import pandas as pd
-from PIL import Image
+from PIL import Image, ImageFile
 
 Image.MAX_IMAGE_PIXELS = 1e9
+# Some benchmark TSVs include recoverable truncated image payloads.
+ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 
 def rescale_img(img, tgt=None):