import numpy as np
import xarray as xr
from dask import array as da
from numpy import pi
from scipy.fftpack import fft, fftfreq, ifft
def da_stack(data):
if data.ndim > 1:
return da.stack([da_stack(d) for d in data])
return da.stack(data)
def relu_energy_second_order_total(image):
return (relu_energy_second_order(image, axis=-1) +
relu_energy_second_order(image, axis=-2))
def relu_energy_second_order(image, axis=-1):
N = image.shape[axis]
fx = fftfreq(N)
filter_x = np.sin(pi * fx) ** 2
axis_to_add = (image.ndim - axis - 1) % image.ndim
filter_x = filter_x[(..., *((np.newaxis, ) * axis_to_add))]
img_filtered = np.real(ifft(fft(image, axis=axis) * filter_x, axis=axis))
img_filtered[img_filtered < 0] = 0
return np.sum(img_filtered ** 2)
# Convenience functions to use with dask for ultimate speedup
# If a dask array is provided, calling this function will automatically
# parallelize the computation of relu_energy_second_order
# Each array, of shape [M, N] will be operated on independently
# And a single number, with no shape (), of dtype float64
# will be returned by the function.
da_relu_energy_second_order_total = da.gufunc(
relu_energy_second_order_total,
signature='(m,n)->()', output_dtypes=np.float64, vectorize=True)
[docs]
def find_in_focus_indices(data, *, search_axis=0, image_dims=2):
"""Find the indices that are in focus for each camera.
The provided data should be something that looks like grayscale
``mcam_data`` with an additional dimension, the one you wish to search
along.
For example, the mcam_data might look something like:
>>> raw_data
<xarray.DataArray 'stack-329162f6' (z: 12, image_y: 6, image_x: 4, y: 2432, x: 4320)>
dask.array<shape=(12, 6, 4, 2432, 4320), dtype=uint8, chunksize=(1, 1, 1, 2432, 4320)>
Coordinates:
* image_y (image_y) int64 0 1 2 3 4 5
* image_x (image_x) int64 0 1 2 3
* y (y) int64 0 1 2 3 4 5 6 ... 2426 2427 2428 2429 2430 2431
* x (x) int64 0 1 2 3 4 5 6 ... 4314 4315 4316 4317 4318 4319
* z (z) float64 0.0 0.5 1.0 1.5 2.0 ... 3.5 4.0 4.5 5.0 5.5
Attributes:
sys_info: Python: 3.6.6 | packaged by conda-forge | (default, Ju...
imagename_format: cam{image_y}_{image_x}.bmp
With ``z`` being the dimension we wish to search along.
The function would then be called, returning a delayed object that consists
of the indices corresponding to the slice in focus for each camera.
Here we ask that the metric only be computed considering every other pixel.
This is an acceptable compromise since it avoids any Bayer Pattern
artifacts and speeds up the computation.
>>> in_focus_z_indices = find_in_focus_indices(raw_data[..., ::2, ::2])
>>> in_focus_z_indices
<xarray.DataArray 'z_stage' (image_y: 6, image_x: 4)>
dask.array<shape=(6, 4), dtype=int64, chunksize=(1, 1)>
Coordinates:
* image_y (image_y) int64 0 1 2 3 4 5
* image_x (image_x) int64 0 1 2 3
Data should be a dask array with the dimension you want to search for
being the leading dimension.
Parameters
----------
data: ArrayLike [N_z, N_cameras_y, N_cameras_x, y, x]
A 5D data volume containing info about all cameras.
search_axis: int ????
Not supported yet, but hopefully we can at some point specify
which axis to search along. This doesn't seem to be useful immediately.
For now this is simply set to the axis of ``N_z``.
image_dims: int
The number of dimensions each image has. For now, this is hardcoded
to 2, grayscale image, but in the future may be expanded to support
more than grayscale images.
Returns
-------
in_focus_indices: ArrayLike [image_y, image_x]
The indices where each camera has the best focus. This is potentially
a delayed object.
"""
if image_dims != 2:
raise NotImplementedError(
"We only consider grayscale images.")
if search_axis != 0:
raise NotImplementedError(
"We only consider the case we search along the 0 dimension.")
if data.shape[-1] in [3, 4]:
raise ValueError(f"You supplied an data with a dimension of {data.shape[-1]}. "
"This likely means that your data is RGB or RGBA. "
"This function only supports grayscale data.")
if isinstance(data, xr.DataArray):
if 'rgb' in data.dims:
raise ValueError("We only support grayscale images in this function.")
da_data = data.data
else:
da_data = data
energy = da_relu_energy_second_order_total(da_data)
in_focus_index = energy.argmax(search_axis)
if isinstance(data, xr.DataArray):
coords = {d: data.coords[d] for d in data.dims[1:3]}
in_focus_index = xr.DataArray(in_focus_index, name=data.dims[0],
dims=data.dims[1:3], coords=coords)
return in_focus_index
# Why does slice_in_focus_images exist if this function exists????
def select_in_focus_images(data):
"""Data should be an xarray with the dimension you want to search for
being the zero dimension"""
in_focus_index = find_in_focus_indices(data.data).compute()
# See vectorized indexing
# Unfortunately, vectorized indexing causing the data to full aggregate
# Which means it is no longer easily parallelizable
# http://xarray.pydata.org/en/stable/indexing.html#vectorized-indexing
# index_z = xr.DataArray(in_focus_index, dims=data.dims[1:-2])
selected_array = np.empty(shape=data.shape[1:-2], dtype=object)
for i in np.ndindex(data.shape[1:-2]):
z = in_focus_index[i]
selected_array[i] = data[(z, *i)]
return selected_array
[docs]
def slice_in_focus_images(data, in_focus_indices):
"""Picks out the images with that are in focus from MCAM Data.
The `data` is assumed to be an xarray with the leading dimension being
the one we wish to slice into. For example, if the stack is a focal
stack where the leading dimension contains measurements at 10 different
stage heights `z`, the array might have the following dimensions
>>> data = mcam_data.new(N_cameras=(10, 6, 4),
dims=['z_stage', 'image_y', 'image_x', 'y', 'x'])
>>> data.shape
(10, 6, 4, 2432, 4320)
in_focus_indices is an array of shape `[N_image_y, N_image_x]`
containing the index that should be sliced in for each camera.
>>> color_data = slice_in_focus_images(color_data, in_focus_z_index)
>>> color_data
<xarray.DataArray 'stack-e599e60' (image_y: 6, image_x: 4, y: 2432, x: 4320, rgb: 3)>
dask.array<shape=(6, 4, 2432, 4320, 3), dtype=uint8, chunksize=(1, 1, 2432, 4320, 3)>
Coordinates:
* image_y (image_y) int64 0 1 2 3 4 5
* image_x (image_x) int64 0 1 2 3
* y (y) int64 0 1 2 3 4 5 6 ... 2426 2427 2428 2429 2430 2431
* x (x) int64 0 1 2 3 4 5 6 ... 4314 4315 4316 4317 4318 4319
bayer_pattern (image_y, image_x) <U4 'bggr' 'bggr' ... 'bggr' 'bggr'
camera_number (image_y, image_x) int64 20 21 23 22 16 ... 6 0 1 3 2
exposure (image_y, image_x) int64 500 500 500 ... 500 500 500
gain (image_y, image_x) int64 1 1 1 1 1 1 1 ... 1 1 1 1 1 1
* rgb (rgb) <U1 'r' 'g' 'b'
acquisition_count (image_y, image_x) int64 5 6 3 3 4 5 6 ... 4 3 3 4 3 3
trigger (image_y, image_x) int64 175 176 173 ... 174 173 173
z (image_y, image_x) float64 2.0 2.5 1.0 ... 1.5 1.0 1.0
Attributes:
sys_info: Python: 3.6.6 | packaged by conda-forge | (default, Ju...
imagename_format: cam{image_y}_{image_x}.bmp
>>> color_data = color_data.persist(scheduler='threads')
>>> color_data
<xarray.DataArray 'stack-e599e60' (image_y: 6, image_x: 4, y: 2432, x: 4320, rgb: 3)>
dask.array<shape=(6, 4, 2432, 4320, 3), dtype=uint8, chunksize=(1, 1, 2432, 4320, 3)>
Coordinates:
* image_y (image_y) int64 0 1 2 3 4 5
* image_x (image_x) int64 0 1 2 3
* y (y) int64 0 1 2 3 4 5 6 ... 2426 2427 2428 2429 2430 2431
* x (x) int64 0 1 2 3 4 5 6 ... 4314 4315 4316 4317 4318 4319
bayer_pattern (image_y, image_x) <U4 'bggr' 'bggr' ... 'bggr' 'bggr'
camera_number (image_y, image_x) int64 20 21 23 22 16 ... 6 0 1 3 2
exposure (image_y, image_x) int64 500 500 500 ... 500 500 500
gain (image_y, image_x) int64 1 1 1 1 1 1 1 ... 1 1 1 1 1 1
* rgb (rgb) <U1 'r' 'g' 'b'
acquisition_count (image_y, image_x) int64 5 6 3 3 4 5 6 ... 4 3 3 4 3 3
trigger (image_y, image_x) int64 175 176 173 ... 174 173 173
z (image_y, image_x) float64 2.0 2.5 1.0 ... 1.5 1.0 1.0
Attributes:
sys_info: Python: 3.6.6 | packaged by conda-forge | (default, Ju...
imagename_format: cam{image_y}_{image_x}.bmp
The use of ``persist`` instead of ``compute`` keeps the data discontinuous
in memory. This is useful since the final array, especially for Gigacam
data might be several GB large. Moving that around needlessly is not
necessary unless you require it to be fully contiguous.
Parameters
----------
data: ArrayLike[N_z, N_cameras_y, N_cameras_x, y, x, [...]]
The mcam_data you wish to slice into.
in_focus_indices: [N_cameras_y, N_cameras_x]
The in-focus indices for each camera.
Returns
-------
in_focus_data: ArrayLikee[N_cameras_y, N_cameras_x, y, x, [...]]
The mcam_data at the desired slice.
"""
z_dim_name = data.dims[0]
# This is such a hack, just because it is dask slicing, it works here
# But it actually forces the dask array to come together. And we can't
# have it come together because we want the computation to stay
# distributed
coords = {**data[in_focus_indices].coords}
# See vectorized indexing
# Unfortunately, vectorized indexing causing the data to full aggregate
# Which means it is no longer easily parallelizable
# http://xarray.pydata.org/en/stable/indexing.html#vectorized-indexing
# index_z = xr.DataArray(in_focus_index, dims=data.dims[1:-2])
selected_array = np.empty(shape=in_focus_indices.data.shape, dtype=object)
for i in np.ndindex(in_focus_indices.data.shape):
z = in_focus_indices.data[i]
selected_array[i] = data.data[(z, *i)]
coords[z_dim_name] = data[z_dim_name][in_focus_indices]
sliced_data = xr.DataArray(da_stack(selected_array), dims=data.dims[1:],
coords=coords, attrs=data.attrs)
return sliced_data