def _group_tensor_into_windowed_tensor(self, x, valid_first_frame):
if self._window_size == 1:
valid_first_frame_pruned = ivy.cast(valid_first_frame[:, 0],
'bool')
else:
valid_first_frame_pruned = ivy.cast(
valid_first_frame[:1 - self._window_size, 0], 'bool')
if ivy.reduce_sum(ivy.cast(valid_first_frame_pruned, 'int32'))[0] == 0:
valid_first_frame_pruned =\
ivy.cast(ivy.one_hot(0, self._sequence_lengths[0] - self._window_size + 1), 'bool')
window_idxs_single = ivy.indices_where(valid_first_frame_pruned)
gather_idxs_list = list()
for w_idx in window_idxs_single:
gather_idxs_list.append(
ivy.expand_dims(
ivy.arange(w_idx[0] + self._window_size, w_idx[0], 1), 0))
gather_idxs = ivy.concatenate(gather_idxs_list, 0)
gather_idxs = ivy.reshape(gather_idxs, (-1, 1))
num_valid_windows_for_seq = ivy.shape(window_idxs_single)[0:1]
return ivy.reshape(
ivy.gather_nd(x, gather_idxs),
ivy.concatenate(
(num_valid_windows_for_seq, ivy.array(
[self._window_size]), ivy.shape(x)[1:]), 0))
def _group_tensor_into_windowed_tensor_simple(self, x, seq_info):
seq_info = self._update_seq_info_for_window(seq_info)
if self._fixed_sequence_length:
return ivy.reshape(ivy.gather_nd(x, ivy.array(self._gather_idxs)),
(self._windows_per_seq, self._window_size) +
x.shape[1:])
else:
num_windows_in_seq = int(
ivy.to_numpy(
ivy.maximum(seq_info.length[0] - self._window_size + 1,
1)))
window_idxs_in_seq = ivy.arange(num_windows_in_seq, 0, 1)
gather_idxs = ivy.tile(
ivy.reshape(ivy.arange(self._window_size, 0, 1),
(1, self._window_size)),
(num_windows_in_seq, 1)) + ivy.expand_dims(
window_idxs_in_seq, -1)
gather_idxs_flat = ivy.reshape(
gather_idxs, (self._window_size * num_windows_in_seq, 1))
return ivy.reshape(ivy.gather_nd(x, gather_idxs_flat),
(num_windows_in_seq, self._window_size) +
x.shape[1:])
def coords_to_voxel_grid(coords,
voxel_shape_spec,
mode='DIMS',
coord_bounds=None,
features=None,
batch_shape=None,
dev_str=None):
"""
Create voxel grid :math:`\mathbf{X}_v\in\mathbb{R}^{x×y×z×(3+N+1)}` from homogeneous co-ordinates
:math:`\mathbf{X}_w\in\mathbb{R}^{num\_coords×4}`. Each voxel contains 3+N+1 values: the mean normalized
co-ordinate inside the voxel for the projected pixels with :math:`0 < x, y, z < 1`, N coordinte features (optional),
and also the number of projected pixels inside the voxel.
Grid resolutions and dimensions are also returned separately for each entry in the batch.
Note that the final batched voxel grid returned uses the maximum grid dimensions across the batch, this means
some returned grids may contain redundant space, with all but the single largest batched grid occupying a subset
of the grid space, originating from the corner of minimum :math:`x,y,z` values.\n
`[reference] <https://en.wikipedia.org/wiki/Voxel>`_
:param coords: Homogeneous co-ordinates *[batch_shape,c,4]*
:type coords: array
:param voxel_shape_spec: Either the number of voxels in x,y,z directions, or the resolutions (metres) in x,y,z
directions, depending on mode. Batched or unbatched. *[batch_shape,3]* or *[3]*
:type voxel_shape_spec: array
:param mode: Shape specification mode, either "DIMS" or "RES"
:type mode: str
:param coord_bounds: Co-ordinate x, y, z boundaries *[batch_shape,6]* or *[6]*
:type coord_bounds: array
:param features: Co-ordinate features *[batch_shape,c,4]*.
E.g. RGB values, low-dimensional features, etc.
Features mapping to the same voxel are averaged.
:type features: array
:param batch_shape: Shape of batch. Inferred from inputs if None.
:type batch_shape: sequence of ints, optional
:param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Same as x if None.
:type dev_str: str, optional
:return: Voxel grid *[batch_shape,x_max,v_max,z_max,3+feature_size+1]*, dimensions *[batch_shape,3]*, resolutions *[batch_shape,3]*, voxel_grid_lower_corners *[batch_shape,3]*
"""
if batch_shape is None:
batch_shape = coords.shape[:-2]
if dev_str is None:
dev_str = _ivy.dev_str(coords)
# shapes as list
batch_shape = list(batch_shape)
num_batch_dims = len(batch_shape)
num_coords_per_batch = coords.shape[-2]
# voxel shape spec as array
if len(voxel_shape_spec) is 3:
# BS x 1 x 3
voxel_shape_spec = _ivy.expand_dims(
_ivy.tile(
_ivy.reshape(_ivy.array(voxel_shape_spec),
[1] * num_batch_dims + [3]), batch_shape + [1]),
-2)
# coord bounds spec as array
if coord_bounds is not None and len(coord_bounds) is 6:
# BS x 1 x 6
coord_bounds = _ivy.expand_dims(
_ivy.tile(
_ivy.reshape(_ivy.array(coord_bounds, dtype_str='float32'),
[1] * num_batch_dims + [6]), batch_shape + [1]),
-2)
# BS x N x 3
coords = coords[..., 0:3]
if coord_bounds is not None:
# BS x 1
x_min = coord_bounds[..., 0:1]
y_min = coord_bounds[..., 1:2]
z_min = coord_bounds[..., 2:3]
x_max = coord_bounds[..., 3:4]
y_max = coord_bounds[..., 4:5]
z_max = coord_bounds[..., 5:6]
# BS x N x 1
x_coords = coords[..., 0:1]
y_coords = coords[..., 1:2]
z_coords = coords[..., 2:3]
x_validity_mask = _ivy.logical_and(x_coords > x_min, x_coords < x_max)
y_validity_mask = _ivy.logical_and(y_coords > y_min, y_coords < y_max)
z_validity_mask = _ivy.logical_and(z_coords > z_min, z_coords < z_max)
# BS x N
full_validity_mask = _ivy.logical_and(
_ivy.logical_and(x_validity_mask, y_validity_mask),
z_validity_mask)[..., 0]
# BS x 1 x 3
bb_mins = coord_bounds[..., 0:3]
bb_maxs = coord_bounds[..., 3:6]
bb_ranges = bb_maxs - bb_mins
else:
# BS x N
full_validity_mask = _ivy.cast(
_ivy.ones(batch_shape + [num_coords_per_batch]), 'bool')
# BS x 1 x 3
bb_mins = _ivy.reduce_min(coords, axis=-2, keepdims=True)
bb_maxs = _ivy.reduce_max(coords, axis=-2, keepdims=True)
bb_ranges = bb_maxs - bb_mins
# get voxel dimensions
if mode is 'DIMS':
# BS x 1 x 3
dims = _ivy.cast(voxel_shape_spec, 'int32')
elif mode is 'RES':
# BS x 1 x 3
res = _ivy.cast(voxel_shape_spec, 'float32')
dims = _ivy.cast(_ivy.ceil(bb_ranges / (res + MIN_DENOMINATOR)),
'int32')
else:
raise Exception(
'Invalid mode selection. Must be either "DIMS" or "RES"')
dims_m_one = _ivy.cast(dims - 1, 'int32')
# BS x 1 x 3
res = bb_ranges / (_ivy.cast(dims, 'float32') + MIN_DENOMINATOR)
# BS x NC x 3
voxel_indices = _ivy.minimum(
_ivy.cast(_ivy.floor((coords - bb_mins) / (res + MIN_DENOMINATOR)),
'int32'), dims_m_one)
# BS x NC x 3
voxel_values = ((coords - bb_mins) % res) / (res + MIN_DENOMINATOR)
feature_size = 0
if features is not None:
feature_size = features.shape[-1]
voxel_values = _ivy.concatenate([voxel_values, features], axis=-1)
# TNVC x len(BS)+1
valid_coord_indices = _ivy.cast(_ivy.indices_where(full_validity_mask),
'int32')
# scalar
total_num_valid_coords = valid_coord_indices.shape[0]
# TNVC x 3
voxel_values_pruned_flat = _ivy.gather_nd(voxel_values,
valid_coord_indices)
voxel_indices_pruned_flat = _ivy.gather_nd(voxel_indices,
valid_coord_indices)
# TNVC x len(BS)+2
if num_batch_dims == 0:
all_indices_pruned_flat = voxel_indices_pruned_flat
else:
batch_indices = valid_coord_indices[..., :-1]
all_indices_pruned_flat = _ivy.concatenate(
[batch_indices] + [voxel_indices_pruned_flat], -1)
# TNVC x 4
voxel_values_pruned_flat =\
_ivy.concatenate((voxel_values_pruned_flat, _ivy.ones([total_num_valid_coords, 1], dev_str=dev_str)), -1)
# get max dims list for scatter
if num_batch_dims > 0:
max_dims = _ivy.reduce_max(_ivy.reshape(dims, batch_shape + [3]),
axis=list(range(num_batch_dims)))
else:
max_dims = _ivy.reshape(dims, batch_shape + [3])
batch_shape_array_list = [_ivy.array(batch_shape, 'int32')
] if num_batch_dims != 0 else []
total_dims_list = _ivy.to_list(
_ivy.concatenate(
batch_shape_array_list +
[max_dims, _ivy.array([4 + feature_size], 'int32')], -1))
# BS x x_max x y_max x z_max x 4
scattered = _ivy.scatter_nd(
all_indices_pruned_flat,
voxel_values_pruned_flat,
total_dims_list,
reduction='replace' if _ivy.backend == 'mxnd' else 'sum')
# BS x x_max x y_max x z_max x 4 + feature_size, BS x 3, BS x 3, BS x 3
return _ivy.concatenate(
(scattered[..., :-1] /
(_ivy.maximum(scattered[..., -1:], 1.) + MIN_DENOMINATOR),
scattered[..., -1:]),
-1), dims[..., 0, :], res[..., 0, :], bb_mins[..., 0, :]
def quantize_to_image(pixel_coords,
final_image_dims,
feat=None,
feat_prior=None,
with_db=False,
pixel_coords_var=1e-3,
feat_var=1e-3,
pixel_coords_prior_var=1e12,
feat_prior_var=1e12,
var_threshold=(1e-3, 1e12),
uniform_pixel_coords=None,
batch_shape=None,
dev_str=None):
"""
Quantize pixel co-ordinates with d feature channels (for depth, rgb, normals etc.), from
images :math:`\mathbf{X}\in\mathbb{R}^{input\_images\_shape×(2+d)}`, which may have been reprojected from a host of
different cameras (leading to non-integer pixel values), to a new quantized pixel co-ordinate image with the same
feature channels :math:`\mathbf{X}\in\mathbb{R}^{h×w×(2+d)}`, and with integer pixel co-ordinates.
Duplicates during the quantization are either probabilistically fused based on variance, or the minimum depth is
chosen when using depth buffer mode.
:param pixel_coords: Pixel co-ordinates *[batch_shape,input_size,2]*
:type pixel_coords: array
:param final_image_dims: Image dimensions of the final image.
:type final_image_dims: sequence of ints
:param feat: Features (i.e. depth, rgb, encoded), default is None. *[batch_shape,input_size,d]*
:type feat: array, optional
:param feat_prior: Prior feature image mean, default is None. *[batch_shape,input_size,d]*
:type feat_prior: array or float to fill with
:param with_db: Whether or not to use depth buffer in rendering, default is false
:type with_db: bool, optional
:param pixel_coords_var: Pixel coords variance *[batch_shape,input_size,2]*
:type pixel_coords_var: array or float to fill with
:param feat_var: Feature variance *[batch_shape,input_size,d]*
:type feat_var: array or float to fill with
:param pixel_coords_prior_var: Pixel coords prior variance *[batch_shape,h,w,2]*
:type pixel_coords_prior_var: array or float to fill with
:param feat_prior_var: Features prior variance *[batch_shape,h,w,3]*
:type feat_prior_var: array or float to fill with
:param var_threshold: Variance threshold, for projecting valid coords and clipping *[batch_shape,2+d,2]*
:type var_threshold: array or sequence of floats to fill with
:param uniform_pixel_coords: Homogeneous uniform (integer) pixel co-ordinate images, inferred from final_image_dims
if None *[batch_shape,h,w,3]*
:type uniform_pixel_coords: array, optional
:param batch_shape: Shape of batch. Assumed no batches if None.
:type batch_shape: sequence of ints, optional
:param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Same as x if None.
:type dev_str: str, optional
:return: Quantized pixel co-ordinates image with d feature channels (for depth, rgb, normals etc.) *[batch_shape,h,w,2+d]*,
maybe the quantized variance, *[batch_shape,h,w,2+d]*, and scatter counter image *[batch_shape,h,w,1]*
"""
# ToDo: make variance fully optional. If not specified,
# then do not compute and scatter during function call for better efficiency.
# config
if batch_shape is None:
batch_shape = pixel_coords.shape[:-2]
if dev_str is None:
dev_str = _ivy.dev_str(pixel_coords)
if feat is None:
d = 0
else:
d = feat.shape[-1]
min_depth_diff = _ivy.array([MIN_DEPTH_DIFF], dev_str=dev_str)
red = 'min' if with_db else 'sum'
# shapes as list
batch_shape = list(batch_shape)
final_image_dims = list(final_image_dims)
num_batch_dims = len(batch_shape)
# variance threshold
if isinstance(var_threshold, tuple) or isinstance(var_threshold, list):
ones = _ivy.ones(batch_shape + [1, 2 + d, 1])
var_threshold = _ivy.concatenate(
(ones * var_threshold[0], ones * var_threshold[1]), -1)
else:
var_threshold = _ivy.reshape(var_threshold,
batch_shape + [1, 2 + d, 2])
# uniform pixel coords
if uniform_pixel_coords is None:
uniform_pixel_coords =\
_ivy_svg.create_uniform_pixel_coords_image(final_image_dims, batch_shape, dev_str=dev_str)
uniform_pixel_coords = uniform_pixel_coords[..., 0:2]
# Extract Values #
feat_prior = _ivy.ones_like(feat) * feat_prior if isinstance(
feat_prior, float) else feat_prior
pixel_coords_var = _ivy.ones_like(pixel_coords) * pixel_coords_var\
if isinstance(pixel_coords_var, float) else pixel_coords_var
feat_var = _ivy.ones_like(feat) * feat_var if isinstance(
feat_var, float) else feat_var
pixel_coords_prior_var = _ivy.ones(batch_shape + final_image_dims + [2]) * pixel_coords_prior_var\
if isinstance(pixel_coords_prior_var, float) else pixel_coords_prior_var
feat_prior_var = _ivy.ones(batch_shape + final_image_dims + [d]) * feat_prior_var\
if isinstance(feat_prior_var, float) else feat_prior_var
# Quantize #
# BS x N x 2
quantized_pixel_coords = _ivy.reshape(
_ivy.cast(_ivy.round(pixel_coords), 'int32'), batch_shape + [-1, 2])
# Combine #
# BS x N x (2+D)
pc_n_feat = _ivy.reshape(_ivy.concatenate((pixel_coords, feat), -1),
batch_shape + [-1, 2 + d])
pc_n_feat_var = _ivy.reshape(
_ivy.concatenate((pixel_coords_var, feat_var), -1),
batch_shape + [-1, 2 + d])
# BS x H x W x (2+D)
prior = _ivy.concatenate((uniform_pixel_coords, feat_prior), -1)
prior_var = _ivy.concatenate((pixel_coords_prior_var, feat_prior_var), -1)
# Validity Mask #
# BS x N x 1
var_validity_mask = \
_ivy.reduce_sum(_ivy.cast(pc_n_feat_var < var_threshold[..., 1], 'int32'), -1, keepdims=True) == 2+d
bounds_validity_mask = _ivy.logical_and(
_ivy.logical_and(quantized_pixel_coords[..., 0:1] >= 0,
quantized_pixel_coords[..., 1:2] >= 0),
_ivy.logical_and(
quantized_pixel_coords[..., 0:1] <= final_image_dims[1] - 1,
quantized_pixel_coords[..., 1:2] <= final_image_dims[0] - 1))
validity_mask = _ivy.logical_and(var_validity_mask, bounds_validity_mask)
# num_valid_indices x len(BS)+2
validity_indices = _ivy.reshape(
_ivy.cast(_ivy.indices_where(validity_mask), 'int32'),
[-1, num_batch_dims + 2])
num_valid_indices = validity_indices.shape[-2]
if num_valid_indices == 0:
return _ivy.concatenate((uniform_pixel_coords[..., 0:2], feat_prior), -1), \
_ivy.concatenate((pixel_coords_prior_var, feat_prior_var), -1),\
_ivy.zeros_like(feat[..., 0:1], dev_str=dev_str)
# Depth Based Scaling #
mean_depth_min = None
mean_depth_range = None
pc_n_feat_wo_depth_range = None
pc_n_feat_wo_depth_min = None
var_vals_range = None
var_vals_min = None
if with_db:
# BS x N x 1
mean_depth = pc_n_feat[..., 2:3]
# BS x 1 x 1
mean_depth_min = _ivy.reduce_min(mean_depth, -2, keepdims=True)
mean_depth_max = _ivy.reduce_max(mean_depth, -2, keepdims=True)
mean_depth_range = mean_depth_max - mean_depth_min
# BS x N x 1
scaled_depth = (mean_depth - mean_depth_min) / (
mean_depth_range * min_depth_diff + MIN_DENOMINATOR)
if d == 1:
# BS x 1 x 1+D
pc_n_feat_wo_depth_min = _ivy.zeros(batch_shape + [1, 0],
dev_str=dev_str)
pc_n_feat_wo_depth_range = _ivy.ones(batch_shape + [1, 0],
dev_str=dev_str)
else:
# feat without depth
# BS x N x 1+D
pc_n_feat_wo_depth = _ivy.concatenate(
(pc_n_feat[..., 0:2], pc_n_feat[..., 3:]), -1)
# find the min and max of each value
# BS x 1 x 1+D
pc_n_feat_wo_depth_max = _ivy.reduce_max(
pc_n_feat_wo_depth, -2, keepdims=True) + 1
pc_n_feat_wo_depth_min = _ivy.reduce_min(
pc_n_feat_wo_depth, -2, keepdims=True) - 1
pc_n_feat_wo_depth_range = pc_n_feat_wo_depth_max - pc_n_feat_wo_depth_min
# BS x N x 1+D
normed_pc_n_feat_wo_depth = (pc_n_feat_wo_depth - pc_n_feat_wo_depth_min) / \
(pc_n_feat_wo_depth_range + MIN_DENOMINATOR)
# combine with scaled depth
# BS x N x 1+D
pc_n_feat_wo_depth_scaled = normed_pc_n_feat_wo_depth + scaled_depth
# BS x N x (2+D)
pc_n_feat = _ivy.concatenate(
(pc_n_feat_wo_depth_scaled[..., 0:2], mean_depth,
pc_n_feat_wo_depth_scaled[..., 2:]), -1)
# scale variance
# BS x 1 x (2+D)
var_vals_max = _ivy.reduce_max(pc_n_feat_var, -2, keepdims=True) + 1
var_vals_min = _ivy.reduce_min(pc_n_feat_var, -2, keepdims=True) - 1
var_vals_range = var_vals_max - var_vals_min
# BS x N x (2+D)
normed_var_vals = (pc_n_feat_var - var_vals_min) / (var_vals_range +
MIN_DENOMINATOR)
pc_n_feat_var = normed_var_vals + scaled_depth
# ready for later reversal with full image dimensions
# BS x 1 x 1 x D
var_vals_min = _ivy.expand_dims(var_vals_min, -2)
var_vals_range = _ivy.expand_dims(var_vals_range, -2)
# Validity Pruning #
# num_valid_indices x (2+D)
pc_n_feat = _ivy.gather_nd(pc_n_feat,
validity_indices[..., 0:num_batch_dims + 1])
pc_n_feat_var = _ivy.gather_nd(pc_n_feat_var,
validity_indices[..., 0:num_batch_dims + 1])
# num_valid_indices x 2
quantized_pixel_coords = _ivy.gather_nd(
quantized_pixel_coords, validity_indices[..., 0:num_batch_dims + 1])
if with_db:
means_to_scatter = pc_n_feat
vars_to_scatter = pc_n_feat_var
else:
# num_valid_indices x (2+D)
vars_to_scatter = 1 / (pc_n_feat_var + MIN_DENOMINATOR)
means_to_scatter = pc_n_feat * vars_to_scatter
# Scatter #
# num_valid_indices x 1
counter = _ivy.ones_like(pc_n_feat[..., 0:1], dev_str=dev_str)
if with_db:
counter *= -1
# num_valid_indices x 2(2+D)+1
values_to_scatter = _ivy.concatenate(
(means_to_scatter, vars_to_scatter, counter), -1)
# num_valid_indices x (num_batch_dims + 2)
all_indices = _ivy.flip(quantized_pixel_coords, -1)
if num_batch_dims > 0:
all_indices = _ivy.concatenate(
(validity_indices[..., :-2], all_indices), -1)
# BS x H x W x (2(2+D) + 1)
quantized_img = _ivy.scatter_nd(
_ivy.reshape(all_indices, [-1, num_batch_dims + 2]),
_ivy.reshape(values_to_scatter, [-1, 2 * (2 + d) + 1]),
batch_shape + final_image_dims + [2 * (2 + d) + 1],
reduction='replace' if _ivy.backend == 'mxnd' else red)
# BS x H x W x 1
quantized_counter = quantized_img[..., -1:]
if with_db:
invalidity_mask = quantized_counter != -1
else:
invalidity_mask = quantized_counter == 0
if with_db:
# BS x H x W x (2+D)
quantized_mean_scaled = quantized_img[..., 0:2 + d]
quantized_var_scaled = quantized_img[..., 2 + d:2 * (2 + d)]
# BS x H x W x 1
quantized_depth_mean = quantized_mean_scaled[..., 2:3]
# BS x 1 x 1 x 1
mean_depth_min = _ivy.expand_dims(mean_depth_min, -2)
mean_depth_range = _ivy.expand_dims(mean_depth_range, -2)
# BS x 1 x 1 x (1+D)
pc_n_feat_wo_depth_min = _ivy.expand_dims(pc_n_feat_wo_depth_min, -2)
pc_n_feat_wo_depth_range = _ivy.expand_dims(pc_n_feat_wo_depth_range,
-2)
# BS x 1 x 1 x (2+D) x 2
var_threshold = _ivy.expand_dims(var_threshold, -3)
# BS x H x W x (1+D)
quantized_mean_wo_depth_scaled = _ivy.concatenate(
(quantized_mean_scaled[..., 0:2], quantized_mean_scaled[..., 3:]),
-1)
quantized_mean_wo_depth_normed = quantized_mean_wo_depth_scaled - (quantized_depth_mean - mean_depth_min) / \
(mean_depth_range * min_depth_diff + MIN_DENOMINATOR)
quantized_mean_wo_depth = quantized_mean_wo_depth_normed * pc_n_feat_wo_depth_range + pc_n_feat_wo_depth_min
prior_wo_depth = _ivy.concatenate((prior[..., 0:2], prior[..., 3:]),
-1)
quantized_mean_wo_depth = _ivy.where(invalidity_mask, prior_wo_depth,
quantized_mean_wo_depth)
# BS x H x W x (2+D)
quantized_mean = _ivy.concatenate(
(quantized_mean_wo_depth[..., 0:2], quantized_depth_mean,
quantized_mean_wo_depth[..., 2:]), -1)
# BS x H x W x (2+D)
quantized_var_normed = quantized_var_scaled - (quantized_depth_mean - mean_depth_min) / \
(mean_depth_range * min_depth_diff + MIN_DENOMINATOR)
quantized_var = _ivy.maximum(
quantized_var_normed * var_vals_range + var_vals_min,
var_threshold[..., 0])
quantized_var = _ivy.where(invalidity_mask, prior_var, quantized_var)
else:
# BS x H x W x (2+D)
quantized_sum_mean_x_recip_var = quantized_img[..., 0:2 + d]
quantized_var_wo_increase = _ivy.where(
invalidity_mask, prior_var,
(1 / (quantized_img[..., 2 + d:2 * (2 + d)] + MIN_DENOMINATOR)))
quantized_var = _ivy.maximum(
quantized_var_wo_increase * quantized_counter,
_ivy.expand_dims(var_threshold[..., 0], -2))
quantized_var = _ivy.where(invalidity_mask, prior_var, quantized_var)
quantized_mean = _ivy.where(
invalidity_mask, prior,
quantized_var_wo_increase * quantized_sum_mean_x_recip_var)
# BS x H x W x (2+D) BS x H x W x (2+D) BS x H x W x 1
return quantized_mean, quantized_var, quantized_counter
def coord_image_to_trimesh(coord_img, validity_mask=None, batch_shape=None, image_dims=None, dev_str=None):
"""
Create trimesh, with vertices and triangle indices, from co-ordinate image.
:param coord_img: Image of co-ordinates *[batch_shape,h,w,3]*
:type coord_img: array
:param validity_mask: Boolean mask of where the coord image contains valid values *[batch_shape,h,w,1]*
:type validity_mask: array, optional
:param batch_shape: Shape of batch. Inferred from inputs if None.
:type batch_shape: sequence of ints, optional
:param image_dims: Image dimensions. Inferred from inputs in None.
:type image_dims: sequence of ints, optional
:param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Same as x if None.
:type dev_str: str, optional
:return: Vertices *[batch_shape,(hxw),3]* amd Trimesh indices *[batch_shape,n,3]*
"""
if dev_str is None:
dev_str = _ivy.dev_str(coord_img)
if batch_shape is None:
batch_shape = _ivy.shape(coord_img)[:-3]
if image_dims is None:
image_dims = _ivy.shape(coord_img)[-3:-1]
# shapes as lists
batch_shape = list(batch_shape)
image_dims = list(image_dims)
# BS x (HxW) x 3
vertices = _ivy.reshape(coord_img, batch_shape + [image_dims[0] * image_dims[1], 3])
if validity_mask is not None:
# BS x H-1 x W-1 x 1
t00_validity = validity_mask[..., 0:image_dims[0] - 1, 0:image_dims[1] - 1, :]
t01_validity = validity_mask[..., 0:image_dims[0] - 1, 1:image_dims[1], :]
t02_validity = validity_mask[..., 1:image_dims[0], 0:image_dims[1] - 1, :]
t10_validity = validity_mask[..., 1:image_dims[0], 1:image_dims[1], :]
t11_validity = t01_validity
t12_validity = t02_validity
# BS x H-1 x W-1 x 1
t0_validity = _ivy.logical_and(t00_validity, _ivy.logical_and(t01_validity, t02_validity))
t1_validity = _ivy.logical_and(t10_validity, _ivy.logical_and(t11_validity, t12_validity))
# BS x (H-1xW-1)
t0_validity_flat = _ivy.reshape(t0_validity, batch_shape + [-1])
t1_validity_flat = _ivy.reshape(t1_validity, batch_shape + [-1])
# BS x 2x(H-1xW-1)
trimesh_index_validity = _ivy.concatenate((t0_validity_flat, t1_validity_flat), -1)
# BS x N
trimesh_valid_indices = _ivy.indices_where(trimesh_index_validity)
# BS x 2x(H-1xW-1) x 3
all_trimesh_indices = create_trimesh_indices_for_image(batch_shape, image_dims, dev_str)
# BS x N x 3
trimesh_indices = _ivy.gather_nd(all_trimesh_indices, trimesh_valid_indices)
else:
# BS x N=2x(H-1xW-1) x 3
trimesh_indices = create_trimesh_indices_for_image(batch_shape, image_dims)
# BS x (HxW) x 3, BS x N x 3
return vertices, trimesh_indices