[fix] saves partial z-stack during mid error#3466
Conversation
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).
There was a problem hiding this comment.
Pull request overview
This PR prevents TIFF export crashes caused by mixed-shape z-stack frames (notably on NumPy < 1.24 where numpy.array() can silently create object arrays), and ensures that when an acquisition fails mid z-stack, already-acquired valid planes are still assembled/saved instead of being discarded.
Changes:
- Add explicit validation in
assembleZCube()to reject empty inputs and inconsistent YX shapes with a clearValueError. - Update
ZStackAcquisitionTask.run()to detect wrong-shape frames during z acquisition, stop acquisition, and assemble/save the partial z-cube collected so far. - Improve GUI export callback robustness by handling exceptions from
future.result()and resetting the acquisition UI to a failed state.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
src/odemis/util/img.py |
Adds early validation in assembleZCube() to prevent mixed-shape/object-dtype cubes. |
src/odemis/util/test/img_test.py |
Adds regression tests for empty input and inconsistent-shape z-stack handling. |
src/odemis/acq/acqmng.py |
Implements partial z-stack assembly on mid-stack failure and adds shape consistency checks per z-level. |
src/odemis/gui/cont/acquisition/cryo_acq.py |
Handles export failures in _on_export_data_done() and updates GUI state accordingly. |
src/odemis/acq/test/acq_test.py |
Adds simulation tests covering full success, mid-stack wrong-shape failure, and first-level failure behavior. |
Comments suppressed due to low confidence (1)
src/odemis/acq/test/acq_test.py:46
from odemis.dataio import tiffis imported but not referenced anywhere in this file. Please remove it (or add the intended usage) to avoid unused-import lint failures.
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)
| import logging | ||
| import os | ||
| import time | ||
| import unittest |
| 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:]) |
|
Warning Rate limit exceeded
You’ve run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (5)
📝 WalkthroughWalkthroughThis PR implements partial z-stack acquisition with graceful error recovery. When acquiring multi-level z-stacks, the task now detects per-zlevel failures and shape inconsistencies, assembles z-cubes from successfully acquired levels only, and returns early with error information instead of failing the entire acquisition. Changes include per-stream error tracking in ZStackAcquisitionTask, input validation in assembleZCube, comprehensive simulation tests covering success and partial-failure scenarios, and improved error handling in the GUI export path. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
src/odemis/gui/cont/acquisition/cryo_acq.py (2)
535-535: ⚡ Quick winAdd type hints for function parameters and return type.
The method signature is missing type hints for the
futureparameter and return type.📝 Proposed fix
- def _on_export_data_done(self, future): + def _on_export_data_done(self, future: futures.Future) -> None:As per coding guidelines: Always use type hints for function parameters and return types in Python code.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/odemis/gui/cont/acquisition/cryo_acq.py` at line 535, The method _on_export_data_done is missing type hints for its parameter and return type; update its signature to annotate future (e.g., future: concurrent.futures.Future[Any] or typing.Any if the exact Future type is unknown) and add an explicit return type of -> None, and add the necessary import(s) (from concurrent.futures import Future or from typing import Any) at the top of the module to satisfy the codebase typing guidelines.
536-538: ⚡ Quick winEnhance docstring to document parameters and error handling behavior.
The docstring should document the
futureparameter and describe both success and error paths, especially since this PR introduces error handling.📝 Suggested enhancement
def _on_export_data_done(self, future: futures.Future) -> None: """ - Called after exporting the data + Called after exporting the data to file. + + Retrieves the exported data from the future, updates the GUI to finished state, + and displays the acquired data. If export fails, logs the error and resets the + GUI to failed state with a user-facing error message. + + :param future: The future containing the exported data result. """As per coding guidelines: Include docstrings for all functions and classes, following the reStructuredText style guide.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/odemis/gui/cont/acquisition/cryo_acq.py` around lines 536 - 538, Update the docstring for the callback method in cryo_acq.py whose docstring currently reads "Called after exporting the data" (the export-completion callback) to follow reStructuredText: add a :param future: description explaining it is the concurrent.futures.Future (or asyncio.Future) returned by the export operation, describe the success path (what attributes/state are expected on success and any side-effects performed), and describe the error path (how exceptions are obtained via future.exception()/future.result(), what exceptions may be raised/handled, and that errors are logged/propagated). Also document any return value and whether the method raises or swallows exceptions so callers know the behavior.src/odemis/util/test/img_test.py (1)
2071-2096: ⚡ Quick winAdd return type hints to the new test methods.
These new function definitions still omit
-> None, which breaks the repo-wide typing rule for Python code.As per coding guidelines, "Always use type hints for function parameters and return types in Python code"Suggested fix
- def test_assemble_zcube_empty_list_raises(self): + def test_assemble_zcube_empty_list_raises(self) -> None: """ assembleZCube() must raise ValueError when given an empty image list. """ with self.assertRaises(ValueError): img.assembleZCube([], []) - def test_assemble_zcube_inconsistent_shapes_raises(self): + def test_assemble_zcube_inconsistent_shapes_raises(self) -> None: """ assembleZCube() must raise ValueError when z-level images have different shapes.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/odemis/util/test/img_test.py` around lines 2071 - 2096, Both new test methods (test_assemble_zcube_empty_list_raises and test_assemble_zcube_inconsistent_shapes_raises) are missing return type hints; add "-> None" to each function definition so they comply with the repo typing rule (e.g., def test_assemble_zcube_empty_list_raises(self) -> None: and def test_assemble_zcube_inconsistent_shapes_raises(self) -> None:).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/odemis/acq/acqmng.py`:
- Around line 246-279: The partial z-stack currently passes
self._zlevels[stream][:len(zstack)] to assembleZCube which assumes the first N z
positions succeeded; instead track the actual z coordinates for successful
frames (e.g., maintain a z_positions list alongside zstack and append the
current z when you append da) and pass that list to assembleZCube so the Z
metadata matches the acquired images (update the code around zstack appends and
replace the assembleZCube call to use the recorded z_positions rather than
slicing self._zlevels).
In `@src/odemis/acq/test/acq_test.py`:
- Line 24: Remove the unused imports `tempfile` and `tiff` from the test module
import list so they no longer appear at the top of the file; locate the import
statements (the line that imports `tempfile` and the line that imports `tiff`)
in src/odemis/acq/test/acq_test.py and delete those lines to clean up unused
dependencies and avoid linter warnings.
- Around line 863-921: The four test helpers (_make_sim_future,
_make_sim_data_array, _make_sim_stream, _make_sim_task) are missing type hints;
update their signatures and returns to follow project typing conventions:
annotate _make_sim_future(result: Any = None) -> concurrent.futures.Future[Any]
(or typing.Future[Any]), _make_sim_data_array(shape: Tuple[int, int] = (64, 64),
dtype: numpy.dtype = numpy.uint16) -> model.DataArray, _make_sim_stream(name:
str = "mock_stream") -> mock.MagicMock (or typing.Any if MagicMock is not
imported for typing), and _make_sim_task(stream_mock: mock.MagicMock, zlevels:
Dict[mock.MagicMock, List[float]]) -> Tuple[ZStackAcquisitionTask,
mock.MagicMock]; add any needed typing imports (typing.Tuple, Dict, List, Any)
at the top of the file. Ensure default values remain the same and do not change
runtime behavior.
In `@src/odemis/util/img.py`:
- Around line 1289-1297: Add a check that images and zlevels have the same
non-zero length immediately after the empty-images guard: verify len(zlevels) >
0 and len(images) == len(zlevels) and raise ValueError with a clear message if
not; do this alongside the existing spatial-shape validation (you can reference
the variables images, zlevels and ref_shape to locate the code) so callers
cannot pass mismatched lists that would later corrupt Z metadata.
---
Nitpick comments:
In `@src/odemis/gui/cont/acquisition/cryo_acq.py`:
- Line 535: The method _on_export_data_done is missing type hints for its
parameter and return type; update its signature to annotate future (e.g.,
future: concurrent.futures.Future[Any] or typing.Any if the exact Future type is
unknown) and add an explicit return type of -> None, and add the necessary
import(s) (from concurrent.futures import Future or from typing import Any) at
the top of the module to satisfy the codebase typing guidelines.
- Around line 536-538: Update the docstring for the callback method in
cryo_acq.py whose docstring currently reads "Called after exporting the data"
(the export-completion callback) to follow reStructuredText: add a :param
future: description explaining it is the concurrent.futures.Future (or
asyncio.Future) returned by the export operation, describe the success path
(what attributes/state are expected on success and any side-effects performed),
and describe the error path (how exceptions are obtained via
future.exception()/future.result(), what exceptions may be raised/handled, and
that errors are logged/propagated). Also document any return value and whether
the method raises or swallows exceptions so callers know the behavior.
In `@src/odemis/util/test/img_test.py`:
- Around line 2071-2096: Both new test methods
(test_assemble_zcube_empty_list_raises and
test_assemble_zcube_inconsistent_shapes_raises) are missing return type hints;
add "-> None" to each function definition so they comply with the repo typing
rule (e.g., def test_assemble_zcube_empty_list_raises(self) -> None: and def
test_assemble_zcube_inconsistent_shapes_raises(self) -> None:).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 1cea2471-8866-4df6-beb3-180d143ae214
📒 Files selected for processing (5)
src/odemis/acq/acqmng.pysrc/odemis/acq/test/acq_test.pysrc/odemis/gui/cont/acquisition/cryo_acq.pysrc/odemis/util/img.pysrc/odemis/util/test/img_test.py
| """ | ||
| import logging | ||
| import os | ||
| import tempfile |
There was a problem hiding this comment.
Remove unused imports.
The imports tempfile (line 24) and tiff (line 45) are not used anywhere in this test file. Removing them keeps dependencies clean.
🧹 Proposed fix
Remove line 24:
-import tempfileRemove line 45:
-from odemis.dataio import tiffAlso applies to: 45-45
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/odemis/acq/test/acq_test.py` at line 24, Remove the unused imports
`tempfile` and `tiff` from the test module import list so they no longer appear
at the top of the file; locate the import statements (the line that imports
`tempfile` and the line that imports `tiff`) in src/odemis/acq/test/acq_test.py
and delete those lines to clean up unused dependencies and avoid linter
warnings.
| 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 |
There was a problem hiding this comment.
Add type hints to helper functions.
All four helper functions lack type hints for parameters and return types. As per coding guidelines, Python functions should include type hints.
📝 Proposed fix with type hints
+from typing import Any, Dict, List, Tuple
+from unittest.mock import MagicMock
+
-def _make_sim_future(result=None):
+def _make_sim_future(result: Any = None) -> Future:
"""
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):
+def _make_sim_data_array(shape: Tuple[int, int] = (64, 64), dtype: numpy.dtype = numpy.uint16) -> model.DataArray:
"""
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"):
+def _make_sim_stream(name: str = "mock_stream") -> MagicMock:
"""
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):
+def _make_sim_task(stream_mock: MagicMock, zlevels: Dict[Any, List[float]]) -> Tuple[ZStackAcquisitionTask, MagicMock]:
"""
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, futureAs per coding guidelines: "Always use type hints for function parameters and return types in Python code."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/odemis/acq/test/acq_test.py` around lines 863 - 921, The four test
helpers (_make_sim_future, _make_sim_data_array, _make_sim_stream,
_make_sim_task) are missing type hints; update their signatures and returns to
follow project typing conventions: annotate _make_sim_future(result: Any = None)
-> concurrent.futures.Future[Any] (or typing.Future[Any]),
_make_sim_data_array(shape: Tuple[int, int] = (64, 64), dtype: numpy.dtype =
numpy.uint16) -> model.DataArray, _make_sim_stream(name: str = "mock_stream") ->
mock.MagicMock (or typing.Any if MagicMock is not imported for typing), and
_make_sim_task(stream_mock: mock.MagicMock, zlevels: Dict[mock.MagicMock,
List[float]]) -> Tuple[ZStackAcquisitionTask, mock.MagicMock]; add any needed
typing imports (typing.Tuple, Dict, List, Any) at the top of the file. Ensure
default values remain the same and do not change runtime behavior.
| :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:] |
There was a problem hiding this comment.
Reject mismatched images and zlevels lengths here as well.
This new validation block still lets callers pass a non-empty images list with the wrong number of zlevels. That either blows up later when zlevels is empty, or worse, silently writes incorrect Z metadata when the lengths differ.
Suggested fix
if not images:
raise ValueError("Cannot assemble z-cube from an empty image list")
+ if len(images) != len(zlevels):
+ raise ValueError(
+ "Cannot assemble z-cube from %d images and %d z-levels"
+ % (len(images), len(zlevels))
+ )
# Validate that all images have the same spatial (YX) dimensions.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/odemis/util/img.py` around lines 1289 - 1297, Add a check that images and
zlevels have the same non-zero length immediately after the empty-images guard:
verify len(zlevels) > 0 and len(images) == len(zlevels) and raise ValueError
with a clear message if not; do this alongside the existing spatial-shape
validation (you can reference the variables images, zlevels and ref_shape to
locate the code) so callers cannot pass mismatched lists that would later
corrupt Z metadata.
pieleric
left a comment
pieleric
left a comment
There was a problem hiding this comment.
Nice way to write defensive code by handling the error at 3 different levels.
The whole code assumes that the issue is because one image was of different shape. Did you really really confirm this? Were you able to reproduce it and observe the shape, or see it in the log? To the best of my knowledge, a camera issue would just return no data at all. I've never seen a camera returning a part of a frame. There would be other ways to end up with a frame of different shape though, such as if the binning was changed by the user.
| 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 |
There was a problem hiding this comment.
We already have something similar. It's a class: create a InstantaneousFuture().
| bad_img = _make_sim_data_array((32, 64)) # truncated height — simulates camera error | ||
|
|
||
| acq_futures = [ | ||
| _make_sim_future(([good_img], None)), |
There was a problem hiding this comment.
InstantaneousFuture([good_img])
| 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. |
There was a problem hiding this comment.
Drop this second line about handling the TODO. If it's done, it's done.
| 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 |
There was a problem hiding this comment.
I don't think that simulates a camera error. In no cases will there be a truncated heigh. (it could happen only if the user manually change the binning while acquiring). Typically, in case of camera error, you just get a timeout error and (maybe) an empty array as data.
|
|
||
| # 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. |
There was a problem hiding this comment.
Drop the "which cannot be written to TIFF", and replace by "instead of generating the expected 3D array".
| # 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. |
There was a problem hiding this comment.
No need to be so specific. Just indicate in case of partial failure, the data might not be the right shape, so better drop it than not being able to create a Z-stack at all.
There was a problem hiding this comment.
I think using "not the right shape" wording here is not ideal as we only care about situations when one of the Z-stack layers is simply None or fully empty and has the (0, 0) shape. We just want to discard those empty arrays. I think this wording is misleading and that's the reason why in line 969 in acq_test.py Karishma also tried to simulate the truncated height.
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).