Ejemplo n.º 1
0
def _rot_mat_to_zxz_euler(rot_mat):
    # BS x 1
    euler_angles_1 = _ivy.acos(rot_mat[..., 2, 2:3])

    gimbal_validity = _ivy.abs(rot_mat[..., 0, 2:3]) > GIMBAL_TOL

    r12 = rot_mat[..., 0, 1:2]
    r11 = rot_mat[..., 0, 0:1]
    gimbal_euler_angles_0 = _ivy.atan2(-r12, r11)
    gimbal_euler_angles_2 = _ivy.zeros_like(gimbal_euler_angles_0)

    # BS x 3
    gimbal_euler_angles = _ivy.concatenate(
        (gimbal_euler_angles_0, euler_angles_1, gimbal_euler_angles_2), -1)

    # BS x 1
    r31 = rot_mat[..., 2, 0:1]
    r32 = rot_mat[..., 2, 1:2]
    r13 = rot_mat[..., 0, 2:3]
    r23 = rot_mat[..., 1, 2:3]
    normal_euler_angles_0 = _ivy.atan2(r31, r32)
    normal_euler_angles_2 = _ivy.atan2(r13, -r23)

    # BS x 3
    normal_euler_angles = _ivy.concatenate(
        (normal_euler_angles_0, euler_angles_1, normal_euler_angles_2), -1)

    return _ivy.where(gimbal_validity, normal_euler_angles,
                      gimbal_euler_angles)
Ejemplo n.º 2
0
def _rot_mat_to_yzy_euler(rot_mat):
    # BS x 1
    euler_angles_1 = _ivy.acos(rot_mat[..., 1, 1:2])

    gimbal_validity = _ivy.abs(rot_mat[..., 1, 0:1]) > GIMBAL_TOL

    r31 = rot_mat[..., 2, 0:1]
    r33 = rot_mat[..., 2, 2:3]
    gimbal_euler_angles_0 = _ivy.atan2(-r31, r33)
    gimbal_euler_angles_2 = _ivy.zeros_like(gimbal_euler_angles_0)

    # BS x 3
    gimbal_euler_angles = _ivy.concatenate(
        (gimbal_euler_angles_0, euler_angles_1, gimbal_euler_angles_2), -1)

    # BS x 1
    r23 = rot_mat[..., 1, 2:3]
    r21 = rot_mat[..., 1, 0:1]
    r32 = rot_mat[..., 2, 1:2]
    r12 = rot_mat[..., 0, 1:2]
    normal_euler_angles_0 = _ivy.atan2(r23, r21)
    normal_euler_angles_2 = _ivy.atan2(r32, r12)

    # BS x 3
    normal_euler_angles = _ivy.concatenate(
        (normal_euler_angles_0, euler_angles_1, normal_euler_angles_2), -1)

    return _ivy.where(gimbal_validity, normal_euler_angles,
                      gimbal_euler_angles)
Ejemplo n.º 3
0
def _rot_mat_to_zyx_euler(rot_mat):
    # BS x 1
    euler_angles_1 = _ivy.asin(rot_mat[..., 0, 2:3])

    gimbal_validity = _ivy.abs(rot_mat[..., 1, 1:2]) > GIMBAL_TOL

    r21 = rot_mat[..., 1, 0:1]
    r22 = rot_mat[..., 1, 1:2]
    gimbal_euler_angles_0 = _ivy.atan2(r21, r22)
    gimbal_euler_angles_2 = _ivy.zeros_like(gimbal_euler_angles_0)

    # BS x 3
    gimbal_euler_angles = _ivy.concatenate(
        (gimbal_euler_angles_0, euler_angles_1, gimbal_euler_angles_2), -1)

    # BS x 1
    r12 = rot_mat[..., 0, 1:2]
    r11 = rot_mat[..., 0, 0:1]
    r23 = rot_mat[..., 1, 2:3]
    r33 = rot_mat[..., 2, 2:3]
    normal_euler_angles_0 = _ivy.atan2(-r12, r11)
    normal_euler_angles_2 = _ivy.atan2(-r23, r33)

    # BS x 3
    normal_euler_angles = _ivy.concatenate(
        (normal_euler_angles_0, euler_angles_1, normal_euler_angles_2), -1)

    return _ivy.where(gimbal_validity, normal_euler_angles,
                      gimbal_euler_angles)
