abc: cb take CameraFrame

software/squid/abc.py

@@ -1,5 +1,6 @@
+import dataclasses
 from abc import ABC, abstractmethod
-from typing import Callable, Optional, Tuple, Sequence
+from typing import Callable, Optional, Tuple, Sequence, List
 import abc
 import enum
 import time
@@ -188,15 +189,15 @@ def zero(self, x: bool, y: bool, z: bool, theta: bool, blocking: bool = True):
 
     @abc.abstractmethod
     def set_limits(
-        self,
-        x_pos_mm: Optional[float] = None,
-        x_neg_mm: Optional[float] = None,
-        y_pos_mm: Optional[float] = None,
-        y_neg_mm: Optional[float] = None,
-        z_pos_mm: Optional[float] = None,
-        z_neg_mm: Optional[float] = None,
-        theta_pos_rad: Optional[float] = None,
-        theta_neg_rad: Optional[float] = None,
+            self,
+            x_pos_mm: Optional[float] = None,
+            x_neg_mm: Optional[float] = None,
+            y_pos_mm: Optional[float] = None,
+            y_neg_mm: Optional[float] = None,
+            z_pos_mm: Optional[float] = None,
+            z_neg_mm: Optional[float] = None,
+            theta_pos_rad: Optional[float] = None,
+            theta_neg_rad: Optional[float] = None,
     ):
         pass
 
@@ -233,11 +234,13 @@ class CameraFrameFormat(enum.Enum):
     RAW = "RAW"
     RGB = "RGB"
 
+
 class CameraGainRange(pydantic.BaseModel):
     min_gain: float
     max_gain: float
     gain_step: float
 
+
 # NOTE(imo): Dataclass because pydantic does not like the np.array since there's no reasonable default
 # we can provide it.
 @dataclass
@@ -248,6 +251,7 @@ class CameraFrame:
     frame_format: CameraFrameFormat
     frame_pixel_format: CameraPixelFormat
 
+
 class AbstractCamera(metaclass=abc.ABCMeta):
     def __init__(self, camera_config: CameraConfig):
         """
@@ -261,9 +265,9 @@ def __init__(self, camera_config: CameraConfig):
         # Frame callbacks is a list of (id, callback) managed by add_frame_callback and remove_frame_callback.
         # Your frame receiving functions should call self._send_frame_to_callbacks(frame), and doesn't need
         # to do more than that.
-        self._frame_callbacks = []
+        self._frame_callbacks: List[Tuple[int, Callable[[CameraFrame], None]]] = []
 
-    def add_frame_callback(self, frame_callback: Callable[[np.ndarray], None]) -> int:
+    def add_frame_callback(self, frame_callback: Callable[[CameraFrame], None]) -> int:
         """
         Adds a new callback that will be called with the receipt of every new frame.  This callback
         should not block for a long time because it will be called in the frame receiving hot path!
@@ -289,16 +293,17 @@ def remove_frame_callback(self, callback_id):
         except ValueError:
             self._log.warning(f"No callback with id={callback_id}, cannot remove it.")
 
-    def _propogate_frame(self, frame):
+    def _propogate_frame(self, raw_frame: CameraFrame):
         """
         Implementations can call this to propogate a new frame to all registered callbacks.  The frame
         will be rotated/cropped/etc based on our config, so the callbacks don't need to do that.
         """
-        cb_frame = control.utils.rotate_and_flip_image(
-            frame, rotate_image_angle=self._config.rotate_image_angle, flip_image=self._config.flip
-        )
-        for cb in self._frame_callbacks:
-            cb(cb_frame)
+        camera_frame = dataclasses.replace(raw_frame,
+                                           frame=control.utils.rotate_and_flip_image(
+                                               raw_frame.frame, rotate_image_angle=self._config.rotate_image_angle,
+                                               flip_image=self._config.flip))
+        for (_, cb) in self._frame_callbacks:
+            cb(camera_frame)
 
     @abc.abstractmethod
     def set_exposure_time(self, exposure_time_ms: float):
@@ -441,7 +446,7 @@ def read_frame(self) -> np.ndarray:
     def _get_frame(self):
         """
         Camera implementations need to provide this _get_frame which should trigger, then get and return a raw
-        frame.  It is used by get_frame to grab a raw frame before doign the required flip/rotate/crop operations
+        frame.  It is used by get_frame to grab a raw frame before doing the required flip/rotate/crop operations
         based on this cameras settings.  See get_frame() above.
         """
         pass
@@ -490,7 +495,7 @@ def get_black_level(self) -> float:
         pass
 
     def set_acquisition_mode(
-        self, acquisition_mode: CameraAcquisitionMode, hw_trigger_fn: Optional[Callable[[None], None]] = None
+            self, acquisition_mode: CameraAcquisitionMode, hw_trigger_fn: Optional[Callable[[None], None]] = None
     ):
         """
         Sets the acquisition mode.  If you are specifying hardware trigger, and an external

software/squid/camera/utils.py

@@ -6,7 +6,7 @@
 
 import squid.logging
 from squid.config import CameraConfig, CameraPixelFormat, CameraVariant
-from squid.abc import AbstractCamera, CameraAcquisitionMode, CameraFrameFormat
+from squid.abc import AbstractCamera, CameraAcquisitionMode, CameraFrameFormat, CameraFrame
 
 _log = squid.logging.get_logger("squid.camera.utils")
 
@@ -267,7 +267,16 @@ def send_trigger(self):
             self._current_frame = np.roll(self._current_frame, 10, axis=0)
 
         self._frame_id += 1
-        self._propogate_frame(self._current_frame)
+
+        frame = CameraFrame(
+            frame_id=self._frame_id,
+            timestamp=time.time(),
+            frame=self._current_frame,
+            frame_format=self.get_frame_format(),
+            frame_pixel_format=self.get_pixel_format(),
+        )
+
+        self._propogate_frame(frame)
 
     @debug_log
     def cancel_exposure(self):