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