diff --git a/json/pose_correction.json b/json/pose_correction.json index 97ab18a0..b7bb8127 100644 --- a/json/pose_correction.json +++ b/json/pose_correction.json @@ -1,10 +1,21 @@ { + "0": { + "start_arr": [ + 0 + ], + "end_arr": [ + -1 + ], + "yaw_arr": [ + -1.2 + ] + }, "1": { "start_arr": [ - 1130 + 0 ], "end_arr": [ - 7300 + -1 ], "yaw_arr": [ -2.0 @@ -18,18 +29,18 @@ -1 ], "yaw_arr": [ - -2.6 + -2.05 ] }, "3": { "start_arr": [ - 600 + 0 ], "end_arr": [ - 13850 + -1 ], "yaw_arr": [ - -3.2 + -4.3 ] }, "4": { @@ -40,7 +51,7 @@ -1 ], "yaw_arr": [ - -2.0 + -2.2 ] }, "5": { @@ -51,7 +62,7 @@ -1 ], "yaw_arr": [ - -2.0 + -2.9 ] }, "6": { @@ -67,13 +78,24 @@ }, "7": { "start_arr": [ - 0, 5100 + 0 ], "end_arr": [ - -1, 10530 + -1 ], "yaw_arr": [ - -2.9, 2 + -2.4 + ] + }, + "8": { + "start_arr": [ + 0 + ], + "end_arr": [ + -1 + ], + "yaw_arr": [ + 0 ] }, "9": { @@ -84,7 +106,7 @@ -1 ], "yaw_arr": [ - -5.3 + -3.9 ] }, "10": { @@ -95,7 +117,7 @@ -1 ], "yaw_arr": [ - -4.5 + -4.75 ] }, "11": { @@ -106,18 +128,18 @@ -1 ], "yaw_arr": [ - -5.8 + -6.2 ] }, "12": { "start_arr": [ - 12050, 13490, 0 + 0, 785 ], "end_arr": [ - -1, -1, -1 + -1, 7234 ], "yaw_arr": [ - -71, 33.5, 176 + -191, 8.5 ] }, "13": { @@ -128,10 +150,10 @@ -1 ], "yaw_arr": [ - -4.3 + -4.5 ] }, - "16": { + "14": { "start_arr": [ 0 ], @@ -139,10 +161,10 @@ -1 ], "yaw_arr": [ - -3.0 + 0 ] }, - "17": { + "15": { "start_arr": [ 0 ], @@ -150,7 +172,29 @@ -1 ], "yaw_arr": [ - -3.95 + 0 + ] + }, + "16": { + "start_arr": [ + 0 + ], + "end_arr": [ + -1 + ], + "yaw_arr": [ + -2.6 + ] + }, + "17": { + "start_arr": [ + 0, 1000 + ], + "end_arr": [ + -1, 7006 + ], + "yaw_arr": [ + -6.5, 2.75 ] }, "18": { @@ -161,7 +205,7 @@ -1 ], "yaw_arr": [ - -2.3 + -2.7 ] }, "19": { @@ -172,7 +216,7 @@ -1 ], "yaw_arr": [ - 0.5 + -22.0 ] }, "20": { @@ -183,7 +227,7 @@ -1 ], "yaw_arr": [ - -4.35 + -5.25 ] }, "21": { @@ -194,7 +238,7 @@ -1 ], "yaw_arr": [ - 1.5 + -3.5 ] }, "22": { @@ -205,7 +249,7 @@ -1 ], "yaw_arr": [ - -6.0 + -6.5 ] } } \ No newline at end of file diff --git a/scripts/3d_legoloam_to_2d_hitl.py b/scripts/3d_legoloam_to_2d_hitl.py deleted file mode 100644 index a8cb59dd..00000000 --- a/scripts/3d_legoloam_to_2d_hitl.py +++ /dev/null @@ -1,239 +0,0 @@ -# import pdb; pdb.set_trace() -import os -from os.path import join -import sys -import argparse -import numpy as np -import time - -from scipy.spatial.transform import Rotation as R -from sklearn.neighbors import KDTree - -from multiprocessing import Pool -import tqdm - -parser = argparse.ArgumentParser() -parser.add_argument('--traj', default="0", - help="number of trajectory, e.g. 1") -parser.add_argument('--option', default="hitl", - help="hitl for hitl SLAM and vis for visualization ") - -def read_bin(bin_path, keep_intensity=True): - OS1_POINTCLOUD_SHAPE = [1024, 128, 3] - num_points = OS1_POINTCLOUD_SHAPE[0]*OS1_POINTCLOUD_SHAPE[1] - bin_np = np.fromfile(bin_path, dtype=np.float32).reshape(num_points, -1) - - if not keep_intensity: - bin_np = bin_np[:, :3] - return bin_np - -# https://stackoverflow.com/questions/5782658/extracting-yaw-from-a-quaternion -def calc_yaw(qw, qx, qy, qz): - yaw = np.arctan2(2.0*(qy*qz + qw*qx), qw*qw - qx*qx - qy*qy + qz*qz) - return yaw - -def get_pose(pose_np, np_n_of_val_scans, vis = False): - _, x, y, z, qw, qx, qy, qz = pose_np.transpose() - if vis: - np_tmp = np.concatenate([x.reshape(-1,1), y.reshape(-1,1), z.reshape(-1,1)], axis=1).reshape(-1, 3) - return np.repeat(np_tmp, repeats=np_n_of_val_scans, axis=0) - else: - yaw = calc_yaw(qw, qx, qy, qz) - np_tmp = np.concatenate([x.reshape(-1,1), y.reshape(-1,1), yaw.reshape(-1,1)], axis=1).reshape(-1, 3) - return np.repeat(np_tmp, repeats=np_n_of_val_scans, axis=0) - -def pose_to_homo_mat(pose): - trans = pose[1:4] - quat = np.array([pose[5], pose[6], pose[7], pose[4]]) # qx, qy, qz, qw - order confirmed - rot_mat = R.from_quat(quat).as_matrix() - homo_mat = np.eye(4, dtype=np.float64) - homo_mat[:3, :3] = rot_mat - homo_mat[:3, 3] = trans - return homo_mat - -def correct_obs(homo_mat, obs): - # correct observation x & y with homo matrix - obs = np.concatenate( (obs, np.ones((len(obs), 1))), axis=-1 ).transpose() - corrected_obs = np.matmul(homo_mat, obs)[:3, :].transpose() - return corrected_obs # (# of scans, 3) - -def get_normal(closest_points): - """ test data - a = [-1, 1, 2, -4, 2, 2, -2, 1, 5, 0, 0, 0, 1, 0, 0, 0, 1, 0, -1, 1, 2, -4, 2, 2, -2, 1, 5] - closest_points = np.array(a, dtype='f').reshape(3, 3, 3) - result: (3, 9, 1) """ - p0,p1,p2 = [closest_points[:,i,:] for i in range(3)] - nv = np.cross(p0-p1, p0-p2) - return nv[:, :2] - -def cov_gen(n_of_bin, np_n_of_val_scans, xxyy_sigma=0.0007, thth_sigma=0.000117, xyth_sigma=0.002, xyyx_sigma=0.000008): - xxyy_unc = np.random.normal(0, xxyy_sigma, n_of_bin-1) - thth_unc = np.random.normal(0, thth_sigma, n_of_bin-1) - xyth_unc = np.random.normal(0, xyth_sigma, n_of_bin-1) - xyyx_unc = np.random.normal(0, xyyx_sigma, n_of_bin-1) - cov_mat = np.concatenate([xxyy_unc, xyyx_unc, xyth_unc, xyyx_unc, xxyy_unc, xyth_unc, xyth_unc, xyth_unc, thth_unc]) - cov_mat = cov_mat.reshape(9, -1).transpose() - cov_mat = np.concatenate([np.zeros(9).reshape(1, 9), cov_mat], axis=0) - return np.repeat(cov_mat, repeats=np_n_of_val_scans, axis=0) - -def process_pc(args): - pose_of_frame, frame, bin_path, outdir = args - - pc_np_org = read_bin(bin_path, False) - - # pc_np = pc_np_org.reshape(1024, 128, 3) # reshape incorrectly - # axis 0 - hortizontal channel (1024), axis 1 - vertical channel (128) - pc_np = pc_np_org.reshape(128, 1024, 3) - pc_np = np.transpose(pc_np, [1, 0, 2]) # result (1024, 128, 3) - pc_np = pc_np[::2, ::8, :] # downsample: 128 -> 16 and 1024 -> 512 - # TEST - np.equal(pc_np_org[::1024], pc_np[0]) - - pc_np_flattened = pc_np.reshape(-1, 3) - tree = KDTree(pc_np_flattened, leaf_size=2) # for an efficient closest points search - - # 1.1) filter Z Channel based on thresholds - z_lower, z_upper = 4.0, 6.0 # 0.7, 2.0 - z_mask = np.logical_or( (pc_np[:, :, 2] < z_lower), (pc_np[:, :, 2] > z_upper) ) - pc_np_filtered = np.copy(pc_np) # copy. o.w. overwrite pc_np - pc_np_filtered[z_mask, :] = np.inf # to help finding the min Euc. dist. - - # 1.2) filter based on range (x-y Euc. dist.) thresholds - r_lower, r_upper = 15.0, 25.0 - pc_np_norm_2d = np.linalg.norm(pc_np_filtered[:, :, :2], axis=-1) # x-y Euc. dist. - r_mask = np.logical_or( (pc_np_norm_2d < r_lower), (pc_np_norm_2d > r_upper) ) - pc_np_filtered[r_mask, :] = np.inf - - # 2) find the min ||p_i|| - pc_np_norm_3d = np.linalg.norm(pc_np_filtered, axis=-1) - pc_np_min_idx = np.argmin(pc_np_norm_3d, axis=-1) - pc_np_min_mask = (pc_np_min_idx != 0) - n_of_val_scans = np.count_nonzero(pc_np_min_idx) - # np_n_of_val_scans[frame] = n_of_val_scans - dummy_idx = np.arange(len(pc_np_min_idx)) - pc_np_min = pc_np[dummy_idx, pc_np_min_idx, :] # pc_np_min = pc_np[dummy_idx, pc_np_min_idx[dummy_idx], :] - pc_np_min = pc_np_min[pc_np_min_mask, :] - - n_of_val_path = join(outdir, "n_of_val_%i.npy"%frame) - np.save(n_of_val_path, n_of_val_scans) - - if ( n_of_val_scans != 0 ): # skips to write if no valid pc in the frame - if (option == "vis"): - obs_xy = correct_obs(pose_to_homo_mat(pose_of_frame), pc_np_min) - else: - obs_xy = correct_obs(pose_to_homo_mat(pose_of_frame), pc_np_min)[:, :2] - - # 3) find two nearest points - _, ind = tree.query(pc_np_min, k=3) - closest_points = pc_np_flattened[ind] - - # 4) find the normal vector to the plane - norm_xy = get_normal(closest_points) - np_bin_single = np.hstack((obs_xy, norm_xy)) - - # Dump frame to npy file - np_bin_path = join(outdir, "np_bin_%i.npy"%frame) - np.save(np_bin_path, np_bin_single) - - print("Processed frame %s" % str(frame)) - -def main(args): - global option - trajectory, option = args.traj, args.option - pose_path = f"/robodata/arthurz/Datasets/CODa_dev/poses/dense/{trajectory}.txt" - bin_dir = f"/robodata/arthurz/Datasets/CODa/3d_comp/os1/{trajectory}/" - outdir = f"./cloud_to_laser/%s" % args.traj - save_dir = os.path.join("./", "HitL") - - if not os.path.exists(outdir): - os.makedirs(outdir) - - if not os.path.exists(save_dir): - os.makedirs(save_dir) - - pose_np = np.loadtxt(pose_path).reshape(-1, 8) - - # number of lidar frames - # n_of_bin_files = 1000 - n_of_bin_files = len([entry for entry in os.listdir(bin_dir) if os.path.isfile(os.path.join(bin_dir, entry))]) - print("\n" + "---"*15) - print(f"\nNumber of lidar frames: {len(pose_np)}\n") - - np_n_of_val_scans = np.zeros(len(pose_np), dtype=int) - - if (option == "vis"): - np_bin = np.zeros((len(pose_np)*1024, 8)) - else: - np_bin = np.zeros((len(pose_np)*1024, 16)) - - pose_of_frame_list, frame_list, bin_path_list, outdir_list = [], [], [], [] - for frame in range(n_of_bin_files): - pose_of_frame_list.append(pose_np[frame, :]) - bin_file = f"3d_comp_os1_{trajectory}_{frame}.bin" - bin_path = os.path.join(bin_dir, bin_file) - bin_path_list.append(bin_path) - frame_list.append(frame) - outdir_list.append(outdir) - - pool = Pool(processes=96) - pool.map(process_pc, zip(pose_of_frame_list, frame_list, bin_path_list, outdir_list), chunksize=32) - - print("\nAggregating all pcs.\n") - # Load in pcs for each frame - start, end = 0, 0 - for frame in range(n_of_bin_files): - np_bin_path = join(outdir, "np_bin_%i.npy"%frame) - n_of_val_path = join(outdir, "n_of_val_%i.npy"%frame) - assert os.path.exists(np_bin_path), "file %s does not exist" % np_bin_path - assert os.path.exists(n_of_val_path), "file %s does not exist" % n_of_val_path - - # Load and assign processed point clouds to unified file - np_bin_single = np.load(np_bin_path) - np_n_of_val_scans[frame] = np.load(n_of_val_path) - end = start + np_n_of_val_scans[frame] - if (option == "vis"): - np_bin[start:end, 3:6] = np_bin_single[:, :3] - np_bin[start:end:, 6:] = np_bin_single[:, 3:5]# normal vector - else: - np_bin[start:end, 3:7] = np_bin_single - - # update start for next iteration - start = end - - # Truncate unused zeros and fill pose and covariance - np_bin = np_bin[:end] - - # Remove all points too close to any of our global poses - """ - Steps: - 1. Create a KD tree with the global observations - 2. Query using global poses to get all points in KD tree within ball radius (MIN_DISTANCE) - 3. Merge the point indices list into a mask - 4. Mask the poses - # """ - - if (option == "vis"): - '''3d pose and 3d obs''' - np_bin[:, :3] = get_pose(pose_np, np_n_of_val_scans, True) # add pose - (x,y,z) - fpath_out = os.path.join(save_dir, trajectory + "_vis" + ".bin") - # save as bin file - np_bin_flattended = np_bin.reshape(-1,) - np_bin_flattended.tofile(fpath_out) - else: - np_bin[:, :3] = get_pose(pose_np, np_n_of_val_scans, False) - np_bin[:, 7:] = cov_gen(n_of_bin=len(pose_np), np_n_of_val_scans=np_n_of_val_scans) - fpath_out = os.path.join(save_dir, trajectory + ".txt") - header = 'StarterMap\n1455656519.379815' - np.savetxt(fpath_out, np_bin, delimiter=',', header=header, comments='') - - -if __name__ == "__main__": - start_time = time.time() - args = parser.parse_args() - main(args) - print("--- Final: %s seconds ---" % (time.time() - start_time)) - - if (args.option == "vis"): - os.system(f"scripts/debug_visualize.py --traj {args.traj}") - -# python -W ignore scripts/3d_legoloam_to_2d_hitl.py --traj 0 --option hitl -# python -W ignore scripts/3d_legoloam_to_2d_hitl.py --traj 0 --option vis \ No newline at end of file diff --git a/scripts/3d_legoloam_to_2d_rviz.py b/scripts/3d_legoloam_to_2d_rviz.py index 57b68934..5f7d7a4a 100644 --- a/scripts/3d_legoloam_to_2d_rviz.py +++ b/scripts/3d_legoloam_to_2d_rviz.py @@ -5,17 +5,15 @@ import argparse import numpy as np import time +import json from scipy.spatial.transform import Rotation as R -from sklearn.neighbors import KDTree - import rospy from sensor_msgs.msg import PointCloud2 from std_msgs.msg import Header from geometry_msgs.msg import PoseStamped -from multiprocessing import Pool -import tqdm +import open3d as o3d # For imports sys.path.append(os.getcwd()) @@ -26,90 +24,261 @@ parser = argparse.ArgumentParser() parser.add_argument('--traj', default="0", help="number of trajectory, e.g. 1") -parser.add_argument('--option', default="hitl", - help="hitl for hitl SLAM and vis for visualization ") +parser.add_argument('--pose_tag', default="0", + help="Identify tag to publish pose") +parser.add_argument('--pub', default=1, + help="Publish global poses to rviz for visualization, set to 0 to stop publishing") +parser.add_argument('--save_to_file', default="-1", + help="Saves entire map to PCD file format, set to -1 if you don't want to save. Otherwise, specify directory to save pcd file to") + + + + +def yaw_to_homo(pose_np, yaw): + trans = pose_np[:, 1:4] + rot_mat = R.from_euler('z', yaw, degrees=True).as_matrix() + tmp = np.expand_dims(np.eye(4, dtype=np.float64), axis=0) + homo_mat = np.repeat(tmp, len(trans), axis=0) + homo_mat[:, :3, :3] = rot_mat + # homo_mat[:, :3, 3 ] = trans + return homo_mat + +def apply_hom_mat(pose_np, homo_mat, non_origin): + _, x, y, z, _, _, _, _ = pose_np.transpose() + x, y, z, ones = [p.reshape(-1, 1) for p in [x, y, z, np.ones(len(x))]] + xyz1 = np.expand_dims(np.hstack((x, y, z, ones)), -1) + + if non_origin: + xyz1_center = np.copy(xyz1) + xyz1_center[:, :2] = xyz1[:, :2] - xyz1[0, :2] + xyz1_center_rotated = np.matmul(homo_mat, xyz1_center)[:, :3, 0] + xyz1_center_rotated[:, :2] = xyz1_center_rotated[:, :2] + xyz1[0, :2].reshape(1, -1) + corrected_pose_np = xyz1_center_rotated + else: + corrected_pose_np = np.matmul(homo_mat, xyz1)[:, :3, 0] + return corrected_pose_np + +def correct_pose(pose_np, trajectory): + dir = './json' + fpath = os.path.join(dir, 'pose_correction.json') + f = open(fpath, "r") + pose_correction = json.load(f) + f.close() + + JSON_NAMES = ["start_arr", "end_arr", "yaw_arr"] + start_arr, end_arr, yaw_arr = [], [], [] + + if trajectory in pose_correction.keys(): + traj_dict = pose_correction[trajectory] + start_arr, end_arr, yaw_arr = traj_dict[JSON_NAMES[0]], traj_dict[JSON_NAMES[1]], traj_dict[JSON_NAMES[2]] + + corrected_pose = np.copy(pose_np) + + # handles multiple rotation + for i in range(len(start_arr)): + start, end, yaw = start_arr[i], end_arr[i], yaw_arr[i] + non_origin = False + if start != 0: + non_origin = True + homo_mat = yaw_to_homo(corrected_pose[:, :], yaw) + corrected_pose[:, 1:4] = apply_hom_mat(corrected_pose[:, :], homo_mat, non_origin) + # import pdb; pdb.set_trace(); + return corrected_pose + +def apply_manual_correction_obs(trajectory, pc_np_filtered, n_pcs, pose_np): + + dir = './json' + fpath = os.path.join(dir, 'pose_correction.json') + f = open(fpath, "r") + pose_correction = json.load(f) + f.close() + + JSON_NAMES = ["start_arr", "end_arr", "yaw_arr"] + start_arr, end_arr, yaw_arr = [], [], [] + + if trajectory in pose_correction.keys(): + traj_dict = pose_correction[trajectory] + start_arr, end_arr, yaw_arr = traj_dict[JSON_NAMES[0]], traj_dict[JSON_NAMES[1]], traj_dict[JSON_NAMES[2]] + + rot_mat = None + manual_corrected_pc_np = None + manual_pose_np = None + + # handles multiple rotation + for i in range(len(start_arr)): + start, end, yaw = start_arr[i], end_arr[i], yaw_arr[i] + r = R.from_euler('z', yaw, degrees=True) + rot_mat = r.as_matrix() + # print(f"Yaw Rotation: {yaw}") + if i == 0: + manual_corrected_pc_np = rot_mat @ np.transpose(pc_np_filtered) + manual_corrected_pc_np = np.transpose(manual_corrected_pc_np) + manual_pose_np = rot_mat @ np.transpose(pose_np) + manual_pose_np = np.transpose(manual_pose_np) + else: + start_sum = np.sum(n_pcs[:start]) + end_sum = np.sum(n_pcs[:end+1]) + 1 + # import pdb; pdb.set_trace() + + # observations + offset = manual_pose_np[start, :2] # (x, y) of offset + temp_np = manual_corrected_pc_np[start_sum:end_sum] + temp_np[:, :2] -= offset + temp_np = np.transpose(rot_mat @ np.transpose(temp_np)) + temp_np[:, :2] += offset + manual_corrected_pc_np[start_sum:end_sum] = temp_np + + # poses + offset = manual_pose_np[start, :2] # (x, y) of offset + temp_np = manual_pose_np[start_sum:end_sum] + temp_np[:, :2] -= offset + temp_np = np.transpose(rot_mat @ np.transpose(temp_np)) + temp_np[:, :2] += offset + manual_pose_np[start_sum:end_sum] = temp_np + + return manual_corrected_pc_np + + +def get_obs(frame_list, outdir): + + obs = np.zeros((len(frame_list)*1024*10, 3)) + n_pcs = np.zeros(len(frame_list), dtype=int) + + start, end = 0, 0 + for i in range(len(frame_list)): # Load in pcs for each frame + frame = frame_list[i] + np_bin_path = join(outdir, "np_bin_%i.npy"%frame) + + # Load and assign processed point clouds to unified file + np_bin_single = np.load(np_bin_path) + n_pcs[i] = len(np_bin_single) + end = start + n_pcs[i] + + obs[start:end] = np_bin_single + + # update start for next iteration + start = end + + obs = obs[:np.sum(n_pcs)] + + return obs, n_pcs + +def find_closest_frame(pose_np, timestamp_np, indir, trajectory): + bin_path_list, frame_list = [], [] + for _, pose in enumerate(pose_np): + pose_ts = pose[0] + closest_lidar_frame = np.searchsorted(timestamp_np, pose_ts, side='left') + frame_list.append(closest_lidar_frame) + bin_path = set_filename_dir(indir, "3d_raw", "os1", trajectory, closest_lidar_frame, include_name=True) + bin_path_list.append(bin_path) + + return bin_path_list, frame_list def main(args): - global option - trajectory, option = args.traj, args.option + + trajectory = args.traj + pub_rviz = int(args.pub) + save_pcd = str(args.save_to_file)!="-1" + pose_tag = args.pose_tag + + # trajectory = traj indir = "/robodata/arthurz/Datasets/CODa_dev" - pose_path = f"{indir}/poses_highfreq/{trajectory}.txt" + pose_path = f"{indir}/poses/{trajectory}.txt" # pose_path = f"{trajectory}.txt" ts_path = f"{indir}/timestamps/{trajectory}_frame_to_ts.txt" - bin_dir = f"{indir}/3d_comp/os1/{trajectory}/" - outdir = f"./cloud_to_laser/%s" % args.traj + outdir = f"./cloud_to_laser/%s" % trajectory if not os.path.exists(outdir): os.makedirs(outdir) - - # Initialize ros publishing - rospy.init_node('bin_publisher', anonymous=True) - pc_pub = rospy.Publisher('/coda/ouster/lidar_packets', PointCloud2, queue_size=10) - pose_pub = rospy.Publisher('/coda/pose', PoseStamped, queue_size=10) + + if pub_rviz: + # Initialize ros publishing + rospy.init_node('bin_publisher', anonymous=True) + pc_pub = rospy.Publisher('/coda/ouster/lidar_packets', PointCloud2, queue_size=10) + pose_pub = rospy.Publisher(f'/coda/pose/{pose_tag}', PoseStamped, queue_size=10) + rate = rospy.Rate(75) pose_np = np.loadtxt(pose_path).reshape(-1, 8) timestamp_np = np.loadtxt(ts_path).reshape(-1) - # number of lidar frames - # n_of_bin_files = 1000 - n_of_bin_files = len([entry for entry in os.listdir(bin_dir) if os.path.isfile(os.path.join(bin_dir, entry))]) print("\n" + "---"*15) print(f"\nNumber of lidar frames: {len(pose_np)}\n") + bin_path_list, frame_list = find_closest_frame(pose_np, timestamp_np, indir, trajectory) + + obs_xyz, n_pcs = get_obs(frame_list, outdir) + obs_xyz_man = apply_manual_correction_obs(trajectory, obs_xyz, n_pcs, pose_np[:, 1:4]) + # Iterate through all poses LIDAR_HEIGHT = 0.8 # meters ZBAND_HEIGHT = 0.5 # meters ZMIN_OFFSET = 1.9 # meters + # ZMIN_OFFSET = 2.5 # meters - rate = rospy.Rate(100) + xyz = np.empty((0,3)) - dense_pose_np = densify_poses_between_ts(pose_np, timestamp_np) - last_pose = None + start = 0 + end = 0 for pose_idx, pose in enumerate(pose_np): + pose_ts = pose[0] - # if pose_idx < 1000: - # continue - print("pose ", pose_idx) closest_lidar_frame = np.searchsorted(timestamp_np, pose_ts, side='left') lidar_ts = timestamp_np[closest_lidar_frame] - bin_path = set_filename_dir(indir, "3d_raw", "os1", trajectory, closest_lidar_frame, include_name=True) - lidar_np = read_bin(bin_path, keep_intensity=False) - - # Filter all point between zmin and zmax, downsample angular to 1/4 original size - lidar_np = lidar_np.reshape(128, 1024, -1) - zmin = ZMIN_OFFSET - LIDAR_HEIGHT - zmax = zmin+ZBAND_HEIGHT - z_mask = np.logical_and(lidar_np[:, :, 2] > zmin, lidar_np[:, :, 2] < zmax) - lidar_np = lidar_np[z_mask].reshape(-1, 3).astype(np.float32) - - LtoG = pose_to_homo(pose) # lidar to global frame - homo_lidar_np = np.hstack((lidar_np, np.ones((lidar_np.shape[0], 1)))) - trans_homo_lidar_np = (LtoG @ homo_lidar_np.T).T - trans_lidar_np = trans_homo_lidar_np[:, :3] #.reshape(128, 1024, -1) - trans_lidar_np = trans_lidar_np.reshape(-1, 3).astype(np.float32) - - pub_pc_to_rviz(trans_lidar_np, pc_pub, lidar_ts) - pub_pose(pose_pub, pose, pose[0]) - if last_pose is None: - last_pose = pose + + pose_np = pose.reshape(1, -1) + pose = correct_pose(pose_np, trajectory) + + + + import pdb; pdb.set_trace() + + + pose = pose.reshape(-1,) + + + # print("pose ", pose_idx) + + end += n_pcs[pose_idx] + trans_lidar_np = obs_xyz_man[start:end, :] + trans_lidar_np = np.float32(trans_lidar_np) + + if save_pcd is not True and pub_rviz: + print(f'Publishing pose {pose_idx}') + pub_pc_to_rviz(trans_lidar_np, pc_pub, lidar_ts) + pub_pose(pose_pub, pose, pose[0]) + # import pdb; pdb.set_trace() else: - print("check is last poses are similar") - # print("last pose ", last_pose[1:]) - # print("last pose ", pose[1:]) - print(np.allclose(last_pose[1:], pose[1:], rtol=1e-5) ) - last_pose = pose - # import pdb; pdb.set_trace() - # rate.sleep() + if len(xyz) == 0: + xyz = trans_lidar_np + else: + xyz = np.vstack([xyz, trans_lidar_np]) + + if pub_rviz: + rate.sleep() + + start = end + + if save_pcd: + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(xyz) + fpath = f"./raw_map/{trajectory}.pcd" + o3d.io.write_point_cloud(fpath, pcd, write_ascii=False) + if __name__ == "__main__": start_time = time.time() args = parser.parse_args() main(args) - print("--- Final: %s seconds ---" % (time.time() - start_time)) + # # main(args) + + # GDC = [0,1,3,4,5,18,19] + + # for i in [0]: + # args.traj = str(i) + # args.pub = '0' + # args.save_to_file = './raw_map' + # main(args) - if (args.option == "vis"): - os.system(f"scripts/debug_visualize.py --traj {args.traj}") -# python -W ignore scripts/3d_legoloam_to_2d_hitl.py --traj 0 --option hitl -# python -W ignore scripts/3d_legoloam_to_2d_hitl.py --traj 0 --option vis \ No newline at end of file + print("--- Final: %s seconds ---" % (time.time() - start_time)) \ No newline at end of file diff --git a/scripts/gen_corrected_pose_and_obs.py b/scripts/gen_corrected_pose_and_obs.py new file mode 100644 index 00000000..80b2a6c2 --- /dev/null +++ b/scripts/gen_corrected_pose_and_obs.py @@ -0,0 +1,205 @@ +import sys, os +import numpy as np +import json +from numpy.linalg import inv +import open3d as o3d + +from scipy.spatial.transform import Rotation as R + +# For imports +sys.path.append(os.getcwd()) +from scripts.pose_to_hitl import save_hitl_input + +''' +def pose_to_homo_stacked(pose_np): + + homo_mat = np.tile(np.eye(4, dtype=np.float64), (len(pose_np), 1, 1)) + trans = pose_np[:, 1:4] + quat = np.array([pose_np[:, 5], pose_np[:, 6], pose_np[:, 7], pose_np[:, 4]]).transpose() + rot_mat = R.from_quat(quat).as_matrix() + + homo_mat[:, :3, :3] = rot_mat + homo_mat[:, :3, 3] = trans + return homo_mat + +def ego_to_global(homo_mat, pose_np): + _, x, y, z, _, _, _, _ = pose_np.transpose() + x, y, z, ones = [p.reshape(-1, 1) for p in [x, y, z, np.ones(len(x))]] + xyz1 = np.expand_dims(np.hstack((x, y, z, ones)), -1) + pose_global = np.matmul(homo_mat, xyz1)[:, :3, 0] + return pose_global + +''' + +def grab_json(trajectory): + dir = './json' + fpath = os.path.join(dir, 'pose_correction.json') + f = open(fpath, "r") + pose_correction = json.load(f) # dictionary + f.close() + + JSON_NAMES = ["start_arr", "end_arr", "yaw_arr"] + start_arr, end_arr, yaw_arr = [], [], [] + + if trajectory in pose_correction.keys(): + traj_dict = pose_correction[trajectory] + start_arr, end_arr, yaw_arr = traj_dict[JSON_NAMES[0]], traj_dict[JSON_NAMES[1]], traj_dict[JSON_NAMES[2]] + return [start_arr, end_arr, yaw_arr] + +def pose_np_to_mat(pose_np): + """ + INPUT + pose_np - (N, 8) Pose Numpy Arrays + OUTPUT + pose_mat - (N, 4, 4) Pose Matrix + """ + pose_mat = np.tile(np.eye(4, dtype=np.float64), (len(pose_np), 1, 1)) # Initialize + + xyz = pose_np[:, 1:4] + quat = pose_np[:, 4:][:, [1, 2, 3, 0]] # (qw, qx, qy, qz) -> (qx, qy, qz, qw) + rot_mat = R.from_quat(quat).as_matrix() + + pose_mat[:, :3, 3] = xyz + pose_mat[:, :3, :3] = rot_mat + return pose_mat + +def pose_mat_to_np(pose_np, pose_mat): + """ + INPUT + pose_np - (N, 8) Pose Numpy Arrays + pose_mat - (N, 4, 4) Pose Matrix + OUTPUT + pose_np_new - (N, 8) Pose Numpy Arrays after transformation + """ + + ts = pose_np[:, 0] # timestamp + pose_np_new = np.zeros((len(pose_np), 8)) + + xyz = pose_mat[:, :3, 3] # (N, 3) xyz matrix + rot_mat = pose_mat[:, :3, :3] # (N, 3, ) rotaiton matrix + r = R.from_matrix(rot_mat) + quat = r.as_quat()[:, [3, 0, 1, 2]] # (N, 4) quaternion matrix | (qx, qy, qz, qw) ->(qw, qx, qy, qz) + + pose_np_new[:, 0] = ts + pose_np_new[:, 1:4] = xyz + pose_np_new[:, 4:] = quat + return pose_np_new + +def get_rotation_mat(mode, angle): + + seq_dict = {"yaw": "z", "roll": "y", "pitch": "x"} + seq = seq_dict[mode] + r = R.from_euler(seq, angle, degrees=True) + rot_mat = r.as_matrix() + return rot_mat + +def apply_affine(pose_mat, start, end, mode='yaw', angle=0, trans=np.zeros((1, 3))): + ''' + INPUT + pose_mat - (N, 4, 4) + start - starting idx + end - ending idx + mode - either yaw, roll or pitch + angle - amount of rotation in [deg] + trans - amount of translation in (x, y, z) + OUTPUT + pose_mat - (N, 4, 4) + ''' + + end = len(pose_mat) if end == -1 else end # to differentiate global and local origin + + pose_mat_copy = pose_mat.copy()[start:end, :, :] + + rot_mat = get_rotation_mat(mode, angle) + + # Find affine matrix to multiply + affine_mat = np.tile(np.eye(4, dtype=np.float64), (len(pose_mat_copy), 1, 1)) + affine_mat[:, :3, :3] = rot_mat + affine_mat[:, :3, 3] = trans.reshape(1, -1) + + # Rotate about global origin (0, 0, 0) - cp_g = A * p_g + if end == len(pose_mat): + pose_mat = np.matmul(affine_mat, pose_mat) + return pose_mat + + # Rotate about local origin - cp_g = (T_g_r)^-1 * A * T_g_r * p_g + else: + rot_mat = R.from_euler('z', 0, degrees=True).as_matrix() + trans = pose_mat_copy[0, :3, 3] # (x,y,z) starting point + homo_mat = np.tile(np.eye(4, dtype=np.float64), (len(pose_mat_copy), 1, 1)) + homo_mat[:, :3, :3] = rot_mat + homo_mat[:, :3, 3] = trans + + T_r_g = homo_mat # Hmogenoues matrix from robot to global + T_g_r = inv(homo_mat) # Inverse of homogenoues matrix from robot to global + + tmp = np.matmul(T_g_r, pose_mat_copy) + tmp = np.matmul(affine_mat, tmp) + tmp = np.matmul(T_r_g, tmp) + pose_mat[start:end, :, :] = tmp + return pose_mat + +# Functions to save files +def saveas_txt(traj, pose_np): + ''' + INPUT + pose_np - (N, 8) standard pose format + ''' + save_dir = f"./poses_cor/" + if not os.path.exists(save_dir): + os.makedirs(save_dir) + f_out = os.path.join(save_dir, str(traj) + ".txt") + fmt = ['%.6f'] + ['%.8f']*7 + np.savetxt(f_out, pose_np, delimiter=' ', fmt=fmt, comments='') + print(f'\n[txt] Trajectory {traj} saved in {save_dir}') + +def saveas_pcd(traj, route, pose_np, obs = False): + save_dir = f"./poses_cor/{route}/" + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pose_np) + if obs == True: + fpath = os.path.join(save_dir, f"{traj}_obs.pcd") + else: + fpath = os.path.join(save_dir, f"{traj}.pcd") + o3d.io.write_point_cloud(fpath, pcd, write_ascii=False) + print(f'\n[PCD] Trajectory {traj} saved in {save_dir}') + + +def main(trajectory, route): + + indir = "/robodata/arthurz/Datasets/CODa_dev" + pose_path = f"{indir}/poses/{trajectory}.txt" + pose_global = np.loadtxt(pose_path).reshape(-1, 8) # original poses from LeGO-LOAM + + start_arr, end_arr, yaw_arr = grab_json(trajectory) + + pose_mat = pose_np_to_mat(pose_global) + + # Step 1) Apply rotation about global origin + start, end, yaw = start_arr.pop(0), end_arr.pop(0), yaw_arr.pop(0) + pose_mat = apply_affine(pose_mat, start, end, mode='yaw', angle=yaw, trans = np.zeros((1, 3))) + + # Step 2) Apply rotation about local origin + while len(start_arr) != 0: + start, end, yaw = start_arr.pop(0), end_arr.pop(0), yaw_arr.pop(0) + pose_mat = apply_affine(pose_mat, start, end, mode='yaw', angle=yaw, trans = np.zeros((1, 3))) + + pose_global = pose_mat_to_np(pose_global, pose_mat) + + + print("\n\nSaving poses.") + saveas_txt(trajectory, pose_global) + saveas_pcd(trajectory, route, pose_global[:, 1:4]) + + print("\n\nSaving observations.") + obs_xyz, _ = save_hitl_input(trajectory, route, hitl_corrected=False) + saveas_pcd(trajectory, route, obs_xyz, obs=True) + +if __name__ == "__main__": + GDC = [0,1,3,4,5,18,19] + GUAD = [2,7,12,16,17,21] + WCP = [6,9,10,11,13,20,22] + route = "GUAD" + for trajectory in [12]: + main(str(trajectory), route) + print("\n\nGenerating all corrected poses done.\n\n") \ No newline at end of file diff --git a/scripts/hitl_phase2_gen.py b/scripts/hitl_phase2_gen.py new file mode 100644 index 00000000..29251520 --- /dev/null +++ b/scripts/hitl_phase2_gen.py @@ -0,0 +1,46 @@ +import numpy as np +import os +import sys + +sys.path.append(os.getcwd()) +from scripts.pose_to_hitl import get_normal + +GDC = [0,1,3,4,5,18,19] +GUAD = [2,7,12,16,17,21] +WCP = [6,9,10,11,13,20,22] + +save_dir = os.path.join("./", "HitL") +n = None +np_bin = np.array([]) + +trajs = GDC[:2] + +print("\nPutting two txt files together.") +for traj in trajs: + dir = "/home/jihwan98/coda/HitL/Corrected_GDC" + fpath = os.path.join(dir, f"{traj}.txt") + np_txt = np.loadtxt(fpath, comments=None, delimiter=',', skiprows=2) + if len(np_bin) == 0: + np_bin = np_txt + else: + np_bin = np.vstack((np_bin, np_txt)) + if n == None: + n = len(np_bin) + +print("Compute Normal.") +np_bin[:, 5:7] = get_normal(np_bin[:, 3:6]) + +print("Save txt.") +fpath_out = os.path.join(save_dir, "GDC" + f"_{trajs[0]}_{trajs[1]}.txt") +header = 'StarterMap\n1455656519.379815' +np.savetxt(fpath_out, np_bin, delimiter=',', header=header, comments='') + +# fpath_out = os.path.join(save_dir, "0" + ".txt") +# header = 'StarterMap\n1455656519.379815' +# np.savetxt(fpath_out, np_bin[:n], delimiter=',', header=header, comments='') + +# fpath_out = os.path.join(save_dir, "1" + ".txt") +# header = 'StarterMap\n1455656519.379815' +# np.savetxt(fpath_out, np_bin[n:], delimiter=',', header=header, comments='') + +print("Global Map done.") \ No newline at end of file diff --git a/scripts/hitl_to_pose.py b/scripts/hitl_to_pose.py new file mode 100644 index 00000000..955efb1c --- /dev/null +++ b/scripts/hitl_to_pose.py @@ -0,0 +1,72 @@ +import os +from os.path import join +import sys +import argparse +import numpy as np +import time + +from scipy.spatial.transform import Rotation as R + +# parser = argparse.ArgumentParser() +# parser.add_argument('--traj', default=0, +# help="number of trajectory, e.g. 1") + +def yaw_to_quat(yaw): + angle = R.from_euler('z', yaw, degrees=False) + quat_np = angle.as_quat() + quat_np = quat_np[:, [3, 0, 1, 2]] # (qx, qy, qz, qw) ->( qw, qx, qy, qz) + return quat_np + +def get_Z(pose_np): + pose_Transpose = pose_np.T[3] + return pose_Transpose + +def create_Output(pose_np, xy_np, quat_np): + timestamp_np = pose_np[:, 0].reshape(-1, 1) + z_np = pose_np[:, 3].reshape(-1, 1) + final_np = np.hstack((timestamp_np, xy_np, z_np, quat_np)) + return final_np + +def main(): + # trajectory = args.traj + base = str(0) + comp = str(1) + dir = f"HitL/hitl_results" + hitl_path = os.path.join(dir, f"hitl_results_trajectory_{base}_{comp}.txt") + pose_dir = f"/robodata/arthurz/Datasets/CODa_dev/poses" + base_path = os.path.join(pose_dir, f"{base}.txt") + comp_path = os.path.join(pose_dir, f"{comp}.txt") + + hitl_np = np.loadtxt(hitl_path, delimiter=" ").reshape(-1, 3) + base_np = np.loadtxt(base_path).reshape(-1, 8) + comp_np = np.loadtxt(comp_path).reshape(-1, 8) + pose_np = np.vstack((base_np, comp_np)) + + # first create xy_np + xy_np = hitl_np[:, :2] + + # then create the quat_np + yaw_Val = hitl_np[:, 2].reshape(-1, 1) + quat_np = yaw_to_quat(yaw_Val) + + # then combine + output_np = create_Output(pose_np, xy_np, quat_np) + + base_np = output_np[:len(base_np)] + base_np = np.around(base_np[:, 0], decimals = 6) + base_np = np.around(base_np[:, 1:], decimals = 8) + comp_np = output_np[len(base_np):] + comp_np = np.around(comp_np[:, 0], decimals = 6) + comp_np = np.around(comp_np[:, 1:], decimals = 8) + + save_dir = os.path.join("HitL/poses_cor") + base_out = os.path.join(save_dir, base + "_cor.txt") + comp_out = os.path.join(save_dir, comp + "_cor.txt") + np.savetxt(base_out, base_np, delimiter=' ', comments='') + np.savetxt(comp_out, comp_np, delimiter=' ', comments='') + +if __name__ == "__main__": + start_time = time.time() + # args = parser.parse_args() + main() + print("--- Final: %s seconds ---" % (time.time() - start_time)) \ No newline at end of file diff --git a/scripts/pose_to_hitl.py b/scripts/pose_to_hitl.py new file mode 100755 index 00000000..bf173cc6 --- /dev/null +++ b/scripts/pose_to_hitl.py @@ -0,0 +1,353 @@ +# import pdb; pdb.set_trace() +import os +import shutil +from os.path import join +import sys +import argparse +import numpy as np +import time +import json + +import open3d as o3d +from scipy.spatial.transform import Rotation as R +from sklearn.preprocessing import normalize +from multiprocessing import Pool + +# For imports +sys.path.append(os.getcwd()) +from helpers.sensors import set_filename_dir + +# parser = argparse.ArgumentParser() +# parser.add_argument('--traj', default="0", +# help="number of trajectory, e.g. 1") + +def read_bin(bin_path, keep_intensity=False): + OS1_POINTCLOUD_SHAPE = [1024, 128, 3] + num_points = OS1_POINTCLOUD_SHAPE[0]*OS1_POINTCLOUD_SHAPE[1] + bin_np = np.fromfile(bin_path, dtype=np.float32).reshape(num_points, -1) + if not keep_intensity: + bin_np = bin_np[:, :3] + return bin_np + +def calc_yaw(qw, qx, qy, qz): + quat_np = np.stack((qx, qy, qz, qw), axis=-1) + euler_rot = R.from_quat(quat_np).as_euler('xyz', degrees=False) + yaw = euler_rot[:, -1] + return yaw + +def yaw_to_homo(pose_np, yaw): + trans = pose_np[:, 1:4] + rot_mat = R.from_euler('z', yaw, degrees=True).as_matrix() + tmp = np.expand_dims(np.eye(4, dtype=np.float64), axis=0) + homo_mat = np.repeat(tmp, len(trans), axis=0) + homo_mat[:, :3, :3] = rot_mat + # homo_mat[:, :3, 3 ] = trans + return homo_mat + +def apply_hom_mat(pose_np, homo_mat, non_origin): + _, x, y, z, _, _, _, _ = pose_np.transpose() + x, y, z, ones = [p.reshape(-1, 1) for p in [x, y, z, np.ones(len(x))]] + xyz1 = np.expand_dims(np.hstack((x, y, z, ones)), -1) + + if non_origin: + xyz1_center = np.copy(xyz1) + xyz1_center[:, :2] = xyz1[:, :2] - xyz1[0, :2] + xyz1_center_rotated = np.matmul(homo_mat, xyz1_center)[:, :3, 0] + xyz1_center_rotated[:, :2] = xyz1_center_rotated[:, :2] + xyz1[0, :2].reshape(1, -1) + corrected_pose_np = xyz1_center_rotated + else: + corrected_pose_np = np.matmul(homo_mat, xyz1)[:, :3, 0] + return corrected_pose_np + +def get_pose(pose_np, n_pcs): + _, x, y, z, qw, qx, qy, qz = pose_np.transpose() + yaw = calc_yaw(qw, qx, qy, qz) + np_tmp = np.concatenate([x.reshape(-1,1), y.reshape(-1,1), yaw.reshape(-1,1)], axis=1).reshape(-1, 3) + return np.repeat(np_tmp, repeats=n_pcs, axis=0) + +def pose_to_homo_mat(pose): + trans = pose[1:4] + quat = np.array([pose[5], pose[6], pose[7], pose[4]]) # qx, qy, qz, qw - order confirmed + rot_mat = R.from_quat(quat).as_matrix() + homo_mat = np.eye(4, dtype=np.float64) + homo_mat[:3, :3] = rot_mat + homo_mat[:3, 3] = trans + return homo_mat + +def correct_obs(homo_mat, obs): + # correct observation x & y with homo matrix + obs = np.concatenate( (obs, np.ones((len(obs), 1))), axis=-1 ).transpose() + corrected_obs = np.matmul(homo_mat, obs)[:3, :].transpose() + return corrected_obs # (# of scans, 3) + +def get_normal(pcs, knn=32): + obs = o3d.geometry.PointCloud() + obs.points = o3d.utility.Vector3dVector(pcs) + search = o3d.geometry.KDTreeSearchParamKNN(knn=knn) + obs.estimate_normals(search_param=search) + obs.orient_normals_consistent_tangent_plane(k=knn) + norm = np.asarray(obs.normals) + norm = normalize(norm[:, :2]) + return norm + +def cov_gen(n_of_bin, n_pcs, xxyy_sigma=0.0007, thth_sigma=0.000117, xyth_sigma=0.002, xyyx_sigma=0.000008): + xxyy_unc = np.random.normal(0, xxyy_sigma, n_of_bin-1) + thth_unc = np.random.normal(0, thth_sigma, n_of_bin-1) + xyth_unc = np.random.normal(0, xyth_sigma, n_of_bin-1) + xyyx_unc = np.random.normal(0, xyyx_sigma, n_of_bin-1) + cov_mat = np.concatenate([xxyy_unc, xyyx_unc, xyth_unc, xyyx_unc, xxyy_unc, xyth_unc, xyth_unc, xyth_unc, thth_unc]) + cov_mat = cov_mat.reshape(9, -1).transpose() + cov_mat = np.concatenate([np.zeros(9).reshape(1, 9), cov_mat], axis=0) + return np.repeat(cov_mat, repeats=n_pcs, axis=0) + +def z_filter(pc_np): + LIDAR_HEIGHT = 0.8 # meters + ZBAND_HEIGHT = 0.3 # meters + # ZMIN_OFFSET = 1.9 # meters + ZMIN_OFFSET = 2.5 # meters + zmin = ZMIN_OFFSET - LIDAR_HEIGHT + zmax = zmin + ZBAND_HEIGHT + z_mask = np.logical_and( (pc_np[:, 2] > zmin), (pc_np[:, 2] < zmax) ) + pc_np_filtered = np.copy(pc_np) # copy original pc. o.w. overwrite pc_np + pc_np_filtered = pc_np_filtered[z_mask] + return pc_np_filtered + +def r_filter(pc_np_filtered): + rmin, rmax = 5.0, 100.0 + pc_np_norm_2d = np.linalg.norm(pc_np_filtered[:, :2], axis=-1) # x-y Euc. dist. + r_mask = np.logical_and( (pc_np_norm_2d > rmin), (pc_np_norm_2d < rmax) ) + pc_np_filtered = pc_np_filtered[r_mask] + return pc_np_filtered + +def process_pc(args): + + pose_of_frame, frame, bin_path, outdir, traj = args + pc_np = read_bin(bin_path) + + pc_np = pc_np.reshape(128, 1024, 3) + pc_np = np.transpose(pc_np, [1, 0, 2]) # result (1024, 128, 3) + pc_np = pc_np[::, ::8, :] # downsample: 1024 -> 512 and 128 -> 32 + pc_np = pc_np.reshape(-1, 3) + + # 1.1) filter Z Channel based on thresholds + pc_np_filtered = z_filter(pc_np) + # 1.2) filter based on range (x-y Euc. dist.) thresholds + pc_np_filtered = r_filter(pc_np_filtered) + + # manual_corrected_pc_np = apply_manual_correction_obs(traj, frame, pc_np_filtered) + # obs_xy = correct_obs(pose_to_homo_mat(pose_of_frame), manual_corrected_pc_np) + obs_xy = correct_obs(pose_to_homo_mat(pose_of_frame), pc_np_filtered) + if len(obs_xy) == 0: + print(f"Issue. Frame: {frame}") + + # Dump frame to npy file + np_bin_path = join(outdir, "np_bin_%i.npy"%frame) + np.save(np_bin_path, obs_xy) + +def find_closest_frame(pose_np, timestamp_np, indir, trajectory): + bin_path_list, frame_list = [], [] + for _, pose in enumerate(pose_np): + pose_ts = pose[0] + closest_lidar_frame = np.searchsorted(timestamp_np, pose_ts, side='left') + frame_list.append(closest_lidar_frame) + bin_path = set_filename_dir(indir, "3d_raw", "os1", trajectory, closest_lidar_frame, include_name=True) + bin_path_list.append(bin_path) + + return bin_path_list, frame_list + +def get_obs(frame_list, outdir): + + obs = np.zeros((len(frame_list)*1024*10, 3)) + n_pcs = np.zeros(len(frame_list), dtype=int) + + start, end = 0, 0 + for i in range(len(frame_list)): # Load in pcs for each frame + frame = frame_list[i] + np_bin_path = join(outdir, "np_bin_%i.npy"%frame) + + # Load and assign processed point clouds to unified file + np_bin_single = np.load(np_bin_path) + n_pcs[i] = len(np_bin_single) + end = start + n_pcs[i] + + obs[start:end] = np_bin_single + + # update start for next iteration + start = end + + obs = obs[:np.sum(n_pcs)] + + return obs, n_pcs + +def save_txt(np_bin, trajectory, route, hitl_corrected=False): + print("Saving txt ...") + + save_dir = os.path.join("./", "HitL", f"Raw_{route}") + # if hitl_corrected: + # save_dir = os.path.join("./", "HitL", f"HitL_corrected_{route}") + # else: + # save_dir = os.path.join("./", "HitL", f"Manual_corrected_{route}") + if not os.path.exists(save_dir): + os.makedirs(save_dir) + fpath_out = os.path.join(save_dir, trajectory + ".txt") + header = 'StarterMap\n1455656519.379815' + np.savetxt(fpath_out, np_bin, delimiter=',', header=header, comments='') + print(f"Saved in {save_dir}") + +def apply_manual_correction_obs(trajectory, pc_np_filtered, n_pcs, pose_np): + + dir = './json' + fpath = os.path.join(dir, 'pose_correction.json') + f = open(fpath, "r") + pose_correction = json.load(f) + f.close() + + JSON_NAMES = ["start_arr", "end_arr", "yaw_arr"] + start_arr, end_arr, yaw_arr = [], [], [] + + if trajectory in pose_correction.keys(): + traj_dict = pose_correction[trajectory] + start_arr, end_arr, yaw_arr = traj_dict[JSON_NAMES[0]], traj_dict[JSON_NAMES[1]], traj_dict[JSON_NAMES[2]] + + rot_mat = None + manual_corrected_pc_np = None + manual_pose_np = None + + # handles multiple rotation + for i in range(len(start_arr)): + start, end, yaw = start_arr[i], end_arr[i], yaw_arr[i] + r = R.from_euler('z', yaw, degrees=True) + rot_mat = r.as_matrix() + # print(f"Yaw Rotation: {yaw}") + if i == 0: + manual_corrected_pc_np = rot_mat @ np.transpose(pc_np_filtered) + manual_corrected_pc_np = np.transpose(manual_corrected_pc_np) + manual_pose_np = rot_mat @ np.transpose(pose_np) + manual_pose_np = np.transpose(manual_pose_np) + else: + start_sum = np.sum(n_pcs[:start]) + end_sum = np.sum(n_pcs[:end+1]) + 1 + # import pdb; pdb.set_trace() + + # observations + offset = manual_pose_np[start, :2] # (x, y) of offset + temp_np = manual_corrected_pc_np[start_sum:end_sum] + temp_np[:, :2] -= offset + temp_np = np.transpose(rot_mat @ np.transpose(temp_np)) + temp_np[:, :2] += offset + manual_corrected_pc_np[start_sum:end_sum] = temp_np + + # poses + offset = manual_pose_np[start, :2] # (x, y) of offset + temp_np = manual_pose_np[start_sum:end_sum] + temp_np[:, :2] -= offset + temp_np = np.transpose(rot_mat @ np.transpose(temp_np)) + temp_np[:, :2] += offset + manual_pose_np[start_sum:end_sum] = temp_np + + return manual_corrected_pc_np + +def save_hitl_input(trajectory, route, hitl_corrected=False): + pose_np = None + # indir = "/robodata/arthurz/Datasets/CODa_dev/poses" + indir = "/home/jihwan98/coda/poses_cor" # Corrected Pose + pose_path = os.path.join(indir, f"{trajectory}.txt") + pose_np = np.loadtxt(pose_path).reshape(-1, 8) + # if hitl_corrected: + # indir = f"/home/jihwan98/coda/HitL/poses_cor/{route}" + # pose_path = os.path.join(indir, f"{trajectory}_cor.txt") + # pose_np = np.loadtxt(pose_path).reshape(-1, 8) + # else: + # indir = f"/home/jihwan98/coda/HitL/poses_manual/{route}" + # pose_path = os.path.join(indir, f"{trajectory}.txt") + # pose_np = np.loadtxt(pose_path).reshape(-1, 8) + # print(f"Length: {len(pose_np)}") + + # Get timestamp array + CODa_dev_dir = "/robodata/arthurz/Datasets/CODa_dev" + # ts_path = f"{CODa_dev_dir}/timestamps/{trajectory}_frame_to_ts.txt" + ts_path = f"{CODa_dev_dir}/timestamps/{trajectory}.txt" + timestamp_np = np.loadtxt(ts_path).reshape(-1) + + outdir = f"./cloud_to_laser/%s" % trajectory + if not os.path.exists(outdir): + os.makedirs(outdir) + + bin_path_list, frame_list = find_closest_frame(pose_np, timestamp_np, CODa_dev_dir, trajectory) + + pose_of_frame_list, outdir_list = [], [] + for frame in range(len(frame_list)): + pose_of_frame_list.append(pose_np[frame, :]) + outdir_list.append(outdir) + + assert len(pose_of_frame_list) == len(frame_list) == len(bin_path_list) == len(outdir_list) + + traj_list = [trajectory] * len(pose_of_frame_list) + + pool = Pool(processes=96) + pool.map(process_pc, zip(pose_of_frame_list, frame_list, bin_path_list, outdir_list, traj_list), chunksize=32) + pool.close() + pool.join() + + print("Pool process done.") + + obs_xyz, n_pcs = get_obs(frame_list, outdir) + + return (obs_xyz, n_pcs) + + # HITL = False + + # if not HITL: + + # print("Applying manual correction and saving") + + # # Rotation correction + # obs_xyz_man = apply_manual_correction_obs(trajectory, obs_xyz, n_pcs, pose_np[:, 1:4]) + # # Trnaslation correction + # # obs_xyz_man[:, :2] = obs_xyz_man[:, :2] + np.array([-3.5, -0.0]) + + # pcd = o3d.geometry.PointCloud() + # pcd.points = o3d.utility.Vector3dVector(obs_xyz) + # fpath = f"./obs_xyz/{route}/{trajectory}_org.pcd" + # o3d.io.write_point_cloud(fpath, pcd, write_ascii=False) + + # pcd = o3d.geometry.PointCloud() + # pcd.points = o3d.utility.Vector3dVector(obs_xyz_man) + # fpath = f"./obs_xyz/{route}/{trajectory}_man.pcd" + # o3d.io.write_point_cloud(fpath, pcd, write_ascii=False) + + # else: + # np_bin = np.zeros((len(obs_xyz), 16)) + # np_bin[:, :3] = get_pose(pose_np, n_pcs) + # np_bin[:, 3:6] = obs_xyz + # print("Normal vectors start.") + # np_bin[:, 5:7] = get_normal(obs_xyz) + # print("Normal vectors done.") + # np_bin[:, 7:] = cov_gen(n_of_bin=len(pose_np), n_pcs=n_pcs) + + # # Round up values + # np_bin[:, :7] = np.around(np_bin[:, :7], decimals = 4) + # np_bin[:, 7:] = np.around(np_bin[:, 7:], decimals = 6) + + # n_unique = len(np.unique(np_bin[:, :2], axis=0)) + # print(f"Unique: {n_unique}") + + # print("\nSaving text") + # save_txt(np_bin, trajectory, route, hitl_corrected) + +if __name__ == "__main__": + + GDC = [0,1,3,4,5,18,19] + GUAD = [2,7,12,16,17,21] + WCP = [6,9,10,11,13,20,22] + + route = "GDC" + route = "GUAD" + # route = "WCP" + + trajs = [0] + for i in trajs: + print(f"\nStarted Trajectory {i}") + trajectory = str(i) + save_hitl_input(trajectory, route, hitl_corrected=False) + print("Done.\n") \ No newline at end of file diff --git a/scripts/vis_map.py b/scripts/vis_map.py index 7ab15dc9..13dfe972 100644 --- a/scripts/vis_map.py +++ b/scripts/vis_map.py @@ -37,7 +37,7 @@ def add_marker(x, y, z, m, color): latitude, longitude, altitude = convert_xyz_to_latlon(x, y, z) # folium.Marker(location=[latitude, longitude]).add_to(m) folium.CircleMarker(location=[latitude, longitude], - radius=1, weight=1, color=color).add_to(m) + radius=0.5, weight=1, color=color).add_to(m) return m def yaw_to_homo(pose_np, yaw): @@ -73,18 +73,20 @@ def correct_pose(pose_np, start_arr, end_arr, yaw_arr): if start != 0: non_origin = True homo_mat = yaw_to_homo(corrected_pose[start:end, :], yaw) + # update poses corrected_pose[start:end, 1:4] = apply_hom_mat(corrected_pose[start:end, :], homo_mat, non_origin) return corrected_pose def find_overlapping_pc(pose_np, pc): tree = KDTree(pose_np, leaf_size=2) # for an efficient closest points search - _, ind = tree.query(pc.reshape(1, -1), k=10) + _, ind = tree.query(pc.reshape(1, -1), k=30) print(ind) def main(): TRAJ_name = list(TRAJECTORY.keys())[1] trajectory_list = TRAJECTORY[TRAJ_name] - # trajectory_list = [7] + + trajectory_list = [2, 12] outdir = './json' fpath = os.path.join(outdir, 'pose_correction.json') @@ -99,9 +101,11 @@ def main(): for i in range(len(trajectory_list)): trajectory = str(trajectory_list[i]) print("---"*10 + f"\nTrajectory {trajectory}") - pose_path = f"/robodata/arthurz/Datasets/CODa_dev/poses/dense/{trajectory}.txt" + pose_path = f"/robodata/arthurz/Datasets/CODa_dev/poses/{trajectory}.txt" pose_np = np.loadtxt(pose_path).reshape(-1, 8) - + + # pose_np = pose_np[:1000] + start_arr, end_arr, yaw_arr = [], [], [] if trajectory in pose_correction.keys(): @@ -117,8 +121,8 @@ def main(): corrected_pose_np = correct_pose(pose_np, start_arr, end_arr, yaw_arr) # import pdb; pdb.set_trace() - # find_overlapping_pc(corrected_pose_np[:, 1:3], corrected_pose_np[5100, 1:3]) - + find_overlapping_pc(corrected_pose_np[:, 1:3], corrected_pose_np[1000, 1:3]) + # corrected_pose_np = corrected_pose_np[1340:5592] for pose in corrected_pose_np: # if cnt > start_arr[1]: # _, x, y, z, _, _, _, _ = pose @@ -134,26 +138,7 @@ def main(): m.save(map_filename) print("-"*20 + "\nMap saved as HTML:", map_filename) -if __name__ == '__main__': - main() - - # if not os.path.exists(outdir): - # os.makedirs(outdir) - # print("json directory created") - - # if not os.path.exists(fpath): - # json_object = json.dumps({}) - # with open(fpath, "w") as f: - # f.write(json_object) - # f.close() - # print("empty json file created") - # traj_dict = {JSON_NAMES[0]: [start], JSON_NAMES[1]: [end], JSON_NAMES[2]: [yaw]} - # pose_correction[int(trajectory)] = traj_dict - # json_object = json.dumps(pose_correction, indent=4) - - # with open(fpath, "w+") as f: - # f.write(json_object) - # f.close() - # print("\njson file successfully saved") \ No newline at end of file +if __name__ == '__main__': + main() \ No newline at end of file