Exemple #1
0
def read(config):
    task_name = config['task_name']
    nb_object = config['object']['nb_object']
    fps = config['animation']['fps']

    se3_path = './output/' + task_name + '/pose/se3_pose.npy'
    demo_list = os.listdir('./data/' + task_name)
    output_dir = './output/' + task_name + '/read_pose'

    se3_dict = np.load(se3_path, allow_pickle=True).item()
    obj = vision.SE3object(np.zeros(6), angle_type='axis')
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    color = ['g', 'r']
    # draw trajectory w.r.t. camera
    for demo_name, se3_traj in se3_dict.items():
        depth_dir = './data/' + task_name + '/' + demo_name + '/depth'
        #mask_dir = './output/'+task_name+'/segment/'+demo_name
        mask_dir = './data/' + task_name + '/' + demo_name + '/mask'
        depth_files = sorted(glob.glob(depth_dir + '/*.npy'))
        mask_files = sorted(glob.glob(mask_dir + '/*.npy'))
        output_path = output_dir + '/' + demo_name
        util.create_dir(output_path, clear=True)

        depth = np.load(depth_files[0])
        mask = np.load(mask_files[0])

        se30 = se3_traj[0, 0, :]
        SE30 = se3_to_SE3(se30)
        T0 = SE30[0:3, 3]
        T_traj = np.expand_dims(T0, 0)
        # modify for multiple object
        for t in range(len(se3_traj)):
            ax.clear()
            for obj_idx in range(nb_object):
                se3 = se3_traj[t, obj_idx, :]
                SE3 = se3_to_SE3(se3)
                T = SE3[0:3, 3]
                T_traj = np.concatenate([T_traj, np.expand_dims(T, 0)], 0)
                ax.plot(T_traj[:, 0],
                        T_traj[:, 1],
                        T_traj[:, 2],
                        color[obj_idx],
                        alpha=0.5,
                        linewidth=4)
                obj.apply_pose(se3)
                obj.plot(ax, scale=0.05, linewidth=3)

            util.set_axes_equal(ax)
            #ax.axis('scaled')
            ax.set_xlabel('x(m)')
            ax.set_ylabel('y(m)')
            ax.set_zlabel('z(m)')
            ax.view_init(elev=-60, azim=-90)
            fig.savefig(output_path + '/%06d.png' % t)

        np.save(output_path + '/pose_traj.npy', se3_traj)
        video_path = output_path + '/object_trajectory.avi'
        util.frame_to_video(output_path, video_path, 'png', fps=fps)()
