From bf2a1cea611d4318f7b488c34a5b01dfd5bb839c Mon Sep 17 00:00:00 2001 From: Manuel Gomes Date: Wed, 13 Jul 2022 15:14:24 +0100 Subject: [PATCH] #41 Visual odometry tool using ICP created --- atlascar2_odom/icp_visual_odometry.py | 202 ++++++++++++++++++++++++++ 1 file changed, 202 insertions(+) create mode 100755 atlascar2_odom/icp_visual_odometry.py diff --git a/atlascar2_odom/icp_visual_odometry.py b/atlascar2_odom/icp_visual_odometry.py new file mode 100755 index 0000000..b69c228 --- /dev/null +++ b/atlascar2_odom/icp_visual_odometry.py @@ -0,0 +1,202 @@ +#!/usr/bin/env python3 + +""" +ICP visual odometry from open3d +""" + +# ------------------------------------------------------------------------------- +# --- IMPORTS +# ------------------------------------------------------------------------------- + +import argparse +from cmath import e +import copy +import json +import math +import os +from collections import OrderedDict +from cv2 import transform +from matplotlib.collections import Collection + +import numpy as np +import open3d as o3d +import cv2 +import tf +from colorama import Style, Fore +from atom_evaluation.utilities import atomicTfFromCalibration +from atom_core.atom import getTransform +from atom_core.dataset_io import addNoiseToInitialGuess, saveResultsJSON +from atom_core.vision import depthToPointCloud + +def drawRegistrationResults(source, target, transformation, initial_transformation): + """ + Visualization tool for open3d + """ + source_temp = copy.deepcopy(source) + source_initial_temp = copy.deepcopy(source) + target_temp = copy.deepcopy(target) + source_temp.paint_uniform_color([0, 1, 1]) + source_initial_temp.paint_uniform_color([0, 0.5, 0.5]) + target_temp.paint_uniform_color([1, 0, 0]) + source_temp.transform(transformation) + source_initial_temp.transform(initial_transformation) + o3d.visualization.draw_geometries([source_initial_temp, source_temp, target_temp]) + +def pickPoints(pcd): + """ + Open a window to pick points in order to align two pointclouds + """ + print("") + print( + "1) Please pick at least three correspondences using [shift + left click]" + ) + print(" Press [shift + right click] to undo point picking") + print("2) After picking points, press q for close the window") + print("3) Press [shift + '] or [shift + «] to increase or decrease the size of the picked points") + vis = o3d.visualization.VisualizerWithEditing() + vis.create_window() + vis.add_geometry(pcd) + vis.run() # user picks points + vis.destroy_window() + print("") + return vis.get_picked_points() + +def ICPTransformations(source_point_cloud, target_point_cloud, threshold, T_init, show_images): + """ + Use ICP to estimate the transformation between two range sensors + """ + reg_p2p = o3d.pipelines.registration.registration_icp( + source_point_cloud, target_point_cloud, threshold, T_init, + o3d.pipelines.registration.TransformationEstimationPointToPoint(), + o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=200000)) + print(reg_p2p) + print("Transformation is:") + print(reg_p2p.transformation) + + if show_images: + drawRegistrationResults(source_point_cloud, target_point_cloud, reg_p2p.transformation, T_init) + + return reg_p2p + + +def saveICPTransforms(dataset, transforms, json_file, od): + """ + Save a JSON file with the data of the ICP visual odometry + """ + parent_link = 'odom' + child_link = 'base_footprint' + frame = parent_link + '-' + child_link + for collection_key, collection in od.items(): + transform = transforms[collection_key] + quat = tf.transformations.quaternion_from_matrix(transform) + dataset['collections'][collection_key]['transforms'][frame]['quat'] = quat + dataset['collections'][collection_key]['transforms'][frame]['trans'] = transform[0:3, 3] + + # Save results to a json file + saveResultsJSON(json_file, dataset) + + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("-json", "--json_file", help="Json file containing input dataset.", type=str, + required=True) + ap.add_argument("-s", "--sensor", help="Sensor used.", type=str, required=True) + ap.add_argument("-si", "--show_images", help="If true the script shows images.", action='store_true', default=False) + + # Save args + args = vars(ap.parse_args()) + json_file = args['json_file'] + show_images = args['show_images'] + + # Read json file + f = open(json_file, 'r') + dataset = json.load(f) + + # Define variables + threshold = 0.1 + od = OrderedDict(sorted(dataset['collections'].items(), key=lambda t: int(t[0]))) + pointclouds = {} + transforms = {} + filename_results_json = os.path.dirname(json_file) + '/dataset_corrected_odom.json' + + print(Fore.YELLOW + '\nRetrieving pointclouds\n' + Style.RESET_ALL) + + for collection_key, collection in od.items(): + print('Reading pointcloud from collection ' + collection_key) + # Retrieve pointclouds + if dataset['calibration_config']['sensors'][args['sensor']]['modality'] == 'lidar3d': + filename = os.path.dirname( + args['json_file']) + '/' + dataset['collections'][collection_key]['data'][args['sensor']]['data_file'] + point_cloud = o3d.io.read_point_cloud(filename) + pointclouds[collection_key] = point_cloud + elif dataset['calibration_config']['sensors'][args['sensor']]['modality'] == 'depth': + filename = os.path.dirname(json_file) + '/' + \ + dataset['collections'][collection_key]['data'][args['sensor']]['data_file'] + img = cv2.imread(filename, cv2.IMREAD_UNCHANGED) + pixels = img.shape[0] * img.shape[1] + idxs_to_trim = range(pixels) + idxs = [idx for idx in idxs_to_trim if not idx % args['sub_sample_depth']] + point_cloud = depthToPointCloud(dataset, collection_key, args['json_file'], args['sensor'], idxs) + pointclouds[collection_key] = point_cloud + else: + print('The sensor ' + args['sensor'] + ' does not have a pointcloud to retrieve, so ICP would not work.\nShutting down...') + exit(0) + + print('Source collection is ' + list(od.keys())[0]) + + od = {'001': None, '002': None, '003': None} + + # Defining source points + source_point_cloud = pointclouds[list(od.keys())[0]] + while True: + source_picked_points = pickPoints(source_point_cloud) + # Conditions of alignment + if not (len(source_picked_points) >= 3): + print('\nYou have chosen less than 3 points in the pointcloud, please redo it.') + continue + else: + break + + for collection_key, collection in od.items(): + target_point_cloud = pointclouds[collection_key] + # While cycle to force the user to choose the correct number of points + while True: + # Align the pointclouds + print('\nAligning pointclouds of collection ' + collection_key) + target_picked_points = pickPoints(target_point_cloud) + + # Conditions of alignment + if not (len(source_picked_points) >= 3 and len(target_picked_points) >= 3): + print('\nYou have chosen less than 3 points in the pointcloud, please redo it.') + continue + elif not (len(source_picked_points) == len(target_picked_points)): + print(f'\nYou have chosen {len(source_picked_points)} and {len(target_picked_points)} points, which needed to be equal, please redo them.') + continue + else: + break + + corr = np.zeros((len(source_picked_points), 2)) + corr[:, 0] = source_picked_points + corr[:, 1] = target_picked_points + + # Estimate rough transformation using correspondences + print("Compute a rough transform using the correspondences given by user") + p2p = o3d.pipelines.registration.TransformationEstimationPointToPoint() + T_target_to_source = p2p.compute_transformation(source_point_cloud, target_point_cloud, + o3d.utility.Vector2iVector(corr)) + + # Aligning using ICP + print('T_target_to_source = \n' + str(T_target_to_source)) + reg_p2p = ICPTransformations(source_point_cloud, target_point_cloud, threshold, T_target_to_source, show_images) + transforms[collection_key] = reg_p2p.transformation + print(reg_p2p.transformation) + + # Saving the dataset + print(Fore.YELLOW + '\nSaving odometry transformations' + Style.RESET_ALL) + saveICPTransforms(dataset, transforms, filename_results_json, od) + + + +if __name__ == '__main__': + main() \ No newline at end of file