Merge branch 'autofocus-time' into 'v3'

Capture timestamp with frame to improve autofocus

Closes #723

See merge request openflexure/openflexure-microscope-server!656

src/openflexure_microscope_server/things/autofocus.py

@@ -345,12 +345,14 @@ class JPEGSharpnessMonitor:
     async def monitor_sharpness(self) -> None:
         """Start monitoring sharpness metrics."""
         self.running = True
-        async for frame in self.camera.lores_mjpeg_stream.frame_async_generator():
-            self._jpeg_times.append(time.time())
+        while self.running:
+            i = await self.camera.lores_mjpeg_stream.next_frame()
+            entry = await self.camera.lores_mjpeg_stream.ringbuffer_entry(i)
+            self._jpeg_times.append(entry.timestamp.timestamp())
 
             # JPEG sharpness metric
             if self.record & SharpnessMethod.JPEG:
-                self._jpeg_sizes.append(len(frame))
+                self._jpeg_sizes.append(len(entry.frame))
 
             # FocusFoM metric
             if self.record & SharpnessMethod.FOCUS_FOM:

src/openflexure_microscope_server/things/camera/__init__.py

@@ -16,7 +16,7 @@ from abc import ABC, abstractmethod
 from copy import deepcopy
 from datetime import datetime
 from types import TracebackType
-from typing import Any, Literal, Mapping, Optional, Self
+from typing import Any, Literal, Mapping, Optional, Self, Union, overload
 
 import numpy as np
 import piexif
@@ -216,6 +216,70 @@ class CaptureGalleryInfo(BaseModel):
     card_type: Literal["Capture"] = "Capture"
 
 
+class MJPEGStreamWithTimestamp(lt.outputs.MJPEGStream):
+    """An MJPEG stream where the frame time can be recorded with the frame.
+
+    This can be used to capture when the image was taken by the sensor if known
+    rather than take the time it is added into the ring buffer.
+    """
+
+    def add_frame(self, frame: bytes, timestamp: Optional[datetime] = None) -> None:
+        """Add a JPEG to the MJPEG stream.
+
+        Modify the standard function to have an option to send in the capture time.
+
+        :param frame: The frame to add
+        :param timestamp: The time the frame was captured.
+
+        :raise ValueError: if the supplied frame does not start with the JPEG
+            start bytes and end with the end bytes.
+        """
+        if not (
+            frame[0] == 0xFF
+            and frame[1] == 0xD8
+            and frame[-2] == 0xFF
+            and frame[-1] == 0xD9
+        ):
+            raise ValueError("Invalid JPEG")
+        with self._lock:
+            entry = self._ringbuffer[(self.last_frame_i + 1) % len(self._ringbuffer)]
+            entry.timestamp = timestamp if timestamp is not None else datetime.now()
+            entry.frame = frame
+            entry.index = self.last_frame_i + 1
+            self._thing_server_interface.start_async_task_soon(
+                self.notify_new_frame, entry.index
+            )
+
+
+class MJPEGStreamDescriptorWithTimestamp(lt.outputs.MJPEGStreamDescriptor):
+    """Subclass the Labthings descriptor to return our own class."""
+
+    @overload
+    def __get__(self, obj: Literal[None], type: type | None = None) -> Self: ...
+
+    @overload
+    def __get__(
+        self, obj: lt.Thing, type: type | None = None
+    ) -> MJPEGStreamWithTimestamp: ...
+
+    def __get__(
+        self,
+        obj: Optional[lt.Thing],
+        type: type[lt.Thing] | None = None,  # noqa: A002
+    ) -> Union[MJPEGStreamWithTimestamp | Self]:
+        """Return our own MJPEG Stream with timestamp."""
+        if obj is None:
+            return self
+        try:
+            return obj.__dict__[self.name]
+        except KeyError:
+            obj.__dict__[self.name] = MJPEGStreamWithTimestamp(
+                **self._kwargs,
+                thing_server_interface=obj._thing_server_interface,
+            )
+            return obj.__dict__[self.name]
+
+
 class BaseCamera(OFMThing, ABC):
     """The base class for all cameras. All cameras must directly inherit from this class.
 
@@ -227,8 +291,8 @@ class BaseCamera(OFMThing, ABC):
 
     _all_background_detectors: Mapping[str, BackgroundDetectAlgorithm] = lt.thing_slot()
 
-    mjpeg_stream = lt.outputs.MJPEGStreamDescriptor()
-    lores_mjpeg_stream = lt.outputs.MJPEGStreamDescriptor()
+    mjpeg_stream = MJPEGStreamDescriptorWithTimestamp()
+    lores_mjpeg_stream = MJPEGStreamDescriptorWithTimestamp()
     _memory_buffer = CameraMemoryBuffer()
     supports_focus_fom: bool = False
 

src/openflexure_microscope_server/things/camera/picamera.py

@@ -24,6 +24,7 @@ import tempfile
 import time
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
+from datetime import datetime
 from threading import RLock
 from types import TracebackType
 from typing import (
@@ -54,7 +55,7 @@ from openflexure_microscope_server.ui import (
     property_control_for,
 )
 
-from . import BaseCamera, CaptureMode, StreamingMode
+from . import BaseCamera, CaptureMode, MJPEGStreamWithTimestamp, StreamingMode
 from . import picamera_recalibrate_utils as recalibrate_utils
 from . import picamera_tuning_file_utils as tf_utils
 
@@ -75,24 +76,36 @@ class MissingCalibrationError(RuntimeError):
 class PicameraStreamOutput(Output):
     """An Output class that sends frames to a stream."""
 
-    def __init__(self, stream: lt.outputs.MJPEGStream) -> None:
+    def __init__(self, stream: MJPEGStreamWithTimestamp, encoder: MJPEGEncoder) -> None:
         """Create an output that puts frames in an MJPEGStream.
 
         :param stream: The labthings MJPEGStream to send frames to.
