├── .gitignore ├── CMakeLists.txt ├── INSTALL.md ├── LICENSE ├── README.md ├── bag └── placeforrosbag.txt ├── config ├── data_generation.json ├── data_params.yaml ├── training_config.json └── vehicle_sim.yaml ├── img └── example.jpg ├── iplanner ├── camera_intrinsic │ └── robot_intrinsic.txt ├── data │ ├── CollectedData │ │ └── placeholder.txt │ ├── TrainingData │ │ └── placeholder.txt │ ├── collect_list.txt │ ├── placeholder.txt │ └── training_list.txt ├── data_generation.py ├── dataloader.py ├── esdf_mapping.py ├── ip_algo.py ├── models │ └── placeholder.txt ├── percept_net.py ├── planner_net.py ├── rosutil.py ├── torchutil.py ├── training_run.py ├── traj_cost.py ├── traj_opt.py ├── traj_viz.py └── tsdf_map.py ├── iplanner_env.yml ├── launch ├── data_collector.launch ├── iplanner.launch └── iplanner_viz.launch ├── package.xml ├── rviz └── default.rviz └── src ├── data_collect_node.py ├── iplanner_node.py └── iplanner_viz.py /.gitignore: -------------------------------------------------------------------------------- 1 | build 2 | 3 | devel 4 | 5 | install 6 | 7 | logs 8 | 9 | models/ 10 | 11 | .catkin_workspace 12 | 13 | .catkin_tools 14 | 15 | .vscode 16 | 17 | *.png 18 | 19 | *.ply 20 | 21 | *.out 22 | 23 | *.pt 24 | 25 | __pycache__ 26 | 27 | data/ 28 | 29 | *~ 30 | 31 | *.bag 32 | 33 | .DS_Store 34 | 35 | wandb/ -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 2.8.3) 2 | project(iplanner_node) 3 | 4 | ## Compile as C++11, supported in ROS Kinetic and newer 5 | # add_compile_options(-std=c++11) 6 | # set(CMAKE_CXX_STANDARD 11) 7 | # list(APPEND CMAKE_PREFIX_PATH "/usr/lib/libtorch") 8 | 9 | ## Find catkin macros and libraries 10 | ## if COMPONENTS list like find_package(catkin REQUIRED COMPONENTS xyz) 11 | ## is used, also find other catkin packages 12 | find_package(catkin REQUIRED COMPONENTS 13 | roscpp 14 | rospy 15 | std_msgs 16 | image_transport 17 | sensor_msgs 18 | geometry_msgs 19 | message_generation 20 | ) 21 | 22 | 23 | # find_package(Torch REQUIRED) 24 | # find_package(OpenCV REQUIRED) 25 | 26 | ## System dependencies are found with CMake's conventions 27 | find_package(Boost REQUIRED COMPONENTS system) 28 | 29 | 30 | 31 | ## Generate messages in the 'msg' folder 32 | # add_message_files( 33 | # FILES 34 | # Msg1.msg 35 | # Msg2.msg 36 | # ) 37 | 38 | ## Generate services in the 'srv' folder 39 | # add_service_files( 40 | # FILES 41 | # Service1.srv 42 | # Service2.srv 43 | # ) 44 | 45 | ## Generate actions in the 'action' folder 46 | # add_action_files( 47 | # FILES 48 | # Action1.action 49 | # Action2.action 50 | # ) 51 | 52 | ## Generate added messages and services with any dependencies listed here 53 | # generate_messages( 54 | # DEPENDENCIES 55 | # geometry_msgs 56 | # std_msgs 57 | # ) 58 | 59 | catkin_package( 60 | # INCLUDE_DIRS include 61 | CATKIN_DEPENDS 62 | roscpp 63 | rospy 64 | std_msgs 65 | geometry_msgs 66 | # DEPENDS system_lib 67 | ) 68 | 69 | catkin_install_python(PROGRAMS src/data_collect_node.py 70 | DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} 71 | ) 72 | 73 | catkin_install_python(PROGRAMS src/iplanner_node.py 74 | DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} 75 | ) 76 | 77 | catkin_install_python(PROGRAMS src/iplanner_viz.py 78 | DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} 79 | ) 80 | -------------------------------------------------------------------------------- /INSTALL.md: -------------------------------------------------------------------------------- 1 | # iPlanner and Imperative Training Installation Guide 2 | 3 | This guide provides step-by-step instructions on how to setup your environment and install all necessary packages required for iPlanner and Imperative Training. 4 | 5 | ## Manual Installation Steps 6 | 7 | If you're opting for a manual installation, follow these steps: 8 | 9 | 1. Create a new conda environment named 'iplanner' with Python 3.8 using the command: 10 | ```bash 11 | conda create -n iplanner python=3.8 12 | ``` 13 | 14 | 2. Activate the 'iplanner' conda environment with the following command: 15 | ```bash 16 | conda activate iplanner 17 | ``` 18 | 19 | 3. Install PyTorch and Torchvision using the command: 20 | ```bash 21 | conda install pytorch torchvision pytorch-cuda=11.8 -c pytorch -c nvidia 22 | ``` 23 | 24 | 4. Install numpy version 1.23.5 using conda: 25 | ```bash 26 | conda install -c anaconda numpy==1.23.5 nbconvert 27 | ``` 28 | 29 | 5. Install rospkg, wandb, and PyYAML from the conda-forge channel using the command: 30 | ```bash 31 | conda install -c conda-forge rospkg wandb pyyaml==6.0 32 | ``` 33 | 34 | 6. Install additional necessary packages using pip3: 35 | ```bash 36 | pip3 install pypose open3d opencv-python rosnumpy 37 | ``` 38 | 39 | ## Installation Using a YAML File 40 | 41 | As an alternative to the manual installation, you can set up your environment and install all necessary packages at once using the `iplanner_env.yaml` file: 42 | 43 | 1. Create and activate the environment from the `iplanner_env.yml` file using the command: 44 | ```bash 45 | conda env create -f iplanner_env.yml 46 | ``` 47 | 48 | If you encounter any issues during the installation process, please open an issue on the project's GitHub page. 49 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Fan Yang 4 | Robotic Systems Lab, ETH Zurich 5 | 6 | Permission is hereby granted, free of charge, to any person obtaining a copy 7 | of this software and associated documentation files (the "Software"), to deal 8 | in the Software without restriction, including without limitation the rights 9 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 10 | copies of the Software, and to permit persons to whom the Software is 11 | furnished to do so, subject to the following conditions: 12 | 13 | The above copyright notice and this permission notice shall be included in all 14 | copies or substantial portions of the Software. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 17 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 18 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 19 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 20 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 21 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 22 | SOFTWARE. 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Imperative Path Planner (iPlanner) 2 | 3 | ## Overview 4 | 5 | Welcome to the [**iPlanner: Imperative Path Planning**](https://arxiv.org/abs/2302.11434) code repository. The iPlanner is trained via an innovative Imperative Learning Approach and exclusively uses front-facing depth images for local path planning. 6 | 7 | A video showing the functionalities of iPlanner is available here: [**Video**](https://youtu.be/-IfjSW0wPBI) 8 | 9 | **Keywords:** Navigation, Local Planning, Imperative Learning 10 | 11 | ### License 12 | 13 | This code is released under the MIT License. 14 | 15 | **Author: Fan Yang
16 | Maintainer: Fan Yang, fanyang1@ethz.ch** 17 | 18 | The iPlanner package has been tested under ROS Noetic on Ubuntu 20.04. This is research code, and any fitness for a particular purpose is disclaimed. 19 | 20 |

21 | Method 22 |

