def get_reward(self):
"""
Get reward based on current state
:return: Reward array
"""
# Goal proximity.
rew = ivy.exp(-0.5 * ivy.reduce_sum((self.xy - self.goal_xy)**2, -1))
# Urchins proximity.
rew = rew * ivy.reduce_prod(
1 - ivy.exp(-30 * ivy.reduce_sum(
(self.xy - self.urchin_xys)**2, -1)), -1)
return ivy.reshape(rew, (1, ))
def _addressing(self, k, beta, g, s, gamma, prev_M, prev_w):
# Sec 3.3.1 Focusing by Content
# Cosine Similarity
k = ivy.expand_dims(k, axis=2)
inner_product = ivy.matmul(prev_M, k)
k_norm = ivy.reduce_sum(k**2, axis=1, keepdims=True)**0.5
M_norm = ivy.reduce_sum(prev_M**2, axis=2, keepdims=True)**0.5
norm_product = M_norm * k_norm
K = ivy.squeeze(inner_product / (norm_product + 1e-8)) # eq (6)
# Calculating w^c
K_amplified = ivy.exp(ivy.expand_dims(beta, axis=1) * K)
w_c = K_amplified / ivy.reduce_sum(K_amplified, axis=1,
keepdims=True) # eq (5)
if self._addressing_mode == 'content': # Only focus on content
return w_c
# Sec 3.3.2 Focusing by Location
g = ivy.expand_dims(g, axis=1)
w_g = g * w_c + (1 - g) * prev_w # eq (7)
s = ivy.concatenate([
s[:, :self._shift_range + 1],
ivy.zeros(
[s.shape[0], self._memory_size -
(self._shift_range * 2 + 1)]), s[:, -self._shift_range:]
],
axis=1)
t = ivy.concatenate([ivy.flip(s, axis=[1]),
ivy.flip(s, axis=[1])],
axis=1)
s_matrix = ivy.stack([
t[:, self._memory_size - i - 1:self._memory_size * 2 - i - 1]
for i in range(self._memory_size)
],
axis=1)
w_ = ivy.reduce_sum(ivy.expand_dims(w_g, axis=1) * s_matrix,
axis=2) # eq (8)
w_sharpen = w_**ivy.expand_dims(gamma, axis=1)
w = w_sharpen / ivy.reduce_sum(w_sharpen, axis=1,
keepdims=True) # eq (9)
return w
def train_step(loss_fn_in, optimizer, ntm, total_seq, target_seq, seq_len, mw,
vw, step, max_grad_norm):
# compute loss
loss, dldv, pred_vals = ivy.execute_with_gradients(
lambda v_: loss_fn_in(v_, total_seq, target_seq, seq_len), ntm.v)
global_norm = ivy.reduce_sum(
ivy.stack([ivy.reduce_sum(grad**2) for grad in dldv.to_flat_list()],
0))**0.5
dldv = dldv.map(lambda x, _: x * max_grad_norm / ivy.maximum(
global_norm, max_grad_norm))
# update variables
ntm.v = optimizer.step(ntm.v, dldv)
return loss, pred_vals
def get_reward(self):
"""
Get reward based on current state
:return: Reward array
"""
# Goal proximity.
x = ivy.reduce_sum(ivy.cos(self.angles), -1)
y = ivy.reduce_sum(ivy.sin(self.angles), -1)
xy = ivy.concatenate([ivy.expand_dims(x, 0),
ivy.expand_dims(y, 0)],
axis=0)
rew = ivy.reshape(
ivy.exp(-1 * ivy.reduce_sum((xy - self.goal_xy)**2, -1)), (-1, ))
return ivy.reduce_mean(rew, axis=0, keepdims=True)
def sphere_coords_to_world_ray_vectors(sphere_coords, inv_rotation_mat, batch_shape=None, image_dims=None):
"""
Calculate world-centric ray vector image :math:`\mathbf{RV}\in\mathbb{R}^{h×w×3}` from camera-centric ego-sphere
polar co-ordinates image :math:`\mathbf{S}_c\in\mathbb{R}^{h×w×3}`. Each ray vector
:math:`\mathbf{rv}_{i,j}\in\mathbb{R}^{3}` is represented as a unit vector from the camera center
:math:`\overset{\sim}{\mathbf{C}}\in\mathbb{R}^{3×1}`, in the world frame. Co-ordinates
:math:`\mathbf{x}_{i,j}\in\mathbb{R}^{3}` along the world ray can then be parameterized as
:math:`\mathbf{x}_{i,j}=\overset{\sim}{\mathbf{C}} + λ\mathbf{rv}_{i,j}`, where :math:`λ` is a scalar who's
magnitude dictates the position of the world co-ordinate along the world ray.
:param sphere_coords: Camera-centric ego-sphere polar co-ordinates image *[batch_shape,h,w,3]*
:type sphere_coords: array
:param inv_rotation_mat: Inverse rotation matrix *[batch_shape,3,3]*
:type inv_rotation_mat: array
: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
:return: World ray vectors *[batch_shape,h,w,3]*
"""
if batch_shape is None:
batch_shape = sphere_coords.shape[:-3]
if image_dims is None:
image_dims = sphere_coords.shape[-3:-1]
# shapes as list
batch_shape = list(batch_shape)
image_dims = list(image_dims)
# BS x H x W x 3
cam_coords = sphere_to_cam_coords(sphere_coords, batch_shape=batch_shape, image_dims=image_dims)[..., 0:3]
vectors = _ivy_pg.transform(cam_coords, inv_rotation_mat, batch_shape, image_dims)
return vectors / (_ivy.reduce_sum(vectors ** 2, -1, keepdims=True) ** 0.5 + MIN_DENOMINATOR)
def test_reduce_sum(x, axis, kd, dtype_str, tensor_fn, dev_str, call):
# smoke test
x = tensor_fn(x, dtype_str, dev_str)
ret = ivy.reduce_sum(x, axis, kd)
# type test
assert ivy.is_array(ret)
# cardinality test
if axis is None:
expected_shape = [1] * len(x.shape) if kd else []
else:
axis_ = [axis] if isinstance(axis, int) else axis
axis_ = [item % len(x.shape) for item in axis_]
expected_shape = list(x.shape)
if kd:
expected_shape = [
1 if i % len(x.shape) in axis_ else item
for i, item in enumerate(expected_shape)
]
else:
[expected_shape.pop(item) for item in axis_]
expected_shape = [1] if expected_shape == [] else expected_shape
assert ret.shape == tuple(expected_shape)
# value test
assert np.allclose(call(ivy.reduce_sum, x),
ivy.numpy.reduce_sum(ivy.to_numpy(x)))
# compilation test
helpers.assert_compilable(ivy.reduce_sum)
def weighted_image_smooth(mean, weights, kernel_dim):
"""
Smooth an image using weight values from a weight image of the same size.
:param mean: Image to smooth *[batch_shape,h,w,d]*
:type mean: array
:param weights: Variance image, with the variance values of each pixel in the image *[batch_shape,h,w,d]*
:type weights: array
:param kernel_dim: The dimension of the kernel
:type kernel_dim: int
:return: Image smoothed based on variance image and smoothing kernel.
"""
# shapes as list
kernel_shape = [kernel_dim, kernel_dim]
dim = mean.shape[-1]
# KW x KW x D
kernel = _ivy.ones(kernel_shape + [dim])
# D
kernel_sum = _ivy.reduce_sum(kernel, [0, 1])[0]
# BS x H x W x D
mean_x_weights = mean * weights
mean_x_weights_sum = _ivy.abs(
_ivy.depthwise_conv2d(mean_x_weights, kernel, 1, "VALID"))
sum_of_weights = _ivy.depthwise_conv2d(weights, kernel, 1, "VALID")
new_mean = mean_x_weights_sum / (sum_of_weights + MIN_DENOMINATOR)
new_weights = sum_of_weights / (kernel_sum + MIN_DENOMINATOR)
# BS x H x W x D, # BS x H x W x D
return new_mean, new_weights
def sphere_signed_distances(sphere_positions, sphere_radii, query_positions):
"""
Return the signed distances of a set of query points from the sphere surfaces.\n
`[reference] <https://www.iquilezles.org/www/articles/distfunctions/distfunctions.htm>`_
:param sphere_positions: Positions of the spheres *[batch_shape,num_spheres,3]*
:type sphere_positions: array
:param sphere_radii: Radii of the spheres *[batch_shape,num_spheres,1]*
:type sphere_radii: array
:param query_positions: Points for which to query the signed distances *[batch_shape,num_points,3]*
:type query_positions: array
:return: The distances of the query points from the closest sphere surface *[batch_shape,num_points,1]*
"""
# BS x NS x 1 x 3
sphere_positions = _ivy.expand_dims(sphere_positions, -2)
# BS x 1 x NP x 3
query_positions = _ivy.expand_dims(query_positions, -3)
# BS x NS x NP x 1
distances_to_centre = _ivy.reduce_sum(
(query_positions - sphere_positions)**2, -1, keepdims=True)**0.5
# BS x NS x NP x 1
all_sdfs = distances_to_centre - _ivy.expand_dims(sphere_radii, -2)
# BS x NP x 1
return _ivy.reduce_min(all_sdfs, -3)
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 _fuse_measurements_with_uncertainty(measurements,
measurement_uncertainties, axis):
measurements_shape = measurements.shape
batch_shape = measurements_shape[0:axis]
num_batch_dims = len(batch_shape)
num_measurements = measurements_shape[axis]
# BS x 1 x RS
sum_of_variances = ivy.reduce_sum(measurement_uncertainties,
axis,
keepdims=True)
prod_of_variances = ivy.reduce_prod(measurement_uncertainties,
axis,
keepdims=True)
# BS x 1 x RS
new_variance = prod_of_variances / sum_of_variances
# dim size list of BS x (dim-1) x RS
batch_slices = [slice(None, None, None)] * num_batch_dims
concat_lists = \
[[measurement_uncertainties[tuple(batch_slices + [slice(0, i, None)])]]
if i != 0 else [] + [measurement_uncertainties[tuple(batch_slices + [slice(i + 1, None, None)])]]
for i in range(num_measurements)]
partial_variances_list = [
ivy.concatenate(concat_list, axis) for concat_list in concat_lists
]
# dim size list of BS x 1 x RS
partial_prod_of_variances_list = \
[ivy.reduce_prod(partial_variance, axis, True) for partial_variance in
partial_variances_list]
# BS x dim x RS
partial_prod_of_variances = ivy.concatenate(
partial_prod_of_variances_list, axis)
# BS x 1 x RS
new_mean = ivy.reduce_sum(
(partial_prod_of_variances * measurements) / sum_of_variances,
axis,
keepdims=True)
# BS x 1 x RS, BS x 1 x RS
return new_mean, new_variance
def _pairwise_distance(x, y):
# BS x NX x 1 x 1
try:
x = _ivy.expand_dims(x, -2)
except:
d = 0
# BS x 1 x NY x 1
y = _ivy.expand_dims(y, -3)
# BS x NX x NY
return _ivy.reduce_sum((x - y)**2, -1)
def render_rays_via_termination_probabilities(ray_term_probs,
features,
render_variance=False):
"""
Render features onto the image plane, given rays sampled at radial depths with readings of
feature values and densities at these sample points.
:param ray_term_probs: The ray termination probabilities *[batch_shape,num_samples_per_ray]*
:type ray_term_probs: array
:param features: Feature values at the sample points *[batch_shape,num_samples_per_ray,feat_dim]*
:type features: array
:param render_variance: Whether to also render the feature variance. Default is False.
:type render_variance: bool, optional
:return: The feature renderings along the rays, computed via the termination probabilities *[batch_shape,feat_dim]*
"""
# BS x NSPR
rendering = ivy.reduce_sum(
ivy.expand_dims(ray_term_probs, -1) * features, -2)
if not render_variance:
return rendering
var = ivy.reduce_sum(
ray_term_probs * (ivy.expand_dims(rendering, -2) - features)**2, -2)
return rendering, var
def inverse_quaternion(quaternion):
"""
Compute inverse quaternion :math:`\mathbf{q}^{-1}.\n
`[reference] <https://github.com/KieranWynn/pyquaternion/blob/446c31cba66b708e8480871e70b06415c3cb3b0f/pyquaternion/quaternion.py#L473>`_
:param quaternion: Quaternion *[batch_shape,4]*
:type quaternion: array
:return: Inverse quaternion *[batch_shape,4]*
"""
# BS x 1
sum_of_squares = _ivy.reduce_sum(quaternion**2, -1)
vector_conjugate = _ivy.concatenate(
(-quaternion[..., 0:3], quaternion[..., -1:]), -1)
return vector_conjugate / (sum_of_squares + MIN_DENOMINATOR)
def ds_pixel_coords_to_world_ray_vectors(ds_pixel_coords, inv_full_mat, camera_center=None, batch_shape=None, image_dims=None):
"""
Calculate world-centric ray vector image :math:`\mathbf{RV}\in\mathbb{R}^{h×w×3}` from homogeneous pixel co-ordinate
image :math:`\mathbf{X}_p\in\mathbb{R}^{h×w×3}`. Each ray vector :math:`\mathbf{rv}_{i,j}\in\mathbb{R}^{3}` is
represented as a unit vector from the camera center :math:`\overset{\sim}{\mathbf{C}}\in\mathbb{R}^{3×1}`, in the
world frame. Co-ordinates :math:`\mathbf{x}_{i,j}\in\mathbb{R}^{3}` along the world ray can then be parameterized as
:math:`\mathbf{x}_{i,j}=\overset{\sim}{\mathbf{C}} + λ\mathbf{rv}_{i,j}`, where :math:`λ` is a scalar who's
magnitude dictates the position of the world co-ordinate along the world ray.
:param ds_pixel_coords: Homogeneous pixel co-ordinates image *[batch_shape,h,w,3]*
:type ds_pixel_coords: array
:param inv_full_mat: Inverse full projection matrix *[batch_shape,3,4]*
:type inv_full_mat: array
:param camera_center: Camera centers, inferred from inv_full_mat if None *[batch_shape,3,1]*
:type camera_center: 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
:return: World ray vectors *[batch_shape,h,w,3]*
"""
if batch_shape is None:
batch_shape = ds_pixel_coords.shape[:-3]
if image_dims is None:
image_dims = ds_pixel_coords.shape[-3:-1]
# shapes as list
batch_shape = list(batch_shape)
image_dims = list(image_dims)
if camera_center is None:
camera_center = inv_ext_mat_to_camera_center(inv_full_mat)
# BS x 1 x 1 x 3
camera_centers_reshaped = _ivy.reshape(camera_center, batch_shape + [1, 1, 3])
# BS x H x W x 3
vectors = ds_pixel_to_world_coords(ds_pixel_coords, inv_full_mat, batch_shape, image_dims)[..., 0:3] \
- camera_centers_reshaped
# BS x H x W x 3
return vectors / (_ivy.reduce_sum(vectors ** 2, -1, keepdims=True) ** 0.5 + MIN_DENOMINATOR)
def rotation_vector_to_quaternion(rot_vector):
"""
Convert rotation vector :math:`\mathbf{θ}_{rv} = [θe_x, θe_y, θe_z]` to quaternion
:math:`\mathbf{q} = [q_i, q_j, q_k, q_r]`.\n
`[reference] <https://en.wikipedia.org/wiki/Rotation_formalisms_in_three_dimensions#Euler_axis_and_angle_(rotation_vector)>`_
:param rot_vector: Rotation vector *[batch_shape,3]*
:type rot_vector: array
:return: Quaternion *[batch_shape,4]*
"""
# BS x 1
theta = (_ivy.reduce_sum(rot_vector**2, axis=-1, keepdims=True))**0.5
# BS x 3
vector = rot_vector / (theta + MIN_DENOMINATOR)
# BS x 4
return axis_angle_to_quaternion(_ivy.concatenate([vector, theta], -1))
def depth_to_radial_depth(depth, inv_calib_mat, uniform_pixel_coords=None, batch_shape=None, image_dims=None):
"""
Get radial depth image :math:`\mathbf{X}_{rd}\in\mathbb{R}^{h×w×1}` from depth image
:math:`\mathbf{X}_d\in\mathbb{R}^{h×w×1}`.\n
:param depth: Depth image *[batch_shape,h,w,1]*
:type depth: array
:param inv_calib_mat: Inverse calibration matrix *[batch_shape,3,3]*
:type inv_calib_mat: array
:param uniform_pixel_coords: Image of homogeneous pixel co-ordinates. Created if None. *[batch_shape,h,w,3]*
:type uniform_pixel_coords: 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
:return: Radial depth image *[batch_shape,h,w,1]*
"""
if batch_shape is None:
batch_shape = depth.shape[:-3]
if image_dims is None:
image_dims = depth.shape[-3:-1]
# shapes as list
batch_shape = list(batch_shape)
image_dims = list(image_dims)
# BS x H x W x 3
if uniform_pixel_coords is None:
uniform_pixel_coords = create_uniform_pixel_coords_image(image_dims, batch_shape, dev_str=_ivy.dev_str(depth))
# BS x H x W x 3
ds_pixel_coords = depth_to_ds_pixel_coords(depth, uniform_pixel_coords, batch_shape, image_dims)
# BS x H x W x 3
cam_coords = ds_pixel_to_cam_coords(ds_pixel_coords, inv_calib_mat, batch_shape, image_dims)[..., 0:3]
# BS x H x W x 1
return _ivy.reduce_sum(cam_coords**2, -1, keepdims=True)**0.5
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 compute_length(query_vals):
start_vals = query_vals[:, 0:-1]
end_vals = query_vals[:, 1:]
dists_sqrd = ivy.maximum((end_vals - start_vals)**2, 1e-12)
distances = ivy.reduce_sum(dists_sqrd, -1)**0.5
return ivy.reduce_mean(ivy.reduce_sum(distances, 1))
def smooth_image_fom_var_image(mean,
var,
kernel_dim,
kernel_scale,
dev_str=None):
"""
Smooth an image using variance values from a variance image of the same size, and a spatial smoothing kernel.
:param mean: Image to smooth *[batch_shape,h,w,d]*
:type mean: array
:param var: Variance image, with the variance values of each pixel in the image *[batch_shape,h,w,d]*
:type var: array
:param kernel_dim: The dimension of the kernel
:type kernel_dim: int
:param kernel_scale: The scale of the kernel along the channel dimension *[d]*
:type kernel_scale: array
: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: Image smoothed based on variance image and smoothing kernel.
"""
if dev_str is None:
dev_str = _ivy.dev_str(mean)
# shapes as list
kernel_shape = [kernel_dim, kernel_dim]
kernel_size = kernel_dim**2
dims = mean.shape[-1]
# KH x KW x 2
uniform_pixel_coords = _ivy_svg.create_uniform_pixel_coords_image(
kernel_shape, dev_str=dev_str)[..., 0:2]
# 2
kernel_central_pixel_coord = _ivy.array([
float(_math.floor(kernel_shape[0] / 2)),
float(_math.floor(kernel_shape[1] / 2))
],
dev_str=dev_str)
# KH x KW x 2
kernel_xy_dists = kernel_central_pixel_coord - uniform_pixel_coords
kernel_xy_dists_sqrd = kernel_xy_dists**2
# KW x KW x D x D
unit_kernel = _ivy.tile(
_ivy.reduce_sum(kernel_xy_dists_sqrd, -1, keepdims=True)**0.5,
(1, 1, dims))
kernel = 1 + unit_kernel * kernel_scale
recip_kernel = 1 / (kernel + MIN_DENOMINATOR)
# D
kernel_sum = _ivy.reduce_sum(kernel, [0, 1])[0]
recip_kernel_sum = _ivy.reduce_sum(recip_kernel, [0, 1])
# BS x H x W x D
recip_var = 1 / (var + MIN_DENOMINATOR)
recip_var_scaled = recip_var + 1
recip_new_var_scaled = _ivy.depthwise_conv2d(recip_var_scaled,
recip_kernel, 1, "VALID")
# This 0.99 prevents float32 rounding errors leading to -ve variances, the true equation would use 1.0
recip_new_var = recip_new_var_scaled - recip_kernel_sum * 0.99
new_var = 1 / (recip_new_var + MIN_DENOMINATOR)
mean_x_recip_var = mean * recip_var
mean_x_recip_var_sum = _ivy.abs(
_ivy.depthwise_conv2d(mean_x_recip_var, recip_kernel, 1, "VALID"))
new_mean = new_var * mean_x_recip_var_sum
new_var = new_var * kernel_size**2 / (kernel_sum + MIN_DENOMINATOR)
# prevent overconfidence from false meas independence assumption
# BS x H x W x D, # BS x H x W x D
return new_mean, new_var
def loss_fn(ntm, v, total_seq, target_seq, seq_len):
output_sequence = ntm(total_seq, v=v)
pred_logits = output_sequence[:, seq_len + 1:, :]
pred_vals = ivy.sigmoid(pred_logits)
return ivy.reduce_sum(ivy.binary_cross_entropy(
pred_vals, target_seq))[0] / pred_vals.shape[0], pred_vals
def _kalman_filter_on_measurement_sequence(
self, prev_fused_val, prev_fused_variance, hole_prior,
hole_prior_var, meas, meas_vars, uniform_sphere_pixel_coords,
agent_rel_poses, agent_rel_pose_covs, agent_rel_mats, batch_size,
num_timesteps):
"""
Perform kalman filter on measurement sequence
:param prev_fused_val: Fused value from previous timestamp *[batch_size, oh, ow, (3+f)]*
:param prev_fused_variance: Fused variance from previous timestamp *[batch_size, oh, ow, (3+f)]*
:param hole_prior: Prior for holes in quantization *[batch_size, oh, ow, (1+f)]*
:param hole_prior_var: Prior variance for holes in quantization *[batch_size, oh, ow, (3+f)]*
:param meas: Measurements *[batch_size, num_timesteps, oh, ow, (3+f)]*
:param meas_vars: Measurement variances *[batch_size, num_timesteps, oh, ow, (3+f)]*
:param uniform_sphere_pixel_coords: Uniform sphere pixel co-ordinates *[batch_size, oh, ow, 3]*
:param agent_rel_poses: Relative poses of agents to the previous step *[batch_size, num_timesteps, 6]*
:param agent_rel_pose_covs: Agent relative pose covariances *[batch_size, num_timesteps, 6, 6]*
:param agent_rel_mats: Relative transformations matrix of agents to the previous step
*[batch_size, num_timesteps, 3, 4]*
:param batch_size: Size of batch
:param num_timesteps: Number of frames
:return: list of *[batch_size, oh, ow, (3+f)]*, list of *[batch_size, oh, ow, (3+f)]*
"""
fused_list = list()
fused_variances_list = list()
for i in range(num_timesteps):
# project prior from previous frame #
# ----------------------------------#
# B x OH x OW x (3+F)
prev_prior = prev_fused_val
prev_prior_variance = prev_fused_variance
# B x 3 x 4
agent_rel_mat = agent_rel_mats[:, i]
# B x 6
agent_rel_pose = agent_rel_poses[:, i]
# B x 6 x 6
agent_rel_pose_cov = agent_rel_pose_covs[:, i]
# B x OH x OW x (3+F) B x OH x OW x (3+F)
fused_projected, fused_projected_variance = self._omni_frame_to_omni_frame_projection(
agent_rel_pose, agent_rel_mat, uniform_sphere_pixel_coords,
prev_prior[..., 0:2], prev_prior[..., 2:3],
prev_prior[..., 3:], agent_rel_pose_cov, prev_prior_variance,
hole_prior, hole_prior_var, batch_size)
# reset prior
# B x OH x OW x (3+F)
prior = fused_projected
prior_var = fused_projected_variance
# per-pixel fusion with measurements #
# -----------------------------------#
# extract slice for frame
# B x OH x OW x (3+F)
measurement = meas[:, i]
measurement_variance = meas_vars[:, i]
# fuse prior and measurement
# B x 2 x OH x OW x (3+F)
prior_and_meas = ivy.concatenate(
(ivy.expand_dims(prior, 1), ivy.expand_dims(measurement, 1)),
1)
prior_and_meas_variance = ivy.concatenate((ivy.expand_dims(
prior_var, 1), ivy.expand_dims(measurement_variance, 1)), 1)
# B x OH x OW x (3+F)
low_var_mask = ivy.reduce_sum(
ivy.cast(
prior_and_meas_variance <
ivy.expand_dims(hole_prior_var, 1) *
self._threshold_var_factor, 'int32'), 1) > 0
# B x 1 x OH x OW x (3+F) B x 1 x OH x OW x (3+F)
# ToDo: handle this properly once re-implemented with a single scatter operation only
# currently depth values are fused even if these are clearly far apart
fused_val_unsmoothed, fused_variance_unsmoothed = \
self._fuse_measurements_with_uncertainty(prior_and_meas, prior_and_meas_variance, 1)
# B x OH x OW x (3+F)
# This prevents accumulating certainty from duplicate re-projections from prior measurements
fused_variance_unsmoothed = ivy.where(
low_var_mask, fused_variance_unsmoothed[:, 0], hole_prior_var)
# B x OH x OW x (3+F)
fused_val = fused_val_unsmoothed[:, 0]
fused_variance = fused_variance_unsmoothed
low_var_mask = fused_variance < hole_prior_var
# B x OH x OW x (3+F) B x OH x OW x (3+F)
fused_val, fused_variance = self.smooth(fused_val, fused_variance,
low_var_mask,
self._smooth_mean,
self._smooth_kernel_size,
True, True, batch_size)
# append to list for returning
# B x OH x OW x (3+F)
fused_list.append(fused_val)
# B x OH x OW x (3+F)
fused_variances_list.append(fused_variance)
# update for next time step
prev_fused_val = fused_val
prev_fused_variance = fused_variance
# list of *[batch_size, oh, ow, (3+f)]*, list of *[batch_size, oh, ow, (3+f)]*
return fused_list, fused_variances_list
def cuboid_signed_distances(cuboid_ext_mats,
cuboid_dims,
query_positions,
batch_shape=None):
"""
Return the signed distances of a set of query points from the cuboid surfaces.\n
`[reference] <https://www.iquilezles.org/www/articles/distfunctions/distfunctions.htm>`_
:param cuboid_ext_mats: Extrinsic matrices of the cuboids *[batch_shape,num_cuboids,3,4]*
:type cuboid_ext_mats: array
:param cuboid_dims: Dimensions of the cuboids, in the order x, y, z *[batch_shape,num_cuboids,3]*
:type cuboid_dims: array
:param query_positions: Points for which to query the signed distances *[batch_shape,num_points,3]*
:type query_positions: array
:param batch_shape: Shape of batch. Assumed no batches if None.
:type batch_shape: sequence of ints, optional
:return: The distances of the query points from the closest cuboid surface *[batch_shape,num_points,1]*
"""
if batch_shape is None:
batch_shape = cuboid_ext_mats.shape[:-3]
# shapes as list
batch_shape = list(batch_shape)
num_batch_dims = len(batch_shape)
batch_dims_for_trans = list(range(num_batch_dims))
num_cuboids = cuboid_ext_mats.shape[-3]
num_points = query_positions.shape[-2]
# BS x 3 x NP
query_positions_trans = _ivy.transpose(
query_positions,
batch_dims_for_trans + [num_batch_dims + 1, num_batch_dims])
# BS x 1 x NP
ones = _ivy.ones_like(query_positions_trans[..., 0:1, :])
# BS x 4 x NP
query_positions_trans_homo = _ivy.concatenate(
(query_positions_trans, ones), -2)
# BS x NCx3 x 4
cuboid_ext_mats_flat = _ivy.reshape(cuboid_ext_mats, batch_shape + [-1, 4])
# BS x NCx3 x NP
rel_query_positions_trans_flat = _ivy.matmul(cuboid_ext_mats_flat,
query_positions_trans_homo)
# BS x NC x 3 x NP
rel_query_positions_trans = _ivy.reshape(
rel_query_positions_trans_flat,
batch_shape + [num_cuboids, 3, num_points])
# BS x NC x NP x 3
rel_query_positions = _ivy.transpose(
rel_query_positions_trans, batch_dims_for_trans +
[num_batch_dims, num_batch_dims + 2, num_batch_dims + 1])
q = _ivy.abs(rel_query_positions) - _ivy.expand_dims(cuboid_dims / 2, -2)
q_max_clipped = _ivy.maximum(q, 1e-12)
# BS x NC x NP x 1
q_min_clipped = _ivy.minimum(_ivy.reduce_max(q, -1, keepdims=True), 0.)
q_max_clipped_len = _ivy.reduce_sum(q_max_clipped**2, -1,
keepdims=True)**0.5
sdfs = q_max_clipped_len + q_min_clipped
# BS x NP x 1
return _ivy.reduce_min(sdfs, -3)
def velocity_from_flow_cam_coords_and_cam_mats(flow_t_to_tm1,
cam_coords_t,
cam_coords_tm1,
cam_tm1_to_t_ext_mat,
delta_t,
uniform_pixel_coords=None,
batch_shape=None,
image_dims=None,
dev_str=None):
"""
Compute relative cartesian velocity from optical flow, camera co-ordinates, and camera extrinsics.
:param flow_t_to_tm1: Optical flow from frame t to t-1 *[batch_shape,h,w,2]*
:type flow_t_to_tm1: array
:param cam_coords_t: Camera-centric homogeneous co-ordinates image in frame t *[batch_shape,h,w,4]*
:type cam_coords_t: array
:param cam_coords_tm1: Camera-centric homogeneous co-ordinates image in frame t-1 *[batch_shape,h,w,4]*
:type cam_coords_tm1: array
:param cam_tm1_to_t_ext_mat: Camera t-1 to camera t extrinsic projection matrix *[batch_shape,3,4]*
:type cam_tm1_to_t_ext_mat: array
:param delta_t: Time difference between frame at timestep t-1 and t *[batch_shape,1]*
:type delta_t: array
:param uniform_pixel_coords: Homogeneous uniform (integer) pixel co-ordinate images, inferred from image_dims if None *[batch_shape,h,w,3]*
:type uniform_pixel_coords: 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: Cartesian velocity measurements relative to the camera *[batch_shape,h,w,3]*
"""
if batch_shape is None:
batch_shape = flow_t_to_tm1.shape[:-3]
if image_dims is None:
image_dims = flow_t_to_tm1.shape[-3:-1]
# shapes as list
batch_shape = list(batch_shape)
image_dims = list(image_dims)
if dev_str is None:
dev_str = _ivy.dev_str(flow_t_to_tm1)
if uniform_pixel_coords is None:
uniform_pixel_coords = _ivy_svg.create_uniform_pixel_coords_image(
image_dims, batch_shape, dev_str)
# Interpolate cam coords from frame t-1
# BS x H x W x 2
warp = uniform_pixel_coords[..., 0:2] + flow_t_to_tm1
# BS x H x W x 4
cam_coords_tm1_interp = _ivy.image.bilinear_resample(cam_coords_tm1, warp)
# Project to frame t
# BS x H x W x 4
cam_coords_t_proj = _ivy_tvg.cam_to_cam_coords(cam_coords_tm1_interp,
cam_tm1_to_t_ext_mat,
batch_shape, image_dims)
# delta co-ordinates
# BS x H x W x 3
delta_cam_coords_t = (cam_coords_t - cam_coords_t_proj)[..., 0:3]
# velocity
# BS x H x W x 3
vel = delta_cam_coords_t / _ivy.reshape(delta_t, batch_shape + [1] * 3)
# Validity mask
# BS x H x W x 1
validity_mask = \
_ivy.reduce_sum(_ivy.cast(warp < _ivy.array([image_dims[1], image_dims[0]], 'float32', dev_str=dev_str),
'int32'), -1, keepdims=True) == 2
# pruned
# BS x H x W x 3, BS x H x W x 1
return _ivy.where(validity_mask, vel,
_ivy.zeros_like(vel, dev_str=dev_str)), validity_mask
def _forward(self, x, prev_state):
prev_read_vector_list = prev_state[1]
controller_input = ivy.concatenate([x] + prev_read_vector_list, axis=1)
controller_output, controller_state = self._controller(ivy.expand_dims(controller_input, -2),
initial_state=prev_state[0])
controller_output = controller_output[..., -1, :]
parameters = self._controller_proj(controller_output)
parameters = ivy.clip(parameters, -self._clip_value, self._clip_value)
head_parameter_list = \
ivy.split(parameters[:, :self._num_parameters_per_head * self._num_heads], self._num_heads,
axis=1)
erase_add_list = ivy.split(parameters[:, self._num_parameters_per_head * self._num_heads:],
2 * self._write_head_num, axis=1)
prev_w_list = prev_state[2]
prev_M = prev_state[4]
w_list = []
for i, head_parameter in enumerate(head_parameter_list):
k = ivy.tanh(head_parameter[:, 0:self._memory_vector_dim])
beta = ivy.softplus(head_parameter[:, self._memory_vector_dim])
g = ivy.sigmoid(head_parameter[:, self._memory_vector_dim + 1])
s = ivy.softmax(
head_parameter[:, self._memory_vector_dim + 2:self._memory_vector_dim +
2 + (self._shift_range * 2 + 1)])
gamma = ivy.softplus(head_parameter[:, -1]) + 1
w = self._addressing(k, beta, g, s, gamma, prev_M, prev_w_list[i])
w_list.append(w)
# Reading (Sec 3.1)
read_w_list = w_list[:self._read_head_num]
if self._step == 0:
usage_indicator = ivy.zeros_like(w_list[0])
else:
usage_indicator = prev_state[3] + ivy.reduce_sum(ivy.concatenate(read_w_list, 0))
read_vector_list = []
for i in range(self._read_head_num):
read_vector = ivy.reduce_sum(ivy.expand_dims(read_w_list[i], axis=2) * prev_M, axis=1)
read_vector_list.append(read_vector)
# Writing (Sec 3.2)
prev_wrtie_w_list = prev_w_list[self._read_head_num:]
w_wr_size = math.ceil(self._memory_size / 2) if self._retroactive_updates else self._memory_size
if self._sequential_writing:
batch_size = ivy.shape(x)[0]
if self._step < w_wr_size:
w_wr_list = [ivy.tile(ivy.cast(ivy.one_hot(
ivy.array([self._step]), w_wr_size), 'float32'),
(batch_size, 1))] * self._write_head_num
else:
batch_idxs = ivy.expand_dims(ivy.arange(batch_size, 0), -1)
mem_idxs = ivy.expand_dims(ivy.argmax(usage_indicator[..., :w_wr_size], -1), -1)
total_idxs = ivy.concatenate((batch_idxs, mem_idxs), -1)
w_wr_list = [ivy.scatter_nd(total_idxs, ivy.ones((batch_size,)),
(batch_size, w_wr_size))] * self._write_head_num
else:
w_wr_list = w_list[self._read_head_num:]
if self._retroactive_updates:
w_ret_list = [self._retroactive_discount * prev_wrtie_w[..., w_wr_size:] +
(1 - self._retroactive_discount) * prev_wrtie_w[..., :w_wr_size]
for prev_wrtie_w in prev_wrtie_w_list]
w_wrtie_list = [ivy.concatenate((w_wr, w_ret), -1) for w_wr, w_ret in zip(w_wr_list, w_ret_list)]
else:
w_wrtie_list = w_wr_list
M = prev_M
for i in range(self._write_head_num):
w = ivy.expand_dims(w_wrtie_list[i], axis=2)
if self._with_erase:
erase_vector = ivy.expand_dims(ivy.sigmoid(erase_add_list[i * 2]), axis=1)
M = M * ivy.ones(ivy.shape(M)) - ivy.matmul(w, erase_vector)
add_vector = ivy.expand_dims(ivy.tanh(erase_add_list[i * 2 + 1]), axis=1)
M = M + ivy.matmul(w, add_vector)
NTM_output = self._output_proj(ivy.concatenate([controller_output] + read_vector_list, axis=1))
NTM_output = ivy.clip(NTM_output, -self._clip_value, self._clip_value)
self._step += 1
return NTM_output, NTMControllerState(
controller_state=controller_state, read_vector_list=read_vector_list, w_list=w_list,
usage_indicator=usage_indicator, M=M)
def _kalman_filter_on_measurement_sequence(
self, prev_fused_val, prev_fused_variance, hole_prior,
hole_prior_var, meas, meas_vars, uniform_sphere_pixel_coords,
agent_rel_poses, agent_rel_pose_covs, agent_rel_mats, batch_size,
num_timesteps):
"""
Perform kalman filter on measurement sequence
:param prev_fused_val: Fused value from previous timestamp *[batch_size, oh, ow, (3+f)]*
:param prev_fused_variance: Fused variance from previous timestamp *[batch_size, oh, ow, (3+f)]*
:param hole_prior: Prior for holes in quantization *[batch_size, oh, ow, (1+f)]*
:param hole_prior_var: Prior variance for holes in quantization *[batch_size, oh, ow, (3+f)]*
:param meas: Measurements *[batch_size, num_timesteps, oh, ow, (3+f)]*
:param meas_vars: Measurement variances *[batch_size, num_timesteps, oh, ow, (3+f)]*
:param uniform_sphere_pixel_coords: Uniform sphere pixel co-ordinates *[batch_size, oh, ow, 3]*
:param agent_rel_poses: Relative poses of agents to the previous step *[batch_size, num_timesteps, 6]*
:param agent_rel_pose_covs: Agent relative pose covariances *[batch_size, num_timesteps, 6, 6]*
:param agent_rel_mats: Relative transformations matrix of agents to the previous step
*[batch_size, num_timesteps, 3, 4]*
:param batch_size: Size of batch
:param num_timesteps: Number of frames
:return: list of *[batch_size, oh, ow, (3+f)]*, list of *[batch_size, oh, ow, (3+f)]*
"""
fused_list = list()
fused_variances_list = list()
for i in range(num_timesteps):
# project prior from previous frame #
# ----------------------------------#
# B x OH x OW x (3+F)
prev_prior = prev_fused_val
prev_prior_variance = prev_fused_variance
# B x 3 x 4
agent_rel_mat = agent_rel_mats[:, i]
# B x 6
agent_rel_pose = agent_rel_poses[:, i]
# B x 6 x 6
agent_rel_pose_cov = agent_rel_pose_covs[:, i]
# B x OH x OW x (3+F) B x OH x OW x (3+F)
fused_projected, fused_projected_variance = self._omni_frame_to_omni_frame_projection(
agent_rel_pose, agent_rel_mat, uniform_sphere_pixel_coords,
prev_prior[..., 0:2], prev_prior[..., 2:3],
prev_prior[..., 3:], agent_rel_pose_cov, prev_prior_variance,
hole_prior, hole_prior_var, batch_size)
# reset prior
# B x OH x OW x (3+F)
prior = fused_projected
prior_var = fused_projected_variance
# per-pixel fusion with measurements #
# -----------------------------------#
# extract slice for frame
# B x OH x OW x (3+F)
measurement = meas[:, i]
measurement_variance = meas_vars[:, i]
# fuse prior and measurement
# B x 2 x OH x OW x (3+F)
prior_and_meas = ivy.concatenate(
(ivy.expand_dims(prior, 1), ivy.expand_dims(measurement, 1)),
1)
prior_and_meas_variance = ivy.concatenate((ivy.expand_dims(
prior_var, 1), ivy.expand_dims(measurement_variance, 1)), 1)
# B x OH x OW x (3+F)
low_var_mask = ivy.reduce_sum(
ivy.cast(
prior_and_meas_variance <
ivy.expand_dims(hole_prior_var, 1) *
self._threshold_var_factor, 'int32'), 1) > 0
# B x 1 x OH x OW x (3+F) B x 1 x OH x OW x (3+F)
if self._with_depth_buffer:
# ToDo: make this more efficient
prior_low_var_mask = ivy.reduce_max(ivy.cast(
prior_var >= hole_prior_var * self._threshold_var_factor,
'int32'),
-1,
keepdims=True) == 0
meas_low_var_mask = ivy.reduce_max(ivy.cast(
measurement_variance >=
hole_prior_var * self._threshold_var_factor, 'int32'),
-1,
keepdims=True) == 0
neiter_low_var_mask = ivy.logical_and(
ivy.logical_not(prior_low_var_mask),
ivy.logical_not(meas_low_var_mask))
prior_w_large = ivy.where(prior_low_var_mask, prior,
ivy.ones_like(prior) * 1e12)
meas_w_large = ivy.where(meas_low_var_mask, measurement,
ivy.ones_like(measurement) * 1e12)
prior_and_meas_w_large = ivy.concatenate((ivy.expand_dims(
prior_w_large, 1), ivy.expand_dims(meas_w_large, 1)), 1)
fused_val_unsmoothed = ivy.reduce_min(prior_and_meas_w_large,
1,
keepdims=True)
fused_val_unsmoothed = ivy.where(neiter_low_var_mask, prior,
fused_val_unsmoothed)
# ToDo: solve this variance correspondence properly, rather than assuming the most certain
fused_variance_unsmoothed = ivy.reduce_min(
prior_and_meas_variance, 1, keepdims=True)
else:
fused_val_unsmoothed, fused_variance_unsmoothed = \
self._fuse_measurements_with_uncertainty(prior_and_meas, prior_and_meas_variance, 1)
# B x OH x OW x (3+F)
# This prevents accumulating certainty from duplicate re-projections from prior measurements
fused_variance_unsmoothed = ivy.where(
low_var_mask, fused_variance_unsmoothed[:, 0], hole_prior_var)
# B x OH x OW x (3+F)
fused_val = fused_val_unsmoothed[:, 0]
fused_variance = fused_variance_unsmoothed
low_var_mask = fused_variance < hole_prior_var
# B x OH x OW x (3+F) B x OH x OW x (3+F)
fused_val, fused_variance = self.smooth(fused_val, fused_variance,
low_var_mask,
self._smooth_mean,
self._smooth_kernel_size,
True, True, batch_size)
# append to list for returning
# B x OH x OW x (3+F)
fused_list.append(fused_val)
# B x OH x OW x (3+F)
fused_variances_list.append(fused_variance)
# update for next time step
prev_fused_val = fused_val
prev_fused_variance = fused_variance
# list of *[batch_size, oh, ow, (3+f)]*, list of *[batch_size, oh, ow, (3+f)]*
return fused_list, fused_variances_list
def closest_mutual_points_along_two_skew_rays(camera_centers,
world_ray_vectors,
batch_shape=None,
image_dims=None,
dev_str=None):
"""
Compute closest mutual homogeneous co-ordinates :math:`\mathbf{x}_{1,i,j}\in\mathbb{R}^{4}` and
:math:`\mathbf{x}_{2,i,j}\in\mathbb{R}^{4}` along two world-centric rays
:math:`\overset{\sim}{\mathbf{C}_1} + λ_1\mathbf{rv}_{1,i,j}` and
:math:`\overset{\sim}{\mathbf{C}_2} + λ_2\mathbf{rv}_{2,i,j}`, for each index aligned pixel between two
world-centric ray vector images :math:`\mathbf{RV}_1\in\mathbb{R}^{h×w×3}` and
:math:`\mathbf{RV}_2\in\mathbb{R}^{h×w×3}`. The function returns two images of closest mutual homogeneous
co-ordinates :math:`\mathbf{X}_1\in\mathbb{R}^{h×w×4}` and :math:`\mathbf{X}_2\in\mathbb{R}^{h×w×4}`,
concatenated together into a single array.\n
`[reference] <https://math.stackexchange.com/questions/1414285/location-of-shortest-distance-between-two-skew-lines-in-3d>`_
second answer in forum
:param camera_centers: Camera center *[batch_shape,2,3,1]*
:type camera_centers: array
:param world_ray_vectors: World ray vectors *[batch_shape,2,h,w,3]*
:type world_ray_vectors: array
: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: Closest mutual points image *[batch_shape,2,h,w,4]*
"""
if batch_shape is None:
batch_shape = world_ray_vectors.shape[:-4]
if image_dims is None:
image_dims = world_ray_vectors.shape[-3:-1]
if dev_str is None:
dev_str = _ivy.dev_str(camera_centers)
# shapes as list
batch_shape = list(batch_shape)
image_dims = list(image_dims)
# BS x 3 x 1
camera_center0 = camera_centers[..., 0, :, :]
camera_center1 = camera_centers[..., 1, :, :]
# BS x 1 x 1 x 3
cam1_to_cam2 = _ivy.reshape(camera_center1 - camera_center0,
batch_shape + [1, 1, 3])
cam2_to_cam1 = _ivy.reshape(camera_center0 - camera_center1,
batch_shape + [1, 1, 3])
# BS x 2 x H x W x 3
ds = world_ray_vectors
# BS x H x W x 3
ds0 = ds[..., 0, :, :, :]
ds1 = ds[..., 1, :, :, :]
n = _ivy.cross(ds0, ds1)
n1 = _ivy.cross(ds0, n)
n2 = _ivy.cross(ds1, n)
# BS x 1 x H x W
t1 = _ivy.expand_dims(
_ivy.reduce_sum(cam1_to_cam2 * n2, -1) /
(_ivy.reduce_sum(ds0 * n2, -1) + MIN_DENOMINATOR), -3)
t2 = _ivy.expand_dims(
_ivy.reduce_sum(cam2_to_cam1 * n1, -1) /
(_ivy.reduce_sum(ds1 * n1, -1) + MIN_DENOMINATOR), -3)
# BS x 2 x H x W
ts = _ivy.expand_dims(_ivy.concatenate((t1, t2), -3), -1)
# BS x 2 x H x W x 3
world_coords = _ivy.reshape(camera_centers[..., 0], batch_shape +
[2, 1, 1, 3]) + ts * world_ray_vectors
# BS x 2 x H x W x 4
return _ivy.concatenate(
(world_coords,
_ivy.ones(batch_shape + [2] + image_dims + [1], dev_str=dev_str)), -1)
def project_cam_coords_with_object_transformations(cam_coords_1,
id_image,
obj_ids,
obj_trans,
cam_1_to_2_ext_mat,
batch_shape=None,
image_dims=None):
"""
Compute velocity image from co-ordinate image, id image, and object transformations.
:param cam_coords_1: Camera-centric homogeneous co-ordinates image in frame t *[batch_shape,h,w,4]*
:type cam_coords_1: array
:param id_image: Image containing per-pixel object ids *[batch_shape,h,w,1]*
:type id_image: array
:param obj_ids: Object ids *[batch_shape,num_obj,1]*
:type obj_ids: array
:param obj_trans: Object transformations for this frame over time *[batch_shape,num_obj,3,4]*
:type obj_trans: array
:param cam_1_to_2_ext_mat: Camera 1 to camera 2 extrinsic projection matrix *[batch_shape,3,4]*
:type cam_1_to_2_ext_mat: array
: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
:return: Relative velocity image *[batch_shape,h,w,3]*
"""
if batch_shape is None:
batch_shape = cam_coords_1.shape[:-3]
if image_dims is None:
image_dims = cam_coords_1.shape[-3:-1]
# shapes as list
batch_shape = list(batch_shape)
image_dims = list(image_dims)
num_batch_dims = len(batch_shape)
# Transform the co-ordinate image by each transformation
# BS x (num_obj x 3) x 4
obj_trans = _ivy.reshape(obj_trans, batch_shape + [-1, 4])
# BS x 4 x H x W
cam_coords_1_ = _ivy.transpose(
cam_coords_1,
list(range(num_batch_dims)) + [i + num_batch_dims for i in [2, 0, 1]])
# BS x 4 x (HxW)
cam_coords_1_ = _ivy.reshape(cam_coords_1_, batch_shape + [4, -1])
# BS x (num_obj x 3) x (HxW)
cam_coords_2_all_obj_trans = _ivy.matmul(obj_trans, cam_coords_1_)
# BS x (HxW) x (num_obj x 3)
cam_coords_2_all_obj_trans = \
_ivy.transpose(cam_coords_2_all_obj_trans, list(range(num_batch_dims)) + [i + num_batch_dims for i in [1, 0]])
# BS x H x W x num_obj x 3
cam_coords_2_all_obj_trans = _ivy.reshape(
cam_coords_2_all_obj_trans, batch_shape + image_dims + [-1, 3])
# Multiplier
# BS x 1 x 1 x num_obj
obj_ids = _ivy.reshape(obj_ids, batch_shape + [1, 1] + [-1])
# BS x H x W x num_obj x 1
multiplier = _ivy.cast(_ivy.expand_dims(obj_ids == id_image, -1),
'float32')
# compute validity mask, for pixels which are on moving objects
# BS x H x W x 1
motion_mask = _ivy.reduce_sum(multiplier, -2) > 0
# make invalid transformations equal to zero
# BS x H x W x num_obj x 3
cam_coords_2_all_obj_trans_w_zeros = cam_coords_2_all_obj_trans * multiplier
# reduce to get only valid transformations
# BS x H x W x 3
cam_coords_2_all_obj_trans = _ivy.reduce_sum(
cam_coords_2_all_obj_trans_w_zeros, -2)
# find cam coords to for zero motion pixels
# BS x H x W x 3
cam_coords_2_wo_motion = _ivy_tvg.cam_to_cam_coords(
cam_coords_1, cam_1_to_2_ext_mat, batch_shape, image_dims)
# BS x H x W x 4
cam_coords_2_all_trans_homo =\
_ivy_mech.make_coordinates_homogeneous(cam_coords_2_all_obj_trans, batch_shape + image_dims)
cam_coords_2 = _ivy.where(motion_mask, cam_coords_2_all_trans_homo,
cam_coords_2_wo_motion)
# return
# BS x H x W x 3, BS x H x W x 1
return cam_coords_2, motion_mask
