Esempio n. 1
0
def _fit_spline(train_points, train_values, order):

    # shapes
    train_points_shape = train_points.shape
    batch_shape = list(train_points_shape[:-2])
    num_batch_dims = len(batch_shape)
    n = train_points_shape[-2]
    pd = train_values.shape[-1]

    # BS x N x 1
    c = train_points

    # BS x N x PD
    f_ = train_values

    # BS x N x N
    matrix_a = _phi(_pairwise_distance(c, c), order)

    # BS x N x 1
    ones = _ivy.ones_like(c[..., :1])

    # BS x N x 2
    matrix_b = _ivy.concatenate([c, ones], -1)

    # BS x 2 x N
    matrix_b_trans = _ivy.transpose(
        matrix_b,
        list(range(num_batch_dims)) + [num_batch_dims + 1, num_batch_dims])

    # BS x N+2 x N
    left_block = _ivy.concatenate([matrix_a, matrix_b_trans], -2)

    # BS x 2 x 2
    lhs_zeros = _ivy.zeros(batch_shape + [2, 2])

    # BS x N+2 x 2
    right_block = _ivy.concatenate([matrix_b, lhs_zeros], -2)

    # BS x N+2 x N+2
    lhs = _ivy.concatenate([left_block, right_block], -1)

    # BS x 2 x PD
    rhs_zeros = _ivy.zeros(batch_shape + [2, pd])

    # BS x N+2 x PD
    rhs = _ivy.concatenate([f_, rhs_zeros], -2)

    # BS x N+2 x PD
    w_v = _ivy.matmul(_ivy.pinv(lhs), rhs)

    # BS x N x PD
    w = w_v[..., :n, :]

    # BS x 2 x PD
    v = w_v[..., n:, :]

    # BS x N x PD,    BS x 2 x PD
    return w, v
Esempio n. 2
0
 def _update_seq_info_for_window(self, seq_info):
     if not ivy.exists(seq_info):
         return
     seq_idx = int(seq_info.seq_idx[0])
     seq_len = int(seq_info.length[0])
     new_len = self._compute_seq_len(seq_idx, seq_len,
                                     self._spec.containers_to_skip)
     seq_info = seq_info.copy()
     seq_info.length = ivy.ones_like(seq_info.length) * new_len
     return seq_info
Esempio n. 3
0
    def __init__(self,
                 img_mean: ivy.Array,
                 cam_rel_mat: ivy.Array,
                 img_var: ivy.Array = None,
                 validity_mask: ivy.Array = None,
                 pose_mean: ivy.Array = None,
                 pose_cov: ivy.Array = None):
        """
        Create esm image measurement container

        :param img_mean: Camera-relative co-ordinates and image features
                            *[batch_size, timesteps, height, width, 3 + feat]*
        :type: img_mean: array
        :param cam_rel_mat: The pose of the camera relative to the current agent pose, in matrix form
                            *[batch_size, timesteps, 3, 4]*
        :type cam_rel_mat: array
        :param img_var: Image depth and feature variance values, assumed all zero if None.
                        *[batch_size, timesteps, height, width, 1 + feat]*
        :type: img_var: array, optional
        :param validity_mask: Validity mask, for which pixels should be considered. Assumed all valid if None
                                *[batch_size, timesteps, height, width, 1]*
        :type validity_mask: array, optional
        :param pose_mean: The pose of the camera relative to the current agent pose, in rotation vector pose form.
                            Inferred from cam_rel_mat if None. *[batch_size, timesteps, 6]*
        :type pose_mean: array, optional
        :param pose_cov: The convariance of the camera relative pose, in rotation vector form. Assumed all zero if None.
                            *[batch_size, timesteps, 6, 6]*
        :type pose_cov: array, optional
        """
        img_mean = _pad_to_batch_n_time_dims(img_mean, 5)
        cam_rel_mat = _pad_to_batch_n_time_dims(cam_rel_mat, 4)
        self['img_mean'] = img_mean
        self['cam_rel_mat'] = cam_rel_mat
        if img_var is None:
            img_var = ivy.zeros_like(img_mean)
        else:
            img_var = _pad_to_batch_n_time_dims(img_var, 5)
        self['img_var'] = img_var
        if validity_mask is None:
            validity_mask = ivy.ones_like(img_mean[..., 0:1])
        else:
            validity_mask = _pad_to_batch_n_time_dims(validity_mask, 5)
        self['validity_mask'] = validity_mask
        if pose_mean is None:
            pose_mean = ivy_mech.mat_pose_to_rot_vec_pose(cam_rel_mat)
        else:
            pose_mean = _pad_to_batch_n_time_dims(pose_mean, 3)
        self['pose_mean'] = pose_mean
        if pose_cov is None:
            pose_cov = ivy.tile(ivy.expand_dims(ivy.zeros_like(pose_mean), -1), (1, 1, 1, 6))
        else:
            pose_cov = _pad_to_batch_n_time_dims(pose_cov, 4)
        self['pose_cov'] = pose_cov