23 | 24 | ## Installation 25 | 26 | #### Dependencies 27 | 28 | To run iPlanner, you need to install [PyTorch](https://pytorch.org/). We recommend using [Anaconda](https://docs.anaconda.com/anaconda/install/index.html) for installation. Check the official website for installation instructions for Anaconda and PyTorch accordingly. 29 | 30 | Please follow the instructions provided in the `INSTALL.md` file to set up your environment and install the necessary packages. You can find the `INSTALL.md` file in the root directory of the project. 31 | 32 | #### Simulation Environment Setup 33 | 34 | Please refer to the Autonomous Exploration Development Environment repository for setting up the Gazebo simulation Environment: [Website](https://www.cmu-exploration.com/), switch to the branch **noetic_rgbd_camera**. 35 | 36 | #### Building 37 | 38 | To build the repository and set up the right Python version for running, use the command below: 39 | 40 | catkin build iplanner_node -DPYTHON_EXECUTABLE=$(which python) 41 | 42 | The Python3 should be the Python version you set up before with Torch and PyPose ready. If using the Anaconda environment, activate the conda env and check 43 | 44 | which python 45 | 46 | 47 | ## Training 48 | Go to the **iplanner** folder 49 | 50 | cd /iplanner 51 | 52 | #### Pre-trained Network and Training Data 53 | Download the pre-trained network weights `plannernet.pt` [here](https://drive.google.com/file/d/1UD11sSlOZlZhzij2gG_OmxbBN4WxVsO_/view?usp=share_link) and put it into the **models** folder. Noted this pre-trained network has not been adapted to real-world data. 54 | 55 | You can also collect data yourself either inside the simulation environment or in the real-world. Launch the **data_collect_node** 56 | 57 | roslaunch iplanner_node data_collector.launch 58 | 59 | Provide the information for the necessary topics listed in `config/data_params.yaml`. The collected data will be put into the folder `data/CollectedData`, and generate folders for different environments that you can specify in `config/data_params.yaml` under **env_name**. 60 | 61 | For each of the environments, the data contains the structure of: 62 | 63 | Environment Data 64 | ├── camera 65 | | ├── camera.png 66 | │ └── split.pt 67 | ├── camera_extrinsic.txt 68 | ├── cloud.ply 69 | ├── color_intrinsic.txt 70 | ├── depth 71 | | ├── depth.png 72 | │ └── split.pt 73 | ├── depth_intrinsic.txt 74 | ├── maps 75 | │ ├── cloud 76 | │ │ └── tsdf1_cloud.txt 77 | │ ├── data 78 | │ │ ├── tsdf1 79 | ├── data 80 | │ │ └── tsdf1_map.txt 81 | │ └── params 82 | │ └── tsdf1_param.txt 83 | └── odom_ground_truth.txt 84 | 85 | You can download the example data we provided using the Google Drive link [here](https://drive.google.com/file/d/1bUN7NV7arMM8ASA2pTJ8hvdkc5N3qoJw/view?usp=sharing). 86 | 87 | #### Generating Training Data 88 | 89 | Navigate to the iplanner folder within your project using the following command: 90 | 91 | cd <>/src/iPlanner/iplanner 92 | 93 | Run the Python script to generate the training data. The environments for which data should be generated are specified in the file `collect_list.txt`. You can modify the data generation parameters in the `config/data_generation.json` file. 94 | 95 | python data_generation.py 96 | 97 | Once you have the training data ready, use the following command to start the training process. You can specify different training parameters in the `config/training_config.json` file. 98 | 99 | python training_run.py 100 | 101 | ## Run iPlanner ROS node 102 | 103 | Launch the simulation environment without the default local planner 104 | 105 | roslaunch vehicle_simulator simulation_env.launch 106 | 107 | Run the iPlanner ROS node without visualization: 108 | 109 | roslaunch iplanner_node iplanner.launch 110 | 111 | Or run the iPlanner ROS node with visualization: 112 | 113 | roslaunch iplanner_node iplanner_viz.launch 114 | 115 | ### Path Following 116 | To ensure the planner executes the planned path correctly, you need to run an independent controller or path follower. Follow the steps below to set up the path follower using the provided launch file from the iplanner repository: 117 | 118 | Download the default iplanner_path_follower into your workspace. Navigate to your workspace's source directory using the following command: 119 | 120 | cd <>/src 121 | 122 | Then clone the repository: 123 | 124 | git clone https://github.com/MichaelFYang/iplanner_path_follow.git 125 | 126 | Compile the path follower using the following command: 127 | 128 | catkin build iplanner_path_follow 129 | 130 | Please note that this repository is a fork of the path following component from [CMU-Exploration](https://www.cmu-exploration.com/). You are welcome to explore and try different controllers or path followers suitable for your specific robot platform. 131 | 132 | ### Waypoint Navigation 133 | To send the waypoint through Rviz, please download the rviz waypoint plugin. Navigate to your workspace's source directory using the following command: 134 | 135 | cd <>/src 136 | 137 | Then clone the repository: 138 | 139 | git clone https://github.com/MichaelFYang/waypoint_rviz_plugin.git 140 | 141 | Compile the waypoint rviz plugin using the following command: 142 | 143 | catkin build waypoint_rviz_plugin 144 | 145 | 146 | ### SmartJoystick 147 | 148 | Press the **LB** button on the joystick, when seeing the output on the screen: 149 | 150 | Switch to Smart Joystick mode ... 151 | 152 | Now the smartjoystick feature is enabled. It takes the joystick command as motion intention and runs the iPlanner in the background for low-level obstacle avoidance. 153 | 154 | ## Config files 155 | 156 | The params file **`data_params.yaml`** is for data collection 157 | 158 | * **vehicle_sim.yaml** The config file contains: 159 | - **`main_freq`** The ROS node running frequency 160 | - **`odom_associate_id`** Depending on different SLAM setup, the odometry base may not be set under robot base frame 161 | 162 | The params file **`vehicle_sim.yaml`** is for iPlanner ROS node 163 | 164 | * **vehicle_sim.yaml** The config file contains: 165 | - **`main_freq`** The ROS node running frequency 166 | - **`image_flap`** Depending on the camera setup, it may require to flip the image upside down or not 167 | - **`crop_size`** The size to crop the incoming camera images 168 | - **`is_fear_act`** Using the predicted collision possibility value to stop 169 | - **`joyGoal_scale`** The max distance of goal sent by joystick in smart joystick model 170 | 171 | 172 | ## Reference 173 | 174 | If you utilize this codebase in your research, we kindly request you to reference our work. You can cite us as follows: 175 | 176 | - Yang, F., Wang, C., Cadena, C., & Hutter, M. (2023). iPlanner: Imperative Path Planning. Robotics: Science and Systems Conference (RSS). Daegu, Republic of Korea, July 2023. 177 | 178 | Bibtex: 179 | 180 | @INPROCEEDINGS{Yang-RSS-23, 181 | AUTHOR = {Fan Yang AND Chen Wang AND Cesar Cadena AND Marco Hutter}, 182 | TITLE = {{iPlanner: Imperative Path Planning}}, 183 | BOOKTITLE = {Proceedings of Robotics: Science and Systems}, 184 | YEAR = {2023}, 185 | ADDRESS = {Daegu, Republic of Korea}, 186 | MONTH = {July}, 187 | DOI = {10.15607/RSS.2023.XIX.064} 188 | } 189 | 190 | 191 | ## Author 192 | 193 | This codebase has been developed and maintained by [Fan Yang](https://github.com/MichaelFYang). Should you have any queries or require further assistance, you may reach out to him at fanyang1@ethz.ch 194 | -------------------------------------------------------------------------------- /bag/placeforrosbag.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leggedrobotics/iPlanner/4a8d823ff9d09c3f626b727e7e00484b38f80d49/bag/placeforrosbag.txt -------------------------------------------------------------------------------- /config/data_generation.json: -------------------------------------------------------------------------------- 1 | { 2 | "folder_name": "CollectedData", 3 | "outfolder_name": "TrainingData", 4 | "voxel_size": 0.05, 5 | "robot_size": 0.3, 6 | "image_type": "depth", 7 | "map_name": "tsdf1", 8 | "is_max_iter": true, 9 | "max_depth_range": 10.0, 10 | "is_flat_ground": true, 11 | "is_visualize": false 12 | } 13 | -------------------------------------------------------------------------------- /config/data_params.yaml: -------------------------------------------------------------------------------- 1 | # images and odom topics 2 | main_freq: 2.5 3 | depth_topic: /rgbd_camera/depth/image 4 | color_topic: /rgbd_camera/color/image 5 | scan_topic: /velodyne_points 6 | odom_topic: /state_estimation 7 | 8 | # camaera info topics 9 | depth_info_topic: /rgbd_camera/depth/camera_info 10 | color_info_topic: /rgbd_camera/color/camera_info 11 | camera_frame_id: rgbd_camera 12 | scan_frame_id: sensor 13 | base_frame_id: vehicle 14 | odom_associate_id: vehicle 15 | 16 | # environment name 17 | env_name: test_env -------------------------------------------------------------------------------- /config/training_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "dataConfig": { 3 | "data-root": "data", 4 | "env-id": "training_list.txt", 5 | "env_type": "TrainingData", 6 | "crop-size": [360, 640], 7 | "max-camera-depth": 10.0 8 | }, 9 | "modelConfig": { 10 | "model-save": "models/plannernet.pt", 11 | "resume": false, 12 | "in-channel": 16, 13 | "knodes": 5, 14 | "goal-step": 5, 15 | "max-episode-length": 25 16 | }, 17 | "trainingConfig": { 18 | "training": true, 19 | "lr": 0.0001, 20 | "factor": 0.1, 21 | "min-lr": 0.000001, 22 | "patience": 4, 23 | "epochs": 50, 24 | "batch-size": 128, 25 | "w-decay": 0.001, 26 | "num-workers": 2, 27 | "gpu-id": 0 28 | }, 29 | "logConfig": { 30 | "log-save": "models/log-", 31 | "test-env-id": 0, 32 | "visual-number": 10 33 | }, 34 | "sensorConfig": { 35 | "camera-tilt": 0.2, 36 | "sensor-offsetX-ANYmal": 0.4, 37 | "fear-ahead-dist": 2.0 38 | } 39 | } 40 | -------------------------------------------------------------------------------- /config/vehicle_sim.yaml: -------------------------------------------------------------------------------- 1 | # planning 2 | main_freq: 15 3 | image_flip: False 4 | conv_dist: 0.25 5 | is_sim: True 6 | # newwork model 7 | model_save: /models/plannernet.pt 8 | crop_size: [360, 640] 9 | uint_type: False 10 | # images and gaol waypoint topics 11 | depth_topic: /rgbd_camera/depth/image 12 | goal_topic: /way_point 13 | path_topic: /path 14 | image_topic: /iplanner_image 15 | camera_tilt: 0.0 16 | depth_max: 10.0 17 | # frame ids 18 | robot_id: vehicle 19 | world_id: map 20 | # fear reaction 21 | is_fear_act: True 22 | buffer_size: 3 23 | angular_thred: 0.5 24 | track_dist: 0.5 25 | # sensor offset 26 | sensor_offset_x: 0.0 27 | sensor_offset_y: 0.0 28 | # smart joystick 29 | joyGoal_scale: 5.0 -------------------------------------------------------------------------------- /img/example.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leggedrobotics/iPlanner/4a8d823ff9d09c3f626b727e7e00484b38f80d49/img/example.jpg -------------------------------------------------------------------------------- /iplanner/camera_intrinsic/robot_intrinsic.txt: -------------------------------------------------------------------------------- 1 | (205.46963709898583, 0.0, 320.5, -0.0, 0.0, 205.46963709898583, 180.5, 0.0, 0.0, 0.0, 1.0, 0.0) -------------------------------------------------------------------------------- /iplanner/data/CollectedData/placeholder.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leggedrobotics/iPlanner/4a8d823ff9d09c3f626b727e7e00484b38f80d49/iplanner/data/CollectedData/placeholder.txt -------------------------------------------------------------------------------- /iplanner/data/TrainingData/placeholder.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leggedrobotics/iPlanner/4a8d823ff9d09c3f626b727e7e00484b38f80d49/iplanner/data/TrainingData/placeholder.txt -------------------------------------------------------------------------------- /iplanner/data/collect_list.txt: -------------------------------------------------------------------------------- 1 | 2n8kARJN3HM 2 | B6ByNegPMKs_tilt 3 | indoor 4 | B6ByNegPMKs 5 | 2n8kARJN3HM_tilt 6 | ur6pFq6Qu1A 7 | ur6pFq6Qu1A_tilt 8 | forest 9 | Vvot9Ly1tCj 10 | Vvot9Ly1tCj_tilt 11 | garage 12 | Matterport 13 | Matterport_tilt 14 | JeFG25nYj2p 15 | JeFG25nYj2p_tilt 16 | vyrNrziPKCB 17 | tunnel 18 | -------------------------------------------------------------------------------- /iplanner/data/placeholder.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leggedrobotics/iPlanner/4a8d823ff9d09c3f626b727e7e00484b38f80d49/iplanner/data/placeholder.txt -------------------------------------------------------------------------------- /iplanner/data/training_list.txt: -------------------------------------------------------------------------------- 1 | 2n8kARJN3HM 2 | B6ByNegPMKs_tilt 3 | indoor 4 | B6ByNegPMKs 5 | 2n8kARJN3HM_tilt 6 | ur6pFq6Qu1A 7 | campus_01 8 | ur6pFq6Qu1A_tilt 9 | forest 10 | Vvot9Ly1tCj 11 | Vvot9Ly1tCj_tilt 12 | garage 13 | Matterport 14 | Matterport_tilt 15 | JeFG25nYj2p 16 | campus_02 17 | JeFG25nYj2p_tilt 18 | vyrNrziPKCB 19 | vyrNrziPKCB_tilt 20 | tunnel -------------------------------------------------------------------------------- /iplanner/data_generation.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import os 11 | import json 12 | from tsdf_map import TSDF_Map 13 | from esdf_mapping import TSDF_Creator, DepthReconstruction 14 | 15 | if __name__ == '__main__': 16 | 17 | root_folder = os.getenv('EXPERIMENT_DIRECTORY', os.getcwd()) 18 | # Load parameters from json file 19 | with open(os.path.join(os.path.dirname(root_folder), 'config', 'data_generation.json')) as json_file: 20 | parameters = json.load(json_file) 21 | 22 | folder_name = parameters.get('folder_name', "CollectedData") 23 | folder_path = os.path.join(*[root_folder, "data"]) 24 | ids_path = os.path.join(folder_path, "collect_list.txt") 25 | 26 | if not folder_name == "": 27 | folder_path = os.path.join(folder_path, folder_name) 28 | env_list = [] 29 | with open(ids_path) as f: 30 | lines = f.readlines() 31 | for line in lines: 32 | env_list.append(line.rstrip()) 33 | print("Env List: ", env_list) 34 | 35 | outfolder_name = parameters.get('outfolder_name', "TrainingData") 36 | output_folder = os.path.join(*[root_folder, "data", outfolder_name]) 37 | 38 | image_type = parameters.get('image_type', "depth") 39 | voxel_size = parameters.get('voxel_size', 0.05) 40 | robot_size = parameters.get('robot_size', 0.3) # the inflated robot radius 41 | map_name = parameters.get('map_name', "tsdf1") 42 | is_max_iter = parameters.get('is_max_iter', True) 43 | max_depth_range = parameters.get('max_depth_range', 10.0) 44 | is_flat_ground = parameters.get('is_flat_ground', True) 45 | is_visualize = parameters.get('is_visualize', False) 46 | 47 | for env_name in env_list: 48 | root_path = os.path.join(*[folder_path, env_name]) 49 | image_path = os.path.join(root_path, 'depth') 50 | 51 | total_data_n = len([name for name in os.listdir(image_path) if os.path.isfile(os.path.join(image_path, name))]) 52 | print("================= Reconstruction of env: %s =================="%(env_name)) 53 | out_path = os.path.join(output_folder, env_name) 54 | 55 | depth_constructor = DepthReconstruction(root_path, out_path, 0, 100, voxel_size*0.9, max_depth_range, is_max_iter) 56 | depth_constructor.depth_map_reconstruction(is_flat_ground=is_flat_ground) 57 | depth_constructor.save_reconstructed_data(image_type=image_type) 58 | avg_height = depth_constructor.avg_height 59 | print("Average Height: ", avg_height) 60 | if is_visualize: 61 | depth_constructor.show_point_cloud() 62 | 63 | # Construct the 2D cost map 64 | tsdf_creator = TSDF_Creator(out_path, voxel_size=voxel_size, robot_size=robot_size, robot_height=avg_height) 65 | tsdf_creator.read_point_from_file("cloud.ply") 66 | data, coord, params = tsdf_creator.create_TSDF_map() 67 | if is_visualize: 68 | tsdf_creator.visualize_cloud(tsdf_creator.obs_pcd) 69 | tsdf_creator.visualize_cloud(tsdf_creator.free_pcd) 70 | 71 | # Save the esdf map 72 | tsdf_map = TSDF_Map() 73 | tsdf_map.DirectLoadMap(data, coord, params) 74 | tsdf_map.SaveTSDFMap(out_path, map_name) 75 | if is_visualize: 76 | tsdf_map.ShowTSDFMap(cost_map=True) -------------------------------------------------------------------------------- /iplanner/dataloader.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import os 11 | import PIL 12 | import torch 13 | import numpy as np 14 | import pypose as pp 15 | from PIL import Image 16 | from pathlib import Path 17 | from random import sample 18 | from operator import itemgetter 19 | from torch.utils.data import Dataset, DataLoader 20 | 21 | torch.set_default_dtype(torch.float32) 22 | 23 | 24 | class MultiEpochsDataLoader(DataLoader): 25 | 26 | def __init__(self, *args, **kwargs): 27 | super().__init__(*args, **kwargs) 28 | self._DataLoader__initialized = False 29 | self.batch_sampler = _RepeatSampler(self.batch_sampler) 30 | self._DataLoader__initialized = True 31 | self.iterator = super().__iter__() 32 | 33 | def __len__(self): 34 | return len(self.batch_sampler.sampler) 35 | 36 | def __iter__(self): 37 | for i in range(len(self)): 38 | yield next(self.iterator) 39 | 40 | 41 | class _RepeatSampler(object): 42 | """ Sampler that repeats forever. 43 | Args: 44 | sampler (Sampler) 45 | """ 46 | 47 | def __init__(self, sampler): 48 | self.sampler = sampler 49 | 50 | def __iter__(self): 51 | while True: 52 | yield from iter(self.sampler) 53 | 54 | class PlannerData(Dataset): 55 | def __init__(self, root, max_episode, goal_step, train, ratio=0.9, max_depth=10.0, sensorOffsetX=0.0, transform=None, is_robot=True): 56 | super().__init__() 57 | self.transform = transform 58 | self.is_robot = is_robot 59 | self.max_depth = max_depth 60 | img_filename_list = [] 61 | img_path = os.path.join(root, "depth") 62 | img_filename_list = [str(s) for s in Path(img_path).rglob('*.png')] 63 | img_filename_list.sort(key = lambda x : int(x.split("/")[-1][:-4])) 64 | 65 | odom_path = os.path.join(root, "odom_ground_truth.txt") 66 | odom_list = [] 67 | offset_T = pp.identity_SE3() 68 | offset_T.tensor()[0] = sensorOffsetX 69 | with open(odom_path) as f: 70 | lines = f.readlines() 71 | for line in lines: 72 | odom = np.fromstring(line[1:-2], dtype=np.float32, sep=', ') 73 | odom = pp.SE3(odom) 74 | if is_robot: 75 | odom = odom @ offset_T 76 | odom_list.append(odom) 77 | 78 | N = len(odom_list) 79 | 80 | self.img_filename = [] 81 | self.odom_list = [] 82 | self.goal_list = [] 83 | 84 | for ahead in range(1, max_episode+1, goal_step): 85 | for i in range(N): 86 | odom = odom_list[i] 87 | goal = odom_list[min(i+ahead, N-1)] 88 | goal = (pp.Inv(odom) @ goal) 89 | # gp = goal.tensor() 90 | # if (gp[0] > 1.0 and gp[1]/gp[0] < 1.2 and gp[1]/gp[0] > -1.2 and torch.norm(gp[:3]) > 1.0): 91 | self.img_filename.append(img_filename_list[i]) 92 | self.odom_list.append(odom.tensor()) 93 | self.goal_list.append(goal.tensor()) 94 | 95 | N = len(self.odom_list) 96 | 97 | indexfile = os.path.join(img_path, 'split.pt') 98 | is_generate_split = True; 99 | if os.path.exists(indexfile): 100 | train_index, test_index = torch.load(indexfile) 101 | if len(train_index)+len(test_index) == N: 102 | is_generate_split = False 103 | else: 104 | print("Data changed! Generate a new split file") 105 | if (is_generate_split): 106 | indices = range(N) 107 | train_index = sample(indices, int(ratio*N)) 108 | test_index = np.delete(indices, train_index) 109 | torch.save((train_index, test_index), indexfile) 110 | 111 | if train == True: 112 | self.img_filename = itemgetter(*train_index)(self.img_filename) 113 | self.odom_list = itemgetter(*train_index)(self.odom_list) 114 | self.goal_list = itemgetter(*train_index)(self.goal_list) 115 | else: 116 | self.img_filename = itemgetter(*test_index)(self.img_filename) 117 | self.odom_list = itemgetter(*test_index)(self.odom_list) 118 | self.goal_list = itemgetter(*test_index)(self.goal_list) 119 | 120 | assert len(self.odom_list) == len(self.img_filename), "odom numbers should match with image numbers" 121 | 122 | 123 | def __len__(self): 124 | return len(self.img_filename) 125 | 126 | def __getitem__(self, idx): 127 | image = Image.open(self.img_filename[idx]) 128 | if self.is_robot: 129 | image = np.array(image.transpose(PIL.Image.ROTATE_180)) 130 | else: 131 | image = np.array(image) 132 | image[~np.isfinite(image)] = 0.0 133 | image = (image / 1000.0).astype("float32") 134 | image[image > self.max_depth] = 0.0 135 | # DEBUG show image 136 | # img = Image.fromarray((image * 255 / np.max(image)).astype('uint8')) 137 | # img.show() 138 | image = Image.fromarray(image) 139 | image = self.transform(image).expand(3, -1, -1) 140 | 141 | return image, self.odom_list[idx], self.goal_list[idx] -------------------------------------------------------------------------------- /iplanner/esdf_mapping.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import os 11 | import cv2 12 | import math 13 | import numpy as np 14 | import open3d as o3d 15 | from scipy import ndimage 16 | from scipy.ndimage import gaussian_filter 17 | from scipy.spatial.transform import Rotation as R 18 | 19 | 20 | class CloudUtils: 21 | @staticmethod 22 | def create_open3d_cloud(points, voxel_size): 23 | pcd = o3d.geometry.PointCloud() 24 | pcd.points = o3d.utility.Vector3dVector(points) 25 | pcd = pcd.voxel_down_sample(voxel_size) 26 | return pcd 27 | 28 | @staticmethod 29 | def extract_cloud_from_image(P_matrix, im, T, min_dist=0.2, max_dist=50, scale=1000.0): 30 | p_inv = np.linalg.inv(P_matrix) 31 | im = im / scale 32 | im[immax_dist] = 1e-3 34 | 35 | T_z = np.concatenate((T, np.expand_dims(1.0/im, axis=0)), axis=0).reshape(4, -1) 36 | P = np.multiply(im.reshape(1, -1), p_inv.dot(T_z)).T[:,:3] 37 | return P 38 | 39 | class CameraUtils: 40 | @staticmethod 41 | def compute_pixel_tensor(x_nums, y_nums): 42 | T = np.zeros([3, x_nums, y_nums]) 43 | for u in range(x_nums): 44 | for v in range(y_nums): 45 | T[:, u, v] = np.array([u, v, 1.0]) 46 | return T 47 | 48 | @staticmethod 49 | def compute_e_matrix(odom, is_flat_ground, cameraR, cameraT): 50 | Rc = R.from_quat(odom[3:]) 51 | if is_flat_ground: 52 | euler = Rc.as_euler('xyz', degrees=False) 53 | euler[1] = 0.0 54 | Rc = R.from_euler("xyz", euler, degrees=False) 55 | Rc = Rc * cameraR 56 | C = (odom[:3] + cameraT).reshape(-1,1) 57 | Rc_t = Rc.as_matrix().T 58 | E = np.concatenate((Rc_t, -Rc_t.dot(C)), axis=1) 59 | E = np.concatenate((E, np.array([0.0, 0.0, 0.0, 1.0]).reshape(1,-1)), axis=0) 60 | return E 61 | 62 | class DataUtils: 63 | @staticmethod 64 | def read_odom_list(odom_path): 65 | odom_list = [] 66 | avg_height = 0.0 67 | with open(odom_path) as f: 68 | lines = f.readlines() 69 | for line in lines: 70 | odom = np.fromstring(line[1:-1], dtype=float, sep=', ') 71 | avg_height = avg_height + odom[2] 72 | odom_list.append(list(odom)) 73 | avg_height = avg_height / len(lines) 74 | return odom_list, avg_height 75 | 76 | @staticmethod 77 | def read_intrinsic(intrinsic_path): 78 | with open(intrinsic_path) as f: 79 | lines = f.readlines() 80 | elems = np.fromstring(lines[0][1:-1], dtype=float, sep=', ') 81 | if len(elems) == 12: 82 | P = np.array(elems).reshape(3, 4) 83 | K = np.concatenate((P, np.array([0.0, 0.0, 0.0, 1.0]).reshape(1,-1)), axis=0) 84 | else: 85 | K = np.array(elems).reshape(4, 4) 86 | return K 87 | 88 | @staticmethod 89 | def read_extrinsic(extrinsic_path): 90 | with open(extrinsic_path) as f: 91 | lines = f.readlines() 92 | elems = np.fromstring(lines[0][1:-1], dtype=float, sep=', ') 93 | CR = R.from_quat(elems[:4]) 94 | CT = np.array(elems[4:]) 95 | return CR, CT 96 | 97 | @staticmethod 98 | def load_point_cloud(path): 99 | return o3d.io.read_point_cloud(path) 100 | 101 | @staticmethod 102 | def prepare_output_folders(out_path, image_type): 103 | depth_im_path = os.path.join(out_path, image_type) 104 | if not os.path.exists(out_path): 105 | os.makedirs(out_path) 106 | os.makedirs(depth_im_path) 107 | os.makedirs(os.path.join(out_path, "maps", "cloud")) 108 | os.makedirs(os.path.join(out_path, "maps", "data")) 109 | os.makedirs(os.path.join(out_path, "maps", "params")) 110 | elif os.path.exists(depth_im_path): # remove existing files 111 | for efile in os.listdir(depth_im_path): 112 | os.remove(os.path.join(depth_im_path, efile)) 113 | return None 114 | 115 | @staticmethod 116 | def load_images(start_id, end_id, root_path, image_type): 117 | im_arr_list = [] 118 | im_path = os.path.join(root_path, image_type) 119 | for idx in range(start_id, end_id): 120 | path = os.path.join(im_path, str(idx) + ".png") 121 | im = cv2.imread(path, cv2.IMREAD_ANYDEPTH).T 122 | im_arr_list.append(im) 123 | print("total number of images for reconstruction: {}".format(len(im_arr_list))) 124 | return im_arr_list 125 | 126 | @staticmethod 127 | def save_images(out_path, im_arr_list, image_type, is_transpose=True): 128 | for idx, img in enumerate(im_arr_list): 129 | if is_transpose: 130 | img = img.T 131 | cv2.imwrite(os.path.join(out_path, image_type, f"{idx}.png"), img) 132 | return None 133 | 134 | @staticmethod 135 | def save_odom_list(out_path, odom_list, start_id, num_images): 136 | with open(os.path.join(out_path, "odom_ground_truth.txt"), 'w') as f: 137 | for i in range(start_id, start_id + num_images): 138 | f.write(str(odom_list[i]) + "\n") 139 | return None 140 | 141 | @staticmethod 142 | def save_extrinsic(out_path, cameraR, cameraT): 143 | with open(os.path.join(out_path, "camera_extrinsic.txt"), 'w') as f: 144 | f.write(str(list(cameraR.as_quat()) + list(cameraT)) + "\n") 145 | return None 146 | 147 | @staticmethod 148 | def save_intrinsic(out_path, K): 149 | with open(os.path.join(out_path, "depth_intrinsic.txt"), 'w') as f: 150 | f.write(str(K.flatten().tolist()) + "\n") 151 | return None 152 | 153 | @staticmethod 154 | def save_point_cloud(out_path, pcd): 155 | o3d.io.write_point_cloud(os.path.join(out_path, "cloud.ply"), pcd) # save point cloud 156 | return None 157 | 158 | class TSDF_Creator: 159 | def __init__(self, input_path, voxel_size, robot_height, robot_size, clear_dist=1.0): 160 | self.initialize_path_and_properties(input_path, voxel_size, robot_height, robot_size, clear_dist) 161 | self.initialize_point_clouds() 162 | 163 | def initialize_path_and_properties(self, input_path, voxel_size, robot_height, robot_size, clear_dist): 164 | self.input_path = input_path 165 | self.is_map_ready = False 166 | self.clear_dist = clear_dist 167 | self.voxel_size = voxel_size 168 | self.robot_height = robot_height 169 | self.robot_size = robot_size 170 | 171 | def initialize_point_clouds(self): 172 | self.obs_pcd = o3d.geometry.PointCloud() 173 | self.free_pcd = o3d.geometry.PointCloud() 174 | 175 | def update_point_cloud(self, P_obs, P_free, is_downsample=False): 176 | self.obs_pcd.points = o3d.utility.Vector3dVector(P_obs) 177 | self.free_pcd.points = o3d.utility.Vector3dVector(P_free) 178 | self.downsample_point_cloud(is_downsample) 179 | self.obs_points = np.asarray(self.obs_pcd.points) 180 | self.free_points = np.asarray(self.free_pcd.points) 181 | 182 | def read_point_from_file(self, file_name, is_filter=True): 183 | file_path = os.path.join(self.input_path, file_name) 184 | pcd_load = DataUtils.load_point_cloud(file_path) 185 | 186 | print("Running terrain analysis...") 187 | obs_p, free_p = self.terrain_analysis(np.asarray(pcd_load.points)) 188 | self.update_point_cloud(obs_p, free_p, is_downsample=True) 189 | 190 | if is_filter: 191 | obs_p = self.filter_cloud(self.obs_points, num_nbs=50, std_ratio=2.0) 192 | self.update_point_cloud(obs_p, free_p) 193 | 194 | self.update_map_params() 195 | 196 | def downsample_point_cloud(self, is_downsample): 197 | if is_downsample: 198 | self.obs_pcd = self.obs_pcd.voxel_down_sample(self.voxel_size) 199 | self.free_pcd = self.free_pcd.voxel_down_sample(self.voxel_size * 0.85) 200 | 201 | def update_map_params(self): 202 | self._handle_no_points() 203 | self._set_map_limits_and_start_coordinates() 204 | self._log_map_initialization() 205 | self.is_map_ready = True 206 | 207 | def terrain_analysis(self, input_points, ground_height=0.25): 208 | obs_points, free_points = self._initialize_point_arrays(input_points) 209 | obs_idx = free_idx = 0 210 | 211 | for p in input_points: 212 | p_height = p[2] 213 | if self._is_obstacle(p_height, ground_height): 214 | obs_points[obs_idx, :] = p 215 | obs_idx += 1 216 | elif self._is_free_space(p_height, ground_height): 217 | free_points[free_idx, :] = p 218 | free_idx += 1 219 | 220 | return obs_points[:obs_idx, :], free_points[:free_idx, :] 221 | 222 | def create_TSDF_map(self, sigma_smooth=2.5): 223 | if not self.is_map_ready: 224 | print("create tsdf map fails, no points received.") 225 | return 226 | 227 | free_map = np.ones([self.num_x, self.num_y]) 228 | obs_map = self._create_obstacle_map() 229 | 230 | # create free place map 231 | free_I = self._index_array_of_points(self.free_points) 232 | free_map = self._create_free_space_map(free_I, free_map, sigma_smooth) 233 | 234 | free_map[obs_map > 0.3] = 1.0 # re-assign obstacles if they are in free space 235 | print("occupancy map generation completed.") 236 | 237 | # Distance Transform 238 | tsdf_array = self._distance_transform_and_smooth(free_map, sigma_smooth) 239 | 240 | viz_points = np.concatenate((self.obs_points, self.free_points), axis=0) 241 | # TODO: Use true terrain analysis module 242 | ground_array = np.ones([self.num_x, self.num_y]) * 0.0 243 | 244 | return [tsdf_array, viz_points, ground_array], [self.start_x, self.start_y], [self.voxel_size, self.clear_dist] 245 | 246 | def filter_cloud(self, points, num_nbs=100, std_ratio=1.0): 247 | pcd = self._convert_to_point_cloud(points) 248 | filtered_pcd = self._remove_statistical_outliers(pcd, num_nbs, std_ratio) 249 | return np.asarray(filtered_pcd.points) 250 | 251 | def visualize_cloud(self, pcd): 252 | o3d.visualization.draw_geometries([pcd]) 253 | 254 | def _convert_to_point_cloud(self, points): 255 | pcd = o3d.geometry.PointCloud() 256 | pcd.points = o3d.utility.Vector3dVector(points) 257 | return pcd 258 | 259 | def _remove_statistical_outliers(self, pcd, num_nbs, std_ratio): 260 | filtered_pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=num_nbs, std_ratio=std_ratio) 261 | return filtered_pcd 262 | 263 | def _create_obstacle_map(self): 264 | obs_map = np.zeros([self.num_x, self.num_y]) 265 | obs_I = self._index_array_of_points(self.obs_points) 266 | for i in obs_I: 267 | obs_map[i[0], i[1]] = 1.0 268 | obs_map = gaussian_filter(obs_map, sigma=self.robot_size / self.voxel_size) 269 | obs_map /= np.max(obs_map + 1e-5) # normalize 270 | return obs_map 271 | 272 | def _create_free_space_map(self, free_I, free_map, sigma_smooth): 273 | for i in free_I: 274 | if i[0] >= 0 and i[0] < self.num_x and i[1] >= 0 and i[1] < self.num_y: 275 | free_map[i[0], i[1]] = 0 276 | free_map = gaussian_filter(free_map, sigma=sigma_smooth) 277 | free_map /= np.max(free_map) # normalize 278 | free_map[free_map < 0.7] = 0 # 0.683% is the probability of standard normal distribution 279 | return free_map 280 | 281 | def _distance_transform_and_smooth(self, free_map, sigma_smooth, is_log=True): 282 | dt_map = ndimage.distance_transform_edt(free_map) 283 | tsdf_array = gaussian_filter(dt_map, sigma=sigma_smooth) 284 | if is_log: 285 | tsdf_array = np.log(tsdf_array + 1.00001) 286 | return tsdf_array 287 | 288 | def _index_array_of_points(self, points): 289 | I = np.round((points[:, :2] - np.array([self.start_x, self.start_y])) / self.voxel_size).astype(int) 290 | return I 291 | 292 | def _initialize_point_arrays(self, input_points): 293 | return np.zeros(input_points.shape), np.zeros(input_points.shape) 294 | 295 | def _is_obstacle(self, p_height, ground_height): 296 | return (p_height > ground_height) and (p_height < self.robot_height * 1.5) 297 | 298 | def _is_free_space(self, p_height, ground_height): 299 | return p_height < ground_height and p_height > - ground_height 300 | 301 | def _handle_no_points(self): 302 | if (self.obs_points.shape[0] == 0): 303 | print("No points received.") 304 | return 305 | 306 | def _set_map_limits_and_start_coordinates(self): 307 | max_x, max_y, _ = np.amax(self.obs_points, axis=0) + self.clear_dist 308 | min_x, min_y, _ = np.amin(self.obs_points, axis=0) - self.clear_dist 309 | self.num_x = np.ceil((max_x - min_x) / self.voxel_size / 10).astype(int) * 10 310 | self.num_y = np.ceil((max_y - min_y) / self.voxel_size / 10).astype(int) * 10 311 | self.start_x = (max_x + min_x) / 2.0 - self.num_x / 2.0 * self.voxel_size 312 | self.start_y = (max_y + min_y) / 2.0 - self.num_y / 2.0 * self.voxel_size 313 | 314 | def _log_map_initialization(self): 315 | print("tsdf map initialized, with size: %d, %d" %(self.num_x, self.num_y)) 316 | 317 | 318 | class DepthReconstruction: 319 | def __init__(self, input_path, out_path, start_id, iters, voxel_size, max_range, is_max_iter=True): 320 | self._initialize_paths(input_path, out_path) 321 | self._initialize_parameters(voxel_size, max_range, is_max_iter) 322 | self._read_camera_params() 323 | 324 | # odom list read 325 | self.odom_list, self._avg_height = DataUtils.read_odom_list(self.input_path + "/odom_ground_truth.txt") 326 | 327 | N = len(self.odom_list) 328 | self.start_id = 0 if self.is_max_iter else start_id 329 | self.end_id = N if self.is_max_iter else min(start_id + iters, N) 330 | 331 | self.is_constructed = False 332 | print("Ready to read depth data.") 333 | 334 | # public methods 335 | def depth_map_reconstruction(self, is_output=False, is_flat_ground=False): 336 | self.im_arr_list = DataUtils.load_images(self.start_id, self.end_id, self.input_path, "depth") 337 | 338 | x_nums, y_nums = self.im_arr_list[0].shape 339 | T = CameraUtils.compute_pixel_tensor(x_nums, y_nums) 340 | pixel_nums = x_nums * y_nums 341 | 342 | print("start reconstruction...") 343 | self.points = np.zeros([(self.end_id - self.start_id + 1) * pixel_nums, 3]) 344 | 345 | for idx, im in enumerate(self.im_arr_list): 346 | odom = self.odom_list[idx + self.start_id].copy() 347 | if is_flat_ground: 348 | odom[2] = self._avg_height 349 | E = CameraUtils.compute_e_matrix(odom, is_flat_ground, self.cameraR, self.cameraT) 350 | P_matrix = self.K.dot(E) 351 | if is_output: 352 | print("Extracting points from image: ", idx + self.start_id) 353 | self.points[idx * pixel_nums: (idx + 1) * pixel_nums, :] = CloudUtils.extract_cloud_from_image( 354 | P_matrix, im, T, max_dist=self.max_range) 355 | 356 | print("creating open3d geometry point cloud...") 357 | self.pcd = CloudUtils.create_open3d_cloud(self.points, self.voxel_size) 358 | self.is_constructed = True 359 | print("construction completed.") 360 | 361 | def show_point_cloud(self): 362 | if not self.is_constructed: 363 | print("no reconstructed cloud") 364 | o3d.visualization.draw_geometries([self.pcd]) # visualize point cloud 365 | 366 | def save_reconstructed_data(self, image_type="depth"): 367 | if not self.is_constructed: 368 | print("save points failed, no reconstructed cloud!") 369 | 370 | print("save output files to: " + self.out_path) 371 | DataUtils.prepare_output_folders(self.out_path, image_type) 372 | 373 | DataUtils.save_images(self.out_path, self.im_arr_list, image_type) 374 | DataUtils.save_odom_list(self.out_path, self.odom_list, self.start_id, len(self.im_arr_list)) 375 | DataUtils.save_extrinsic(self.out_path, self.cameraR, self.cameraT) 376 | DataUtils.save_intrinsic(self.out_path, self.K) 377 | DataUtils.save_point_cloud(self.out_path, self.pcd) # save point cloud 378 | print("saved cost map data.") 379 | 380 | @property 381 | def avg_height(self): 382 | return self._avg_height 383 | 384 | # private methods 385 | def _initialize_paths(self, input_path, out_path): 386 | self.input_path = input_path 387 | self.out_path = out_path 388 | 389 | def _initialize_parameters(self, voxel_size, max_range, is_max_iter): 390 | self.voxel_size = voxel_size 391 | self.is_max_iter = is_max_iter 392 | self.max_range = max_range 393 | 394 | def _read_camera_params(self): 395 | # Get Camera Parameters 396 | self.K = DataUtils.read_intrinsic(self.input_path + "/depth_intrinsic.txt") 397 | self.cameraR, self.cameraT = DataUtils.read_extrinsic(self.input_path + "/camera_extrinsic.txt") 398 | 399 | 400 | 401 | -------------------------------------------------------------------------------- /iplanner/ip_algo.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import PIL 11 | import math 12 | import torch 13 | import torchvision.transforms as transforms 14 | 15 | from iplanner import traj_opt 16 | 17 | class IPlannerAlgo: 18 | def __init__(self, args): 19 | super(IPlannerAlgo, self).__init__() 20 | self.config(args) 21 | 22 | self.depth_transform = transforms.Compose([ 23 | transforms.Resize(tuple(self.crop_size)), 24 | transforms.ToTensor()]) 25 | 26 | net, _ = torch.load(self.model_save, map_location=torch.device("cpu")) 27 | self.net = net.cuda() if torch.cuda.is_available() else net 28 | 29 | self.traj_generate = traj_opt.TrajOpt() 30 | return None 31 | 32 | def config(self, args): 33 | self.model_save = args.model_save 34 | self.crop_size = args.crop_size 35 | self.sensor_offset_x = args.sensor_offset_x 36 | self.sensor_offset_y = args.sensor_offset_y 37 | self.is_traj_shift = False 38 | if math.hypot(self.sensor_offset_x, self.sensor_offset_y) > 1e-1: 39 | self.is_traj_shift = True 40 | return None 41 | 42 | 43 | def plan(self, image, goal_robot_frame): 44 | img = PIL.Image.fromarray(image) 45 | img = self.depth_transform(img).expand(1, 3, -1, -1) 46 | if torch.cuda.is_available(): 47 | img = img.cuda() 48 | goal_robot_frame = goal_robot_frame.cuda() 49 | with torch.no_grad(): 50 | keypoints, fear = self.net(img, goal_robot_frame) 51 | if self.is_traj_shift: 52 | batch_size, _, dims = keypoints.shape 53 | keypoints = torch.cat((torch.zeros(batch_size, 1, dims, device=keypoints.device, requires_grad=False), keypoints), axis=1) 54 | keypoints[..., 0] += self.sensor_offset_x 55 | keypoints[..., 1] += self.sensor_offset_y 56 | traj = self.traj_generate.TrajGeneratorFromPFreeRot(keypoints , step=0.1) 57 | 58 | return keypoints, traj, fear, img 59 | -------------------------------------------------------------------------------- /iplanner/models/placeholder.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/leggedrobotics/iPlanner/4a8d823ff9d09c3f626b727e7e00484b38f80d49/iplanner/models/placeholder.txt -------------------------------------------------------------------------------- /iplanner/percept_net.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import torch.nn as nn 11 | 12 | def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1): 13 | """3x3 convolution with padding""" 14 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 15 | padding=dilation, groups=groups, bias=False, dilation=dilation) 16 | 17 | 18 | def conv1x1(in_planes, out_planes, stride=1): 19 | """1x1 convolution""" 20 | return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) 21 | 22 | 23 | class BasicBlock(nn.Module): 24 | expansion = 1 25 | 26 | def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, 27 | base_width=64, dilation=1, norm_layer=None): 28 | super(BasicBlock, self).__init__() 29 | if norm_layer is None: 30 | norm_layer = nn.BatchNorm2d 31 | if groups != 1 or base_width != 64: 32 | raise ValueError('BasicBlock only supports groups=1 and base_width=64') 33 | if dilation > 1: 34 | raise NotImplementedError("Dilation > 1 not supported in BasicBlock") 35 | # Both self.conv1 and self.downsample layers downsample the input when stride != 1 36 | self.conv1 = conv3x3(inplanes, planes, stride) 37 | self.relu = nn.ReLU(inplace=True) 38 | self.conv2 = conv3x3(planes, planes) 39 | self.downsample = downsample 40 | self.stride = stride 41 | 42 | def forward(self, x): 43 | identity = x 44 | 45 | out = self.conv1(x) 46 | out = self.relu(out) 47 | 48 | out = self.conv2(out) 49 | 50 | if self.downsample is not None: 51 | identity = self.downsample(x) 52 | 53 | out += identity 54 | out = self.relu(out) 55 | 56 | return out 57 | 58 | 59 | class PerceptNet(nn.Module): 60 | 61 | def __init__(self, layers, block=BasicBlock, groups=1, 62 | width_per_group=64, replace_stride_with_dilation=None, norm_layer=None): 63 | 64 | super(PerceptNet, self).__init__() 65 | if norm_layer is None: 66 | norm_layer = nn.BatchNorm2d 67 | self._norm_layer = norm_layer 68 | 69 | self.inplanes = 64 70 | self.dilation = 1 71 | if replace_stride_with_dilation is None: 72 | # each element in the tuple indicates if we should replace 73 | # the 2x2 stride with a dilated convolution instead 74 | replace_stride_with_dilation = [False, False, False] 75 | if len(replace_stride_with_dilation) != 3: 76 | raise ValueError("replace_stride_with_dilation should be None " 77 | "or a 3-element tuple, got {}".format(replace_stride_with_dilation)) 78 | self.groups = groups 79 | self.base_width = width_per_group 80 | self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, 81 | bias=False) 82 | self.relu = nn.ReLU(inplace=True) 83 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 84 | self.layer1 = self._make_layer(block, 64, layers[0]) 85 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2, 86 | dilate=replace_stride_with_dilation[0]) 87 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2, 88 | dilate=replace_stride_with_dilation[1]) 89 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2, 90 | dilate=replace_stride_with_dilation[2]) 91 | 92 | for m in self.modules(): 93 | if isinstance(m, nn.Conv2d): 94 | nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') 95 | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): 96 | nn.init.constant_(m.weight, 1) 97 | nn.init.constant_(m.bias, 0) 98 | 99 | 100 | def _make_layer(self, block, planes, blocks, stride=1, dilate=False): 101 | norm_layer = self._norm_layer 102 | downsample = None 103 | previous_dilation = self.dilation 104 | if dilate: 105 | self.dilation *= stride 106 | stride = 1 107 | if stride != 1 or self.inplanes != planes * block.expansion: 108 | downsample = nn.Sequential( 109 | conv1x1(self.inplanes, planes * block.expansion, stride), 110 | norm_layer(planes * block.expansion), 111 | ) 112 | 113 | layers = [] 114 | layers.append(block(self.inplanes, planes, stride, downsample, self.groups, 115 | self.base_width, previous_dilation, norm_layer)) 116 | self.inplanes = planes * block.expansion 117 | for _ in range(1, blocks): 118 | layers.append(block(self.inplanes, planes, groups=self.groups, 119 | base_width=self.base_width, dilation=self.dilation, 120 | norm_layer=norm_layer)) 121 | 122 | return nn.Sequential(*layers) 123 | 124 | def _forward_impl(self, x): 125 | x = self.conv1(x) 126 | x = self.relu(x) 127 | x = self.maxpool(x) 128 | 129 | x = self.layer1(x) 130 | x = self.layer2(x) 131 | x = self.layer3(x) 132 | x = self.layer4(x) 133 | 134 | return x 135 | 136 | def forward(self, x): 137 | return self._forward_impl(x) -------------------------------------------------------------------------------- /iplanner/planner_net.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import torch 11 | from percept_net import PerceptNet 12 | import torch.nn as nn 13 | 14 | 15 | class PlannerNet(nn.Module): 16 | def __init__(self, encoder_channel=64, k=5): 17 | super().__init__() 18 | self.encoder = PerceptNet(layers=[2, 2, 2, 2]) 19 | self.decoder = Decoder(512, encoder_channel, k) 20 | 21 | def forward(self, x, goal): 22 | x = self.encoder(x) 23 | x, c = self.decoder(x, goal) 24 | return x, c 25 | 26 | 27 | class Decoder(nn.Module): 28 | def __init__(self, in_channels, goal_channels, k=5): 29 | super().__init__() 30 | self.k = k 31 | self.relu = nn.ReLU(inplace=True) 32 | self.fg = nn.Linear(3, goal_channels) 33 | self.sigmoid = nn.Sigmoid() 34 | 35 | self.conv1 = nn.Conv2d((in_channels + goal_channels), 512, kernel_size=5, stride=1, padding=1) 36 | self.conv2 = nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=0); 37 | 38 | self.fc1 = nn.Linear(256 * 128, 1024) 39 | self.fc2 = nn.Linear(1024, 512) 40 | self.fc3 = nn.Linear(512, k*3) 41 | 42 | self.frc1 = nn.Linear(1024, 128) 43 | self.frc2 = nn.Linear(128, 1) 44 | 45 | def forward(self, x, goal): 46 | # compute goal encoding 47 | goal = self.fg(goal[:, 0:3]) 48 | goal = goal[:, :, None, None].expand(-1, -1, x.shape[2], x.shape[3]) 49 | # cat x with goal in channel dim 50 | x = torch.cat((x, goal), dim=1) 51 | # compute x 52 | x = self.relu(self.conv1(x)) # size = (N, 512, x.H/32, x.W/32) 53 | x = self.relu(self.conv2(x)) # size = (N, 512, x.H/60, x.W/60) 54 | x = torch.flatten(x, 1) 55 | 56 | f = self.relu(self.fc1(x)) 57 | 58 | x = self.relu(self.fc2(f)) 59 | x = self.fc3(x) 60 | x = x.reshape(-1, self.k, 3) 61 | 62 | c = self.relu(self.frc1(f)) 63 | c = self.sigmoid(self.frc2(c)) 64 | 65 | return x, c 66 | -------------------------------------------------------------------------------- /iplanner/rosutil.py: -------------------------------------------------------------------------------- 1 | # ============================================================================== 2 | # Copyright <2019> 3 | # Refer to: https://github.com/wang-chen/interestingness_ros/blob/master/script/rosutil.py 4 | # ============================================================================== 5 | 6 | import os 7 | import rospy 8 | import torch 9 | import numpy as np 10 | 11 | class ROSArgparse(): 12 | def __init__(self, relative=None): 13 | self.relative = relative 14 | 15 | def add_argument(self, name, default, type=None, help=None): 16 | name = os.path.join(self.relative, name) 17 | if rospy.has_param(name): 18 | rospy.loginfo('Get param %s', name) 19 | else: 20 | rospy.logwarn('Couldn\'t find param: %s, Using default: %s', name, default) 21 | value = rospy.get_param(name, default) 22 | variable = name[name.rfind('/')+1:].replace('-','_') 23 | if isinstance(value, str): 24 | exec('self.%s=\'%s\''%(variable, value)) 25 | else: 26 | exec('self.%s=%s'%(variable, value)) 27 | 28 | def parse_args(self): 29 | return self 30 | 31 | 32 | def msg_to_torch(data, shape=np.array([-1])): 33 | return torch.from_numpy(data).view(shape.tolist()) 34 | 35 | 36 | def torch_to_msg(tensor): 37 | return [tensor.view(-1).cpu().numpy(), tensor.shape] 38 | -------------------------------------------------------------------------------- /iplanner/torchutil.py: -------------------------------------------------------------------------------- 1 | # ============================================================================== 2 | # Copyright <2019> 3 | # Refer to: https://github.com/wang-chen/interestingness/blob/tro/torchutil.py 4 | # ============================================================================== 5 | 6 | import cv2 7 | import time 8 | import math 9 | import torch 10 | import random 11 | import torch.fft 12 | import collections 13 | import torchvision 14 | from torch import nn 15 | from itertools import repeat 16 | import torch.nn.functional as F 17 | import torchvision.transforms.functional as TF 18 | 19 | 20 | def _ntuple(n): 21 | def parse(x): 22 | if isinstance(x, collections.abc.Iterable): 23 | return x 24 | return tuple(repeat(x, n)) 25 | return parse 26 | 27 | _single = _ntuple(1) 28 | _pair = _ntuple(2) 29 | _triple = _ntuple(3) 30 | _quadruple = _ntuple(4) 31 | 32 | 33 | class Timer: 34 | def __init__(self): 35 | self.start_time = time.time() 36 | 37 | def tic(self): 38 | self.start() 39 | 40 | def show(self, prefix="", output=True): 41 | duration = time.time()-self.start_time 42 | if output: 43 | print(prefix+"%fs" % duration) 44 | return duration 45 | 46 | def toc(self, prefix=""): 47 | self.end() 48 | print(prefix+"%fs = %fHz" % (self.duration, 1/self.duration)) 49 | return self.duration 50 | 51 | def start(self): 52 | torch.cuda.synchronize() 53 | self.start_time = time.time() 54 | 55 | def end(self): 56 | torch.cuda.synchronize() 57 | self.duration = time.time()-self.start_time 58 | self.start() 59 | return self.duration 60 | 61 | 62 | class MovAvg(nn.Module): 63 | def __init__(self, window_size=3): 64 | super(MovAvg, self).__init__() 65 | assert(window_size>=1) 66 | self.window_size = window_size 67 | weight = torch.arange(1, window_size+1).type('torch.FloatTensor') 68 | self.register_buffer('weight', torch.zeros(1,1,window_size)) 69 | self.weight.data = (weight/weight.sum()).view(1,1,-1) 70 | self.nums = [] 71 | 72 | def append(self, point): 73 | if len(self.nums) == 0: 74 | self.nums = [point]*self.window_size 75 | else: 76 | self.nums.append(point) 77 | self.nums.pop(0) 78 | return F.conv1d(torch.tensor(self.nums, dtype=torch.float).view(1,1,-1), self.weight).view(-1) 79 | 80 | 81 | class ConvLoss(nn.Module): 82 | def __init__(self, input_size, kernel_size, stride, in_channels=3, color=1): 83 | super(ConvLoss, self).__init__() 84 | self.color, input_size, kernel_size, stride = color, _pair(input_size), _pair(kernel_size), _pair(stride) 85 | self.conv = nn.Conv2d(in_channels, 1, kernel_size=kernel_size, stride=stride, bias=False) 86 | self.conv.weight.data = torch.ones(self.conv.weight.size()).cuda()/self.conv.weight.numel() 87 | self.width = (input_size[0] - kernel_size[0]) // stride[0] + 1 88 | self.hight = (input_size[0] - kernel_size[1]) // stride[1] + 1 89 | self.pool = nn.MaxPool2d((self.width, self.hight)) 90 | 91 | def forward(self, x, y): 92 | loss = self.conv((x-y).abs()) 93 | value, index = loss.view(-1).max(dim=0) 94 | w = (index//self.width)*self.conv.stride[0] 95 | h = (index%self.width)*self.conv.stride[1] 96 | x[:,:,w:w+self.conv.kernel_size[0],h] -= self.color 97 | x[:,:,w:w+self.conv.kernel_size[0],h+self.conv.kernel_size[1]-1] -= self.color 98 | x[:,:,w,h:h+self.conv.kernel_size[1]] -= self.color 99 | x[:,:,w+self.conv.kernel_size[0]-1,h:h+self.conv.kernel_size[1]] -= self.color 100 | return value 101 | 102 | 103 | class CosineSimilarity(nn.Module): 104 | ''' 105 | Averaged Cosine Similarity for 3-D tensor(C, H, W) over channel dimension 106 | Input Shape: 107 | x: tensor(N, C, H, W) 108 | y: tensor(B, C, H, W) 109 | Output Shape: 110 | o: tensor(N, B) 111 | ''' 112 | def __init__(self, eps=1e-7): 113 | super(CosineSimilarity, self).__init__() 114 | self.eps = eps 115 | 116 | def forward(self, x, y): 117 | N, C, H, W = x.size() 118 | B, c, h, w = y.size() 119 | assert(C==c and H==h and W==w) 120 | x, y = x.view(N,1,C,H*W), y.view(B,C,H*W) 121 | xx, yy = x.norm(dim=-1), y.norm(dim=-1) 122 | xx[xx self.patience: 277 | self.cooldown_counter = self.cooldown 278 | self.num_bad_epochs = 0 279 | return self._reduce_lr(epoch) 280 | 281 | def _reduce_lr(self, epoch): 282 | for i, param_group in enumerate(self.optimizer.param_groups): 283 | old_lr = float(param_group['lr']) 284 | new_lr = max(old_lr * self.factor, self.min_lrs[i]) 285 | if old_lr - new_lr > self.eps: 286 | param_group['lr'] = new_lr 287 | if self.verbose: 288 | print('Epoch {:5d}: reducing learning rate' 289 | ' of group {} to {:.4e}.'.format(epoch, i, new_lr)) 290 | return False 291 | else: 292 | return True 293 | 294 | 295 | class CorrelationSimilarity(nn.Module): 296 | ''' 297 | Correlation Similarity for multi-channel 2-D tensor(C, H, W) via FFT 298 | args: input_size: tuple(H, W) --> size of last two dimensions 299 | Input Shape: 300 | x: tensor(B, C, H, W) 301 | y: tensor(N, C, H, W) 302 | Output Shape: 303 | o: tensor(B, N) --> maximum similarity for (x_i, y_j) {i\in [0,B), j\in [0,N)} 304 | i: tensor(B, N, 2) --> 2-D translation between x_i and y_j 305 | ''' 306 | def __init__(self, input_size): 307 | super(CorrelationSimilarity, self).__init__() 308 | self.input_size = input_size = _pair(input_size) 309 | assert(input_size[-1]!=1) # FFT2 is wrong if last dimension is 1 310 | self.N = math.sqrt(input_size[0]*input_size[1]) 311 | self.fft_args = {'s': input_size, 'dim':[-2,-1], 'norm': 'ortho'} 312 | self.max = nn.MaxPool2d(kernel_size=input_size) 313 | 314 | def forward(self, x, y): 315 | X = torch.fft.rfftn(x, **self.fft_args).unsqueeze(1) 316 | Y = torch.fft.rfftn(y, **self.fft_args) 317 | g = torch.fft.irfftn((X.conj()*Y).sum(2), **self.fft_args)*self.N 318 | xx = x.view(x.size(0),-1).norm(dim=-1).view(x.size(0), 1, 1) 319 | yy = y.view(y.size(0),-1).norm(dim=-1).view(1, y.size(0), 1) 320 | g = g.view(x.size(0), y.size(0),-1)/xx/yy 321 | values, indices = torch.max(g, dim=-1) 322 | indices = torch.stack((indices // self.input_size[1], indices % self.input_size[1]), dim=-1) 323 | values[values>+1] = +1 # prevent from overflow of 1 324 | values[values<-1] = -1 # prevent from overflow of -1 325 | assert((values>+1).sum()==0 and (values<-1).sum()==0) 326 | return values, indices 327 | 328 | 329 | class Correlation(nn.Module): 330 | ''' 331 | Correlation Similarity for multi-channel 2-D patch via FFT 332 | args: input_size: tuple(H, W) --> size of last two dimensions 333 | Input Shape: 334 | x: tensor(B, C, H, W) 335 | y: tensor(B, C, H, W) 336 | Output Shape: 337 | o: tensor(B) 338 | if accept_translation is False, output is the same with cosine similarity 339 | ''' 340 | def __init__(self, input_size, accept_translation=True): 341 | super(Correlation, self).__init__() 342 | self.accept_translation = accept_translation 343 | input_size = _pair(input_size) 344 | assert(input_size[-1]!=1) # FFT2 is wrong if last dimension is 1 345 | self.N = math.sqrt(input_size[0]*input_size[1]) 346 | self.fft_args = {'s': input_size, 'dim':[-2,-1], 'norm': 'ortho'} 347 | self.max = nn.MaxPool2d(kernel_size=input_size) 348 | 349 | def forward(self, x, y): 350 | X = torch.fft.rfftn(x, **self.fft_args) 351 | Y = torch.fft.rfftn(y, **self.fft_args) 352 | g = torch.fft.irfftn((X.conj()*Y).sum(2), **self.fft_args)*self.N 353 | xx = x.view(x.size(0),-1).norm(dim=-1) 354 | yy = y.view(y.size(0),-1).norm(dim=-1) 355 | if self.accept_translation is True: 356 | return self.max(g).view(-1)/xx/yy 357 | else: 358 | return g[:,0,0].view(-1)/xx/yy 359 | 360 | 361 | class CorrelationLoss(Correlation): 362 | ''' 363 | Correlation Similarity for multi-channel 2-D patch via FFT 364 | args: input_size: tuple(H, W) --> size of last two dimensions 365 | Input Shape: 366 | x: tensor(B, C, H, W) 367 | y: tensor(B, C, H, W) 368 | Output Shape: 369 | o: tensor(1) if 'reduce' is True 370 | o: tensor(B) if 'reduce' is not True 371 | ''' 372 | def __init__(self, input_size, reduce = True, accept_translation=True): 373 | super(CorrelationLoss, self).__init__(input_size, accept_translation) 374 | self.reduce = reduce 375 | 376 | def forward(self, x, y): 377 | loss = (1 - super(CorrelationLoss, self).forward(x, y))/2 378 | if self.reduce is True: 379 | return loss.mean() 380 | else: 381 | return loss 382 | 383 | 384 | def rolls2d(inputs, shifts, dims=[-2,-1]): 385 | ''' 386 | shifts: list of tuple/ints for 2-D/1-D roll 387 | dims: along which dimensions to shift 388 | inputs: tensor(N, C, H, W); shifts has to be int tensor 389 | if shifts: tensor(B, N, 2) 390 | output: tensor(B, N, C, H, W) 391 | if shifts: tensor(N, 2) 392 | output: tensor(N, C, H, W) 393 | ''' 394 | shift_size = shifts.size() 395 | N, C, H, W = inputs.size() 396 | assert(shift_size[-1]==2 and N==shift_size[1]) 397 | if len(shift_size) == 2: 398 | return torch.stack([inputs[i].roll(shifts[i].tolist(), dims) for i in range(N)], dim=0) 399 | elif len(shift_size) == 3: 400 | B = shift_size[0] 401 | o = torch.stack([inputs[i].roll(shifts[j,i].tolist(), dims) for j in range(B) for i in range(N)], dim=0) 402 | return o.view(B, N, C, H, W) 403 | 404 | 405 | def count_parameters(model): 406 | return sum(p.numel() for p in model.parameters() if p.requires_grad) 407 | 408 | 409 | def show_batch(batch, name='video', waitkey=1): 410 | min_v = torch.min(batch) 411 | range_v = torch.max(batch) - min_v 412 | if range_v > 0: 413 | batch = (batch - min_v) / range_v 414 | else: 415 | batch = torch.zeros(batch.size()) 416 | grid = torchvision.utils.make_grid(batch, padding=0).cpu() 417 | img = grid.numpy()[::-1].transpose((1, 2, 0)) 418 | cv2.imshow(name, img) 419 | cv2.waitKey(waitkey) 420 | return img 421 | 422 | 423 | def show_batch_origin(batch, name='video', waitkey=1): 424 | grid = torchvision.utils.make_grid(batch).cpu() 425 | img = grid.numpy()[::-1].transpose((1, 2, 0)) 426 | cv2.imshow(name, img) 427 | cv2.waitKey(waitkey) 428 | return img 429 | -------------------------------------------------------------------------------- /iplanner/training_run.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import os 11 | import tqdm 12 | import time 13 | import torch 14 | import json 15 | import wandb 16 | import random 17 | import argparse 18 | import torch.optim as optim 19 | from datetime import datetime 20 | import torch.nn.functional as F 21 | import torchvision.transforms as transforms 22 | 23 | from planner_net import PlannerNet 24 | from dataloader import PlannerData 25 | from torch.utils.data import DataLoader 26 | from torchutil import EarlyStopScheduler 27 | from traj_cost import TrajCost 28 | from traj_viz import TrajViz 29 | 30 | torch.set_default_dtype(torch.float32) 31 | 32 | class PlannerNetTrainer: 33 | def __init__(self): 34 | self.root_folder = os.getenv('EXPERIMENT_DIRECTORY', os.getcwd()) 35 | self.load_config() 36 | self.parse_args() 37 | self.prepare_model() 38 | self.prepare_data() 39 | if self.args.training == True: 40 | self.init_wandb() 41 | else: 42 | print("Testing Mode") 43 | 44 | def init_wandb(self): 45 | # Convert to string in the format you prefer 46 | date_time_str = datetime.now().strftime("%d-%m-%Y-%H-%M-%S") 47 | # using Wandb Core 48 | wandb.require("core") 49 | # Initialize wandb 50 | self.wandb_run = wandb.init( 51 | # set the wandb project where this run will be logged 52 | project="imperative-path-planning", 53 | # Set the run name to current date and time 54 | name=date_time_str + "adamW", 55 | config={ 56 | "learning_rate": self.args.lr, 57 | "architecture": "PlannerNet", # Replace with your actual architecture 58 | "dataset": self.args.data_root, # Assuming this holds the dataset name 59 | "epochs": self.args.epochs, 60 | "goal_step": self.args.goal_step, 61 | "max_episode": self.args.max_episode, 62 | "fear_ahead_dist": self.args.fear_ahead_dist, 63 | } 64 | ) 65 | 66 | def load_config(self): 67 | with open(os.path.join(os.path.dirname(self.root_folder), 'config', 'training_config.json')) as json_file: 68 | self.config = json.load(json_file) 69 | 70 | def prepare_model(self): 71 | self.net = PlannerNet(self.args.in_channel, self.args.knodes) 72 | if self.args.resume == True or not self.args.training: 73 | self.net, self.best_loss = torch.load(self.args.model_save, map_location=torch.device("cpu")) 74 | print("Resume training from best loss: {}".format(self.best_loss)) 75 | else: 76 | self.best_loss = float('Inf') 77 | 78 | if torch.cuda.is_available(): 79 | print("Available GPU list: {}".format(list(range(torch.cuda.device_count())))) 80 | print("Runnin on GPU: {}".format(self.args.gpu_id)) 81 | self.net = self.net.cuda(self.args.gpu_id) 82 | 83 | self.optimizer = optim.AdamW(self.net.parameters(), lr=self.args.lr, weight_decay=self.args.w_decay) 84 | self.scheduler = EarlyStopScheduler(self.optimizer, factor=self.args.factor, verbose=True, min_lr=self.args.min_lr, patience=self.args.patience) 85 | 86 | def prepare_data(self): 87 | ids_path = os.path.join(self.args.data_root, self.args.env_id) 88 | with open(ids_path) as f: 89 | self.env_list = [line.rstrip() for line in f.readlines()] 90 | 91 | depth_transform = transforms.Compose([ 92 | transforms.Resize((self.args.crop_size)), 93 | transforms.ToTensor()]) 94 | 95 | total_img_data = 0 96 | track_id = 0 97 | test_env_id = min(self.args.test_env_id, len(self.env_list)-1) 98 | 99 | self.train_loader_list = [] 100 | self.val_loader_list = [] 101 | self.traj_cost_list = [] 102 | self.traj_viz_list = [] 103 | 104 | for env_name in tqdm.tqdm(self.env_list): 105 | if not self.args.training and track_id != test_env_id: 106 | track_id += 1 107 | continue 108 | is_anymal_frame = False 109 | sensorOffsetX = 0.0 110 | camera_tilt = 0.0 111 | if 'anymal' in env_name: 112 | is_anymal_frame = True 113 | sensorOffsetX = self.args.sensor_offsetX_ANYmal 114 | camera_tilt = self.args.camera_tilt 115 | elif 'tilt' in env_name: 116 | camera_tilt = self.args.camera_tilt 117 | data_path = os.path.join(*[self.args.data_root, self.args.env_type, env_name]) 118 | 119 | train_data = PlannerData(root=data_path, 120 | train=True, 121 | transform=depth_transform, 122 | sensorOffsetX=sensorOffsetX, 123 | is_robot=is_anymal_frame, 124 | goal_step=self.args.goal_step, 125 | max_episode=self.args.max_episode, 126 | max_depth=self.args.max_camera_depth) 127 | 128 | total_img_data += len(train_data) 129 | train_loader = DataLoader(train_data, batch_size=self.args.batch_size, shuffle=True, num_workers=2) 130 | self.train_loader_list.append(train_loader) 131 | 132 | val_data = PlannerData(root=data_path, 133 | train=False, 134 | transform=depth_transform, 135 | sensorOffsetX=sensorOffsetX, 136 | is_robot=is_anymal_frame, 137 | goal_step=self.args.goal_step, 138 | max_episode=self.args.max_episode, 139 | max_depth=self.args.max_camera_depth) 140 | 141 | val_loader = DataLoader(val_data, batch_size=self.args.batch_size, shuffle=True, num_workers=2) 142 | self.val_loader_list.append(val_loader) 143 | 144 | # Load Map and Trajectory Class 145 | map_name = "tsdf1" 146 | traj_cost = TrajCost(self.args.gpu_id) 147 | traj_cost.SetMap(data_path, map_name) 148 | 149 | self.traj_cost_list.append(traj_cost) 150 | self.traj_viz_list.append(TrajViz(data_path, map_name=map_name, cameraTilt=camera_tilt)) 151 | track_id += 1 152 | 153 | print("Data Loading Completed!") 154 | print("Number of image: %d | Number of goal-image pairs: %d"%(total_img_data, total_img_data * (int)(self.args.max_episode / self.args.goal_step))) 155 | 156 | return None 157 | 158 | def MapObsLoss(self, preds, fear, traj_cost, odom, goal, step=0.1): 159 | waypoints = traj_cost.opt.TrajGeneratorFromPFreeRot(preds, step=step) 160 | loss1, fear_labels = traj_cost.CostofTraj(waypoints, odom, goal, ahead_dist=self.args.fear_ahead_dist) 161 | loss2 = F.binary_cross_entropy(fear, fear_labels) 162 | return loss1+loss2, waypoints 163 | 164 | def train_epoch(self, epoch): 165 | loss_sum = 0.0 166 | env_num = len(self.train_loader_list) 167 | 168 | # Zip the lists and convert to a list of tuples 169 | combined = list(zip(self.train_loader_list, self.traj_cost_list)) 170 | # Shuffle the combined list 171 | random.shuffle(combined) 172 | 173 | # Iterate through shuffled pairs 174 | for env_id, (loader, traj_cost) in enumerate(combined): 175 | train_loss, batches = 0, len(loader) 176 | 177 | enumerater = tqdm.tqdm(enumerate(loader)) 178 | for batch_idx, inputs in enumerater: 179 | if torch.cuda.is_available(): 180 | image = inputs[0].cuda(self.args.gpu_id) 181 | odom = inputs[1].cuda(self.args.gpu_id) 182 | goal = inputs[2].cuda(self.args.gpu_id) 183 | self.optimizer.zero_grad() 184 | preds, fear = self.net(image, goal) 185 | 186 | loss, _ = self.MapObsLoss(preds, fear, traj_cost, odom, goal) 187 | 188 | loss.backward() 189 | self.optimizer.step() 190 | train_loss += loss.item() 191 | enumerater.set_description("Epoch: %d in Env: (%d/%d) - train loss: %.4f on %d/%d" % (epoch, env_id+1, env_num, train_loss/(batch_idx+1), batch_idx, batches)) 192 | 193 | loss_sum += train_loss/(batch_idx+1) 194 | wandb.log({"Running Loss": train_loss/(batch_idx+1)}) 195 | 196 | loss_sum /= env_num 197 | 198 | return loss_sum 199 | 200 | def train(self): 201 | # Convert to string in the format you prefer 202 | date_time_str = datetime.now().strftime("%d-%m-%Y-%H-%M-%S") 203 | 204 | self.args.log_save += (date_time_str + ".txt") 205 | open(self.args.log_save, 'w').close() 206 | 207 | for epoch in range(self.args.epochs): 208 | start_time = time.time() 209 | train_loss = self.train_epoch(epoch) 210 | val_loss = self.evaluate(is_visualize=False) 211 | duration = (time.time() - start_time) / 60 # minutes 212 | 213 | self.log_message("Epoch: %d | Training Loss: %f | Val Loss: %f | Duration: %f" % (epoch, train_loss, val_loss, duration)) 214 | # Log metrics to wandb 215 | wandb.log({"Avg Training Loss": train_loss, "Validation Loss": val_loss, "Duration (min)": duration}) 216 | 217 | if val_loss < self.best_loss: 218 | self.log_message("Save model of epoch %d" % epoch) 219 | torch.save((self.net, val_loss), self.args.model_save) 220 | self.best_loss = val_loss 221 | self.log_message("Current val loss: %.4f" % self.best_loss) 222 | self.log_message("Epoch: %d model saved | Current Min Val Loss: %f" % (epoch, val_loss)) 223 | 224 | self.log_message("------------------------------------------------------------------------") 225 | if self.scheduler.step(val_loss): 226 | self.log_message('Early Stopping!') 227 | break 228 | 229 | # Close wandb run at the end of training 230 | self.wandb_run.finish() 231 | 232 | def log_message(self, message): 233 | with open(self.args.log_save, 'a') as f: 234 | f.writelines(message) 235 | f.write('\n') 236 | print(message) 237 | 238 | def evaluate(self, is_visualize=False): 239 | self.net.eval() 240 | test_loss = 0 # Declare and initialize test_loss 241 | total_batches = 0 # Count total number of batches 242 | with torch.no_grad(): 243 | for _, (val_loader, traj_cost, traj_viz) in enumerate(zip(self.val_loader_list, self.traj_cost_list, self.traj_viz_list)): 244 | preds_viz = [] 245 | wp_viz = [] 246 | for batch_idx, inputs in enumerate(val_loader): 247 | total_batches += 1 # Increment total number of batches 248 | if torch.cuda.is_available(): 249 | image = inputs[0].cuda(self.args.gpu_id) 250 | odom = inputs[1].cuda(self.args.gpu_id) 251 | goal = inputs[2].cuda(self.args.gpu_id) 252 | 253 | preds, fear = self.net(image, goal) 254 | loss, waypoints = self.MapObsLoss(preds, fear, traj_cost, odom, goal) 255 | test_loss += loss.item() 256 | 257 | if is_visualize and len(preds_viz) < self.args.visual_number: 258 | if batch_idx == 0: 259 | image_viz = image 260 | odom_viz = odom 261 | goal_viz = goal 262 | fear_viz = fear 263 | else: 264 | image_viz = torch.cat((image_viz, image), dim=0) 265 | odom_viz = torch.cat((odom_viz, odom), dim=0) 266 | goal_viz = torch.cat((goal_viz, goal), dim=0) 267 | fear_viz = torch.cat((fear_viz, fear), dim=0) 268 | preds_viz.extend(preds.tolist()) 269 | wp_viz.extend(waypoints.tolist()) 270 | 271 | if is_visualize: 272 | max_n = min(len(wp_viz), self.args.visual_number) 273 | preds_viz = torch.tensor(preds_viz[:max_n]) 274 | wp_viz = torch.tensor(wp_viz[:max_n]) 275 | odom_viz = odom_viz[:max_n].cpu() 276 | goal_viz = goal_viz[:max_n].cpu() 277 | fear_viz = fear_viz[:max_n, :].cpu() 278 | image_viz = image_viz[:max_n].cpu() 279 | # visual trajectory and images 280 | traj_viz.VizTrajectory(preds_viz, wp_viz, odom_viz, goal_viz, fear_viz) 281 | traj_viz.VizImages(preds_viz, wp_viz, odom_viz, goal_viz, fear_viz, image_viz) 282 | 283 | return test_loss / total_batches # Compute mean test_loss 284 | 285 | def parse_args(self): 286 | parser = argparse.ArgumentParser(description='Training script for PlannerNet') 287 | 288 | # dataConfig 289 | parser.add_argument("--data-root", type=str, default=os.path.join(self.root_folder, self.config['dataConfig'].get('data-root')), help="dataset root folder") 290 | parser.add_argument('--env-id', type=str, default=self.config['dataConfig'].get('env-id'), help='environment id list') 291 | parser.add_argument('--env_type', type=str, default=self.config['dataConfig'].get('env_type'), help='the dataset type') 292 | parser.add_argument('--crop-size', nargs='+', type=int, default=self.config['dataConfig'].get('crop-size'), help='image crop size') 293 | parser.add_argument('--max-camera-depth', type=float, default=self.config['dataConfig'].get('max-camera-depth'), help='maximum depth detection of camera, unit: meter') 294 | 295 | # modelConfig 296 | parser.add_argument("--model-save", type=str, default=os.path.join(self.root_folder, self.config['modelConfig'].get('model-save')), help="model save point") 297 | parser.add_argument('--resume', type=str, default=self.config['modelConfig'].get('resume')) 298 | parser.add_argument('--in-channel', type=int, default=self.config['modelConfig'].get('in-channel'), help='goal input channel numbers') 299 | parser.add_argument("--knodes", type=int, default=self.config['modelConfig'].get('knodes'), help="number of max nodes predicted") 300 | parser.add_argument("--goal-step", type=int, default=self.config['modelConfig'].get('goal-step'), help="number of frames betwen goals") 301 | parser.add_argument("--max-episode", type=int, default=self.config['modelConfig'].get('max-episode-length'), help="maximum episode frame length") 302 | 303 | # trainingConfig 304 | parser.add_argument('--training', type=str, default=self.config['trainingConfig'].get('training')) 305 | parser.add_argument("--lr", type=float, default=self.config['trainingConfig'].get('lr'), help="learning rate") 306 | parser.add_argument("--factor", type=float, default=self.config['trainingConfig'].get('factor'), help="ReduceLROnPlateau factor") 307 | parser.add_argument("--min-lr", type=float, default=self.config['trainingConfig'].get('min-lr'), help="minimum lr for ReduceLROnPlateau") 308 | parser.add_argument("--patience", type=int, default=self.config['trainingConfig'].get('patience'), help="patience of epochs for ReduceLROnPlateau") 309 | parser.add_argument("--epochs", type=int, default=self.config['trainingConfig'].get('epochs'), help="number of training epochs") 310 | parser.add_argument("--batch-size", type=int, default=self.config['trainingConfig'].get('batch-size'), help="number of minibatch size") 311 | parser.add_argument("--w-decay", type=float, default=self.config['trainingConfig'].get('w-decay'), help="weight decay of the optimizer") 312 | parser.add_argument("--num-workers", type=int, default=self.config['trainingConfig'].get('num-workers'), help="number of workers for dataloader") 313 | parser.add_argument("--gpu-id", type=int, default=self.config['trainingConfig'].get('gpu-id'), help="GPU id") 314 | 315 | # logConfig 316 | parser.add_argument("--log-save", type=str, default=os.path.join(self.root_folder, self.config['logConfig'].get('log-save')), help="train log file") 317 | parser.add_argument('--test-env-id', type=int, default=self.config['logConfig'].get('test-env-id'), help='the test env id in the id list') 318 | parser.add_argument('--visual-number', type=int, default=self.config['logConfig'].get('visual-number'), help='number of visualized trajectories') 319 | 320 | # sensorConfig 321 | parser.add_argument('--camera-tilt', type=float, default=self.config['sensorConfig'].get('camera-tilt'), help='camera tilt angle for visualization only') 322 | parser.add_argument('--sensor-offsetX-ANYmal', type=float, default=self.config['sensorConfig'].get('sensor-offsetX-ANYmal'), help='anymal front camera sensor offset in X axis') 323 | parser.add_argument("--fear-ahead-dist", type=float, default=self.config['sensorConfig'].get('fear-ahead-dist'), help="fear lookahead distance") 324 | 325 | self.args = parser.parse_args() 326 | 327 | 328 | def main(): 329 | trainer = PlannerNetTrainer() 330 | if trainer.args.training == True: 331 | trainer.train() 332 | trainer.evaluate(is_visualize=True) 333 | 334 | if __name__ == "__main__": 335 | main() 336 | -------------------------------------------------------------------------------- /iplanner/traj_cost.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import torch 11 | import pypose as pp 12 | from tsdf_map import TSDF_Map 13 | from traj_opt import TrajOpt 14 | import torch.nn.functional as F 15 | 16 | torch.set_default_dtype(torch.float32) 17 | 18 | class TrajCost: 19 | def __init__(self, gpu_id=0): 20 | self.tsdf_map = TSDF_Map(gpu_id) 21 | self.opt = TrajOpt() 22 | self.is_map = False 23 | return None 24 | 25 | def TransformPoints(self, odom, points): 26 | batch_size, num_p, _ = points.shape 27 | world_ps = pp.identity_SE3(batch_size, num_p, device=points.device, requires_grad=points.requires_grad) 28 | world_ps.tensor()[:, :, 0:3] = points 29 | world_ps = pp.SE3(odom[:, None, :]) @ pp.SE3(world_ps) 30 | return world_ps 31 | 32 | def SetMap(self, root_path, map_name): 33 | self.tsdf_map.ReadTSDFMap(root_path, map_name) 34 | self.is_map = True 35 | return 36 | 37 | def CostofTraj(self, waypoints, odom, goal, ahead_dist, alpha=0.5, beta=1.0, gamma=2.0, delta=5.0, obstalce_thred=0.5): 38 | batch_size, num_p, _ = waypoints.shape 39 | if self.is_map: 40 | world_ps = self.TransformPoints(odom, waypoints) 41 | norm_inds, _ = self.tsdf_map.Pos2Ind(world_ps) 42 | # Obstacle Cost 43 | cost_grid = self.tsdf_map.cost_array.T.expand(batch_size, 1, -1, -1) 44 | oloss_M = F.grid_sample(cost_grid, norm_inds[:, None, :, :], mode='bicubic', padding_mode='border', align_corners=False).squeeze(1).squeeze(1) 45 | oloss_M = oloss_M.to(torch.float32) 46 | oloss = torch.mean(torch.sum(oloss_M, axis=1)) 47 | 48 | # Terrian Height loss 49 | height_grid = self.tsdf_map.ground_array.T.expand(batch_size, 1, -1, -1) 50 | hloss_M = F.grid_sample(height_grid, norm_inds[:, None, :, :], mode='bicubic', padding_mode='border', align_corners=False).squeeze(1).squeeze(1) 51 | hloss_M = torch.abs(waypoints[:, :, 2] - hloss_M) 52 | hloss = torch.mean(torch.sum(hloss_M, axis=1)) 53 | 54 | # Goal Cost 55 | gloss = torch.norm(goal[:, :3] - waypoints[:, -1, :], dim=1) 56 | gloss = torch.mean(torch.log(gloss + 1.0)) 57 | # gloss = torch.mean(gloss) 58 | 59 | # Motion Loss 60 | desired_wp = self.opt.TrajGeneratorFromPFreeRot(goal[:, None, 0:3], step=1.0/(num_p-1)) 61 | desired_ds = torch.norm(desired_wp[:, 1:num_p, :] - desired_wp[:, 0:num_p-1, :], dim=2) 62 | wp_ds = torch.norm(waypoints[:, 1:num_p, :] - waypoints[:, 0:num_p-1, :], dim=2) 63 | mloss = torch.abs(desired_ds - wp_ds) 64 | mloss = torch.sum(mloss, axis=1) 65 | mloss = torch.mean(mloss) 66 | 67 | # Fear labels 68 | goal_dists = torch.cumsum(wp_ds, dim=1, dtype=wp_ds.dtype) 69 | floss_M = torch.clone(oloss_M)[:, 1:] 70 | floss_M[goal_dists > ahead_dist] = 0.0 71 | fear_labels = torch.max(floss_M, 1, keepdim=True)[0] 72 | fear_labels = (fear_labels > obstalce_thred).to(torch.float32) 73 | 74 | return alpha*oloss + beta*hloss + gamma*mloss + delta*gloss, fear_labels -------------------------------------------------------------------------------- /iplanner/traj_opt.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import torch 11 | 12 | torch.set_default_dtype(torch.float32) 13 | 14 | class CubicSplineTorch: 15 | # Reference: https://stackoverflow.com/questions/61616810/how-to-do-cubic-spline-interpolation-and-integration-in-pytorch 16 | def __init__(self): 17 | return None 18 | 19 | def h_poly(self, t): 20 | alpha = torch.arange(4, device=t.device, dtype=t.dtype) 21 | tt = t[:, None, :]**alpha[None, :, None] 22 | A = torch.tensor([ 23 | [1, 0, -3, 2], 24 | [0, 1, -2, 1], 25 | [0, 0, 3, -2], 26 | [0, 0, -1, 1] 27 | ], dtype=t.dtype, device=t.device) 28 | return A @ tt 29 | 30 | def interp(self, x, y, xs): 31 | m = (y[:, 1:, :] - y[:, :-1, :]) / torch.unsqueeze(x[:, 1:] - x[:, :-1], 2) 32 | m = torch.cat([m[:, None, 0], (m[:, 1:] + m[:, :-1]) / 2, m[:, None, -1]], 1) 33 | idxs = torch.searchsorted(x[0, 1:], xs[0, :]) 34 | dx = x[:, idxs + 1] - x[:, idxs] 35 | hh = self.h_poly((xs - x[:, idxs]) / dx) 36 | hh = torch.transpose(hh, 1, 2) 37 | out = hh[:, :, 0:1] * y[:, idxs, :] 38 | out = out + hh[:, :, 1:2] * m[:, idxs] * dx[:,:,None] 39 | out = out + hh[:, :, 2:3] * y[:, idxs + 1, :] 40 | out = out + hh[:, :, 3:4] * m[:, idxs + 1] * dx[:,:,None] 41 | return out 42 | 43 | class TrajOpt: 44 | def __init__(self): 45 | self.cs_interp = CubicSplineTorch() 46 | return None 47 | 48 | def TrajGeneratorFromPFreeRot(self, preds, step): 49 | # Points is in se3 50 | batch_size, num_p, dims = preds.shape 51 | points_preds = torch.cat((torch.zeros(batch_size, 1, dims, device=preds.device, requires_grad=preds.requires_grad), preds), axis=1) 52 | num_p = num_p + 1 53 | xs = torch.arange(0, num_p-1+step, step, device=preds.device) 54 | xs = xs.repeat(batch_size, 1) 55 | x = torch.arange(num_p, device=preds.device, dtype=preds.dtype) 56 | x = x.repeat(batch_size, 1) 57 | waypoints = self.cs_interp.interp(x, points_preds, xs) 58 | return waypoints # R3 -------------------------------------------------------------------------------- /iplanner/traj_viz.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import os 11 | import cv2 12 | import copy 13 | import torch 14 | import numpy as np 15 | import pypose as pp 16 | import open3d as o3d 17 | from tsdf_map import TSDF_Map 18 | from scipy.spatial.transform import Rotation as R 19 | import open3d.visualization.rendering as rendering 20 | from typing import List, Optional 21 | 22 | IMAGE_WIDTH = 640 23 | IMAGE_HEIGHT = 360 24 | MESH_SIZE = 0.5 25 | 26 | class TrajViz: 27 | def __init__(self, root_path: str, map_name: str = "tsdf1", cameraTilt: float = 0.0, robot_name: str = "robot"): 28 | """ 29 | Initialize TrajViz class. 30 | """ 31 | self.is_map = False 32 | if map_name: 33 | self.tsdf_map = TSDF_Map() 34 | self.tsdf_map.ReadTSDFMap(root_path, map_name) 35 | self.is_map = True 36 | intrinsic_path = os.path.join(root_path, "depth_intrinsic.txt") 37 | else: 38 | intrinsic_path = os.path.join(root_path, robot_name + "_intrinsic.txt") 39 | self.SetCamera(intrinsic_path) 40 | self.camera_tilt = cameraTilt 41 | 42 | def TransformPoints(self, odom: torch.Tensor, points: torch.Tensor) -> pp.SE3: 43 | """ 44 | Transform points in the trajectory. 45 | """ 46 | batch_size, num_p, _ = points.shape 47 | world_ps = pp.identity_SE3(batch_size, num_p, device=points.device) 48 | world_ps.tensor()[:, :, 0:3] = points 49 | world_ps = pp.SE3(odom[:, None, :]) @ pp.SE3(world_ps) 50 | return world_ps 51 | 52 | def SetCamera(self, intrinsic_path: str, img_width: int = IMAGE_WIDTH, img_height: int = IMAGE_HEIGHT): 53 | """ 54 | Set camera parameters. 55 | """ 56 | with open(intrinsic_path) as f: 57 | lines = f.readlines() 58 | elems = np.fromstring(lines[0][1:-2], dtype=float, sep=', ') 59 | K = np.array(elems).reshape(-1, 4) 60 | self.camera = o3d.camera.PinholeCameraIntrinsic(img_width, img_height, K[0,0], K[1,1], K[0,2], K[1,2]) 61 | 62 | def VizTrajectory(self, preds: torch.Tensor, waypoints: torch.Tensor, odom: torch.Tensor, goal: torch.Tensor, fear: torch.Tensor, cost_map: bool = True, visual_height: float = 0.5, mesh_size: float = MESH_SIZE) -> None: 63 | """ 64 | Visualize the trajectory. 65 | 66 | Parameters: 67 | preds (tensor): Predicted trajectory points 68 | waypoints (tensor): Original waypoints 69 | odom (tensor): Odometry data 70 | goal (tensor): Goal position 71 | fear (tensor): Fear value per trajectory 72 | cost_map (bool): If True, visualize cost map; otherwise, visualize point cloud 73 | visual_height (float): Height offset for visualization 74 | mesh_size (float): Size of mesh objects in visualization 75 | """ 76 | if not self.is_map: 77 | print("Cost map is missing.") 78 | return 79 | 80 | batch_size, _, _ = waypoints.shape 81 | 82 | # Transform to map frame 83 | preds_ws = self.TransformPoints(odom, preds) 84 | wp_ws = self.TransformPoints(odom, waypoints) 85 | goal_ws = pp.SE3(odom) @ pp.SE3(goal) 86 | 87 | # Convert to positions 88 | preds_ws = preds_ws.tensor()[:, :, 0:3].cpu().detach().numpy() 89 | wp_ws = wp_ws.tensor()[:, :, 0:3].cpu().detach().numpy() 90 | goal_ws = goal_ws.tensor()[:, 0:3].cpu().detach().numpy() 91 | 92 | visual_list = [] 93 | if cost_map: 94 | visual_list.append(self.tsdf_map.pcd_tsdf) 95 | else: 96 | visual_list.append(self.tsdf_map.pcd_viz) 97 | visual_height /= 5.0 98 | 99 | # Visualize trajectory 100 | traj_pcd = o3d.geometry.PointCloud() 101 | wp_ws = wp_ws.reshape(-1, 3) 102 | wp_ws[:, 2] += visual_height 103 | traj_pcd.points = o3d.utility.Vector3dVector(wp_ws[:, 0:3]) 104 | traj_pcd.paint_uniform_color([0.99, 0.1, 0.1]) 105 | visual_list.append(traj_pcd) 106 | 107 | # Start and goal markers 108 | mesh_sphere = o3d.geometry.TriangleMesh.create_sphere(mesh_size/1.5) # start points 109 | small_sphere = o3d.geometry.TriangleMesh.create_sphere(mesh_size/3.0) # successful trajectory points 110 | small_sphere_fear = o3d.geometry.TriangleMesh.create_sphere(mesh_size/3.0) # unsuccessful trajectory points 111 | mesh_box = o3d.geometry.TriangleMesh.create_box(mesh_size, mesh_size, mesh_size) # end points 112 | 113 | # Set mesh colors 114 | mesh_box.paint_uniform_color([1.0, 0.64, 0.0]) 115 | small_sphere.paint_uniform_color([0.4, 1.0, 0.1]) 116 | small_sphere_fear.paint_uniform_color([1.0, 0.4, 0.1]) 117 | 118 | lines = [] 119 | points = [] 120 | for i in range(batch_size): 121 | lines.append([2*i, 2*i+1]) 122 | gp = goal_ws[i, :] 123 | op = odom.numpy()[i, :] 124 | op[2] += visual_height 125 | gp[2] += visual_height 126 | points.append(gp[:3].tolist()) 127 | points.append(op[:3].tolist()) 128 | 129 | # Add visualization 130 | visual_list.append(copy.deepcopy(mesh_sphere).translate((op[0], op[1], op[2]))) 131 | visual_list.append(copy.deepcopy(mesh_box).translate((gp[0]-mesh_size/2.0, gp[1]-mesh_size/2.0, gp[2]-mesh_size/2.0))) 132 | for j in range(preds_ws[i].shape[0]): 133 | kp = preds_ws[i, j, :] 134 | if fear[i, :] > 0.5: 135 | visual_list.append(copy.deepcopy(small_sphere_fear).translate((kp[0], kp[1], kp[2] + visual_height))) 136 | else: 137 | visual_list.append(copy.deepcopy(small_sphere).translate((kp[0], kp[1], kp[2] + visual_height))) 138 | # Set line from odom to goal 139 | colors = [[0.99, 0.99, 0.1] for i in range(len(lines))] 140 | line_set = o3d.geometry.LineSet(o3d.utility.Vector3dVector(points), o3d.utility.Vector2iVector(lines)) 141 | line_set.colors = o3d.utility.Vector3dVector(colors) 142 | visual_list.append(line_set) 143 | 144 | o3d.visualization.draw_geometries(visual_list) 145 | return 146 | 147 | def VizImages(self, preds: torch.Tensor, waypoints: torch.Tensor, odom: torch.Tensor, goal: torch.Tensor, fear: torch.Tensor, images: torch.Tensor, visual_offset: float = 0.5, mesh_size: float = 0.3, is_shown: bool = True) -> List[np.ndarray]: 148 | """ 149 | Visualize images of the trajectory. 150 | 151 | Parameters: 152 | preds (tensor): Predicted Key points 153 | waypoints (tensor): Trajectory waypoints 154 | odom (tensor): Odometry data 155 | goal (tensor): Goal position 156 | fear (tensor): Fear value per waypoint 157 | images (tensor): Image data 158 | visual_offset (float): Offset for visualizing images 159 | mesh_size (float): Size of mesh objects in images 160 | is_shown (bool): If True, show images; otherwise, return image list 161 | """ 162 | batch_size, _, _ = waypoints.shape 163 | 164 | preds_ws = self.TransformPoints(odom, preds) 165 | wp_ws = self.TransformPoints(odom, waypoints) 166 | 167 | if goal.shape[-1] != 7: 168 | pp_goal = pp.identity_SE3(batch_size, device=goal.device) 169 | pp_goal.tensor()[:, 0:3] = goal 170 | goal = pp_goal.tensor() 171 | 172 | goal_ws = pp.SE3(odom) @ pp.SE3(goal) 173 | 174 | # Detach to CPU 175 | preds_ws = preds_ws.tensor()[:, :, 0:3].cpu().detach().numpy() 176 | wp_ws = wp_ws.tensor()[:, :, 0:3].cpu().detach().numpy() 177 | goal_ws = goal_ws.tensor()[:, 0:3].cpu().detach().numpy() 178 | 179 | # Adjust height 180 | preds_ws[:, :, 2] -= visual_offset 181 | wp_ws[:, :, 2] -= visual_offset 182 | goal_ws[:, 2] -= visual_offset 183 | 184 | # Set material shader 185 | mtl = o3d.visualization.rendering.MaterialRecord() 186 | mtl.base_color = [1.0, 1.0, 1.0, 0.3] 187 | mtl.shader = "defaultUnlit" 188 | 189 | # Set meshes 190 | small_sphere = o3d.geometry.TriangleMesh.create_sphere(mesh_size/20.0) # trajectory points 191 | mesh_sphere = o3d.geometry.TriangleMesh.create_sphere(mesh_size/5.0) # successful predict points 192 | mesh_sphere_fear = o3d.geometry.TriangleMesh.create_sphere(mesh_size/5.0) # unsuccessful predict points 193 | mesh_box = o3d.geometry.TriangleMesh.create_box(mesh_size, mesh_size, mesh_size) # end points 194 | 195 | # Set colors 196 | small_sphere.paint_uniform_color([0.99, 0.2, 0.1]) # green 197 | mesh_sphere.paint_uniform_color([0.4, 1.0, 0.1]) 198 | mesh_sphere_fear.paint_uniform_color([1.0, 0.64, 0.0]) 199 | mesh_box.paint_uniform_color([1.0, 0.64, 0.1]) 200 | 201 | # Init open3D render 202 | render = rendering.OffscreenRenderer(self.camera.width, self.camera.height) 203 | render.scene.set_background([0.0, 0.0, 0.0, 1.0]) # RGBA 204 | render.scene.scene.enable_sun_light(False) 205 | 206 | # compute veretx normals 207 | small_sphere.compute_vertex_normals() 208 | mesh_sphere.compute_vertex_normals() 209 | mesh_sphere_fear.compute_vertex_normals() 210 | mesh_box.compute_vertex_normals() 211 | 212 | wp_start_idx = 1 213 | cv_img_list = [] 214 | 215 | for i in range(batch_size): 216 | # Add geometries 217 | gp = goal_ws[i, :] 218 | 219 | # Add goal marker 220 | goal_mesh = copy.deepcopy(mesh_box).translate((gp[0]-mesh_size/2.0, gp[1]-mesh_size/2.0, gp[2]-mesh_size/2.0)) 221 | render.scene.add_geometry("goal_mesh", goal_mesh, mtl) 222 | 223 | # Add predictions 224 | for j in range(preds_ws[i, :, :].shape[0]): 225 | kp = preds_ws[i, j, :] 226 | if fear[i, :] > 0.5: 227 | kp_mesh = copy.deepcopy(mesh_sphere_fear).translate((kp[0], kp[1], kp[2])) 228 | else: 229 | kp_mesh = copy.deepcopy(mesh_sphere).translate((kp[0], kp[1], kp[2])) 230 | render.scene.add_geometry("keypose"+str(j), kp_mesh, mtl) 231 | 232 | # Add trajectory 233 | for k in range(wp_start_idx, wp_ws[i, :, :].shape[0]): 234 | wp = wp_ws[i, k, :] 235 | wp_mesh = copy.deepcopy(small_sphere).translate((wp[0], wp[1], wp[2])) 236 | render.scene.add_geometry("waypoint"+str(k), wp_mesh, mtl) 237 | 238 | # Set cameras 239 | self.CameraLookAtPose(odom[i, :], render, self.camera_tilt) 240 | 241 | # Project to image 242 | img_o3d = np.asarray(render.render_to_image()) 243 | mask = (img_o3d < 10).all(axis=2) 244 | 245 | # Attach image 246 | c_img = images[i, :, :].expand(3, -1, -1) 247 | c_img = c_img.cpu().detach().numpy().transpose(1, 2, 0) 248 | c_img = (c_img * 255 / np.max(c_img)).astype('uint8') 249 | img_o3d[mask, :] = c_img[mask, :] 250 | img_cv2 = cv2.cvtColor(img_o3d, cv2.COLOR_RGBA2BGRA) 251 | cv_img_list.append(img_cv2) 252 | 253 | # Visualize image 254 | if is_shown: 255 | cv2.imshow("Preview window", img_cv2) 256 | cv2.waitKey() 257 | 258 | # Clear render geometry 259 | render.scene.clear_geometry() 260 | 261 | return cv_img_list 262 | 263 | def CameraLookAtPose(self, odom: torch.Tensor, render: o3d.visualization.rendering.OffscreenRenderer, tilt: float) -> None: 264 | """ 265 | Set camera to look at at current odom pose. 266 | 267 | Parameters: 268 | odom (tensor): Odometry data 269 | render (OffscreenRenderer): Renderer object 270 | tilt (float): Tilt angle for camera 271 | """ 272 | unit_vec = pp.identity_SE3(device=odom.device) 273 | unit_vec.tensor()[0] = 1.0 274 | tilt_vec = [0, 0, 0] 275 | tilt_vec.extend(list(R.from_euler('y', tilt, degrees=False).as_quat())) 276 | tilt_vec = torch.tensor(tilt_vec, device=odom.device, dtype=odom.dtype) 277 | target_pose = pp.SE3(odom) @ pp.SE3(tilt_vec) @ unit_vec 278 | camera_up = [0, 0, 1] # camera orientation 279 | eye = pp.SE3(odom) 280 | eye = eye.tensor()[0:3].cpu().detach().numpy() 281 | target = target_pose.tensor()[0:3].cpu().detach().numpy() 282 | render.scene.camera.set_projection(self.camera.intrinsic_matrix, 0.1, 100.0, self.camera.width, self.camera.height) 283 | render.scene.camera.look_at(target, eye, camera_up) 284 | return 285 | -------------------------------------------------------------------------------- /iplanner/tsdf_map.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import open3d as o3d 11 | import numpy as np 12 | import torch 13 | import os 14 | 15 | torch.set_default_dtype(torch.float32) 16 | 17 | class TSDF_Map: 18 | def __init__(self, gpu_id=0): 19 | self.map_init = False 20 | if torch.cuda.is_available(): 21 | self.device = torch.device("cuda:" + str(gpu_id)) 22 | else: 23 | self.device = torch.device("cpu") 24 | self.pcd_tsdf = o3d.geometry.PointCloud() 25 | self.pcd_viz = o3d.geometry.PointCloud() 26 | 27 | # def DirectLoadMap(self, tsdf_array, viz_points, start_xy, voxel_size, clear_dist): 28 | def DirectLoadMap(self, data, coord, params): 29 | self.voxel_size = params[0] 30 | self.clear_dist = params[1] 31 | self.start_x, self.start_y = coord 32 | self.tsdf_array = torch.tensor(data[0], device=self.device) 33 | self.num_x, self.num_y = self.tsdf_array.shape 34 | # visualization points 35 | self.viz_points = data[1] 36 | self.pcd_viz.points = o3d.utility.Vector3dVector(self.viz_points) 37 | self.ground_array = torch.tensor(data[2], device=self.device) 38 | # set cost map 39 | self.SetUpCostArray() 40 | # update pcd instance 41 | xv, yv = np.meshgrid(np.linspace(0, self.num_x * self.voxel_size, self.num_x), np.linspace(0, self.num_y * self.voxel_size, self.num_y), indexing="ij") 42 | T = np.concatenate((np.expand_dims(xv, axis=0), np.expand_dims(yv, axis=0)), axis=0) 43 | T = np.concatenate((T, np.expand_dims(self.cost_array.cpu(), axis=0)), axis=0) 44 | self.pcd_tsdf.points = o3d.utility.Vector3dVector(T.reshape(3, -1).T) 45 | 46 | self.map_init = True; 47 | 48 | def ShowTSDFMap(self, cost_map=True): # not run with cuda 49 | if not self.map_init: 50 | print("Error: cannot show map, map has not been init yet!") 51 | return; 52 | if cost_map: 53 | o3d.visualization.draw_geometries([self.pcd_tsdf]) 54 | else: 55 | o3d.visualization.draw_geometries([self.pcd_viz]) 56 | return 57 | 58 | def Pos2Ind(self, points): 59 | # points [torch shapes [num_p, 3]] 60 | start_xy = torch.tensor([self.start_x, self.start_y], dtype=torch.float64, device=points.device).expand(1, 1, -1) 61 | H = (points.tensor()[:, :, 0:2] - start_xy) / self.voxel_size 62 | # H = torch.round(H) 63 | mask = torch.logical_and((H > 0).all(axis=2), (H < torch.tensor([self.num_x, self.num_y], device=points.device)[None,None,:]).all(axis=2)) 64 | return self.NormInds(H), H[mask, :] 65 | 66 | def NormInds(self, H): 67 | norm_matrix = torch.tensor([self.num_x/2.0, self.num_y/2.0], dtype=torch.float64, device=H.device) 68 | H = (H - norm_matrix) / norm_matrix 69 | return H 70 | 71 | def DeNormInds(self, NH): 72 | norm_matrix = torch.tensor([self.num_x/2.0, self.num_y/2.0], dtype=torch.float64, device=NH.device) 73 | NH = NH * norm_matrix + norm_matrix 74 | return NH 75 | 76 | def SaveTSDFMap(self, root_path, map_name): 77 | if not self.map_init: 78 | print("Error: map has not been init yet!") 79 | return; 80 | map_path = os.path.join(*[root_path, "maps", "data", map_name + "_map.txt"]) 81 | ground_path = os.path.join(*[root_path, "maps", "data", map_name + "_ground.txt"]) 82 | params_path = os.path.join(*[root_path, "maps", "params", map_name + "_param.txt"]) 83 | cloud_path = os.path.join(*[root_path, "maps", "cloud", map_name + "_cloud.txt"]) 84 | # save datas 85 | np.savetxt(map_path, self.tsdf_array.cpu()) 86 | np.savetxt(ground_path, self.ground_array.cpu()) 87 | np.savetxt(cloud_path, self.viz_points) 88 | params = [str(self.voxel_size), str(self.start_x), str(self.start_y), str(self.clear_dist)] 89 | with open(params_path, 'w') as f: 90 | for param in params: 91 | f.write(param) 92 | f.write('\n') 93 | print("TSDF Map saved.") 94 | 95 | def SetUpCostArray(self): 96 | self.cost_array = self.tsdf_array 97 | 98 | def ReadTSDFMap(self, root_path, map_name): 99 | map_path = os.path.join(*[root_path, "maps", "data", map_name + "_map.txt"]) 100 | ground_path = os.path.join(*[root_path, "maps", "data", map_name + "_ground.txt"]) 101 | params_path = os.path.join(*[root_path, "maps", "params", map_name + "_param.txt"]) 102 | cloud_path = os.path.join(*[root_path, "maps", "cloud", map_name + "_cloud.txt"]) 103 | # open params file 104 | with open(params_path) as f: 105 | content = f.readlines() 106 | self.voxel_size = float(content[0]) 107 | self.start_x = float(content[1]) 108 | self.start_y = float(content[2]) 109 | self.clear_dist = float(content[3]) 110 | self.tsdf_array = torch.tensor(np.loadtxt(map_path), device=self.device) 111 | self.viz_points = np.loadtxt(cloud_path) 112 | self.ground_array = torch.tensor(np.loadtxt(ground_path), device=self.device) 113 | 114 | self.num_x, self.num_y = self.tsdf_array.shape 115 | # visualization points 116 | self.pcd_viz.points = o3d.utility.Vector3dVector(self.viz_points) 117 | # opne map array 118 | self.SetUpCostArray() 119 | # update pcd instance 120 | xv, yv = np.meshgrid(np.linspace(0, self.num_x * self.voxel_size, self.num_x), np.linspace(0, self.num_y * self.voxel_size, self.num_y), indexing="ij") 121 | T = np.concatenate((np.expand_dims(xv, axis=0), np.expand_dims(yv, axis=0)), axis=0) 122 | T = np.concatenate((T, np.expand_dims(self.cost_array.cpu().detach().numpy(), axis=0)), axis=0) 123 | wps = T.reshape(3, -1).T + np.array([self.start_x, self.start_y, 0.0]) 124 | self.pcd_tsdf.points = o3d.utility.Vector3dVector(wps) 125 | 126 | self.map_init = True; 127 | return -------------------------------------------------------------------------------- /iplanner_env.yml: -------------------------------------------------------------------------------- 1 | name: iplanner 2 | channels: 3 | - pytorch 4 | - nvidia 5 | - anaconda 6 | - conda-forge 7 | - defaults 8 | dependencies: 9 | - _libgcc_mutex=0.1=main 10 | - _openmp_mutex=5.1=1_gnu 11 | - appdirs=1.4.4=pyh9f0ad1d_0 12 | - attrs=22.1.0=py38h06a4308_0 13 | - beautifulsoup4=4.11.1=py38h06a4308_0 14 | - blas=1.0=mkl 15 | - bleach=4.1.0=pyhd3eb1b0_0 16 | - brotlipy=0.7.0=py38h27cfd23_1003 17 | - bzip2=1.0.8=h7b6447c_0 18 | - ca-certificates=2023.5.7=hbcca054_0 19 | - catkin_pkg=0.5.2=pyhd8ed1ab_0 20 | - certifi=2023.5.7=pyhd8ed1ab_0 21 | - cffi=1.15.1=py38h5eee18b_3 22 | - charset-normalizer=2.0.4=pyhd3eb1b0_0 23 | - click=8.1.3=unix_pyhd8ed1ab_2 24 | - cryptography=39.0.1=py38h9ce1e76_0 25 | - cuda-cudart=11.8.89=0 26 | - cuda-cupti=11.8.87=0 27 | - cuda-libraries=11.8.0=0 28 | - cuda-nvrtc=11.8.89=0 29 | - cuda-nvtx=11.8.86=0 30 | - cuda-runtime=11.8.0=0 31 | - defusedxml=0.7.1=pyhd3eb1b0_0 32 | - distro=1.8.0=pyhd8ed1ab_0 33 | - docker-pycreds=0.4.0=py_0 34 | - docutils=0.20.1=py38h578d9bd_0 35 | - entrypoints=0.4=py38h06a4308_0 36 | - ffmpeg=4.3=hf484d3e_0 37 | - filelock=3.9.0=py38h06a4308_0 38 | - freetype=2.12.1=h4a9f257_0 39 | - giflib=5.2.1=h5eee18b_3 40 | - gitdb=4.0.10=pyhd8ed1ab_0 41 | - gitpython=3.1.31=pyhd8ed1ab_0 42 | - gmp=6.2.1=h295c915_3 43 | - gmpy2=2.1.2=py38heeb90bb_0 44 | - gnutls=3.6.15=he1e5248_0 45 | - icu=58.2=he6710b0_3 46 | - idna=3.4=py38h06a4308_0 47 | - importlib_resources=5.2.0=pyhd3eb1b0_1 48 | - intel-openmp=2023.1.0=hdb19cb5_46305 49 | - jpeg=9e=h5eee18b_1 50 | - jsonschema=4.16.0=py38h06a4308_0 51 | - jupyter_client=7.4.8=py38h06a4308_0 52 | - jupyter_core=5.1.1=py38h06a4308_0 53 | - jupyterlab_pygments=0.1.2=py_0 54 | - lame=3.100=h7b6447c_0 55 | - lcms2=2.12=h3be6417_0 56 | - ld_impl_linux-64=2.38=h1181459_1 57 | - lerc=3.0=h295c915_0 58 | - libcublas=11.11.3.6=0 59 | - libcufft=10.9.0.58=0 60 | - libcufile=1.6.1.9=0 61 | - libcurand=10.3.2.106=0 62 | - libcusolver=11.4.1.48=0 63 | - libcusparse=11.7.5.86=0 64 | - libdeflate=1.17=h5eee18b_0 65 | - libffi=3.4.4=h6a678d5_0 66 | - libgcc-ng=11.2.0=h1234567_1 67 | - libgomp=11.2.0=h1234567_1 68 | - libiconv=1.16=h7f8727e_2 69 | - libidn2=2.3.4=h5eee18b_0 70 | - libnpp=11.8.0.86=0 71 | - libnvjpeg=11.9.0.86=0 72 | - libpng=1.6.39=h5eee18b_0 73 | - libprotobuf=3.20.3=he621ea3_0 74 | - libsodium=1.0.18=h7b6447c_0 75 | - libstdcxx-ng=11.2.0=h1234567_1 76 | - libtasn1=4.19.0=h5eee18b_0 77 | - libtiff=4.5.0=h6a678d5_2 78 | - libunistring=0.9.10=h27cfd23_0 79 | - libwebp=1.2.4=h11a3e52_1 80 | - libwebp-base=1.2.4=h5eee18b_1 81 | - libxml2=2.9.14=h74e7548_0 82 | - libxslt=1.1.35=h4e12654_0 83 | - lxml=4.9.1=py38h1edc446_0 84 | - lz4-c=1.9.4=h6a678d5_0 85 | - markupsafe=2.1.1=py38h7f8727e_0 86 | - mistune=0.8.4=py38h7b6447c_1000 87 | - mkl=2023.1.0=h6d00ec8_46342 88 | - mkl-service=2.4.0=py38h5eee18b_1 89 | - mkl_fft=1.3.6=py38h417a72b_1 90 | - mkl_random=1.2.2=py38h417a72b_1 91 | - mpc=1.1.0=h10f8cd9_1 92 | - mpfr=4.0.2=hb69a4c5_1 93 | - mpmath=1.2.1=py38h06a4308_0 94 | - nbclient=0.5.13=py38h06a4308_0 95 | - nbconvert=6.5.4=py38h06a4308_0 96 | - nbformat=5.7.0=py38h06a4308_0 97 | - ncurses=6.4=h6a678d5_0 98 | - nest-asyncio=1.5.6=py38h06a4308_0 99 | - nettle=3.7.3=hbbd107a_1 100 | - networkx=2.8.4=py38h06a4308_1 101 | - numpy=1.23.5=py38hf6e8229_1 102 | - numpy-base=1.23.5=py38h060ed82_1 103 | - openh264=2.1.1=h4ff587b_0 104 | - openssl=1.1.1t=h7f8727e_0 105 | - packaging=22.0=py38h06a4308_0 106 | - pandocfilters=1.5.0=pyhd3eb1b0_0 107 | - pathtools=0.1.2=py_1 108 | - pillow=9.4.0=py38h6a678d5_0 109 | - pip=23.0.1=py38h06a4308_0 110 | - pkgutil-resolve-name=1.3.10=py38h06a4308_0 111 | - platformdirs=2.5.2=py38h06a4308_0 112 | - protobuf=3.20.3=py38h6a678d5_0 113 | - psutil=5.9.0=py38h5eee18b_0 114 | - pycparser=2.21=pyhd3eb1b0_0 115 | - pygments=2.11.2=pyhd3eb1b0_0 116 | - pyopenssl=23.0.0=py38h06a4308_0 117 | - pyparsing=3.0.9=pyhd8ed1ab_0 118 | - pyrsistent=0.18.0=py38heee7806_0 119 | - pysocks=1.7.1=py38h06a4308_0 120 | - python=3.8.16=h7a1cb2a_3 121 | - python-dateutil=2.8.2=pyhd3eb1b0_0 122 | - python-fastjsonschema=2.16.2=py38h06a4308_0 123 | - python_abi=3.8=2_cp38 124 | - pytorch=2.0.1=py3.8_cuda11.8_cudnn8.7.0_0 125 | - pytorch-cuda=11.8=h7e8668a_5 126 | - pytorch-mutex=1.0=cuda 127 | - pyyaml=6.0=py38h0a891b7_4 128 | - pyzmq=23.2.0=py38h6a678d5_0 129 | - readline=8.2=h5eee18b_0 130 | - requests=2.29.0=py38h06a4308_0 131 | - rospkg=1.5.0=pyhd8ed1ab_0 132 | - sentry-sdk=1.21.1=pyhd8ed1ab_0 133 | - setproctitle=1.2.2=py38h0a891b7_2 134 | - setuptools=67.8.0=py38h06a4308_0 135 | - six=1.16.0=pyhd3eb1b0_1 136 | - smmap=3.0.5=pyh44b312d_0 137 | - soupsieve=2.3.2.post1=py38h06a4308_0 138 | - sqlite=3.41.2=h5eee18b_0 139 | - sympy=1.11.1=py38h06a4308_0 140 | - tbb=2021.8.0=hdb19cb5_0 141 | - tinycss2=1.2.1=py38h06a4308_0 142 | - tk=8.6.12=h1ccaba5_0 143 | - torchtriton=2.0.0=py38 144 | - torchvision=0.15.2=py38_cu118 145 | - tornado=6.2=py38h5eee18b_0 146 | - traitlets=5.7.1=py38h06a4308_0 147 | - typing_extensions=4.5.0=py38h06a4308_0 148 | - urllib3=1.26.15=py38h06a4308_0 149 | - wandb=0.15.3=pyhd8ed1ab_0 150 | - webencodings=0.5.1=py38_1 151 | - wheel=0.38.4=py38h06a4308_0 152 | - xz=5.4.2=h5eee18b_0 153 | - yaml=0.2.5=h7f98852_2 154 | - zeromq=4.3.4=h2531618_0 155 | - zipp=3.11.0=py38h06a4308_0 156 | - zlib=1.2.13=h5eee18b_0 157 | - zstd=1.5.5=hc292b87_0 158 | - pip: 159 | - bagpy==0.5 160 | - configargparse==1.5.3 161 | - contourpy==1.0.7 162 | - cycler==0.11.0 163 | - dash==2.10.1 164 | - dash-core-components==2.0.0 165 | - dash-html-components==2.0.0 166 | - dash-table==5.0.0 167 | - decorator==5.1.1 168 | - flask==2.2.3 169 | - fonttools==4.39.4 170 | - itsdangerous==2.1.2 171 | - jinja2==3.0.3 172 | - kiwisolver==1.4.4 173 | - matplotlib==3.7.1 174 | - open3d==0.17.0 175 | - opencv-python==4.7.0.72 176 | - pandas==2.0.2 177 | - pexpect==4.8.0 178 | - plotly==5.14.1 179 | - pluggy==1.0.0 180 | - pycryptodomex==3.18.0 181 | - pypose==0.4.3 182 | - pyquaternion==0.9.9 183 | - pytest==7.1.2 184 | - rosnumpy==0.0.5.2 185 | - scipy==1.10.1 186 | - tenacity==8.2.2 187 | - wcwidth==0.2.6 188 | - werkzeug==2.2.3 189 | prefix: /home/fanyang/anaconda3/envs/iplanner 190 | -------------------------------------------------------------------------------- /launch/data_collector.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /launch/iplanner.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | -------------------------------------------------------------------------------- /launch/iplanner_viz.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | -------------------------------------------------------------------------------- /package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | iplanner_node 4 | 1.0.0 5 | A ROS warper for Imperative Planner (iPlanner) package 6 | 7 | Fan Yang 8 | 9 | BSD-3 10 | 11 | catkin 12 | roscpp 13 | rospy 14 | std_msgs 15 | sensor_msgs 16 | geometry_msgs 17 | image_transport 18 | python-numpy 19 | 20 | roscpp 21 | rospy 22 | std_msgs 23 | geometry_msgs 24 | python-numpy 25 | 26 | python-numpy 27 | roscpp 28 | rospy 29 | std_msgs 30 | geometry_msgs 31 | sensor_msgs 32 | 33 | -------------------------------------------------------------------------------- /rviz/default.rviz: -------------------------------------------------------------------------------- 1 | Panels: 2 | - Class: rviz/Displays 3 | Help Height: 0 4 | Name: Displays 5 | Property Tree Widget: 6 | Expanded: ~ 7 | Splitter Ratio: 0.729411780834198 8 | Tree Height: 419 9 | - Class: rviz/Selection 10 | Name: Selection 11 | - Class: rviz/Tool Properties 12 | Expanded: ~ 13 | Name: Tool Properties 14 | Splitter Ratio: 0.5886790156364441 15 | - Class: rviz/Views 16 | Expanded: 17 | - /Current View1 18 | Name: Views 19 | Splitter Ratio: 0.5 20 | - Class: rviz/Time 21 | Name: Time 22 | SyncMode: 0 23 | SyncSource: OverallMap 24 | Preferences: 25 | PromptSaveOnExit: true 26 | Toolbars: 27 | toolButtonStyle: 2 28 | Visualization Manager: 29 | Class: "" 30 | Displays: 31 | - Alpha: 1 32 | Class: rviz/Axes 33 | Enabled: true 34 | Length: 1 35 | Name: Vehicle 36 | Radius: 0.11999999731779099 37 | Reference Frame: vehicle 38 | Show Trail: false 39 | Value: true 40 | - Class: rviz/Image 41 | Enabled: true 42 | Image Topic: /iplanner_image 43 | Max Value: 1 44 | Median window: 5 45 | Min Value: 0 46 | Name: iPlannerImage 47 | Normalize Range: true 48 | Queue Size: 2 49 | Transport Hint: raw 50 | Unreliable: false 51 | Value: true 52 | - Alpha: 0.20000000298023224 53 | Autocompute Intensity Bounds: true 54 | Autocompute Value Bounds: 55 | Max Value: 0.9582680463790894 56 | Min Value: -0.043718695640563965 57 | Value: true 58 | Axis: Z 59 | Channel Name: intensity 60 | Class: rviz/PointCloud2 61 | Color: 255; 255; 255 62 | Color Transformer: FlatColor 63 | Decay Time: 200 64 | Enabled: false 65 | Invert Rainbow: false 66 | Max Color: 255; 255; 255 67 | Min Color: 0; 0; 0 68 | Name: RegScan 69 | Position Transformer: XYZ 70 | Queue Size: 10 71 | Selectable: true 72 | Size (Pixels): 1.5 73 | Size (m): 0.05000000074505806 74 | Style: Points 75 | Topic: /registered_scan 76 | Unreliable: false 77 | Use Fixed Frame: true 78 | Use rainbow: true 79 | Value: false 80 | - Alpha: 0.20000000298023224 81 | Autocompute Intensity Bounds: true 82 | Autocompute Value Bounds: 83 | Max Value: 0.9582680463790894 84 | Min Value: -0.043718695640563965 85 | Value: true 86 | Axis: Z 87 | Channel Name: intensity 88 | Class: rviz/PointCloud2 89 | Color: 255; 255; 255 90 | Color Transformer: FlatColor 91 | Decay Time: 10 92 | Enabled: false 93 | Invert Rainbow: false 94 | Max Color: 255; 255; 255 95 | Min Color: 0; 0; 0 96 | Name: SensorScan 97 | Position Transformer: XYZ 98 | Queue Size: 10 99 | Selectable: true 100 | Size (Pixels): 2 101 | Size (m): 0.10000000149011612 102 | Style: Points 103 | Topic: /sensor_scan 104 | Unreliable: false 105 | Use Fixed Frame: true 106 | Use rainbow: true 107 | Value: false 108 | - Alpha: 1 109 | Autocompute Intensity Bounds: false 110 | Autocompute Value Bounds: 111 | Max Value: 10 112 | Min Value: -10 113 | Value: true 114 | Axis: Z 115 | Channel Name: intensity 116 | Class: rviz/PointCloud2 117 | Color: 255; 0; 0 118 | Color Transformer: Intensity 119 | Decay Time: 0 120 | Enabled: false 121 | Invert Rainbow: true 122 | Max Color: 255; 255; 255 123 | Max Intensity: 0.10000000149011612 124 | Min Color: 0; 0; 0 125 | Min Intensity: -0.10000000149011612 126 | Name: TerrainMap 127 | Position Transformer: XYZ 128 | Queue Size: 10 129 | Selectable: true 130 | Size (Pixels): 2 131 | Size (m): 0.03999999910593033 132 | Style: Points 133 | Topic: /terrain_map 134 | Unreliable: false 135 | Use Fixed Frame: true 136 | Use rainbow: true 137 | Value: false 138 | - Alpha: 1 139 | Autocompute Intensity Bounds: false 140 | Autocompute Value Bounds: 141 | Max Value: 10 142 | Min Value: -10 143 | Value: true 144 | Axis: Z 145 | Channel Name: intensity 146 | Class: rviz/PointCloud2 147 | Color: 255; 0; 0 148 | Color Transformer: Intensity 149 | Decay Time: 0 150 | Enabled: false 151 | Invert Rainbow: true 152 | Max Color: 255; 255; 255 153 | Max Intensity: 0.20000000298023224 154 | Min Color: 0; 0; 0 155 | Min Intensity: -0.20000000298023224 156 | Name: TerrainMapExt 157 | Position Transformer: XYZ 158 | Queue Size: 10 159 | Selectable: true 160 | Size (Pixels): 4 161 | Size (m): 0.03999999910593033 162 | Style: Points 163 | Topic: /terrain_map_ext 164 | Unreliable: false 165 | Use Fixed Frame: true 166 | Use rainbow: true 167 | Value: false 168 | - Alpha: 1 169 | Buffer Length: 1 170 | Class: rviz/Path 171 | Color: 25; 255; 0 172 | Enabled: true 173 | Head Diameter: 0.30000001192092896 174 | Head Length: 0.20000000298023224 175 | Length: 0.30000001192092896 176 | Line Style: Billboards 177 | Line Width: 0.05000000074505806 178 | Name: Path 179 | Offset: 180 | X: 0 181 | Y: 0 182 | Z: 0 183 | Pose Color: 255; 85; 255 184 | Pose Style: None 185 | Queue Size: 10 186 | Radius: 0.029999999329447746 187 | Shaft Diameter: 0.10000000149011612 188 | Shaft Length: 0.10000000149011612 189 | Topic: /path 190 | Unreliable: false 191 | Value: true 192 | - Alpha: 1 193 | Class: rviz/PointStamped 194 | Color: 204; 41; 204 195 | Enabled: true 196 | History Length: 1 197 | Name: Waypoint 198 | Queue Size: 10 199 | Radius: 0.5 200 | Topic: /way_point 201 | Unreliable: false 202 | Value: true 203 | - Alpha: 0.30000001192092896 204 | Autocompute Intensity Bounds: true 205 | Autocompute Value Bounds: 206 | Max Value: 10 207 | Min Value: -10 208 | Value: true 209 | Axis: Z 210 | Channel Name: intensity 211 | Class: rviz/PointCloud2 212 | Color: 255; 255; 255 213 | Color Transformer: FlatColor 214 | Decay Time: 0 215 | Enabled: true 216 | Invert Rainbow: false 217 | Max Color: 255; 255; 255 218 | Min Color: 0; 0; 0 219 | Name: OverallMap 220 | Position Transformer: XYZ 221 | Queue Size: 10 222 | Selectable: true 223 | Size (Pixels): 1.5 224 | Size (m): 0.019999999552965164 225 | Style: Flat Squares 226 | Topic: /overall_map 227 | Unreliable: false 228 | Use Fixed Frame: true 229 | Use rainbow: true 230 | Value: true 231 | - Alpha: 0.20000000298023224 232 | Autocompute Intensity Bounds: true 233 | Autocompute Value Bounds: 234 | Max Value: 10 235 | Min Value: -10 236 | Value: true 237 | Axis: Z 238 | Channel Name: intensity 239 | Class: rviz/PointCloud2 240 | Color: 0; 170; 255 241 | Color Transformer: FlatColor 242 | Decay Time: 0 243 | Enabled: false 244 | Invert Rainbow: false 245 | Max Color: 255; 255; 255 246 | Min Color: 0; 0; 0 247 | Name: ExploredAreas 248 | Position Transformer: XYZ 249 | Queue Size: 10 250 | Selectable: true 251 | Size (Pixels): 3 252 | Size (m): 0.009999999776482582 253 | Style: Points 254 | Topic: /explored_areas 255 | Unreliable: false 256 | Use Fixed Frame: true 257 | Use rainbow: true 258 | Value: false 259 | - Alpha: 1 260 | Autocompute Intensity Bounds: true 261 | Autocompute Value Bounds: 262 | Max Value: 10 263 | Min Value: -10 264 | Value: true 265 | Axis: Z 266 | Channel Name: intensity 267 | Class: rviz/PointCloud2 268 | Color: 255; 255; 255 269 | Color Transformer: Intensity 270 | Decay Time: 0 271 | Enabled: true 272 | Invert Rainbow: false 273 | Max Color: 255; 255; 255 274 | Min Color: 0; 0; 0 275 | Name: Trajectory 276 | Position Transformer: XYZ 277 | Queue Size: 10 278 | Selectable: true 279 | Size (Pixels): 10 280 | Size (m): 0.009999999776482582 281 | Style: Points 282 | Topic: /trajectory 283 | Unreliable: false 284 | Use Fixed Frame: true 285 | Use rainbow: true 286 | Value: true 287 | - Alpha: 1 288 | Buffer Length: 1 289 | Class: rviz/Path 290 | Color: 204; 0; 0 291 | Enabled: true 292 | Head Diameter: 0.30000001192092896 293 | Head Length: 0.20000000298023224 294 | Length: 0.30000001192092896 295 | Line Style: Billboards 296 | Line Width: 0.05000000074505806 297 | Name: FearPath 298 | Offset: 299 | X: 0 300 | Y: 0 301 | Z: 0 302 | Pose Color: 255; 85; 255 303 | Pose Style: None 304 | Queue Size: 10 305 | Radius: 0.029999999329447746 306 | Shaft Diameter: 0.10000000149011612 307 | Shaft Length: 0.10000000149011612 308 | Topic: /iplanner_path_fear 309 | Unreliable: false 310 | Value: true 311 | - Alpha: 0.20000000298023224 312 | Autocompute Intensity Bounds: true 313 | Autocompute Value Bounds: 314 | Max Value: 2.655569553375244 315 | Min Value: -2.58453106880188 316 | Value: true 317 | Axis: Z 318 | Channel Name: intensity 319 | Class: rviz/PointCloud2 320 | Color: 255; 255; 255 321 | Color Transformer: AxisColor 322 | Decay Time: 0 323 | Enabled: true 324 | Invert Rainbow: false 325 | Max Color: 255; 255; 255 326 | Min Color: 0; 0; 0 327 | Name: DepthCamera 328 | Position Transformer: XYZ 329 | Queue Size: 10 330 | Selectable: true 331 | Size (Pixels): 3 332 | Size (m): 0.019999999552965164 333 | Style: Flat Squares 334 | Topic: /rgbd_camera/depth/points 335 | Unreliable: false 336 | Use Fixed Frame: true 337 | Use rainbow: true 338 | Value: true 339 | - Class: rviz/Image 340 | Enabled: true 341 | Image Topic: /rgbd_camera/color/image 342 | Max Value: 1 343 | Median window: 5 344 | Min Value: 0 345 | Name: ColorImage 346 | Normalize Range: true 347 | Queue Size: 2 348 | Transport Hint: raw 349 | Unreliable: false 350 | Value: true 351 | - Alpha: 1 352 | Buffer Length: 1 353 | Class: rviz/Path 354 | Color: 239; 41; 41 355 | Enabled: true 356 | Head Diameter: 0.30000001192092896 357 | Head Length: 0.20000000298023224 358 | Length: 0.30000001192092896 359 | Line Style: Billboards 360 | Line Width: 0.05000000074505806 361 | Name: CollisionPath 362 | Offset: 363 | X: 0 364 | Y: 0 365 | Z: 0 366 | Pose Color: 255; 85; 255 367 | Pose Style: None 368 | Queue Size: 10 369 | Radius: 0.029999999329447746 370 | Shaft Diameter: 0.10000000149011612 371 | Shaft Length: 0.10000000149011612 372 | Topic: /path_fear 373 | Unreliable: false 374 | Value: true 375 | Enabled: true 376 | Global Options: 377 | Background Color: 0; 0; 0 378 | Default Light: true 379 | Fixed Frame: map 380 | Frame Rate: 15 381 | Name: root 382 | Tools: 383 | - Class: rviz/MoveCamera 384 | - Class: rviz/WaypointTool 385 | Topic: waypoint 386 | Value: true 387 | Views: 388 | Current: 389 | Angle: 0 390 | Class: rviz/TopDownOrtho 391 | Enable Stereo Rendering: 392 | Stereo Eye Separation: 0.05999999865889549 393 | Stereo Focal Distance: 1 394 | Swap Stereo Eyes: false 395 | Value: false 396 | Invert Z Axis: false 397 | Name: Current View 398 | Near Clip Distance: 0.009999999776482582 399 | Scale: 46.7901496887207 400 | Target Frame: vehicle 401 | X: 0 402 | Y: 0 403 | Saved: ~ 404 | Window Geometry: 405 | ColorImage: 406 | collapsed: false 407 | Displays: 408 | collapsed: false 409 | Height: 1349 410 | Hide Left Dock: false 411 | Hide Right Dock: true 412 | QMainWindow State: 000000ff00000000fd00000004000000000000030d000004ebfc020000000bfb0000001200530065006c0065006300740069006f006e000000003d000001b80000005c00fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000a3fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afb0000001a00690050006c0061006e006e006500720049006d006100670065010000003d000002740000001600fffffffc000002b700000271000000e60100001cfa000000000100000002fb000000140043006f006c006f00720049006d0061006700650100000000ffffffff0000008300fffffffb000000100044006900730070006c00610079007301000000000000030d0000015600fffffffb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261fb0000001400430061006d0065007200610049006e006900740100000474000000160000000000000000fb0000000a0049006d00610067006503000004000000035600000199000000fd0000000100000294000004ebfc0200000003fb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a00560069006500770073000000003d000004eb000000a400fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000490000000a9fc0100000001fb0000000a00560069006500770073030000004e00000080000002e100000197000000030000073d0000003efc0100000002fb0000000800540069006d006500000000000000073d0000041800fffffffb0000000800540069006d00650100000000000004500000000000000000000006ed000004eb00000004000000040000000800000008fc0000000100000002000000010000000a0054006f006f006c00730100000000ffffffff0000000000000000 413 | Selection: 414 | collapsed: false 415 | Time: 416 | collapsed: false 417 | Tool Properties: 418 | collapsed: false 419 | Views: 420 | collapsed: true 421 | Width: 2560 422 | X: 2560 423 | Y: 0 424 | iPlannerImage: 425 | collapsed: false 426 | -------------------------------------------------------------------------------- /src/data_collect_node.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | 11 | import rospy 12 | import open3d as o3d 13 | import message_filters 14 | from sensor_msgs.msg import PointCloud2 15 | from nav_msgs.msg import Odometry 16 | from sensor_msgs.msg import Image 17 | from sensor_msgs.msg import CameraInfo 18 | import numpy as np 19 | import ros_numpy 20 | from scipy.spatial.transform import Rotation 21 | # OpenCV2 for saving an image 22 | import sys 23 | import cv2 24 | import tf 25 | import os 26 | import rospkg 27 | 28 | rospack = rospkg.RosPack() 29 | pack_path = rospack.get_path('iplanner_node') 30 | planner_path = os.path.join(pack_path,'iplanner') 31 | sys.path.append(pack_path) 32 | sys.path.append(planner_path) 33 | 34 | from iplanner.rosutil import ROSArgparse 35 | 36 | class DataCollector: 37 | def __init__(self, args): 38 | super(DataCollector, self).__init__() 39 | self.__config(args) 40 | 41 | # internal values 42 | self.__tf_listener = tf.TransformListener() 43 | rospy.sleep(2.5) # wait for tf listener to be ready 44 | 45 | self.__odom_extrinsic = None 46 | self.__cv2_img_depth = np.ndarray([640, 360]) 47 | self.__cv2_img_cam = np.ndarray([640, 360]) 48 | self.__odom_list = [] 49 | self.__pcd = o3d.geometry.PointCloud() 50 | self.__init_check_dics = {"color": 0, "depth": 0, "scan_extrinsic": 0, "camera_extrinsic": 0, "odometry_extrinsic": 0} 51 | 52 | depth_sub = message_filters.Subscriber(self.__depth_topic, Image) 53 | image_sub = message_filters.Subscriber(self.__color_topic, Image) 54 | odom_sub = message_filters.Subscriber(self.__odom_topic, Odometry) 55 | scan_sub = message_filters.Subscriber(self.__scan_topic, PointCloud2) 56 | 57 | time_sync_thred = 0.01 # second 58 | ts = message_filters.ApproximateTimeSynchronizer([image_sub, depth_sub, scan_sub, odom_sub], 50, time_sync_thred) 59 | ts.registerCallback(self.__syncCallback) 60 | 61 | # camera info subscriber 62 | rospy.Subscriber(self.__depth_info_topic, CameraInfo, self.__depthInfoCallback, (self.__depth_intrinc_path)) 63 | rospy.Subscriber(self.__color_info_topic, CameraInfo, self.__colorInfoCallback, (self.__color_intrinc_path)) 64 | 65 | print("deleting previous files, if any ...") 66 | folder_list = ["depth", "camera", "scan"] 67 | for folder_name in folder_list: 68 | dir_path = os.path.join(*[self.__root_path, folder_name]) 69 | if not os.path.exists(dir_path): 70 | os.makedirs(dir_path) 71 | for efile in os.listdir(dir_path): 72 | os.remove(os.path.join(dir_path, efile)) 73 | 74 | self.__odom_file_name = os.path.join(self.__root_path, "odom_ground_truth.txt") 75 | open(self.__odom_file_name, 'w').close() # clear txt file 76 | return 77 | 78 | def __config(self, args): 79 | self.__main_freq = args.main_freq 80 | self.__root_path = args.root_path 81 | self.__depth_topic = args.depth_topic 82 | self.__color_topic = args.color_topic 83 | self.__odom_topic = args.odom_topic 84 | self.__scan_topic = args.scan_topic 85 | self.__depth_info_topic = args.depth_info_topic 86 | self.__color_info_topic = args.color_info_topic 87 | self.__camera_frame_id = args.camera_frame_id 88 | self.__scan_frame_id = args.scan_frame_id 89 | self.__base_frame_id = args.base_frame_id 90 | self.__odom_associate_id = args.odom_associate_id 91 | 92 | self.__depth_intrinc_path = os.path.join(self.__root_path, "depth_intrinsic.txt") 93 | self.__color_intrinc_path = os.path.join(self.__root_path, "color_intrinsic.txt") 94 | self.__camera_extrinc_path = os.path.join(self.__root_path, "camera_extrinsic.txt") 95 | self.__scan_extrinc_path = os.path.join(self.__root_path, "scan_extrinsic.txt") 96 | # print(self.__depth_intrinc_path) 97 | return 98 | 99 | def spin(self): 100 | r = rospy.Rate(self.__main_freq) 101 | time_step = 0 102 | last_odom = [] 103 | odom_file = open(self.__odom_file_name, 'w') 104 | while not rospy.is_shutdown(): 105 | # listen to camera extrinsic 106 | if (self.__init_check_dics["camera_extrinsic"] == 0): 107 | try: 108 | (pos, ori) = self.__tf_listener.lookupTransform(self.__base_frame_id, self.__camera_frame_id, rospy.Time(0)) 109 | self.__writeExtrinstic(pos, ori, self.__camera_extrinc_path, "camera_extrinsic") 110 | except (tf.Exception, tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException): 111 | print("Wait to get camera extrinsic.") 112 | if (self.__init_check_dics["scan_extrinsic"] == 0): 113 | try: 114 | (pos, ori) = self.__tf_listener.lookupTransform(self.__base_frame_id, self.__scan_frame_id, rospy.Time(0)) 115 | self.__writeExtrinstic(pos, ori, self.__scan_extrinc_path, "scan_extrinsic") 116 | except (tf.Exception, tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException): 117 | print("Wait to get scan extrinsic.") 118 | # listen to odom extrinsic 119 | if (self.__init_check_dics["odometry_extrinsic"] == 0): 120 | try: 121 | (pos, ori) = self.__tf_listener.lookupTransform(self.__base_frame_id, self.__odom_associate_id, rospy.Time(0)) 122 | self.__odom_extrinsic = np.eye(4) 123 | self.__odom_extrinsic[:3, :3] = Rotation.from_quat(ori).as_matrix() 124 | self.__odom_extrinsic[:3,3] = np.array(pos) 125 | self.__init_check_dics["odometry_extrinsic"] = 1 126 | print("odometry_extrinsic" + " save succeed.") 127 | except (tf.Exception, tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException): 128 | print("Wait to get odometry extrinsic.") 129 | 130 | # save images, depths and odom 131 | if len(self.__odom_list) > 0: 132 | if len(last_odom) > 0 and np.linalg.norm(np.array(last_odom) - np.array(self.__odom_list)) < 1e-2: 133 | self.__odom_list.clear() 134 | r.sleep() 135 | continue 136 | odom_file.writelines(str(self.__odom_list)) 137 | odom_file.write('\n') 138 | file_name_depth = os.path.join(*[self.__root_path, "depth", str(time_step) + ".png"]) 139 | file_name_camera = os.path.join(*[self.__root_path, "camera", str(time_step) + ".png"]) 140 | file_name_scan = os.path.join(*[self.__root_path, "scan", str(time_step) + ".ply"]) 141 | self.__saveDepthImage(file_name_depth, self.__cv2_img_depth) 142 | cv2.imwrite(file_name_camera, self.__cv2_img_cam) 143 | o3d.io.write_point_cloud(file_name_scan, self.__pcd) 144 | time_step = time_step + 1 145 | print("save current idx: %d", time_step) 146 | last_odom = self.__odom_list.copy() 147 | r.sleep() 148 | 149 | odom_file.close() 150 | rospy.spin() 151 | return 152 | 153 | def __syncCallback(self, image, depth, scan, odom): 154 | if self.__init_check_dics["odometry_extrinsic"] == 0: 155 | return 156 | self.__cv2_img_cam = ros_numpy.numpify(image) 157 | self.__cv2_img_depth = ros_numpy.numpify(depth) 158 | self.__updateScanPoints(scan, self.__pcd) 159 | 160 | self.__cv2_img_depth[~np.isfinite(self.__cv2_img_depth)] = 0.0 161 | 162 | pos = odom.pose.pose.position 163 | ori = odom.pose.pose.orientation 164 | if not self.__odom_associate_id == self.__base_frame_id: 165 | trans = np.eye(4) 166 | trans[:3,:3] = Rotation.from_quat([ori.x, ori.y, ori.z, ori.w]).as_matrix() 167 | trans[:3,3] = np.array([pos.x, pos.y, pos.z]) 168 | trans = trans @ self.__odom_extrinsic 169 | ori = Rotation.from_matrix(trans[:3,:3]).as_quat().tolist() 170 | self.__odom_list = trans[:3,3].tolist() 171 | self.__odom_list.extend(ori) 172 | else: 173 | self.__odom_list = [pos.x, pos.y, pos.z, ori.x, ori.y, ori.z, ori.w] 174 | return 175 | 176 | 177 | def __saveDepthImage(self, path, image, scale=1000.0): 178 | image[~np.isfinite(image)] = 0.0 179 | if not image.dtype.name == 'uint16': 180 | image = image * scale 181 | image = np.uint16(image) 182 | cv2.imwrite(path, image) 183 | return 184 | 185 | def __updateScanPoints(self, pc_msg, pcd): 186 | pc_np = ros_numpy.point_cloud2.pointcloud2_to_xyz_array(pc_msg, remove_nans=True) 187 | pcd.clear() 188 | pcd.points = o3d.utility.Vector3dVector(pc_np) 189 | 190 | 191 | def __depthInfoCallback(self, depth_info, args): 192 | if (self.__init_check_dics["depth"] == 1): 193 | return 194 | # TODO: save intrinsic matrix in path 195 | path = args 196 | P = depth_info.P 197 | open(path, 'w').close() # clear txt file 198 | fc = open(path, 'w') 199 | fc.writelines(str(P)) 200 | fc.close() 201 | self.__init_check_dics["depth"] = 1 202 | return 203 | 204 | def __colorInfoCallback(self, color_info, args): 205 | if (self.__init_check_dics["color"] == 1): 206 | return 207 | # TODO: save intrinsic matrix in path 208 | path = args 209 | P = color_info.P 210 | open(path, 'w').close() # clear txt file 211 | fc = open(path, 'w') 212 | fc.writelines(str(P)) 213 | fc.close() 214 | self.__init_check_dics["color"] = 1 215 | return 216 | 217 | def __writeExtrinstic(self, pos, ori, path, name): 218 | extric_list = [ori[0], ori[1], ori[2], ori[3], pos[0], pos[1], pos[2]] 219 | open(path, 'w').close() # clear txt file 220 | fc = open(path, 'w') 221 | fc.writelines(str(extric_list)) 222 | fc.close() 223 | self.__init_check_dics[name] = 1 224 | print(name + " save succeed.") 225 | return 226 | 227 | 228 | if __name__ == '__main__': 229 | # init global valuables for callback funcs 230 | # init ros node 231 | node_name = "data_collect_node" 232 | rospy.init_node(node_name, anonymous=False) 233 | 234 | parser = ROSArgparse(relative=node_name) 235 | parser.add_argument('main_freq', type=int, default=5, help="Main frequency of the path planner.") 236 | parser.add_argument('depth_topic', type=str, default='/rgbd_camera/depth/image', help="Topic for depth image.") 237 | parser.add_argument('color_topic', type=str, default='/rgbd_camera/color/image', help="Topic for color image.") 238 | parser.add_argument('odom_topic', type=str, default='/state_estimation', help='Topic for odometry data.') 239 | parser.add_argument('scan_topic', type=str, default='/velodyne_points', help='Topic for lidar point cloud data.') 240 | parser.add_argument('env_name', type=str, default='empty', help='Name of the environment, also used as the folder name for data.') 241 | parser.add_argument('depth_info_topic', type=str, default='/rgbd_camera/depth/camera_info', help='Topic for depth camera information.') 242 | parser.add_argument('color_info_topic', type=str, default='/rgbd_camera/color/camera_info', help='Topic for color camera information.') 243 | parser.add_argument('camera_frame_id', type=str, default='camera', help='Frame ID for the camera.') 244 | parser.add_argument('scan_frame_id', type=str, default='sensor', help='Frame ID for the lidar sensor.') 245 | parser.add_argument('base_frame_id', type=str, default='vehicle', help='Base frame ID.') 246 | parser.add_argument('odom_associate_id', type=str, default='vehicle', help='Frame ID associated with odometry data.') 247 | 248 | rospack = rospkg.RosPack() 249 | pack_path = rospack.get_path("iplanner_node") 250 | args = parser.parse_args() 251 | args.root_path = os.path.join(*[pack_path, "iplanner", "data", "CollectedData", args.env_name]) 252 | 253 | node = DataCollector(args) 254 | node.spin() -------------------------------------------------------------------------------- /src/iplanner_node.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import os 11 | import PIL 12 | import sys 13 | import torch 14 | import rospy 15 | import rospkg 16 | import tf 17 | import time 18 | from std_msgs.msg import Float32, Int16 19 | import numpy as np 20 | from sensor_msgs.msg import Image, Joy 21 | from nav_msgs.msg import Path 22 | from geometry_msgs.msg import PoseStamped, PointStamped 23 | import ros_numpy 24 | 25 | rospack = rospkg.RosPack() 26 | pack_path = rospack.get_path('iplanner_node') 27 | planner_path = os.path.join(pack_path,'iplanner') 28 | sys.path.append(pack_path) 29 | sys.path.append(planner_path) 30 | 31 | from iplanner.ip_algo import IPlannerAlgo 32 | from iplanner.rosutil import ROSArgparse 33 | 34 | class iPlannerNode: 35 | def __init__(self, args): 36 | super(iPlannerNode, self).__init__() 37 | self.config(args) 38 | 39 | # init planner algo class 40 | self.iplanner_algo = IPlannerAlgo(args=args) 41 | self.tf_listener = tf.TransformListener() 42 | 43 | rospy.sleep(2.5) # wait for tf listener to be ready 44 | 45 | self.image_time = rospy.get_rostime() 46 | self.is_goal_init = False 47 | self.ready_for_planning = False 48 | 49 | # planner status 50 | self.planner_status = Int16() 51 | self.planner_status.data = 0 52 | self.is_goal_processed = False 53 | self.is_smartjoy = False 54 | 55 | # fear reaction 56 | self.fear_buffter = 0 57 | self.is_fear_reaction = False 58 | # process time 59 | self.timer_data = Float32() 60 | 61 | rospy.Subscriber(self.image_topic, Image, self.imageCallback) 62 | rospy.Subscriber(self.goal_topic, PointStamped, self.goalCallback) 63 | rospy.Subscriber("/joy", Joy, self.joyCallback, queue_size=10) 64 | 65 | timer_topic = '/ip_timer' 66 | status_topic = '/ip_planner_status' 67 | 68 | # planning status topics 69 | self.timer_pub = rospy.Publisher(timer_topic, Float32, queue_size=10) 70 | self.status_pub = rospy.Publisher(status_topic, Int16, queue_size=10) 71 | 72 | self.path_pub = rospy.Publisher(self.path_topic, Path, queue_size=10) 73 | self.fear_path_pub = rospy.Publisher(self.path_topic + "_fear", Path, queue_size=10) 74 | 75 | rospy.loginfo("iPlanner Ready.") 76 | 77 | 78 | def config(self, args): 79 | self.main_freq = args.main_freq 80 | self.model_save = args.model_save 81 | self.image_topic = args.depth_topic 82 | self.goal_topic = args.goal_topic 83 | self.path_topic = args.path_topic 84 | self.frame_id = args.robot_id 85 | self.world_id = args.world_id 86 | self.uint_type = args.uint_type 87 | self.image_flip = args.image_flip 88 | self.conv_dist = args.conv_dist 89 | self.depth_max = args.depth_max 90 | # fear reaction 91 | self.is_fear_act = args.is_fear_act 92 | self.buffer_size = args.buffer_size 93 | self.ang_thred = args.angular_thred 94 | self.track_dist = args.track_dist 95 | self.joyGoal_scale = args.joyGoal_scale 96 | return 97 | 98 | def spin(self): 99 | r = rospy.Rate(self.main_freq) 100 | while not rospy.is_shutdown(): 101 | if self.ready_for_planning and self.is_goal_init: 102 | # main planning starts 103 | cur_image = self.img.copy() 104 | start = time.time() 105 | # Network Planning 106 | self.preds, self.waypoints, fear_output, _ = self.iplanner_algo.plan(cur_image, self.goal_rb) 107 | end = time.time() 108 | self.timer_data.data = (end - start) * 1000 109 | self.timer_pub.publish(self.timer_data) 110 | # check goal less than converage range 111 | if (np.sqrt(self.goal_rb[0][0]**2 + self.goal_rb[0][1]**2) < self.conv_dist) and self.is_goal_processed and (not self.is_smartjoy): 112 | self.ready_for_planning = False 113 | self.is_goal_init = False 114 | # planner status -> Success 115 | if self.planner_status.data == 0: 116 | self.planner_status.data = 1 117 | self.status_pub.publish(self.planner_status) 118 | 119 | rospy.loginfo("Goal Arrived") 120 | self.fear = torch.tensor([[0.0]], device=fear_output.device) 121 | if self.is_fear_act: 122 | self.fear = fear_output 123 | is_track_ahead = self.isForwardTraking(self.waypoints) 124 | self.fearPathDetection(self.fear, is_track_ahead) 125 | if self.is_fear_reaction: 126 | rospy.logwarn_throttle(2.0, "current path prediction is invaild.") 127 | # planner status -> Fails 128 | if self.planner_status.data == 0: 129 | self.planner_status.data = -1 130 | self.status_pub.publish(self.planner_status) 131 | self.pubPath(self.waypoints, self.is_goal_init) 132 | r.sleep() 133 | rospy.spin() 134 | 135 | def pubPath(self, waypoints, is_goal_init=True): 136 | path = Path() 137 | fear_path = Path() 138 | if is_goal_init: 139 | for p in waypoints.squeeze(0): 140 | pose = PoseStamped() 141 | pose.pose.position.x = p[0] 142 | pose.pose.position.y = p[1] 143 | pose.pose.position.z = p[2] 144 | path.poses.append(pose) 145 | # add header 146 | path.header.frame_id = fear_path.header.frame_id = self.frame_id 147 | path.header.stamp = fear_path.header.stamp = self.image_time 148 | # publish fear path 149 | if self.is_fear_reaction: 150 | fear_path.poses = path.poses.copy() 151 | path.poses = path.poses[:1] 152 | # publish path 153 | self.fear_path_pub.publish(fear_path) 154 | self.path_pub.publish(path) 155 | return 156 | 157 | def fearPathDetection(self, fear, is_forward): 158 | if fear > 0.5 and is_forward: 159 | if not self.is_fear_reaction: 160 | self.fear_buffter = self.fear_buffter + 1 161 | elif self.is_fear_reaction: 162 | self.fear_buffter = self.fear_buffter - 1 163 | if self.fear_buffter > self.buffer_size: 164 | self.is_fear_reaction = True 165 | elif self.fear_buffter <= 0: 166 | self.is_fear_reaction = False 167 | return None 168 | 169 | def isForwardTraking(self, waypoints): 170 | xhead = np.array([1.0, 0]) 171 | phead = None 172 | for p in waypoints.squeeze(0): 173 | if torch.norm(p[0:2]).item() > self.track_dist: 174 | phead = np.array([p[0].item(), p[1].item()]) 175 | phead /= np.linalg.norm(phead) 176 | break 177 | if phead is None or phead.dot(xhead) > 1.0 - self.ang_thred: 178 | return True 179 | return False 180 | 181 | def joyCallback(self, joy_msg): 182 | if joy_msg.buttons[4] > 0.9: 183 | rospy.loginfo("Switch to Smart Joystick mode ...") 184 | self.is_smartjoy = True 185 | # reset fear reaction 186 | self.fear_buffter = 0 187 | self.is_fear_reaction = False 188 | if self.is_smartjoy: 189 | if np.sqrt(joy_msg.axes[3]**2 + joy_msg.axes[4]**2) < 1e-3: 190 | # reset fear reaction 191 | self.fear_buffter = 0 192 | self.is_fear_reaction = False 193 | self.ready_for_planning = False 194 | self.is_goal_init = False 195 | else: 196 | joy_goal = PointStamped() 197 | joy_goal.header.frame_id = self.frame_id 198 | joy_goal.point.x = joy_msg.axes[4] * self.joyGoal_scale 199 | joy_goal.point.y = joy_msg.axes[3] * self.joyGoal_scale 200 | joy_goal.point.z = 0.0 201 | joy_goal.header.stamp = rospy.Time.now() 202 | self.goal_pose = joy_goal 203 | self.is_goal_init = True 204 | self.is_goal_processed = False 205 | return 206 | 207 | def goalCallback(self, msg): 208 | rospy.loginfo("Recevied a new goal") 209 | self.goal_pose = msg 210 | self.is_smartjoy = False 211 | self.is_goal_init = True 212 | self.is_goal_processed = False 213 | # reset fear reaction 214 | self.fear_buffter = 0 215 | self.is_fear_reaction = False 216 | # reste planner status 217 | self.planner_status.data = 0 218 | return 219 | 220 | def imageCallback(self, msg): 221 | # rospy.loginfo("Received image %s: %d"%(msg.header.frame_id, msg.header.seq)) 222 | self.image_time = msg.header.stamp 223 | frame = ros_numpy.numpify(msg) 224 | frame[~np.isfinite(frame)] = 0 225 | if self.uint_type: 226 | frame = frame / 1000.0 227 | frame[frame > self.depth_max] = 0.0 228 | # DEBUG - Visual Image 229 | # img = PIL.Image.fromarray((frame * 255 / np.max(frame[frame>0])).astype('uint8')) 230 | # img.show() 231 | if self.image_flip: 232 | frame = PIL.Image.fromarray(frame) 233 | self.img = np.array(frame.transpose(PIL.Image.ROTATE_180)) 234 | else: 235 | self.img = frame 236 | 237 | if self.is_goal_init: 238 | goal_robot_frame = self.goal_pose; 239 | if not self.goal_pose.header.frame_id == self.frame_id: 240 | try: 241 | goal_robot_frame.header.stamp = self.tf_listener.getLatestCommonTime(self.goal_pose.header.frame_id, 242 | self.frame_id) 243 | goal_robot_frame = self.tf_listener.transformPoint(self.frame_id, goal_robot_frame) 244 | except (tf.Exception, tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException): 245 | rospy.logerr("Fail to transfer the goal into base frame.") 246 | return 247 | goal_robot_frame = torch.tensor([goal_robot_frame.point.x, goal_robot_frame.point.y, goal_robot_frame.point.z], dtype=torch.float32)[None, ...] 248 | self.goal_rb = goal_robot_frame 249 | else: 250 | return 251 | self.ready_for_planning = True 252 | self.is_goal_processed = True 253 | return 254 | 255 | if __name__ == '__main__': 256 | 257 | node_name = "iplanner_node" 258 | rospy.init_node(node_name, anonymous=False) 259 | 260 | parser = ROSArgparse(relative=node_name) 261 | parser.add_argument('main_freq', type=int, default=5, help="Main frequency of the path planner.") 262 | parser.add_argument('model_save', type=str, default='/models/plannernet.pt', help="Path to the saved model.") 263 | parser.add_argument('crop_size', type=tuple, default=[360,640], help='Size to crop the image to.') 264 | parser.add_argument('uint_type', type=bool, default=False, help="Determines if the image is in uint type.") 265 | parser.add_argument('depth_topic', type=str, default='/rgbd_camera/depth/image', help='Topic for depth image.') 266 | parser.add_argument('goal_topic', type=str, default='/way_point', help='Topic for goal waypoints.') 267 | parser.add_argument('path_topic', type=str, default='/path', help='Topic for iPlanner path.') 268 | parser.add_argument('robot_id', type=str, default='base', help='TF frame ID for the robot.') 269 | parser.add_argument('world_id', type=str, default='odom', help='TF frame ID for the world.') 270 | parser.add_argument('depth_max', type=float, default=10.0, help='Maximum depth distance in the image.') 271 | parser.add_argument('image_flip', type=bool, default=True, help='Indicates if the image is flipped.') 272 | parser.add_argument('conv_dist', type=float, default=0.5, help='Convergence range to the goal.') 273 | parser.add_argument('is_fear_act', type=bool, default=True, help='Indicates if fear action is enabled.') 274 | parser.add_argument('buffer_size', type=int, default=10, help='Buffer size for fear reaction.') 275 | parser.add_argument('angular_thred', type=float, default=0.3, help='Angular threshold for turning.') 276 | parser.add_argument('track_dist', type=float, default=0.5, help='Look-ahead distance for path tracking.') 277 | parser.add_argument('joyGoal_scale', type=float, default=0.5, help='Scale for joystick goal distance.') 278 | parser.add_argument('sensor_offset_x', type=float, default=0.0, help='Sensor offset on the X-axis.') 279 | parser.add_argument('sensor_offset_y', type=float, default=0.0, help='Sensor offset on the Y-axis.') 280 | 281 | args = parser.parse_args() 282 | args.model_save = planner_path + args.model_save 283 | 284 | 285 | node = iPlannerNode(args) 286 | 287 | node.spin() 288 | -------------------------------------------------------------------------------- /src/iplanner_viz.py: -------------------------------------------------------------------------------- 1 | # ====================================================================== 2 | # Copyright (c) 2023 Fan Yang 3 | # Robotic Systems Lab, ETH Zurich 4 | # All rights reserved. 5 | 6 | # This source code is licensed under the MIT license found in the 7 | # LICENSE file in the root directory of this source tree. 8 | # ====================================================================== 9 | 10 | import os 11 | import PIL 12 | import sys 13 | import torch 14 | import rospy 15 | import rospkg 16 | import tf 17 | import numpy as np 18 | from sensor_msgs.msg import Image, Joy 19 | from std_msgs.msg import Int16 20 | from nav_msgs.msg import Path 21 | from geometry_msgs.msg import PoseStamped, PointStamped 22 | import ros_numpy 23 | 24 | rospack = rospkg.RosPack() 25 | pack_path = rospack.get_path('iplanner_node') 26 | planner_path = os.path.join(pack_path,'iplanner') 27 | sys.path.append(pack_path) 28 | sys.path.append(planner_path) 29 | 30 | from iplanner.ip_algo import IPlannerAlgo 31 | from iplanner.rosutil import ROSArgparse 32 | from iplanner import traj_viz 33 | 34 | class iPlannerNode: 35 | def __init__(self, args): 36 | super(iPlannerNode, self).__init__() 37 | self.config(args) 38 | 39 | self.iplanner_algo = IPlannerAlgo(args) 40 | self.traj_viz = traj_viz.TrajViz(os.path.join(*[planner_path, 'camera_intrinsic']), 41 | map_name=None, 42 | cameraTilt=args.camera_tilt, 43 | robot_name="robot") 44 | self.tf_listener = tf.TransformListener() 45 | 46 | rospy.sleep(2.5) # wait for tf listener to be ready 47 | 48 | self.image_time = rospy.get_rostime() 49 | self.is_goal_init = False 50 | self.ready_for_planning = False 51 | 52 | # planner status 53 | self.planner_status = Int16() 54 | self.planner_status.data = 0 55 | self.is_goal_processed = False 56 | self.is_smartjoy = False 57 | 58 | # fear reaction 59 | self.fear_buffter = 0 60 | self.is_fear_reaction = False 61 | 62 | rospy.Subscriber(self.depth_topic, Image, self.imageCallback) 63 | rospy.Subscriber(self.goal_topic, PointStamped, self.goalCallback) 64 | rospy.Subscriber("/joy", Joy, self.joyCallback, queue_size=10) 65 | 66 | status_topic = '/iplanner_status' 67 | 68 | # planning status topics 69 | self.status_pub = rospy.Publisher(status_topic, Int16, queue_size=10) 70 | # image visualizer 71 | self.img_pub = rospy.Publisher(self.image_topic, Image, queue_size=10) 72 | self.path_pub = rospy.Publisher(self.path_topic, Path, queue_size=10) 73 | self.fear_path_pub = rospy.Publisher(self.path_topic + "_fear", Path, queue_size=10) 74 | 75 | rospy.loginfo("iPlanner Ready.") 76 | 77 | 78 | def config(self, args): 79 | self.main_freq = args.main_freq 80 | self.depth_topic = args.depth_topic 81 | self.goal_topic = args.goal_topic 82 | self.path_topic = args.path_topic 83 | self.image_topic = args.image_topic 84 | self.camera_tilt = args.camera_tilt 85 | self.frame_id = args.robot_id 86 | self.world_id = args.world_id 87 | self.uint_type = args.uint_type 88 | self.image_flip = args.image_flip 89 | self.conv_dist = args.conv_dist 90 | self.depth_max = args.depth_max 91 | # fear reaction 92 | self.is_fear_act = args.is_fear_act 93 | self.buffer_size = args.buffer_size 94 | self.ang_thred = args.angular_thred 95 | self.track_dist = args.track_dist 96 | self.joyGoal_scale = args.joyGoal_scale 97 | return 98 | 99 | def spin(self): 100 | r = rospy.Rate(self.main_freq) 101 | while not rospy.is_shutdown(): 102 | if self.ready_for_planning and self.is_goal_init: 103 | # network planning 104 | cur_image = self.img.copy() 105 | self.preds, self.waypoints, fear_output, img_process = self.iplanner_algo.plan(cur_image, self.goal_rb) 106 | # check goal less than converage range 107 | goal_np = self.goal_rb[0, :].cpu().detach().numpy() 108 | if (np.sqrt(goal_np[0]**2 + goal_np[1]**2) < self.conv_dist) and self.is_goal_processed and (not self.is_smartjoy): 109 | self.ready_for_planning = False 110 | self.is_goal_init = False 111 | # planner status -> Success 112 | if self.planner_status.data == 0: 113 | self.planner_status.data = 1 114 | self.status_pub.publish(self.planner_status) 115 | 116 | rospy.loginfo("Goal Arrived") 117 | self.fear = torch.tensor([[0.0]], device=fear_output.device) 118 | if self.is_fear_act: 119 | self.fear = fear_output 120 | is_track_ahead = self.isForwardTraking(self.waypoints) 121 | self.fearPathDetection(self.fear, is_track_ahead) 122 | if self.is_fear_reaction: 123 | rospy.logwarn_throttle(2.0, "current path prediction is invaild.") 124 | # planner status -> Fails 125 | if self.planner_status.data == 0: 126 | self.planner_status.data = -1 127 | self.status_pub.publish(self.planner_status) 128 | self.pubPath(self.waypoints, self.is_goal_init) 129 | # visualize image 130 | self.pubRenderImage(self.preds, self.waypoints, self.odom, self.goal_rb, self.fear, img_process) 131 | r.sleep() 132 | rospy.spin() 133 | 134 | def pubPath(self, waypoints, is_goal_init=True): 135 | path = Path() 136 | fear_path = Path() 137 | if is_goal_init: 138 | for p in waypoints.squeeze(0): 139 | pose = PoseStamped() 140 | pose.pose.position.x = p[0] 141 | pose.pose.position.y = p[1] 142 | pose.pose.position.z = p[2] 143 | path.poses.append(pose) 144 | # add header 145 | path.header.frame_id = fear_path.header.frame_id = self.frame_id 146 | path.header.stamp = fear_path.header.stamp = self.image_time 147 | # publish fear path 148 | if self.is_fear_reaction: 149 | fear_path.poses = path.poses.copy() 150 | path.poses = path.poses[:1] 151 | # publish path 152 | self.fear_path_pub.publish(fear_path) 153 | self.path_pub.publish(path) 154 | return 155 | 156 | def fearPathDetection(self, fear, is_forward): 157 | if fear > 0.5 and is_forward: 158 | if not self.is_fear_reaction: 159 | self.fear_buffter = self.fear_buffter + 1 160 | elif self.is_fear_reaction: 161 | self.fear_buffter = self.fear_buffter - 1 162 | if self.fear_buffter > self.buffer_size: 163 | self.is_fear_reaction = True 164 | elif self.fear_buffter <= 0: 165 | self.is_fear_reaction = False 166 | return None 167 | 168 | def isForwardTraking(self, waypoints): 169 | xhead = np.array([1.0, 0]) 170 | phead = None 171 | for p in waypoints.squeeze(0): 172 | if torch.norm(p[0:2]).item() > self.track_dist: 173 | phead = np.array([p[0].item(), p[1].item()]) 174 | phead /= np.linalg.norm(phead) 175 | break 176 | if phead is None or phead.dot(xhead) > 1.0 - self.ang_thred: 177 | return True 178 | return False 179 | 180 | def joyCallback(self, joy_msg): 181 | if joy_msg.buttons[4] > 0.9: 182 | rospy.loginfo("Switch to Smart Joystick mode ...") 183 | self.is_smartjoy = True 184 | # reset fear reaction 185 | self.fear_buffter = 0 186 | self.is_fear_reaction = False 187 | if self.is_smartjoy: 188 | if np.sqrt(joy_msg.axes[3]**2 + joy_msg.axes[4]**2) < 1e-3: 189 | # reset fear reaction 190 | self.fear_buffter = 0 191 | self.is_fear_reaction = False 192 | self.ready_for_planning = False 193 | self.is_goal_init = False 194 | else: 195 | joy_goal = PointStamped() 196 | joy_goal.header.frame_id = self.frame_id 197 | joy_goal.point.x = joy_msg.axes[4] * self.joyGoal_scale 198 | joy_goal.point.y = joy_msg.axes[3] * self.joyGoal_scale 199 | joy_goal.point.z = 0.0 200 | joy_goal.header.stamp = rospy.Time.now() 201 | self.goal_pose = joy_goal 202 | self.is_goal_init = True 203 | self.is_goal_processed = False 204 | return 205 | 206 | def goalCallback(self, msg): 207 | rospy.loginfo("Recevied a new goal") 208 | self.goal_pose = msg 209 | self.is_smartjoy = False 210 | self.is_goal_init = True 211 | self.is_goal_processed = False 212 | # reset fear reaction 213 | self.fear_buffter = 0 214 | self.is_fear_reaction = False 215 | # reste planner status 216 | self.planner_status.data = 0 217 | return 218 | 219 | def pubRenderImage(self, preds, waypoints, odom, goal, fear, image): 220 | if torch.cuda.is_available(): 221 | odom = odom.cuda() 222 | goal = goal.cuda() 223 | image = self.traj_viz.VizImages(preds, waypoints, odom, goal, fear, image, is_shown=False)[0] 224 | ros_img = ros_numpy.msgify(Image, image, encoding='8UC4') 225 | self.img_pub.publish(ros_img) 226 | return None 227 | 228 | def imageCallback(self, msg): 229 | # rospy.loginfo("Received image %s: %d"%(msg.header.frame_id, msg.header.seq)) 230 | self.image_time = msg.header.stamp 231 | frame = ros_numpy.numpify(msg) 232 | frame[~np.isfinite(frame)] = 0 233 | if self.uint_type: 234 | frame = frame / 1000.0 235 | frame[frame > self.depth_max] = 0.0 236 | # DEBUG - Visual Image 237 | # img = PIL.Image.fromarray((frame * 255 / np.max(frame[frame>0])).astype('uint8')) 238 | # img.show() 239 | if self.image_flip: 240 | frame = PIL.Image.fromarray(frame) 241 | self.img = np.array(frame.transpose(PIL.Image.ROTATE_180)) 242 | else: 243 | self.img = frame 244 | # get odom from TF for camera image visualization 245 | try: 246 | (odom, ori) = self.tf_listener.lookupTransform(self.world_id, self.frame_id, rospy.Time(0)) 247 | odom.extend(ori) 248 | self.odom = torch.tensor(odom, dtype=torch.float32).unsqueeze(0) 249 | except (tf.Exception, tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException): 250 | rospy.logerr("Fail to get odomemrty from tf.") 251 | return 252 | 253 | if self.is_goal_init: 254 | goal_robot_frame = self.goal_pose 255 | if not self.goal_pose.header.frame_id == self.frame_id: 256 | try: 257 | goal_robot_frame.header.stamp = self.tf_listener.getLatestCommonTime(self.goal_pose.header.frame_id, 258 | self.frame_id) 259 | goal_robot_frame = self.tf_listener.transformPoint(self.frame_id, goal_robot_frame) 260 | except (tf.Exception, tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException): 261 | rospy.logerr("Fail to transfer the goal into robot frame.") 262 | return 263 | self.goal_rb = torch.tensor([goal_robot_frame.point.x, 264 | goal_robot_frame.point.y, 265 | goal_robot_frame.point.z], dtype=torch.float32)[None, ...] 266 | else: 267 | return 268 | self.ready_for_planning = True 269 | self.is_goal_processed = True 270 | return 271 | 272 | if __name__ == '__main__': 273 | 274 | node_name = "iplanner_node" 275 | rospy.init_node(node_name, anonymous=False) 276 | 277 | parser = ROSArgparse(relative=node_name) 278 | parser.add_argument('main_freq', type=int, default=5, help="Main frequency of the path planner.") 279 | parser.add_argument('model_save', type=str, default='/models/plannernet.pt', help="Path to the saved model.") 280 | parser.add_argument('crop_size', type=tuple, default=[360,640], help='Dimensions to crop the image to.') 281 | parser.add_argument('uint_type', type=bool, default=False, help="Flag to indicate if the image is in uint type.") 282 | parser.add_argument('depth_topic', type=str, default='/rgbd_camera/depth/image', help='ROS topic for depth image.') 283 | parser.add_argument('goal_topic', type=str, default='/way_point', help='ROS topic for goal waypoints.') 284 | parser.add_argument('path_topic', type=str, default='/iplanner_path', help='ROS topic for the iPlanner path.') 285 | parser.add_argument('image_topic', type=str, default='/path_image', help='ROS topic for iPlanner image view.') 286 | parser.add_argument('camera_tilt', type=float, default=0.0, help='Tilt angle of the camera.') 287 | parser.add_argument('robot_id', type=str, default='base', help='TF frame ID for the robot.') 288 | parser.add_argument('world_id', type=str, default='odom', help='TF frame ID for the world.') 289 | parser.add_argument('depth_max', type=float, default=10.0, help='Maximum depth distance in the image.') 290 | parser.add_argument('is_sim', type=bool, default=True, help='Flag to indicate if the system is in a simulation setting.') 291 | parser.add_argument('image_flip', type=bool, default=True, help='Flag to indicate if the image is flipped.') 292 | parser.add_argument('conv_dist', type=float, default=0.5, help='Convergence range to the goal.') 293 | parser.add_argument('is_fear_act', type=bool, default=True, help='Flag to indicate if fear action is enabled.') 294 | parser.add_argument('buffer_size', type=int, default=10, help='Buffer size for fear reaction.') 295 | parser.add_argument('angular_thred', type=float, default=1.0, help='Angular threshold for turning.') 296 | parser.add_argument('track_dist', type=float, default=0.5, help='Look-ahead distance for path tracking.') 297 | parser.add_argument('joyGoal_scale', type=float, default=0.5, help='Scale for joystick goal distance.') 298 | parser.add_argument('sensor_offset_x', type=float, default=0.0, help='Sensor offset on the X-axis.') 299 | parser.add_argument('sensor_offset_y', type=float, default=0.0, help='Sensor offset on the Y-axis.') 300 | 301 | args = parser.parse_args() 302 | args.model_save = planner_path + args.model_save 303 | 304 | node = iPlannerNode(args) 305 | 306 | node.spin() 307 | --------------------------------------------------------------------------------