[fix] saves partial z-stack during mid error

Root cause of the original bug: a camera communication error could return
an image with the wrong shape.  On NumPy < 1.24, numpy.array() on
mixed-shape arrays silently produces an object-dtype array, which cannot
be written to TIFF (WriteDirectory() → AssertionError: 0).

Author: K4rishma

src/odemis/acq/acqmng.py

@@ -166,6 +166,7 @@ def run(self):
         """
         remaining_t = self.estimate_total_duration()
         acquired_data = []
+        exp = None
         # iterate through streams
         for stream in self._streams:
             zstack = []
@@ -201,6 +202,7 @@ def run(self):
 
             else:
                 # for each stream, iterate through zlevels
+                z_exp = None  # tracks failure within this stream's z-stack
                 for i, z in enumerate(self._zlevels[stream]):
                     # move the focuser
                     self._actuator_f = stream.focuser.moveAbs({"z": z})
@@ -218,9 +220,10 @@ def run(self):
                     try:
                         # acquire this single stream, and get the data
                         self._single_acqui_f = acquire([stream], self._settings_obs)
-                        data, exp = self._single_acqui_f.result()
-                        if exp:
-                            return acquired_data, exp
+                        data, z_exp = self._single_acqui_f.result()
+                        if z_exp:
+                            # exit the z-loop; partial z-stack will be assembled below
+                            break
                         # check if cancellation happened while the acquiring future is working
                         if self._future_state == CANCELLED:
                             raise CancelledError()
@@ -234,20 +237,57 @@ def run(self):
                     except CancelledError:
                         raise
                     except Exception as e:
-                        logging.exception("The acquisition failed at the %s-th zlevel of the stream %s, because %s" % (
-                        i + 1, stream, e))
-                        # TODO handle zstack assembling in case of error
-                        return acquired_data, e
+                        logging.exception("The acquisition failed at the %s-th zlevel of the stream %s",
+                                          i + 1, stream.name.value)
+                        z_exp = e
+                        break  # exit the z-loop; partial z-stack will be assembled below
 
                     # only if there is data acquired
                     if data:
-                        zstack.append(data[0])
+                        da = data[0]
+                        # Validate spatial shape consistency across z-levels.
+                        # A camera communication error can return an image of the wrong size.
+                        # With NumPy < 1.24, numpy.array() on mixed-shape images silently
+                        # creates an object-dtype array, which cannot be written to TIFF.
+                        if zstack and da.shape[-2:] != zstack[0].shape[-2:]:
+                            logging.error(
+                                "Z-level %d of stream %s has shape %s, inconsistent with "
+                                "first z-level shape %s; stopping z-stack acquisition",
+                                i, stream.name.value, da.shape[-2:], zstack[0].shape[-2:]
+                            )
+                            z_exp = ValueError(
+                                "Z-level %d image shape %s is inconsistent with first "
+                                "z-level shape %s" % (i, da.shape[-2:], zstack[0].shape[-2:])
+                            )
+                            break
+                        zstack.append(da)
+
                     # update the remaining time
                     remaining_t -= stream.estimateAcquisitionTime()
                     self._main_future.set_end_time(time.time() + remaining_t)
 
-                zcube = assembleZCube(zstack, self._zlevels[stream])
-                acquired_data.append(zcube)
+                # Assemble whatever z-levels were acquired (partial or complete).
+                # This also handles the TODO: save partial z-stack data on failure.
+                if zstack:
+                    if len(zstack) < len(self._zlevels[stream]):
+                        logging.warning(
+                            "Partial z-stack for stream %s: assembling %d of %d levels",
+                            stream.name.value, len(zstack), len(self._zlevels[stream])
+                        )
+                    try:
+                        zcube = assembleZCube(zstack, self._zlevels[stream][:len(zstack)])
+                        acquired_data.append(zcube)
+                    except Exception as e:
+                        logging.exception("Failed to assemble z-stack for stream %s",
+                                          stream.name.value)
+                        if z_exp is None:
+                            z_exp = e
+                elif z_exp is None:
+                    logging.warning("All z-levels for stream %s returned empty data, skipping",
+                                    stream.name.value)
+
+                if z_exp is not None:
+                    return acquired_data, z_exp
 
         # state that the future has finished
         with self._future_lock:

src/odemis/acq/test/acq_test.py

@@ -23,6 +23,7 @@
 import os
 import time
 import unittest
+from concurrent.futures import Future
 from concurrent.futures._base import CancelledError
 from unittest import mock
 
@@ -33,13 +34,14 @@
 import odemis.acq.stream as stream
 from odemis import model
 from odemis.acq import acqmng
-from odemis.acq.acqmng import SettingsObserver, acquireZStack
+from odemis.acq.acqmng import SettingsObserver, ZStackAcquisitionTask, acquireZStack
 from odemis.acq.leech import ProbeCurrentAcquirer
 from odemis.acq.move import MicroscopePostureManager, FM_IMAGING, SEM_IMAGING, LOADING
 from odemis.driver import xt_client
 from odemis.driver.test.xt_client_test import CONFIG_FIB_SEM, CONFIG_FIB_SCANNER, CONFIG_DETECTOR
 from odemis.util import testing
 from odemis.util.comp import generate_zlevels
+from odemis.dataio import tiff
 
 logging.getLogger().setLevel(logging.DEBUG)
 
@@ -856,5 +858,153 @@ def test_settings_observer_metadata_with_zstack(self):
         self.assertEqual(data[0].metadata[model.MD_EXTRA_SETTINGS]
                          ["Camera"]["exposureTime"], [0.023, "s"])
 
+
+def _make_sim_future(result=None):
+    """
+    Return a stdlib Future already completed with result.
+
+    :param result: value to store in the future
+    :return: completed concurrent.futures.Future
+    """
+    f = Future()
+    f.set_result(result)
+    return f
+
+
+def _make_sim_data_array(shape=(64, 64), dtype=numpy.uint16):
+    """
+    Return a minimal 2-D DataArray suitable as a z-level image.
+
+    :param shape: 2-tuple (height, width)
+    :param dtype: NumPy dtype for the pixel data
+    :return: model.DataArray with pixel-size and position metadata
+    """
+    md = {
+        model.MD_DIMS: "YX",
+        model.MD_PIXEL_SIZE: (1e-7, 1e-7),
+        model.MD_POS: (0.0, 0.0),
+    }
+    return model.DataArray(numpy.zeros(shape, dtype=dtype), md)
+
+
+def _make_sim_stream(name="mock_stream"):
+    """
+    Build a MagicMock that satisfies the interface used by ZStackAcquisitionTask.
+
+    :param name: human-readable name for the stream mock
+    :return: unittest.mock.MagicMock mimicking a Stream
+    """
+    s = mock.MagicMock()
+    s.name.value = name
+    s.estimateAcquisitionTime.return_value = 0.0
+    s.focuser.moveAbs.return_value = _make_sim_future(None)
+    return s
+
+
+def _make_sim_task(stream_mock, zlevels):
+    """
+    Construct a ZStackAcquisitionTask with a mock ProgressiveFuture.
+
+    Both guessActuatorMoveDuration (called in __init__) and
+    estimate_total_duration (called inside run()) are patched to avoid
+    the need for real actuator hardware.
+
+    :param stream_mock: mock Stream object
+    :param zlevels: dict mapping stream_mock to list of z positions
+    :return: (task, mock_future) tuple ready to call task.run() on
+    """
+    future = mock.MagicMock()
+    with mock.patch("odemis.acq.acqmng.guessActuatorMoveDuration", return_value=0.0):
+        task = ZStackAcquisitionTask(future, [stream_mock], zlevels, settings_obs=None)
+    task.estimate_total_duration = mock.MagicMock(return_value=1.0)
+    return task, future
+
+
+class TestZStackPartialFailureSim(unittest.TestCase):
+    """
+    Simulation tests (no hardware) for the fix that saves partial z-stack data
+    when a camera error occurs during an acquisition.
+
+    Root cause of the original bug: a camera communication error could return
+    an image with the wrong shape.  On NumPy < 1.24, numpy.array() on
+    mixed-shape arrays silently produces an object-dtype array, which cannot
+    be written to TIFF (WriteDirectory() → AssertionError: 0).
+
+    Shape validation in assembleZCube() and ZStackAcquisitionTask.run(),
+    plus partial z-stack assembly is implemented instead of discarding data on failure.
+    """
+
+    def test_full_success_returns_zcube(self):
+        """
+        When all z-levels succeed, run() returns a single ZYX DataArray and no exception.
+        """
+        n = 3
+        zlevels_list = [i * 1e-6 for i in range(n)]
+        s = _make_sim_stream("fluo")
+        task, _ = _make_sim_task(s, {s: zlevels_list})
+
+        good_img = _make_sim_data_array((64, 64))
+        acq_futures = [_make_sim_future(([good_img], None)) for _ in range(n)]
+
+        with mock.patch("odemis.acq.acqmng.acquire", side_effect=acq_futures):
+            data, exp = task.run()
+
+        self.assertIsNone(exp)
+        self.assertEqual(len(data), 1)
+        self.assertEqual(data[0].shape, (n, 64, 64))
+        self.assertNotEqual(data[0].dtype, object)
+
+    def test_wrong_shape_mid_zstack_saves_partial(self):
+        """
+        When a z-level image has a wrong spatial shape (truncated camera read),
+        run() must stop, assemble only the valid z-levels, and report the error.
+
+        This is the primary regression test for the TIFF-crash bug.
+        """
+        zlevels_list = [0.0e-6, 1.0e-6, 2.0e-6]
+        s = _make_sim_stream("fluo")
+        task, _ = _make_sim_task(s, {s: zlevels_list})
+
+        good_img = _make_sim_data_array((64, 64))
+        bad_img = _make_sim_data_array((32, 64))  # truncated height — simulates camera error
+
+        acq_futures = [
+            _make_sim_future(([good_img], None)),
+            _make_sim_future(([good_img], None)),
+            _make_sim_future(([bad_img],  None)),  # 3rd level: wrong shape
+        ]
+
+        with mock.patch("odemis.acq.acqmng.acquire", side_effect=acq_futures):
+            data, exp = task.run()
+
+        # Partial data must be saved
+        self.assertEqual(len(data), 1)
+        zcube = data[0]
+        # Must NOT be an object-dtype array (the original bug)
+        self.assertNotEqual(zcube.dtype, object)
+        # Only the 2 valid levels are included
+        self.assertEqual(zcube.shape, (2, 64, 64))
+        # Error must be reported
+        self.assertIsNotNone(exp)
+
+    def test_first_zlevel_fails_returns_empty_data(self):
+        """
+        When the very first z-level fails, no z-cube can be assembled.
+        run() must return an empty data list and the exception.
+        """
+        zlevels_list = [0.0e-6, 1.0e-6]
+        s = _make_sim_stream("fluo")
+        task, _ = _make_sim_task(s, {s: zlevels_list})
+
+        hw_error = IOError("Camera connection lost")
+
+        with mock.patch("odemis.acq.acqmng.acquire",
+                        return_value=_make_sim_future(([], hw_error))):
+            data, exp = task.run()
+
+        self.assertEqual(len(data), 0)
+        self.assertIs(exp, hw_error)
+
+
 if __name__ == "__main__":
     unittest.main()

src/odemis/gui/cont/acquisition/cryo_acq.py

@@ -536,9 +536,14 @@ def _on_export_data_done(self, future):
         """
         Called after exporting the data
         """
+        try:
+            data = future.result()
+        except Exception:
+            logging.exception("Failed to save acquired data to file")
+            self._reset_acquisition_gui(text="Failed to save data (see log panel).", state=ST_FAILED)
+            return
         self._reset_acquisition_gui(state=ST_FINISHED)
         self._update_acquisition_time()
-        data = future.result()
         self._display_acquired_data(data)
 
     def _create_cryo_filename(self, filename: str, acq_type: Optional[str] = None) -> str:

src/odemis/util/img.py

@@ -1286,7 +1286,22 @@ def assembleZCube(images, zlevels):
         :param images:  (list of DataArray of shape YX) list of z ordered images
         :param zlevels:  (list of float) list of focus positions
         :return: (DataArray of shape ZYX) the data array of the xyz cube
+        :raises ValueError: if images is empty or the images have inconsistent spatial shapes
         """