+        :param encoder: The encoder that is encoding the frames. This is needed to get
+            the full timestamp.
         """
         Output.__init__(self)
         self.stream = stream
+        self.encoder = encoder
+        # The difference between unix time and the CPU time
+        self.cpu_dt = time.time() - time.monotonic()
 
     def outputframe(
         self,
         frame: bytes,
         _keyframe: Optional[bool] = True,
-        _timestamp: Optional[int] = None,
+        timestamp: Optional[int] = None,
         _packet: Any = None,
         _audio: bool = False,
     ) -> None:
         """Add a frame to the stream's ringbuffer."""
-        self.stream.add_frame(frame)
+        if timestamp is None or self.encoder.firsttimestamp is None:
+            frame_timestamp = None
+        else:
+            # Reconstruct full CPU time of the frame in us
+            timestamp += self.encoder.firsttimestamp
+            # Convert to datetime
+            frame_timestamp = datetime.fromtimestamp(self.cpu_dt + timestamp / 1e6)
+        self.stream.add_frame(frame, frame_timestamp)
 
 
 class PiCamera2StreamingMode(StreamingMode):
@@ -487,14 +500,16 @@ class StreamingPiCamera2(BaseCamera, ABC):
                 picam.configure(stream_config)
                 LOGGER.info("Starting picamera MJPEG stream...")
                 stream_name = "lores" if mode_info.use_lores_as_preview else "main"
+                encoder = MJPEGEncoder(self.mjpeg_bitrate)
                 picam.start_recording(
-                    MJPEGEncoder(self.mjpeg_bitrate),
-                    PicameraStreamOutput(self.mjpeg_stream),
+                    encoder,
+                    PicameraStreamOutput(self.mjpeg_stream, encoder),
                     name=stream_name,
                 )
+                encoder = MJPEGEncoder(100000000)
                 picam.start_encoder(
-                    MJPEGEncoder(100000000),
-                    PicameraStreamOutput(self.lores_mjpeg_stream),
+                    encoder,
+                    PicameraStreamOutput(self.lores_mjpeg_stream, encoder),
                     name="lores",
                 )
             except Exception as e:

src/openflexure_microscope_server/things/camera/simulation.py

@@ -12,6 +12,7 @@ import io
 import logging
 import re
 import time
+from datetime import datetime
 from threading import Thread
 from types import TracebackType
 from typing import Optional, Self, overload
@@ -467,9 +468,10 @@ class SimulatedCamera(BaseCamera):
             last_frame_t = time.time()
 
             frame = self.generate_frame()
-            self.mjpeg_stream.add_frame(_frame2bytes(frame))
+            timestamp = datetime.fromtimestamp(last_frame_t)
+            self.mjpeg_stream.add_frame(_frame2bytes(frame), timestamp)
             ds_frame = frame.resize((320, 240), resample=Image.Resampling.NEAREST)
-            self.lores_mjpeg_stream.add_frame(_frame2bytes(ds_frame))
+            self.lores_mjpeg_stream.add_frame(_frame2bytes(ds_frame), timestamp)
 
     @lt.action
     def discard_frames(self) -> None: