Fix auto-exposure and AWB

The previous "auto calibrate" button for the camera relied on a
combination of the built-in autoexposure/AWB and
some JPEG-based tweaks to get the exposure settings right.
This often led to confusing results, e.g. oscillating between green and
 pink images.

 I have rewritten the auto-exposure code to adjust gain and shutter
 speed based on the raw image.  This seems to be
 very reliable, at the expense of being quite slow.  That's a price I'm
 happy to pay.

 I have replaced both the camera's auto white balance and my
 JPEG-based AWB hack with a single-shot AWB method that
 looks at the raw image.

openflexure_microscope/api/default_extensions/picamera_autocalibrate/extension.py

@@ -1,6 +1,8 @@
 import logging
 from contextlib import contextmanager
 
+import labthings.fields as fields
+import numpy as np
 import picamerax
 from flask import abort
 from labthings import find_component
@@ -13,7 +15,11 @@ from openflexure_microscope.microscope import Microscope
 from .recalibrate_utils import (
     auto_expose_and_freeze_settings,
     flat_lens_shading_table,
+    get_channel_percentiles,
     recalibrate_camera,
+    adjust_shutter_and_gain_from_raw,
+    adjust_white_balance_from_raw,
+    lst_from_camera,
 )
 
 
@@ -61,48 +67,71 @@ class LSTExtension(BaseExtension):
             "/delete_lens_shading_table",
             endpoint="delete_lens_shading_table",
         )
+        self.add_view(
+            GetRawChannelPercentilesView,
+            "/get_raw_channel_percentiles",
+            endpoint="get_raw_channel_percentiles",
+        )
+        self.add_view(
+            AutoExposureFromRawView,
+            "/auto_exposure_from_raw",
+            endpoint="auto_exposure_from_raw",
+        )
+        self.add_view(
+            AutoWhiteBalanceFromRawView,
+            "/auto_white_balance_from_raw",
+            endpoint="auto_white_balance_from_raw",
+        )
         
-    def recalibrate(self, microscope: Microscope):
+                
+
+
+def find_microscope():
+    """Locate the microscope and raise a sensible error if it's missing."""
+    microscope = find_component("org.openflexure.microscope")
+
+    if not microscope:
+        abort(
+            503, "No microscope connected. Unable to use camera calibration functions."
+        )
+
+    if not hasattr(microscope.camera, "picamera"):
+        abort(
+            503, "The PiCamera calibration plugin requires a Raspberry Pi camera."
+        )
+
+    return microscope
+
+
+class RecalibrateView(ActionView):
+    args = {
+        "skip_auto_exposure": fields.Bool(
+            missing=False, 
+            description="Whether to "
+        )
+    }
+    def post(self, args):
         """Reset the camera's settings.
 
         This generates new gains, exposure time, and lens shading
         table such that the background is as uniform as possible
         with a gray level of 230.  It takes a little while to run.
         """
-        with pause_stream(microscope.camera) as scamera:
+        microscope = find_microscope()
-            if hasattr(scamera, "picamera"):
-                picamera_obj: picamerax.PiCamera = getattr(scamera, "picamera")
-                auto_expose_and_freeze_settings(picamera_obj)
-                recalibrate_camera(picamera_obj)
-                microscope.save_settings()
-            else:
-                raise RuntimeError(
-                    "Recalibrate can only be used with a Raspberry Pi camera"
-                )
-
-
-class RecalibrateView(ActionView):
-    def post(self):
-        microscope = find_component("org.openflexure.microscope")
-
-        if not microscope:
-            abort(503, "No microscope connected. Unable to recalibrate.")
-
         logging.info("Starting microscope recalibration...")
-
+        picamera = microscope.camera.picamera
-        return self.extension.recalibrate(microscope)
+        if not args.get("skip_auto_exposure"):
+            adjust_shutter_and_gain_from_raw(picamera)
+            adjust_white_balance_from_raw(picamera)
+        lst = lst_from_camera(picamera)
+        with pause_stream(microscope.camera) as scamera:
+            scamera.picamera.lens_shading_table = lst
+        microscope.save_settings()
 
 
 class FlattenLSTView(ActionView):
     def post(self):
-        microscope = find_component("org.openflexure.microscope")
+        microscope = find_microscope()
-
-        if not microscope:
-            abort(
-                503,
-                "No microscope connected. Unable to flatten the lens shading table.",
-            )
-
         with pause_stream(microscope.camera) as scamera:
             flat_lst = flat_lens_shading_table(scamera.camera)
             scamera.camera.lens_shading_table = flat_lst