Esempio n. 4
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)
Esempio n. 5
0
def _get_dummy_obs(batch_size, num_frames, num_cams, image_dims, num_feature_channels, dev_str='cpu', ones=False,
                   empty=False):

    uniform_pixel_coords =\
        ivy_vision.create_uniform_pixel_coords_image(image_dims, [batch_size, num_frames], dev_str=dev_str)

    img_meas = dict()
    for i in range(num_cams):
        validity_mask = ivy.ones([batch_size, num_frames] + image_dims + [1], dev_str=dev_str)
        if ones:
            img_mean = ivy.concatenate((uniform_pixel_coords[..., 0:2], ivy.ones(
                [batch_size, num_frames] + image_dims + [1 + num_feature_channels], dev_str=dev_str)), -1)
            img_var = ivy.ones(
                     [batch_size, num_frames] + image_dims + [3 + num_feature_channels], dev_str=dev_str)*1e-3
            pose_mean = ivy.zeros([batch_size, num_frames, 6], dev_str=dev_str)
            pose_cov = ivy.ones([batch_size, num_frames, 6, 6], dev_str=dev_str)*1e-3
        else:
            img_mean = ivy.concatenate((uniform_pixel_coords[..., 0:2], ivy.random_uniform(
                1e-3, 1, [batch_size, num_frames] + image_dims + [1 + num_feature_channels], dev_str=dev_str)), -1)
            img_var = ivy.random_uniform(
                     1e-3, 1, [batch_size, num_frames] + image_dims + [3 + num_feature_channels], dev_str=dev_str)
            pose_mean = ivy.random_uniform(1e-3, 1, [batch_size, num_frames, 6], dev_str=dev_str)
            pose_cov = ivy.random_uniform(1e-3, 1, [batch_size, num_frames, 6, 6], dev_str=dev_str)
        if empty:
            img_var = ivy.ones_like(img_var) * 1e12
            validity_mask = ivy.zeros_like(validity_mask)
        img_meas['dummy_cam_{}'.format(i)] =\
            {'img_mean': img_mean,
             'img_var': img_var,
             'validity_mask': validity_mask,
             'pose_mean': pose_mean,
             'pose_cov': pose_cov,
             'cam_rel_mat': ivy.identity(4, batch_shape=[batch_size, num_frames], dev_str=dev_str)[..., 0:3, :]}

    if ones:
        control_mean = ivy.zeros([batch_size, num_frames, 6], dev_str=dev_str)
        control_cov = ivy.ones([batch_size, num_frames, 6, 6], dev_str=dev_str)*1e-3
    else:
        control_mean = ivy.random_uniform(1e-3, 1, [batch_size, num_frames, 6], dev_str=dev_str)
        control_cov = ivy.random_uniform(1e-3, 1, [batch_size, num_frames, 6, 6], dev_str=dev_str)
    return Container({'img_meas': img_meas,
                      'control_mean': control_mean,
                      'control_cov': control_cov,
                      'agent_rel_mat': ivy.identity(4, batch_shape=[batch_size, num_frames],
                                                    dev_str=dev_str)[..., 0:3, :]})
