<filename>utils/process_data.py<gh_stars>0
import os
import pdb
import h5py
import pickle
import numpy as np
from scipy.io import loadmat
import cv2
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw
import csv
import bisect
import matplotlib as mpl
import matplotlib.cm as cm
import tensorflow as tf

# from bilateral_filter import bilateral_filter
from tools import *

def resave_imu_data():
    dataset = '/freezer/nyudepthV2_raw'
    seqs = os.listdir(dataset)
    for seq in seqs:
        seq_dir = dataset + '/' + seq
        for data in os.listdir(seq_dir):
            if data[0] == 'a':
                imu_data_path = seq_dir + '/' + data
                resave_imu_data_path = seq_dir + '/' + data[:-4] + '.txt'
                call_resave_imu(imu_data_path, resave_imu_data_path)


def call_resave_imu(orig_path, resave_path):
    command = './resave_imu ' + orig_path + ' ' + resave_path
    os.system(command)


def collect_acc_data(folder):
    data_list = []
    for file in os.listdir(folder):
        if file[0] == 'a' and file[-1] == 't':
            data_list.append(folder + '/' + file)

    return sorted(data_list)


def get_acc_timestamp(path):
    return float(path.split('-')[1])


def read_acc_data(file_path):
    timestamp = get_acc_timestamp(file_path)
    file = open(file_path, 'r')
    data = file.read().split(',')
    for i in range(len(data)):
        data[i] = float(data[i])

    data.insert(0, timestamp)
    return data


def plot_acc_data(folder):
    acc_path = collect_acc_data(folder)
    x = []
    y = []
    z = []
    for path in acc_path:
        data = read_acc_data(path)
        x.append(data[1])
        y.append(data[2])
        z.append(data[3])
    plt.plot(x)
    plt.plot(y)
    plt.plot(z)
    plt.show()


def compute_acc_vel_pos(acc_data_1, acc_data_2, v_1, p_1):
    t1 = acc_data_1[0]
    t2 = acc_data_2[0]
    t_delta = t2 - t1

    acc_xyz_1 = np.array(acc_data_1[1:4])
    acc_xyz_2 = np.array(acc_data_2[1:4])
    a_avg = (acc_xyz_1 + acc_xyz_2) / 2.
    v_2 = v_1 + a_avg * t_delta
    p_2 = p_1 + v_1 * t_delta + a_avg * t_delta * t_delta / 2.

    # pdb.set_trace()
    return v_2, p_2


def plot_imu_traj(folder):
    acc_path = collect_acc_data(folder)

    p_x = []
    p_y = []
    p_z = []

    v_cur = np.array([0., 0., 0.])
    p_cur = np.array([0., 0., 0.])
    N = len(acc_path)
    for idx in range(N-1):
        p_x.append(p_cur[0])
        p_y.append(p_cur[1])
        p_z.append(p_cur[2])
        acc_1 = read_acc_data(acc_path[idx])
        acc_2 = read_acc_data(acc_path[idx + 1])
        v_cur, p_cur = compute_acc_vel_pos(acc_1, acc_2, v_cur, p_cur)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    #ax.scatter(p_x[:], p_y[:], p_z[:0])
    ax.plot(p_x[:-1], p_y[:-1], p_z[:-1])

    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')

    plt.show()

    # plt.plot(p_x)
    # plt.plot(p_y)
    # plt.plot(p_z)
    # plt.show()

def plot_trajectory(data_file_name):
    data = open(data_file_name,"rb")
    poses_log = pickle.load(data)

    poses_mat_log = []
    import torch
    for i in range(len(poses_log.keys())):
        pose = poses_log[i]
        pose = np.expand_dims(pose, axis=0)
        pose = np.expand_dims(pose, axis=0)
        pose_mat = transformation_from_parameters(torch.tensor(pose[:, :, :3]).float(), torch.tensor(pose[:, :, 3:]).float(), False)
        poses_mat_log.append(pose_mat.numpy())

    xyzs = np.array(dump_xyz(poses_mat_log))

    xs = []
    ys = []
    zs = []
    for i in range(xyzs.shape[0]):
        xs.append(xyzs[i][0])
        ys.append(xyzs[i][1])
        zs.append(xyzs[i][2])

        plt.plot(xs, ys)
        plt.savefig(
            '/home/jiatian/dataset/rs_eval/pose/' + str(i).zfill(6) + '.jpg')

