hilti2022 experiment

GTO.py

import copy
import time

import open3d as o3d
import os
from collections import defaultdict
import numpy as np
import numpy.linalg as LA
import gtsam
import tqdm
import utilities as U
import transforms3d as t3d
import pickle
from scipy import spatial

pcd_dict = defaultdict(dict)

submission_list = {
    'exp01': 'exp01_construction_ground_level.txt',
    'exp02': 'exp02_construction_multilevel.txt',
    'exp03': 'exp03_construction_stairs.txt',
    'exp04': 'exp04_construction_upper_level.txt',
    'exp05': 'exp05_construction_upper_level_2.txt',
    'exp06': 'exp06_construction_upper_level_3.txt',
    'exp07': 'exp07_long_corridor.txt',
    'exp09': 'exp09_cupola.txt',
    'exp11': 'exp11_lower_gallery.txt',
    'exp15': 'exp15_attic_to_upper_gallery.txt',
    'exp21': 'exp21_outside_building.txt'
}


def registration_o3d(source, target, initial):
    # coarse
    source_coarse = source.voxel_down_sample(0.4)
    target_coarse = target.voxel_down_sample(0.4)
    # 估计法向量
    source_coarse.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=0.8, max_nn=30))
    target_coarse.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=0.8, max_nn=30))

    # registration
    max_correspondence_distance_fine = 0.4  # 0.03
    icp = o3d.pipelines.registration.registration_icp(
        source_coarse, target_coarse, max_correspondence_distance_fine, initial,
        o3d.pipelines.registration.TransformationEstimationPointToPlane(),
        o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=200))
    # print(idx, icp.fitness)

    # refine
    # coarse
    source_refine = source.voxel_down_sample(0.05)
    target_refine = target.voxel_down_sample(0.05)
    # 估计法向量
    source_refine.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=1.0, max_nn=30))
    target_refine.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=1.0, max_nn=30))

    # registration
    max_correspondence_distance_fine = 0.05  # 0.03
    icp = o3d.pipelines.registration.registration_icp(
        source_refine, target_refine, max_correspondence_distance_fine, icp.transformation,
        o3d.pipelines.registration.TransformationEstimationPointToPlane(),
        o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=200))
    # print(idx, icp.fitness)

    information_icp = o3d.pipelines.registration.get_information_matrix_from_point_clouds(
        source, target, max_correspondence_distance_fine,
        icp.transformation)

    return icp, information_icp