Esempio n. 6
0
def make_coordinates_homogeneous(coords, batch_shape=None):
    """
    Append ones to array of non-homogeneous co-ordinates to make them homogeneous.

    :param coords: Array of non-homogeneous co-ordinates *[batch_shape,d]*
    :type coords: array
    :param batch_shape: Shape of batch. Inferred from inputs if None.
    :type batch_shape: sequence of ints, optional
    :return: Array of Homogeneous co-ordinates *[batch_shape,d+1]*
    """

    if batch_shape is None:
        batch_shape = coords.shape[:-1]

    # shapes as list
    batch_shape = list(batch_shape)

    # BS x 1
    ones = _ivy.ones_like(coords[..., 0:1])

    # BS x (D+1)
    return _ivy.concatenate((coords, ones), -1)
Esempio n. 7
0
def sample_spline_path(anchor_points, anchor_vals, sample_points, order=3):
    """
    Sample spline path, given sample locations for path defined by the anchor locations and points.
    `[reference] <https://github.com/tensorflow/addons/blob/v0.11.2/tensorflow_addons/image/interpolate_spline.py>`_

    :param anchor_points: Anchor locations between 0-1 (regular spacing not necessary) *[batch_shape,num_anchors,1]*
    :type anchor_points: array
    :param anchor_vals: Anchor points along the spline path, in path space *[batch_shape,num_anchors,path_dim]*
    :type anchor_vals: array
    :param sample_points: Sample locations between 0-1 *[batch_shape,num_samples,1]*
    :type sample_points: array
    :param order: Order of the spline path interpolation
    :type order: float
    :return: Spline path sampled at sample_locations, giving points in path space *[batch_shape,num_samples,path_dim]*
    """

    # BS x N x PD,    BS x 2 x PD
    w, v = _fit_spline(anchor_points, anchor_vals, order)

    # Kernel term

    # BS x NS x N
    pairwise_dists = _pairwise_distance(sample_points, anchor_points)
    phi_pairwise_dists = _phi(pairwise_dists, order)

    # BS x NS x PD
    rbf_term = _ivy.matmul(phi_pairwise_dists, w)

    # Polynomial / linear term.

    # BS x NS x 2
    query_points_pad = _ivy.concatenate(
        [sample_points, _ivy.ones_like(sample_points[..., :1])], -1)

    # BS x NS x PD
    linear_term = _ivy.matmul(query_points_pad, v)
    return rbf_term + linear_term
Esempio n. 8
0
File: esm.py Progetto: wx-b/memory
    def _convert_images_to_omni_observations(self, measurements,
                                             uniform_sphere_pixel_coords,
                                             holes_prior, batch_size,
                                             num_timesteps, num_cams,
                                             image_dims):
        """
        Convert image to omni-directional measurements

        :param measurements: perspective captured images and relative poses container
        :param uniform_sphere_pixel_coords: Uniform  sphere pixel coords *[batch_size, num_timesteps, oh, ow, 3]*
        :param holes_prior: Prior for quantization holes *[batch_size, num_timesteps, oh, ow, 1+f]*
        :param batch_size: Size of batch
        :param num_timesteps: Number of frames
        :param num_cams: Number of cameras
        :param image_dims: Image dimensions
        :return: *[batch_size, n, oh, ow, 3+f]*    *[batch_size, n, oh, ow, 3+f]*
        """

        # coords from all scene cameras wrt world

        images_list = list()
        images_var_list = list()
        cam_rel_poses_list = list()
        cam_rel_poses_cov_list = list()
        cam_rel_mats_list = list()
        validity_mask_list = list()
        for key, item in measurements.to_iterator():
            if key == 'img_mean':
                # B x N x 1 x H x W x (3+f)
                images_list.append(ivy.expand_dims(item, 2))
            elif key == 'img_var':
                # B x N x 1 x H x W x (3+f)
                images_var_list.append(ivy.expand_dims(item, 2))
            elif key == 'pose_mean':
                # B x N x 1 x 6
                cam_rel_poses_list.append(ivy.expand_dims(item, 2))
            elif key == 'pose_cov':
                # B x N x 1 x 6 x 6
                cam_rel_poses_cov_list.append(ivy.expand_dims(item, 2))
            elif key == 'cam_rel_mat':
                # B x N x 1 x 3 x 4
                cam_rel_mats_list.append(ivy.expand_dims(item, 2))
            elif key == 'validity_mask':
                validity_mask_list.append(ivy.expand_dims(item, 2))
            else:
                raise Exception('Invalid image key: {}'.format(key))

        # B x N x C x H x W x (3+f)
        images = ivy.concatenate(images_list, 2)

        # B x N x C x H x W x (3+f)
        var_to_project = ivy.concatenate(images_var_list, 2)

        # B x N x C x 6
        cam_to_cam_poses = ivy.concatenate(cam_rel_poses_list, 2)

        # B x N x C x 3 x 4
        cam_to_cam_mats = ivy.concatenate(cam_rel_mats_list, 2)

        # B x N x C x 6 x 6
        cam_to_cam_pose_covs = ivy.concatenate(cam_rel_poses_cov_list, 2)

        # B x N x C x 1
        validity_masks = ivy.concatenate(validity_mask_list, 2) > 0

        # B x N x OH x OW x (3+f)
        holes_prior_var = ivy.ones(
            [batch_size, num_timesteps] + self._sphere_img_dims +
            [3 + self._feat_dim],
            dev_str=self._dev_str) * 1e12

        # reset invalid regions to prior

        # B x N x C x H x W x (3+f)
        images = ivy.where(
            validity_masks, images,
            ivy.concatenate(
                (images[..., 0:2],
                 ivy.zeros_like(images[..., 2:], dev_str=self._dev_str)), -1))

        # B x N x C x H x W x (3+f)
        var_to_project = ivy.where(
            validity_masks, var_to_project,
            ivy.ones_like(var_to_project, dev_str=self._dev_str) * 1e12)

        # B x N x OH x OW x (3+f)    # B x N x OH x OW x (3+f)
        return self._frame_to_omni_frame_projection(
            cam_to_cam_poses, cam_to_cam_mats, uniform_sphere_pixel_coords,
            images[..., 0:3], images[..., 3:], cam_to_cam_pose_covs,
            var_to_project, holes_prior, holes_prior_var, batch_size,
            num_timesteps, num_cams, image_dims)