def delete_folder(folder, num):
    dirlist = sorted(os.listdir(folder))
    num_total = len(dirlist)
    for idx in range(num):
        datapath = folder + '/' + dirlist[idx]
        os.remove(datapath)
    for idx in range(num_total - num, num_total):
        datapath = folder + '/' + dirlist[idx]
        os.remove(datapath)

def vis_rgb_pose_image(folder_1, folder_2):
    dirlist_1 = sorted(os.listdir(folder_1))
    dirlist_2 = sorted(os.listdir(folder_2))
    for idx in range(len(dirlist_1)):
        data_1 = Image.open(folder_1 + '/' + dirlist_1[idx]).convert('RGB')
        data_2 = Image.open(folder_2 + '/' + dirlist_2[idx]).convert('RGB')

        image_show = np.hstack((np.array(data_1), np.array(data_2)))

        img = Image.fromarray(image_show)
        img.save('/home/jiatian/dataset/rs_eval/comp/' + dirlist_2[idx])

def read_csv(path, folder):
    dict_gt = {}
    ts = []
    with open(path, mode='r') as csv_file:
        csv_reader = csv.DictReader(csv_file)
        line_count = 0
        for row in csv_reader:
            if line_count == 0:
                print(f'Column names are {", ".join(row)}')
                line_count += 1
            line_count += 1
            dict_gt[row['#timestamp']] = row
            ts.append(row['#timestamp'])
        print(f'Processed {line_count} lines.')

    dict_img = {}
    dirlist = sorted(os.listdir(folder))
    match = 0
    unmatch = 0
    ts_valid = []
    for img_path in dirlist:
        key = img_path[:-4]
        ts_valid.append(key)
        try:
            print(dict_gt[key])
            match += 1
            dict_img[key] = dict_gt[key]
        except:
            unmatch += 1
            idx = bisect.bisect_left(ts, key)
            if idx == len(ts):
                idx -= 1
            dict_img[key] = dict_gt[ts[idx]]
            print('ERROR', key)
    print('UNMATCH', unmatch)

    # plot_xyz(ts_valid, dict_img)
    calAxisAngle(ts_valid, dict_img)

def calAxisAngle(ts_valid, dict_img):
    nums = len(ts_valid)
    dict_pose = {}
    for i in range(nums-2):
        ts_cur = ts_valid[i]
        ts_next = ts_valid[i+2]
        data_cur = dict_img[ts_cur]
        data_next = dict_img[ts_next]
        diff_tx = float(data_next[' p_RS_R_x [m]']) - float(data_cur[' p_RS_R_x [m]'])
        diff_ty = float(data_next[' p_RS_R_y [m]']) - float(data_cur[' p_RS_R_y [m]'])
        diff_tz = float(data_next[' p_RS_R_z [m]']) - float(data_cur[' p_RS_R_z [m]'])
        diff_qw = float(data_next[' q_RS_w []']) - float(data_cur[' q_RS_w []'])
        diff_qx = float(data_next[' q_RS_x []']) - float(data_cur[' q_RS_x []'])
        diff_qy = float(data_next[' q_RS_y []']) - float(data_cur[' q_RS_y []'])
        diff_qz = float(data_next[' q_RS_z []']) - float(data_cur[' q_RS_z []'])
        diff_norm = np.linalg.norm([diff_qw, diff_qx, diff_qy, diff_qz])
        diff_qw = diff_qw / (diff_norm + 1e-7)
        diff_qx = diff_qx / (diff_norm + 1e-7)
        diff_qy = diff_qy / (diff_norm + 1e-7)
        diff_qz = diff_qz / (diff_norm + 1e-7)
        angle = 2 * np.arccos(diff_qw)
        rx = diff_qx * angle / np.sqrt(1 - diff_qw * diff_qw)
        ry = diff_qy * angle / np.sqrt(1 - diff_qw * diff_qw)
        rz = diff_qz * angle / np.sqrt(1 - diff_qw * diff_qw)
        dict_pose[ts_cur] = [rx, ry, rz, diff_tx, diff_ty, diff_tz]

    dict_file_name = '/home/jiatian/dataset/euroc/V1_01_easy/mav0/cam0_pose.pkl'
    f = open(dict_file_name, "wb")
    pickle.dump(dict_pose, f)
    f.close()

