dls-dodal 1.33.0__py3-none-any.whl → 1.35.0__py3-none-any.whl

dodal/devices/oav/oav_parameters.py

@@ -1,15 +1,10 @@
 import json
 import xml.etree.ElementTree as et
 from collections import ChainMap
+from dataclasses import dataclass
 from typing import Any
 from xml.etree.ElementTree import Element
 
-from dodal.devices.oav.oav_errors import (
-    OAVError_BeamPositionNotFound,
-    OAVError_ZoomLevelNotFound,
-)
-from dodal.log import LOGGER
-
 # GDA currently assumes this aspect ratio for the OAV window size.
 # For some beamline this doesn't affect anything as the actual OAV aspect ratio
 # matches. Others need to take it into account to rescale the values stored in
@@ -109,106 +104,57 @@ class OAVParameters:
         return self.max_tip_distance / micronsPerPixel
 
 
-class OAVConfigParams:
-    """
-    The OAV parameters which may update depending on settings such as the zoom level.
-    """
+@dataclass
+class ZoomParams:
+    microns_per_pixel: tuple[float, float]
+    crosshair: tuple[int, int]
 
-    def __init__(
-        self,
-        zoom_params_file,
-        display_config,
-    ):
-        self.zoom_params_file: str = zoom_params_file
-        self.display_config: str = display_config
-
-    def update_on_zoom(self, value, xsize, ysize, *args, **kwargs):
-        xsize, ysize = int(xsize), int(ysize)
-        if isinstance(value, str) and value.endswith("x"):
-            value = value.strip("x")
-        zoom = float(value)
-        self.load_microns_per_pixel(zoom, xsize, ysize)
-        self.beam_centre_i, self.beam_centre_j = self.get_beam_position_from_zoom(
-            zoom, xsize, ysize
-        )
-
-    def load_microns_per_pixel(self, zoom: float, xsize: int, ysize: int) -> None:
-        """
-        Loads the microns per x pixel and y pixel for a given zoom level. These are
-        currently generated by GDA, though hyperion could generate them in future.
-        """
-        tree = et.parse(self.zoom_params_file)
-        self.micronsPerXPixel = self.micronsPerYPixel = None
-        root = tree.getroot()
-        levels = root.findall(".//zoomLevel")
-        for node in levels:
-            if _get_element_as_float(node, "level") == zoom:
-                self.micronsPerXPixel = (
-                    _get_element_as_float(node, "micronsPerXPixel")
-                    * DEFAULT_OAV_WINDOW[0]
-                    / xsize
-                )
-                self.micronsPerYPixel = (
-                    _get_element_as_float(node, "micronsPerYPixel")
-                    * DEFAULT_OAV_WINDOW[1]
-                    / ysize
-                )
-        if self.micronsPerXPixel is None or self.micronsPerYPixel is None:
-            raise OAVError_ZoomLevelNotFound(
-                f"""
-                Could not find the micronsPer[X,Y]Pixel parameters in
-                {self.zoom_params_file} for zoom level {zoom}.
-                """
-            )
 
-    def get_beam_position_from_zoom(
-        self, zoom: float, xsize: int, ysize: int
-    ) -> tuple[int, int]:
-        """
-        Extracts the beam location in pixels `xCentre` `yCentre`, for a requested zoom \
-        level. The beam location is manually inputted by the beamline operator on GDA \
-        by clicking where on screen a scintillator lights up, and stored in the \
-        display.configuration file.
-        """
-        crosshair_x_line = None
-        crosshair_y_line = None
-        with open(self.display_config) as f:
-            file_lines = f.readlines()
-            for i in range(len(file_lines)):
-                if file_lines[i].startswith("zoomLevel = " + str(zoom)):
-                    crosshair_x_line = file_lines[i + 1]
-                    crosshair_y_line = file_lines[i + 2]
-                    break
-
-        if crosshair_x_line is None or crosshair_y_line is None:
-            raise OAVError_BeamPositionNotFound(
-                f"Could not extract beam position at zoom level {zoom}"
-            )
+class OAVConfig:
+    """ Read the OAV config files and return a dictionary of {'zoom_level': ZoomParams}\
+    with information about microns per pixels and crosshairs.
+    """
 