+        if not images:
+            raise ValueError("Cannot assemble z-cube from an empty image list")
+
+        # Validate that all images have the same spatial (YX) dimensions.
+        # With NumPy < 1.24, numpy.array() on a list of arrays with different shapes
+        # silently creates an object-dtype array, which cannot be written to TIFF.
+        ref_shape = images[0].shape[-2:]
+        for idx, im in enumerate(images[1:], start=1):
+            if im.shape[-2:] != ref_shape:
+                raise ValueError(
+                    "Z-stack image %d has shape %s, inconsistent with first image shape %s"
+                    % (idx, im.shape[-2:], ref_shape)
+                )
+
         # images is a list of 3 dim data arrays.
         # Will fail on purpose if the images contain more than 2 dimensions
         ret = numpy.array([im.reshape(im.shape[-2:]) for im in images])

src/odemis/util/test/img_test.py

@@ -2068,6 +2068,33 @@ def test_assemble_zcube_z_order(self):
         self.assertGreater(output_rev_z.metadata[model.MD_PIXEL_SIZE][2], 0)
         numpy.testing.assert_array_equal(output_da_after, output_rev_z)
 
+    def test_assemble_zcube_empty_list_raises(self):
+        """
+        assembleZCube() must raise ValueError when given an empty image list.
+        """
+        with self.assertRaises(ValueError):
+            img.assembleZCube([], [])
+
+    def test_assemble_zcube_inconsistent_shapes_raises(self):
+        """
+        assembleZCube() must raise ValueError when z-level images have different shapes.
+
+        On NumPy < 1.24, numpy.array() on mixed-shape arrays silently creates an
+        object-dtype array, which cannot be written to TIFF.  The fix detects this
+        early and raises explicitly.
+        """
+        images = [
+            model.DataArray(numpy.zeros(self.size, dtype=numpy.uint16), self.md),
+            model.DataArray(numpy.zeros(self.size, dtype=numpy.uint16), self.md),
+            model.DataArray(numpy.zeros((self.size[0] // 2, self.size[1]), dtype=numpy.uint16), self.md),
+        ]
+        zlevels = self.z_list[:3]
+
+        with self.assertRaises(ValueError) as ctx:
+            img.assembleZCube(images, zlevels)
+
+        self.assertIn("shape", str(ctx.exception).lower())
+
 
 class TestFloodFill(unittest.TestCase):