from copy import deepcopy
import numpy as np
import xarray as xr
from owl.mcam_data import get_stack_dimension
from owl.util import ndrange
from owl.util.io_utils import check_RAM
logical_and = np.logical_and
logical_or = np.logical_or
logical_not = np.logical_not
# See https://github.com/numpy/numpy/issues/12139
def logical_nor(x, y, out=None):
out = logical_or(x, y, out=out)
return logical_not(out, out=out)
def logical_nand(x, y, out=None):
out = logical_and(x, y, out=out)
return logical_not(out, out=out)
[docs]
def hdr_combine(images, exposures,
image_range=(None, None),
invalid_relative_range=(0.1, 0.1),
target_exposure=None):
"""Combine a stack of images taken with different exposures.
If a given pixel is valid in multiple images, the weighted average of the
intensity in all the images will be taken into account in the final
value returned.
Pixels are combined on a per pixel basis. This means that the images may
be of any shape, so long as the shape is consistent between images.
Parameters
----------
images: list of array-like objects
An array containing the image data. These can be multi-dimensional
arrays so long as they all have the same shape.
exposures: list of exposures
List of the exposures taken for all the images.
image_range: tuple of floats (min, max)
The range of the image. If not provided, limits are taken to be
be those from ``skimage.util.dtype.dtype_limits`` with
``clip_negative=True``.
invalid_relative_range: tuple of floats (min_invalid, max_invalid)
Pixels within a certain fraction of the minimum and maximum are
considered to be invalid with the exception of the image taken with
minimal and maximal exposure.
target_exposure: float
If this is set the outputted image will be scaled to match this exposure value.
If None is given it will be set to the lowest exposure value from the give list.
Examples
--------
>>> image_low_exposure = np.asarray([[0.6, 0.6],
... [0.01, 0.01]])
>>> image_high_exposure = np.asarray([[1.0, 1.0],
... [0.16, 0.16]])
>>> exposures = [1, 10]
>>> hdr_combine([image_low_exposure, image_high_exposure], exposures)
array([[0.6 , 0.6 ],
[0.016, 0.016]])
Notice that in this case, the pixels on the first row were correctly
exposed during the short exposure. The pixels on the second row were
correctly exposed during the longer exposure. HDR combine selects
the pixels that were correctly exposed in both cases and creates a
composite from the input images.
"""
if image_range is None or image_range == (None, None):
from skimage.util.dtype import dtype_limits
image_range = dtype_limits(images[0], clip_negative=True)
imagetype_delta = image_range[1] - image_range[0]
underexposed_threshold = image_range[0] + invalid_relative_range[0] * imagetype_delta
overexposed_threshold = image_range[1] - invalid_relative_range[0] * imagetype_delta
if len(images) != len(exposures):
raise ValueError(
'Must provide the same number of images as exposures.')
# This honestly simplifies the logic below because we can assume
# one image has minimum exposure, and a distinct image has maximum exposure
num_images = len(images)
if num_images == 0:
return None
if num_images == 1:
return images[0].astype('float64', copy=True)
# All images assumed to have the same shape.
shape = images[0].shape
# For the first image, we don't care if we assign overexposed pixels.
final_image = np.zeros(shape, dtype='float64')
final_exposure = np.zeros(shape, dtype='float64')
# Sort the images and exposure pairs by the exposure
# Longest exposure now at the end
images_exposures = iter(sorted(zip(images, exposures),
key=lambda i: i[1]))
# cache the minimum exposure image
first_image, exposure_min = next(images_exposures)
if target_exposure is None:
target_exposure = exposure_min
not_underexposed = first_image > underexposed_threshold
not_overexposed = first_image < overexposed_threshold
# cache this result
# first_overexposed = logical_not(not_overexposed)
valid_pixels = np.logical_and(not_underexposed, not_overexposed)
any_not_underexposed = not_underexposed
any_not_overexposed = not_overexposed
# += and = are same here
# no need to do this * (min_exposure / this_exposure)
final_image[valid_pixels] = first_image[valid_pixels]
final_exposure[valid_pixels] += exposure_min
any_valid_pixels = valid_pixels.copy()
for image, exposure in images_exposures:
not_underexposed = image > underexposed_threshold
not_overexposed = image < overexposed_threshold
logical_or(any_not_underexposed, not_underexposed, out=any_not_underexposed)
logical_or(any_not_overexposed, not_overexposed, out=any_not_overexposed)
valid_pixels = np.logical_and(not_underexposed, not_overexposed)
logical_or(any_valid_pixels, valid_pixels, out=any_valid_pixels)
final_image[valid_pixels] += image[valid_pixels]
final_exposure[valid_pixels] += exposure
invalid_pixels = logical_not(any_valid_pixels)
underexposed = logical_not(not_underexposed)
valid_underexposed = logical_and(underexposed, invalid_pixels)
final_image[valid_underexposed] += image[valid_underexposed]
final_exposure[valid_underexposed] += exposure
logical_or(any_valid_pixels, valid_underexposed, out=any_valid_pixels)
invalid_pixels = logical_not(any_valid_pixels)
# valid_overexposed = logical_and(first_overexposed, invalid_pixels)
# final_image[valid_overexposed] += first_image[valid_overexposed]
# final_exposure[valid_overexposed] += 1
# logical_or(any_valid_pixels, valid_overexposed, out=any_valid_pixels)
# This kinda catches an edge case, where the provided images might have
# been garbled
final_image[invalid_pixels] += first_image[invalid_pixels]
final_exposure[invalid_pixels] += exposure_min
logical_or(any_valid_pixels, invalid_pixels, out=any_valid_pixels)
logical_not(any_valid_pixels, out=invalid_pixels)
# assert not np.any(invalid_pixels)
final_image /= final_exposure
final_image *= target_exposure
return final_image
def _tqdm(iterable, *args, **kwargs):
return iterable
def get_hdr_dataset(dataset_stack,
*,
image_range=(None, None),
invalid_relative_range=(0.1, 0.1),
tonemap_gamma=1.,
target_exposure=None,
tqdm=None):
if tqdm is None:
tqdm = _tqdm
data_hdr = np.zeros_like(dataset_stack.images.data[0])
stack_dim = get_stack_dimension(dataset_stack)
dataset_slice = dataset_stack.isel({stack_dim: 0}).drop_vars(stack_dim)
coords = dataset_slice.images.coords
dataset_hdr = xr.DataArray(data=data_hdr, coords=coords).to_dataset(name='images')
# using a single slice removes issue with the matching the stack dimension
for key, value in dataset_slice.drop_vars(['images']).items():
dataset_hdr[key] = deepcopy(value)
for i in tqdm(ndrange(dataset_stack.images.shape[1:3])):
images_da = dataset_stack.images[:, i[0], i[1]]
exposures = np.asarray(dataset_stack.exposure[:, i[0], i[1]])
images = images_da.data
if target_exposure is None:
target_exposure = min(exposures)
hdr_image = hdr_combine(
images, exposures,
image_range=image_range,
invalid_relative_range=invalid_relative_range,
target_exposure=target_exposure)
# Tone mapping formula found here: https://en.wikipedia.org/wiki/Tone_mapping
# V_out = A * V_in ** gamma
# where hdr_image.max() = A ** (-1/gamma) so that the V_out is between 0 and 1
A = hdr_image.max() ** -tonemap_gamma
dataset_hdr.images.data[i][:] = 255 * A * hdr_image ** tonemap_gamma
dataset_hdr.exposure.data[i] = target_exposure
return dataset_hdr
def acquire_hdr_dataset_stack(mcam, exposures,
*,
selection_slice=None,
tqdm=None,
):
if tqdm is None:
tqdm = _tqdm
if selection_slice in (None, Ellipsis):
# xarray doesn't deal well with using None as a selector
selection_slice = np.s_[:, :]
if isinstance(selection_slice, slice):
selection_slice = (selection_slice,)
if len(selection_slice) == 1:
selection_slice = selection_slice + (slice(None, None, None),)
if len(selection_slice) > 2:
raise ValueError("Selection slice must be of length 2 or less. "
f"Got {len(selection_slice)}.")
selection = np.zeros(mcam.N_cameras, dtype=bool)
selection[selection_slice] = True
N_cameras_selection = selection.sum()
nbytes = mcam.image_nbytes * N_cameras_selection
# make sure exposure stack can fit in RAM
check_RAM(len(exposures), nbytes, hugepages=True)
sorted_exposures = sorted(exposures)
N_exposures = len(sorted_exposures)
# Capture an image to get a sample dataset to build the stack from
sample_dataset = mcam.acquire_selection(selection).isel({
'image_y': selection_slice[0],
'image_x': selection_slice[1],
})
dataset_stack = sample_dataset.drop_vars(['images', 'exposure'])
# expand dataset_stack to correct shape.
stacked_vars = [
'software_timestamp',
'acquisition_index',
'acquisition_count',
'latest_acquisition_index',
]
for variable_name in stacked_vars:
variable = dataset_stack[variable_name].expand_dims({
'frame_number': np.arange(N_exposures, dtype='int'),
})
dataset_stack[variable_name] = variable
# must perform a deep copy after expanding dimensions to have dataset_stack be writable
# https://github.com/pydata/xarray/issues/2891
dataset_stack = dataset_stack.copy(
# This works around our linter that helps us avoid
# doing deep copies for our datasets
deep=True
)
dataset_stack['images'] = (
('frame_number',) + sample_dataset['images'].dims,
np.zeros(
(N_exposures,) + sample_dataset['images'].shape,
dtype=sample_dataset['images'].dtype,
)
)
# While the rule would be to typically expand the dimension,
# thus having a 3D array for the exposure, historically
# this has not been the case with the HDR stack
# We thus only have a single dimension for the exposure, that is
# 'frame_number'
dataset_stack['exposure'] = (
('frame_number',) + sample_dataset['images'].dims[:2],
np.zeros(
(N_exposures,) + sample_dataset['images'].shape[:2],
dtype=sample_dataset['exposure'].dtype,
)
)
starting_exposure = mcam.exposure
fill_vars = stacked_vars + ['images']
for i, exposure in tqdm(enumerate(sorted_exposures), total=N_exposures):
mcam.exposure = exposure
dataset = mcam.acquire_selection(selection).isel({
'image_y': selection_slice[0],
'image_x': selection_slice[1],
})
for variable_name in fill_vars:
dataset_stack[variable_name].data[i] = dataset[variable_name].data
# Manually fill in the exposure since it has different dimensions
# than the dataset exposure
dataset_stack['exposure'].data[i] = mcam.exposure
mcam.exposure = starting_exposure
return dataset_stack