Source code for owl.calibration.color_correction

import math
from time import sleep

import cv2
import numpy as np
import xarray as xr
from numpy.polynomial.polynomial import polygrid2d, polyval2d

from owl.color.colorconv import get_converted_data, get_converted_dataset

from ..mcam_data import bayer_dataset_to_single_channel, get_valid_data
from ..util import ndrange


def tqdm(*_, **__):
    return _[0]


def _make_llt_inv(L_matrix):
    return np.linalg.inv(L_matrix @ L_matrix.T)


def _make_L_matrix(led_values):
    L_matrix = np.ones(shape=(len(led_values), 2)).T
    L_matrix[0, :] = led_values
    return L_matrix


def _calculate_chunk_means(dataset, chunk_size, mono_sensor):
    chunks = (dataset.images.shape[2] // chunk_size,
              dataset.images.shape[3] // chunk_size)

    red_data = bayer_dataset_to_single_channel(dataset, 'red').images.data
    green_data = bayer_dataset_to_single_channel(dataset, 'green').images.data
    blue_data = bayer_dataset_to_single_channel(dataset, 'blue').images.data

    channel_chunk_size = chunk_size // 2
    data_shape = red_data.shape
    new_shape = (
        data_shape[:2] +
        (chunks[0], channel_chunk_size) +
        (chunks[1], channel_chunk_size)
    )

    data = np.empty(shape=(3,) + new_shape, dtype=np.uint8)

    red_data = red_data[..., :chunks[0] * channel_chunk_size,
                        :chunks[1] * channel_chunk_size]
    green_data = green_data[..., :chunks[0] * channel_chunk_size,
                            :chunks[1] * channel_chunk_size]
    blue_data = blue_data[..., :chunks[0] * channel_chunk_size,
                          :chunks[1] * channel_chunk_size]

    data[0] = red_data.reshape(new_shape)
    data[1] = green_data.reshape(new_shape)
    data[2] = blue_data.reshape(new_shape)

    data_means = data.mean(axis=(-1, -3))

    if mono_sensor:
        data_means = data_means.mean(axis=0, keepdims=True)

    return data_means


def _apply_sensor_corrections_to_means(data_means, response_matrix):
    sensor_corrections = create_average_response_correction(response_matrix)
    for i in np.ndindex(data_means.shape[1:]):
        chunk_index = (slice(None),) + i
        sensor_index = i[:2]
        chunk_data = np.concatenate((np.array(data_means[chunk_index]), np.ones(1)))
        data_means[chunk_index] = (sensor_corrections[sensor_index] @ chunk_data)[:-1]


def _calculate_response_columns_rgb(red_means_stack, green_means_stack,
                                    blue_means_stack, led_values):
    index_max = red_means_stack.shape[1:3]
    r_column = np.empty(shape=index_max + (3,), dtype=float)
    b_column = np.empty(shape=index_max + (3,), dtype=float)
    led_matrix = _make_L_matrix(led_values)
    LLt_inv = _make_llt_inv(led_matrix)
    for j, i in ndrange(index_max):
        P_red = red_means_stack[:, j, i]
        R_red = P_red @ led_matrix.T @ LLt_inv

        P_green = green_means_stack[:, j, i]
        R_green = P_green @ led_matrix.T @ LLt_inv

        P_blue = blue_means_stack[:, j, i]
        R_blue = P_blue @ led_matrix.T @ LLt_inv

        r_column[j, i] = np.array([R_red[0], R_green[0], R_blue[0]]).T
        b_column[j, i] = np.array([R_red[1], R_green[1], R_blue[1]]).T

    return r_column, b_column


def _calculate_response_columns_mono(means_stack, led_values):
    index_max = means_stack.shape[1:3]
    r = np.empty(shape=index_max + (1,), dtype=float)
    b = np.empty(shape=index_max + (1,), dtype=float)
    led_matrix = _make_L_matrix(led_values)
    LLt_inv = _make_llt_inv(led_matrix)
    for j, i in ndrange(index_max):
        P_mono = means_stack[:, j, i]
        R_mono = P_mono @ led_matrix.T @ LLt_inv
        r[j, i] = R_mono[0]
        b[j, i] = R_mono[1]
    return r, b


def _define_response_column(mcam, target_channel, *,
                            illumination,
                            number_led_values,
                            led_type='rgb',
                            mono_sensor=False,
                            tqdm=tqdm):
    color = np.zeros(3)
    color[target_channel] = 1

    led_list = getattr(illumination, led_type + '_leds')
    led_value = illumination.find_max_brightness(len(led_list),
                                                 color_ratio=color)[target_channel]
    illumination_color = color * led_value
    illumination.color = illumination_color
    illumination.clear()
    dataset = mcam.acquire_full_field_of_view()
    if dataset['images'].max() >= 255:
        raise ValueError('Sensors are saturated without illumination. '
                         'Reduce either exposure time or ambient lighting.')
    illumination.fill_array(led_type=led_type)
    dataset = mcam.acquire_full_field_of_view()
    illumination.clear()
    # Define the led value that saturates the sensor
    while dataset['images'].max() >= 255:
        led_value /= 2
        illumination_color = color * led_value
        illumination.color = illumination_color
        illumination.fill_array(led_type=led_type)
        dataset = mcam.acquire_full_field_of_view()
    # Ensure max_pixel_value is a numpy array since we want to do math on it
    # otherwise we were getting strange errors with xarray
    # https://github.com/pydata/xarray/issues/9424
    max_pixel_value = np.asarray(dataset['images'].max())
    saturated_led_value = led_value * 255 / max_pixel_value
    led_values = np.linspace(0, .9 * saturated_led_value, number_led_values)
    if mono_sensor:
        means_stack = _create_means_stack_mono(led_values, color, illumination,
                                               mcam, led_type, tqdm)
        r_column, b_column = _calculate_response_columns_mono(means_stack,
                                                              led_values)
    else:
        (red_means_stack,
         green_means_stack,
         blue_means_stack) = _create_means_stack_rgb(led_values, color, illumination,
                                                     mcam, led_type, tqdm)
        r_column, b_column = _calculate_response_columns_rgb(red_means_stack,
                                                             green_means_stack,
                                                             blue_means_stack,
                                                             led_values)
    return r_column, b_column


def _create_means_stack_rgb(led_values, color, illumination, mcam, led_type, tqdm):
    red_means_list = []
    green_means_list = []
    blue_means_list = []
    #  Avoid sensor saturation and led indefinitely increasing
    #  need to use max or a localized mean to avoid localized saturation
    for led_value in tqdm(led_values):
        illumination_color = color * led_value
        illumination.color = illumination_color
        illumination.fill_array(led_type=led_type)
        dataset = mcam.acquire_full_field_of_view()
        illumination.clear()

        red_means = bayer_dataset_to_single_channel(dataset,
                                                    'red').images.data.mean(axis=(-1, -2))
        green_means = bayer_dataset_to_single_channel(dataset,
                                                      'green').images.data.mean(axis=(-1, -2))
        blue_means = bayer_dataset_to_single_channel(dataset,
                                                     'blue').images.data.mean(axis=(-1, -2))

        red_means_list.append(red_means)
        green_means_list.append(green_means)
        blue_means_list.append(blue_means)

    stack_height = len(red_means_list)
    array_shape = red_means_list[0].shape
    red_means_stack = np.reshape(red_means_list, (stack_height,) + array_shape)
    green_means_stack = np.reshape(green_means_list, (stack_height,) + array_shape)
    blue_means_stack = np.reshape(blue_means_list, (stack_height,) + array_shape)
    return red_means_stack, green_means_stack, blue_means_stack


def _create_means_stack_mono(led_values, color, illumination, mcam, led_type, tqdm):
    means_list = []
    #  Avoid sensor saturation and led indefinitely increasing
    #  need to use max or a localized mean to avoid localized saturation
    for led_value in tqdm(led_values):
        illumination_color = color * led_value
        illumination.color = illumination_color
        illumination.fill_array(led_type=led_type)
        dataset = mcam.acquire_full_field_of_view()
        illumination.clear()

        means = dataset.images.data.mean(axis=(-1, -2))

        means_list.append(means)

    stack_height = len(means_list)
    array_shape = means_list[0].shape
    means_stack = np.reshape(means_list, (stack_height,) + array_shape)
    return means_stack


def _define_response_column_rails(mcam, illumination, *,
                                  number_led_values,
                                  mono_sensor=False,
                                  tqdm=tqdm):

    illumination.clear()
    dataset = mcam.acquire_full_field_of_view()
    if dataset['images'].max() >= 255:
        raise ValueError('Sensors are saturated without illumination. '
                         'Reduce either exposure time or ambient lighting.')
    brightness_fraction = 1
    illumination.set_brightness(brightness_fraction)
    dataset = mcam.acquire_full_field_of_view()
    illumination.clear()
    # Define the led value that saturates the sensor
    while dataset['images'].max() >= 255:
        brightness_fraction /= 2
        illumination.set_brightness(brightness_fraction)
        dataset = mcam.acquire_full_field_of_view()
    max_pixel_value = np.asarray(dataset['images']).max()
    saturated_brightness_fraction = min(1, brightness_fraction * 255 / max_pixel_value)
    brightness_fractions = np.linspace(0, .9 * saturated_brightness_fraction, number_led_values)
    if mono_sensor:
        means_stack = _create_means_stack_mono_rails(brightness_fractions,
                                                     illumination,
                                                     mcam,
                                                     tqdm,)
        r_column, b_column = _calculate_response_columns_mono(means_stack,
                                                              brightness_fractions)
    else:
        (red_means_stack,
         green_means_stack,
         blue_means_stack) = _create_means_stack_rgb_rails(brightness_fractions,
                                                           illumination,
                                                           mcam,
                                                           tqdm)
        r_column, b_column = _calculate_response_columns_rgb(red_means_stack,
                                                             green_means_stack,
                                                             blue_means_stack,
                                                             brightness_fractions)
    return r_column, b_column


def _create_means_stack_rgb_rails(led_values, illumination, mcam, tqdm,
                                  illumination_mode=None, color_ratio=None):
    red_means_list = []
    green_means_list = []
    blue_means_list = []
    #  Avoid sensor saturation and led indefinitely increasing
    #  need to use max or a localized mean to avoid localized saturation
    for led_value in tqdm(led_values):
        if illumination_mode is not None and color_ratio is not None:
            illumination.set_brightness(led_value,
                                        illumination_mode=illumination_mode,
                                        color_ratio=color_ratio)
        else:
            illumination.set_brightness(led_value)
            # need to wait for light to turn on
            sleep(1)
        dataset = mcam.acquire_full_field_of_view()

        red_means = bayer_dataset_to_single_channel(dataset,
                                                    'red').images.data.mean(axis=(-1, -2))
        green_means = bayer_dataset_to_single_channel(dataset,
                                                      'green').images.data.mean(axis=(-1, -2))
        blue_means = bayer_dataset_to_single_channel(dataset,
                                                     'blue').images.data.mean(axis=(-1, -2))

        red_means_list.append(red_means)
        green_means_list.append(green_means)
        blue_means_list.append(blue_means)

    illumination.clear()
    stack_height = len(red_means_list)
    array_shape = red_means_list[0].shape
    red_means_stack = np.reshape(red_means_list, (stack_height,) + array_shape)
    green_means_stack = np.reshape(green_means_list, (stack_height,) + array_shape)
    blue_means_stack = np.reshape(blue_means_list, (stack_height,) + array_shape)
    return red_means_stack, green_means_stack, blue_means_stack


def _create_means_stack_mono_rails(led_values, illumination, mcam, tqdm,
                                   illumination_mode=None, color_ratio=None):
    means_list = []
    #  Avoid sensor saturation and led indefinitely increasing
    #  need to use max or a localized mean to avoid localized saturation
    for led_value in tqdm(led_values):
        if illumination_mode is not None and color_ratio is not None:
            illumination.set_brightness(led_value,
                                        illumination_mode=illumination_mode,
                                        color_ratio=color_ratio)
        else:
            illumination.set_brightness(led_value)
            # need to wait for light to turn on
            sleep(1)
        dataset = mcam.acquire_full_field_of_view()

        means = dataset.images.data.mean(axis=(-1, -2))

        means_list.append(means)

    illumination.clear()
    stack_height = len(means_list)
    array_shape = means_list[0].shape
    means_stack = np.reshape(means_list, (stack_height,) + array_shape)
    return means_stack


def define_response_matrix(mcam,
                           illumination,
                           number_led_values=15,
                           led_type='rgb',
                           lighting_channels=(0, 1, 2),
                           mono_sensor=False,
                           tqdm=tqdm):
    """Define the photometric response of each sensor with a 4x4 matrix.

    Parameters
    ----------
    mcam : owl.instruments.MCAM
        Connection to the MCAM unit.
    illumination : string ('reflection', 'transmission')
        The illumination device to use to illuminate the sensors.
    number_led_values : int
        The number of led values for each led channel to define the response matrix.
    led_type : str, optional
        The kind of LED being characterized (e.g. ``'rgb'``).
    lighting_channels : tuple of int, optional
        The illumination channels to iterate over when probing responses.
    mono_sensor : bool, optional
        If True, treat the sensors as monochrome rather than 4-channel.
    tqdm : optional
        Progress-bar callable used to wrap the inner sweep.

    Returns
    -------
    response_matrix : numpy array
        The matrix that describes the sensor's response to the leds. It is a
        M x N x 4 x 4 array of float where M and N are the shape of the sensor array.

    """
    if mono_sensor:
        N_rows = 2
    else:
        N_rows = 4
    N_columns = 4
    response_matrix = np.zeros(mcam.N_cameras + (N_rows, N_columns), dtype=float)
    b = np.zeros(mcam.N_cameras + (N_rows - 1,))
    for target_channel in lighting_channels:
        (response_column,
         b_temp) = _define_response_column(mcam=mcam,
                                           target_channel=target_channel,
                                           illumination=illumination,
                                           number_led_values=number_led_values,
                                           led_type=led_type,
                                           mono_sensor=mono_sensor,
                                           tqdm=tqdm)
        response_matrix[:, :, :-1, target_channel] = response_column
        b += b_temp
    b /= N_columns - 1
    response_matrix[:, :, :-1, -1] = b
    response_matrix[:, :, -1, -1] = 1

    return response_matrix


def define_response_matrix_rails(mcam,
                                 illumination,
                                 number_led_values=15,
                                 mono_sensor=False,
                                 tqdm=tqdm):
    """Define the photometric response of each sensor with a 4x4 matrix.

    Parameters
    ----------
    mcam : owl.instruments.MCAM
        Connection to the MCAM unit.
    illumination : owl.instruments.Fluorescence
        Object to control the illumination source.
    number_led_values : int
        The number of led values for each led channel to define the response matrix.
    mono_sensor : bool, optional
        If True, treat the sensors as monochrome rather than 4-channel.
    tqdm : optional
        Progress-bar callable used to wrap the inner sweep.

    Returns
    -------
    response_matrix : numpy array
        The matrix that describes the sensor's response to the leds. It is a
        M x N x 4 x 4 array of float where M and N are the shape of the sensor array.

    """
    if mono_sensor:
        N_rows = 2
    else:
        N_rows = 4
    N_columns = 4
    response_matrix = np.zeros(mcam.N_cameras + (N_rows, N_columns), dtype=float)
    (response_column,
     b) = _define_response_column_rails(mcam=mcam,
                                        illumination=illumination,
                                        number_led_values=number_led_values,
                                        mono_sensor=mono_sensor,
                                        tqdm=tqdm,
                                        )
    response_matrix[:, :, :-1, 0] = response_column
    response_matrix[:, :, :-1, -1] = b
    response_matrix[:, :, -1, -1] = 1

    return response_matrix


def create_response_means_stack(response_matrix,
                                mcam,
                                illumination_type,
                                lighting_channels,
                                chunk_size=40,
                                *,
                                led_type='rgb',
                                number_led_values=5,
                                mono_sensor=False,
                                tqdm=tqdm):
    """Create a binned pixel response that is used to create pixel corrections.

    Pixel response is taken of the sensor corrected values

    Parameters
    ----------
    response_matrix : numpy array
        M x N x 4 x 4 array of float where M and N are the shape of the sensor array.
    mcam : owl.instruments.MCAM
        Connection to the MCAM unit.
    illumination_type : string ('reflection', 'transmission')
        Which illumination board to use to illuminate the sensors.
    lighting_channels : sequence of int
        The illumination color channels exercised during the sweep.
    chunk_size : int
        Size in pixels of the subarrays to break the images into to determine
        the local response of the sensors.
    led_type : str, optional
        The kind of LED being driven (e.g. ``'rgb'``).
    number_led_values : int, optional
        Number of LED brightness steps to sample across the sweep.
    mono_sensor : bool, optional
        If True, treat the sensors as monochrome rather than 4-channel.
    tqdm : optional
        Progress-bar callable used to wrap the LED sweep loop.

    Returns
    -------
    response_means_stack : xarray dataset
        Binned pixel values produce by a series of different led values.

    """
    color = define_white_light(response_matrix)

    illumination_color = np.array((0, 0, 0))
    illumination = getattr(mcam, illumination_type + '_illumination')
    illumination.color = illumination_color
    illumination.fill_array(led_type=led_type)
    mcam.acquire_full_field_of_view()
    illumination.clear()
    means_chunk_list = []
    led_factors_list = []
    unit_pixel_response = (response_matrix @ np.array(color + (1,)))[..., :3]
    saturation_factor = (255 / unit_pixel_response).min()
    for i in tqdm(range(number_led_values)):
        led_factor = saturation_factor * i / number_led_values
        illumination_color[lighting_channels,] = np.array(color)[lighting_channels,] * led_factor
        illumination.color = illumination_color
        illumination.fill_array(led_type=led_type)
        dataset = mcam.acquire_full_field_of_view()
        illumination.clear()
        data_means = _calculate_chunk_means(dataset, chunk_size=chunk_size, mono_sensor=mono_sensor)
        _apply_sensor_corrections_to_means(data_means,
                                           response_matrix)
        means_chunk_list.append(data_means)
        led_factors_list.append(led_factor)

    chunk_stack_height = len(means_chunk_list)
    led_stack = np.reshape(led_factors_list, chunk_stack_height)
    led_index = np.arange(chunk_stack_height)
    dims = ['led_index', 'color_channels', 'image_y', 'image_x', 'chunk_y', 'chunk_x']
    chunk_array_shape = means_chunk_list[0].shape
    chunk_y_index = (np.arange(chunk_array_shape[3]) + .5) * chunk_size
    chunk_x_index = (np.arange(chunk_array_shape[4]) + .5) * chunk_size
    means_chunk_stack = np.reshape(means_chunk_list,
                                   (chunk_stack_height,) + chunk_array_shape)
    means_dataarray = xr.DataArray(means_chunk_stack,
                                   dims=dims,
                                   coords=[led_index,
                                           np.arange(chunk_array_shape[0]),
                                           np.arange(chunk_array_shape[1]),
                                           np.arange(chunk_array_shape[2]),
                                           chunk_y_index,
                                           chunk_x_index])
    led_factors_dataarray = xr.DataArray(led_stack,
                                         dims=['led_index'],
                                         coords=[led_index])
    led_factors_dataarray['led_color'] = str(color)
    response_means_stack = xr.Dataset({'data_means': means_dataarray,
                                       'led_factors': led_factors_dataarray})
    return response_means_stack


def create_response_means_stack_rails(response_matrix,
                                      mcam,
                                      illumination,
                                      chunk_size=40,
                                      *,
                                      number_led_values=5,
                                      mono_sensor=False,
                                      tqdm=tqdm):

    means_chunk_list = []
    led_factors_list = []
    brightness_fraction = 1
    illumination.set_brightness(brightness_fraction=brightness_fraction)
    dataset = mcam.acquire_full_field_of_view()
    illumination.clear()
    # Define the led value that saturates the sensor
    while dataset['images'].max() >= 255:
        brightness_fraction /= 2
        illumination.set_brightness(brightness_fraction=brightness_fraction)
        dataset = mcam.acquire_full_field_of_view()
    max_pixel_value = np.asarray(dataset['images']).max()
    saturated_brightness_fraction = min(1, brightness_fraction * 255 / max_pixel_value)
    brightness_fractions = np.linspace(0, .9 * saturated_brightness_fraction, number_led_values)
    for brightness_fraction in tqdm(brightness_fractions):
        illumination.set_brightness(brightness_fraction=brightness_fraction)
        # need to wait for light to turn on
        sleep(1)
        dataset = mcam.acquire_full_field_of_view()
        data_means = _calculate_chunk_means(dataset, chunk_size=chunk_size, mono_sensor=mono_sensor)
        _apply_sensor_corrections_to_means(data_means,
                                           response_matrix)
        means_chunk_list.append(data_means)
        led_factors_list.append(brightness_fraction)

    illumination.clear()
    chunk_stack_height = len(means_chunk_list)
    led_stack = np.reshape(led_factors_list, chunk_stack_height)
    led_index = np.arange(chunk_stack_height)
    dims = ['led_index', 'color_channels', 'image_y', 'image_x', 'chunk_y', 'chunk_x']
    chunk_array_shape = means_chunk_list[0].shape
    chunk_y_index = (np.arange(chunk_array_shape[3]) + .5) * chunk_size
    chunk_x_index = (np.arange(chunk_array_shape[4]) + .5) * chunk_size
    means_chunk_stack = np.reshape(means_chunk_list,
                                   (chunk_stack_height,) + chunk_array_shape)
    means_dataarray = xr.DataArray(means_chunk_stack,
                                   dims=dims,
                                   coords=[led_index,
                                           np.arange(chunk_array_shape[0]),
                                           np.arange(chunk_array_shape[1]),
                                           np.arange(chunk_array_shape[2]),
                                           chunk_y_index,
                                           chunk_x_index])
    led_factors_dataarray = xr.DataArray(led_stack,
                                         dims=['led_index'],
                                         coords=[led_index])
    response_means_stack = xr.Dataset({'data_means': means_dataarray,
                                       'led_factors': led_factors_dataarray})
    return response_means_stack


def _polynomial_coefficients(chunk_values, deg):
    #  need to normalize x, y coords for best fitting
    y = (np.arange(chunk_values.shape[0]) + .5) / chunk_values.shape[0]
    x = (np.arange(chunk_values.shape[1]) + .5) / chunk_values.shape[1]
    X, Y = np.array(np.meshgrid(x, y))
    #  create a pseudo-vandermonde matrix
    vander = np.zeros(shape=(X.size, math.factorial(deg + 1)), dtype=float)
    for r, index in enumerate(ndrange(X.shape)):
        c = 0
        # hardcode (4, 4) array shape to match array shape of photometric response
        # this limits us to at max 3rd degree polynomials, but this should be fine
        coefficient_array_shape = (4, 4)
        for x_deg, y_deg in ndrange(coefficient_array_shape):
            if x_deg + y_deg <= deg:
                vander[r, c] = X[index] ** x_deg * Y[index] ** y_deg
                c += 1
    chunk_values_vector = chunk_values.reshape((vander.shape[0],))
    coefficients = np.linalg.lstsq(vander, chunk_values_vector, rcond=None)[0]
    """ reorganize coefficients into an array shape to be accepted by polyval2d
        [[x0y0, x0y1, x0y2, x0y3],
         [x1y0, x1y1, x1y2, x1y3],
         [x2y0, x2y1, x2y2, x2y3],
         [x3y0, x3y1, x3y2, x3y3]] """
    coefficients_array = np.zeros(shape=coefficient_array_shape, dtype=float)
    c = 0
    for j, i in ndrange(coefficients_array.shape):
        if j + i <= deg:
            coefficients_array[j, i] = coefficients[c]
            c += 1
    return coefficients_array


def _define_chunk_responses(color_means, led_values):
    response_offset_array = np.zeros(shape=color_means.shape[1:], dtype=float)
    response_coefficient_array = np.zeros(shape=color_means.shape[1:], dtype=float)
    L_matrix = _make_L_matrix(led_values)
    LLt_inv = _make_llt_inv(L_matrix)
    for chunk_index in ndrange(color_means.shape[1:]):
        y, x, y_chunk, x_chunk = chunk_index
        chunk_pixel_means = color_means[:, y, x, y_chunk, x_chunk]
        chunk_response = chunk_pixel_means.T @ L_matrix.T @ LLt_inv
        response_coefficient_array[chunk_index] = chunk_response[0]
        response_offset_array[chunk_index] = chunk_response[1]
    return response_coefficient_array, response_offset_array


def create_pixel_polynomial_coefficients(response_means_stack, deg=2,
                                         *, tqdm=tqdm):
    """Create polynomial coefficients that describe the pixels response to the led.

    Parameters
    ----------
    response_means_stack : xarray dataset
        Binned pixel values produce by a series of different led values.
    deg : int
        The degree of polynomial used to describe surfaces. The maximum
        accepted value is 3.
    tqdm : optional
        Progress-bar callable used to wrap the per-sensor loop.

    Returns
    -------
    sensor_polyco_coefficient : numpy array
        An MxNx4x4 array of polynomial coefficients describing the pixel
        coefficient corrections where M and N correspond to the shape of the image array.
    sensor_polyco_offset : numpy array
        An MxNx4x4 array of polynomial coefficients describing the pixel
        offset corrections where M and N correspond to the shape of the image array.

    """
    if deg > 3:
        raise ValueError('deg value must be an integer less than 4.')
    led_factors = response_means_stack.led_factors.data
    # polynomial coefficient arrays are set to be shape (4, 4) to match the
    # photometric response array shapes
    number_color_channels = response_means_stack.data_means.shape[1]
    camera_polyco_shape = (response_means_stack.data_means.shape[2:4] +
                           (4, 4) + (number_color_channels,))
    sensor_polyco_coefficient = np.zeros(shape=camera_polyco_shape, dtype=float)
    sensor_polyco_offset = np.zeros(shape=camera_polyco_shape, dtype=float)

    def _get_surface_mean(matrix):
        # assumes a 2D matrix and that the surface goes from (0, 0) to (1, 1)
        mean = 0
        for i, j in np.ndindex(matrix.shape):
            mean += matrix[i, j] / ((i + 1) * (j + 1))
        return mean

    for color_channel in range(number_color_channels):
        (chunks_response_coefficient_array,
         chunks_response_offset_array) = _define_chunk_responses(
            response_means_stack.data_means.data[:, color_channel],
            led_factors)
        for camera_index in tqdm(ndrange(camera_polyco_shape[:2]), desc='Per camera coefficient'):
            sensor_polyco_coefficient[camera_index][..., color_channel] = _polynomial_coefficients(
                chunks_response_coefficient_array[camera_index],
                deg=deg)
            sensor_polyco_coefficient[camera_index][..., color_channel] /= _get_surface_mean(
                sensor_polyco_coefficient[camera_index][..., color_channel]
            )
            sensor_polyco_offset[camera_index][..., color_channel] = _polynomial_coefficients(
                chunks_response_offset_array[camera_index],
                deg=deg)
            sensor_polyco_offset[camera_index][0, 0, color_channel] -= _get_surface_mean(
                sensor_polyco_offset[camera_index][..., color_channel]
            )
    return sensor_polyco_coefficient, sensor_polyco_offset


def _create_single_pixel_correction(polynomial_coefficients, X, Y):
    corrections = polyval2d(X, Y, polynomial_coefficients)
    return corrections


def _create_array_pixel_correction(polynomial_coefficients,
                                   X, Y, tqdm=tqdm):
    array_shape = polynomial_coefficients.shape[:2]
    corrections = np.zeros(shape=array_shape + X.shape, dtype=np.float32)
    for camera_index in tqdm(ndrange(array_shape)):
        corrections[camera_index] = _create_single_pixel_correction(
            polynomial_coefficients[camera_index],
            X, Y)
    return corrections


[docs] def create_pixel_corrections(coefficient_polynomial_coefficients, offset_polynomial_coefficients, *, image_shape, tqdm=tqdm): """ Create a coefficient and offset pixel correction based on the polynomial coefficients. Parameters ---------- coefficient_polynomial_coefficients : numpy array An MxNx4x4 array of polynomial coefficients describing the pixel coefficient corrections. offset_polynomial_coefficients : numpy array An MxNx4x4 array of polynomial coefficients describing the pixel offset corrections. image_shape : tuple The desired shape of the images in pixel (y_pixels, x_pixels). tqdm : optional Progress-bar callable used to wrap the per-camera loop. Returns ------- coefficient_corrections : numpy array Array of pixel correction coefficients of shape M x N x image_shape. offset_corrections : numpy array Array of pixel correction offsets of shape M x N x image_shape. """ channel_number = coefficient_polynomial_coefficients.shape[-1] array_shape = coefficient_polynomial_coefficients.shape[:2] corrections_shape = array_shape + image_shape + (channel_number,) coefficient_corrections = np.zeros(shape=corrections_shape, dtype=np.float32) offset_corrections = np.zeros(shape=corrections_shape, dtype=np.float32) y = (np.arange(image_shape[0], dtype=np.float32) + .5) / (image_shape[0]) x = (np.arange(image_shape[1], dtype=np.float32) + .5) / (image_shape[1]) X, Y = np.array(np.meshgrid(x, y), dtype=np.float32) for channel_index in ndrange(channel_number): tmp_coefficient = _create_array_pixel_correction( coefficient_polynomial_coefficients[..., channel_index], X, Y, tqdm=tqdm) tmp_offset = _create_array_pixel_correction( offset_polynomial_coefficients[..., channel_index], X, Y, tqdm=tqdm) for array_index in ndrange(coefficient_corrections.shape[:2]): tmp_coefficient[array_index] = 1 / tmp_coefficient[array_index] tmp_offset[array_index] = -tmp_offset[array_index] coefficient_corrections[..., channel_index] = tmp_coefficient[..., None] offset_corrections[..., channel_index] = tmp_offset[..., None] return coefficient_corrections, offset_corrections
def _remove_unused_channels(response_matrix): # the response matrix is a image_y, image_x, 4, 4 if on a rgb sensor mcam # or image_y, image_x, 2, 4 on a mono sensor mcam used_channels = [] for c in range(response_matrix.shape[-1]): if np.sum(response_matrix[..., c]) != 0: used_channels.append(c) # the last column will always be added because the bottom right entry is always one. return response_matrix[..., used_channels] def _create_response_correction(current_response, desired_response): # 20231023 Jed - we have changed fluorescence response matrices to be 4x4, # but only one column is populated. This causes the correction pipeline to # make corrects that alter the color of the images. To avoid this, we remove # the unused columns to make a 4x2 matrix for each response matrix so that # the correction is just to unify those column values. current_response = _remove_unused_channels(current_response) desired_response = _remove_unused_channels(desired_response) correction_array_dimension = current_response.shape[2] corrections = np.zeros( shape=(current_response.shape[:2] + (correction_array_dimension, correction_array_dimension)), dtype=current_response.dtype) if current_response.shape[-2:] == (4, 2): # Using the pinv method on a 4x2 creates a 4x4 matrix with the first 3 rows equivalent. # This cause the images to become grayscale after white light corrections. To avoid this # we instead use the inverse of the responses times the desired response (element by # element) and place the results in the matrix diagonal so that it can continue to be # applied through matrix multiplication. adjustments = desired_response[:3, 0] / current_response[..., :3, 0] corrections[..., 0, 0] = adjustments[..., 0] corrections[..., 1, 1] = adjustments[..., 1] corrections[..., 2, 2] = adjustments[..., 2] corrections[..., :-1, -1] = desired_response[:1, -1] - current_response[..., :1, -1] else: for i in ndrange(current_response.shape[:2]): corrections[i] = desired_response @ np.linalg.pinv(current_response[i]) return corrections def create_photometric_corrections(current_response, scale=None): # To maintain the brightness of the pixels after the correction we do not want to create the # correction to the identity matrix but rather to the identity matrix times some constant. # To find an appropriate constant we use the root square mean of the response of each of the # sensors responses to the unified illumination (all ones) and then take the mean through # all sensors. if scale is None: channel_responses = current_response.sum(axis=-1) scale = np.sqrt((channel_responses[..., :-1] ** 2).mean(axis=-1)).mean() # 20231023 Jed - we have changed fluorescence response matrices to be 4x4, # but only one column is populated. This causes the correction pipeline to # make corrects that alter the color of the images. To avoid this, we remove # the unused columns to make a 4x2 matrix for each response matrix so that # the correction is just to unify those column values. current_response = _remove_unused_channels(current_response) if current_response.shape[-1] == current_response.shape[-2]: # RGB sensor, RGB Response measurement desired_response = np.eye(current_response.shape[-1]) * scale else: # Mono sensor desired_response = np.zeros_like(current_response[0, 0]) desired_response[:-1, :-1] = scale / (desired_response.shape[-1] - 1) desired_response[-1, -1] = 1 photometric_corrections = _create_response_correction(current_response, desired_response) return photometric_corrections def create_monochrome_corrections(current_response, gray_ratios=(1, 1, 1)): ratio_sum = sum(gray_ratios) desired_response = np.eye(4) desired_response[:, 0] = gray_ratios[0] / ratio_sum desired_response[:, 1] = gray_ratios[1] / ratio_sum desired_response[:, 2] = gray_ratios[2] / ratio_sum monochrome_corrections = _create_response_correction(current_response, desired_response) return monochrome_corrections def create_average_response_correction(current_response, target_sensor=None, white_balance=False): if target_sensor is None: desired_response = current_response.mean(axis=(0, 1)) else: desired_response = current_response[target_sensor] if white_balance: average_response_correction = create_photometric_corrections(current_response, scale=None) else: average_response_correction = _create_response_correction(current_response, desired_response) return average_response_correction def define_white_light(response_matrix): """Gives normalized to unit vector led values to produce white light. Parameters ---------- response_matrix : numpy array Array of float that model the responses of the sensors. The array is shape (image_y, image_x, 4, 4). Returns ------- led_values : tuple The unit vector values of led values to produce white light. """ pixels = np.ones(shape=response_matrix.shape[-2] - 1) * 128 led_values = calculate_led_for_desired_pixel(response_matrix, pixels=pixels) led_values = led_values / np.sqrt((led_values ** 2).sum()) return tuple(led_values)
[docs] def calculate_led_for_desired_pixel(response_matrix, pixels): """Gives led values that will produce the given pixel ratio. Parameters ---------- response_matrix : numpy array Array of float that model the responses of the sensors. The array is shape (image_y, image_x, 4, 4). pixels : tuple Ratio of the desired pixel values. It must be length 3. Returns ------- led_values : tuple The values of led values to produce near the desired pixel values on white paper. """ pix_vec = np.concatenate((np.array(pixels), np.ones(1)))[..., None] if pix_vec.shape[0] != response_matrix.shape[-2]: raise ValueError('Value `pixels` must match shape of given response matrix. ' f'Given `pixels` has length of {len(pixels)} while response ' f'matrix has shape {response_matrix.shape}. Expected a pixel ' f'value of length {response_matrix.shape[-2] - 1}') led_values_shape = response_matrix.shape[:2] + (response_matrix.shape[-1] - 1,) led_values = np.zeros(shape=led_values_shape, dtype=float) for i in ndrange(response_matrix.shape[:2]): led_values[i] = (np.linalg.pinv(response_matrix[i]) @ pix_vec)[:-1, 0] led_values = led_values.mean(axis=(0, 1)) return led_values
def get_corrected_data(dataset, *, white_balance): response_matrix = dataset.photometric_response.data sensor_corrections = create_sensor_corrections(response_matrix, white_balance=white_balance) return get_converted_data(dataset, sensor_corrections) def get_corrected_dataset(dataset, white_balance): """Apply photometric response to the dataset images variable Parameters ---------- dataset : xarray Dataset MCAM data containing image data as well as additional metadata. Must have `photometric_response` included in the metadata for the corrections to be applied. Image data should not be bayered. white_balance : bool If color correction should be applied to whiten image. If white light illumination was used then this should be True, otherwise it should be False. Returns ------- corrected_dataset : xarray Dataset """ if 'photometric_response' not in dataset: return dataset response_matrix = dataset.photometric_response.data sensor_corrections = create_sensor_corrections(response_matrix, white_balance=white_balance) corrected_dataset = get_converted_dataset(dataset, sensor_corrections) response_shape = dataset.photometric_response.data.shape[-2:] photometric_response_data = np.zeros_like( corrected_dataset.photometric_response ) if response_shape == (4, 2): photometric_response_data[...] = np.asarray([ [1, 0], # red [1, 0], # green [1, 0], # blue [0, 1]], # noqa dtype=photometric_response_data.dtype ) else: photometric_response_data[...] = np.eye( N=response_shape[-2], # Number of rows in the output. M=response_shape[-1], # Number of columns in the output. dtype=photometric_response_data.dtype ) corrected_dataset['photometric_response'] = xr.Variable( corrected_dataset['photometric_response'].dims, photometric_response_data, ) return corrected_dataset def get_corrected_dataset_stack(dataset_stack, white_balance=False): response_matrix = dataset_stack.photometric_response.data sensor_corrections = create_sensor_corrections(np.asarray(response_matrix), white_balance=white_balance) corrected_dataset_stack = get_converted_dataset( dataset_stack, sensor_corrections, ) response_shape = dataset_stack.photometric_response.shape[-2:] new_response_matrix = np.zeros_like(response_matrix) if response_shape[-2:] == (4, 2): new_response_matrix[:] = np.asarray([ [1, 0], # red [1, 0], # green [1, 0], # blue [0, 1]], # noqa dtype=response_matrix.dtype ) else: new_response_matrix[:] = np.eye( N=response_shape[-2], # Number of rows in the output. M=response_shape[-1], # Number of columns in the output. dtype=response_matrix.dtype ) dims = dataset_stack.photometric_response.dims corrected_dataset_stack['photometric_response'] = xr.DataArray( data=new_response_matrix, name='photometric_response', dims=dims, ) return corrected_dataset_stack def create_sensor_corrections(response_matrix, white_balance=False): return create_average_response_correction( response_matrix, white_balance=white_balance, ) def get_pixel_corrected_dataset(dataset): """Apply pixel correction to the dataset images variable Parameters ---------- dataset : xarray Dataset MCAM data containing image data as well as additional metadata. Must have `pixel_response_coefficient` and `pixel_response_offset` included in the metadata for the corrections to be applied. Image data should not be bayered and should already have sensor corrections applied. Returns ------- pixel_corrected_dataset : xarray Dataset """ (coefficient_polynomial_coefficients, offset_polynomial_coefficients) = get_crop_binned_pixel_polynomial(dataset) images_variable = dataset.images.variable image_shape = images_variable.shape[2:4] y = (np.arange(image_shape[0], dtype=np.float32) + .5) / (image_shape[0]) x = (np.arange(image_shape[1], dtype=np.float32) + .5) / (image_shape[1]) images_corrected = np.empty_like(images_variable) N_cameras = dataset.images.shape[:2] for i in np.ndindex(N_cameras): coefficient_corrections = 1 / polygrid2d( y, x, coefficient_polynomial_coefficients[i].mean(axis=-1).T )[..., np.newaxis] offset_corrections = polygrid2d( y, x, offset_polynomial_coefficients[i].mean(axis=-1).T )[..., np.newaxis] image = np.asarray(images_variable[i], dtype='float32') images_corrected[i] = np.clip(( image * coefficient_corrections - offset_corrections ), 0, 255) pixel_corrected_dataset = dataset.copy(deep=False) pixel_corrected_dataset['images'] = (dataset.images.dims, images_corrected) new_coefficient = np.zeros_like(pixel_corrected_dataset.pixel_response_coefficient) new_coefficient[..., 0, 0] = 1 pixel_corrected_dataset['pixel_response_coefficient'] = ( pixel_corrected_dataset.pixel_response_coefficient.dims, new_coefficient, ) new_offset = np.zeros_like(pixel_corrected_dataset.pixel_response_offset) pixel_corrected_dataset['pixel_response_offset'] = ( pixel_corrected_dataset.pixel_response_offset.dims, new_offset, ) return pixel_corrected_dataset def correct_for_new_external_light(mcam): illumination_types = ('transmission', 'reflection') # make sure all light are off for illumination_type in illumination_types: light = getattr(mcam, illumination_types + '_illumination') if light is not None: light.clear() dataset = mcam.acquire_full_field_of_view() red_offset = bayer_dataset_to_single_channel(dataset, 'red').images.data.mean(axis=(-1, -2)) green_offset = bayer_dataset_to_single_channel(dataset, 'green').images.data.mean(axis=(-1, -2)) blue_offset = bayer_dataset_to_single_channel(dataset, 'blue').images.data.mean(axis=(-1, -2)) for illumination_type in illumination_types: response_matrix_key = illumination_type + '_photometric_response' if response_matrix_key in dataset: response_matrix = dataset[response_matrix_key] response_matrix[..., 0, 3] = red_offset response_matrix[..., 1, 3] = green_offset response_matrix[..., 2, 3] = blue_offset def get_crop_binned_pixel_polynomial(dataset): coefficient_polynomial_coefficients = \ dataset['pixel_response_coefficient'].data offset_polynomial_coefficients = \ dataset['pixel_response_offset'].data photometric_start_pixel_y = dataset['photometric_start_pixel_y'].data photometric_end_pixel_y = dataset['photometric_end_pixel_y'].data photometric_start_pixel_x = dataset['photometric_start_pixel_x'].data photometric_end_pixel_x = dataset['photometric_end_pixel_x'].data # Get shift factors calibrated_image_height = photometric_end_pixel_y - photometric_start_pixel_y calibrated_image_width = photometric_end_pixel_x - photometric_start_pixel_x shift_y = (photometric_start_pixel_y - dataset.y.data[0]) / calibrated_image_height shift_x = (photometric_start_pixel_x - dataset.x.data[0]) / calibrated_image_width # Get scale factors binning_y = dataset.y.data[1] - dataset.y.data[0] binning_x = dataset.x.data[1] - dataset.x.data[0] image_height = dataset.y.data[-1] - dataset.y.data[0] + binning_y image_width = dataset.x.data[-1] - dataset.x.data[0] + binning_x scale_y = image_height / calibrated_image_height scale_x = image_width / calibrated_image_width coefficient_polynomial_coefficients = shift_and_scale_array_polynomial_matrix( coefficient_polynomial_coefficients, (shift_y, shift_x), (scale_y, scale_x)) offset_polynomial_coefficients = shift_and_scale_array_polynomial_matrix( offset_polynomial_coefficients, (shift_y, shift_x), (scale_y, scale_x)) return coefficient_polynomial_coefficients, offset_polynomial_coefficients def shift_and_scale_array_polynomial_matrix(polynomial_matrix, shift, scale): # shift is proportional to the full image shape # movement to right and down (cropping is negative) polynomial_matrix = _shift_polynomial_matrix( polynomial_matrix, shift[0], shift[1]) polynomial_matrix = _scale_polynomial_matrix( polynomial_matrix, scale[0], scale[1]) return polynomial_matrix def _scale_polynomial_matrix(polynomial_matrix, scale_y, scale_x): # This matrix must be a collection of at least 3x3 matrices that hold 2nd order 2d polynomials. # The values are stored as such: # [ c_00 c_01 c_02] # [ c_10 c_11 0 ] # [ c_20 0 0 ] # which represents the following polynomial: # C[y, x] = c_00 + y * c_01 + y^2 * c_02 + x * c_10 + x^2 * c_20 + x * y * c_11 # we replace y with y * scale_y and x with x * scale_x and collect the values to compute the # polynomial for the shifted correction: # C[y, x] = c_00 + y * scale_y * c_01 + y^2 * scale_y^2 * c_02 # + x * scale_x * c_10 + x^2 * scale_x^2 * c_20 # + x * scale_x * y * scale_y * c_11 scale_y = scale_y[..., None, None] scale_x = scale_x[..., None, None] polynomial_matrix_scaled = polynomial_matrix.copy() polynomial_matrix_scaled[..., :, 1, :] *= scale_y polynomial_matrix_scaled[..., :, 2, :] *= scale_y ** 2 polynomial_matrix_scaled[..., 1, :, :] *= scale_x polynomial_matrix_scaled[..., 2, :, :] *= scale_x ** 2 return polynomial_matrix_scaled def _shift_polynomial_matrix(polynomial_matrix, shift_y, shift_x): # This matrix must be a collection of at least 3x3 matrices that hold 2nd order 2d polynomials. # The values are stored as such: # [ c_00 c_01 c_02] # [ c_10 c_11 0 ] # [ c_20 0 0 ] # which represents the following polynomial: # C[y, x] = c_00 + y * c_01 + y^2 * c_02 + x * c_10 + x^2 * c_20 + x * y * c_11 # we replace y with y - z and x with x - w and collect the values to compute the # polynomial for the shifted correction: # C_shift[y, x] = c_02 z^2 - c_01 z + c_00 + c_20 w^2 + c_11 w z - c_10 w # + y (-c_11 w - 2 c_02 z + c_01) + c_02 y^2 + x (-c_11 z - 2 c_20 w + c_10) # + c_20 x^2 + c_11 y x z = shift_y[..., None] w = shift_x[..., None] c_00 = polynomial_matrix[..., 0, 0, :] c_01 = polynomial_matrix[..., 0, 1, :] c_02 = polynomial_matrix[..., 0, 2, :] c_10 = polynomial_matrix[..., 1, 0, :] c_20 = polynomial_matrix[..., 2, 0, :] c_11 = polynomial_matrix[..., 1, 1, :] c_00_shifted = c_02 * z**2 - c_01 * z + c_00 + c_20 * w**2 + c_11 * w * z - c_10 * w c_01_shifted = -c_11 * w - 2 * c_02 * z + c_01 c_02_shifted = c_02 c_10_shifted = -c_11 * z - 2 * c_20 * w + c_10 c_20_shifted = c_20 c_11_shifted = c_11 polynomial_matrix_shifted = np.zeros(polynomial_matrix.shape, dtype=polynomial_matrix.dtype) polynomial_matrix_shifted[..., 0, 0, :] = c_00_shifted polynomial_matrix_shifted[..., 0, 1, :] = c_01_shifted polynomial_matrix_shifted[..., 0, 2, :] = c_02_shifted polynomial_matrix_shifted[..., 1, 0, :] = c_10_shifted polynomial_matrix_shifted[..., 2, 0, :] = c_20_shifted polynomial_matrix_shifted[..., 1, 1, :] = c_11_shifted return polynomial_matrix_shifted def chunk_data(dataset, chunk_size): """ Chunk data into smaller parts and calculate the mean along specified axes. Parameters ---------- dataset: Dataset The dataset to chunk and extract data from. chunk_size: int The size of each chunk to create. Returns ------- background: ndarray An array containing the mean values of the chunks along specified axes. """ background = np.asarray(dataset.images) new_shape = () chunk_indices = () for d in dataset.images.dims: if d in ('y', 'x'): new_shape += (dataset.sizes[d] // chunk_size, chunk_size) chunk_indices += (len(new_shape) - 1,) else: new_shape += (dataset.sizes[d],) background = np.ascontiguousarray(background.reshape(new_shape)) background = background.mean(axis=chunk_indices) return background def _tqdm(x, *args, **kwargs): return x def subtract_background(dataset, chunk_size, tqdm=None): """ Subtract background from dataset images. Parameters ---------- dataset : xarray Dataset MCAM data containing image data as well as additional metadata. chunk_size : int The size of the chunks to use for background subtraction. Must be a divisor of the image x and y dimensions. tqdm : callable, optional A progress bar function to use for tracking progress. Returns ------- xarray Dataset The dataset with the background subtracted from the images. """ if tqdm is None: tqdm = _tqdm image_shape = dataset.sizes['y'], dataset.sizes['x'] assert image_shape[0] % chunk_size == 0, \ (f"Dataset size in 'y' dimension ({image_shape[0]}) is " f"not divisible by chunk size ({chunk_size}).") assert image_shape[1] % chunk_size == 0, \ (f"Dataset size in 'x' dimension ({image_shape[1]}) is " f"not divisible by chunk size ({chunk_size}).") dataset.images.data = np.asarray(dataset.images.data) chunked_data = chunk_data(dataset, chunk_size=chunk_size) valid_data = get_valid_data(dataset) for i in tqdm(ndrange(chunked_data.shape[:-2])): if not valid_data[i]: continue background = cv2.resize(chunked_data[i], image_shape, interpolation=cv2.INTER_LINEAR) dataset.images.data[i] = np.clip( dataset.images.data[i].astype(np.float32) - background.astype(np.float32), 0, 255, ).astype(dataset.images.dtype) return dataset