def plot_xyz(ts, dict):
    xs = []
    ys = []
    zs = []
    nums = len(ts)
    for i in range(nums):
        time = ts[i]
        xs.append(float(dict[time][' p_RS_R_x [m]']))
        ys.append(float(dict[time][' p_RS_R_y [m]']))
        zs.append(float(dict[time][' p_RS_R_z [m]']))
        # ' q_RS_w []' ' q_RS_x []' ' q_RS_y []' ' q_RS_z []'

        plt.plot(xs, ys)
        plt.savefig(
            '/home/jiatian/dataset/euroc/V1_01_easy/mav0/cam0/traj' + str(i).zfill(6) + '.jpg')

if __name__ == "__main__":
    read_csv('/home/jiatian/dataset/euroc/V1_01_easy/mav0/state_groundtruth_estimate0/data.csv', '/home/jiatian/dataset/euroc/V1_01_easy/mav0/cam0/data')
    # vis_rgb_pose_image('/home/jiatian/dataset/recordvi/402-000-undistort', '/home/jiatian/dataset/rs_eval/pose')
    # plot_trajectory('/home/jiatian/dataset/rs_eval/pose/poses_log.pickle')
    # delete_folder('/home/jiatian/dataset/realsense/recordvi-4-02-00/402-004-undistort', 300)
<filename>utils/process_data.py<gh_stars>0
import os
import pdb
import h5py
import pickle
import numpy as np
from scipy.io import loadmat
import cv2
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw
import csv
import bisect
import matplotlib as mpl
import matplotlib.cm as cm
import tensorflow as tf

# from bilateral_filter import bilateral_filter
from tools import *

def resave_imu_data():
    dataset = '/freezer/nyudepthV2_raw'
    seqs = os.listdir(dataset)
    for seq in seqs:
        seq_dir = dataset + '/' + seq
        for data in os.listdir(seq_dir):
            if data[0] == 'a':
                imu_data_path = seq_dir + '/' + data
                resave_imu_data_path = seq_dir + '/' + data[:-4] + '.txt'
                call_resave_imu(imu_data_path, resave_imu_data_path)


def call_resave_imu(orig_path, resave_path):
    command = './resave_imu ' + orig_path + ' ' + resave_path
    os.system(command)


def collect_acc_data(folder):
    data_list = []
    for file in os.listdir(folder):
        if file[0] == 'a' and file[-1] == 't':
            data_list.append(folder + '/' + file)

    return sorted(data_list)


def get_acc_timestamp(path):
    return float(path.split('-')[1])


def read_acc_data(file_path):
    timestamp = get_acc_timestamp(file_path)
    file = open(file_path, 'r')
    data = file.read().split(',')
    for i in range(len(data)):
        data[i] = float(data[i])

    data.insert(0, timestamp)
    return data


def plot_acc_data(folder):
    acc_path = collect_acc_data(folder)
    x = []
    y = []
    z = []
    for path in acc_path:
        data = read_acc_data(path)
        x.append(data[1])
        y.append(data[2])
        z.append(data[3])
    plt.plot(x)
    plt.plot(y)
    plt.plot(z)
    plt.show()


def compute_acc_vel_pos(acc_data_1, acc_data_2, v_1, p_1):
    t1 = acc_data_1[0]
    t2 = acc_data_2[0]
    t_delta = t2 - t1

    acc_xyz_1 = np.array(acc_data_1[1:4])
    acc_xyz_2 = np.array(acc_data_2[1:4])
    a_avg = (acc_xyz_1 + acc_xyz_2) / 2.
    v_2 = v_1 + a_avg * t_delta
    p_2 = p_1 + v_1 * t_delta + a_avg * t_delta * t_delta / 2.

    # pdb.set_trace()
    return v_2, p_2