Esempio n. 9
0
def main(f=None):

    # Framework Setup #
    # ----------------#

    # choose random framework

    f = choose_random_framework() if f is None else f
    set_framework(f)

    # Orientation #
    # ------------#

    # rotation representations

    # 3
    rot_vec = ivy.array([0., 1., 0.])

    # 3 x 3
    rot_mat = ivy_mech.rot_vec_to_rot_mat(rot_vec)

    # 3
    euler_angles = ivy_mech.rot_mat_to_euler(rot_mat, 'zyx')

    # 4
    quat = ivy_mech.euler_to_quaternion(euler_angles)

    # 4
    axis_and_angle = ivy_mech.quaternion_to_axis_angle(quat)

    # 3
    rot_vec_again = axis_and_angle[..., :-1] * axis_and_angle[..., -1:]

    # Pose #
    # -----#

    # pose representations

    # 3
    position = ivy.ones_like(rot_vec)

    # 6
    rot_vec_pose = ivy.concatenate((position, rot_vec), 0)

    # 3 x 4
    mat_pose = ivy_mech.rot_vec_pose_to_mat_pose(rot_vec_pose)

    # 6
    euler_pose = ivy_mech.mat_pose_to_euler_pose(mat_pose)

    # 7
    quat_pose = ivy_mech.euler_pose_to_quaternion_pose(euler_pose)

    # 6
    rot_vec_pose_again = ivy_mech.quaternion_pose_to_rot_vec_pose(quat_pose)

    # Position #
    # ---------#

    # conversions of positional representation

    # 3
    cartesian_coord = ivy.random_uniform(0., 1., (3, ))

    # 3
    polar_coord = ivy_mech.cartesian_to_polar_coords(cartesian_coord)

    # 3
    cartesian_coord_again = ivy_mech.polar_to_cartesian_coords(polar_coord)

    # cartesian co-ordinate frame-of-reference transformations

    # 3 x 4
    trans_mat = ivy.random_uniform(0., 1., (3, 4))

    # 4
    cartesian_coord_homo = ivy_mech.make_coordinates_homogeneous(
        cartesian_coord)

    # 3
    trans_cartesian_coord = ivy.matmul(
        trans_mat, ivy.expand_dims(cartesian_coord_homo, -1))[:, 0]

    # 4
    trans_cartesian_coord_homo = ivy_mech.make_coordinates_homogeneous(
        trans_cartesian_coord)

    # 4 x 4
    trans_mat_homo = ivy_mech.make_transformation_homogeneous(trans_mat)

    # 3 x 4
    inv_trans_mat = ivy.inv(trans_mat_homo)[0:3]

    # 3
    cartesian_coord_again = ivy.matmul(
        inv_trans_mat, ivy.expand_dims(trans_cartesian_coord_homo, -1))[:, 0]

    # message
    print('End of Run Through Demo!')
