datago 2025.10.2__tar.gz → 2025.12.1__tar.gz

{datago-2025.10.2 → datago-2025.12.1}/Cargo.lock

@@ -111,6 +111,16 @@ version = "0.7.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
 
+[[package]]
+name = "assert-json-diff"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12"
+dependencies = [
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "async-channel"
 version = "1.9.0"
@@ -464,7 +474,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d067ad48b8650848b989a59a86c6c36a995d02d2bf778d45c3c5d57bc2718f02"
 dependencies = [
  "smallvec",
- "target-lexicon 0.12.16",
+ "target-lexicon",
 ]
 
 [[package]]
@@ -613,7 +623,7 @@ dependencies = [
 
 [[package]]
 name = "datago"
-version = "2025.10.2"
+version = "2025.12.1"
 dependencies = [
  "async-compression",
  "async-tar",
@@ -644,8 +654,27 @@ dependencies = [
  "tokio-util",
  "url",
  "walkdir",
+ "wiremock",
+]
+
+[[package]]
+name = "deadpool"
+version = "0.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fb84100978c1c7b37f09ed3ce3e5f843af02c2a2c431bae5b19230dad2c1b490"
+dependencies = [
+ "async-trait",
+ "deadpool-runtime",
+ "num_cpus",
+ "tokio",
 ]
 
+[[package]]
+name = "deadpool-runtime"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b"
+
 [[package]]
 name = "dirs-next"
 version = "2.0.0"
@@ -1118,6 +1147,12 @@ version = "1.10.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
 
+[[package]]
+name = "httpdate"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9"
+
 [[package]]
 name = "hyper"
 version = "1.6.0"
@@ -1131,6 +1166,7 @@ dependencies = [
  "http",
  "http-body",
  "httparse",
+ "httpdate",
  "itoa",
  "pin-project-lite",
  "smallvec",
@@ -2026,9 +2062,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3"
-version = "0.24.1"
+version = "0.22.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "17da310086b068fbdcefbba30aeb3721d5bb9af8db4987d6735b2183ca567229"
+checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884"
 dependencies = [
  "cfg-if",
  "indoc",
@@ -2044,19 +2080,19 @@ dependencies = [
 
 [[package]]
 name = "pyo3-build-config"
-version = "0.24.1"
+version = "0.22.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e27165889bd793000a098bb966adc4300c312497ea25cf7a690a9f0ac5aa5fc1"
+checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38"
 dependencies = [
  "once_cell",
- "target-lexicon 0.13.2",
+ "target-lexicon",
 ]
 
 [[package]]
 name = "pyo3-ffi"
-version = "0.24.1"
+version = "0.22.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "05280526e1dbf6b420062f3ef228b78c0c54ba94e157f5cb724a609d0f2faabc"
+checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636"
 dependencies = [
  "libc",
  "pyo3-build-config",
@@ -2064,9 +2100,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros"
-version = "0.24.1"
+version = "0.22.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5c3ce5686aa4d3f63359a5100c62a127c9f15e8398e5fdeb5deef1fed5cd5f44"
+checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453"
 dependencies = [
  "proc-macro2",
  "pyo3-macros-backend",
@@ -2076,9 +2112,9 @@ dependencies = [
 
 [[package]]
 name = "pyo3-macros-backend"
-version = "0.24.1"
+version = "0.22.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f4cf6faa0cbfb0ed08e89beb8103ae9724eb4750e3a78084ba4017cbe94f3855"
+checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe"
 dependencies = [
  "heck",
  "proc-macro2",
@@ -2753,12 +2789,6 @@ version = "0.12.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
 
-[[package]]
-name = "target-lexicon"
-version = "0.13.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e502f78cdbb8ba4718f566c418c52bc729126ffd16baee5baa718cf25dd5a69a"
-
 [[package]]
 name = "tempfile"
 version = "3.20.0"
@@ -3454,6 +3484,30 @@ dependencies = [
  "memchr",
 ]
 
+[[package]]
+name = "wiremock"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a2b8b99d4cdbf36b239a9532e31fe4fb8acc38d1897c1761e161550a7dc78e6a"
+dependencies = [
+ "assert-json-diff",
+ "async-trait",
+ "base64",
+ "deadpool",
+ "futures",
+ "http",
+ "http-body-util",
+ "hyper",
+ "hyper-util",
+ "log",
+ "once_cell",
+ "regex",
+ "serde",
+ "serde_json",
+ "tokio",
+ "url",
+]
+
 [[package]]
 name = "wit-bindgen-rt"
 version = "0.33.0"

{datago-2025.10.2 → datago-2025.12.1}/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "datago"
 edition = "2021"
-version = "2025.10.2"
+version = "2025.12.1"
 readme = "README.md"
 
 [lib]
@@ -24,7 +24,7 @@ kanal = "0.1"
 clap = { version = "4.5.27", features = ["derive"] }
 tokio = { version = "1.43.1", features = ["rt-multi-thread", "macros"] }
 prettytable-rs = "0.10.0"
-pyo3 = { version = "0.24.1", features = ["extension-module"] }
+pyo3 = { version = "0.22", features = ["extension-module"] }
 threadpool = "1.8.1"
 openssl = { version = "0.10", features = ["vendored"] }
 walkdir = "2.5.0"
@@ -46,6 +46,7 @@ fast_image_resize = { version ="5.1.3", features=["image"]}
 
 [dev-dependencies]
 tempfile = "3.13.0"
+wiremock = "0.6.0"
 
 [profile.release]
 opt-level = 3

{datago-2025.10.2 → datago-2025.12.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datago
-Version: 2025.10.2
+Version: 2025.12.1
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
@@ -8,7 +8,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
 License-File: LICENSE
 Summary: A high performance dataloader for Python, written in Rust
-Author: Benjamin Lefaudeux
+Author: Benjamin Lefaudeux, Roman Frigg
 Author-email: Photoroom <team@photoroom.com>
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM

{datago-2025.10.2 → datago-2025.12.1}/pyproject.toml

@@ -2,7 +2,7 @@
 name = "datago"
 dynamic = ["version"]
 authors = [
-  { name = "Benjamin Lefaudeux" },
+  { name = "Benjamin Lefaudeux, Roman Frigg" },
   { name = "Photoroom", email = "team@photoroom.com" }
 ]
 description = "A high performance dataloader for Python, written in Rust"

{datago-2025.10.2 → datago-2025.12.1}/python/benchmark_db.py

@@ -6,7 +6,7 @@ import typer
 from benchmark_defaults import IMAGE_CONFIG
 from datago import DatagoClient  # type: ignore
 from PIL import Image
-from raw_types import raw_array_to_numpy, raw_array_to_pil_image
+from raw_types import raw_array_to_numpy
 from tqdm import tqdm
 
 
@@ -58,21 +58,7 @@ def benchmark(
     for _ in tqdm(range(limit), dynamic_ncols=True):
         sample = client.get_sample()
         if sample.id:
-            # Bring the masks and image to PIL
-            if hasattr(sample, "image"):
-                img = raw_array_to_pil_image(sample.image)
-
-            if hasattr(sample, "masks"):
-                for _, mask_buffer in sample.masks.items():
-                    mask = raw_array_to_pil_image(mask_buffer)
-
-            if (
-                hasattr(sample, "additional_images")
-                and "masked_image" in sample.additional_images
-            ):
-                masked_image = raw_array_to_pil_image(
-                    sample.AdditionalImages["masked_image"]
-                )
+            # Images are already PIL by default
 
             # Bring the latents to numpy
             if hasattr(sample, "latents"):

{datago-2025.10.2 → datago-2025.12.1}/python/benchmark_filesystem.py

@@ -9,26 +9,34 @@ from tqdm import tqdm
 
 
 def benchmark(
-    root_path: str = typer.Option(os.getenv("DATAGO_TEST_FILESYSTEM", ""), help="The source to test out"),
+    root_path: str = typer.Option(
+        os.getenv("DATAGO_TEST_FILESYSTEM", ""), help="The source to test out"
+    ),
     limit: int = typer.Option(2000, help="The number of samples to test on"),
-    crop_and_resize: bool = typer.Option(False, help="Crop and resize the images on the fly"),
+    crop_and_resize: bool = typer.Option(
+        False, help="Crop and resize the images on the fly"
+    ),
     compare_torch: bool = typer.Option(True, help="Compare against torch dataloader"),
     num_workers: int = typer.Option(os.cpu_count(), help="Number of workers to use"),
     sweep: bool = typer.Option(False, help="Sweep over the number of workers"),
 ):
     if sweep:
-        results = {}
-        for num_workers in range(2, (os.cpu_count() or 2), 16):
-            results[num_workers] = benchmark(root_path, limit, crop_and_resize, compare_torch, num_workers, False)
+        results_sweep = {}
+        for num_workers in range(2, (os.cpu_count() * 2 or 2), 2):
+            results_sweep[num_workers] = benchmark(
+                root_path, limit, crop_and_resize, compare_torch, num_workers, False
+            )
 
         # Save results to a json file
 
         with open("benchmark_results_filesystem.json", "w") as f:
-            json.dump(results, f, indent=2)
+            json.dump(results_sweep, f, indent=2)
 
-        return results
+        return results_sweep
 
-    print(f"Running benchmark for {root_path} - {limit} samples - {num_workers} workers")
+    print(
+        f"Running benchmark for {root_path} - {limit} samples - {num_workers} workers"
+    )
 
     # This setting is not exposed in the config, but an env variable can be used instead
     os.environ["DATAGO_MAX_TASKS"] = str(num_workers)
@@ -59,6 +67,11 @@ def benchmark(
     for sample in tqdm(datago_dataset, desc="Datago", dynamic_ncols=True):
         assert sample["id"] != ""
         img = sample["image"]
+
+        if count < limit - 1:
+            del img
+            img = None  # Help with memory pressure
+
         count += 1
 
     assert count == limit, f"Expected {limit} samples, got {count}"
@@ -80,7 +93,9 @@ def benchmark(
         transform = (
             transforms.Compose(
                 [
-                    transforms.Resize((1024, 1024), interpolation=transforms.InterpolationMode.LANCZOS),
+                    transforms.Resize(
+                        (1024, 1024), interpolation=transforms.InterpolationMode.LANCZOS
+                    ),
                 ]
             )
             if crop_and_resize
@@ -88,7 +103,9 @@ def benchmark(
         )
 
         # Create the ImageFolder dataset
-        dataset = datasets.ImageFolder(root=root_path, transform=transform, allow_empty=True)
+        dataset = datasets.ImageFolder(
+            root=root_path, transform=transform, allow_empty=True
+        )
 
         # Create a DataLoader to allow for multiple workers
         # Use available CPU count for num_workers
@@ -107,6 +124,8 @@ def benchmark(
             n_images += len(batch)
             if n_images > limit:
                 break
+
+            del batch  # Help with memory pressure, same as above
         fps = n_images / (time.time() - start)
         results["torch"] = {"fps": fps, "count": n_images}
         print(f"Torch - FPS {fps:.2f} - workers {num_workers}")

{datago-2025.10.2 → datago-2025.12.1}/python/dataset.py

@@ -1,7 +1,7 @@
 from datago import DatagoClient, initialize_logging
 import json
 from typing import Dict, Any
-from raw_types import raw_array_to_pil_image, raw_array_to_numpy
+from raw_types import raw_array_to_numpy
 
 
 class DatagoIterDataset:
@@ -29,11 +29,13 @@ class DatagoIterDataset:
             return json.loads(item)
 
         if isinstance(item, dict):
-            # recurvisely convert the dictionary
+            # recursively convert the dictionary
             return {k: DatagoIterDataset.to_python_types(v, k) for k, v in item.items()}
 
         elif "image" in key:
-            return raw_array_to_pil_image(item)
+            # The Rust-side returns PythonImagePayload objects that are callable
+            # Call them to get the actual PIL image
+            return item()
         elif "latent" in key:
             return raw_array_to_numpy(item)
 

datago-2025.12.1/python/raw_types.py

@@ -0,0 +1,40 @@
+from PIL import Image
+from typing import Optional, Union
+import numpy as np
+from datago import ImagePayload
+
+
+def raw_array_to_numpy(raw_array: ImagePayload) -> Optional[np.ndarray]:
+    if len(raw_array.data) == 0:
+        return None
+
+    # Generic numpy-serialized array
+    try:
+        return np.load(raw_array.data, allow_pickle=False)
+    except ValueError:
+        # Do not try to handle these, return None and we'll handle it in the caller
+        print("Could not deserialize numpy array")
+        return None
+
+
+def decode_image_payload(payload: ImagePayload) -> Image.Image:
+    """
+    Decode an ImagePayload (encoded image) into a PIL Image.
+    This is the proper way to decode encoded images for API users.
+    """
+    import io
+
+    return Image.open(io.BytesIO(payload.data))
+
+
+def get_image_mode(image_or_payload: Union[ImagePayload, Image.Image]) -> str:
+    """
+    Helper function to get the mode of an image, whether it's a PIL Image or ImagePayload.
+    For ImagePayload objects (encoded images), we need to decode them first.
+    """
+    if hasattr(image_or_payload, "mode"):
+        # It's a PIL Image
+        return image_or_payload.mode
+    else:
+        # It's an ImagePayload (encoded image), decode it first
+        return decode_image_payload(image_or_payload).mode

{datago-2025.10.2 → datago-2025.12.1}/python/test_datago_client.py

@@ -115,7 +115,13 @@ class TestDatagoClient:
                 assert sample.source == "filesystem"
                 assert sample.image.width > 0
                 assert sample.image.height > 0
-                assert len(sample.image.data) > 0
+
+                # Check the payload path
+                payload = sample.image.get_payload()
+                assert len(payload.data) > 0
+                assert payload.width == sample.image.width
+                assert payload.height == sample.image.height
+                assert payload.channels == 3
 
     def test_client_with_image_transformations(self):
         """Test client with image transformation configuration."""
@@ -149,7 +155,7 @@ class TestDatagoClient:
             assert sample is not None
             assert sample.image.width <= 64
             assert sample.image.height <= 64
-            assert sample.image.channels == 3  # RGB8
+            assert sample.image.mode == "RGB"
 
     def test_client_with_image_encoding(self):
         """Test client with image encoding enabled."""
@@ -177,8 +183,7 @@ class TestDatagoClient:
             sample = client.get_sample()
 
             assert sample is not None
-            assert sample.image.channels == -1  # Encoded images have channels = -1
-            assert len(sample.image.data) > 0
+            assert sample.image.mode == "RGB"
 
     def test_random_sampling(self):
         """Test that random sampling produces different results."""

{datago-2025.10.2 → datago-2025.12.1}/python/test_datago_db.py

@@ -3,7 +3,9 @@ import pytest
 import os
 import json
 
-from raw_types import raw_array_to_pil_image, raw_array_to_numpy, get_image_mode, decode_image_payload
+from raw_types import (
+    decode_image_payload,
+)
 from dataset import DatagoIterDataset
 
 
@@ -62,9 +64,11 @@ def test_caption_and_image():
         assert img.height > 0
         assert img.width > 0
 
-        assert img.height <= img.original_height
-        assert img.width <= img.original_width
-        assert img.channels == channels
+        payload = img.get_payload()
+        assert img.height <= payload.original_height
+        assert img.width <= payload.original_width
+        assert img.mode == "RGB" if channels == 3 else "L"
+        assert payload.channels == channels
 
     for i, sample in enumerate(dataset):
         assert sample.source != ""
@@ -80,14 +84,11 @@ def test_caption_and_image():
         check_image(sample.masks["segmentation_mask"], 1)
 
         # Check the image decoding
-        assert raw_array_to_pil_image(sample.image).mode == "RGB", "Image should be RGB"
-        assert (
-            raw_array_to_pil_image(sample.additional_images["masked_image"]).mode
-            == "RGB"
-        ), "Image should be RGB"
-        assert raw_array_to_pil_image(sample.masks["segmentation_mask"]).mode == "L", (
-            "Mask should be L"
+        assert sample.image.mode == "RGB", "Image should be RGB"
+        assert sample.additional_images["masked_image"].mode == "RGB", (
+            "Image should be RGB"
         )
+        assert sample.masks["segmentation_mask"].mode == "L", "Mask should be L"
 
         if i > N_SAMPLES:
             break
@@ -146,36 +147,13 @@ def test_jpeg_compression():
 
     sample = next(iter(dataset))
 
-    # When images are encoded, channels is set to -1 to signal encoded format
-    assert sample.image.channels == -1, "Image should be encoded (channels == -1)"
-    assert (
-        sample.additional_images["masked_image"].channels == -1
-    ), "Additional image should be encoded"
-    assert (
-        sample.masks["segmentation_mask"].channels == -1
-    ), "Mask should be encoded"
-
-    # Test that raw_array_to_pil_image returns ImagePayload for encoded images
-    image_result = raw_array_to_pil_image(sample.image)
-    assert not hasattr(image_result, 'mode'), "Should return ImagePayload, not PIL Image"
-    assert hasattr(image_result, 'data'), "Should have data attribute"
-    assert hasattr(image_result, 'channels'), "Should have channels attribute"
-    assert image_result.channels == -1, "Should be encoded ImagePayload"
-
-    # Test proper decoding using decode_image_payload
-    decoded_image = decode_image_payload(image_result)
-    assert hasattr(decoded_image, 'mode'), "Decoded image should be PIL Image"
-    assert decoded_image.mode == "RGB", "Image should decode to RGB"
-    assert decoded_image.size == (sample.image.width, sample.image.height), "Size should match"
-
-    # Test additional images and masks
-    additional_result = raw_array_to_pil_image(sample.additional_images["masked_image"])
-    decoded_additional = decode_image_payload(additional_result)
-    assert decoded_additional.mode == "RGB", "Additional image should decode to RGB"
-
-    mask_result = raw_array_to_pil_image(sample.masks["segmentation_mask"])
-    decoded_mask = decode_image_payload(mask_result)
-    assert decoded_mask.mode == "L", "Mask should decode to L"
+    # Check that the image is properly accessible through PIL, but that it is encoded
+    assert sample.image.mode == "RGB", "Image should be RGB"
+    assert sample.additional_images["masked_image"].mode == "RGB", "Image should be RGB"
+    assert sample.masks["segmentation_mask"].mode == "L", "Mask should be L"
+
+    # Check that the image is encoded, as JPG PIL
+    # TODO: @blefaudeux
 
 
 def test_png_compression():
@@ -185,23 +163,11 @@ def test_png_compression():
     # Don't specify encode_format - should default to PNG
     dataset = DatagoIterDataset(client_config, return_python_types=False)
 
-    sample = next(iter(dataset))
-
-    # When images are encoded, channels is set to -1 to signal encoded format
-    assert sample.image.channels == -1, "Image should be encoded (channels == -1)"
-
-    # Test that raw_array_to_pil_image returns ImagePayload for encoded images
-    image_result = raw_array_to_pil_image(sample.image)
-    assert not hasattr(image_result, 'mode'), "Should return ImagePayload, not PIL Image"
-    assert hasattr(image_result, 'data'), "Should have data attribute"
-    assert hasattr(image_result, 'channels'), "Should have channels attribute"
-    assert image_result.channels == -1, "Should be encoded ImagePayload"
+    _sample = next(iter(dataset))
 
-    # Test proper decoding using decode_image_payload
-    decoded_image = decode_image_payload(image_result)
-    assert hasattr(decoded_image, 'mode'), "Decoded image should be PIL Image"
-    assert decoded_image.mode == "RGB", "Image should decode to RGB"
-    assert decoded_image.size == (sample.image.width, sample.image.height), "Size should match"
+    # Check that the image is encoded, as JPG PIL
+    # TODO: @blefaudeux
+    # same as above
 
 
 def test_original_image():
@@ -212,14 +178,9 @@ def test_original_image():
     dataset = DatagoIterDataset(client_config, return_python_types=False)
 
     sample = next(iter(dataset))
-
-    assert raw_array_to_pil_image(sample.image).mode == "RGB", "Image should be RGB"
-    assert (
-        raw_array_to_pil_image(sample.additional_images["masked_image"]).mode == "RGB"
-    ), "Image should be RGB"
-    assert raw_array_to_pil_image(sample.masks["segmentation_mask"]).mode == "L", (
-        "Mask should be L"
-    )
+    payload = sample.image.get_payload()
+    assert payload.original_height == payload.height == sample.image.height
+    assert payload.original_width == payload.width == sample.image.width
 
 
 def test_duplicate_state():

{datago-2025.10.2 → datago-2025.12.1}/python/test_datago_edge_cases.py

@@ -90,11 +90,13 @@ class TestDatagoEdgeCases:
             sample = client.get_sample()
 
             assert sample is not None
-            assert sample.image.original_width == 2000
-            assert sample.image.original_height == 2000
+            image_payload = sample.image.get_payload()
+            assert image_payload is not None
+            assert image_payload.original_width == 2000
+            assert image_payload.original_height == 2000
             # Should be resized
-            assert sample.image.width <= 512
-            assert sample.image.height <= 512
+            assert image_payload.width <= 512
+            assert image_payload.height <= 512
 
     def test_very_small_images(self):
         """Test handling of very small images."""

{datago-2025.10.2 → datago-2025.12.1}/python/test_datago_filesystem.py

@@ -54,21 +54,23 @@ def test_get_sample_filesystem(pre_encode_images: bool, rgb16: bool, rgba: bool)
 
         client = DatagoClient(json.dumps(client_config))
         count = 0
-        for i in range(limit):
+        for _ in range(limit):
             data = client.get_sample()
             if not data:
                 break
             count += 1
             assert data.id != ""
-            assert data.image.width == 100
-            assert data.image.height == 100
+
+            image_payload = data.image.get_payload()
+            assert image_payload.width == 100
+            assert image_payload.height == 100
 
             if rgb16:
-                assert data.image.bit_depth == 8
+                assert image_payload.bit_depth == 8
 
             # Open the image in python scope and check properties
             if pre_encode_images:
-                test_image = Image.open(BytesIO(data.image.data))
+                test_image = Image.open(BytesIO(bytes(image_payload.data)))
                 assert test_image.width == 100
                 assert test_image.height == 100
                 assert test_image.mode == "RGB"