def plot_imu_traj(folder):
    acc_path = collect_acc_data(folder)

    p_x = []
    p_y = []
    p_z = []

    v_cur = np.array([0., 0., 0.])
    p_cur = np.array([0., 0., 0.])
    N = len(acc_path)
    for idx in range(N-1):
        p_x.append(p_cur[0])
        p_y.append(p_cur[1])
        p_z.append(p_cur[2])
        acc_1 = read_acc_data(acc_path[idx])
        acc_2 = read_acc_data(acc_path[idx + 1])
        v_cur, p_cur = compute_acc_vel_pos(acc_1, acc_2, v_cur, p_cur)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    #ax.scatter(p_x[:], p_y[:], p_z[:0])
    ax.plot(p_x[:-1], p_y[:-1], p_z[:-1])

    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')

    plt.show()

    # plt.plot(p_x)
    # plt.plot(p_y)
    # plt.plot(p_z)
    # plt.show()

def plot_trajectory(data_file_name):
    data = open(data_file_name,"rb")
    poses_log = pickle.load(data)

    poses_mat_log = []
    import torch
    for i in range(len(poses_log.keys())):
        pose = poses_log[i]
        pose = np.expand_dims(pose, axis=0)
        pose = np.expand_dims(pose, axis=0)
        pose_mat = transformation_from_parameters(torch.tensor(pose[:, :, :3]).float(), torch.tensor(pose[:, :, 3:]).float(), False)
        poses_mat_log.append(pose_mat.numpy())

    xyzs = np.array(dump_xyz(poses_mat_log))

    xs = []
    ys = []
    zs = []
    for i in range(xyzs.shape[0]):
        xs.append(xyzs[i][0])
        ys.append(xyzs[i][1])
        zs.append(xyzs[i][2])

        plt.plot(xs, ys)
        plt.savefig(
            '/home/jiatian/dataset/rs_eval/pose/' + str(i).zfill(6) + '.jpg')

def delete_folder(folder, num):
    dirlist = sorted(os.listdir(folder))
    num_total = len(dirlist)
    for idx in range(num):
        datapath = folder + '/' + dirlist[idx]
        os.remove(datapath)
    for idx in range(num_total - num, num_total):
        datapath = folder + '/' + dirlist[idx]
        os.remove(datapath)

def vis_rgb_pose_image(folder_1, folder_2):
    dirlist_1 = sorted(os.listdir(folder_1))
    dirlist_2 = sorted(os.listdir(folder_2))
    for idx in range(len(dirlist_1)):
        data_1 = Image.open(folder_1 + '/' + dirlist_1[idx]).convert('RGB')
        data_2 = Image.open(folder_2 + '/' + dirlist_2[idx]).convert('RGB')

        image_show = np.hstack((np.array(data_1), np.array(data_2)))

        img = Image.fromarray(image_show)
        img.save('/home/jiatian/dataset/rs_eval/comp/' + dirlist_2[idx])

def read_csv(path, folder):
    dict_gt = {}
    ts = []
    with open(path, mode='r') as csv_file:
        csv_reader = csv.DictReader(csv_file)
        line_count = 0
        for row in csv_reader:
            if line_count == 0:
                print(f'Column names are {", ".join(row)}')
                line_count += 1
            line_count += 1
            dict_gt[row['#timestamp']] = row
            ts.append(row['#timestamp'])
        print(f'Processed {line_count} lines.')

    dict_img = {}
    dirlist = sorted(os.listdir(folder))
    match = 0
    unmatch = 0
    ts_valid = []
    for img_path in dirlist:
        key = img_path[:-4]
        ts_valid.append(key)
        try:
            print(dict_gt[key])
            match += 1
            dict_img[key] = dict_gt[key]
        except:
            unmatch += 1
            idx = bisect.bisect_left(ts, key)
            if idx == len(ts):
                idx -= 1
            dict_img[key] = dict_gt[ts[idx]]
            print('ERROR', key)
    print('UNMATCH', unmatch)

    # plot_xyz(ts_valid, dict_img)
    calAxisAngle(ts_valid, dict_img)

