auto exposure algorithm for Imager (#612)

Author: rickwierenga

docs/user_guide/02_analytical/plate-reading/cytation5.ipynb

@@ -21,13 +21,6 @@
     "%autoreload 2"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "code",
    "execution_count": 2,
@@ -60,13 +53,6 @@
     "await pr.setup(use_cam=True)"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "code",
    "execution_count": 5,
@@ -425,6 +411,51 @@
     "plt.imshow(ims[0], cmap=\"gray\", vmin=0, vmax=255)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Autoexposure"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Two autoexposure functions are available in the PLR library:\n",
+    "- `max_pixel_at_fraction`: the value of the highest pixel in the image is a fraction of the maximum possible value (e.g. highest value is 50% of max, which would be 255/2 = 127.5 in the case of an 8 bit image)\n",
+    "- `fraction_overexposed`: the fraction of pixels at the cap (eg 255 for an 8 bit image) should be a certain fraction of the total number of pixels (e.g. 0.5% of pixels should be at the cap, so 0.005 * total_pixels). This is useful for images that are not well illuminated, as it ensures that a certain fraction of pixels is overexposed, which can help with image quality.\n",
+    "\n",
+    "You can also define your own autoexposure function.\n",
+    "\n",
+    "The `AutoExposure` dataclass is used to configure the autoexposure settings, including the evaluation function, maximum number of rounds, and low and high exposure time limits."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pylabrobot.plate_reading.imager import Imager, max_pixel_at_fraction, fraction_overexposed\n",
+    "from pylabrobot.plate_reading.standard import AutoExposure\n",
+    "\n",
+    "ims = await pr.capture(\n",
+    "  exposure_time=AutoExposure(\n",
+    "    # evaluate_exposure=fraction_overexposed(fraction=0.005, margin=0.005/10),\n",
+    "    evaluate_exposure=max_pixel_at_fraction(fraction=0.90, margin=0.05),\n",
+    "    max_rounds=15,\n",
+    "    low=1,\n",
+    "    high=100\n",
+    "  ),\n",
+    "  well=(2, 2),\n",
+    "  mode=ImagingMode.PHASE_CONTRAST,\n",
+    "  objective=Objective.O_20X_PL_FL_Phase,\n",
+    "  focal_height=1.8,  # focal height must be specified when using auto exposure\n",
+    "  gain=20  # gain must be specified when using auto exposure\n",
+    ")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

pylabrobot/plate_reading/biotek_backend.py

@@ -788,7 +788,7 @@ async def led_off(self):
   async def set_focus(self, focal_position: FocalPosition):
     """focus position in mm"""
 
-    if focal_position == "auto":
+    if focal_position == "machine-auto":
       await self.auto_focus()
       return
 
@@ -938,7 +938,7 @@ async def set_auto_exposure(self, auto_exposure: Literal["off", "once", "continu
     )
 
   async def set_exposure(self, exposure: Exposure):
-    """exposure (integration time) in ms, or "auto" """
+    """exposure (integration time) in ms, or "machine-auto" """
 
     if exposure == self._exposure:
       logger.debug("Exposure time is already set to %s", exposure)
@@ -949,9 +949,9 @@ async def set_exposure(self, exposure: Exposure):
 
     # either set auto exposure to continuous, or turn off
     if isinstance(exposure, str):
-      if exposure == "auto":
+      if exposure == "machine-auto":
         await self.set_auto_exposure("continuous")
-        self._exposure = "auto"
+        self._exposure = "machine-auto"
         return
       raise ValueError("exposure must be a number or 'auto'")
     self.cam.ExposureAuto.SetValue(PySpin.ExposureAuto_Off)
@@ -980,15 +980,15 @@ async def select(self, row: int, column: int):
     await self.set_position(0, 0)
 
   async def set_gain(self, gain: Gain):
-    """gain of unknown units, or "auto" """
+    """gain of unknown units, or "machine-auto" """
     if self.cam is None:
       raise ValueError("Camera not initialized. Run setup(use_cam=True) first.")
 
     if gain == self._gain:
       logger.debug("Gain is already set to %s", gain)
       return
 
-    if not (gain == "auto" or 0 <= gain <= 30):
+    if not (gain == "machine-auto" or 0 <= gain <= 30):
       raise ValueError("gain must be between 0 and 30 (inclusive), or 'auto'")
 
     nodemap = self.cam.GetNodeMap()
@@ -999,14 +999,14 @@ async def set_gain(self, gain: Gain):
       raise RuntimeError("unable to set automatic gain")
     node = (
       PySpin.CEnumEntryPtr(node_gain_auto.GetEntryByName("Continuous"))
-      if gain == "auto"
+      if gain == "machine-auto"
       else PySpin.CEnumEntryPtr(node_gain_auto.GetEntryByName("Off"))
     )
     if not PySpin.IsReadable(node):
       raise RuntimeError("unable to set automatic gain (enum entry retrieval)")
     node_gain_auto.SetIntValue(node.GetValue())
 
-    if not gain == "auto":
+    if not gain == "machine-auto":
       node_gain = PySpin.CFloatPtr(nodemap.GetNode("Gain"))
       if (
         not PySpin.IsReadable(node_gain)
@@ -1164,8 +1164,8 @@ async def capture(
     speed: 211 ms ± 331 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
 
     Args:
-      exposure_time: exposure time in ms, or `"auto"`
-      focal_height: focal height in mm, or `"auto"`
+      exposure_time: exposure time in ms, or `"machine-auto"`
+      focal_height: focal height in mm, or `"machine-auto"`
       coverage: coverage of the well, either `"full"` or a tuple of `(num_rows, num_columns)`.
         Around `center_position`.
       center_position: center position of the well, in mm from the center of the selected well. If

pylabrobot/plate_reading/imager.py

@@ -1,8 +1,10 @@
-from typing import List, Optional, Tuple, Union, cast
+import math
+from typing import Awaitable, Callable, List, Literal, Optional, Tuple, Union, cast
 
 from pylabrobot.machines import Machine
 from pylabrobot.plate_reading.backend import ImagerBackend
 from pylabrobot.plate_reading.standard import (
+  AutoExposure,
   Exposure,
   FocalPosition,
   Gain,
@@ -13,6 +15,11 @@
 )
 from pylabrobot.resources import Plate, Resource, Well
 
+try:
+  import numpy as np
+except ImportError:
+  np = None  # type: ignore[assignment]
+
 
 class Imager(Resource, Machine):
   """Microscope"""
@@ -52,14 +59,78 @@ def get_plate(self) -> Plate:
       raise NoPlateError("There is no plate in the plate reader.")
     return cast(Plate, self.children[0])
 
+  async def _capture_auto_exposure(
+    self,
+    well: Union[Well, Tuple[int, int]],
+    mode: ImagingMode,
+    objective: Objective,
+    auto_exposure: AutoExposure,
+    focal_height: float,
+    gain: float,
+    **backend_kwargs,
+  ) -> List[Image]:
+    """
+    Capture an image with auto exposure.
+
+    This function will iteratively adjust the exposure time until a good exposure is found.
+    It uses the provided `evaluate_exposure` function to determine if the exposure is good, too high, or too low.
+    It uses a weighted binary search to find the optimal exposure time. The search is weighted by exposure time,
+    meaning that instead of splitting the range in half, we split the range at the point that equalizes the integral
+    of the exposure time on both sides (this works out to be equal to the root mean square of the endpoints).
+    """
+
+    if focal_height == "auto":
+      raise ValueError("Focal height must be specified for auto exposure")
+    if gain == "auto":
+      raise ValueError("Gain must be specified for auto exposure")
+
+    def _rms_split(low: float, high: float) -> float:
+      """Split point that equalizes ∫t dt on both sides (RMS of endpoints)."""
+      if low == high:
+        return low
+      return math.sqrt((low**2 + high**2) / 2)
+
+    low, high = auto_exposure.low, auto_exposure.high
+
+    rounds = 0
+    while high - low > 1e-3:
+      if auto_exposure.max_rounds is not None and rounds >= auto_exposure.max_rounds:
+        raise ValueError("Exceeded maximum number of rounds")
+      rounds += 1
+
+      p = _rms_split(low, high)
+      ims = await self.capture(
+        well=well,
+        mode=mode,
+        objective=objective,
+        exposure_time=p,
+        focal_height=focal_height,
+        gain=gain,
+        **backend_kwargs,
+      )
+      assert len(ims) == 1, "Expected exactly one image to be returned"
+      im = ims[0]
+      result = await auto_exposure.evaluate_exposure(im)
+
+      if result == "good":
+        return ims
+      if result == "lower":
+        high = p
+      elif result == "higher":
+        low = p
+      else:
+        raise ValueError(f"Unexpected evaluation result: {result}")
+
+    raise RuntimeError("Failed to find a good exposure time.")
+
   async def capture(
     self,
     well: Union[Well, Tuple[int, int]],
     mode: ImagingMode,
     objective: Objective,
-    exposure_time: Exposure = "auto",
-    focal_height: FocalPosition = "auto",
-    gain: Gain = "auto",
+    exposure_time: Union[Exposure, AutoExposure] = "machine-auto",
+    focal_height: FocalPosition = "machine-auto",
+    gain: Gain = "machine-auto",
     **backend_kwargs,
   ) -> List[Image]:
     if isinstance(well, tuple):
@@ -70,6 +141,19 @@ async def capture(
         raise ValueError(f"Well {well} not in plate {well.parent}")
       row, column = divmod(idx, cast(Plate, well.parent).num_items_x)
 
+    if isinstance(exposure_time, AutoExposure):
+      assert focal_height != "machine-auto", "Focal height must be specified for auto exposure"
+      assert gain != "machine-auto", "Gain must be specified for auto exposure"
+      return await self._capture_auto_exposure(
+        well=well,
+        mode=mode,
+        objective=objective,
+        auto_exposure=exposure_time,
+        focal_height=focal_height,
+        gain=gain,
+        **backend_kwargs,
+      )
+
     return await self.backend.capture(
       row=row,
       column=column,
@@ -81,3 +165,57 @@ async def capture(
       plate=self.get_plate(),
       **backend_kwargs,
     )
+
+
+def max_pixel_at_fraction(
+  fraction: float, margin: float
+) -> Callable[[Image], Awaitable[Literal["higher", "lower", "good"]]]:
+  """The maximum pixel value in a given image should be a fraction of the maximum possible pixel value (eg 255 for 8-bit images).
+
+  Args:
+    fraction: the desired fraction of the actual maximum pixel value over the theoretically maximum pixel value (e.g. 0.8 for 80%). If it is an 8-bit image, the maximum value would be 0.8 * 255 = 204.
+    margin: the margin of error that is accepted. A fraction of the theoretical maximum pixel value, e.g. 0.05 for 5%, so the maximum pixel value should be between 0.75 * 255 and 0.85 * 255.
+  """
+
+  if np is None:
+    raise ImportError("numpy is required for max_pixel_at_fraction")
+
+  async def evaluate_exposure(im) -> Literal["higher", "lower", "good"]:
+    array = np.array(im, dtype=np.float32)
+    value = np.max(array) - (255.0 * fraction)
+    margin_value = 255.0 * margin
+    if abs(value) <= margin_value:
+      return "good"
+    # lower the exposure time if the max pixel value is too high
+    return "lower" if value > 0 else "higher"
+
+  return evaluate_exposure
+
+
+def fraction_overexposed(
+  fraction: float, margin: float, max_pixel_value: int = 255
+) -> Callable[[Image], Awaitable[Literal["higher", "lower", "good"]]]:
+  """A certain fraction of pixels in the image should be overexposed (e.g. 0.5%).
+
+  This is useful for images that are not well illuminated, as it ensures that a certain fraction of pixels is overexposed, which can help with image quality.
+
+  Args:
+    fraction: the desired fraction of pixels that should be overexposed (e.g. 0.005 for 0.5%). Overexposed is defined as pixels with a value greater than the maximum pixel value (e.g. 255 for 8-bit images). You can customize this number if needed.
+    margin: the margin of error for the fraction of pixels that should be overexposed (e.g. 0.001 for 0.1%, so the fraction of overexposed pixels should be between 0.004 and 0.006).
+    max_pixel_value: the maximum pixel value for the image (e.g. 255 for 8-bit images). You can override it to change the definition of "overexposed" pixels.
+  """
+
+  if np is None:
+    raise ImportError("numpy is required for fraction_overexposed")
+
+  async def evaluate_exposure(im) -> Literal["higher", "lower", "good"]:
+    # count the number of pixels that are overexposed
+    arr = np.asarray(im, dtype=np.uint8)
+    actual_fraction = np.count_nonzero(arr > max_pixel_value) / arr.size
+    lower_bound, upper_bound = fraction - margin, fraction + margin
+    if lower_bound <= actual_fraction <= upper_bound:
+      return "good"
+    # too many saturated pixels -> shorten exposure
+    return "lower" if (actual_fraction - fraction) > 0 else "higher"
+
+  return evaluate_exposure

pylabrobot/plate_reading/standard.py

@@ -1,5 +1,6 @@
 import enum
-from typing import List, Literal, Union
+from dataclasses import dataclass
+from typing import Awaitable, Callable, List, Literal, Union
 
 Image = List[List[float]]
 
@@ -89,6 +90,14 @@ class NoPlateError(Exception):
   pass
 
 
-Exposure = Union[float, Literal["auto"]]
-FocalPosition = Union[float, Literal["auto"]]
-Gain = Union[float, Literal["auto"]]
+@dataclass
+class AutoExposure:
+  evaluate_exposure: Callable[[Image], Awaitable[Literal["higher", "lower", "good"]]]
+  max_rounds: int
+  low: float
+  high: float
+
+
+Exposure = Union[float, Literal["machine-auto"]]
+FocalPosition = Union[float, Literal["machine-auto"]]
+Gain = Union[float, Literal["machine-auto"]]