import sys
from datetime import UTC, datetime
import numpy as np
import xarray as xr
from owl import __version__
from ..memory_allocator import zeros_aligned
from ..util import sys_info
[docs]
def new_rgb_dataset(N_cameras=(6, 4), image_shape=(2432, 4320, 3),
dtype=np.uint8,
dims=('image_y', 'image_x', 'y', 'x', 'rgb'),
coords=None, array=None):
"""Create a xarray.Dataset containing RGB MCAM data.
By default, this calls ``new``, but passes the parameters
`dims=('image_y', 'image_x', 'y', 'x', 'rgb')`
and
`coords={'rgb': ['r', 'g', 'b']}`
Returns
-------
mcam_dataset : xr.Dataset
An ``xarray.Dataset`` with the last dimension having coordinates of
``'rgb'``.
See Also
--------
new_dataset, new_rgba_dataset
"""
if coords is None:
coords = {}
coords['rgb'] = ['r', 'g', 'b']
return new_dataset(
N_cameras=N_cameras, image_shape=image_shape, dtype=dtype,
dims=dims, coords=coords, array=array)
[docs]
def new_rgba_dataset(N_cameras=(6, 4), image_shape=(2432, 4320, 4),
dtype=np.uint8,
dims=('image_y', 'image_x', 'y', 'x', 'rgba'),
coords=None, array=None):
"""Create a new ``xr.Dataset`` containing the MCAM data with RGBA colors.
By default, this calls ``new``, but passes the parameters
`dims=('image_y', 'image_x', 'y', 'x', 'rgba')`
and
`coords={'rgba': ['r', 'g', 'b', 'a']}`
Returns
-------
mcam_dataset : xr.Dataset
An ``xarray.Dataset`` with the last dimension having coordinates of
``'rgba'``.
See Also
--------
new_dataset, new_rgb_dataset
"""
if coords is None:
coords = {}
coords['rgba'] = ['r', 'g', 'b', 'a']
return new_dataset(
N_cameras=N_cameras, image_shape=image_shape, dtype=dtype,
dims=dims, coords=coords, array=array)
[docs]
def new_dataset(
N_cameras: (int, int)=(9, 6),
image_shape: (int, int)=(3120, 4096),
dtype=np.uint8,
dims=('image_y', 'image_x', 'y', 'x'),
coords=None,
*,
array=None,
add_ml_index=True,
):
"""Create a new xarray.Dataset for MCAM data.
Returns an xarray object that contains MCAM data along with the bare
minimum metadata. The dataset holds an ``mcam_data`` ``xr.DataArray``
of size `(*N_cameras, *image_shape)`
with dimensions ``dims`` and coordinates ``coords``.
If a coordinate is not provided, a new coordinate is set with `np.arange`.
Parameters
----------
N_cameras: tuple of length 2
Number of cameras that exist in the MCAM.
image_shape: tuple
The size of the images returned by the individual cameras.
dtype:
The dtype of the array that should be allocated.
dims:
A list or tuple containing the names of the dimensions.
coords:
A dictionary containing the information for the coordinates.
array:
if provided, ``N_cameras``, ``image_shape`` and ``dtype`` will be
overwritten to match array. The array will be used as the data
in the returned xarray object.
add_ml_index: bool
If True, a random machine learning index is added to the dataset.
Returns
-------
mcam_dataset: xr.Dataset
The dataset containing an ``'images'`` array with associated
dimensions.
"""
if coords is None:
coords = {}
else:
# We don't want to edit the user's coordinates
# We will be adding to this dictionary later
coords = coords.copy()
if array is None:
if len(N_cameras) < 2:
raise ValueError("N_cameras should a vector of length at least 2.")
if len(image_shape) < 2:
raise ValueError("Images should be at least 2D.")
shape = tuple(N_cameras) + tuple(image_shape)
array = zeros_aligned(shape=shape, dtype=dtype)
data_vars = {
'images': ((dims), array),
# Some basic metdadata variables we want to add
'__sys_version__': ((), sys.version),
'__owl_sys_info__': ((), str(sys_info())),
}
coords['__owl_version__'] = ((), __version__)
# Don't do any checking for dims and default coords
# let dataarray complain
default_coords = {
d: range(n)
for (d, n) in zip(dims, array.shape)
if d not in coords
}
coords.update(default_coords)
dataset = xr.Dataset(data_vars=data_vars, coords=coords)
if add_ml_index:
# create random_array for ML image indexing
dataset = add_ml_index_array(dataset)
# a new dataset always has a product_line of "unknown" until a
# serial number is added which differentiates the product line
if 'product_line' not in dataset:
dataset['product_line'] = "unknown"
return dataset
def add_ml_index_array(dataset, ml_index_seed=None):
# create random_array for ML image indexing
if ml_index_seed is None:
ml_index_seed = int(datetime.now(UTC).timestamp() * 1E6)
rng = np.random.default_rng(ml_index_seed)
dimensions = ('timelapse_index', 'channel', 'image_y', 'image_x')
full_shape = tuple()
dims = tuple()
for d in dimensions:
if d in dataset.dims:
dims += (d,)
full_shape += (dataset.sizes[d],)
ml_index = rng.integers(0, np.iinfo(np.uint32).max, size=full_shape, dtype=np.uint32)
dataset['ml_index'] = (dims, ml_index)
dataset['ml_index_seed'] = ml_index_seed
return dataset