def test_gradcheck(self, dev_str, dtype_str, call):
if call is not helpers.torch_call:
# ivy gradcheck method not yet implemented
pytest.skip()
input_ = ivy.variable(ivy.cast(ivy.random_uniform(shape=(2, 3, 4, 4), dev_str=dev_str), 'float64'))
kernel = ivy.variable(ivy.cast(ivy.random_uniform(shape=(3, 3), dev_str=dev_str), 'float64'))
assert gradcheck(top_hat, (input_, kernel), raise_exception=True)
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 create_uniform_pixel_coords_image(image_dims, batch_shape=None, normalized=False, dev_str=None):
"""
Create image of homogeneous integer :math:`xy` pixel co-ordinates :math:`\mathbf{X}\in\mathbb{Z}^{h×w×3}`, stored
as floating point values. The origin is at the top-left corner of the image, with :math:`+x` rightwards, and
:math:`+y` downwards. The final homogeneous dimension are all ones. In subsequent use of this image, the depth of
each pixel can be represented using this same homogeneous representation, by simply scaling each 3-vector by the
depth value. The final dimension therefore always holds the depth value, while the former two dimensions hold depth
scaled pixel co-ordinates.\n
`[reference] <localhost:63342/ivy/docs/source/references/mvg_textbook.pdf#page=172>`_
deduction from top of page 154, section 6.1, equation 6.1
:param image_dims: Image dimensions.
:type image_dims: sequence of ints.
:param batch_shape: Shape of batch. Assumed no batch dimensions if None.
:type batch_shape: sequence of ints, optional
:param normalized: Whether to normalize x-y pixel co-ordinates to the range 0-1.
:type normalized: bool
:param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc.
:type dev_str: str, optional
:return: Image of homogeneous pixel co-ordinates *[batch_shape,height,width,3]*
"""
# shapes as lists
batch_shape = [] if batch_shape is None else batch_shape
batch_shape = list(batch_shape)
image_dims = list(image_dims)
# other shape specs
num_batch_dims = len(batch_shape)
flat_shape = [1] * num_batch_dims + image_dims + [3]
tile_shape = batch_shape + [1] * 3
# H x W x 1
pixel_x_coords = _ivy.cast(_ivy.reshape(_ivy.tile(_ivy.arange(image_dims[1], dev_str=dev_str), [image_dims[0]]),
(image_dims[0], image_dims[1], 1)), 'float32')
if normalized:
pixel_x_coords = pixel_x_coords / (float(image_dims[1]) + MIN_DENOMINATOR)
# W x H x 1
pixel_y_coords_ = _ivy.cast(_ivy.reshape(_ivy.tile(_ivy.arange(image_dims[0], dev_str=dev_str), [image_dims[1]]),
(image_dims[1], image_dims[0], 1)), 'float32')
# H x W x 1
pixel_y_coords = _ivy.transpose(pixel_y_coords_, (1, 0, 2))
if normalized:
pixel_y_coords = pixel_y_coords / (float(image_dims[0]) + MIN_DENOMINATOR)
# H x W x 1
ones = _ivy.ones_like(pixel_x_coords, dev_str=dev_str)
# BS x H x W x 3
return _ivy.tile(_ivy.reshape(_ivy.concatenate((pixel_x_coords, pixel_y_coords, ones), -1),
flat_shape), tile_shape)
def main(interactive=True, try_use_sim=True, f=None):
# config
this_dir = os.path.dirname(os.path.realpath(__file__))
f = choose_random_framework(excluded=['numpy']) if f is None else f
set_framework(f)
sim = Simulator(interactive, try_use_sim)
lr = 0.5
num_anchors = 3
num_sample_points = 100
# spline start
anchor_points = ivy.cast(
ivy.expand_dims(ivy.linspace(0, 1, 2 + num_anchors), -1), 'float32')
query_points = ivy.cast(
ivy.expand_dims(ivy.linspace(0, 1, num_sample_points), -1), 'float32')
# learnable parameters
robot_start_config = ivy.array(ivy.cast(sim.robot_start_config, 'float32'))
robot_target_config = ivy.array(
ivy.cast(sim.robot_target_config, 'float32'))
learnable_anchor_vals = ivy.variable(
ivy.cast(
ivy.transpose(
ivy.linspace(robot_start_config, robot_target_config,
2 + num_anchors)[..., 1:-1], (1, 0)), 'float32'))
# optimizer
optimizer = ivy.SGD(lr=lr)
# optimize
it = 0
colliding = True
clearance = 0
joint_query_vals = None
while colliding:
total_cost, grads, joint_query_vals, link_positions, sdf_vals = ivy.execute_with_gradients(
lambda xs: compute_cost_and_sdfs(xs[
'w'], anchor_points, robot_start_config, robot_target_config,
query_points, sim),
Container({'w': learnable_anchor_vals}))
colliding = ivy.reduce_min(sdf_vals[2:]) < clearance
sim.update_path_visualization(
link_positions, sdf_vals,
os.path.join(this_dir, 'msp_no_sim', 'path_{}.png'.format(it)))
learnable_anchor_vals = optimizer.step(
Container({'w': learnable_anchor_vals}), grads)['w']
it += 1
sim.execute_motion(joint_query_vals)
sim.close()
unset_framework()
def test_lstm(b_t_ic_hc_otf_sctv, dtype_str, tensor_fn, dev_str, call):
# smoke test
b, t, input_channels, hidden_channels, output_true_flat, state_c_true_val = b_t_ic_hc_otf_sctv
x = ivy.cast(ivy.linspace(ivy.zeros([b, t]), ivy.ones([b, t]), input_channels), 'float32')
init_h = ivy.ones([b, hidden_channels])
init_c = ivy.ones([b, hidden_channels])
kernel = ivy.variable(ivy.ones([input_channels, 4*hidden_channels]))*0.5
recurrent_kernel = ivy.variable(ivy.ones([hidden_channels, 4*hidden_channels]))*0.5
output, state_c = ivy.lstm_update(x, init_h, init_c, kernel, recurrent_kernel)
# type test
assert ivy.is_array(output)
assert ivy.is_array(state_c)
# cardinality test
assert output.shape == (b, t, hidden_channels)
assert state_c.shape == (b, hidden_channels)
# value test
output_true = np.tile(np.asarray(output_true_flat).reshape((b, t, 1)), (1, 1, hidden_channels))
state_c_true = np.ones([b, hidden_channels]) * state_c_true_val
output, state_c = call(ivy.lstm_update, x, init_h, init_c, kernel, recurrent_kernel)
assert np.allclose(output, output_true, atol=1e-6)
assert np.allclose(state_c, state_c_true, atol=1e-6)
# compilation test
if call in [helpers.torch_call]:
# this is not a backend implemented function
pytest.skip()
helpers.assert_compilable(ivy.lstm_update)
def test_linear_layer(bs_ic_oc_target, with_v, dtype_str, tensor_fn, dev_str,
call):
# smoke test
batch_shape, input_channels, output_channels, target = bs_ic_oc_target
x = ivy.cast(
ivy.linspace(ivy.zeros(batch_shape), ivy.ones(batch_shape),
input_channels), 'float32')
if with_v:
np.random.seed(0)
wlim = (6 / (output_channels + input_channels))**0.5
w = ivy.variable(
ivy.array(
np.random.uniform(-wlim, wlim,
(output_channels, input_channels)),
'float32'))
b = ivy.variable(ivy.zeros([output_channels]))
v = Container({'w': w, 'b': b})
else:
v = None
linear_layer = ivy.Linear(input_channels, output_channels, v=v)
ret = linear_layer(x)
# type test
assert ivy.is_array(ret)
# cardinality test
assert ret.shape == tuple(batch_shape + [output_channels])
# value test
if not with_v:
return
assert np.allclose(call(linear_layer, x), np.array(target))
# compilation test
if call is helpers.torch_call:
# pytest scripting does not **kwargs
return
helpers.assert_compilable(linear_layer)
def test_lstm_layer(b_t_ic_hc_otf_sctv, with_v, with_initial_state, dtype_str,
tensor_fn, dev_str, call):
# smoke test
b, t, input_channels, hidden_channels, output_true_flat, state_c_true_val = b_t_ic_hc_otf_sctv
x = ivy.cast(
ivy.linspace(ivy.zeros([b, t]), ivy.ones([b, t]), input_channels),
'float32')
if with_initial_state:
init_h = ivy.ones([b, hidden_channels])
init_c = ivy.ones([b, hidden_channels])
initial_state = ([init_h], [init_c])
else:
initial_state = None
if with_v:
kernel = ivy.variable(
ivy.ones([input_channels, 4 * hidden_channels]) * 0.5)
recurrent_kernel = ivy.variable(
ivy.ones([hidden_channels, 4 * hidden_channels]) * 0.5)
v = Container({
'input': {
'layer_0': {
'w': kernel
}
},
'recurrent': {
'layer_0': {
'w': recurrent_kernel
}
}
})
else:
v = None
lstm_layer = ivy.LSTM(input_channels, hidden_channels, v=v)
output, (state_h, state_c) = lstm_layer(x, initial_state=initial_state)
# type test
assert ivy.is_array(output)
assert ivy.is_array(state_h[0])
assert ivy.is_array(state_c[0])
# cardinality test
assert output.shape == (b, t, hidden_channels)
assert state_h[0].shape == (b, hidden_channels)
assert state_c[0].shape == (b, hidden_channels)
# value test
if not with_v or not with_initial_state:
return
output_true = np.tile(
np.asarray(output_true_flat).reshape((b, t, 1)),
(1, 1, hidden_channels))
state_c_true = np.ones([b, hidden_channels]) * state_c_true_val
output, (state_h, state_c) = call(lstm_layer,
x,
initial_state=initial_state)
assert np.allclose(output, output_true, atol=1e-6)
assert np.allclose(state_c, state_c_true, atol=1e-6)
# compilation test
if call in [helpers.torch_call]:
# this is not a backend implemented function
pytest.skip()
helpers.assert_compilable(ivy.lstm_update)
def test_sgd_optimizer(bs_ic_oc_target, with_v, dtype_str, tensor_fn, dev_str,
call):
# smoke test
if call is helpers.np_call:
# NumPy does not support gradients
pytest.skip()
batch_shape, input_channels, output_channels, target = bs_ic_oc_target
x = ivy.cast(
ivy.linspace(ivy.zeros(batch_shape), ivy.ones(batch_shape),
input_channels), 'float32')
if with_v:
np.random.seed(0)
wlim = (6 / (output_channels + input_channels))**0.5
w = ivy.variable(
ivy.array(
np.random.uniform(-wlim, wlim,
(output_channels, input_channels)),
'float32'))
b = ivy.variable(ivy.zeros([output_channels]))
v = Container({'w': w, 'b': b})
else:
v = None
linear_layer = ivy.Linear(input_channels, output_channels, v=v)
def loss_fn(v_):
out = linear_layer(x, v=v_)
return ivy.reduce_mean(out)[0]
# optimizer
optimizer = ivy.SGD()
# train
loss_tm1 = 1e12
loss = None
grads = None
for i in range(10):
loss, grads = ivy.execute_with_gradients(loss_fn, linear_layer.v)
linear_layer.v = optimizer.step(linear_layer.v, grads)
assert loss < loss_tm1
loss_tm1 = loss
# type test
assert ivy.is_array(loss)
assert isinstance(grads, ivy.Container)
# cardinality test
if call is helpers.mx_call:
# mxnet slicing cannot reduce dimension to zero
assert loss.shape == (1, )
else:
assert loss.shape == ()
# value test
assert ivy.reduce_max(ivy.abs(grads.b)) > 0
assert ivy.reduce_max(ivy.abs(grads.w)) > 0
# compilation test
if call is helpers.torch_call:
# pytest scripting does not **kwargs
return
helpers.assert_compilable(loss_fn)
def test_module_w_none_attribute(bs_ic_oc, dev_str, call):
# smoke test
if call is helpers.np_call:
# NumPy does not support gradients
pytest.skip()
batch_shape, input_channels, output_channels = bs_ic_oc
x = ivy.cast(
ivy.linspace(ivy.zeros(batch_shape), ivy.ones(batch_shape),
input_channels), 'float32')
module = ModuleWithNoneAttribute()
def _se_to_mask(se: ivy.Array) -> ivy.Array:
se_h, se_w = se.shape
se_flat = ivy.reshape(se, (-1,))
num_feats = se_h * se_w
i_s = ivy.expand_dims(ivy.arange(num_feats, dev_str=ivy.dev_str(se)), -1)
y_s = i_s % se_h
x_s = i_s // se_h
indices = ivy.concatenate((i_s, ivy.zeros_like(i_s, dtype_str='int32'), x_s, y_s), -1)
out = ivy.scatter_nd(
indices, ivy.cast(se_flat >= 0, ivy.dtype_str(se)), (num_feats, 1, se_h, se_w), dev_str=ivy.dev_str(se))
return out
def test_jit(self, dev_str, dtype_str, call):
op = top_hat
if call in [helpers.jnp_call, helpers.torch_call]:
# compiled jax tensors do not have device_buffer attribute, preventing device info retrieval,
# pytorch scripting does not support .type() casting, nor Union or Numbers for type hinting
pytest.skip()
op_compiled = ivy.compile_fn(op)
input_ = ivy.cast(ivy.random_uniform(shape=(1, 2, 7, 7), dev_str=dev_str), dtype_str)
kernel = ivy.ones((3, 3), dev_str=dev_str, dtype_str=dtype_str)
actual = call(op_compiled, input_, kernel)
expected = call(op, input_, kernel)
assert np.allclose(actual, expected)
def test_module_training(bs_ic_oc, dev_str, call):
# smoke test
if call is helpers.np_call:
# NumPy does not support gradients
pytest.skip()
batch_shape, input_channels, output_channels = bs_ic_oc
x = ivy.cast(
ivy.linspace(ivy.zeros(batch_shape), ivy.ones(batch_shape),
input_channels), 'float32')
module = TrainableModule(input_channels, output_channels)
def loss_fn(v_):
out = module(x, v=v_)
return ivy.reduce_mean(out)[0]
# train
loss_tm1 = 1e12
loss = None
grads = None
for i in range(10):
loss, grads = ivy.execute_with_gradients(loss_fn, module.v)
module.v = ivy.gradient_descent_update(module.v, grads, 1e-3)
assert loss < loss_tm1
loss_tm1 = loss
# type test
assert ivy.is_array(loss)
assert isinstance(grads, ivy.Container)
# cardinality test
if call is helpers.mx_call:
# mxnet slicing cannot reduce dimension to zero
assert loss.shape == (1, )
else:
assert loss.shape == ()
# value test
assert ivy.reduce_max(ivy.abs(grads.linear0.b)) > 0
assert ivy.reduce_max(ivy.abs(grads.linear0.w)) > 0
assert ivy.reduce_max(ivy.abs(grads.linear1.b)) > 0
assert ivy.reduce_max(ivy.abs(grads.linear1.w)) > 0
assert ivy.reduce_max(ivy.abs(grads.linear2.b)) > 0
assert ivy.reduce_max(ivy.abs(grads.linear2.w)) > 0
# compilation test
if call is helpers.torch_call:
# pytest scripting does not support **kwargs
return
helpers.assert_compilable(loss_fn)
def _update_path_visualization_pyrep(self, multi_spline_points, multi_spline_sdf_vals):
if len(self._spline_paths) > 0:
for spline_path_segs in self._spline_paths:
for spline_path_seg in spline_path_segs:
spline_path_seg.remove()
self._spline_paths.clear()
for spline_points, sdf_vals in zip(multi_spline_points, multi_spline_sdf_vals):
sdf_flags_0 = sdf_vals[1:, 0] > 0
sdf_flags_1 = sdf_vals[:-1, 0] > 0
sdf_borders = sdf_flags_1 != sdf_flags_0
borders_indices = ivy.indices_where(sdf_borders)
if borders_indices.shape[0] != 0:
to_concat = (ivy.array([0], 'int32'), ivy.cast(borders_indices, 'int32')[:, 0],
ivy.array([-1], 'int32'))
else:
to_concat = (ivy.array([0], 'int32'), ivy.array([-1], 'int32'))
border_indices = ivy.concatenate(to_concat, 0)
num_groups = border_indices.shape[0] - 1
spline_path = list()
for i in range(num_groups):
border_idx_i = int(ivy.to_numpy(border_indices[i]).item())
border_idx_ip1 = int(ivy.to_numpy(border_indices[i + 1]).item())
if i < num_groups - 1:
control_group = spline_points[border_idx_i:border_idx_ip1]
sdf_group = sdf_vals[border_idx_i:border_idx_ip1]
else:
control_group = spline_points[border_idx_i:]
sdf_group = sdf_vals[border_idx_i:]
num_points = control_group.shape[0]
orientation_zeros = np.zeros((num_points, 3))
color = (0.2, 0.8, 0.2) if sdf_group[-1] > 0 else (0.8, 0.2, 0.2)
control_poses = np.concatenate((ivy.to_numpy(control_group), orientation_zeros), -1)
spline_path_seg = CartesianPath.create(show_orientation=False, show_position=False, line_size=8,
path_color=color)
spline_path_seg.insert_control_points(control_poses.tolist())
spline_path.append(spline_path_seg)
self._spline_paths.append(spline_path)
def persp_angles_to_focal_lengths(persp_angles, image_dims, dev_str=None):
"""
Compute focal lengths :math:`f_x, f_y` from perspective angles :math:`θ_x, θ_y`.\n
`[reference] <localhost:63342/ivy/docs/source/references/mvg_textbook.pdf#page=172>`_
deduction from page 154, section 6.1, figure 6.1
:param persp_angles: Perspective angles *[batch_shape,2]*
:type persp_angles: array
:param image_dims: Image dimensions.
:type image_dims: sequence of ints
: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: Focal lengths *[batch_shape,2]*
"""
if dev_str is None:
dev_str = _ivy.dev_str(persp_angles)
# shapes as list
image_dims = list(image_dims)
# BS x 2
return -_ivy.flip(_ivy.cast(_ivy.array(image_dims, dev_str=dev_str), 'float32'), -1) /\
(2 * _ivy.tan(persp_angles / 2) + MIN_DENOMINATOR)
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 test_lstm_layer_training(b_t_ic_hc_otf_sctv, with_v, dtype_str, tensor_fn,
dev_str, call):
# smoke test
if call is helpers.np_call:
# NumPy does not support gradients
pytest.skip()
# smoke test
b, t, input_channels, hidden_channels, output_true_flat, state_c_true_val = b_t_ic_hc_otf_sctv
x = ivy.cast(
ivy.linspace(ivy.zeros([b, t]), ivy.ones([b, t]), input_channels),
'float32')
if with_v:
kernel = ivy.variable(
ivy.ones([input_channels, 4 * hidden_channels]) * 0.5)
recurrent_kernel = ivy.variable(
ivy.ones([hidden_channels, 4 * hidden_channels]) * 0.5)
v = Container({
'input': {
'layer_0': {
'w': kernel
}
},
'recurrent': {
'layer_0': {
'w': recurrent_kernel
}
}
})
else:
v = None
lstm_layer = ivy.LSTM(input_channels, hidden_channels, v=v)
def loss_fn(v_):
out, (state_h, state_c) = lstm_layer(x, v=v_)
return ivy.reduce_mean(out)[0]
# train
loss_tm1 = 1e12
loss = None
grads = None
for i in range(10):
loss, grads = ivy.execute_with_gradients(loss_fn, lstm_layer.v)
lstm_layer.v = ivy.gradient_descent_update(lstm_layer.v, grads, 1e-3)
assert loss < loss_tm1
loss_tm1 = loss
# type test
assert ivy.is_array(loss)
assert isinstance(grads, ivy.Container)
# cardinality test
if call is helpers.mx_call:
# mxnet slicing cannot reduce dimension to zero
assert loss.shape == (1, )
else:
assert loss.shape == ()
# value test
for key, val in grads.to_iterator():
assert ivy.reduce_max(ivy.abs(val)) > 0
# compilation test
if call is helpers.torch_call:
# pytest scripting does not **kwargs
return
helpers.assert_compilable(loss_fn)
def rasterize_triangles(pixel_coords_triangles, image_dims, batch_shape=None, dev_str=None):
"""
Rasterize image-projected triangles
based on: https://www.scratchapixel.com/lessons/3d-basic-rendering/rasterization-practical-implementation/rasterization-stage
and: https://www.scratchapixel.com/lessons/3d-basic-rendering/rasterization-practical-implementation/rasterization-practical-implementation
:param pixel_coords_triangles: Projected image-space triangles to be rasterized
*[batch_shape,input_size,3,3]*
:type pixel_coords_triangles: array
:param image_dims: Image dimensions.
:type image_dims: sequence of ints
: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: Rasterized triangles
"""
if batch_shape is None:
batch_shape = []
if dev_str is None:
dev_str = _ivy.dev_str(pixel_coords_triangles)
# shapes as list
batch_shape = list(batch_shape)
num_batch_dims = len(batch_shape)
image_dims = list(image_dims)
input_image_dims = pixel_coords_triangles.shape[num_batch_dims:-2]
input_image_dims_prod = _reduce(_mul, input_image_dims, 1)
# BS x 3 x 2
pixel_xy_coords = pixel_coords_triangles[..., 0:2]
# BS x 3 x 1
pixel_x_coords = pixel_coords_triangles[..., 0:1]
pixel_y_coords = pixel_coords_triangles[..., 1:2]
# 1
x_min = _ivy.reshape(_ivy.reduce_min(pixel_x_coords, keepdims=True), (-1,))
x_max = _ivy.reshape(_ivy.reduce_max(pixel_x_coords, keepdims=True), (-1,))
x_range = x_max - x_min
y_min = _ivy.reshape(_ivy.reduce_min(pixel_y_coords, keepdims=True), (-1,))
y_max = _ivy.reshape(_ivy.reduce_max(pixel_y_coords, keepdims=True), (-1,))
y_range = y_max - y_min
# 2
bbox = _ivy.concatenate((x_range, y_range), 0)
img_bbox_list = [int(item) for item in _ivy.to_list(_ivy.concatenate((y_range + 1, x_range + 1), 0))]
# BS x 2
v0 = pixel_xy_coords[..., 0, :]
v1 = pixel_xy_coords[..., 1, :]
v2 = pixel_xy_coords[..., 2, :]
tri_centres = (v0 + v1 + v2) / 3
# BS x 1
v0x = v0[..., 0:1]
v0y = v0[..., 1:2]
v1x = v1[..., 0:1]
v1y = v1[..., 1:2]
v2x = v2[..., 0:1]
v2y = v2[..., 1:2]
# BS x BBX x BBY x 2
uniform_sample_coords = _ivy_svg.create_uniform_pixel_coords_image(img_bbox_list, batch_shape)[..., 0:2]
P = _ivy.round(uniform_sample_coords + tri_centres - bbox / 2)
# BS x BBX x BBY x 1
Px = P[..., 0:1]
Py = P[..., 1:2]
v0v1_edge_func = ((Px - v0x) * (v1y - v0y) - (Py - v0y) * (v1x - v0x)) >= 0
v1v2_edge_func = ((Px - v1x) * (v2y - v1y) - (Py - v1y) * (v2x - v1x)) >= 0
v2v0_edge_func = ((Px - v2x) * (v0y - v2y) - (Py - v2y) * (v0x - v2x)) >= 0
edge_func = _ivy.logical_and(_ivy.logical_and(v0v1_edge_func, v1v2_edge_func), v2v0_edge_func)
batch_indices_list = list()
for i, batch_dim in enumerate(batch_shape):
# get batch shape
batch_dims_before = batch_shape[:i]
num_batch_dims_before = len(batch_dims_before)
batch_dims_after = batch_shape[i + 1:]
num_batch_dims_after = len(batch_dims_after)
# [batch_dim]
batch_indices = _ivy.arange(batch_dim, dtype_str='int32', dev_str=dev_str)
# [1]*num_batch_dims_before x batch_dim x [1]*num_batch_dims_after x 1 x 1
reshaped_batch_indices = _ivy.reshape(batch_indices, [1] * num_batch_dims_before + [batch_dim] +
[1] * num_batch_dims_after + [1, 1])
# BS x N x 1
tiled_batch_indices = _ivy.tile(reshaped_batch_indices, batch_dims_before + [1] + batch_dims_after +
[input_image_dims_prod * 9, 1])
batch_indices_list.append(tiled_batch_indices)
# BS x N x (num_batch_dims + 2)
all_indices = _ivy.concatenate(
batch_indices_list + [_ivy.cast(_ivy.flip(_ivy.reshape(P, batch_shape + [-1, 2]), -1),
'int32')], -1)
# offset uniform images
return _ivy.cast(_ivy.flip(_ivy.scatter_nd(_ivy.reshape(all_indices, [-1, num_batch_dims + 2]),
_ivy.reshape(_ivy.cast(edge_func, 'int32'), (-1, 1)),
batch_shape + image_dims + [1],
reduction='replace' if _ivy.backend == 'mxnd' else 'sum'), -3), 'bool')
def __init__(self, interactive, try_use_sim):
super().__init__(interactive, try_use_sim)
# ivy robot
rel_body_points = ivy.array([[0., 0., 0.], [-0.15, 0., -0.15],
[-0.15, 0., 0.15], [0.15, 0., -0.15],
[0.15, 0., 0.15]])
self.ivy_drone = RigidMobile(rel_body_points)
# initialize scene
if self.with_pyrep:
self._spherical_vision_sensor.remove()
for i in range(6):
self._vision_sensors[i].remove()
self._vision_sensor_bodies[i].remove()
[ray.remove() for ray in self._vision_sensor_rays[i]]
drone_start_pos = np.array([-1.15, -1.028, 0.6])
target_pos = np.array([1.025, 1.125, 0.6])
self._drone.set_position(drone_start_pos)
self._drone.set_orientation(
np.array([-90, 50, -180]) * np.pi / 180)
self._target.set_position(target_pos)
self._target.set_orientation(
np.array([-90, 50, -180]) * np.pi / 180)
self._default_camera.set_position(
np.array([-3.2835, -0.88753, 1.3773]))
self._default_camera.set_orientation(
np.array([-151.07, 70.079, -120.45]) * np.pi / 180)
input(
'\nScene initialized.\n\n'
'The simulator visualizer can be translated and rotated by clicking either the left mouse button or the wheel, '
'and then dragging the mouse.\n'
'Scrolling the mouse wheel zooms the view in and out.\n\n'
'You can click on any object either in the scene or the left hand panel, '
'then select the box icon with four arrows in the top panel of the simulator, '
'and then drag the object around dynamically.\n'
'Starting to drag and then holding ctrl allows you to also drag the object up and down.\n'
'Clicking the top icon with a box and two rotating arrows similarly allows rotation of the object.\n\n'
'Once you have aranged the scene as desired, press enter in the terminal to continue with the demo...\n'
)
# primitive scene
self.setup_primitive_scene()
# public objects
drone_starting_inv_ext_mat = ivy.array(
np.reshape(self._drone.get_matrix(), (3, 4)), 'float32')
drone_start_rot_vec_pose = ivy_mech.mat_pose_to_rot_vec_pose(
drone_starting_inv_ext_mat)
self.drone_start_pose = drone_start_rot_vec_pose
target_inv_ext_mat = ivy.array(
np.reshape(self._target.get_matrix(), (3, 4)), 'float32')
target_rot_vec_pose = ivy_mech.mat_pose_to_rot_vec_pose(
target_inv_ext_mat)
self.drone_target_pose = target_rot_vec_pose
# spline path
drone_start_to_target_poses = ivy.transpose(
ivy.linspace(self.drone_start_pose, self.drone_target_pose,
100), (1, 0))
drone_start_to_target_inv_ext_mats = ivy_mech.rot_vec_pose_to_mat_pose(
drone_start_to_target_poses)
drone_start_to_target_positions =\
ivy.transpose(self.ivy_drone.sample_body(drone_start_to_target_inv_ext_mats), (1, 0, 2))
initil_sdf_vals = ivy.reshape(
self.sdf(
ivy.reshape(
ivy.cast(drone_start_to_target_positions, 'float32'),
(-1, 3))), (-1, 100, 1))
self.update_path_visualization(drone_start_to_target_positions,
initil_sdf_vals, None)
# wait for user input
self._user_prompt(
'\nInitialized scene with a drone and a target position to reach.'
'\nPress enter in the terminal to use method ivy_robot.interpolate_spline_points '
'to plan a spline path which reaches the target whilst avoiding the objects in the scene...\n'
)
else:
# primitive scene
self.setup_primitive_scene_no_sim()
# public objects
self.drone_start_pose = ivy.array(
[-1.1500, -1.0280, 0.6000, 0.7937, 1.7021, 1.7021])
self.drone_target_pose = ivy.array(
[1.0250, 1.1250, 0.6000, 0.7937, 1.7021, 1.7021])
# message
print(
'\nInitialized dummy scene with a drone and a target position to reach.'
'\nClose the visualization window to use method ivy_robot.interpolate_spline_points '
'to plan a spline path which reaches the target whilst avoiding the objects in the scene...\n'
)
# plot scene before rotation
if interactive:
plt.imshow(
mpimg.imread(
os.path.join(
os.path.dirname(os.path.realpath(__file__)),
'dsp_no_sim', 'path_0.png')))
plt.show()
print('\nOptimizing spline path...\n')
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 _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 sample_links(self, joint_angles, link_num=None, samples_per_metre=25, batch_shape=None):
"""
Sample links of the robot at uniformly distributed cartesian positions.
:param joint_angles: Joint angles of the robot *[batch_shape,num_joints]*
:type joint_angles: array
:param link_num: Link number for which to compute matrices up to. Default is the last link.
:type link_num: int, optional
:param samples_per_metre: Number of samples per metre of robot link
:type samples_per_metre: int
:param batch_shape: Shape of batch. Inferred from inputs if None.
:type batch_shape: sequence of ints, optional
:return: The sampled link cartesian positions *[batch_shape,total_sampling_chain_length,3]*
"""
if link_num is None:
link_num = self._num_joints
if batch_shape is None:
batch_shape = joint_angles.shape[:-1]
batch_shape = list(batch_shape)
num_batch_dims = len(batch_shape)
batch_dims_for_trans = list(range(num_batch_dims))
# BS x NJ x 4 x 4
link_matrices = self.compute_link_matrices(joint_angles, link_num, batch_shape)
# BS x LN+1 x 3
link_positions = link_matrices[..., 0:3, -1]
# BS x LN x 3
segment_starts = link_positions[..., :-1, :]
segment_ends = link_positions[..., 1:, :]
# LN
segment_sizes = _ivy.cast(_ivy.ceil(
self._link_lengths[0:link_num] * samples_per_metre), 'int32')
# list of segments
segments_list = list()
for link_idx in range(link_num):
segment_size = segment_sizes[link_idx]
# BS x 1 x 3
segment_start = segment_starts[..., link_idx:link_idx + 1, :]
segment_end = segment_ends[..., link_idx:link_idx + 1, :]
# BS x segment_size x 3
segment = _ivy.linspace(segment_start, segment_end, segment_size, axis=-2)[..., 0, :, :]
if link_idx == link_num - 1 or segment_size == 1:
segments_list.append(segment)
else:
segments_list.append(segment[..., :-1, :])
# BS x total_robot_chain_length x 3
all_segments = _ivy.concatenate(segments_list, -2)
# BS x total_robot_chain_length x 4
all_segments_homo = _ivy_mech.make_coordinates_homogeneous(all_segments)
# 4 x BSxtotal_robot_chain_length
all_segments_homo_trans = _ivy.reshape(_ivy.transpose(
all_segments_homo, [num_batch_dims + 1] + batch_dims_for_trans + [num_batch_dims]), (4, -1))
# 3 x BSxtotal_robot_chain_length
transformed_trans = _ivy.matmul(self._base_inv_ext_mat[..., 0:3, :], all_segments_homo_trans)
# BS x total_robot_chain_length x 3
return _ivy.transpose(_ivy.reshape(
transformed_trans, [3] + batch_shape + [-1]),
[i+1 for i in batch_dims_for_trans] + [num_batch_dims+1] + [0])
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 main(batch_size=32,
num_train_steps=31250,
compile_flag=True,
num_bits=8,
seq_len=28,
ctrl_output_size=100,
memory_size=128,
memory_vector_dim=28,
overfit_flag=False,
interactive=True,
f=None):
f = choose_random_framework() if f is None else f
set_framework(f)
# train config
lr = 1e-3 if not overfit_flag else 1e-2
batch_size = batch_size if not overfit_flag else 1
num_train_steps = num_train_steps if not overfit_flag else 150
max_grad_norm = 50
# logging config
vis_freq = 250 if not overfit_flag else 1
# optimizer
optimizer = ivy.Adam(lr=lr)
# ntm
ntm = NTM(input_dim=num_bits + 1,
output_dim=num_bits,
ctrl_output_size=ctrl_output_size,
ctrl_layers=1,
memory_size=memory_size,
memory_vector_dim=memory_vector_dim,
read_head_num=1,
write_head_num=1)
# compile loss fn
total_seq_example = ivy.random_uniform(shape=(batch_size, 2 * seq_len + 1,
num_bits + 1))
target_seq_example = total_seq_example[:, 0:seq_len, :-1]
if compile_flag:
loss_fn_maybe_compiled = ivy.compile_fn(
lambda v, ttl_sq, trgt_sq, sq_ln: loss_fn(ntm, v, ttl_sq, trgt_sq,
sq_ln),
dynamic=False,
example_inputs=[
ntm.v, total_seq_example, target_seq_example, seq_len
])
else:
loss_fn_maybe_compiled = lambda v, ttl_sq, trgt_sq, sq_ln: loss_fn(
ntm, v, ttl_sq, trgt_sq, sq_ln)
# init
input_seq_m1 = ivy.cast(
ivy.random_uniform(0., 1., (batch_size, seq_len, num_bits)) > 0.5,
'float32')
mw = None
vw = None
for i in range(num_train_steps):
# sequence to copy
if not overfit_flag:
input_seq_m1 = ivy.cast(
ivy.random_uniform(0., 1.,
(batch_size, seq_len, num_bits)) > 0.5,
'float32')
target_seq = input_seq_m1
input_seq = ivy.concatenate(
(input_seq_m1, ivy.zeros((batch_size, seq_len, 1))), -1)
eos = ivy.ones((batch_size, 1, num_bits + 1))
output_seq = ivy.zeros_like(input_seq)
total_seq = ivy.concatenate((input_seq, eos, output_seq), -2)
# train step
loss, pred_vals = train_step(loss_fn_maybe_compiled, optimizer, ntm,
total_seq, target_seq, seq_len, mw, vw,
ivy.array(i + 1,
'float32'), max_grad_norm)
# log
print('step: {}, loss: {}'.format(i, ivy.to_numpy(loss).item()))
# visualize
if i % vis_freq == 0:
target_to_vis = (ivy.to_numpy(target_seq[0] * 255)).astype(
np.uint8)
target_to_vis = np.transpose(
cv2.resize(target_to_vis, (560, 160),
interpolation=cv2.INTER_NEAREST), (1, 0))
pred_to_vis = (ivy.to_numpy(pred_vals[0] * 255)).astype(np.uint8)
pred_to_vis = np.transpose(
cv2.resize(pred_to_vis, (560, 160),
interpolation=cv2.INTER_NEAREST), (1, 0))
img_to_vis = np.concatenate((pred_to_vis, target_to_vis), 0)
img_to_vis = cv2.resize(img_to_vis, (1120, 640),
interpolation=cv2.INTER_NEAREST)
img_to_vis[0:60, -200:] = 0
img_to_vis[5:55, -195:-5] = 255
cv2.putText(img_to_vis, 'step {}'.format(i), (935, 42),
cv2.FONT_HERSHEY_SIMPLEX, 1.2, tuple([0] * 3), 2)
img_to_vis[0:60, 0:200] = 0
img_to_vis[5:55, 5:195] = 255
cv2.putText(img_to_vis, 'prediction', (7, 42),
cv2.FONT_HERSHEY_SIMPLEX, 1.2, tuple([0] * 3), 2)
img_to_vis[320:380, 0:130] = 0
img_to_vis[325:375, 5:125] = 255
cv2.putText(img_to_vis, 'target', (7, 362),
cv2.FONT_HERSHEY_SIMPLEX, 1.2, tuple([0] * 3), 2)
if interactive:
cv2.imshow('prediction_and_target', img_to_vis)
if overfit_flag:
cv2.waitKey(1)
else:
cv2.waitKey(100)
cv2.destroyAllWindows()
def test_smoke(self, dev_str, dtype_str, call):
kernel = ivy.cast(ivy.random_uniform(shape=(3, 3), dev_str=dev_str),
dtype_str)
assert call(_se_to_mask, kernel) is not None
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
