From 415c2f7d6bc218ac2d328fe2fbcab86f195aba29 Mon Sep 17 00:00:00 2001
From: Nico Kemnitz <nico@zetta.ai>
Date: Fri, 20 Feb 2026 12:36:58 +0100
Subject: [PATCH] feat(zettasets): Fall back to zettasets images

---
 deepem/data/dataset/multi_zettaset.py | 29 +++++++++++++++++++--------
 deepem/data/dataset/zettaset.py       | 27 ++++++++++++++++++-------
 2 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/deepem/data/dataset/multi_zettaset.py b/deepem/data/dataset/multi_zettaset.py
index 4b31734..2378329 100644
--- a/deepem/data/dataset/multi_zettaset.py
+++ b/deepem/data/dataset/multi_zettaset.py
@@ -5,6 +5,7 @@
 import re
 
 from cloudvolume import CloudVolume, Bbox
+from cloudvolume.exceptions import InfoUnavailableError
 
 from zettasets.dataset import Dataset as Zettaset
 from zettasets.sample import Sample
@@ -199,14 +200,26 @@ def convert_array(arr: ArrayLike) -> np.ndarray:
         else (0, 0, 0)
     )
     image_bbox = Bbox(bbox.minpt - xyz_padding, bbox.maxpt + xyz_padding)
-    
-    # Image
-    vol = CloudVolume(  # pylint: disable=unsubscriptable-object
-        sample.src_image_path,
-        mip=resolution,
-        fill_missing=True,
-        bounded=False,
-    )[image_bbox.to_slices()]
+
+    # Image: try src_image_path first for backward compatibility,
+    #        fall back to sample's own image volume
+    try:
+        if sample.src_image_path is None:
+            raise ValueError("No src_image_path available")
+        vol = CloudVolume(  # pylint: disable=unsubscriptable-object
+            sample.src_image_path,
+            mip=resolution,
+            fill_missing=True,
+            bounded=False,
+        )[image_bbox.to_slices()]
+    except (InfoUnavailableError, ValueError) as e:
+        print(f"\tWarning: src_image_path failed ({e}), falling back to sample image volume")
+        vol = CloudVolume(  # pylint: disable=unsubscriptable-object
+            sample.volumes["image"].cloudpath,
+            mip=resolution,
+            fill_missing=True,
+            bounded=False,
+        )[image_bbox.to_slices()]
     dset["input"] = convert_array(vol) / 255.0
     print(f"\tinput: {dset['input'].shape}")
 
diff --git a/deepem/data/dataset/zettaset.py b/deepem/data/dataset/zettaset.py
index 8151b9c..004428d 100644
--- a/deepem/data/dataset/zettaset.py
+++ b/deepem/data/dataset/zettaset.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from cloudvolume import CloudVolume, Bbox
+from cloudvolume.exceptions import InfoUnavailableError
 import numpy as np
 from numpy.typing import ArrayLike
 import re
@@ -90,13 +91,25 @@ def convert_array(arr: ArrayLike) -> np.ndarray:
     )
     image_bbox = Bbox(bbox.minpt - xyz_padding, bbox.maxpt + xyz_padding)
 
-    # Image
-    vol = CloudVolume(  # pylint: disable=unsubscriptable-object
-        sample.src_image_path,
-        mip=resolution,
-        fill_missing=True,
-        bounded=False,
-    )[image_bbox.to_slices()]
+    # Image: try src_image_path first for backward compatibility,
+    #        fall back to sample's own image volume
+    try:
+        if sample.src_image_path is None:
+            raise ValueError("No src_image_path available")
+        vol = CloudVolume(  # pylint: disable=unsubscriptable-object
+            sample.src_image_path,
+            mip=resolution,
+            fill_missing=True,
+            bounded=False,
+        )[image_bbox.to_slices()]
+    except (InfoUnavailableError, ValueError) as e:
+        print(f"\tWarning: src_image_path failed ({e}), falling back to sample image volume")
+        vol = CloudVolume(  # pylint: disable=unsubscriptable-object
+            sample.volumes["image"].cloudpath,
+            mip=resolution,
+            fill_missing=True,
+            bounded=False,
+        )[image_bbox.to_slices()]
     dset["input"] = convert_array(vol)
     dset["input"] = (dset["input"] / 255.).astype('float32')
     print(f"input: {dset['input'].shape}")