Merge branch 'prevent-broken-frames-ending-scans' into 'v3'

Prevent broken frames ending scans

Closes #431

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

src/openflexure_microscope_server/things/camera/__init__.py

@@ -9,6 +9,7 @@ See repository root for licensing information.
 from __future__ import annotations
 from typing import Literal, Optional, Tuple, Any
 import json
+import io
 import time
 import logging
 
@@ -231,7 +232,7 @@ class BaseCamera(lt.Thing):
         self,
         stream_name: Literal["main", "lores", "raw", "full"] = "main",
         wait: Optional[float] = 5,
-    ) -> NDArray:
+    ) -> ArrayModel:
         """Acquire one image from the camera and return as an array."""
         raise NotImplementedError(
             "CameraThings must define their own capture_array method"
@@ -241,7 +242,7 @@ class BaseCamera(lt.Thing):
     """The downsampling factor when calling capture_downsampled_array."""
 
     @lt.thing_action
-    def capture_downsampled_array(self) -> NDArray:
+    def capture_downsampled_array(self) -> ArrayModel:
         """Acquire one image from the camera, downsample, and return as an array.
 
         * The array is downsamples by the thing property `downsampled_array_factor`.
@@ -270,7 +271,12 @@ class BaseCamera(lt.Thing):
         portal: lt.deps.BlockingPortal,
         stream_name: Literal["main", "lores"] = "main",
     ) -> JPEGBlob:
-        """Acquire one image from the preview stream and return as an array.
+        """Acquire one image from the preview stream and return as blob of JPEG data.
+
+        Note: in rare cases the JPEG stream may be broken. This can cause an OS error
+        when loading the image. If loading with PIL, as long as the header data is
+        complete, this error will not be raised until the data is accessed. Consider
+        using ``grab_jpeg_as_array`` instead.
 
         This differs from ``capture_jpeg`` in that it does not pause the MJPEG
         preview stream. Instead, we simply return the next frame from that
@@ -283,6 +289,33 @@ class BaseCamera(lt.Thing):
         frame = portal.call(stream.grab_frame)
         return JPEGBlob.from_bytes(frame)
 
+    @lt.thing_action
+    def grab_as_array(
+        self,
+        portal: lt.deps.BlockingPortal,
+        stream_name: Literal["main", "lores"] = "main",
+    ) -> ArrayModel:
+        """Acquire one image from the preview stream and return as an array.
+
+        It works like ``grab_jpeg`` but reliably handles broken streams. Prefer using
+        this method over directly grabbing the frame and converting to a numpy array
+        via PIL.
+
+        This differs from ``capture_array`` in that it does not pause the MJPEG
+        preview stream.
+        """
+        stream = (
+            self.lores_mjpeg_stream if stream_name == "lores" else self.mjpeg_stream
+        )
+        tries = 0
+        while tries < 3:
+            try:
+                frame = portal.call(stream.grab_frame)
+                return np.asarray(Image.open(io.BytesIO(frame)))
+            except OSError:
+                tries += 1
+        raise OSError("Could not open frames from MJPEG stream.")
+
     @lt.thing_action
     def grab_jpeg_size(
         self,
@@ -567,8 +600,7 @@ class BaseCamera(lt.Thing):
     @lt.thing_action
     def image_is_sample(self, portal: lt.deps.BlockingPortal) -> tuple[bool, str]:
         """Label the current image as either background or sample."""
-        current_image = self.grab_jpeg(portal)
+        current_image = self.grab_as_array(portal)
-        current_image = np.array(Image.open(current_image.open()))
         return self.active_detector.image_is_sample(current_image)
 
     @lt.thing_action
@@ -581,8 +613,7 @@ class BaseCamera(lt.Thing):
         future images to the distribution, to determine if each pixel is
         foreground or background.
         """
-        background = self.grab_jpeg(portal)
+        background = self.grab_as_array(portal)
-        background = np.array(Image.open(background.open()))
         self.active_detector.set_background(background)
         # Manually save settings as the setter is not called.
         self.save_settings()

src/openflexure_microscope_server/things/smart_scan.py

@@ -17,7 +17,6 @@ from subprocess import SubprocessError
 from fastapi import HTTPException
 from fastapi.responses import FileResponse
 import numpy as np
-from PIL import Image
 
 import labthings_fastapi as lt
 
@@ -234,8 +233,7 @@ class SmartScanThing(lt.Thing):
 
         :returns: (dx, dy) - the x and y displacements in steps
         """
-        test_jpg = self._cam.grab_jpeg()
+        test_image = self._cam.grab_as_array()
-        test_image = np.array(Image.open(test_jpg.open()))
 
         test_image_res = list(test_image.shape)
         csm_image_res = [int(i) for i in self._csm.image_resolution]

tests/test_camera.py

@@ -0,0 +1,66 @@
+"""Use the Simulation camera to test base camera functionality."""
+
+import tempfile
+from contextlib import contextmanager
+
+import pytest
+from fastapi.testclient import TestClient
+import numpy as np
+from PIL import Image
+
+import labthings_fastapi as lt
+
+from openflexure_microscope_server.things.camera.simulation import SimulatedCamera
+
+
+@contextmanager
+def camera_server(camera: SimulatedCamera) -> lt.ThingClient:
+    """Add the camera to a ThingServer and start a TestClient application.
+
+    The test client application is needed for the camera to have a blocking portal.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        server = lt.ThingServer(settings_folder=tmpdir)
+        server.add_thing(camera, "/camera/")
+        with TestClient(server.app):
+            yield server
+
+
+def test_handle_broken_frame():
+    """Monkey patch the the mjpeg steam so 1 in 5 frames are broken, then test operation.
+
+    This simulates the very occasional broken frames that can occur when grabbing
+    directly from the MJPEG stream.
+    """
+    camera = SimulatedCamera()
+
+    # Money patch the mjpeg_stream grab_frame to break 1 in 5 frames.
+    frame_number = 0
+    original_grabber = camera.mjpeg_stream.grab_frame
+
+    async def flaky_grabber():
+        """Break 1 in 5 frames."""
+        # Use a non-local variable to know the frame count.
+        nonlocal frame_number
+        frame = await original_grabber()
+        if frame_number % 5 == 2:
+            # Make a weird broken frame
+            frame = frame[:2000] + frame[:2000]
+        frame_number += 1
+        return frame
+
+    camera.mjpeg_stream.grab_frame = flaky_grabber
+    with camera_server(camera):
+        portal = lt.get_blocking_portal(camera)
+
+        # Check that this does cause broken frames.
+        with pytest.raises(OSError, match="broken data stream when reading image file"):
+            for _i in range(15):
+                jpeg = camera.grab_jpeg(portal)
+                np.asarray(Image.open(jpeg.open()))
+
+        # Check that grab_as_array handles the broken frames and completes without
+        # the same error.
+        for _i in range(15):
+            array = camera.grab_as_array(portal)
+            assert isinstance(array, np.ndarray)

tests/test_cameras.py

@@ -1,4 +1,8 @@
-"""Generic tests for cameras, specific camera hardware reqipres hardware specific tests."""
+"""Tests for camera classes. These tests focus on checking the camera APIs are equivalent.
+
+Tests for specific camera hardware are in the hardware-specific-tests directory. Tests
+on camera functionality using the simulation camera are in "test_camera".
+"""
 
 import pytest