def gtsam_graph_construction(odom_list, odom_ts_list, cov_list=None):
    # 读取图定义
    # loop_constraints = pickle.load(open('./data/exp01.graph', 'rb'))
    loop_constraints = pickle.load(open('./data/{}.graph'.format(SEQ), 'rb'))

    ODOMETRY_NOISE = gtsam.noiseModel.Diagonal.Sigmas(np.array([1e-5, 1e-5, 1e-5, 1e-3, 1e-3, 1e-3]) * 0.1)
    # PRIOR_NOISE = gtsam.noiseModel.Diagonal.Sigmas(np.array([0.1, 0.1, np.pi * np.pi, 1e4, 1e4, 1e4]))
    PRIOR_NOISE = gtsam.noiseModel.Diagonal.Sigmas(np.array([1e-9, 1e-9, 1e-9, 1e-8, 1e-8, 1e-8]))
    graph = gtsam.NonlinearFactorGraph()
    initial = gtsam.Values()

    priorMean = gtsam.Pose3(odom_list[0])  # prior at origin
    graph.add(gtsam.PriorFactorPose3(0, priorMean, PRIOR_NOISE))
    T_delta = []

    # 依次添加odom factor
    for idx in range(1, len(odom_list)):
        # TODO 这里BetweenFactor odometry是T__(k-1)__o__(k) = (T__world__to__k-1)^-1 (T__world__to__k)
        T__odom__o__t_1 = odom_list[idx - 1]
        T__odom__o__t = odom_list[idx]
        T__t_1__o__t = np.matmul(U.inverse_se3(T__odom__o__t_1), T__odom__o__t)

        # # 检测关键帧
        delta_t = LA.norm(T__t_1__o__t[:3, 3])
        T_delta.append(delta_t)
        # if delta_t < 0.1:
        #     continue
        # T_delta.append(utilities.xyzrpy_from_se3_matrix(T__t_1__o__t)[4])

        # add edge
        odometry = gtsam.Pose3(T__t_1__o__t)
        # TODO 这里用协方差
        # cov_diag = cov_list[idx].reshape(-1)[[0, 7, 14, 21, 28, 35]] * 100
        # print(list(cov_diag))
        graph.add(gtsam.BetweenFactorPose3(
            idx - 1, idx,
            odometry,
            # gtsam.noiseModel.Diagonal.Sigmas(cov_diag)
            ODOMETRY_NOISE
        ))

        # add node
        if idx == 1:
            # 第一个节点 加入0时刻odometry
            initial.insert(0, gtsam.Pose3(T__odom__o__t_1))
        initial.insert(idx, gtsam.Pose3(T__odom__o__t))

    # 添加回环factor
    for idx, (pair, (T__target__o__source, cov_diag)) in enumerate(loop_constraints.items()):
        # if idx < 145:
        # if idx in [6, 7, 8]:
        #     continue
        target_idx, source_idx = pair
        # 添加边 噪声是icp的噪声
        graph.add(
            gtsam.BetweenFactorPose3(
                target_idx, source_idx,
                gtsam.Pose3(T__target__o__source),
                # gtsam.noiseModel.Diagonal.Sigmas(list(cov_diag[3:]) + list(cov_diag[:3]))
                # 这里的cov应该是 roll pitch yaw x y z
                gtsam.noiseModel.Diagonal.Sigmas(cov_diag)
            ))
        print(pair)

    # optimize using Levenberg-Marquardt optimization
    params = gtsam.LevenbergMarquardtParams()
    params.setVerbosity('SUMMARY')
    params.setVerbosityLM('SUMMARY')

    optimizer = gtsam.LevenbergMarquardtOptimizer(graph, initial, params)

    tic = time.time()
    result = optimizer.optimize()
    toc = time.time()
    print('Optimization Time:', toc - tic)
    # print("\nFinal Result:\n{}".format(result))

    # # # 5. Calculate and print marginal covariances for all variables
    # marginals = gtsam.Marginals(graph, result)

    assert len(odom_ts_list) == result.size()
    EST_pose_list = []
    # 保存优化之后的odometry和对应的gt
    for key_id in range(result.size()):
        T__est__odom__o__imu = result.atPose3(key_id).matrix()
        EST_pose_list.append((
            odom_ts_list[key_id],  # TODO 时间戳!
            T__est__odom__o__imu
        ))
    # TUM格式
    tum_array = U.pose_list_to_TUM_ndarray(EST_pose_list)
    # 保存结果
    rst_seq_idx = SEQ.find('exp')
    rst_seq = root.split('/')[-1][rst_seq_idx:rst_seq_idx + 5]
    np.savetxt('/home/hvt/dataset/hilti/result/pgo/{}'.format(submission_list[rst_seq]), tum_array, fmt="%.9f")

    # 优化之前的结果
    RAW_pose_list = []
    for idx, pose in enumerate(odom_list):
        RAW_pose_list.append((
            odom_ts_list[idx],  # TODO 时间戳!
            pose
        ))
    tum_array = U.pose_list_to_TUM_ndarray(RAW_pose_list)
    np.savetxt('/tmp/test_gtsam_{}_odometry.txt'.format(SEQ), tum_array, fmt="%.9f")

    # 显示最后的地图
    submap = o3d.geometry.PointCloud()
    for key_id in tqdm.tqdm(range(result.size())):
        if key_id % 10 != 0:
            continue
        T__est__odom__o__imu = result.atPose3(key_id).matrix()
        pcd = o3d.io.read_point_cloud(pcd_dict[odom_ts_list[key_id]][0])
        submap += pcd.transform(T__est__odom__o__imu)
        submap = submap.voxel_down_sample(0.05)
    o3d.visualization.draw_geometries([submap])


# SEQ = 'exp01'
# SEQ = 'exp03'
# SEQ = 'cuhk'
SEQ = 'rotating'