Ejemplo n.º 4
0
Archivo: esm.py Proyecto: wx-b/memory
    def smooth(self, fused_val, fused_variance, low_var_mask, smooth_mean,
               smooth_kernel_size, variance_mode, fix_low_var_pixels,
               batch_size):
        variance_mode = True
        var_for_smoothing = fused_variance

        # image smoothing
        if smooth_mean:

            # pad borders
            pad_size = int(smooth_kernel_size / 2)
            fused_val_pad = ivy_vision.pad_omni_image(fused_val, pad_size,
                                                      self._sphere_img_dims)
            fused_variance_pad = ivy_vision.pad_omni_image(
                var_for_smoothing, pad_size, self._sphere_img_dims)
            expanded_sphere_img_dims = [
                item + 2 * pad_size for item in self._sphere_img_dims
            ]

            # reshape for dilation and erosion
            fused_val_pad_flat = ivy.reshape(
                fused_val_pad[..., 2:],
                [batch_size] + expanded_sphere_img_dims + [1 + self._feat_dim])
            fused_var_pad_flat = ivy.reshape(
                fused_variance_pad[..., 2:],
                [batch_size] + expanded_sphere_img_dims + [1 + self._feat_dim])

            if variance_mode:
                smoothed_fused_val_flat, smoothed_fused_var_flat = \
                    ivy_vision.smooth_image_fom_var_image(fused_val_pad_flat, fused_var_pad_flat, smooth_kernel_size,
                                                          ivy.array([1.] * (1 + self._feat_dim), dev_str=self._dev_str))
            else:
                smoothed_fused_val_flat, smoothed_fused_var_flat = \
                    ivy_vision.weighted_image_smooth(fused_val_pad_flat, 1 - fused_var_pad_flat, smooth_kernel_size)

            # reshape to image dims
            smoothed_fused_val = ivy.reshape(
                smoothed_fused_val_flat,
                [batch_size] + self._sphere_img_dims + [1 + self._feat_dim])
            smoothed_fused_val = ivy.concatenate(
                (fused_val[..., 0:2], smoothed_fused_val), -1)

            # replace temporary zeros with their prior values
            # This ensures that the smoothing operation only changes the values for regions of high variance
            if fix_low_var_pixels:
                fused_val = ivy.where(low_var_mask, fused_val,
                                      smoothed_fused_val)
            else:
                fused_val = smoothed_fused_val

        return fused_val, fused_variance
Ejemplo n.º 5
0
Archivo: esm.py Proyecto: wx-b/memory
    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