-        beam_centre_i = (
-            int(crosshair_x_line.split(" = ")[1]) * xsize / DEFAULT_OAV_WINDOW[0]
-        )
-        beam_centre_j = (
-            int(crosshair_y_line.split(" = ")[1]) * ysize / DEFAULT_OAV_WINDOW[1]
-        )
-        LOGGER.info(f"Beam centre: {beam_centre_i, beam_centre_j}")
-        return int(beam_centre_i), int(beam_centre_j)
+    def __init__(self, zoom_params_file: str, display_config_file: str):
+        self.zoom_params = self._get_zoom_params(zoom_params_file)
+        self.display_config = self._get_display_config(display_config_file)
 
-    def calculate_beam_distance(
-        self, horizontal_pixels: int, vertical_pixels: int
-    ) -> tuple[int, int]:
-        """
-        Calculates the distance between the beam centre and the given (horizontal, vertical).
-
-        Args:
-            horizontal_pixels (int): The x (camera coordinates) value in pixels.
-            vertical_pixels (int): The y (camera coordinates) value in pixels.
-        Returns:
-            The distance between the beam centre and the (horizontal, vertical) point in pixels as a tuple
-            (horizontal_distance, vertical_distance).
-        """
+    def _get_display_config(self, display_config_file: str):
+        with open(display_config_file) as f:
+            file_lines = f.readlines()
+        return file_lines
 
-        return (
-            self.beam_centre_i - horizontal_pixels,
-            self.beam_centre_j - vertical_pixels,
-        )
+    def _get_zoom_params(self, zoom_params_file: str):
+        tree = et.parse(zoom_params_file)
+        root = tree.getroot()
+        return root.findall(".//zoomLevel")
+
+    def _read_zoom_params(self) -> dict:
+        um_per_pix = {}
+        for node in self.zoom_params:
+            zoom = str(_get_element_as_float(node, "level"))
+            um_pix_x = _get_element_as_float(node, "micronsPerXPixel")
+            um_pix_y = _get_element_as_float(node, "micronsPerYPixel")
+            um_per_pix[zoom] = (um_pix_x, um_pix_y)
+        return um_per_pix
+
+    def _read_display_config(self) -> dict:
+        crosshairs = {}
+        for i in range(len(self.display_config)):
+            if self.display_config[i].startswith("zoomLevel"):
+                zoom = self.display_config[i].split(" = ")[1].strip()
+                x = int(self.display_config[i + 1].split(" = ")[1])
+                y = int(self.display_config[i + 2].split(" = ")[1])
+                crosshairs[zoom] = (x, y)
+        return crosshairs
+
+    def get_parameters(self) -> dict[str, ZoomParams]:
+        config = {}
+        um_xy = self._read_zoom_params()
+        bc_xy = self._read_display_config()
+        for zoom_key in list(bc_xy.keys()):
+            config[zoom_key] = ZoomParams(
+                microns_per_pixel=um_xy[zoom_key],
+                crosshair=bc_xy[zoom_key],
+            )
+        return config

dodal/devices/oav/oav_to_redis_forwarder.py

@@ -1,21 +1,20 @@
 import asyncio
-import io
-import pickle
-import uuid
 from collections.abc import Awaitable, Callable
 from datetime import timedelta
+from enum import Enum
+from uuid import uuid4
 
-import numpy as np
 from aiohttp import ClientResponse, ClientSession
 from bluesky.protocols import Flyable, Stoppable
 from ophyd_async.core import (
     AsyncStatus,
+    DeviceVector,
     StandardReadable,
+    observe_value,
     soft_signal_r_and_setter,
     soft_signal_rw,
 )
-from ophyd_async.epics.signal import epics_signal_r
-from PIL import Image
+from ophyd_async.epics.core import epics_signal_r
 from redis.asyncio import StrictRedis
 
 from dodal.log import LOGGER