# root = '/home/hvt/dataset/acsc_v2x/VILM/lamppost/CUHK_lamppost_road_mapping_GTO/pcds'
root = '/tmp/rotating'
files = os.listdir(root)
for file in sorted(files):
    ts, file_type = float(file[:file.rfind('.')]), file[file.rfind('.') + 1:]
    if file_type == 'pcd':
        pcd_dict[ts][0] = os.path.join(root, file)
    elif file_type == 'odom':
        pcd_dict[ts][1] = np.loadtxt(os.path.join(root, file))
    elif file_type == 'cov':
        pcd_dict[ts][2] = np.loadtxt(os.path.join(root, file))
    else:
        raise Exception('Not Support')
# print(pcd_dict)


# # TODO 确保odometry是连续递增的
# # 读取odometry
# odom_ts_list = list(pcd_dict.keys())
# odom_list = [pcd_dict[ts][1] for ts in pcd_dict.keys()]
# cov_list = [pcd_dict[ts][2] for ts in pcd_dict.keys()]
# print(odom_list, odom_ts_list, cov_list)
# gtsam_graph_construction(odom_list, odom_ts_list, cov_list)
# exit()

adj_table = set()

graph_definition = dict()

# # TODO 处理成关键帧
#
#
# # TODO 应该首先根据odometry判断共视关系，大部分视角有overlap的scan再做registration
# pose_xyz = []
# for ts, item in pcd_dict.items():
#     pose_xyz.append(item[1][:3, 3])
# pose_kdtree = spatial.cKDTree(np.vstack(pose_xyz))
# print(pose_kdtree.query([0, 0, 0]))
# exit()

for target_idx in range(100, len(pcd_dict), 20):
    submap = o3d.geometry.PointCloud()
    for source_idx in range(100, len(pcd_dict), 1):
        # 最近若干帧之内的不匹配
        if abs(target_idx - source_idx) < 100:
            continue
        source_item = pcd_dict[list(pcd_dict.keys())[source_idx]]
        target_item = pcd_dict[list(pcd_dict.keys())[target_idx]]

        source = o3d.io.read_point_cloud(source_item[0])
        target = o3d.io.read_point_cloud(target_item[0])

        # 求初始值
        # target = T * source
        initial = U.INV(target_item[1]) @ source_item[1]
        # target_est = copy.copy(source).transform(initial)
        # source.paint_uniform_color([1, 0.706, 0])
        # target.paint_uniform_color([0, 0.294, 1])
        # target_est.paint_uniform_color([1, 0, 0])
        # o3d.visualization.draw_geometries([target_est, target])
        # continue

        result_icp, result_information = registration_o3d(source, target, initial=initial)
        # print(result_information, result_information.shape)
        if LA.norm(result_information) == 0 or result_information[5, 5] < 1e-9:
            print('Error: invalid information matrix', target_idx, source_idx)
            # print(result_information)
            continue
        try:
            cov_diag = LA.inv(result_information).reshape(-1)[[0, 7, 14, 21, 28, 35]]
        except:
            print('Error: invalid information matrix', target_idx, source_idx)
            # print(result_information)
            continue
        # 注意这里cov可能有很高的值
        if 1e-3 > cov_diag[5] > 1e-9:
            print(source_idx, target_idx)
            print(list(cov_diag))
            print(result_icp.fitness)
            T__target__o__source = result_icp.transformation
            source_est = copy.copy(source).transform(T__target__o__source)

            source.paint_uniform_color([1, 0.706, 0])
            target.paint_uniform_color([0, 0.294, 1])
            source_est.paint_uniform_color([1, 0, 0])
            # o3d.visualization.draw_geometries([source_est, target], height=500)

            submap += source_est
            submap = submap.voxel_down_sample(0.1)
            # if source_idx % 100 == 0:
            o3d.visualization.draw_geometries([submap], height=500)
            # 合并到graph definition
            graph_definition[(target_idx, source_idx)] = [T__target__o__source, cov_diag]
            # print(graph_definition)
            # pickle.dump(graph_definition, open('./data/exp01.graph', 'wb'))
            pickle.dump(graph_definition, open('./data/{}.graph'.format(SEQ), 'wb'))

    o3d.visualization.draw_geometries([submap], height=500)

odom_saver.py

#!/usr/bin/env python2.7
# coding=utf-8

