diff --git a/software/control/core/job_processing.py b/software/control/core/job_processing.py index bcf183df0..8c68b3a76 100644 --- a/software/control/core/job_processing.py +++ b/software/control/core/job_processing.py @@ -150,13 +150,40 @@ def _acquire_file_lock(lock_path: str, context: str = ""): ) from exc +@dataclass +class TiffWriteResult: + """Result from a SaveImageJob, containing frame info for viewer notification.""" + + filepath: str + fov: int + time_point: int + z_index: int + channel_name: str + region_id: int = 0 + + class SaveImageJob(Job): _log: ClassVar = squid.logging.get_logger("SaveImageJob") - def run(self) -> bool: + def _make_result(self) -> TiffWriteResult: + info = self.capture_info + filepath = utils_acquisition.get_image_filepath( + info.save_directory, info.file_id, info.configuration.name, np.uint16 + ) + return TiffWriteResult( + filepath=filepath, + fov=info.fov, + time_point=info.time_point or 0, + z_index=info.z_index, + channel_name=info.configuration.name, + region_id=info.region_id, + ) + + def run(self) -> TiffWriteResult: from control.core.io_simulation import is_simulation_enabled, simulated_tiff_write image = self.image_array() + result = self._make_result() # Simulated disk I/O mode - encode to buffer, throttle, discard if is_simulation_enabled(): @@ -164,10 +191,11 @@ def run(self) -> bool: self._log.debug( f"SaveImageJob {self.job_id}: simulated write of {bytes_written} bytes " f"(image shape={image.shape})" ) - return True + return result is_color = len(image.shape) > 2 - return self.save_image(image, self.capture_info, is_color) + self.save_image(image, self.capture_info, is_color) + return result def save_image(self, image: np.array, info: CaptureInfo, is_color: bool): # NOTE(imo): We silently fall back to individual image saving here. We should warn or do something. @@ -235,7 +263,30 @@ class SaveOMETiffJob(Job): _log: ClassVar = squid.logging.get_logger("SaveOMETiffJob") acquisition_info: Optional[AcquisitionInfo] = field(default=None) - def run(self) -> bool: + def _make_result(self) -> TiffWriteResult: + info = self.capture_info + # Use the actual OME-TIFF path with page index for correct plane reading + # OME-TIFF stack layout: (T, Z, C, Y, X) — page = t * (Z * C) + z * C + c + ome_folder = ome_tiff_writer.ome_output_folder(self.acquisition_info, info) + base_name = ome_tiff_writer.ome_base_name(info) + ome_path = os.path.join(ome_folder, base_name + ".ome.tiff") + t = info.time_point or 0 + z = info.z_index + c = info.configuration_idx + n_z = self.acquisition_info.total_z_levels + n_c = self.acquisition_info.total_channels + page_idx = t * (n_z * n_c) + z * n_c + c + filepath = f"{ome_path}#{page_idx}" + return TiffWriteResult( + filepath=filepath, + fov=info.fov, + time_point=t, + z_index=z, + channel_name=info.configuration.name, + region_id=info.region_id, + ) + + def run(self) -> TiffWriteResult: if self.acquisition_info is None: raise ValueError( "SaveOMETiffJob.run() requires acquisition_info but it is None. " @@ -275,10 +326,10 @@ def run(self) -> bool: f"SaveOMETiffJob {self.job_id}: simulated write of {bytes_written} bytes " f"(image shape={image.shape})" ) - return True + return self._make_result() self._save_ome_tiff(image, self.capture_info) - return True + return self._make_result() def _save_ome_tiff(self, image: np.ndarray, info: CaptureInfo) -> None: # with reference to Talley's https://github.com/pymmcore-plus/pymmcore-plus/blob/main/src/pymmcore_plus/mda/handlers/_ome_tiff_writer.py and Christoph's https://forum.image.sc/t/how-to-create-an-image-series-ome-tiff-from-python/42730/7 diff --git a/software/control/core/multi_point_utils.py b/software/control/core/multi_point_utils.py index 067ce06a6..590c70dc2 100644 --- a/software/control/core/multi_point_utils.py +++ b/software/control/core/multi_point_utils.py @@ -125,3 +125,6 @@ class MultiPointControllerFunctions: # Zarr frame written callback - called when subprocess completes writing a frame # Args: (fov, time_point, z_index, channel_name, region_idx) signal_zarr_frame_written: Callable[[int, int, int, str, int], None] = lambda *a, **kw: None + # TIFF frame written callback - called when subprocess completes writing a frame + # Args: (filepath, fov, time_point, z_index, channel_name, region_id) + signal_tiff_frame_written: Callable[[str, int, int, int, str, int], None] = lambda *a, **kw: None diff --git a/software/control/core/multi_point_worker.py b/software/control/core/multi_point_worker.py index fa5fabb71..b8515717c 100644 --- a/software/control/core/multi_point_worker.py +++ b/software/control/core/multi_point_worker.py @@ -33,7 +33,7 @@ from squid.abc import AbstractCamera, CameraFrame, CameraFrameFormat import squid.logging import control.core.job_processing -from control.core.job_processing import ZarrWriteResult +from control.core.job_processing import TiffWriteResult, ZarrWriteResult from control.core.job_processing import ( CaptureInfo, SaveImageJob, @@ -833,6 +833,12 @@ def _summarize_job_result(self, job_result: JobResult) -> bool: elif isinstance(job_result.result, ZarrWriteResult): r = job_result.result self.callbacks.signal_zarr_frame_written(r.fov, r.time_point, r.z_index, r.channel_name, r.region_idx) + # Handle TiffWriteResult - notify viewer that frame is written + elif isinstance(job_result.result, TiffWriteResult): + r = job_result.result + self.callbacks.signal_tiff_frame_written( + r.filepath, r.fov, r.time_point, r.z_index, r.channel_name, r.region_id + ) return True def _handle_downsampled_view_result(self, result: DownsampledViewResult) -> None: diff --git a/software/control/gui_hcs.py b/software/control/gui_hcs.py index aea0c843b..73b7d1f28 100644 --- a/software/control/gui_hcs.py +++ b/software/control/gui_hcs.py @@ -252,6 +252,7 @@ def __init__( signal_slack_timepoint_notification=self._signal_slack_timepoint_notification_fn, signal_slack_acquisition_finished=self._signal_slack_acquisition_finished_fn, signal_zarr_frame_written=self._signal_zarr_frame_written_fn, + signal_tiff_frame_written=self._signal_tiff_frame_written_fn, ), scan_coordinates=scan_coordinates, laser_autofocus_controller=laser_autofocus_controller, @@ -362,30 +363,9 @@ def _signal_new_image_fn(self, frame: squid.abc.CameraFrame, info: CaptureInfo): frame.frame, info.position.x_mm, info.position.y_mm, info.z_index, napri_layer_name ) - # NDViewer push-based API: register image - # Compute flat FOV index from region and fov within region - region_offset = self._ndviewer_region_fov_offset.get(info.region_id) - if region_offset is None: - # This should not happen if start_acquisition was called correctly - self.log.warning( - f"Unknown region_id '{info.region_id}' in NDViewer registration. " - f"Available: {list(self._ndviewer_region_fov_offset.keys())}. Skipping." - ) - return - flat_fov_idx = region_offset + info.fov - - if self._ndviewer_mode in (NDViewerMode.ZARR_6D, NDViewerMode.ZARR_5D): - # Zarr modes: notification happens via signal_zarr_frame_written callback - # when the subprocess completes writing, not here (too early). - pass - else: - # TIFF mode: register with filepath (synchronous write, notification is correct here) - filepath = control.utils_acquisition.get_image_filepath( - info.save_directory, info.file_id, info.configuration.name, frame.frame.dtype - ) - self.ndviewer_register_image.emit( - info.time_point, flat_fov_idx, info.z_index, info.configuration.name, filepath - ) + # NDViewer notification for both TIFF and Zarr modes happens via post-write + # callbacks (signal_tiff_frame_written / signal_zarr_frame_written) to avoid + # race conditions with async file writes. Do not emit here. def _signal_current_configuration_fn(self, config: AcquisitionChannel): self.signal_current_configuration.emit(config) @@ -449,6 +429,19 @@ def _signal_zarr_frame_written_fn( flat_fov = fov self.ndviewer_notify_zarr_frame.emit(time_point, flat_fov, z_index, channel_name, 0) + def _signal_tiff_frame_written_fn( + self, filepath: str, fov: int, time_point: int, z_index: int, channel_name: str, region_id: int + ): + """Called when subprocess completes writing a TIFF frame. + + This is the correct time to notify the viewer - after data is on disk. + """ + region_offset = self._ndviewer_region_fov_offset.get(region_id) + if region_offset is None: + return + flat_fov_idx = region_offset + fov + self.ndviewer_register_image.emit(time_point, flat_fov_idx, z_index, channel_name, filepath) + # ------------------------------------------------------------------------- # Helper methods for Zarr FOV path building # -------------------------------------------------------------------------