def compare(config, LOAD=False):
    '''
    vicon data is required!
    execute 'read_bag' first
    '''

    task_name = config['task_name']
    fps = config['animation']['fps']
    object_list = util.load_txt('./configure/%s_objects.txt' % task_name)
    nb_object = len(object_list)

    data_dir = './data/' + task_name
    pose_dir = './output/' + task_name + '/read_pose'
    output_dir = './output/' + task_name + '/compare_traj'
    util.create_dir(output_dir, clear=False)

    ## load data
    demo_list = os.listdir('./data/' + task_name)
    for demo in demo_list:
        rgb_demo_dir = data_dir + '/' + demo + '/rgb'
        cam_demo_dir = data_dir + '/' + demo + '/vicon/k_zed'
        output_demo_dir = output_dir + '/' + demo
        util.create_dir(output_demo_dir + '/traj', clear=True)
        util.create_dir(output_demo_dir + '/plot', clear=True)

        # g_vc1 : vicon to camera(vicon_pose) = cam pose by vicon
        # g_c1c2 : camera(vicon_pose) to camera center(center of image plane)
        # g_c2o1 : camera center to object(vision coordinates) = object pose by vision
        # g_o1o2 : object(vision coordinates) to object(vicon coordinates)
        # g_vo2_gt : vicon to object(vicon coordinates) = object pose by vicon
        assert len(object_list) == 1  ## to do : multiple object
        obj_demo_dir = data_dir + '/' + demo + '/vicon/' + object_list[
            0]  ## to do : multiple object
        se3_c2o1 = np.load(pose_dir + '/' + demo +
                           '/pose_traj.npy')[:,
                                             0, :]  ## to do : multiple object
        se3_vc1 = load_vicon(cam_demo_dir)
        se3_vo2_gt = load_vicon(obj_demo_dir)

        # make vicon orientation = [0,0,0]
        demo_len = 50  #len(se3_c2o1) #################################
        g_vo2_gt_0 = se3_to_SE3(se3_vo2_gt[0, :])
        R0 = g_vo2_gt_0[0:3, 0:3]
        T0 = np.zeros(3)
        g0_inv = RT_to_SE3(np.transpose(R0), T0)
        for t in range(demo_len):
            g_vo2_gt_t = se3_to_SE3(se3_vo2_gt[t, :])
            new_g_vo2_gt_t = np.matmul(g_vo2_gt_t, g0_inv)
            new_se3_vo2_gt_t = SE3_to_se3(new_g_vo2_gt_t)
            se3_vo2_gt[t, :] = new_se3_vo2_gt_t

        x0 = np.random.rand(12)
        #x0 = np.random.rand(6)
        optimize_len = int(len(se3_c2o1))

        def objective_fn(x0):
            # g_vc1 : vicon to camera(vicon_pose) = se3_cam
            # g_c1c2 : camera(vicon_pose) to camera center(center of image plane)
            # g_c2o1 : camera center to object(vision coordinates) = se3_vision
            # g_o1o2 : object(vision coordinates) to object(vicon coordinates)

            g_c1c2 = se3_to_SE3(x0[0:6])
            g_o1o2 = se3_to_SE3(x0[6:12])
            #g_c1c2 = se3_to_SE3(np.zeros(6))
            #g_o1o2 = se3_to_SE3(x0)

            loss = 0
            for t in range(optimize_len):
                g_c2o1_t = se3_to_SE3(se3_c2o1[t, :])
                g_vc1_t = se3_to_SE3(se3_vc1[t, :])
                g_vo2 = np.matmul(
                    g_vc1_t, np.matmul(g_c1c2, np.matmul(g_c2o1_t, g_o1o2)))
                g_vo2 = np.matmul(g_vc1_t, np.matmul(g_c2o1_t, g_o1o2))

                g_vo2_gt = se3_to_SE3(se3_vo2_gt[t, :])
                se3_vo2 = SE3_to_se3(g_vo2)
                #loss += np.sum(np.square(se3_vo2[:]-se3_vo2_gt[t,:])) #----------------
                loss += np.sum(np.square(g_vo2[0:3, 3] - g_vo2_gt[0:3, 3]))
                #loss += np.sum(np.square(g_vo2-g_vo2_gt)) #----------------------
            return loss

        print(colored('initial_loss:' + str(objective_fn(x0)), 'blue'))

        if LOAD:
            result = np.load(output_demo_dir + '/optimization.npy',
                             allow_pickle=True).item()
            cmd = input('Optimization more? [y/n]')
            if cmd == 'y':
                x0 = result.x
                result = minimize(objective_fn,
                                  x0,
                                  method='BFGS',
                                  options={'disp': True})
            else:
                pass
        else:
            #'''
            result = minimize(objective_fn,
                              x0,
                              method='BFGS',
                              options={'disp': True})
            #'''
            '''
            result = minimize(objective_fn, 
                    x0, 
                    method='BFGS', 
                    tol=1e-8,
                    options={'gtol': 1e-8, 'disp': True})
            '''

            np.save(output_demo_dir + '/optimization.npy', result)
        print(colored('optimized_loss:' + str(objective_fn(result.x)), 'blue'))
        g_c1c2 = se3_to_SE3(result.x[0:6])
        g_o1o2 = se3_to_SE3(result.x[6:12])
        #g_c1c2 = se3_to_SE3(np.zeros(6))
        #g_o1o2 = se3_to_SE3(result.x)

        obj_vicon = vision.SE3object(np.zeros(6), angle_type='axis')
        obj_vision = vision.SE3object(np.zeros(6), angle_type='axis')
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')

        loss = 0
        position_error = []
        rotation_error = []
        vicon_traj = []
        vision_traj = []
        vicon_euler = []
        vision_euler = []

        T_vo2 = np.zeros((0, 3))
        T_vo2_gt = np.zeros((0, 3))

        g_c2o1_0 = se3_to_SE3(se3_c2o1[0, :])
        g_vc1_0 = se3_to_SE3(se3_vc1[0, :])
        g_vo2_0 = np.matmul(g_vc1_0,
                            np.matmul(g_c1c2, np.matmul(g_c2o1_0, g_o1o2)))
        g_vo2_0_gt = se3_to_SE3(se3_vo2_gt[0, :])

        R_target = g_vo2_0_gt[0:3, 0:3]
        R_ori = g_vo2_0[0:3, 0:3]
        T_ori = g_vo2_0[0:3, 3]
        SO3_align = np.matmul(R_target, np.transpose(R_ori))
        SE3_align = np.matmul(inv_SE3(g_vo2_0), g_vo2_0_gt)

        for t in range(demo_len):
            g_c2o1_t = se3_to_SE3(se3_c2o1[t, :])
            g_vc1_t = se3_to_SE3(se3_vc1[t, :])
            g_vo2_t = np.matmul(g_vc1_t,
                                np.matmul(g_c1c2, np.matmul(g_c2o1_t, g_o1o2)))
            #g_vo2_t[0:3,0:3] = np.matmul(SO3_align, g_vo2_t[0:3,0:3])
            #g_vo2_t = np.matmul(g_vo2_t, SE3_align)

            se3_vo2_t_gt = SE3_to_se3(se3_to_SE3(se3_vo2_gt[t, :]))
            g_vo2_t_gt = se3_to_SE3(se3_vo2_t_gt)
            se3_vo2_t = SE3_to_se3(g_vo2_t)

            T_vo2 = np.concatenate(
                [T_vo2, np.expand_dims(g_vo2_t[0:3, 3], 0)], 0)
            T_vo2_gt = np.concatenate(
                [T_vo2_gt, np.expand_dims(g_vo2_t_gt[0:3, 3], 0)], 0)

            ax.clear()
            obj_vicon.apply_pose(se3_vo2_t)
            obj_vision.apply_pose(se3_vo2_t_gt)
            obj_vicon.plot(ax, scale=0.015, linewidth=3)
            obj_vision.plot(ax, scale=0.015, linewidth=3)
            ax.plot(T_vo2_gt[:t, 0],
                    T_vo2_gt[:t, 1],
                    T_vo2_gt[:t, 2],
                    '--',
                    color='r',
                    alpha=0.5,
                    linewidth=4)
            ax.plot(T_vo2[:, 0],
                    T_vo2[:, 1],
                    T_vo2[:, 2],
                    color='g',
                    alpha=0.5,
                    linewidth=3)

            util.set_axes_equal(ax)
            ax.set_xlabel('x(m)')
            ax.set_ylabel('y(m)')
            ax.set_zlabel('z(m)')
            fig.savefig(output_demo_dir + '/traj/%05d.png' % t)

            ## rescale (||w|| = 1)
            #se3_vo2_t[3:6] =  (1./np.sqrt(np.sum(np.square(se3_vo2_t[3:6]))))*se3_vo2_t[3:6]
            #se3_vo2_t_gt[3:6] = (1./np.sqrt(np.sum(np.square(se3_vo2_t_gt[3:6]))))*se3_vo2_t_gt[3:6]

            loss += np.sqrt(np.sum(np.square(se3_vo2_t_gt - se3_vo2_t)))
            position_error.append(
                np.expand_dims(
                    np.sqrt(
                        np.sum(np.square(g_vo2_t_gt[0:3, 3] -
                                         g_vo2_t[0:3, 3]))), 0))
            rotation_error.append(
                np.expand_dims(
                    np.sqrt(
                        np.sum(np.square(se3_vo2_t_gt[3:6] - se3_vo2_t[3:6]))),
                    0))

            vicon_traj.append(np.expand_dims(se3_vo2_t_gt, 0))
            vision_traj.append(np.expand_dims(se3_vo2_t, 0))
            vicon_euler.append(
                np.expand_dims(R_to_euler(g_vo2_t_gt[0:3, 0:3]), 0))
            vision_euler.append(
                np.expand_dims(R_to_euler(g_vo2_t[0:3, 0:3]), 0))

            if t == 0:
                total_translation = 0
                total_rotation = 0
                prev_g_vo = g_vo2_t_gt
                prev_se3_vo = se3_vo2_t_gt
            else:
                total_translation += np.sqrt(
                    np.sum(np.square(g_vo2_t_gt[0:3, 3] - prev_g_vo[0:3, 3])))
                total_rotation += np.sqrt(
                    np.sum(np.square(se3_vo2_t_gt[3:6] - prev_se3_vo[3:6])))
                prev_g_vo = g_vo2_t_gt
                prev_se3_vo = se3_vo2_t_gt
        plt.close()

        ## save loss
        loss = loss / demo_len
        position_error = np.sum(
            position_error) / demo_len  #np.concatenate(position_error,0)
        rotation_error = np.sum(
            rotation_error) / demo_len  #np.concatenate(rotation_error,0)

        vicon_traj = np.concatenate(vicon_traj, 0)
        vision_traj = np.concatenate(vision_traj, 0)
        vicon_euler = np.concatenate(vicon_euler, 0)
        vision_euler = np.concatenate(vision_euler, 0)

        np.savetxt(output_demo_dir + '/loss.txt', [loss])
        np.savetxt(output_demo_dir + '/position_error.txt', [position_error])
        np.savetxt(output_demo_dir + '/rotation_error.txt', [rotation_error])
        np.savetxt(output_demo_dir + '/total_translation.txt',
                   [total_translation])
        np.savetxt(output_demo_dir + '/total_rotation.txt', [total_rotation])
        np.savetxt(output_demo_dir + '/vicon_traj.txt', vicon_traj)
        np.savetxt(output_demo_dir + '/vision_traj.txt', vision_traj)

        ## save plot
        fig = plt.figure()
        ymins = []
        ymaxs = []
        axes = []
        scales = []
        for i in range(3):
            ax = plt.subplot(3, 1, i + 1)
            ax.plot(np.arange(demo_len),
                    vicon_traj[:, i],
                    '--',
                    color='r',
                    alpha=0.5,
                    linewidth=4)
            ax.plot(np.arange(demo_len),
                    vision_traj[:, i],
                    color='g',
                    alpha=0.5,
                    linewidth=3)
            ymin, ymax = ax.get_ylim()
            ymins.append(ymin)
            ymaxs.append(ymax)
            axes.append(ax)
            scales.append(ymax - ymin)
        ymin = min(ymins)
        ymax = max(ymaxs)
        scale_max = max(scales)
        for ax, ymin, ymax in zip(axes, ymins, ymaxs):
            center = (ymin + ymax) / 2
            ax.set_ylim([center - scale_max / 2, center + scale_max / 2])
        fig.savefig(output_demo_dir + '/v_component.png')
        plt.close()

        fig = plt.figure()
        ymins = []
        ymaxs = []
        axes = []
        scales = []
        for i in range(3):
            ax = plt.subplot(3, 1, i + 1)
            ax.plot(np.arange(demo_len),
                    vicon_traj[:, i + 3],
                    '--',
                    color='r',
                    alpha=0.5,
                    linewidth=4)
            ax.plot(np.arange(demo_len),
                    vision_traj[:, i + 3],
                    color='g',
                    alpha=0.5,
                    linewidth=3)
            ymin, ymax = ax.get_ylim()
            ymins.append(ymin)
            ymaxs.append(ymax)
            axes.append(ax)
            scales.append(ymax - ymin)
        ymin = min(ymins)
        ymax = max(ymaxs)
        scale_max = max(scales)
        for ax, ymin, ymax in zip(axes, ymins, ymaxs):
            center = (ymin + ymax) / 2
            ax.set_ylim([center - scale_max / 2, center + scale_max / 2])
        fig.savefig(output_demo_dir + '/w_component.png')
        plt.close()

        ## translation and rotation
        fig = plt.figure()
        ymins = []
        ymaxs = []
        axes = []
        scales = []
        for i in range(3):
            ax = plt.subplot(3, 1, i + 1)
            ax.plot(np.arange(demo_len),
                    T_vo2_gt[:, i],
                    '--',
                    color='r',
                    alpha=0.5,
                    linewidth=4)
            ax.plot(np.arange(demo_len),
                    T_vo2[:, i],
                    color='g',
                    alpha=0.5,
                    linewidth=3)
            ymin, ymax = ax.get_ylim()
            ymins.append(ymin)
            ymaxs.append(ymax)
            axes.append(ax)
            scales.append(ymax - ymin)
        ymin = min(ymins)
        ymax = max(ymaxs)
        scale_max = max(scales)
        for ax, ymin, ymax in zip(axes, ymins, ymaxs):
            center = (ymin + ymax) / 2
            ax.set_ylim([center - scale_max / 2, center + scale_max / 2])
        fig.savefig(output_demo_dir + '/translation_component.png')
        plt.close()

        fig = plt.figure()
        ymins = []
        ymaxs = []
        axes = []
        scales = []
        for i in range(3):
            ax = plt.subplot(3, 1, i + 1)
            ax.plot(np.arange(demo_len),
                    vicon_euler[:, i],
                    '--',
                    color='r',
                    alpha=0.5,
                    linewidth=4)
            ax.plot(np.arange(demo_len),
                    vision_euler[:, i],
                    color='g',
                    alpha=0.5,
                    linewidth=3)
            ymin, ymax = ax.get_ylim()
            ymins.append(ymin)
            ymaxs.append(ymax)
            axes.append(ax)
            scales.append(ymax - ymin)
        ymin = min(ymins)
        ymax = max(ymaxs)
        scale_max = max(scales)
        for ax, ymin, ymax in zip(axes, ymins, ymaxs):
            center = (ymin + ymax) / 2
            ax.set_ylim([center - scale_max / 2, center + scale_max / 2])
        fig.savefig(output_demo_dir + '/rotation_component.png')
        plt.close()