@@ -30,6 +29,22 @@ async def get_next_jpeg(response: ClientResponse) -> bytes:
             return line + await response.content.readuntil(JPEG_STOP_BYTE)
 
 
+class Source(Enum):
+    FULL_SCREEN = 0
+    ROI = 1
+
+
+class OAVSource(StandardReadable):
+    def __init__(
+        self,
+        prefix: str,
+        oav_name: str,
+    ):
+        self.url = epics_signal_r(str, f"{prefix}MJPG_URL_RBV")
+        self.oav_name = oav_name
+        super().__init__()
+
+
 class OAVToRedisForwarder(StandardReadable, Flyable, Stoppable):
     """Forwards OAV image data to redis. To use call:
 
@@ -41,6 +56,9 @@ class OAVToRedisForwarder(StandardReadable, Flyable, Stoppable):
 
     DATA_EXPIRY_DAYS = 7
 
+    # This timeout is the maximum time that the forwarder can be streaming for
+    TIMEOUT = 30
+
     def __init__(
         self,
         prefix: str,
@@ -59,59 +77,80 @@ class OAVToRedisForwarder(StandardReadable, Flyable, Stoppable):
             redis_db: int           which redis database to connect to, defaults to 0
             name: str               the name of this device
         """
-        self.stream_url = epics_signal_r(str, f"{prefix}MJPG:MJPG_URL_RBV")
+        self.counter = epics_signal_r(int, f"{prefix}CAM:ArrayCounter_RBV")
 
-        with self.add_children_as_readables():
-            self.uuid, self.uuid_setter = soft_signal_r_and_setter(str)
+        self.sources = DeviceVector(
+            {
+                Source.ROI.value: OAVSource(f"{prefix}MJPG:", "roi"),
+                Source.FULL_SCREEN.value: OAVSource(f"{prefix}XTAL:", "fullscreen"),
+            }
+        )
+        self.selected_source = soft_signal_rw(int)
 
         self.forwarding_task = None
         self.redis_client = StrictRedis(
             host=redis_host, password=redis_password, db=redis_db
         )
 
-        self._stop_flag = False
+        self._stop_flag = asyncio.Event()
 
         self.sample_id = soft_signal_rw(int, initial_value=0)
 
-        # The uuid that images are being saved under, this should be monitored for
-        # callbacks to correlate the data
-        self.uuid, self.uuid_setter = soft_signal_r_and_setter(str)
+        with self.add_children_as_readables():
+            # The uuid that images are being saved under, this should be monitored for
+            # callbacks to correlate the data
+            self.uuid, self.uuid_setter = soft_signal_r_and_setter(str)
 
         super().__init__(name=name)
 
-    async def _get_frame_and_put_to_redis(self, response: ClientResponse):
-        """Converts the data that comes in as a jpeg byte stream into a numpy array of
-        RGB values, pickles this array then writes it to redis.
-        """
+    async def _get_frame_and_put_to_redis(
+        self, redis_uuid: str, response: ClientResponse
+    ):
+        """Stores the raw bytes of the jpeg image in redis. Murko ultimately wants a
+        pickled numpy array of pixel values but raw byes are more space efficient. There
+        may be better ways of doing this, see https://github.com/DiamondLightSource/mx-bluesky/issues/592"""
         jpeg_bytes = await get_next_jpeg(response)
-        self.uuid_setter(image_uuid := str(uuid.uuid4()))
-        img = Image.open(io.BytesIO(jpeg_bytes))
-        image_data = pickle.dumps(np.asarray(img))
-        sample_id = str(await self.sample_id.get_value())
-        await self.redis_client.hset(sample_id, image_uuid, image_data)  # type: ignore
-        await self.redis_client.expire(sample_id, timedelta(days=self.DATA_EXPIRY_DAYS))
-        LOGGER.debug(f"Sent frame to redis key {sample_id} with uuid {image_uuid}")
+        self.uuid_setter(redis_uuid)
+        sample_id = await self.sample_id.get_value()
+        redis_key = f"murko:{sample_id}:raw"
+        await self.redis_client.hset(redis_key, redis_uuid, jpeg_bytes)  # type: ignore
+        await self.redis_client.expire(redis_key, timedelta(days=self.DATA_EXPIRY_DAYS))
 
     async def _open_connection_and_do_function(
-        self, function_to_do: Callable[[ClientResponse, str | None], Awaitable]
+        self, function_to_do: Callable[[ClientResponse, OAVSource], Awaitable]
     ):
