diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py index ea5650cc..8722a01c 100644 --- a/src/aare/daq/daq.py +++ b/src/aare/daq/daq.py @@ -11,6 +11,7 @@ from typing import List, Tuple, Optional, Callable import cv2 import numpy as np from aareDB import SampleEventType +from jfjoch_client.exceptions import NotFoundException from jfjoch_client import ScanResult, ScanResultImagesInner @@ -1229,6 +1230,56 @@ class AareDAQ: centre_of_mass=None, ) + @staticmethod + def _scan_result_image_count(scan_result: ScanResult | None) -> int | None: + if scan_result is None or scan_result.images is None: + return None + return len(scan_result.images) + + def _upload_raster_diffraction_preview( + self, + *, + sample_id: int, + filename: str, + image_id: int, + scan_result: ScanResult | None, + request: RasterGridRequest, + ) -> None: + image_count = self._scan_result_image_count(scan_result) + + if image_count is not None and not (0 <= image_id < image_count): + logger.warning( + "Skipping diffraction image upload because image id is outside scan result range", + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(request), + { + "diffraction_image_id": image_id, + "image_count": image_count, + }, + ), + ) + return + + try: + diffraction_image = self.__jfjoch.get_diffraction_image(image_id) + except NotFoundException: + logger.warning( + "JFJoch diffraction preview image was not found after raster; continuing without upload", + extra=merge_log_context( + sample_log_context(self.sample), + raster_request_log_context(request), + { + "diffraction_image_id": image_id, + "image_count": image_count, + }, + ), + exc_info=True, + ) + return + + self.__aare.upload_jpg(sample_id, filename, diffraction_image) + @staticmethod def _grid_image_id_from_centre_offset( *, @@ -1398,8 +1449,14 @@ class AareDAQ: y_mm=y, request=request, ) - diffraction_image = self.__jfjoch.get_diffraction_image(diffraction_image_id) - self.__aare.upload_jpg(self.sample.db_id, diffraction_image_filename, diffraction_image) + + self._upload_raster_diffraction_preview( + sample_id=self.sample.db_id, + filename=diffraction_image_filename, + image_id=diffraction_image_id, + scan_result=scan_result, + request=request, + ) self.__aare.send_sample_event( self.sample, event_type=SampleEventType.RASTERED,