@@ -111,14 +140,86 @@ class FlattenLSTView(ActionView):
 
 class DeleteLSTView(ActionView):
     def post(self):
-        microscope = find_component("org.openflexure.microscope")
+        microscope = find_microscope()
-
-        if not microscope:
-            abort(
-                503,
-                "No microscope connected. Unable to flatten the lens shading table.",
-            )
-
         with pause_stream(microscope.camera) as scamera:
             scamera.camera.lens_shading_table = None
         microscope.save_settings()
+
+
+class AutoExposureFromRawView(ActionView):
+    args = {
+        "target_white_level": fields.Int(
+            missing=700,
+            example=700,
+            description=(
+                "The pixel value (10-bit format) that we aim for when adjusting shutter/gain."
+            ),
+        ),
+        "max_iterations": fields.Int(
+            missing=20,
+            description=(
+                "The number of adjustments to the camera's settings to make before giving up."
+            ),
+        ),
+        "tolerance": fields.Float(
+            missing=0.05,
+            example=0.05,
+            description=(
+                "We stop adjusting when we get within this fraction of the target "
+                "value.  It is a number between 0 and 1, usually 0.01--0.1."
+            ),
+        ),
+        "percentile": fields.Float(
+            missing=99.9,
+            example=99.9,
+            description=(
+                "A float between 0 and 100 setting the centile to use "
+                "to measure the white point of the image.  A value "
+                "of 99.9 allows 0.1% of the pixels to be erroneously "
+                "bright - this helps stability in low light."
+            ),
+        ),
+    }
+
+    def post(self, args):
+        camera = find_component("org.openflexure.microscope").camera.picamera
+        adjust_shutter_and_gain_from_raw(
+            camera,
+            **args  # I'm relying on the schema to fill in missing values
+        )
+
+
+class AutoWhiteBalanceFromRawView(ActionView):
+    args = {
+        "percentile": fields.Float(
+            missing=99.9,
+            example=99.9,
+            description=(
+                "A float between 0 and 100 setting the centile to use "
+                "to measure the white point of the image.  A value "
+                "of 99.9 allows 0.1% of the pixels to be erroneously "
+                "bright - this helps stability in low light."
+            ),
+        )
+    }
+
+    def post(self, args):
+        camera = find_component("org.openflexure.microscope").camera.picamera
+        adjust_white_balance_from_raw(
+            camera,
+            **args  # I'm relying on the schema to fill in missing values
+        )
+
+
+class GetRawChannelPercentilesView(ActionView):
+    args = {
+        "percentile": fields.Float(
+            example=99.9,
+            description="A float between 0 and 100 setting the centile to calculate",
+        )
+    }
+    schema = fields.List(fields.Integer)
+
+    def post(self, args):
+        camera = find_component("org.openflexure.microscope").camera.picamera
+        return get_channel_percentiles(camera, args["percentile"])

openflexure_microscope/api/default_extensions/picamera_autocalibrate/recalibrate_utils.py

@@ -44,14 +44,66 @@ def adjust_exposure_to_setpoint(camera: PiCamera, setpoint: int):
     print("done")
 
 
+def adjust_exposure_and_gain(
+    camera: PiCamera, target_white_level: int = 700, max_iterations: int = 99
+):
+    """Adjust exposure and analog gain based on raw images.
+    
+    This routine is slow but effective.  It uses raw images, so we
+    are not affected by white balance or digital gain.
+    """
+    # Start by setting fully manual exposure.
+    camera.exposure_mode = "off"
+    camera.iso = 0  # We must set ISO=0 (auto) or we can't set gain
+    camera.analog_gain = 1
+    camera.digital_gain = 1
+    camera.shutter_speed = 1  # This is not valid - we'll get the minimum
+    time.sleep(0.5)
+
+    # We start with very low exposure settings and work up in coarse steps
+    # until either the brightness is high enough, or we can't increase the
+    # shutter speed any more.
+    iterations = 0
+    while True:
+        iterations += 1
+
+        with PiBayerArray(camera) as a:
+            camera.capture(a, format="jpeg", bayer=True)
+            max_brightness = np.max(a.array)
+
+        if max_brightness > target_white_level:
+            break
+
+        previous_shutter_speed = camera.shutter_speed
+        camera.shutter_speed = previous_shutter_speed * 2
+        time.sleep(0.5)
+        current_shutter_speed = camera.shutter_speed
+        # NB we save this in a variable in case it changes between
+        # here and the next loop
+        logging.info(
+            f"Max brightness {max_brightness} < {target_white_level}. "
+            f"Attempting to double exposure from "
+            f"{previous_shutter_speed} to {current_shutter_speed}"
+        )
+        if current_shutter_speed == previous_shutter_speed:
+            logging.info(f"Shutter speed has saturated, stopping.")
+            break
+
+    if max_brightness > target_white_level:
+        logging.info("Brightness is high enough, fine-tuning")
+        camera.shutter_speed = int(
+            current_shutter_speed * target_white_level / max_brightness
+        )
+
+
 def auto_expose_and_freeze_settings(camera: PiCamera):
     """Freeze the settings after auto-exposing to white illumination"""
     logging.info("Allowing the camera to auto-expose")
     if "greyworld" in camera.AWB_MODES:
-        print("Calibrating with greyworld AWB")
+        logging.info("Calibrating with greyworld AWB")
         camera.awb_mode = "greyworld"
     else:
-        print("Calibrating with auto AWB")
+        logging.warning("Calibrating with auto AWB as greyworld is missing")
         camera.awb_mode = "auto"
     camera.exposure_mode = "auto"
     camera.iso = (
@@ -77,6 +129,112 @@ def auto_expose_and_freeze_settings(camera: PiCamera):
     adjust_exposure_to_setpoint(camera, 215)
 
 
+def adjust_shutter_and_gain_from_raw(
+    camera: PiCamera,
+    target_white_level: int = 700,
+    max_iterations: int = 20,
+    tolerance: float = 0.05,
+    percentile: float = 99.9,
+) -> float:
+    """Adjust exposure and analog gain based on raw images.
+    
+    This routine is slow but effective.  It uses raw images, so we
+    are not affected by white balance or digital gain.
+    """
+    # Start by setting fully manual exposure.
+    camera.exposure_mode = "off"
+    camera.iso = 0  # We must set ISO=0 (auto) or we can't set gain
+    camera.analog_gain = 1
+    camera.digital_gain = 1
+    camera.shutter_speed = 1  # This is not valid - we'll get the minimum
+    time.sleep(0.5)
+
+    # We start with very low exposure settings and work up in coarse steps
+    # until either the brightness is high enough, or we can't increase the
+    # shutter speed any more.
+    iterations = 0
+    shutter_speed_maximised = False
+    shutter_speed_increments = [10, 2, None]
+    shutter_speed_increment = shutter_speed_increments.pop(0)
+    while iterations < max_iterations:
+        iterations += 1
+        max_brightness = np.max(get_channel_percentiles(camera, percentile))
+        tested_shutter_speed = camera.shutter_speed
+        tested_analog_gain = camera.analog_gain
+        logging.info(
+            f"Brightness: {max_brightness: >5.0f}, "
+            f"Target: {target_white_level: >5.0f}, "
+            f"Gain: {float(tested_analog_gain): >4.1f}, "
+            f"Shutter: {float(tested_shutter_speed): >7.0f}"
+        )
+
+        if abs(max_brightness - target_white_level) < target_white_level * tolerance:
+            logging.info(
+                f"Brightness has converged to within {tolerance * 100 :.0f}% "
+                f"after {iterations} iterations."
+            )
+            break
+
+        # if not shutter_speed_maximised:  # Start by adjusting shutter speed
+        if max_brightness > target_white_level // 2:
+            shutter_speed_increment = (
+                None
+            )  # If we're within a factor of 2, trigger fine-tuning
+        if shutter_speed_increment is None:
+            logging.info("Fine-tuning shutter speed")
+            new_shutter_speed = int(
+                tested_shutter_speed * target_white_level / max_brightness
+            )
+            camera.shutter_speed = new_shutter_speed
+        else:
+            if max_brightness > target_white_level:
+                logging.info("Reducing shutter speed")
+                camera.shutter_speed = int(
+                    tested_shutter_speed / shutter_speed_increment
+                )
+            else:
+                logging.info("Increasing shutter speed")
+                camera.shutter_speed = tested_shutter_speed * shutter_speed_increment
+        time.sleep(0.5)
+
+        # Check whether the shutter speed is still going up - if not, we've hit a maximum
+        current_shutter_speed = camera.shutter_speed
+        if current_shutter_speed == tested_shutter_speed:
+            camera.analog_gain *= 2
+            if camera.analog_gain == tested_analog_gain:
+                logging.info(f"Shutter speed and gain are both maxed out.  Giving up!")
+                break
+            logging.info(
+                f"Shutter speed has maxed out at {current_shutter_speed}, doubled gain."
+            )
+
+    max_brightness = np.max(get_channel_percentiles(camera, percentile))
+    logging.info(f"Brightness {max_brightness} after {iterations} tries.")
+    return max_brightness
+
+
+def adjust_white_balance_from_raw(
+    camera: PiCamera, percentile: float = 99
+) -> (float, float):
+    """Adjust the white balance in a single shot, based on the raw image.
+    
+    NB if ``channels_from_raw_image`` is broken, this will go haywire.
+    We should probably have better logic to verify the channels really
+    are BGGR...
+    """
+    blue, g1, g2, red = get_channel_percentiles(camera, percentile)
+    green = (g1 + g2) / 2.0
+    new_awb_gains = (green / red, green / blue)
+    logging.info(
+        f"Raw white point is R: {red} G: {green} B: {blue}, "
+        f"setting AWB gains to ({new_awb_gains[0]:.2f}, "
+        f"{new_awb_gains[1]:.2f})."
+    )
+    camera.awb_mode = "off"
+    camera.awb_gains = new_awb_gains
+    return new_awb_gains
+
+
 def channels_from_bayer_array(bayer_array: np.ndarray) -> np.ndarray:
     """Given the 'array' from a PiBayerArray, return the 4 channels."""
     bayer_pattern: List[Tuple[int, int]] = [(0, 0), (0, 1), (1, 0), (1, 1)]
@@ -95,6 +253,14 @@ def channels_from_bayer_array(bayer_array: np.ndarray) -> np.ndarray:
     return channels
 
 
+def get_channel_percentiles(camera: PiCamera, percentile: float) -> np.ndarray:
+    """Calculate the brightness percentile of the pixels in each channel"""
+    with PiBayerArray(camera) as output:
+        camera.capture(output, format="jpeg", bayer=True)
+        channels = channels_from_bayer_array(output.array)
+        return np.percentile(channels, percentile, axis=(1, 2))
+
+
 def lst_from_channels(channels: np.ndarray) -> np.ndarray:
     """Given the 4 Bayer colour channels from a white image, generate a LST."""
     full_resolution: np.ndarray = np.array(
@@ -164,6 +330,18 @@ def lst_from_channels(channels: np.ndarray) -> np.ndarray:
     lens_shading_table: np.ndarray = gains.astype(np.uint8)
     return lens_shading_table[::-1, :, :].copy()
 
+def lst_from_camera(camera: PiCamera) -> np.ndarray:
+    """Acquire a raw image and use it to calculate a lens shading table."""
+    with PiBayerArray(camera) as a:
+        camera.capture(a, format="jpeg", bayer=True)
+        raw_image = a.array.copy()
+
+    # Now we need to calculate a lens shading table that would make this flat.
+    # raw_image is a 3D array, with full resolution and 3 colour channels.  No
+    # de-mosaicing has been done, so 2/3 of the values are zero (3/4 for R and B
+    # channels, 1/2 for green because there's twice as many green pixels).
+    channels = channels_from_bayer_array(raw_image)
+    return lst_from_channels(channels)
 
 def recalibrate_camera(camera: PiCamera):
     """Reset the lens shading table and exposure settings.
@@ -178,16 +356,7 @@ def recalibrate_camera(camera: PiCamera):
     camera.lens_shading_table = flat_lens_shading_table(camera)
     _ = rgb_image(camera)  # for some reason the camera won't work unless I do this!
 
-    with PiBayerArray(camera) as a:
+    lens_shading_table = lst_from_camera(camera)
-        camera.capture(a, format="jpeg", bayer=True)
-        raw_image = a.array.copy()
-
-    # Now we need to calculate a lens shading table that would make this flat.
-    # raw_image is a 3D array, with full resolution and 3 colour channels.  No
-    # de-mosaicing has been done, so 2/3 of the values are zero (3/4 for R and B
-    # channels, 1/2 for green because there's twice as many green pixels).
-    channels = channels_from_bayer_array(raw_image)
-    lens_shading_table = lst_from_channels(channels)
 
     camera.lens_shading_table = lens_shading_table
     _ = rgb_image(camera)