-        stream_url = await self.stream_url.get_value()
+        source_idx = await self.selected_source.get_value()
+        LOGGER.info(
+            f"Forwarding data from sample {await self.sample_id.get_value()} and OAV {source_idx}"
+        )
+        source = self.sources[source_idx]
+        stream_url = await source.url.get_value()
         async with ClientSession() as session:
             async with session.get(stream_url) as response:
-                await function_to_do(response, stream_url)
+                await function_to_do(response, source)
+
+    async def _stream_to_redis(self, response: ClientResponse, source: OAVSource):
+        """Uses the update of the frame counter as a trigger to pull an image off the OAV
+        and into redis.
 
-    async def _stream_to_redis(self, response, _):
-        while not self._stop_flag:
-            await self._get_frame_and_put_to_redis(response)
-            await asyncio.sleep(0.01)
+        The frame counter is continually increasing on the timescales we store data and
+        so can be used as a uuid. If the OAV is updating too quickly we may drop frames
+        but in this case a best effort on getting as many frames as possible is sufficient.
+        """
+        done_status = AsyncStatus(
+            asyncio.wait_for(self._stop_flag.wait(), timeout=self.TIMEOUT)
+        )
+        async for frame_count in observe_value(self.counter, done_status=done_status):
+            redis_uuid = f"{source.oav_name}-{frame_count}-{uuid4()}"
+            await self._get_frame_and_put_to_redis(redis_uuid, response)
 
-    async def _confirm_mjpg_stream(self, response, stream_url):
+    async def _confirm_mjpg_stream(self, response: ClientResponse, source: OAVSource):
         if response.content_type != "multipart/x-mixed-replace":
-            raise ValueError(f"{stream_url} is not an MJPG stream")
+            raise ValueError(f"{await source.url.get_value()} is not an MJPG stream")
 
     @AsyncStatus.wrap
     async def kickoff(self):