def calAxisAngle(ts_valid, dict_img):
    nums = len(ts_valid)
    dict_pose = {}
    for i in range(nums-2):
        ts_cur = ts_valid[i]
        ts_next = ts_valid[i+2]
        data_cur = dict_img[ts_cur]
        data_next = dict_img[ts_next]
        diff_tx = float(data_next[' p_RS_R_x [m]']) - float(data_cur[' p_RS_R_x [m]'])
        diff_ty = float(data_next[' p_RS_R_y [m]']) - float(data_cur[' p_RS_R_y [m]'])
        diff_tz = float(data_next[' p_RS_R_z [m]']) - float(data_cur[' p_RS_R_z [m]'])
        diff_qw = float(data_next[' q_RS_w []']) - float(data_cur[' q_RS_w []'])
        diff_qx = float(data_next[' q_RS_x []']) - float(data_cur[' q_RS_x []'])
        diff_qy = float(data_next[' q_RS_y []']) - float(data_cur[' q_RS_y []'])
        diff_qz = float(data_next[' q_RS_z []']) - float(data_cur[' q_RS_z []'])
        diff_norm = np.linalg.norm([diff_qw, diff_qx, diff_qy, diff_qz])
        diff_qw = diff_qw / (diff_norm + 1e-7)
        diff_qx = diff_qx / (diff_norm + 1e-7)
        diff_qy = diff_qy / (diff_norm + 1e-7)
        diff_qz = diff_qz / (diff_norm + 1e-7)
        angle = 2 * np.arccos(diff_qw)
        rx = diff_qx * angle / np.sqrt(1 - diff_qw * diff_qw)
        ry = diff_qy * angle / np.sqrt(1 - diff_qw * diff_qw)
        rz = diff_qz * angle / np.sqrt(1 - diff_qw * diff_qw)
        dict_pose[ts_cur] = [rx, ry, rz, diff_tx, diff_ty, diff_tz]

    dict_file_name = '/home/jiatian/dataset/euroc/V1_01_easy/mav0/cam0_pose.pkl'
    f = open(dict_file_name, "wb")
    pickle.dump(dict_pose, f)
    f.close()

def plot_xyz(ts, dict):
    xs = []
    ys = []
    zs = []
    nums = len(ts)
    for i in range(nums):
        time = ts[i]
        xs.append(float(dict[time][' p_RS_R_x [m]']))
        ys.append(float(dict[time][' p_RS_R_y [m]']))
        zs.append(float(dict[time][' p_RS_R_z [m]']))
        # ' q_RS_w []' ' q_RS_x []' ' q_RS_y []' ' q_RS_z []'

        plt.plot(xs, ys)
        plt.savefig(
            '/home/jiatian/dataset/euroc/V1_01_easy/mav0/cam0/traj' + str(i).zfill(6) + '.jpg')

if __name__ == "__main__":
    read_csv('/home/jiatian/dataset/euroc/V1_01_easy/mav0/state_groundtruth_estimate0/data.csv', '/home/jiatian/dataset/euroc/V1_01_easy/mav0/cam0/data')
    # vis_rgb_pose_image('/home/jiatian/dataset/recordvi/402-000-undistort', '/home/jiatian/dataset/rs_eval/pose')
    # plot_trajectory('/home/jiatian/dataset/rs_eval/pose/poses_log.pickle')
    # delete_folder('/home/jiatian/dataset/realsense/recordvi-4-02-00/402-004-undistort', 300)
# from bilateral_filter import bilateral_filter # pdb.set_trace() #ax.scatter(p_x[:], p_y[:], p_z[:0]) # plt.plot(p_x) # plt.plot(p_y) # plt.plot(p_z) # plt.show() # plot_xyz(ts_valid, dict_img) # ' q_RS_w []' ' q_RS_x []' ' q_RS_y []' ' q_RS_z []' # vis_rgb_pose_image('/home/jiatian/dataset/recordvi/402-000-undistort', '/home/jiatian/dataset/rs_eval/pose') # plot_trajectory('/home/jiatian/dataset/rs_eval/pose/poses_log.pickle') # delete_folder('/home/jiatian/dataset/realsense/recordvi-4-02-00/402-004-undistort', 300)