Ejemplo n.º 6
0
Archivo: esm.py Proyecto: 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)
Ejemplo n.º 7
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
Ejemplo n.º 8
0
def rot_mat_to_quaternion(rot_mat):
    """
    Convert rotation matrix :math:`\mathbf{R}\in\mathbb{R}^{3×3}` to quaternion
    :math:`\mathbf{q} = [q_i, q_j, q_k, q_r]`.\n
    `[reference] <http://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToQuaternion/>`_

    :param rot_mat: Rotation matrix *[batch_shape,3,3]*
    :type rot_mat: array
    :return: Quaternion *[batch_shape,4]*
    """

    # BS x 1 x 1
    tr = rot_mat[..., 0:1, 0:1] + rot_mat[..., 1:2, 1:2] + rot_mat[..., 2:3,
                                                                   2:3]

    # if tf > 0
    # BS x 1 x 1
    s_1 = ((tr + 1)**0.5) * 2
    qw_1 = 0.25 * s_1
    qx_1 = (rot_mat[..., 2:3, 1:2] -
            rot_mat[..., 1:2, 2:3]) / (s_1 + MIN_DENOMINATOR)
    qy_1 = (rot_mat[..., 0:1, 2:3] -
            rot_mat[..., 2:3, 0:1]) / (s_1 + MIN_DENOMINATOR)
    qz_1 = (rot_mat[..., 1:2, 0:1] -
            rot_mat[..., 0:1, 1:2]) / (s_1 + MIN_DENOMINATOR)

    # BS x 4 x 1
    quat_1 = _ivy.concatenate((qx_1, qy_1, qz_1, qw_1), -2)

    # elif (m[:,0,0] > m[:,1,1]) and (m[:,0,0] > m[:,2,2])
    # BS x 1 x 1
    s_2 = ((1 + rot_mat[..., 0:1, 0:1] - rot_mat[..., 1:2, 1:2] -
            rot_mat[..., 2:3, 2:3])**0.5) * 2
    qw_2 = (rot_mat[..., 2:3, 1:2] -
            rot_mat[..., 1:2, 2:3]) / (s_2 + MIN_DENOMINATOR)
    qx_2 = 0.25 * s_2
    qy_2 = (rot_mat[..., 0:1, 1:2] +
            rot_mat[..., 1:2, 0:1]) / (s_2 + MIN_DENOMINATOR)
    qz_2 = (rot_mat[..., 0:1, 2:3] +
            rot_mat[..., 2:3, 0:1]) / (s_2 + MIN_DENOMINATOR)

    # BS x 4 x 1
    quat_2 = _ivy.concatenate((qx_2, qy_2, qz_2, qw_2), -2)

    # elif m[:,1,1] > m[:,2,2]
    # BS x 1 x 1
    s_3 = ((1 + rot_mat[..., 1:2, 1:2] - rot_mat[..., 0:1, 0:1] -
            rot_mat[..., 2:3, 2:3])**0.5) * 2
    qw_3 = (rot_mat[..., 0:1, 2:3] -
            rot_mat[..., 2:3, 0:1]) / (s_3 + MIN_DENOMINATOR)
    qx_3 = (rot_mat[..., 0:1, 1:2] +
            rot_mat[..., 1:2, 0:1]) / (s_3 + MIN_DENOMINATOR)
    qy_3 = 0.25 * s_3
    qz_3 = (rot_mat[..., 1:2, 2:3] +
            rot_mat[..., 2:3, 1:2]) / (s_3 + MIN_DENOMINATOR)

    # BS x 4 x 1
    quat_3 = _ivy.concatenate((qx_3, qy_3, qz_3, qw_3), -2)

    # else
    # BS x 1 x 1
    s_4 = ((1 + rot_mat[..., 2:3, 2:3] - rot_mat[..., 0:1, 0:1] -
            rot_mat[..., 1:2, 1:2])**0.5) * 2
    qw_4 = (rot_mat[..., 1:2, 0:1] -
            rot_mat[..., 0:1, 1:2]) / (s_4 + MIN_DENOMINATOR)
    qx_4 = (rot_mat[..., 0:1, 2:3] +
            rot_mat[..., 2:3, 0:1]) / (s_4 + MIN_DENOMINATOR)
    qy_4 = (rot_mat[..., 1:2, 2:3] +
            rot_mat[..., 2:3, 1:2]) / (s_4 + MIN_DENOMINATOR)
    qz_4 = 0.25 * s_4

    # BS x 4 x 1
    quat_4 = _ivy.concatenate((qx_4, qy_4, qz_4, qw_4), -2)
    quat_3_or_other = _ivy.where(
        rot_mat[..., 1:2, 1:2] > rot_mat[..., 2:3, 2:3], quat_3, quat_4)
    quat_2_or_other = \
        _ivy.where(_ivy.logical_and((rot_mat[..., 0:1, 0:1] > rot_mat[..., 1:2, 1:2]),
                                    (rot_mat[..., 0:1, 0:1] > rot_mat[..., 2:3, 2:3])), quat_2, quat_3_or_other)

    # BS x 4
    return _ivy.where(tr > 0, quat_1, quat_2_or_other)[..., 0]
Ejemplo n.º 9
0
def velocity_from_cam_coords_id_image_and_object_trans(cam_coords_t,
                                                       id_image,
                                                       obj_ids,
                                                       obj_trans,
                                                       delta_t,
                                                       batch_shape=None,
                                                       image_dims=None,
                                                       dev_str=None):
    """
    Compute velocity image from co-ordinate image, id image, and object transformations.

    :param cam_coords_t: Camera-centric homogeneous co-ordinates image in frame t *[batch_shape,h,w,4]*
    :type cam_coords_t: 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 delta_t: Time difference between frame at timestep t-1 and t *[batch_shape,1]*
    :type delta_t: array
    :param batch_shape: Shape of batch. Inferred from inputs if None.
    :type batch_shape: sequence of ints, optional
    :param image_dims: Image dimensions. Inferred from inputs in None.
    :type image_dims: sequence of ints, optional
    :param dev_str: device on which to create the array 'cuda:0', 'cuda:1', 'cpu' etc. Same as x if None.
    :type dev_str: str, optional
    :return: Relative velocity image *[batch_shape,h,w,3]*
    """

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

    if image_dims is None:
        image_dims = cam_coords_t.shape[-3:-1]

    if dev_str is None:
        dev_str = _ivy.dev_str(cam_coords_t)

    # shapes as list
    batch_shape = list(batch_shape)
    image_dims = list(image_dims)

    # get co-ordinate re-projections

    # BS x H x W x 4
    cam_coords_t_all_trans, motion_mask =\
        project_cam_coords_with_object_transformations(cam_coords_t, id_image, obj_ids, obj_trans,
                                                       _ivy.identity(4, batch_shape=batch_shape)[..., 0:3, :],
                                                       batch_shape, image_dims)

    # BS x H x W x 4
    cam_coords_t_all_trans = \
        _ivy.where(motion_mask, cam_coords_t_all_trans, _ivy.zeros_like(cam_coords_t_all_trans, dev_str=dev_str))

    # compute velocities

    # BS x H x W x 3
    vel = (cam_coords_t[..., 0:3] - cam_coords_t_all_trans[..., 0:3]) / delta_t

    # prune velocities

    # BS x H x W x 3
    return _ivy.where(motion_mask, vel, _ivy.zeros_like(vel, dev_str=dev_str))