Esempio n. 10
0
def cuboid_signed_distances(cuboid_ext_mats,
                            cuboid_dims,
                            query_positions,
                            batch_shape=None):
    """
    Return the signed distances of a set of query points from the cuboid surfaces.\n
    `[reference] <https://www.iquilezles.org/www/articles/distfunctions/distfunctions.htm>`_

    :param cuboid_ext_mats: Extrinsic matrices of the cuboids *[batch_shape,num_cuboids,3,4]*
    :type cuboid_ext_mats: array
    :param cuboid_dims: Dimensions of the cuboids, in the order x, y, z *[batch_shape,num_cuboids,3]*
    :type cuboid_dims: array
    :param query_positions: Points for which to query the signed distances *[batch_shape,num_points,3]*
    :type query_positions: array
    :param batch_shape: Shape of batch. Assumed no batches if None.
    :type batch_shape: sequence of ints, optional
    :return: The distances of the query points from the closest cuboid surface *[batch_shape,num_points,1]*
    """

    if batch_shape is None:
        batch_shape = cuboid_ext_mats.shape[:-3]

    # shapes as list
    batch_shape = list(batch_shape)
    num_batch_dims = len(batch_shape)
    batch_dims_for_trans = list(range(num_batch_dims))
    num_cuboids = cuboid_ext_mats.shape[-3]
    num_points = query_positions.shape[-2]

    # BS x 3 x NP
    query_positions_trans = _ivy.transpose(
        query_positions,
        batch_dims_for_trans + [num_batch_dims + 1, num_batch_dims])

    # BS x 1 x NP
    ones = _ivy.ones_like(query_positions_trans[..., 0:1, :])

    # BS x 4 x NP
    query_positions_trans_homo = _ivy.concatenate(
        (query_positions_trans, ones), -2)

    # BS x NCx3 x 4
    cuboid_ext_mats_flat = _ivy.reshape(cuboid_ext_mats, batch_shape + [-1, 4])

    # BS x NCx3 x NP
    rel_query_positions_trans_flat = _ivy.matmul(cuboid_ext_mats_flat,
                                                 query_positions_trans_homo)

    # BS x NC x 3 x NP
    rel_query_positions_trans = _ivy.reshape(
        rel_query_positions_trans_flat,
        batch_shape + [num_cuboids, 3, num_points])

    # BS x NC x NP x 3
    rel_query_positions = _ivy.transpose(
        rel_query_positions_trans, batch_dims_for_trans +
        [num_batch_dims, num_batch_dims + 2, num_batch_dims + 1])
    q = _ivy.abs(rel_query_positions) - _ivy.expand_dims(cuboid_dims / 2, -2)
    q_max_clipped = _ivy.maximum(q, 1e-12)

    # BS x NC x NP x 1
    q_min_clipped = _ivy.minimum(_ivy.reduce_max(q, -1, keepdims=True), 0.)
    q_max_clipped_len = _ivy.reduce_sum(q_max_clipped**2, -1,
                                        keepdims=True)**0.5
    sdfs = q_max_clipped_len + q_min_clipped

    # BS x NP x 1
    return _ivy.reduce_min(sdfs, -3)
Esempio n. 11
0
    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
