Removed old recalibrate utility

from __future__ import print_function

from openflexure_microscope.camera.pi import StreamingCamera

from openflexure_microscope import config

import numpy as np

import time

import os

import logging, sys

logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)

def lens_shading_correction_from_rgb(rgb_array, binsize=64):

    """Calculate a correction to a lens shading table from an RGB image.

    Returns:

        a floating-point table of gains that should multiply the current

        lens shading table.

    """

    full_resolution = rgb_array.shape[:2]

    table_resolution = [(r // binsize) + 1 for r in full_resolution]

    lens_shading = np.zeros([4] + table_resolution, dtype=np.float)

    for i in range(3):

        # We simplify life by dealing with only one channel at a time.

        image_channel = rgb_array[:,:,i]

        iw, ih = image_channel.shape

        ls_channel = lens_shading[int(i*1.6),:,:] # NB there are *two* green channels

        lw, lh = ls_channel.shape

        # The lens shading table is rounded **up** in size to 1/64th of the size of

        # the image.  Rather than handle edge images separately, I'm just going to

        # pad the image by copying edge pixels, so that it is exactly 32 times the

        # size of the lens shading table (NB 32 not 64 because each channel is only

        # half the size of the full image - remember the Bayer pattern...  This

        # should give results very close to 6by9's solution, albeit considerably 

        # less computationally efficient!

        padded_image_channel = np.pad(image_channel, 

                                      [(0, lw*binsize - iw), (0, lh*binsize - ih)],

                                      mode="edge") # Pad image to the right and bottom

        assert padded_image_channel.shape == (lw*binsize, lh*binsize), "padding problem"

        # Next, fill the shading table (except edge pixels).  Please excuse the

        # for loop - I know it's not fast but this code needn't be!

        box = 3 # We average together a square of this side length for each pixel.

        # NB this isn't quite what 6by9's program does - it averages 3 pixels

        # horizontally, but not vertically.

        for dx in np.arange(box) - box//2:

            for dy in np.arange(box) - box//2:

                ls_channel[:,:] += padded_image_channel[binsize//2+dx::binsize,binsize//2+dy::binsize]

        ls_channel /= box**2

        # Everything is normalised relative to the centre value.  I follow 6by9's

        # example and average the central 64 pixels in each channel.

        channel_centre = np.mean(image_channel[iw//2-4:iw//2+4, ih//2-4:ih//2+4])

        print("channel {} centre brightness {}".format(i, channel_centre))

        ls_channel /= channel_centre

        # NB the central pixel should now be *approximately* 1.0 (may not be exactly

        # due to different averaging widths between the normalisation & shading table)

        # For most sensible lenses I'd expect that 1.0 is the maximum value.

        # NB ls_channel should be a "view" of the whole lens shading array, so we don't

        # need to update the big array here.

        print("min {}, max {}".format(ls_channel.min(), ls_channel.max()))

    # What we actually want to calculate is the gains needed to compensate for the 

    # lens shading - that's 1/lens_shading_table_float as we currently have it.

    lens_shading[2, ...] = lens_shading[1, ...] # Duplicate the green channels

    gains = 1.0/lens_shading  # 32 is unity gain

    return gains

def gains_to_lst(gains):

    """Given a lens shading gains table (where no gain=1.0), convert to 8-bit."""

    lst = gains / np.min(gains)*32 # minimum gain is 32 (= unity gain)

    lst[lst > 255] = 255 # clip at 255

    return lst.astype(np.uint8)

def generate_lens_shading_table_closed_loop(output_fname="microscopelst.npy",

                                            n_iterations=5,

                                            images_to_average=5):

    """Reset the camera's parameters, and recalibrate the lens shading to get unifrom images.

    This function requires the microscope to be set up with a blank, uniformly 

    illuminated field of view.  When it runs, it first auto-exposes, then fixes 

    the gains/shutter speed and resets the lens shading correction to a unity

    gain.  Rather than take a single raw image and calibrate from that (as 

    done in the open loop version, which is a more or less direct Python port 

    of 6by9's C code), we do it incrementally.  Each iteration (of a default 5)

    consists of acquiring a processed RGB image, then adjusting the lens shading

    table to make it uniform.  It seems that doing this 3-5 times gives much 

    better results than just doing it once.

    At the end, all camera settings are saved into the output file, where they

    can be used to set up a microscope with `load_microscope`.

    """

    print("Regenerating the camera settings, including lens shading.")

    print("This will only work if the camera is looking at something uniform and white.")

    # Start by loading the raw image from the Pi camera.  This creates a ``picamera.PiBayerArray``.

    openflexurerc = config.load_config()

    with StreamingCamera(config=openflexurerc) as cam:

        lens_shading_table = np.zeros(cam.camera._lens_shading_table_shape(), dtype=np.uint8) + 32

        gains = np.ones_like(lens_shading_table, dtype=np.float)

        max_res = cam.camera.MAX_RESOLUTION

    # Open the microscope and start with flat (i.e. no) lens shading correction.

        cam.start_preview()

        def get_rgb_image():  # shorthand for taking an RGB image

            return cam.array(use_video_port=True, resize=(max_res[0]//2, max_res[1]//2))

        # Adjust the shutter speed until the brightest pixels are giving a set value (say 220)

        for i in range(3):

            cam.camera.shutter_speed = int(cam.camera.shutter_speed * 150.0 / np.max(get_rgb_image()))

            time.sleep(1)

        for i in range(n_iterations):

            print("Optimising lens shading, pass {}/{}".format(i+1, n_iterations))

            # Take an RGB (i.e. processed) image, and calculate the change needed in the shading table

            images = [] #averaging to reduce noise

            for j in range(images_to_average):

                images.append(get_rgb_image())

            rgb_image = np.mean(images, axis=0, dtype=np.float)

            incremental_gains = lens_shading_correction_from_rgb(rgb_image, 64//2)

            gains *= incremental_gains

            # Apply this change (actually apply a bit less than the change)

            cam.stop_stream_recording()

            cam.camera.lens_shading_table = gains_to_lst(gains*32)

            cam.start_stream_recording()

            time.sleep(2)

        # Fix the AWB gains so the image is neutral

        channel_means = np.mean(np.mean(get_rgb_image(), axis=0, dtype=np.float), axis=0)

        old_gains = cam.camera.awb_gains

        cam.camera.awb_gains = (channel_means[1]/channel_means[0] * old_gains[0], channel_means[1]/channel_means[2]*old_gains[1])

        time.sleep(1)

        # Adjust shutter speed to make the image bright but not saturated

        for i in range(3):

            cam.camera.shutter_speed = int(cam.camera.shutter_speed * 230.0 / np.max(get_rgb_image()))

            time.sleep(1)

        # Storing shading table

        lens_shading_table = cam.camera.lens_shading_table

        # Saving shading table to disk

        output_path = os.path.join(config.USER_CONFIG_DIR, output_fname)

        np.save(output_path, lens_shading_table)

        print("Lens shading table written to {}".format(output_path))

        cam.config = {'shading_table_path': output_path}

        settings = cam.config

        for k in settings:

            print("{}: {}".format(k, settings[k]))

        logging.debug("Merging config...")

        config.merge_config(settings, safe=True, backup=True)

if __name__ == '__main__':

    generate_lens_shading_table_closed_loop()