Ejemplo n.º 10
0
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
Ejemplo n.º 11
0
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
Ejemplo n.º 12
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
Ejemplo n.º 13
0
def main(interactive=True, f=None):

    global INTERACTIVE
    INTERACTIVE = interactive

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

    # choose random framework
    f = choose_random_framework() if f is None else f
    set_framework(f)

    # Camera Geometry #
    # ----------------#

    # intrinsics

    # common intrinsic params
    img_dims = [512, 512]
    pp_offsets = ivy.array([dim / 2 - 0.5 for dim in img_dims], 'float32')
    cam_persp_angles = ivy.array([60 * np.pi / 180] * 2, 'float32')

    # ivy cam intrinsics container
    intrinsics = ivy_vision.persp_angles_and_pp_offsets_to_intrinsics_object(
        cam_persp_angles, pp_offsets, img_dims)

    # extrinsics

    # 3 x 4
    cam1_inv_ext_mat = ivy.array(np.load(data_dir + '/cam1_inv_ext_mat.npy'),
                                 'float32')
    cam2_inv_ext_mat = ivy.array(np.load(data_dir + '/cam2_inv_ext_mat.npy'),
                                 'float32')

    # full geometry

    # ivy cam geometry container
    cam1_geom = ivy_vision.inv_ext_mat_and_intrinsics_to_cam_geometry_object(
        cam1_inv_ext_mat, intrinsics)
    cam2_geom = ivy_vision.inv_ext_mat_and_intrinsics_to_cam_geometry_object(
        cam2_inv_ext_mat, intrinsics)
    cam_geoms = [cam1_geom, cam2_geom]

    # Camera Geometry Check #
    # ----------------------#

    # assert camera geometry shapes

    for cam_geom in cam_geoms:

        assert cam_geom.intrinsics.focal_lengths.shape == (2, )
        assert cam_geom.intrinsics.persp_angles.shape == (2, )
        assert cam_geom.intrinsics.pp_offsets.shape == (2, )
        assert cam_geom.intrinsics.calib_mats.shape == (3, 3)
        assert cam_geom.intrinsics.inv_calib_mats.shape == (3, 3)

        assert cam_geom.extrinsics.cam_centers.shape == (3, 1)
        assert cam_geom.extrinsics.Rs.shape == (3, 3)
        assert cam_geom.extrinsics.inv_Rs.shape == (3, 3)
        assert cam_geom.extrinsics.ext_mats_homo.shape == (4, 4)
        assert cam_geom.extrinsics.inv_ext_mats_homo.shape == (4, 4)

        assert cam_geom.full_mats_homo.shape == (4, 4)
        assert cam_geom.inv_full_mats_homo.shape == (4, 4)

    # Image Data #
    # -----------#

    # load images

    # h x w x 3
    color1 = ivy.array(
        cv2.imread(data_dir + '/rgb1.png').astype(np.float32) / 255)
    color2 = ivy.array(
        cv2.imread(data_dir + '/rgb2.png').astype(np.float32) / 255)

    # h x w x 1
    depth1 = ivy.array(
        np.reshape(
            np.frombuffer(
                cv2.imread(data_dir + '/depth1.png', -1).tobytes(),
                np.float32), img_dims + [1]))
    depth2 = ivy.array(
        np.reshape(
            np.frombuffer(
                cv2.imread(data_dir + '/depth2.png', -1).tobytes(),
                np.float32), img_dims + [1]))

    # depth scaled pixel coords

    # h x w x 3
    u_pix_coords = ivy_vision.create_uniform_pixel_coords_image(img_dims)
    ds_pixel_coords1 = u_pix_coords * depth1
    ds_pixel_coords2 = u_pix_coords * depth2

    # depth limits
    depth_min = ivy.reduce_min(ivy.concatenate((depth1, depth2), 0))
    depth_max = ivy.reduce_max(ivy.concatenate((depth1, depth2), 0))
    depth_limits = [depth_min, depth_max]

    # show images
    show_rgb_and_depth_images(color1, color2, depth1, depth2, depth_limits)

    # Flow and Depth Triangulation #
    # -----------------------------#

    # required mat formats
    cam1to2_full_mat_homo = ivy.matmul(cam2_geom.full_mats_homo,
                                       cam1_geom.inv_full_mats_homo)
    cam1to2_full_mat = cam1to2_full_mat_homo[..., 0:3, :]
    full_mats_homo = ivy.concatenate(
        (ivy.expand_dims(cam1_geom.full_mats_homo,
                         0), ivy.expand_dims(cam2_geom.full_mats_homo, 0)), 0)
    full_mats = full_mats_homo[..., 0:3, :]

    # flow
    flow1to2 = ivy_vision.flow_from_depth_and_cam_mats(ds_pixel_coords1,
                                                       cam1to2_full_mat)

    # depth again
    depth1_from_flow = ivy_vision.depth_from_flow_and_cam_mats(
        flow1to2, full_mats)

    # show images
    show_flow_and_depth_images(depth1, flow1to2, depth1_from_flow,
                               depth_limits)

    # Inverse Warping #
    # ----------------#

    # inverse warp rendering
    warp = u_pix_coords[..., 0:2] + flow1to2
    color2_warp_to_f1 = ivy.reshape(ivy.bilinear_resample(color2, warp),
                                    color1.shape)

    # projected depth scaled pixel coords 2
    ds_pixel_coords1_wrt_f2 = ivy_vision.ds_pixel_to_ds_pixel_coords(
        ds_pixel_coords1, cam1to2_full_mat)

    # projected depth 2
    depth1_wrt_f2 = ds_pixel_coords1_wrt_f2[..., -1:]

    # inverse warp depth
    depth2_warp_to_f1 = ivy.reshape(ivy.bilinear_resample(depth2, warp),
                                    depth1.shape)

    # depth validity
    depth_validity = ivy.abs(depth1_wrt_f2 - depth2_warp_to_f1) < 0.01

    # inverse warp rendering with mask
    color2_warp_to_f1_masked = ivy.where(depth_validity, color2_warp_to_f1,
                                         ivy.zeros_like(color2_warp_to_f1))

    # show images
    show_inverse_warped_images(depth1_wrt_f2, depth2_warp_to_f1,
                               depth_validity, color1, color2_warp_to_f1,
                               color2_warp_to_f1_masked, depth_limits)

    # Forward Warping #
    # ----------------#

    # forward warp rendering
    ds_pixel_coords1_proj = ivy_vision.ds_pixel_to_ds_pixel_coords(
        ds_pixel_coords2,
        ivy.inv(cam1to2_full_mat_homo)[..., 0:3, :])
    depth1_proj = ds_pixel_coords1_proj[..., -1:]
    ds_pixel_coords1_proj = ds_pixel_coords1_proj[..., 0:2] / depth1_proj
    features_to_render = ivy.concatenate((depth1_proj, color2), -1)

    # without depth buffer
    f1_forward_warp_no_db, _, _ = ivy_vision.quantize_to_image(
        ivy.reshape(ds_pixel_coords1_proj, (-1, 2)),
        img_dims,
        ivy.reshape(features_to_render, (-1, 4)),
        ivy.zeros_like(features_to_render),
        with_db=False)

    # with depth buffer
    f1_forward_warp_w_db, _, _ = ivy_vision.quantize_to_image(
        ivy.reshape(ds_pixel_coords1_proj, (-1, 2)),
        img_dims,
        ivy.reshape(features_to_render, (-1, 4)),
        ivy.zeros_like(features_to_render),
        with_db=False if ivy.get_framework() == 'mxnd' else True)

    # show images
    show_forward_warped_images(depth1, color1, f1_forward_warp_no_db,
                               f1_forward_warp_w_db, depth_limits)

    # message
    print('End of Run Through Demo!')