Esempio n. 12
0
def render_implicit_features_and_depth(network_fn,
                                       rays_o,
                                       rays_d,
                                       near,
                                       far,
                                       samples_per_ray,
                                       render_variance=False,
                                       batch_size_per_query=512 * 64,
                                       inter_feat_fn=None,
                                       v=None):
    """
    Render an rgb-d image, given an implicit rgb and density function conditioned on xyz data.

    :param network_fn: the implicit function.
    :type network_fn: callable
    :param rays_o: The camera center *[batch_shape,3]*
    :type rays_o: array
    :param rays_d: The rays in world space *[batch_shape,ray_batch_shape,3]*
    :type rays_d: array
    :param near: The near clipping plane values *[batch_shape,ray_batch_shape]*
    :type near: array
    :param far: The far clipping plane values *[batch_shape,ray_batch_shape]*
    :type far: array
    :param samples_per_ray: The number of stratified samples to use along each ray
    :type samples_per_ray: int
    :param render_variance: Whether to also render the feature variance. Default is False.
    :type render_variance: bool, optional
    :param batch_size_per_query: The maximum batch size for querying the implicit network. Default is 1024.
    :type batch_size_per_query: int, optional
    :param inter_feat_fn: Function to extract interpolated features from world-coords *[batch_shape,ray_batch_shape,3]*
    :type inter_feat_fn: callable, optional
    :param v: The container of trainable variables for the implicit model. default is to use internal variables.
    :type v: ivy Container of variables
    :return: The rendered feature *[batch_shape,ray_batch_shape,feat]* and depth *[batch_shape,ray_batch_shape,1]* values
    """

    # shapes
    batch_shape = list(near.shape)
    flat_batch_size = np.prod(batch_shape)
    num_sections = math.ceil(flat_batch_size * samples_per_ray /
                             batch_size_per_query)

    # Compute 3D query points

    # BS x SPR x 1
    z_vals = ivy.expand_dims(stratified_sample(near, far, samples_per_ray), -1)

    # BS x 1 x 3
    rays_d = ivy.expand_dims(rays_d, -2)
    rays_o = ivy.broadcast_to(rays_o, rays_d.shape)

    # BS x SPR x 3
    pts = rays_o + rays_d * z_vals

    # (BSxSPR) x 3
    pts_flat = ivy.reshape(pts, (flat_batch_size * samples_per_ray, 3))

    # batch
    # ToDo: use a more general batchify function, from ivy core

    # num_sections size list of BSPQ x 3
    pts_split = [
        pts_flat[i * batch_size_per_query:min(
            (i + 1) * batch_size_per_query, flat_batch_size * samples_per_ray)]
        for i in range(num_sections)
    ]
    if inter_feat_fn is not None:
        # (BSxSPR) x IF
        features = ivy.reshape(inter_feat_fn(pts),
                               (flat_batch_size * samples_per_ray, -1))
        # num_sections size list of BSPQ x IF
        feats_split =\
            [features[i * batch_size_per_query:min((i + 1) * batch_size_per_query, flat_batch_size*samples_per_ray)]
             for i in range(num_sections)]
    else:
        feats_split = [None] * num_sections

    # Run network

    # num_sections size list of tuple of (BSPQ x OF, BSPQ)
    feats_n_densities = [
        network_fn(pt, f, v=v) for pt, f in zip(pts_split, feats_split)
    ]

    # BS x SPR x OF
    feat = ivy.reshape(
        ivy.concatenate([item[0] for item in feats_n_densities], 0),
        batch_shape + [samples_per_ray, -1])

    # FBS x SPR
    densities = ivy.reshape(
        ivy.concatenate([item[1] for item in feats_n_densities], 0),
        batch_shape + [samples_per_ray])

    # BS x (SPR+1)
    z_vals_w_terminal = ivy.concatenate(
        (z_vals[..., 0], ivy.ones_like(z_vals[..., -1:, 0]) * 1e10), -1)

    # BS x SPR
    depth_diffs = z_vals_w_terminal[..., 1:] - z_vals_w_terminal[..., :-1]
    ray_term_probs = ray_termination_probabilities(densities, depth_diffs)

    # BS x OF
    feat = ivy.clip(
        render_rays_via_termination_probabilities(ray_term_probs, feat,
                                                  render_variance), 0., 1.)

    # BS x 1
    depth = render_rays_via_termination_probabilities(ray_term_probs, z_vals,
                                                      render_variance)
    if render_variance:
        # BS x OF, BS x OF, BS x 1, BS x 1
        return feat[0], feat[1], depth[0], depth[1]
    # BS x OF, BS x 1
    return feat, depth
Esempio n. 13
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