import math
import os
import threading
import warnings
from concurrent.futures import ThreadPoolExecutor
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
try:
# Fused Highway SIMD block mean + subtraction for uint8 frames
# (owl/calibration/_background_simd.cpp): a plate that takes the numpy
# path below seconds takes this a fraction of one. Built into every wheel
# via hatch_build.py SIMD_MODULES.
from owl.calibration import _background_simd
except ImportError:
# A source checkout whose compiled extensions have not been (re)built yet:
# fall back to numpy/OpenCV and nudge the developer to rebuild the same way
# the rest of the repo installs (the editable, no-build-isolation install).
_background_simd = None
warnings.warn(
"owl.calibration._background_simd is not built; falling back to the "
"slower numpy background subtraction. Rebuild the compiled extensions "
"with `pip install -e . --no-build-isolation` (the editable install the "
"rest of the repo uses).",
RuntimeWarning,
stacklevel=2,
)
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
# Cap on the worker threads ``chunk_data`` and ``subtract_background`` use.
# Both are memory-bound, and they run inside an assay or the GUI, which have
# other work to keep cores for: the SIMD kernels reach the memory bus limit by
# four threads (14 GB/s of frames on a 12-core machine, against 5 for one), so
# nothing past this would pay anyway.
_MAX_BACKGROUND_THREADS = 4
# Frames handed to the SIMD kernel per call, so a progress bar still moves.
_SIMD_BATCH_FRAMES = 32
def _background_threads(threads):
if threads is None:
return max(1, min(os.cpu_count() or 1, _MAX_BACKGROUND_THREADS))
return max(1, min(int(threads), _MAX_BACKGROUND_THREADS))
def _simd_frames(images):
# The ``(frames, y, x)`` view the SIMD kernels take, or None when they do
# not apply: not built, not uint8, or not contiguous (the kernels work in
# place, so a copy would be the wrong array to subtract from).
if _background_simd is None or images.dtype != np.uint8 or not images.flags.c_contiguous:
return None
return images.reshape(-1, *images.shape[-2:])
def _frame_block_mean(frame, chunk_size, out):
# Average every ``chunk_size`` square block. Accumulating in float32 is
# exact for uint8 blocks up to 256 pixels a side (the sum stays below
# float32's 2**24 integer limit) and within a relative 1e-7 beyond that;
# numpy sums pairwise, which keeps it there.
blocks = frame.reshape(
frame.shape[0] // chunk_size, chunk_size,
frame.shape[1] // chunk_size, chunk_size,
)
np.mean(blocks, axis=(1, 3), dtype=np.float32, out=out)
def chunk_data(dataset, chunk_size, *, threads=None):
"""
Estimate a coarse background as the mean of ``chunk_size`` square blocks.
Each image's ``(y, x)`` plane is cut into non-overlapping
``chunk_size x chunk_size`` blocks and every block is reduced to one
value. The images are memory-bound work, so frames are reduced on a
thread pool.
uint8 images take the fused SIMD kernel (``_background_simd``); other
dtypes take the numpy path, which computes the same thing.
Parameters
----------
dataset: Dataset
The dataset to chunk and extract data from.
chunk_size: int
The size of each chunk to create. Must divide both image dimensions.
threads: int, optional
Worker threads to reduce frames on, at most ``_MAX_BACKGROUND_THREADS``.
Returns
-------
background: ndarray
float32 block means, shaped like the images with ``y`` and ``x``
replaced by ``y // chunk_size`` and ``x // chunk_size``.
"""
images = np.asarray(dataset.images)
dims = dataset.images.dims
if dims[-2:] != ('y', 'x'):
raise ValueError(f"Expected the images to end in ('y', 'x') dimensions, got {dims}.")
height, width = images.shape[-2:]
if height % chunk_size or width % chunk_size:
raise ValueError(
f"Chunk size {chunk_size} must divide the image shape {(height, width)}."
)
leading_shape = images.shape[:-2]
threads = _background_threads(threads)
flat = _simd_frames(images)
if flat is not None:
means = _background_simd.block_mean(flat, chunk_size, None, threads)
return means.reshape(leading_shape + means.shape[-2:])
background = np.empty(
leading_shape + (height // chunk_size, width // chunk_size), dtype=np.float32,
)
frames = list(ndrange(leading_shape))
def reduce_frame(index):
_frame_block_mean(images[index], chunk_size, background[index])
threads = min(threads, max(1, len(frames)))
if threads == 1:
for index in frames:
reduce_frame(index)
else:
with ThreadPoolExecutor(threads) as executor:
# Consume the iterator so a worker's exception is raised here.
for _ in executor.map(reduce_frame, frames):
pass
return background
def _tqdm(x, *args, **kwargs):
return x
class _FrameBackgroundSubtractor:
"""Subtract the upsampled block mean from frames of one shape, in place.
The obvious per-frame pipeline (``cv2.resize`` the block means to a
full-resolution float32 plane, ceil it, cast it, subtract it) moves four
full planes through memory per frame and does not speed up with threads,
because it is memory bound. This does the separable bilinear resize in
two halves: the horizontal pass once per frame on the tiny block grid
with ``cv2.resize``, and the vertical pass one block row at a time as a
``(rows, 2) @ (2, width)`` matrix product, so the strip, its ceil, its
integer cast and the subtraction all stay in a core's cache. The frame
itself is read and written exactly once.
The vertical weights follow ``cv2.INTER_LINEAR``'s half-pixel convention
(source row ``(r + 0.5) / scale - 0.5``, clamped at both edges), which is
also what the viewer's shader samples, so the result matches a
``cv2.resize`` of the whole plane to float32 rounding.
"""
def __init__(self, image_shape, dtype):
self._height, self._width = image_shape
self._dtype = np.dtype(dtype)
self._local = threading.local()
self._strips = {}
def _vertical_strips(self, chunk_rows):
# Output rows grouped by the source row below them: for each group the
# slice of output rows and its ``(rows, 2)`` weights on that source
# row and the next. The last source row is duplicated by the caller
# so the bottom group can always blend with weight zero on "the next".
strips = self._strips.get(chunk_rows)
if strips is not None:
return strips
rows = np.arange(self._height, dtype=np.float64)
source = np.clip((rows + 0.5) * (chunk_rows / self._height) - 0.5, 0, chunk_rows - 1)
lower = np.floor(source).astype(np.intp)
upper_weight = (source - lower).astype(np.float32)
weights = np.stack([1 - upper_weight, upper_weight], axis=1)
boundaries = np.flatnonzero(np.diff(lower)) + 1
starts = np.concatenate([[0], boundaries])
stops = np.concatenate([boundaries, [self._height]])
strips = tuple(
(slice(int(start), int(stop)), int(lower[start]),
np.ascontiguousarray(weights[start:stop]))
for start, stop in zip(starts, stops)
)
self._strips[chunk_rows] = strips
return strips
def _buffers(self, rows):
buffers = vars(self._local).get('buffers')
if buffers is None or buffers[0].shape[0] < rows:
buffers = (
np.empty((rows, self._width), dtype=np.float32),
np.empty((rows, self._width), dtype=self._dtype),
)
self._local.buffers = buffers
return buffers
def __call__(self, frame, chunk_means):
"""Subtract the bilinear upsample of ``chunk_means`` from ``frame`` in place."""
chunk_rows = chunk_means.shape[0]
strips = self._vertical_strips(chunk_rows)
# Horizontal pass: the vertical scale is 1 so rows are copied exactly.
wide = np.empty((chunk_rows + 1, self._width), dtype=np.float32)
cv2.resize(
chunk_means, (self._width, chunk_rows), dst=wide[:chunk_rows],
interpolation=cv2.INTER_LINEAR,
)
wide[chunk_rows] = wide[chunk_rows - 1]
float_strip, integer_strip = self._buffers(
max(weights.shape[0] for _, _, weights in strips))
is_integer = np.issubdtype(frame.dtype, np.integer)
for rows, source_row, weights in strips:
count = weights.shape[0]
background = float_strip[:count]
np.matmul(weights, wide[source_row:source_row + 2], out=background)
target = frame[rows]
if is_integer:
# ``int(frame - background)`` truncated towards zero, clipped
# at zero, is ``frame - ceil(background)`` saturated at zero
# for an integer frame, so the subtraction can be a saturating
# integer one with no float copy of the frame.
np.ceil(background, out=background)
integer = integer_strip[:count]
np.copyto(integer, background, casting='unsafe')
if target.flags.c_contiguous:
cv2.subtract(target, integer, dst=target)
else:
target[...] = cv2.subtract(target, integer)
else:
np.clip(target - background, 0, 255, out=target)
def subtract_background(dataset, chunk_size, tqdm=None, *, threads=None):
"""
Subtract background from dataset images.
The background is ``chunk_data``'s block mean, bilinearly upsampled to the
image, and the subtraction saturates at zero. The images are modified in
place. uint8 images take the fused SIMD kernel (``_background_simd``),
which estimates and subtracts each frame's background in one pass; other
dtypes take the numpy/OpenCV path, which computes the same thing.
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.
threads : int, optional
Worker threads for both the estimate and the subtraction.
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)
images = dataset.images.data
threads = _background_threads(threads)
valid_data = get_valid_data(dataset)
leading_shape = images.shape[:-2]
leading_dims = dataset.images.dims[:-2]
if 'image_y' in leading_dims and 'image_x' in leading_dims:
image_axes = leading_dims.index('image_y'), leading_dims.index('image_x')
valid_frames = np.array([
valid_data[index[image_axes[0]], index[image_axes[1]]]
for index in ndrange(leading_shape)
], dtype=bool).reshape(leading_shape)
else:
valid_frames = np.ones(leading_shape, dtype=bool)
flat = _simd_frames(images)
if flat is not None:
# Estimate and subtract each frame's background in one fused pass, a
# batch of frames per call so the progress bar keeps moving.
flat_valid = valid_frames.reshape(-1)
batches = range(0, len(flat), _SIMD_BATCH_FRAMES)
for start in tqdm(batches, total=len(batches)):
stop = start + _SIMD_BATCH_FRAMES
_background_simd.subtract_block_background(
flat[start:stop], chunk_size, flat_valid[start:stop], threads,
)
return dataset
chunked_data = chunk_data(dataset, chunk_size=chunk_size, threads=threads)
frames = [index for index in ndrange(leading_shape) if valid_frames[index]]
subtract = _FrameBackgroundSubtractor(image_shape, images.dtype)
def subtract_frame(index):
subtract(images[index], chunked_data[index])
threads = min(threads, max(1, len(frames)))
if threads == 1:
for index in tqdm(frames):
subtract_frame(index)
else:
with ThreadPoolExecutor(threads) as executor:
for _ in tqdm(executor.map(subtract_frame, frames), total=len(frames)):
pass
return dataset