Exemple #3
0
            obj_color[obj] = colors[i]

        for t in range(data_len):
            ax.clear()
            for obj in object_names:
                se3_object = se3_objects[obj]
                
                se3 = se3_dict[obj][t,:]
                SE3 = se3_to_SE3(se3)
                T = SE3[0:3,3]
                se3_traj[obj] = np.concatenate([se3_traj[obj], np.expand_dims(T,0)],0)
                
                ax.plot(se3_traj[obj][:,0],se3_traj[obj][:,1],se3_traj[obj][:,2], 
                        obj_color[obj], alpha = 0.5, linewidth = 4)
                se3_object.apply_pose(se3)
                se3_object.plot(ax, scale = 0.05, linewidth = 3) 

            util.set_axes_equal(ax)
            #ax.axis('scaled')
            ax.set_xlabel('x(m)')
            ax.set_ylabel('y(m)')
            ax.set_zlabel('z(m)')
            ax.view_init(elev = -60, azim = -90)
            fig.savefig(output_dir+'/%s.png'%str(t).zfill(10))
        video_path = output_dir+'/object_trajectory.avi'
        util.frame_to_video(output_dir, video_path, 'png', fps=fps)()




Exemple #4
0
def compare(config):
    '''
    vicon data is required!
    execute 'read_bag' first
    '''

    task_name = config['data_name']
    nb_object = config['object']['nb_object']
    supervision = config['pose']['supervision']
    fps = config['animation']['fps']

    data_dir = './data/' + task_name
    pose_dir = './output/read_pose/' + task_name
    vicon_dir = './output/vicon/' + task_name
    output_dir = './output/compare_traj/' + task_name

    util.create_dir(output_dir, clear=True)

    demo_list = os.listdir('./data/' + task_name)
    for demo in demo_list:
        data_demo_dir = data_dir + '/' + demo
        vicon_demo_dir = vicon_dir + '/' + demo
        cam_pos0 = np.load(data_demo_dir + '/camera_position0.npy',
                           allow_pickle=True)
        cam_pos1 = np.load(data_demo_dir + '/camera_position1.npy',
                           allow_pickle=True)
        robot_pos0 = np.load(data_demo_dir + '/pioneer_position0.npy',
                             allow_pickle=True)
        robot_pos1 = np.load(data_demo_dir + '/pioneer_position1.npy',
                             allow_pickle=True)
        cut_point = np.loadtxt(data_demo_dir + '/cut.txt')

        util.create_dir(output_dir + '/' + demo + '/plot', clear=True)
        util.create_dir(output_dir + '/' + demo + '/pickle', clear=True)
        ## open vicon trajectory
        vision_traj_list = np.load(pose_dir + '/' + demo + '/pose_traj.npy')
        obj_files = sorted(glob.glob(vicon_demo_dir + '/*_pose.npy'))
        assert len(obj_files) == 1
        vicon_traj_list = []
        for obj_file in obj_files:
            vicon_traj_loaded = np.load(obj_file, allow_pickle=True)
            #for i in range(1,7):
            #    vicon_traj_loaded[:,i] = medfilt(vicon_traj_loaded[:,i],11)

            vicon_traj_list.append(vicon_traj_loaded)

        ## open cam time
        cam_start_time = cam_pos0[0]
        cam_finish_time = cam_pos1[0]
        ## cut point
        total_img = cut_point[0]
        beginning_cut = cut_point[1]
        endding_cut = cut_point[2]
        ## modify cam start time
        play_time = cam_finish_time - cam_start_time
        cam_start_time = cam_start_time + play_time * (beginning_cut /
                                                       (total_img - 1))
        cam_finish_time = cam_finish_time - play_time * (endding_cut /
                                                         (total_img - 1))

        for i in range(nb_object):
            vicon_traj = vicon_traj_list[i]
            vision_traj = vision_traj_list[:, i, :]
            vicon_time = vicon_traj[:, 0]

            vicon_start_idx = 0
            vicon_finish_idx = len(vicon_time)

            for i in range(len(vicon_time)):
                if vicon_time[i] < cam_start_time:
                    vicon_start_idx = i
                if vicon_time[i] < cam_finish_time:
                    vicon_finish_idx = i
            vicon_start_idx += 1

            clipped_vicon_time = vicon_time[vicon_start_idx:vicon_finish_idx]
            clipped_vicon = vicon_traj[vicon_start_idx:vicon_finish_idx, 1:7]

            ### align time
            len_vision = len(vision_traj)
            vision_time = np.arange(
                clipped_vicon_time[0], clipped_vicon_time[-1],
                (clipped_vicon_time[-1] - clipped_vicon_time[0]) / len_vision)

            fig = plt.figure()
            ax = fig.add_subplot(111, projection='3d')
            vicon_plot = np.zeros((0, 3))
            vicon_se3 = np.zeros((0, 6))
            vision_plot = np.zeros((0, 3))
            vision_se3 = np.zeros((0, 6))

            obj_vicon = vision.SE3object(np.zeros(6), angle_type='axis')
            obj_vision = vision.SE3object(np.zeros(6), angle_type='axis')
            v_t = 1

            for t in range(len(vision_traj)):  # ah
                for tt in range(v_t - 1, len(clipped_vicon_time)):
                    if clipped_vicon_time[tt] > vision_time[t]:
                        break
                v_t = tt
                #print('vision_time:%04.04f'%vision_time[t])
                #print('vicon_time:%04.04f'%clipped_vicon_time[v_t])
                #print('\n')

                ## time
                vicon_T = se3_to_SE3(clipped_vicon[v_t, :])
                vision_T = se3_to_SE3(vision_traj[t, :])
                vicon_plot = np.concatenate(
                    [vicon_plot,
                     np.expand_dims(vicon_T[0:3, 3], 0)], 0)
                vision_plot = np.concatenate(
                    [vision_plot,
                     np.expand_dims(vision_T[0:3, 3], 0)], 0)

                vicon_se3_element = SE3_to_se3(
                    se3_to_SE3(clipped_vicon[v_t, :]))
                vicon_se3 = np.concatenate(
                    [vicon_se3,
                     np.expand_dims(vicon_se3_element, 0)], 0)
                vision_se3 = np.concatenate(
                    [vision_se3,
                     np.expand_dims(SE3_to_se3(vision_T), 0)], 0)

            plt.close('all')
            ## align translation
            x0 = np.random.rand(12)
            optimize_len = int(len(vision_plot))

            ########################## for task3 (occlusion)
            #if task_name == 'task3':
            #    optimize_len = 100
            def objective_fn(x0):
                SE30 = se3_to_SE3(x0[0:6])
                SE31 = se3_to_SE3(x0[6:12])
                loss = 0
                for t in range(optimize_len):
                    transformed = np.matmul(
                        np.matmul(SE30, se3_to_SE3(vision_se3[t, :])), SE31)
                    transformed_se3 = SE3_to_se3(transformed)
                    loss += np.sum(
                        np.square(transformed_se3[:] - vicon_se3[t, :]))

                    #transformed = un_homo( np.matmul(np.matmul(SE30, to_homo(vision_plot[t,:])),SE31))
                    #loss += np.sum(np.square(transformed-vicon_plot[t,:]))
                return loss

            print(demo)
            print('initial_loss:' + str(objective_fn(x0)))
            if False:
                result = np.load(output_dir + '/' + demo +
                                 '/optimization.npy').item()
            else:
                #'''
                result = minimize(objective_fn,
                                  x0,
                                  method='BFGS',
                                  tol=1e-7,
                                  options={
                                      'gtol': 1e-6,
                                      'disp': True
                                  })

                #'''
                '''
                result = minimize(objective_fn, 
                                    x0, 
                                    method='Nelder-Mead',
                                    options={'maxiter':1})
                '''
                np.save(output_dir + '/' + demo + '/optimization.npy', result)
            print('optimized_loss:' + str(objective_fn(result.x)))
            print(result.x)

            SE30 = se3_to_SE3(result.x[0:6])
            SE31 = se3_to_SE3(result.x[6:12])
            ### align orientation
            #'''

            vision_SE30 = np.matmul(
                np.matmul(SE30, se3_to_SE3(vision_se3[0, :])), SE31)
            vicon_SE30 = se3_to_SE3(vicon_se3[0, :])
            R_target = vicon_SE30[0:3, 0:3]
            R_ori = vision_SE30[0:3, 0:3]
            T_ori = vision_SE30[0:3, 3]

            SO3_align = np.matmul(R_target, np.transpose(R_ori))

            #####################################################
            fig = plt.figure()
            ax = fig.add_subplot(111, projection='3d')
            transformed_vision_se3 = np.zeros((0, 6))
            transformed_vision_SE3 = np.zeros((0, 4, 4))
            transformed_vision_plot = np.zeros((0, 3))

            loss = 0
            position_error = []
            rotation_error = []
            total_position = []
            total_rotation = []
            for t in range(len(vicon_se3)):
                ## time
                vicon_se3_t = vicon_se3[t, :]
                vision_se3_t = vision_se3[t, :]
                vicon_T_t = se3_to_SE3(vicon_se3[t, :])
                vision_T_t = se3_to_SE3(vision_se3[t, :])
                ############################### not alignment!!!!!!!!!

                #if supervision == 'never':
                #    se3_rc = np.asarray([0.333653152, -0.19545771, -0.41434446, -0.16426212, 0.4854613, -0.28981152],dtype = np.float32)
                #    g_rc = se3_to_SE3(se3_rc)
                #    _, g_vr = util.load_vicon(data_demo_dir+'/camera_position0.npy')
                #    SE3 = np.matmul(g_vr, g_rc)

                vision_T_t = np.matmul(np.matmul(SE30, vision_T_t), SE31)
                #vision_T_t[0:3,0:3] = np.matmul(SO3_align, vision_T_t[0:3,0:3])
                vision_se3_t = SE3_to_se3(vision_T_t)
                #IPython.embed()

                transformed_vision_se3 = np.concatenate([
                    transformed_vision_se3,
                    np.expand_dims(SE3_to_se3(vision_T_t), 0)
                ], 0)
                transformed_vision_plot = np.concatenate([
                    transformed_vision_plot,
                    np.expand_dims(vision_T_t[0:3, 3], 0)
                ], 0)
                transformed_vision_SE3 = np.concatenate(
                    [transformed_vision_SE3,
                     np.expand_dims(vision_T_t, 0)], 0)

                loss += np.sqrt(
                    np.sum(
                        np.square(
                            SE3_to_se3(vicon_T_t) - SE3_to_se3(vision_T_t))))
                position_error.append(
                    np.expand_dims(
                        np.sqrt(
                            np.sum(
                                np.square(vicon_T_t[0:3, 3] -
                                          vision_T_t[0:3, 3]))), 0))
                rotation_error.append(
                    np.expand_dims(
                        np.sqrt(
                            np.sum(
                                np.square(vicon_se3_t[3:6] -
                                          vision_se3_t[3:6]))), 0))
                total_position.append(np.expand_dims(vicon_T_t[0:3, 3], 0))
                total_rotation.append(np.expand_dims(vicon_se3_t[3:6], 0))

                ax.clear()
                obj_vicon.apply_pose(vicon_se3_t)
                obj_vision.apply_pose(vision_se3_t)
                obj_vicon.plot(ax, scale=0.015, linewidth=3)
                obj_vision.plot(ax, scale=0.015, linewidth=3)
                ##
                ax.plot(vicon_plot[:t, 0],
                        vicon_plot[:t, 1],
                        vicon_plot[:t, 2],
                        '--',
                        color='r',
                        alpha=0.5,
                        linewidth=4)
                ax.plot(transformed_vision_plot[:, 0],
                        transformed_vision_plot[:, 1],
                        transformed_vision_plot[:, 2],
                        color='g',
                        alpha=0.5,
                        linewidth=3)
                ###
                util.set_axes_equal(ax)
                ax.set_xlabel('x(m)')
                ax.set_ylabel('y(m)')
                ax.set_zlabel('z(m)')
                fig.savefig(output_dir + '/' + demo + '/%05d.png' % t)
            plt.close()

            loss = loss / len(vision_traj)
            position_error = np.sum(position_error) / len(
                vision_traj)  #np.concatenate(position_error,0)
            rotation_error = np.sum(rotation_error) / len(
                vision_traj)  #np.concatenate(rotation_error,0)

            total_position = np.concatenate(total_position, 0)
            total_rotation = np.concatenate(total_rotation, 0)

            np.savetxt(output_dir + '/' + demo + '/loss.txt', [loss])
            np.savetxt(output_dir + '/' + demo + '/position_error.txt',
                       [position_error])
            np.savetxt(output_dir + '/' + demo + '/rotation_error.txt',
                       [rotation_error])
            np.savetxt(output_dir + '/' + demo + '/total_position.txt',
                       total_position)
            np.savetxt(output_dir + '/' + demo + '/total_rotation.txt',
                       total_rotation)

            #################### save data for jigang
            pickle_dict = {
                'time': np.transpose(vision_time),
                'se3': transformed_vision_se3,
                'SE3': transformed_vision_SE3
            }
            io.savemat(output_dir + '/' + demo + '/pickle/trajectories.mat',
                       pickle_dict)

            #####################################################
            if len(vision_time) > len(transformed_vision_se3):
                vision_time = vision_time[:len(transformed_vision_se3)]

            fig = plt.figure()
            ax1 = fig.add_subplot(311)
            ax1.plot(clipped_vicon_time, clipped_vicon[:, 0], '--')
            ax1.plot(vision_time, transformed_vision_se3[:, 0])
            ax1.set_title('se3[0:3]')

            ax2 = fig.add_subplot(312)
            ax2.plot(clipped_vicon_time, clipped_vicon[:, 1], '--')
            ax2.plot(vision_time, transformed_vision_se3[:, 1])

            ax3 = fig.add_subplot(313)
            ax3.plot(clipped_vicon_time, clipped_vicon[:, 2], '--')
            ax3.plot(vision_time, transformed_vision_se3[:, 2])
            fig.savefig(output_dir + '/' + demo + '/plot/xyz.png')
            #####################################################
            fig = plt.figure()
            ax1 = fig.add_subplot(311)
            ax1.plot(clipped_vicon_time, clipped_vicon[:, 3], '--')
            ax1.plot(vision_time, transformed_vision_se3[:, 3])
            ax1.set_title('se3[3:6]')

            ax2 = fig.add_subplot(312)
            ax2.plot(clipped_vicon_time, clipped_vicon[:, 4], '--')
            ax2.plot(vision_time, transformed_vision_se3[:, 4])

            ax3 = fig.add_subplot(313)
            ax3.plot(clipped_vicon_time, clipped_vicon[:, 5], '--')
            ax3.plot(vision_time, transformed_vision_se3[:, 5])
            fig.savefig(output_dir + '/' + demo + '/plot/wxwywz.png')
            plt.close()
            #####################################################
            diff_vicon = (
                clipped_vicon[1:, :] - clipped_vicon[:-1, :]) / np.expand_dims(
                    clipped_vicon_time[1:] - clipped_vicon_time[:-1], 1)
            diff_vision = (transformed_vision_se3[1:, :] -
                           transformed_vision_se3[:-1, :]) / np.expand_dims(
                               vision_time[1:] - vision_time[:-1], 1)

            fig = plt.figure()
            ax1 = fig.add_subplot(311)
            ax1.plot(clipped_vicon_time[:-1], diff_vicon[:, 0], '--')
            ax1.plot(vision_time[:-1], diff_vision[:, 0])
            ax1.set_title('diff_se3[0:3]')
            ax1.set_ylim([-0.5, 0.5])

            ax2 = fig.add_subplot(312)
            ax2.plot(clipped_vicon_time[:-1], diff_vicon[:, 1], '--')
            ax2.plot(vision_time[:-1], diff_vision[:, 1])
            ax2.set_ylim([-0.5, 0.5])

            ax3 = fig.add_subplot(313)
            ax3.plot(clipped_vicon_time[:-1], diff_vicon[:, 2], '--')
            ax3.plot(vision_time[:-1], diff_vision[:, 2])
            ax3.set_ylim([-0.5, 0.5])

            fig.savefig(output_dir + '/' + demo + '/plot/diff_xyz.png')
            plt.close()
            #####################################################
            fig = plt.figure()
            ax1 = fig.add_subplot(311)
            ax1.plot(clipped_vicon_time[:-1], diff_vicon[:, 3], '--')
            ax1.plot(vision_time[:-1], diff_vision[:, 3])
            ax1.set_title('diff_se3[0:3]')
            ax1.set_ylim([-0.5, 0.5])

            ax2 = fig.add_subplot(312)
            ax2.plot(clipped_vicon_time[:-1], diff_vicon[:, 4], '--')
            ax2.plot(vision_time[:-1], diff_vision[:, 4])
            ax2.set_ylim([-0.5, 0.5])

            ax3 = fig.add_subplot(313)
            ax3.plot(clipped_vicon_time[:-1], diff_vicon[:, 5], '--')
            ax3.plot(vision_time[:-1], diff_vision[:, 5])
            ax3.set_ylim([-0.5, 0.5])

            fig.savefig(output_dir + '/' + demo + '/plot/diff_wxwywz.png')
            plt.close()
            #####################################################

        video_path = output_dir + '/' + demo + '/se3_compare.avi'
        util.frame_to_video(output_dir + '/' + demo,
                            video_path,
                            'png',
                            fps=fps)()