from __future__ import print_function, division, absolute_import
import rospy
from sensor_msgs.msg import PointCloud2
from nav_msgs.msg import Odometry

import cv2
from cv_bridge import CvBridge, CvBridgeError
import message_filters

import numpy as np
import ros_numpy

import transforms3d

import os
import sys

import open3d as o3d

frame_id = 0
odom_list = []


def get_RT_matrix(odom_msg):
    quaternion = np.asarray(
        [odom_msg.pose.pose.orientation.w, odom_msg.pose.pose.orientation.x, odom_msg.pose.pose.orientation.y,
         odom_msg.pose.pose.orientation.z])
    translation = np.asarray(
        [odom_msg.pose.pose.position.x, odom_msg.pose.pose.position.y, odom_msg.pose.pose.position.z])

    rotation = transforms3d.quaternions.quat2mat(quaternion)

    # print(T_qua2rota)
    RT_matrix = np.eye(4)
    RT_matrix[:3, :3] = rotation
    RT_matrix[:3, 3] = translation.T

    # RT_matrix = np.matmul(R_axes, RT_matrix)

    return RT_matrix


def callback(odom_msg, pc_msg):
    global frame_id, odom_list

    pc_msg.fields = [pc_msg.fields[0], pc_msg.fields[1], pc_msg.fields[2],
                     pc_msg.fields[4], pc_msg.fields[5], pc_msg.fields[6],
                     pc_msg.fields[3], pc_msg.fields[7]]

    pc_array = ros_numpy.numpify(pc_msg)
    if len(pc_array.shape) == 2:
        pc = np.zeros((pc_array.shape[0] * pc_array.shape[1], 4))
    else:
        pc = np.zeros((pc_array.shape[0], 4))
    # 解析点格式
    pc[:, 0] = pc_array['x'].reshape(-1)
    pc[:, 1] = pc_array['y'].reshape(-1)
    pc[:, 2] = pc_array['z'].reshape(-1)
    pc[:, 3] = pc_array['intensity'].reshape(-1)

    if not os.path.exists(ROOT):
        os.mkdir(ROOT)

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(pc[np.where(~np.isnan(pc[:, 0]))][:, :3])
    o3d.io.write_point_cloud(
        os.path.join(ROOT, '{}.pcd'.format(pc_msg.header.stamp.to_sec())),
        pcd
    )

    if odom_msg.pose.pose.position.x == 0 and odom_msg.pose.pose.position.y == 0:
        print('wrong pose')
        return
    odom_file = os.path.join(ROOT, '{}.odom'.format(odom_msg.header.stamp.to_sec()))
    np.savetxt(odom_file, get_RT_matrix(odom_msg))

    cov_file = os.path.join(ROOT, '{}.cov'.format(odom_msg.header.stamp.to_sec()))
    np.savetxt(cov_file, np.array(odom_msg.pose.covariance).reshape(6, 6))

    print("saving {:03}....".format(frame_id))
    frame_id += 1


if __name__ == '__main__':

    rospy.init_node("save_pcd_odom", anonymous=True)
    rospy.loginfo("Start....")

    if len(sys.argv) <= 1:
        rospy.logerr('Useage: rosrun [pakage] save_pcd_odom_KITTI.py PATH')
        exit(1)
    else:
        ROOT = sys.argv[1]
        # ROOT = '/home/hvt/dataset/TJP/KITTI/city_2'
        # ROOT = '/home/hvt/dataset/TJP/KITTI/rtk2'
        # ROOT = '/home/hvt/dataset/TJP/KITTI/bridge'
        # ROOT = '/home/hvt/dataset/HYY/lidar_ins_calibration/KITTI/sync'

    # odom_sub = message_filters.Subscriber("/gnss_odom", Odometry)
    # points_sub = message_filters.Subscriber("/rslidar_points", PointCloud2)
    odom_sub = message_filters.Subscriber("/Odometry", Odometry)
    points_sub = message_filters.Subscriber("/cloud_registered_body", PointCloud2)

    ts = message_filters.ApproximateTimeSynchronizer([odom_sub, points_sub], 100, slop=0.0001, allow_headerless=False)

    ts.registerCallback(callback)

    rospy.spin()