-        self._stop_flag = False
+        self._stop_flag.clear()
         await self._open_connection_and_do_function(self._confirm_mjpg_stream)
         self.forwarding_task = asyncio.create_task(
             self._open_connection_and_do_function(self._stream_to_redis)
@@ -125,5 +164,8 @@ class OAVToRedisForwarder(StandardReadable, Flyable, Stoppable):
     @AsyncStatus.wrap
     async def stop(self, success=True):
         if self.forwarding_task:
-            self._stop_flag = True
+            LOGGER.info(
+                f"Stopping forwarding for source id {await self.selected_source.get_value()}"
+            )
+            self._stop_flag.set()
             await self.forwarding_task

dodal/devices/oav/pin_image_recognition/__init__.py

@@ -4,14 +4,15 @@ import time
 import numpy as np
 from numpy.typing import NDArray
 from ophyd_async.core import (
+    Array1D,
     AsyncStatus,
-    HintedSignal,
     StandardReadable,
+    StandardReadableFormat,
     observe_value,
     soft_signal_r_and_setter,
     soft_signal_rw,
 )
-from ophyd_async.epics.signal import epics_signal_r
+from ophyd_async.epics.core import epics_signal_r
 
 from dodal.devices.oav.pin_image_recognition.utils import (
     ARRAY_PROCESSING_FUNCTIONS_MAP,
@@ -22,7 +23,8 @@ from dodal.devices.oav.pin_image_recognition.utils import (
 )
 from dodal.log import LOGGER
 
-Tip = tuple[int | None, int | None]
+# Tip position in x, y pixel coordinates
+Tip = Array1D[np.int32]
 
 
 class InvalidPinException(Exception):
@@ -45,7 +47,7 @@ class PinTipDetection(StandardReadable):
     no tip is found after this time it will not error but instead return {INVALID_POSITION}.
     """
 
-    INVALID_POSITION = (None, None)
+    INVALID_POSITION = np.array([np.iinfo(np.int32).min, np.iinfo(np.int32).min])
 
     def __init__(self, prefix: str, name: str = ""):
         self._prefix: str = prefix
@@ -84,16 +86,16 @@ class PinTipDetection(StandardReadable):
                 self.triggered_top_edge,
                 self.triggered_bottom_edge,
             ],
-            wrapper=HintedSignal,
+            format=StandardReadableFormat.HINTED_SIGNAL,
         )
 
         super().__init__(name=name)
 
     def _set_triggered_values(self, results: SampleLocation):
-        tip = (results.tip_x, results.tip_y)
-        if tip == self.INVALID_POSITION:
+        if results.tip_x is None or results.tip_y is None:
             raise InvalidPinException
         else:
+            tip = np.array([results.tip_x, results.tip_y])
             self._tip_setter(tip)
         self._top_edge_setter(results.edge_top)
         self._bottom_edge_setter(results.edge_bottom)

dodal/devices/oav/{grid_overlay.py → snapshots/grid_overlay.py}

@@ -1,14 +1,8 @@
-# type: ignore # OAV will soon be ophyd-async, see https://github.com/DiamondLightSource/dodal/issues/716
 from enum import Enum
 from functools import partial
-from os.path import join as path_join
 
-from ophyd import Component, Signal
 from PIL import Image, ImageDraw
 
-from dodal.devices.areadetector.plugins.MJPG import MJPG
-from dodal.log import LOGGER
-
 
 class Orientation(Enum):
     horizontal = 0
@@ -20,7 +14,7 @@ def _add_parallel_lines_to_image(
     start_x: int,
     start_y: int,
     line_length: int,
-    spacing: int,
+    spacing: float,
     num_lines: int,
     orientation=Orientation.horizontal,
 ):
@@ -38,7 +32,7 @@ def _add_parallel_lines_to_image(
         start_x (int): The x coordinate (in pixels) of the start of the initial line.
         start_y (int): The y coordinate (in pixels) of the start of the initial line.
         line_length (int): The length of each of the parallel lines in pixels.
-        spacing (int): The spacing, in pixels, between each parallel line. Strictly, \
+        spacing (float): The spacing, in pixels, between each parallel line. Strictly, \
             there are spacing-1 pixels between each line
         num_lines (int): The total number of parallel lines to draw.
         orientation (Orientation): The orientation (horizontal or vertical) of the \
@@ -46,13 +40,13 @@ def _add_parallel_lines_to_image(
     lines = [
         (
             (
-                (start_x, start_y + i * spacing),
-                (start_x + line_length, start_y + i * spacing),
+                (start_x, start_y + int(i * spacing)),
+                (start_x + line_length, start_y + int(i * spacing)),
             )
             if orientation == Orientation.horizontal
             else (
-                (start_x + i * spacing, start_y),
-                (start_x + i * spacing, start_y + line_length),
+                (start_x + int(i * spacing), start_y),
+                (start_x + int(i * spacing), start_y + line_length),
             )
         )
         for i in range(num_lines)
@@ -76,7 +70,7 @@ def add_grid_border_overlay_to_image(
     image: Image.Image,
     top_left_x: int,
     top_left_y: int,
-    box_width: int,
+    box_width: float,
     num_boxes_x: int,
     num_boxes_y: int,
 ):
@@ -84,16 +78,16 @@ def add_grid_border_overlay_to_image(
         image,
         start_x=top_left_x,
         start_y=top_left_y,
-        line_length=num_boxes_y * box_width,
-        spacing=num_boxes_x * box_width,
+        line_length=int(num_boxes_y * box_width),
+        spacing=int(num_boxes_x * box_width),
         num_lines=2,
     )
     _add_horizontal_parallel_lines_to_image(
         image,
         start_x=top_left_x,
         start_y=top_left_y,
-        line_length=num_boxes_x * box_width,
-        spacing=num_boxes_y * box_width,
+        line_length=int(num_boxes_x * box_width),
+        spacing=int(num_boxes_y * box_width),
         num_lines=2,
     )
 
@@ -102,60 +96,23 @@ def add_grid_overlay_to_image(
     image: Image.Image,
     top_left_x: int,
     top_left_y: int,
-    box_width: int,
+    box_width: float,
     num_boxes_x: int,
     num_boxes_y: int,
 ):
     _add_vertical_parallel_lines_to_image(
         image,
-        start_x=top_left_x + box_width,
+        start_x=int(top_left_x + box_width),
         start_y=top_left_y,
-        line_length=num_boxes_y * box_width,
+        line_length=int(num_boxes_y * box_width),
         spacing=box_width,
         num_lines=num_boxes_x - 1,
     )
     _add_horizontal_parallel_lines_to_image(
         image,
         start_x=top_left_x,
-        start_y=top_left_y + box_width,
-        line_length=num_boxes_x * box_width,
+        start_y=int(top_left_y + box_width),
+        line_length=int(num_boxes_x * box_width),
         spacing=box_width,
         num_lines=num_boxes_y - 1,
     )
-
-
-class SnapshotWithGrid(MJPG):
-    top_left_x = Component(Signal)
-    top_left_y = Component(Signal)
-    box_width = Component(Signal)
-    num_boxes_x = Component(Signal)
-    num_boxes_y = Component(Signal)
-
-    last_path_outer = Component(Signal)
-    last_path_full_overlay = Component(Signal)
-
-    def post_processing(self, image: Image.Image):
-        # Save an unmodified image with no suffix
-        self._save_image(image)
-
-        top_left_x = self.top_left_x.get()
-        top_left_y = self.top_left_y.get()
-        box_width = self.box_width.get()
-        num_boxes_x = self.num_boxes_x.get()
-        num_boxes_y = self.num_boxes_y.get()
-        assert isinstance(filename_str := self.filename.get(), str)
-        assert isinstance(directory_str := self.directory.get(), str)
-        add_grid_border_overlay_to_image(
-            image, top_left_x, top_left_y, box_width, num_boxes_x, num_boxes_y
-        )
-        path = path_join(directory_str, f"{filename_str}_outer_overlay.png")
-        self.last_path_outer.put(path)
-        LOGGER.info(f"Saving grid outer edge at {path}")
-        image.save(path)
-        add_grid_overlay_to_image(
-            image, top_left_x, top_left_y, box_width, num_boxes_x, num_boxes_y
-        )
-        path = path_join(directory_str, f"{filename_str}_grid_overlay.png")
-        self.last_path_full_overlay.put(path)
-        LOGGER.info(f"Saving full grid overlay at {path}")
-        image.save(path)

dodal/devices/oav/snapshots/snapshot_with_beam_centre.py

@@ -0,0 +1,64 @@
+from ophyd_async.core import Reference, SignalR
+from PIL import Image, ImageDraw
+
+from dodal.devices.areadetector.plugins.MJPG import MJPG
+
+CROSSHAIR_LENGTH_PX = 20
+CROSSHAIR_OUTLINE_COLOUR = "Black"
+CROSSHAIR_FILL_COLOUR = "White"
+
+
+def draw_crosshair(image: Image.Image, beam_x: int, beam_y: int):
+    draw = ImageDraw.Draw(image)
+    OUTLINE_WIDTH = 1
+    HALF_LEN = CROSSHAIR_LENGTH_PX / 2
+    draw.rectangle(
+        [
+            beam_x - OUTLINE_WIDTH,
+            beam_y - HALF_LEN - OUTLINE_WIDTH,
+            beam_x + OUTLINE_WIDTH,
+            beam_y + HALF_LEN + OUTLINE_WIDTH,
+        ],
+        fill=CROSSHAIR_OUTLINE_COLOUR,
+    )
+    draw.rectangle(
+        [
+            beam_x - HALF_LEN - OUTLINE_WIDTH,
+            beam_y - OUTLINE_WIDTH,
+            beam_x + HALF_LEN + OUTLINE_WIDTH,
+            beam_y + OUTLINE_WIDTH,
+        ],
+        fill=CROSSHAIR_OUTLINE_COLOUR,
+    )
+    draw.line(
+        ((beam_x, beam_y - HALF_LEN), (beam_x, beam_y + HALF_LEN)),
+        fill=CROSSHAIR_FILL_COLOUR,
+    )
+    draw.line(
+        ((beam_x - HALF_LEN, beam_y), (beam_x + HALF_LEN, beam_y)),
+        fill=CROSSHAIR_FILL_COLOUR,
+    )
+
+
+class SnapshotWithBeamCentre(MJPG):
+    """A child of MJPG which, when triggered, draws an outlined crosshair at the beam
+    centre in the image and saves the image to disk."""
+
+    def __init__(
+        self,
+        prefix: str,
+        beam_x_signal: SignalR,
+        beam_y_signal: SignalR,
+        name: str = "",
+    ) -> None:
+        with self.add_children_as_readables():
+            self._beam_centre_i_ref = Reference(beam_x_signal)
+            self._beam_centre_j_ref = Reference(beam_y_signal)
+        super().__init__(prefix, name)
+
+    async def post_processing(self, image: Image.Image):
+        beam_x = await self._beam_centre_i_ref().get_value()
+        beam_y = await self._beam_centre_j_ref().get_value()
+        draw_crosshair(image, beam_x, beam_y)
+
+        await self._save_image(image)

dodal/devices/oav/snapshots/snapshot_with_grid.py

@@ -0,0 +1,57 @@
+from os.path import join as path_join
+
+from ophyd_async.core import soft_signal_rw
+from PIL.Image import Image
+
+from dodal.devices.areadetector.plugins.MJPG import IMG_FORMAT, MJPG, asyncio_save_image
+from dodal.devices.oav.snapshots.grid_overlay import (
+    add_grid_border_overlay_to_image,
+    add_grid_overlay_to_image,
+)
+from dodal.log import LOGGER
+
+
+class SnapshotWithGrid(MJPG):
+    def __init__(self, prefix: str, name: str = "") -> None:
+        with self.add_children_as_readables():
+            self.top_left_x = soft_signal_rw(float)
+            self.top_left_y = soft_signal_rw(float)
+            self.box_width = soft_signal_rw(float)
+            self.num_boxes_x = soft_signal_rw(int)
+            self.num_boxes_y = soft_signal_rw(int)
+
+            self.last_path_outer = soft_signal_rw(str)
+            self.last_path_full_overlay = soft_signal_rw(str)
+
+        super().__init__(prefix, name)
+
+    async def post_processing(self, image: Image):
+        # Save an unmodified image with no suffix
+        await self._save_image(image)
+
+        top_left_x = await self.top_left_x.get_value()
+        top_left_y = await self.top_left_y.get_value()
+        box_width = await self.box_width.get_value()
+        num_boxes_x = await self.num_boxes_x.get_value()
+        num_boxes_y = await self.num_boxes_y.get_value()
+
+        assert isinstance(filename_str := await self.filename.get_value(), str)
+        assert isinstance(directory_str := await self.directory.get_value(), str)
+
+        add_grid_border_overlay_to_image(
+            image, int(top_left_x), int(top_left_y), box_width, num_boxes_x, num_boxes_y
+        )
+
+        path = path_join(directory_str, f"{filename_str}_outer_overlay.{IMG_FORMAT}")
+        await self.last_path_outer.set(path, wait=True)
+        LOGGER.info(f"Saving grid outer edge at {path}")
+        await asyncio_save_image(image, path)
+
+        add_grid_overlay_to_image(
+            image, int(top_left_x), int(top_left_y), box_width, num_boxes_x, num_boxes_y
+        )
+
+        path = path_join(directory_str, f"{filename_str}_grid_overlay.{IMG_FORMAT}")
+        await self.last_path_full_overlay.set(path, wait=True)
+        LOGGER.info(f"Saving full grid overlay at {path}")
+        await asyncio_save_image(image, path)