├── .github └── workflows │ └── build.yml ├── .gitignore ├── CITATION.cff ├── CMakeLists.txt ├── README.md ├── config └── preset_decimation_4.0_depth_step_100.json ├── docker ├── Dockerfile.noetic ├── init_workspace.sh ├── requirements.txt ├── ros_entrypoint.sh └── run_docker.sh ├── docs ├── COLOR_THRESHOLD.md ├── DOCKER.md ├── LEARN_MORE.md └── RUN.md ├── images ├── rope.png ├── thresholder.png ├── trackdlo.png ├── trackdlo1.gif ├── trackdlo2.gif ├── trackdlo3.gif └── trackdlo4.gif ├── launch ├── evaluation.launch ├── realsense_node.launch ├── trackdlo.launch ├── trackdlo_eval.launch └── visualize_output.launch ├── package.xml ├── rviz └── tracking.rviz ├── trackdlo ├── include │ ├── evaluator.h │ ├── trackdlo.h │ └── utils.h └── src │ ├── evaluator.cpp │ ├── initialize.py │ ├── run_evaluation.cpp │ ├── trackdlo.cpp │ ├── trackdlo_node.cpp │ ├── utils.cpp │ └── utils.py └── utils ├── collect_pointcloud.py ├── color_picker.py ├── mask.py ├── simulate_occlusion.py ├── simulate_occlusion_eval.py ├── tracking_result_img_from_pointcloud_topic.py └── tracking_test.py /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: build 2 | 3 | on: 4 | push: 5 | branches: [ 'master' ] 6 | 7 | jobs: 8 | build-linux: 9 | strategy: 10 | max-parallel: 5 11 | matrix: 12 | os: [ubuntu-latest, ubuntu-22.04, ubuntu-20.04] 13 | 14 | runs-on: ${{matrix.os}} 15 | 16 | steps: 17 | - name: Checkout TrackDLO repository 18 | uses: actions/checkout@v3 19 | - name: Build, test, push base docker (Ubuntu 20.04) 20 | run: cd docker && docker build -t rmdlo-trackdlo:noetic -f Dockerfile.noetic .. 21 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | .vscode/ 3 | data/output/*.png 4 | data/output/*/*.json 5 | data/*.bag 6 | data/ 7 | test.cpp -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- 1 | cff-version: 1.2.0 2 | authors: 3 | - family-names: Xiang 4 | given-names: Jingyi 5 | - family-names: Dinkel 6 | given-names: Holly 7 | - family-names: Zhao 8 | given-names: Harry 9 | - family-names: Gao 10 | given-names: Naixiang 11 | - family-names: Coltin 12 | given-names: Brian 13 | - family-names: Smith 14 | given-names: Trey 15 | - family-names: Bretl 16 | given-names: Timothy 17 | title: "TrackDLO: Tracking Deformable Linear Objects Under Occlusion with Motion Coherence" 18 | version: 1.0.0 19 | url: "https://ieeexplore.ieee.org/document/10214157" 20 | date-released: '2023-08-09' 21 | preferred-citation: 22 | authors: 23 | - family-names: Xiang 24 | given-names: Jingyi 25 | - family-names: Dinkel 26 | given-names: Holly 27 | - family-names: Zhao 28 | given-names: Harry 29 | - family-names: Gao 30 | given-names: Naixiang 31 | - family-names: Coltin 32 | given-names: Brian 33 | - family-names: Smith 34 | given-names: Trey 35 | - family-names: Bretl 36 | given-names: Timothy 37 | title: "TrackDLO: Tracking Deformable Linear Objects Under Occlusion with Motion Coherence" 38 | url: "https://ieeexplore.ieee.org/document/10214157" 39 | year: 2023 40 | collection-title: "IEEE Robotics and Automation Letters" 41 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.0.2) 2 | project(trackdlo) 3 | 4 | set(CMAKE_CXX_STANDARD 17) 5 | set(CMAKE_CXX_STANDARD_REQUIRED ON) 6 | 7 | add_compile_options(-O3) # -std=c++11 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 | sensor_msgs 15 | std_msgs 16 | cv_bridge 17 | image_transport 18 | pcl_conversions 19 | pcl_ros 20 | ) 21 | 22 | add_definitions(${PCL_DEFINITIONS}) 23 | 24 | find_package(OpenCV REQUIRED) 25 | find_package(Eigen3 3.3 REQUIRED NO_MODULE) 26 | find_package(PCL 1.8 REQUIRED COMPONENTS common io filters visualization features kdtree) 27 | include_directories(include SYSTEM PUBLIC 28 | ${catkin_INCLUDE_DIRS} 29 | ${Eigen_INCLUDE_DIRS} 30 | ${PCL_INCLUDE_DIRS} 31 | ${OpenCV_INCLUDE_DIRS} 32 | ) 33 | 34 | catkin_package( 35 | # INCLUDE_DIRS include 36 | # LIBRARIES tracking_ros 37 | # CATKIN_DEPENDS abb_egm_hardware_interface abb_egm_state_controller abb_rws_service_provider abb_rws_state_publisher controller_manager joint_state_controller velocity_controllers 38 | # DEPENDS system_lib 39 | ) 40 | 41 | add_executable( 42 | trackdlo trackdlo/src/trackdlo_node.cpp trackdlo/src/trackdlo.cpp trackdlo/src/utils.cpp 43 | ) 44 | target_link_libraries(trackdlo 45 | ${catkin_LIBRARIES} 46 | ${PCL_LIBRARIES} 47 | ${OpenCV_LIBS} 48 | Eigen3::Eigen 49 | ) 50 | # target_compile_options(trackdlo PRIVATE -O3 -Wall -Wextra -Wconversion -Wshadow -g) 51 | 52 | add_executable( 53 | evaluation trackdlo/src/run_evaluation.cpp trackdlo/src/trackdlo.cpp trackdlo/src/utils.cpp trackdlo/src/evaluator.cpp 54 | ) 55 | target_link_libraries(evaluation 56 | ${catkin_LIBRARIES} 57 | ${PCL_LIBRARIES} 58 | ${OpenCV_LIBS} 59 | Eigen3::Eigen 60 | ) 61 | # target_compile_options(evaluation PRIVATE -O3 -Wall -Wextra -Wconversion -Wshadow -g) 62 | 63 | # add_executable( 64 | # cv_test trackdlo/src/test.cpp 65 | # ) 66 | # target_link_libraries(cv_test 67 | # ${catkin_LIBRARIES} 68 | # ${PCL_LIBRARIES} 69 | # ${OpenCV_LIBS} 70 | # Eigen3::Eigen 71 | # ) 72 | # target_compile_options(cv_test PRIVATE -O3 -Wall -Wextra -Wconversion -Wshadow -g) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |

2 | 3 |

4 | 5 |

6 | 7 | 8 | 9 |

10 | 11 | ## TrackDLO: Tracking Deformable Linear Objects Under Occlusion with Motion Coherence 12 | **IEEE Robotics and Automation Letters (IEEE Xplore: https://ieeexplore.ieee.org/document/10214157)** 13 | 14 | This repository contains the TrackDLO Robot Operating System (ROS) package. The TrackDLO ROS package is an implementation of our paper, *TrackDLO: Tracking Deformable Linear Objects Under Occlusion with Motion Coherence*, by Jingyi Xiang, Holly Dinkel, Harry Zhao, Naixiang Gao, Brian Coltin, Trey Smith, and Timothy Bretl. The TrackDLO algorithm is implemented in C++. 15 | 16 |

17 | 18 |

19 |

20 | 21 |

22 | 23 | ## Get Started 24 | 25 | The [requirements and run instructions](https://github.com/RMDLO/trackdlo/blob/master/docs/RUN.md) page provides information on package installation and usage. 26 | 27 | ## Learn More 28 | 29 | The [supplementary video](https://www.youtube.com/watch?v=MxqNJsen5eg&t) showcases experiments and failure cases. 30 | 31 | [![supplementary video](https://img.youtube.com/vi/MxqNJsen5eg/0.jpg)](https://www.youtube.com/watch?v=MxqNJsen5eg) 32 | 33 | The [supplemental documentation](https://github.com/RMDLO/trackdlo/blob/master/docs/LEARN_MORE.md) contains details about the initialization process and parameter tuning. 34 | 35 | ## Bibtex 36 | 37 | ```bash 38 | @ARTICLE{ 39 | xiang2023trackdlo, 40 | author={Xiang, Jingyi and Dinkel, Holly and Zhao, Harry and Gao, Naixiang and Coltin, Brian and Smith, Trey and Bretl, Timothy}, 41 | journal={IEEE Robotics and Automation Letters}, 42 | title={TrackDLO: Tracking Deformable Linear Objects Under Occlusion With Motion Coherence}, 43 | year={2023}, 44 | volume={8}, 45 | number={10}, 46 | pages={6179-6186}, 47 | doi={10.1109/LRA.2023.3303710} 48 | } 49 | ``` 50 | -------------------------------------------------------------------------------- /config/preset_decimation_4.0_depth_step_100.json: -------------------------------------------------------------------------------- 1 | { 2 | "device": { 3 | "fw version": "05.13.00.50", 4 | "name": "Intel RealSense D435", 5 | "product line": "D400" 6 | }, 7 | "parameters": { 8 | "aux-param-autoexposure-setpoint": "1536", 9 | "aux-param-colorcorrection1": "0.298828", 10 | "aux-param-colorcorrection10": "-0", 11 | "aux-param-colorcorrection11": "-0", 12 | "aux-param-colorcorrection12": "-0", 13 | "aux-param-colorcorrection2": "0.293945", 14 | "aux-param-colorcorrection3": "0.293945", 15 | "aux-param-colorcorrection4": "0.114258", 16 | "aux-param-colorcorrection5": "-0", 17 | "aux-param-colorcorrection6": "-0", 18 | "aux-param-colorcorrection7": "-0", 19 | "aux-param-colorcorrection8": "-0", 20 | "aux-param-colorcorrection9": "-0", 21 | "aux-param-depthclampmax": "65536", 22 | "aux-param-depthclampmin": "0", 23 | "aux-param-disparityshift": "0", 24 | "controls-autoexposure-auto": "True", 25 | "controls-autoexposure-manual": "8500", 26 | "controls-color-autoexposure-auto": "True", 27 | "controls-color-autoexposure-manual": "166", 28 | "controls-color-backlight-compensation": "0", 29 | "controls-color-brightness": "0", 30 | "controls-color-contrast": "50", 31 | "controls-color-gain": "64", 32 | "controls-color-gamma": "300", 33 | "controls-color-hue": "0", 34 | "controls-color-power-line-frequency": "3", 35 | "controls-color-saturation": "64", 36 | "controls-color-sharpness": "50", 37 | "controls-color-white-balance-auto": "True", 38 | "controls-color-white-balance-manual": "4600", 39 | "controls-depth-gain": "16", 40 | "controls-laserpower": "150", 41 | "controls-laserstate": "on", 42 | "ignoreSAD": "0", 43 | "param-amplitude-factor": "0", 44 | "param-autoexposure-setpoint": "1536", 45 | "param-censusenablereg-udiameter": "9", 46 | "param-censusenablereg-vdiameter": "9", 47 | "param-censususize": "9", 48 | "param-censusvsize": "9", 49 | "param-depthclampmax": "65536", 50 | "param-depthclampmin": "0", 51 | "param-depthunits": "100", 52 | "param-disableraucolor": "0", 53 | "param-disablesadcolor": "0", 54 | "param-disablesadnormalize": "0", 55 | "param-disablesloleftcolor": "0", 56 | "param-disableslorightcolor": "0", 57 | "param-disparitymode": "0", 58 | "param-disparityshift": "0", 59 | "param-lambdaad": "618", 60 | "param-lambdacensus": "15", 61 | "param-leftrightthreshold": "18", 62 | "param-maxscorethreshb": "1443", 63 | "param-medianthreshold": "789", 64 | "param-minscorethresha": "96", 65 | "param-neighborthresh": "12", 66 | "param-raumine": "2", 67 | "param-rauminn": "1", 68 | "param-rauminnssum": "6", 69 | "param-raumins": "3", 70 | "param-rauminw": "3", 71 | "param-rauminwesum": "7", 72 | "param-regioncolorthresholdb": "0.109589", 73 | "param-regioncolorthresholdg": "0.572407", 74 | "param-regioncolorthresholdr": "0.0176125", 75 | "param-regionshrinku": "4", 76 | "param-regionshrinkv": "0", 77 | "param-robbinsmonrodecrement": "6", 78 | "param-robbinsmonroincrement": "21", 79 | "param-rsmdiffthreshold": "1.21875", 80 | "param-rsmrauslodiffthreshold": "0.28125", 81 | "param-rsmremovethreshold": "0.488095", 82 | "param-scanlineedgetaub": "8", 83 | "param-scanlineedgetaug": "200", 84 | "param-scanlineedgetaur": "279", 85 | "param-scanlinep1": "55", 86 | "param-scanlinep1onediscon": "326", 87 | "param-scanlinep1twodiscon": "134", 88 | "param-scanlinep2": "235", 89 | "param-scanlinep2onediscon": "506", 90 | "param-scanlinep2twodiscon": "206", 91 | "param-secondpeakdelta": "222", 92 | "param-texturecountthresh": "0", 93 | "param-texturedifferencethresh": "2466", 94 | "param-usersm": "1", 95 | "param-zunits": "100" 96 | }, 97 | "schema version": 1, 98 | "viewer": { 99 | "stream-depth-format": "Z16", 100 | "stream-fps": "30", 101 | "stream-height": "720", 102 | "stream-width": "1280" 103 | } 104 | } -------------------------------------------------------------------------------- /docker/Dockerfile.noetic: -------------------------------------------------------------------------------- 1 | FROM ros:noetic-robot 2 | 3 | # Copyright (c) 2023, UNIVERSITY OF ILLINOIS URBANA-CHAMPAIGN. All rights reserved. 4 | 5 | # To build: 6 | # docker build -t rmdlo-trackdlo:noetic -f Dockerfile.noetic .. 7 | 8 | # Set up directories and copy installation targets 9 | ENV HOME=/root 10 | ENV CATKIN_WS=${HOME}/tracking_ws 11 | ENV SRC=${CATKIN_WS}/src 12 | ENV DEVEL=${CATKIN_WS}/devel 13 | COPY . ${SRC}/trackdlo 14 | COPY ./docker/requirements.txt /tmp/requirements.txt 15 | COPY ./docker/init_workspace.sh /tmp/init_workspace.sh 16 | ENV DEBIAN_FRONTEND=noninteractive 17 | 18 | RUN apt-get update && apt-get -y --no-install-recommends install \ 19 | # Install system and development components 20 | apt-utils \ 21 | software-properties-common \ 22 | build-essential \ 23 | cmake \ 24 | git \ 25 | python3-pip \ 26 | python3-catkin-tools \ 27 | python3-rosbag \ 28 | libeigen3-dev \ 29 | libpcl-dev \ 30 | libopencv-dev \ 31 | python3-opencv \ 32 | # Install required ROS components 33 | ros-noetic-desktop-full \ 34 | ros-noetic-catkin \ 35 | ros-noetic-cv-bridge \ 36 | ros-noetic-pcl-conversions \ 37 | ros-noetic-pcl-ros \ 38 | ros-noetic-geometry-msgs \ 39 | ros-noetic-message-filters \ 40 | ros-noetic-rospy \ 41 | ros-noetic-sensor-msgs \ 42 | ros-noetic-std-msgs \ 43 | ros-noetic-tf \ 44 | ros-noetic-vision-msgs \ 45 | ros-noetic-visualization-msgs \ 46 | ros-noetic-rviz \ 47 | # Optionall install RealSense2 package 48 | ros-noetic-realsense2-camera \ 49 | && apt-get -y autoremove \ 50 | && apt-get clean 51 | 52 | # Install required Python components 53 | RUN python3 -m pip install -r /tmp/requirements.txt 54 | 55 | # Copy and set up the ros_entrypoint.sh 56 | COPY ./docker/ros_entrypoint.sh /ros_entrypoint.sh 57 | RUN chmod +x /ros_entrypoint.sh 58 | 59 | # Make sure the catkin workspace is initialized 60 | RUN . /opt/ros/noetic/setup.sh && \ 61 | /tmp/init_workspace.sh 62 | 63 | # Set the entrypoint to ros_entrypoint.sh 64 | ENTRYPOINT ["/ros_entrypoint.sh"] 65 | 66 | # Default command 67 | CMD ["bash"] 68 | 69 | # Set Display variables for GUI applications 70 | ENV DISPLAY=:0 71 | ENV TERM=xterm 72 | # Some QT-Apps do not show controls without this 73 | ENV QT_X11_NO_MITSHM=1 74 | -------------------------------------------------------------------------------- /docker/init_workspace.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Copyright (c) 2023, UNIVERSITY OF ILLINOIS URBANA-CHAMPAIGN. All rights reserved. 4 | 5 | # Stop in case of any error. 6 | set -e 7 | 8 | # Create catkin workspace. 9 | mkdir -p ${CATKIN_WS}/src 10 | cd ${CATKIN_WS}/src 11 | catkin init 12 | cd .. 13 | catkin build 14 | source devel/setup.bash -------------------------------------------------------------------------------- /docker/requirements.txt: -------------------------------------------------------------------------------- 1 | numpy>=1.22 2 | scipy>=1.10.0 3 | opencv_python==4.10.0.84 4 | Pillow>=10.0.1 5 | rosnumpy 6 | scikit-image==0.20.0 7 | pyrealsense2==2.54.1.5217 8 | -------------------------------------------------------------------------------- /docker/ros_entrypoint.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Source ROS distribution and Catkin Workspace 4 | source /opt/ros/noetic/setup.bash 5 | source ${DEVEL}/setup.bash 6 | 7 | exec "$@" 8 | -------------------------------------------------------------------------------- /docker/run_docker.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Copyright (c) 2023, UNIVERSITY OF ILLINOIS URBANA-CHAMPAIGN. All rights reserved. 4 | 5 | CONTAINER_NAME=$1 6 | if [[ -z "${CONTAINER_NAME}" ]]; then 7 | CONTAINER_NAME=trackdlo 8 | fi 9 | 10 | # Specify a mapping between a host directory and a directory in the 11 | # docker container. Change this to access a different directory. 12 | HOST_DIR=$2 13 | if [[ -z "${HOST_DIR}" ]]; then 14 | HOST_DIR=`realpath ${PWD}/..` 15 | fi 16 | 17 | CONTAINER_DIR=$3 18 | if [[ -z "${CONTAINER_DIR}" ]]; then 19 | CONTAINER_DIR=/root/tracking_ws/src/trackdlo 20 | fi 21 | 22 | echo "Container name : ${CONTAINER_NAME}" 23 | echo "Host directory : ${HOST_DIR}" 24 | echo "Container directory: ${CONTAINER_DIR}" 25 | TRACK_ID=`docker ps -aqf "name=^/${CONTAINER_NAME}$"` 26 | if [ -z "${TRACK_ID}" ]; then 27 | echo "Creating new trackdlo docker container." 28 | xhost +local:root 29 | docker run -it --privileged --network=host -v ${HOST_DIR}:${CONTAINER_DIR}:rw -v /tmp/.X11-unix:/tmp/.X11-unix:rw --env="DISPLAY" --name=${CONTAINER_NAME} rmdlo-trackdlo:noetic bash 30 | else 31 | echo "Found trackdlo docker container: ${TRACK_ID}." 32 | # Check if the container is already running and start if necessary. 33 | if [ -z `docker ps -qf "name=^/${CONTAINER_NAME}$"` ]; then 34 | xhost +local:${TRACK_ID} 35 | echo "Starting and attaching to ${CONTAINER_NAME} container..." 36 | docker start ${TRACK_ID} 37 | docker attach ${TRACK_ID} 38 | else 39 | echo "Found running ${CONTAINER_NAME} container, attaching bash..." 40 | docker exec -it ${TRACK_ID} bash 41 | fi 42 | fi -------------------------------------------------------------------------------- /docs/COLOR_THRESHOLD.md: -------------------------------------------------------------------------------- 1 | ## Color-Segmenting Objects to Use with TrackDLO 2 | 3 | TrackDLO uses color thresholding to obtain the DLO segmentation mask. The [`color_picker.py`](https://github.com/RMDLO/trackdlo/blob/master/utils/color_picker.py) script is provided to simplify tuning the color threshold values using a simple interface with a track bar. This document provides a tutorial of intended use for this script. 4 | 5 | First, save an image to use for color thresholding. We use an Intel RealSense d435 camera. The RealSense ROS package publishes raw RGB images to the `/camera/color/image_raw` topic. Raw `.png` images can be extracted from this topic using the [image_view](http://wiki.ros.org/image_view#image_view.2Fdiamondback.image_saver) ROS package: 6 | 7 | Terminal 1: 8 | 9 | ``` 10 | roslaunch trackdlo realsense_node.launch 11 | ``` 12 | 13 | Terminal 2: 14 | 15 | ``` 16 | rosrun image_view image_saver image:=/camera/color/image_raw 17 | ``` 18 | 19 | This is an example of a rope image ([`rope.png`](https://github.com/RMDLO/trackdlo/blob/master/images/rope.png)) from our camera: 20 | 21 |

22 | 23 |

24 | 25 | Next, run the user interface (replace `images/rope.png` with the relative or full path to the image of the DLO). 26 | 27 | ``` 28 | cd /path/to/trackdlo/root/folder && python utils/color_picker.py images/rope.png 29 | ``` 30 | 31 | The user interface should appear as shown below. Use the track bars to tune the HSV lower and upper limits to segment the background. 32 | 33 |

34 | 35 |

36 | 37 | Once the HSV threshold values are determined with this tool, modify the default values according to the instructions below: 38 | * Monochrome DLO: set the lower and upper HSV value ranges in the [`trackdlo.launch`](https://github.com/RMDLO/trackdlo/blob/master/launch/trackdlo.launch) file. 39 | * Multi-Colored DLO: set `multi_color_dlo` to `true` in [`trackdlo.launch`](https://github.com/RMDLO/trackdlo/blob/master/launch/trackdlo.launch), then modify the `color_thresholding` function in [`initialize.py`](https://github.com/RMDLO/trackdlo/blob/master/trackdlo/src/initialize.py) and [`trackdlo_node.cpp`](https://github.com/RMDLO/trackdlo/blob/master/trackdlo/src/trackdlo_node.cpp) to customize DLO segmentation. -------------------------------------------------------------------------------- /docs/DOCKER.md: -------------------------------------------------------------------------------- 1 | ## TrackDLO with Docker 2 | 3 | TrackDLO can be tested inside of a [Docker](https://www.docker.com/) container to isolate all software and dependency changes from the host system. This document describes how to create and run a Docker image that contains a complete ROS and system environment for TrackDLO. 4 | 5 | The current configuration was tested on an x86 host computer running Ubuntu 20.04 with Docker 24.0.1. 6 | 7 | ### Steps 8 | 9 | 1. **Download TrackDLO** 10 | ```bash 11 | git clone https://github.com/RMDLO/trackdlo.git trackdlo 12 | ``` 13 | 14 | 2. **Build the Docker Image** 15 | ```bash 16 | cd trackdlo/docker 17 | docker build -t rmdlo-trackdlo:noetic -f Dockerfile.noetic .. 18 | ``` 19 | 20 | This will take several minutes and require connection to the internet. This command will install all dependencies and build the TrackDLO ROS workspace within the image. 21 | 22 | 3. **Connect a Camera** 23 | Docker will not recognize a USB device that is plugged in after the container is started. 24 | 25 | 4. **Run the Container** 26 | ``` 27 | ./run_docker.sh [name] [host dir] [container dir] 28 | ``` 29 | Optional Parameters: 30 | - `name` specifies the name of the image. By default, it is `trackdlo`. Multiple containers can be created from the same image by changing this parameter. 31 | - `host dir` and `container dir` map a directory on the host machine to a location inside the container. This enables sharing code and data between the two systems. By default, the `run_docker.sh` bash script maps the directory containing trackdlo to `/root/tracking_ws/src/trackdlo` in the container. 32 | 33 | Only the first call of this script with a given name will create a container. Subsequent executions will attach to the running container to enable running multiple terminal sessions in a single container. In subsequent calls, please check that ROS and the built catkin workspace are properly sourced. 34 | 35 | ``` 36 | source /ros_entrypoint.sh 37 | ``` 38 | 39 | *Note:* Since the Docker container binds directly to the host's network, it will see `roscore` even if running outside the docker container. 40 | 41 | For more information about using ROS with docker, see the ![ROS tutorial](http://wiki.ros.org/docker/Tutorials/Docker). -------------------------------------------------------------------------------- /docs/LEARN_MORE.md: -------------------------------------------------------------------------------- 1 | ## TrackDLO 2 | 3 | The TrackDLO algorithm estimates the shape of a Deformable Linear Object (DLO) under occlusion from a sequence of RGB-D images for use in manipulation tasks. TrackDLO runs in real-time and requires no prior state information as input. The algorithm improves on previous approaches by addressing three common scenarios which cause tracking failure: tip occlusion, mid-section occlusion, and self-intersection. This is achieved through the application of Motion Coherence Theory to impute the spatial velocity of occluded nodes; the use of a geodesic distance function to better track self-intersecting DLOs; and the introduction of a non-Gaussian kernel which only penalizes lower-order spatial displacement derivatives to better reflect DLO physics. The source code and benchmarking dataset are publicly released in this repository. 4 | 5 | ## Initialization 6 | 7 | We adapt the algorithm introduced in [Deformable One-Dimensional Object Detection for Routing and Manipulation](https://ieeexplore.ieee.org/abstract/document/9697357) to allow complicated initial DLO configurations such as self-crossing and minor occlusion at initialization. 8 | 9 | **Initialization under minor occlusion:** 10 |

11 | 12 |

13 | 14 | **Initialization under complicated DLO topology:** 15 |

16 | 17 |

18 | 19 | ## Parameter Tuning 20 | * $\beta$ and $\lambda$: MCT (Motion Coherence Theory) weights. The larger they are, the more rigid the object becomes. Desirable $\beta$ and $\lambda$ values should be as large as possible, but not too large for the object deformation to be reflected in the tracking results. 21 | 22 | * $\alpha$: The alignment strength between registered visible node positions and estimated visible node positions. Small $\alpha$ could lead to failure in length preservation while large $\alpha$ could lead to jittery movement between frames. 23 | 24 | * $\mu$: Ranging from 0 to 1, large $\mu$ indicates the segmented DLO point cloud contains a large amount of outliers. 25 | 26 | * $k_{\mathrm{vis}}$: The strength of visibility information's effect on membership probability computation. When properly set, $k_{\mathrm{vis}}$ helps improve the performance of tracking under occlusion. However, the larger $k_{\mathrm{vis}}$ becomes, the longer it takes for tracking results to "catch up" with the DLO shape changes. -------------------------------------------------------------------------------- /docs/RUN.md: -------------------------------------------------------------------------------- 1 | # Installation and Run Instructions 2 | 3 | This guide contains information about required dependencies and how to run the TrackDLO ROS package. 4 | 5 | ## Minimum Requirements 6 | 7 | Installation and execution of TrackDLO was verified with the below dependencies on an Ubuntu 20.04 system with ROS Noetic. 8 | 9 | * [ROS Noetic](http://wiki.ros.org/noetic/Installation) 10 | * [Eigen3](https://eigen.tuxfamily.org/index.php?title=Main_Page) (Our version: 3.3.7) 11 | * [Point Cloud Library](https://pointclouds.org/) (Our version: 1.10.0) 12 | * [OpenCV](https://opencv.org/releases/) 13 | * [NumPy](https://numpy.org/install/) 14 | * [SciPy](https://scipy.org/install/) 15 | * [scikit-image](https://scikit-image.org/) 16 | * [Pillow](https://pillow.readthedocs.io/en/stable/installation.html) 17 | * [ROS Numpy](https://pypi.org/project/rosnumpy/) 18 | 19 | We also provide Docker files for compatibility with other system configurations. Refer to the [DOCKER.md](https://github.com/RMDLO/trackdlo/blob/master/docs/DOCKER.md) for more information on using docker, and see the docker [requirements.txt](https://github.com/RMDLO/trackdlo/blob/master/docker/requirements.txt) file for a list of the tested versions of TrackDLO package dependencies. 20 | 21 | ## Other Requirements 22 | 23 | We used an Intel RealSense d435 camera in all of the experiments performed in our paper. We used the [librealsense](https://github.com/IntelRealSense/librealsense) and [realsense-ros](https://github.com/IntelRealSense/realsense-ros/tree/ros1-legacy) packages for testing with the RealSense D435 camera and for obtaining the [camera configuration file](https://github.com/RMDLO/trackdlo/blob/master/config/preset_decimation_4.0_depth_step_100.json). 24 | 25 | ## Installation 26 | 27 | The repository is organized into the following directories: 28 | 29 | ``` 30 | ├── trackdlo/ # contains TrackDLO class and corresponding ROS node 31 | ├── launch/ # contains ROS launch files for the camera and RViz 32 | ├── rviz/ # contains Rviz configurations for visualization 33 | ├── config/ # contains camera configurations 34 | └── utils/ # contains scripts used for testing and evaluation 35 | ``` 36 | 37 | First, [create a ROS workspace](http://wiki.ros.org/catkin/Tutorials/create_a_workspace). Next, `cd YOUR_TRACKING_ROS_WORKSPACE/src`. Clone the TrackDLO repository into this workspace and build the package: 38 | 39 | ```bash 40 | git clone https://github.com/RMDLO/trackdlo.git 41 | catkin build trackdlo 42 | source ../devel/setup.bash 43 | ``` 44 | 45 | All configurable parameters for the TrackDLO algorithm are in [`launch/trackdlo.launch`](https://github.com/RMDLO/trackdlo/blob/master/launch/trackdlo.launch). Rebuilding the package is not required for any parameter modifications to take effect. However, `catkin build` is required after modifying any C++ files. Remember that `source /devel/setup.bash` is required in every terminal running TrackDLO ROS nodes. 46 | 47 | ## Usage 48 | 49 | To run TrackDLO, the three ROS topic names below should be modified in `launch/trackdlo.launch` to match the ROS topic names published by the user's RGB-D camera: 50 | * `camera_info_topic`: a CameraInfo topic that contains the camera's projection matrix 51 | * `rgb_topic`: a RGB image stream topic 52 | * `depth_topic`: a depth image stream topic 53 | 54 | **Note: TrackDLO assumes the RGB image and its corresponding depth image are ALIGNED and SYNCHRONIZED. This means depth_image(i, j) should contain the depth value of pixel rgb_image(i, j) and the two images should be published with the same ROS timestamp.** 55 | 56 | TrackDLO uses color thresholding in Hue, Saturation, and Value (HSV) color space to obtain the DLO segmentation mask. A tutorial on how to obtain the HSV limits is provided in [`COLOR_THRESHOLD.md`](https://github.com/RMDLO/trackdlo/blob/master/docs/COLOR_THRESHOLD.md) 57 | 58 | Other useful parameters: 59 | * `num_of_nodes`: the number of nodes initialized for the DLO 60 | * `result_frame_id`: the tf frame the tracking results (point cloud and marker array) will be published to 61 | * `visualize_initialization_process`: if set to `true`, OpenCV windows will appear to visualize the results of each step in initialization. This is helpful for debugging in the event of initialization failures. 62 | 63 | Once all parameters in `launch/trackdlo.launch` are set to proper values, run TrackDLO with the following steps: 64 | 1. Launch the RGB-D camera node 65 | 2. Launch RViz to visualize all published topics: 66 | ```bash 67 | roslaunch trackdlo visualize_output.launch 68 | ``` 69 | 3. Launch the TrackDLO node to publish tracking results: 70 | ```bash 71 | roslaunch trackdlo trackdlo.launch 72 | ``` 73 | 74 | The TrackDLO node outputs the following: 75 | * `/trackdlo/results_marker`: the tracking result with nodes visualized with spheres and edges visualized with cylinders in MarkerArray format, 76 | * `/trackdlo/results_pc`: the tracking results in PointCloud2 format 77 | * `/trackdlo/results_img`: the tracking results projected onto the received input RGB image in RGB Image format 78 | 79 | ## Run TrackDLO with a RealSense D435 camera: 80 | This package was tested using an Intel RealSense D435 camera. The exact camera configurations used are provided in `/config/preset_decimation_4.0_depth_step_100.json` and can be loaded into the camera using the launch files from `realsense-ros`. Run the following commands to start the RealSense camera and the tracking node: 81 | 1. Launch an RViz window visualizing the color image, mask, and tracking result (in both the image and the 3D pointcloud) with 82 | ```bash 83 | roslaunch trackdlo realsense_node.launch 84 | ``` 85 | 2. Launch the TrackDLO tracking node and publish messages containing the estimated node positions defining the object shape with 86 | ```bash 87 | roslaunch trackdlo trackdlo.launch 88 | ``` 89 | 90 | ## Run TrackDLO with Recorded ROS Bag Data: 91 | 1. Download one of the provided rosbag experiment `.bag` files [here](https://drive.google.com/drive/folders/1YjX-xfbNfm_G9FYbdw1voYxmd9VA-Aho?usp=sharing). Note: the file sizes are large--each will require 5-10GB of storage! 92 | 2. Open a new terminal and run 93 | ```bash 94 | roslaunch trackdlo visualize_output.launch bag:=True 95 | ``` 96 | 3. In another terminal, run the below command to start the tracking algorithm with parameters for the `stationary.bag`, `perpendicular_motion.bag`, and `parallel_motion.bag` files used for quantitative evaluation in the TrackDLO paper. 97 | ```bash 98 | roslaunch trackdlo trackdlo_eval.launch 99 | ``` 100 | If testing any of the other provided `.bag` files, run the below command: 101 | ```bash 102 | roslaunch trackdlo trackdlo.launch 103 | ``` 104 | 4. Open a third ternimal and run the below command to replay the `.bag` file and publish its topics: 105 | ```bash 106 | rosbag play --clock /path/to/filename.bag 107 | ``` 108 | Occlusion can also be injected using our provided `simulate_occlusion_eval.py` script. Run the below command and draw bounding rectangles for the occlusion mask in the graphical user interface that appears: 109 | ```bash 110 | rosrun trackdlo simulate_occlusion_eval.py 111 | ``` 112 | 113 | ## Data: 114 | 115 | The ROS bag files used in our paper and the supplementary video can be found [here](https://drive.google.com/file/d/1C7uM515fHXnbsEyx5X38xZUXzBI99mxg/view?usp=drive_link). The `experiment` folder is organized into the following directories: 116 | 117 | ``` 118 | ├── rope/ # contains the three experiments performed with a rope in the supplementary video 119 | ├── rubber_tubing/ # contains the three experiments performed with a rope in the supplementary video 120 | ├── failure_cases/ # contains the three failure cases shown in the supplementary video 121 | └── quantitative_eval/ # contains the bag files used for quantitative evaluation in our paper 122 | ``` 123 | 124 | ### Notes on Running the Bag Files 125 | 126 | * The rope and the rubber tubing require different hsv thresholding values. Both of these objects use the `hsv_threshold_upper_limit` default value = `130 255 255` however the rope uses the `hsv_threshold_lower_limit` default value = `90 90 30` and the rubber tubing uses the `hsv_threshold_upper_limit` default value = `100 200 60`. 127 | * For `.bag` files in `rope/`, `rubber_tubing/`, and `failure_cases/`, the `camera_info_topic` is published under `/camera/aligned_depth_to_color/camera_info`. For `.bag` files in `quantitative_eval/`, the `camera_info_topic` is published under `/camera/color/camera_info`. 128 | -------------------------------------------------------------------------------- /images/rope.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RMDLO/trackdlo/19f92952c6c864f948ee5c05d5c5e3627e190a8f/images/rope.png -------------------------------------------------------------------------------- /images/thresholder.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RMDLO/trackdlo/19f92952c6c864f948ee5c05d5c5e3627e190a8f/images/thresholder.png -------------------------------------------------------------------------------- /images/trackdlo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RMDLO/trackdlo/19f92952c6c864f948ee5c05d5c5e3627e190a8f/images/trackdlo.png -------------------------------------------------------------------------------- /images/trackdlo1.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RMDLO/trackdlo/19f92952c6c864f948ee5c05d5c5e3627e190a8f/images/trackdlo1.gif -------------------------------------------------------------------------------- /images/trackdlo2.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RMDLO/trackdlo/19f92952c6c864f948ee5c05d5c5e3627e190a8f/images/trackdlo2.gif -------------------------------------------------------------------------------- /images/trackdlo3.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RMDLO/trackdlo/19f92952c6c864f948ee5c05d5c5e3627e190a8f/images/trackdlo3.gif -------------------------------------------------------------------------------- /images/trackdlo4.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/RMDLO/trackdlo/19f92952c6c864f948ee5c05d5c5e3627e190a8f/images/trackdlo4.gif -------------------------------------------------------------------------------- /launch/evaluation.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | -------------------------------------------------------------------------------- /launch/realsense_node.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | -------------------------------------------------------------------------------- /launch/trackdlo.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | -------------------------------------------------------------------------------- /launch/trackdlo_eval.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | -------------------------------------------------------------------------------- /launch/visualize_output.launch: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | -------------------------------------------------------------------------------- /package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | trackdlo 4 | 1.0.0 5 | TrackDLO ROS package 6 | Jingyi Xiang 7 | Jingyi Xiang 8 | Holly Dinkel 9 | MIT 10 | 11 | roscpp 12 | roscpp 13 | catkin 14 | 15 | pcl_conversions 16 | pcl_ros 17 | libpcl-all-dev 18 | pcl_conversions 19 | pcl_ros 20 | libpcl-all 21 | 22 | 23 | -------------------------------------------------------------------------------- /rviz/tracking.rviz: -------------------------------------------------------------------------------- 1 | Panels: 2 | - Class: rviz/Displays 3 | Help Height: 0 4 | Name: Displays 5 | Property Tree Widget: 6 | Expanded: ~ 7 | Splitter Ratio: 0.5 8 | Tree Height: 321 9 | - Class: rviz/Selection 10 | Name: Selection 11 | - Class: rviz/Tool Properties 12 | Expanded: 13 | - /2D Pose Estimate1 14 | - /2D Nav Goal1 15 | - /Publish Point1 16 | Name: Tool Properties 17 | Splitter Ratio: 0.5886790156364441 18 | - Class: rviz/Views 19 | Expanded: 20 | - /Current View1 21 | Name: Views 22 | Splitter Ratio: 0.5 23 | - Class: rviz/Time 24 | Name: Time 25 | SyncMode: 0 26 | SyncSource: PointCloud2 27 | Preferences: 28 | PromptSaveOnExit: true 29 | Toolbars: 30 | toolButtonStyle: 2 31 | Visualization Manager: 32 | Class: "" 33 | Displays: 34 | - Alpha: 0.5 35 | Cell Size: 1 36 | Class: rviz/Grid 37 | Color: 160; 160; 164 38 | Enabled: true 39 | Line Style: 40 | Line Width: 0.029999999329447746 41 | Value: Lines 42 | Name: Grid 43 | Normal Cell Count: 0 44 | Offset: 45 | X: 0 46 | Y: 0 47 | Z: 0 48 | Plane: XY 49 | Plane Cell Count: 10 50 | Reference Frame: 51 | Value: true 52 | - Class: rviz/Image 53 | Enabled: true 54 | Image Topic: /camera/color/image_raw 55 | Max Value: 1 56 | Median window: 5 57 | Min Value: 0 58 | Name: Image 59 | Normalize Range: true 60 | Queue Size: 2 61 | Transport Hint: raw 62 | Unreliable: false 63 | Value: true 64 | - Class: rviz/Image 65 | Enabled: true 66 | Image Topic: /trackdlo/results_img 67 | Max Value: 1 68 | Median window: 5 69 | Min Value: 0 70 | Name: Image 71 | Normalize Range: true 72 | Queue Size: 2 73 | Transport Hint: raw 74 | Unreliable: false 75 | Value: true 76 | - Class: rviz/MarkerArray 77 | Enabled: true 78 | Marker Topic: /trackdlo/results_marker 79 | Name: MarkerArray 80 | Namespaces: 81 | {} 82 | Queue Size: 100 83 | Value: true 84 | - Class: rviz/MarkerArray 85 | Enabled: false 86 | Marker Topic: /trackdlo/init_nodes_markers 87 | Name: MarkerArray 88 | Namespaces: 89 | {} 90 | Queue Size: 100 91 | Value: false 92 | - Alpha: 1 93 | Autocompute Intensity Bounds: true 94 | Autocompute Value Bounds: 95 | Max Value: 10 96 | Min Value: -10 97 | Value: true 98 | Axis: Z 99 | Channel Name: intensity 100 | Class: rviz/PointCloud2 101 | Color: 255; 255; 255 102 | Color Transformer: RGB8 103 | Decay Time: 0 104 | Enabled: false 105 | Invert Rainbow: false 106 | Max Color: 255; 255; 255 107 | Min Color: 0; 0; 0 108 | Name: PointCloud2 109 | Position Transformer: XYZ 110 | Queue Size: 2 111 | Selectable: true 112 | Size (Pixels): 2 113 | Size (m): 0.009999999776482582 114 | Style: Points 115 | Topic: /camera/depth/color/points 116 | Unreliable: false 117 | Use Fixed Frame: true 118 | Use rainbow: true 119 | Value: false 120 | - Alpha: 1 121 | Autocompute Intensity Bounds: true 122 | Autocompute Value Bounds: 123 | Max Value: 10 124 | Min Value: -10 125 | Value: true 126 | Axis: Z 127 | Channel Name: intensity 128 | Class: rviz/PointCloud2 129 | Color: 255; 255; 255 130 | Color Transformer: Intensity 131 | Decay Time: 0 132 | Enabled: true 133 | Invert Rainbow: false 134 | Max Color: 255; 255; 255 135 | Min Color: 0; 0; 0 136 | Name: PointCloud2 137 | Position Transformer: XYZ 138 | Queue Size: 10 139 | Selectable: true 140 | Size (Pixels): 5 141 | Size (m): 0.009999999776482582 142 | Style: Points 143 | Topic: /trackdlo/filtered_pointcloud 144 | Unreliable: false 145 | Use Fixed Frame: true 146 | Use rainbow: true 147 | Value: true 148 | Enabled: true 149 | Global Options: 150 | Background Color: 85; 87; 83 151 | Default Light: true 152 | Fixed Frame: base_link 153 | Frame Rate: 30 154 | Name: root 155 | Tools: 156 | - Class: rviz/Interact 157 | Hide Inactive Objects: true 158 | - Class: rviz/MoveCamera 159 | - Class: rviz/Select 160 | - Class: rviz/FocusCamera 161 | - Class: rviz/Measure 162 | - Class: rviz/SetInitialPose 163 | Theta std deviation: 0.2617993950843811 164 | Topic: /initialpose 165 | X std deviation: 0.5 166 | Y std deviation: 0.5 167 | - Class: rviz/SetGoal 168 | Topic: /move_base_simple/goal 169 | - Class: rviz/PublishPoint 170 | Single click: true 171 | Topic: /clicked_point 172 | Value: true 173 | Views: 174 | Current: 175 | Class: rviz/Orbit 176 | Distance: 0.7319697141647339 177 | Enable Stereo Rendering: 178 | Stereo Eye Separation: 0.05999999865889549 179 | Stereo Focal Distance: 1 180 | Swap Stereo Eyes: false 181 | Value: false 182 | Field of View: 0.7853981852531433 183 | Focal Point: 184 | X: 0.49177086353302 185 | Y: -0.018738387152552605 186 | Z: -0.0052645024843513966 187 | Focal Shape Fixed Size: true 188 | Focal Shape Size: 0.05000000074505806 189 | Invert Z Axis: false 190 | Name: Current View 191 | Near Clip Distance: 0.009999999776482582 192 | Pitch: 0.7697963714599609 193 | Target Frame: 194 | Yaw: 0.0008482933044433594 195 | Saved: ~ 196 | Window Geometry: 197 | Displays: 198 | collapsed: false 199 | Height: 2002 200 | Hide Left Dock: false 201 | Hide Right Dock: true 202 | Image: 203 | collapsed: false 204 | QMainWindow State: 000000ff00000000fd0000000400000000000005ad000006ccfc020000000afb0000001200530065006c0065006300740069006f006e00000001e10000009b000000b000fffffffb0000001e0054006f006f006c002000500072006f007000650072007400690065007302000001ed000001df00000185000000b0fb000000120056006900650077007300200054006f006f02000001df000002110000018500000122fb000000200054006f006f006c002000500072006f0070006500720074006900650073003203000002880000011d000002210000017afb000000100044006900730070006c006100790073010000006e000001af0000018200fffffffb0000002000730065006c0065006300740069006f006e00200062007500660066006500720200000138000000aa0000023a00000294fb00000014005700690064006500530074006500720065006f02000000e6000000d2000003ee0000030bfb0000000c004b0069006e0065006300740200000186000001060000030c00000261fb0000000a0049006d0061006700650100000229000002820000002600fffffffb0000000a0049006d00610067006501000004b7000002830000002600ffffff000000010000010f0000035afc0200000003fb0000001e0054006f006f006c002000500072006f00700065007200740069006500730100000041000000780000000000000000fb0000000a00560069006500770073000000003d0000035a0000013200fffffffb0000001200530065006c0065006300740069006f006e010000025a000000b200000000000000000000000200000490000000a9fc0100000001fb0000000a00560069006500770073030000004e00000080000002e1000001970000000300000cc00000005efc0100000002fb0000000800540069006d0065010000000000000cc0000006dc00fffffffb0000000800540069006d0065010000000000000450000000000000000000000707000006cc00000004000000040000000800000008fc0000000100000002000000010000000a0054006f006f006c00730100000000ffffffff0000000000000000 205 | Selection: 206 | collapsed: false 207 | Time: 208 | collapsed: false 209 | Tool Properties: 210 | collapsed: false 211 | Views: 212 | collapsed: false 213 | Width: 3264 214 | X: 144 215 | Y: 54 216 | -------------------------------------------------------------------------------- /trackdlo/include/evaluator.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "trackdlo.h" 4 | 5 | #ifndef EVALUATOR_H 6 | #define EVALUATOR_H 7 | 8 | using Eigen::MatrixXd; 9 | using cv::Mat; 10 | 11 | class evaluator 12 | { 13 | public: 14 | evaluator (); 15 | evaluator (int length, int trial, int pct_occlusion, std::string alg, int bag_file, std::string save_location, 16 | double start_record_at, double exit_at, double wait_before_occlusion, double bag_rate, int num_of_nodes); 17 | MatrixXd get_ground_truth_nodes (Mat rgb_img, pcl::PointCloud cloud_xyz); 18 | MatrixXd sort_pts (MatrixXd Y_0, MatrixXd head); 19 | double calc_min_distance (MatrixXd A, MatrixXd B, MatrixXd E, MatrixXd& closest_pt_on_AB_to_E); 20 | double get_piecewise_error (MatrixXd Y_track, MatrixXd Y_true); 21 | double compute_and_save_error (MatrixXd Y_track, MatrixXd Y_true); 22 | void set_start_time (std::chrono::steady_clock::time_point cur_time); 23 | double pct_occlusion (); 24 | std::chrono::steady_clock::time_point start_time (); 25 | double recording_start_time (); 26 | double exit_time (); 27 | int length (); 28 | double wait_before_occlusion (); 29 | double rate (); 30 | double compute_error (MatrixXd Y_track, MatrixXd Y_true); 31 | void increment_image_counter (); 32 | int image_counter (); 33 | 34 | private: 35 | int length_; 36 | int trial_; 37 | int pct_occlusion_; 38 | std::string alg_; 39 | int bag_file_; 40 | std::vector errors_; 41 | std::string save_location_; 42 | std::chrono::steady_clock::time_point start_time_; 43 | double start_record_at_; 44 | double exit_at_; 45 | double wait_before_occlusion_; 46 | bool cleared_file_; 47 | double bag_rate_; 48 | int image_counter_; 49 | int num_of_nodes_; 50 | }; 51 | 52 | #endif -------------------------------------------------------------------------------- /trackdlo/include/trackdlo.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | #include 10 | #include 11 | 12 | #include 13 | #include 14 | #include 15 | #include 16 | #include 17 | #include 18 | #include 19 | 20 | #include 21 | #include 22 | 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include 29 | #include 30 | #include 31 | #include 32 | #include 33 | #include 34 | #include 35 | 36 | #include 37 | #include 38 | #include 39 | #include 40 | #include 41 | 42 | #include 43 | #include 44 | #include 45 | 46 | 47 | #ifndef TRACKDLO_H 48 | #define TRACKDLO_H 49 | 50 | using Eigen::MatrixXd; 51 | using cv::Mat; 52 | 53 | class trackdlo 54 | { 55 | public: 56 | // default constructor 57 | trackdlo(); 58 | trackdlo(int num_of_nodes); 59 | // fancy constructor 60 | trackdlo(int num_of_nodes, 61 | double visibility_threshold, 62 | double beta, 63 | double lambda, 64 | double alpha, 65 | double k_vis, 66 | double mu, 67 | int max_iter, 68 | double tol, 69 | double beta_pre_proc, 70 | double lambda_pre_proc, 71 | double lle_weight); 72 | 73 | double get_sigma2(); 74 | MatrixXd get_tracking_result(); 75 | MatrixXd get_guide_nodes(); 76 | std::vector get_correspondence_pairs(); 77 | void initialize_geodesic_coord (std::vector geodesic_coord); 78 | void initialize_nodes (MatrixXd Y_init); 79 | void set_sigma2 (double sigma2); 80 | 81 | bool cpd_lle (MatrixXd X_orig, 82 | MatrixXd& Y, 83 | double& sigma2, 84 | double beta, 85 | double lambda, 86 | double lle_weight, 87 | double mu, 88 | int max_iter = 30, 89 | double tol = 0.0001, 90 | bool include_lle = true, 91 | std::vector correspondence_priors = {}, 92 | double alpha = 0, 93 | std::vector visible_nodes = {}, 94 | double k_vis = 0, 95 | double visibility_threshold = 0.01); 96 | 97 | void tracking_step (MatrixXd X_orig, 98 | std::vector visible_nodes, 99 | std::vector visible_nodes_extended, 100 | MatrixXd proj_matrix, 101 | int img_rows, 102 | int img_cols); 103 | 104 | private: 105 | MatrixXd Y_; 106 | MatrixXd guide_nodes_; 107 | double sigma2_; 108 | double beta_; 109 | double beta_pre_proc_; 110 | double lambda_; 111 | double lambda_pre_proc_; 112 | double alpha_; 113 | double k_vis_; 114 | double mu_; 115 | int max_iter_; 116 | double tol_; 117 | double lle_weight_; 118 | 119 | std::vector geodesic_coord_; 120 | std::vector correspondence_priors_; 121 | double visibility_threshold_; 122 | 123 | std::vector get_nearest_indices (int k, int M, int idx); 124 | MatrixXd calc_LLE_weights (int k, MatrixXd X); 125 | std::vector traverse_geodesic (std::vector geodesic_coord, const MatrixXd guide_nodes, 126 | const std::vector visible_nodes, int alignment); 127 | std::vector traverse_euclidean (std::vector geodesic_coord, const MatrixXd guide_nodes, 128 | const std::vector visible_nodes, int alignment, int alignment_node_idx = -1); 129 | 130 | }; 131 | 132 | #endif -------------------------------------------------------------------------------- /trackdlo/include/utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "trackdlo.h" 4 | 5 | #ifndef UTILS_H 6 | #define UTILS_H 7 | 8 | using Eigen::MatrixXd; 9 | using cv::Mat; 10 | 11 | void signal_callback_handler(int signum); 12 | 13 | template void print_1d_vector (const std::vector& vec) { 14 | for (auto item : vec) { 15 | std::cout << item << " "; 16 | // std::cout << item << std::endl; 17 | } 18 | std::cout << std::endl; 19 | } 20 | 21 | double pt2pt_dis_sq (MatrixXd pt1, MatrixXd pt2); 22 | double pt2pt_dis (MatrixXd pt1, MatrixXd pt2); 23 | 24 | void reg (MatrixXd pts, MatrixXd& Y, double& sigma2, int M, double mu = 0, int max_iter = 50); 25 | void remove_row(MatrixXd& matrix, unsigned int rowToRemove); 26 | MatrixXd sort_pts (MatrixXd Y_0); 27 | 28 | std::vector line_sphere_intersection (MatrixXd point_A, MatrixXd point_B, MatrixXd sphere_center, double radius); 29 | 30 | visualization_msgs::MarkerArray MatrixXd2MarkerArray (MatrixXd Y, 31 | std::string marker_frame, 32 | std::string marker_ns, 33 | std::vector node_color, 34 | std::vector line_color, 35 | double node_scale = 0.01, 36 | double line_scale = 0.005, 37 | std::vector visible_nodes = {}, 38 | std::vector occluded_node_color = {}, 39 | std::vector occluded_line_color = {}); 40 | 41 | visualization_msgs::MarkerArray MatrixXd2MarkerArray (std::vector Y, 42 | std::string marker_frame, 43 | std::string marker_ns, 44 | std::vector node_color, 45 | std::vector line_color, 46 | double node_scale = 0.01, 47 | double line_scale = 0.005, 48 | std::vector visible_nodes = {}, 49 | std::vector occluded_node_color = {}, 50 | std::vector occluded_line_color = {}); 51 | 52 | MatrixXd cross_product (MatrixXd vec1, MatrixXd vec2); 53 | double dot_product (MatrixXd vec1, MatrixXd vec2); 54 | 55 | #endif -------------------------------------------------------------------------------- /trackdlo/src/evaluator.cpp: -------------------------------------------------------------------------------- 1 | #include "../include/utils.h" 2 | #include "../include/trackdlo.h" 3 | #include "../include/evaluator.h" 4 | 5 | #include 6 | #include 7 | #include 8 | 9 | using Eigen::MatrixXd; 10 | using Eigen::RowVectorXd; 11 | using cv::Mat; 12 | 13 | evaluator::evaluator () {} 14 | evaluator::evaluator (int length, int trial, int pct_occlusion, std::string alg, int bag_file, std::string save_location, 15 | double start_record_at, double exit_at, double wait_before_occlusion, double bag_rate, int num_of_nodes) 16 | { 17 | length_ = length; 18 | trial_ = trial; 19 | pct_occlusion_ = pct_occlusion; 20 | alg_ = alg; 21 | bag_file_ = bag_file; 22 | save_location_ = save_location; 23 | start_record_at_ = start_record_at; 24 | exit_at_ = exit_at; 25 | wait_before_occlusion_ = wait_before_occlusion; 26 | cleared_file_ = false; 27 | bag_rate_ = bag_rate; 28 | num_of_nodes_ = num_of_nodes; 29 | image_counter_ = 0; 30 | } 31 | 32 | int evaluator::image_counter () { 33 | return image_counter_; 34 | } 35 | 36 | void evaluator::increment_image_counter () { 37 | image_counter_ += 1; 38 | } 39 | 40 | void evaluator::set_start_time (std::chrono::steady_clock::time_point cur_time) { 41 | start_time_ = cur_time; 42 | } 43 | 44 | std::chrono::steady_clock::time_point evaluator::start_time () { 45 | return start_time_; 46 | } 47 | 48 | double evaluator::pct_occlusion () { 49 | return pct_occlusion_; 50 | } 51 | 52 | double evaluator::recording_start_time () { 53 | return start_record_at_; 54 | } 55 | 56 | double evaluator::exit_time () { 57 | return exit_at_; 58 | } 59 | 60 | int evaluator::length () { 61 | return length_; 62 | } 63 | 64 | double evaluator::wait_before_occlusion () { 65 | return wait_before_occlusion_; 66 | } 67 | 68 | double evaluator::rate () { 69 | return bag_rate_; 70 | } 71 | 72 | MatrixXd evaluator::sort_pts (MatrixXd Y_0, MatrixXd head) { 73 | int N = Y_0.rows(); 74 | MatrixXd Y_0_sorted = MatrixXd::Zero(N, 3); 75 | std::vector Y_0_sorted_vec = {}; 76 | std::vector selected_node(N, false); 77 | selected_node[0] = true; 78 | int last_visited_b = 0; 79 | 80 | MatrixXd G = MatrixXd::Zero(N, N); 81 | for (int i = 0; i < N; i ++) { 82 | for (int j = 0; j < N; j ++) { 83 | G(i, j) = (Y_0.row(i) - Y_0.row(j)).squaredNorm(); 84 | } 85 | } 86 | 87 | int reverse = 0; 88 | int counter = 0; 89 | int reverse_on = 0; 90 | int insertion_counter = 0; 91 | 92 | while (counter < N-1) { 93 | double minimum = INFINITY; 94 | int a = 0; 95 | int b = 0; 96 | 97 | for (int m = 0; m < N; m ++) { 98 | if (selected_node[m] == true) { 99 | for (int n = 0; n < N; n ++) { 100 | if ((!selected_node[n]) && (G(m, n) != 0.0)) { 101 | if (minimum > G(m, n)) { 102 | minimum = G(m, n); 103 | a = m; 104 | b = n; 105 | } 106 | } 107 | } 108 | } 109 | } 110 | 111 | if (counter == 0) { 112 | Y_0_sorted_vec.push_back(Y_0.row(a)); 113 | Y_0_sorted_vec.push_back(Y_0.row(b)); 114 | } 115 | else { 116 | if (last_visited_b != a) { 117 | reverse += 1; 118 | reverse_on = a; 119 | insertion_counter = 1; 120 | } 121 | 122 | if (reverse % 2 == 1) { 123 | auto it = find(Y_0_sorted_vec.begin(), Y_0_sorted_vec.end(), Y_0.row(a)); 124 | Y_0_sorted_vec.insert(it, Y_0.row(b)); 125 | } 126 | else if (reverse != 0) { 127 | auto it = find(Y_0_sorted_vec.begin(), Y_0_sorted_vec.end(), Y_0.row(reverse_on)); 128 | Y_0_sorted_vec.insert(it + insertion_counter, Y_0.row(b)); 129 | insertion_counter += 1; 130 | } 131 | else { 132 | Y_0_sorted_vec.push_back(Y_0.row(b)); 133 | } 134 | } 135 | 136 | last_visited_b = b; 137 | selected_node[b] = true; 138 | counter += 1; 139 | } 140 | 141 | if (pt2pt_dis(Y_0_sorted_vec[0], head) > 0.08) { 142 | std::reverse(Y_0_sorted_vec.begin(), Y_0_sorted_vec.end()); 143 | } 144 | 145 | // copy to Y_0_sorted 146 | for (int i = 0; i < N; i ++) { 147 | Y_0_sorted.row(i) = Y_0_sorted_vec[i]; 148 | } 149 | 150 | return Y_0_sorted; 151 | } 152 | 153 | MatrixXd evaluator::get_ground_truth_nodes (Mat rgb_img, pcl::PointCloud cloud_xyz) { 154 | Mat mask_blue, mask_red_1, mask_red_2, mask_red, mask_yellow, mask_markers, mask, mask_rgb; 155 | Mat cur_image_hsv, mask_without_occlusion_block; 156 | 157 | // convert color 158 | cv::cvtColor(rgb_img, cur_image_hsv, cv::COLOR_BGR2HSV); 159 | 160 | std::vector lower_blue = {90, 80, 80}; 161 | std::vector upper_blue = {130, 255, 255}; 162 | 163 | std::vector lower_red_1 = {130, 60, 50}; 164 | std::vector upper_red_1 = {255, 255, 255}; 165 | 166 | std::vector lower_red_2 = {0, 60, 50}; 167 | std::vector upper_red_2 = {10, 255, 255}; 168 | 169 | std::vector lower_yellow = {15, 100, 80}; 170 | std::vector upper_yellow = {40, 255, 255}; 171 | 172 | // filter blue 173 | cv::inRange(cur_image_hsv, cv::Scalar(lower_blue[0], lower_blue[1], lower_blue[2]), cv::Scalar(upper_blue[0], upper_blue[1], upper_blue[2]), mask_blue); 174 | 175 | // filter red 176 | cv::inRange(cur_image_hsv, cv::Scalar(lower_red_1[0], lower_red_1[1], lower_red_1[2]), cv::Scalar(upper_red_1[0], upper_red_1[1], upper_red_1[2]), mask_red_1); 177 | cv::inRange(cur_image_hsv, cv::Scalar(lower_red_2[0], lower_red_2[1], lower_red_2[2]), cv::Scalar(upper_red_2[0], upper_red_2[1], upper_red_2[2]), mask_red_2); 178 | 179 | // filter yellow 180 | cv::inRange(cur_image_hsv, cv::Scalar(lower_yellow[0], lower_yellow[1], lower_yellow[2]), cv::Scalar(upper_yellow[0], upper_yellow[1], upper_yellow[2]), mask_yellow); 181 | 182 | // combine red mask 183 | cv::bitwise_or(mask_red_1, mask_red_2, mask_red); 184 | // combine overall mask 185 | cv::bitwise_or(mask_red, mask_blue, mask_without_occlusion_block); 186 | cv::bitwise_or(mask_yellow, mask_without_occlusion_block, mask_without_occlusion_block); 187 | cv::bitwise_or(mask_red, mask_yellow, mask_markers); 188 | 189 | // simple blob detector 190 | std::vector keypoints_markers; 191 | // std::vector keypoints_blue; 192 | cv::SimpleBlobDetector::Params blob_params; 193 | blob_params.filterByColor = false; 194 | blob_params.filterByArea = true; 195 | blob_params.minArea = 10; 196 | blob_params.filterByCircularity = false; 197 | blob_params.filterByInertia = true; 198 | blob_params.filterByConvexity = false; 199 | cv::Ptr detector = cv::SimpleBlobDetector::create(blob_params); 200 | // detect 201 | detector->detect(mask_markers, keypoints_markers); 202 | // detector->detect(mask_blue, keypoints_blue); 203 | 204 | pcl::PointCloud cur_nodes_xyz; 205 | 206 | for (cv::KeyPoint key_point : keypoints_markers) { 207 | auto keypoint_pc = cloud_xyz(static_cast(key_point.pt.x), static_cast(key_point.pt.y)); 208 | 209 | if (bag_file_ == 2) { 210 | if (keypoint_pc.x < -0.15 || keypoint_pc.y < -0.15 || keypoint_pc.z < 0.58) { 211 | continue; 212 | } 213 | } 214 | else if (bag_file_ == 1) { 215 | if ((keypoint_pc.x < 0.0 && keypoint_pc.y < 0.05) || keypoint_pc.z < 0.58 || 216 | keypoint_pc.x < -0.2 || (keypoint_pc.x < 0.1 && keypoint_pc.y < -0.05)) { 217 | continue; 218 | } 219 | } 220 | else { 221 | if (keypoint_pc.z < 0.58) { 222 | continue; 223 | } 224 | } 225 | 226 | cur_nodes_xyz.push_back(keypoint_pc); 227 | } 228 | 229 | // the node set returned is not sorted 230 | return cur_nodes_xyz.getMatrixXfMap().topRows(3).transpose().cast(); 231 | } 232 | 233 | double evaluator::calc_min_distance (MatrixXd A, MatrixXd B, MatrixXd E, MatrixXd& closest_pt_on_AB_to_E) { 234 | MatrixXd AB = B - A; 235 | MatrixXd AE = E - A; 236 | 237 | double distance = cross_product(AE, AB).norm() / AB.norm(); 238 | closest_pt_on_AB_to_E = A + AB*dot_product(AE, AB) / dot_product(AB, AB); 239 | 240 | MatrixXd AP = closest_pt_on_AB_to_E - A; 241 | if (dot_product(AP, AB) < 0 || dot_product(AP, AB) > dot_product(AB, AB)) { 242 | MatrixXd BE = E - B; 243 | double distance_AE = sqrt(dot_product(AE, AE)); 244 | double distance_BE = sqrt(dot_product(BE, BE)); 245 | if (distance_AE > distance_BE) { 246 | distance = distance_BE; 247 | closest_pt_on_AB_to_E = B.replicate(1, 1); 248 | } 249 | else { 250 | distance = distance_AE; 251 | closest_pt_on_AB_to_E = A.replicate(1, 1); 252 | } 253 | } 254 | 255 | return distance; 256 | } 257 | 258 | double evaluator::get_piecewise_error (MatrixXd Y_track, MatrixXd Y_true) { 259 | double total_distances_to_curve = 0.0; 260 | std::vector closest_pts_on_Y_true = {}; 261 | 262 | for (int idx = 0; idx < Y_track.rows(); idx ++) { 263 | double dist = -1; 264 | MatrixXd closest_pt = MatrixXd::Zero(1, 3); 265 | 266 | for (int i = 0; i < Y_true.rows()-1; i ++) { 267 | MatrixXd closest_pt_i = MatrixXd::Zero(1, 3); 268 | double dist_i = calc_min_distance(Y_true.row(i), Y_true.row(i+1), Y_track.row(idx), closest_pt_i); 269 | if (dist == -1 || dist_i < dist) { 270 | dist = dist_i; 271 | closest_pt = closest_pt_i.replicate(1, 1); 272 | } 273 | } 274 | 275 | total_distances_to_curve += dist; 276 | closest_pts_on_Y_true.push_back(closest_pt); 277 | } 278 | 279 | // double error_frame = total_distances_to_curve / num_of_nodes_; 280 | double error_frame = total_distances_to_curve / Y_track.rows(); 281 | 282 | return error_frame; 283 | } 284 | 285 | double evaluator::compute_and_save_error (MatrixXd Y_track, MatrixXd Y_true) { 286 | // compute error 287 | double E1 = get_piecewise_error(Y_track, Y_true); 288 | double E2 = get_piecewise_error(Y_true, Y_track); 289 | 290 | double cur_frame_error = (E1 + E2) / 2; 291 | errors_.push_back(cur_frame_error); 292 | 293 | std::string dir; 294 | // 0 -> statinary.bag; 1 -> with_gripper_perpendicular.bag; 2 -> with_gripper_parallel.bag 295 | if (bag_file_ == 0) { 296 | dir = save_location_ + alg_ + "_" + std::to_string(trial_) + "_" + std::to_string(pct_occlusion_) + "_stationary_error.txt"; 297 | } 298 | else if (bag_file_ == 1) { 299 | dir = save_location_ + alg_ + "_" + std::to_string(trial_) + "_" + std::to_string(pct_occlusion_) + "_perpendicular_motion_error.txt"; 300 | } 301 | else if (bag_file_ == 2) { 302 | dir = save_location_ + alg_ + "_" + std::to_string(trial_) + "_" + std::to_string(pct_occlusion_) + "_parallel_motion_error.txt"; 303 | } 304 | else if (bag_file_ == 4) { 305 | dir = save_location_ + alg_ + "_" + std::to_string(trial_) + "_" + std::to_string(pct_occlusion_) + "_short_rope_folding_error.txt"; 306 | } 307 | else if (bag_file_ == 5) { 308 | dir = save_location_ + alg_ + "_" + std::to_string(trial_) + "_" + std::to_string(pct_occlusion_) + "_short_rope_stationary_error.txt"; 309 | } 310 | else { 311 | throw std::invalid_argument("Invalid bag file ID!"); 312 | } 313 | 314 | double time_diff; 315 | time_diff = std::chrono::duration_cast(std::chrono::steady_clock::now() - start_time_).count(); 316 | time_diff = time_diff / 1000.0 * bag_rate_; 317 | 318 | if (cleared_file_ = false) { 319 | std::ofstream error_list (dir); 320 | error_list << std::to_string(time_diff - start_record_at_) + " " + std::to_string(cur_frame_error) + "\n"; 321 | error_list.close(); 322 | cleared_file_ = true; 323 | } 324 | else { 325 | std::ofstream error_list (dir, std::fstream::app); 326 | error_list << std::to_string(time_diff - start_record_at_) + " " + std::to_string(cur_frame_error) + "\n"; 327 | error_list.close(); 328 | } 329 | 330 | return cur_frame_error; 331 | } 332 | 333 | double evaluator::compute_error (MatrixXd Y_track, MatrixXd Y_true) { 334 | // compute error 335 | double E1 = get_piecewise_error(Y_track, Y_true); 336 | double E2 = get_piecewise_error(Y_true, Y_track); 337 | 338 | double cur_frame_error = (E1 + E2) / 2; 339 | 340 | return cur_frame_error; 341 | } -------------------------------------------------------------------------------- /trackdlo/src/initialize.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import rospy 4 | import ros_numpy 5 | from sensor_msgs.msg import PointCloud2, PointField, Image, CameraInfo 6 | import sensor_msgs.point_cloud2 as pcl2 7 | import std_msgs.msg 8 | import message_filters 9 | 10 | import struct 11 | import time 12 | import cv2 13 | import numpy as np 14 | 15 | from visualization_msgs.msg import MarkerArray 16 | from scipy import interpolate 17 | 18 | from utils import extract_connected_skeleton, ndarray2MarkerArray 19 | 20 | proj_matrix = None 21 | def camera_info_callback (info): 22 | global proj_matrix 23 | proj_matrix = np.array(list(info.P)).reshape(3, 4) 24 | print('Received camera projection matrix:') 25 | print(proj_matrix) 26 | camera_info_sub.unregister() 27 | 28 | def color_thresholding (hsv_image, cur_depth): 29 | global lower, upper 30 | 31 | mask_dlo = cv2.inRange(hsv_image.copy(), lower, upper).astype('uint8') 32 | 33 | # tape green 34 | lower_green = (58, 130, 50) 35 | upper_green = (90, 255, 89) 36 | mask_green = cv2.inRange(hsv_image.copy(), lower_green, upper_green).astype('uint8') 37 | 38 | # combine masks 39 | mask = cv2.bitwise_or(mask_green.copy(), mask_dlo.copy()) 40 | 41 | # filter mask base on depth values 42 | mask[cur_depth < 0.57*1000] = 0 43 | 44 | return mask, mask_green 45 | 46 | def remove_duplicate_rows(array): 47 | _, idx = np.unique(array, axis=0, return_index=True) 48 | data = array[np.sort(idx)] 49 | rospy.set_param('/init_tracker/num_of_nodes', len(data)) 50 | return data 51 | 52 | def callback (rgb, depth): 53 | global lower, upper 54 | 55 | print("Initializing...") 56 | 57 | # process rgb image 58 | cur_image = ros_numpy.numpify(rgb) 59 | hsv_image = cv2.cvtColor(cur_image.copy(), cv2.COLOR_RGB2HSV) 60 | 61 | # process depth image 62 | cur_depth = ros_numpy.numpify(depth) 63 | 64 | if not multi_color_dlo: 65 | # color thresholding 66 | mask = cv2.inRange(hsv_image, lower, upper) 67 | else: 68 | # color thresholding 69 | mask, mask_tip = color_thresholding(hsv_image, cur_depth) 70 | try: 71 | start_time = time.time() 72 | mask = cv2.cvtColor(mask.copy(), cv2.COLOR_GRAY2BGR) 73 | 74 | # returns the pixel coord of points (in order). a list of lists 75 | img_scale = 1 76 | extracted_chains = extract_connected_skeleton(visualize_initialization_process, mask, img_scale=img_scale, seg_length=8, max_curvature=25) 77 | 78 | all_pixel_coords = [] 79 | for chain in extracted_chains: 80 | all_pixel_coords += chain 81 | print('Finished extracting chains. Time taken:', time.time()-start_time) 82 | 83 | all_pixel_coords = np.array(all_pixel_coords) * img_scale 84 | all_pixel_coords = np.flip(all_pixel_coords, 1) 85 | 86 | pc_z = cur_depth[tuple(map(tuple, all_pixel_coords.T))] / 1000.0 87 | fx = proj_matrix[0, 0] 88 | fy = proj_matrix[1, 1] 89 | cx = proj_matrix[0, 2] 90 | cy = proj_matrix[1, 2] 91 | pixel_x = all_pixel_coords[:, 1] 92 | pixel_y = all_pixel_coords[:, 0] 93 | # if the first mask value is not in the tip mask, reverse the pixel order 94 | if multi_color_dlo: 95 | pixel_value1 = mask_tip[pixel_y[-1],pixel_x[-1]] 96 | if pixel_value1 == 255: 97 | pixel_x, pixel_y = pixel_x[::-1], pixel_y[::-1] 98 | 99 | pc_x = (pixel_x - cx) * pc_z / fx 100 | pc_y = (pixel_y - cy) * pc_z / fy 101 | extracted_chains_3d = np.vstack((pc_x, pc_y)) 102 | extracted_chains_3d = np.vstack((extracted_chains_3d, pc_z)) 103 | extracted_chains_3d = extracted_chains_3d.T 104 | 105 | # do not include those without depth values 106 | extracted_chains_3d = extracted_chains_3d[((extracted_chains_3d[:, 0] != 0) | (extracted_chains_3d[:, 1] != 0) | (extracted_chains_3d[:, 2] != 0))] 107 | 108 | if multi_color_dlo: 109 | depth_threshold = 0.57 # m 110 | extracted_chains_3d = extracted_chains_3d[extracted_chains_3d[:, 2] > depth_threshold] 111 | 112 | # tck, u = interpolate.splprep(extracted_chains_3d.T, s=0.001) 113 | tck, u = interpolate.splprep(extracted_chains_3d.T, s=0.0005) 114 | # 1st fit, less points 115 | u_fine = np.linspace(0, 1, 300) # <-- num fit points 116 | x_fine, y_fine, z_fine = interpolate.splev(u_fine, tck) 117 | spline_pts = np.vstack((x_fine, y_fine, z_fine)).T 118 | 119 | # 2nd fit, higher accuracy 120 | num_true_pts = int(np.sum(np.sqrt(np.sum(np.square(np.diff(spline_pts, axis=0)), axis=1))) * 1000) 121 | u_fine = np.linspace(0, 1, num_true_pts) # <-- num true points 122 | x_fine, y_fine, z_fine = interpolate.splev(u_fine, tck) 123 | spline_pts = np.vstack((x_fine, y_fine, z_fine)).T 124 | 125 | nodes = spline_pts[np.linspace(0, num_true_pts-1, num_of_nodes).astype(int)] 126 | 127 | init_nodes = remove_duplicate_rows(nodes) 128 | # results = ndarray2MarkerArray(init_nodes, result_frame_id, [1, 150/255, 0, 0.75], [0, 1, 0, 0.75]) 129 | results = ndarray2MarkerArray(init_nodes, result_frame_id, [0, 149/255, 203/255, 0.75], [0, 149/255, 203/255, 0.75]) 130 | results_pub.publish(results) 131 | 132 | # add color 133 | pc_rgba = struct.unpack('I', struct.pack('BBBB', 255, 40, 40, 255))[0] 134 | pc_rgba_arr = np.full((len(init_nodes), 1), pc_rgba) 135 | pc_colored = np.hstack((init_nodes, pc_rgba_arr)).astype(object) 136 | pc_colored[:, 3] = pc_colored[:, 3].astype(int) 137 | 138 | header.stamp = rospy.Time.now() 139 | converted_points = pcl2.create_cloud(header, fields, pc_colored) 140 | pc_pub.publish(converted_points) 141 | except: 142 | rospy.logerr("Failed to extract splines.") 143 | rospy.signal_shutdown('Stopping initialization.') 144 | 145 | if __name__=='__main__': 146 | rospy.init_node('init_tracker', anonymous=True) 147 | 148 | global new_messages, lower, upper 149 | new_messages=False 150 | 151 | num_of_nodes = rospy.get_param('/init_tracker/num_of_nodes') 152 | multi_color_dlo = rospy.get_param('/init_tracker/multi_color_dlo') 153 | camera_info_topic = rospy.get_param('/init_tracker/camera_info_topic') 154 | rgb_topic = rospy.get_param('/init_tracker/rgb_topic') 155 | depth_topic = rospy.get_param('/init_tracker/depth_topic') 156 | result_frame_id = rospy.get_param('/init_tracker/result_frame_id') 157 | visualize_initialization_process = rospy.get_param('/init_tracker/visualize_initialization_process') 158 | 159 | hsv_threshold_upper_limit = rospy.get_param('/init_tracker/hsv_threshold_upper_limit') 160 | hsv_threshold_lower_limit = rospy.get_param('/init_tracker/hsv_threshold_lower_limit') 161 | 162 | upper_array = hsv_threshold_upper_limit.split(' ') 163 | lower_array = hsv_threshold_lower_limit.split(' ') 164 | upper = (int(upper_array[0]), int(upper_array[1]), int(upper_array[2])) 165 | lower = (int(lower_array[0]), int(lower_array[1]), int(lower_array[2])) 166 | 167 | camera_info_sub = rospy.Subscriber(camera_info_topic, CameraInfo, camera_info_callback) 168 | rgb_sub = message_filters.Subscriber(rgb_topic, Image) 169 | depth_sub = message_filters.Subscriber(depth_topic, Image) 170 | 171 | # header 172 | header = std_msgs.msg.Header() 173 | header.stamp = rospy.Time.now() 174 | header.frame_id = result_frame_id 175 | fields = [PointField('x', 0, PointField.FLOAT32, 1), 176 | PointField('y', 4, PointField.FLOAT32, 1), 177 | PointField('z', 8, PointField.FLOAT32, 1), 178 | PointField('rgba', 12, PointField.UINT32, 1)] 179 | pc_pub = rospy.Publisher ('/trackdlo/init_nodes', PointCloud2, queue_size=10) 180 | results_pub = rospy.Publisher ('/trackdlo/init_nodes_markers', MarkerArray, queue_size=10) 181 | 182 | ts = message_filters.TimeSynchronizer([rgb_sub, depth_sub], 10) 183 | ts.registerCallback(callback) 184 | 185 | rospy.spin() -------------------------------------------------------------------------------- /trackdlo/src/run_evaluation.cpp: -------------------------------------------------------------------------------- 1 | #include "../include/trackdlo.h" 2 | #include "../include/utils.h" 3 | #include "../include/evaluator.h" 4 | 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | using Eigen::MatrixXd; 11 | using Eigen::RowVectorXd; 12 | using cv::Mat; 13 | 14 | int bag_file; 15 | int trial; 16 | std::string alg; 17 | std::string bag_dir; 18 | std::string save_location; 19 | int pct_occlusion; 20 | double start_record_at; 21 | double exit_at; 22 | double wait_before_occlusion; 23 | double bag_rate; 24 | int num_of_nodes; 25 | bool save_images; 26 | bool save_errors; 27 | 28 | int callback_count = 0; 29 | evaluator tracking_evaluator; 30 | MatrixXd head_node = MatrixXd::Zero(1, 3); 31 | 32 | MatrixXd proj_matrix(3, 4); 33 | ros::Publisher corners_arr_pub; 34 | bool initialized = false; 35 | 36 | // sensor_msgs::ImagePtr Callback(const sensor_msgs::ImageConstPtr& image_msg, const sensor_msgs::PointCloud2ConstPtr& pc_msg, const sensor_msgs::PointCloud2ConstPtr& result_msg) { 37 | sensor_msgs::ImagePtr Callback(const sensor_msgs::ImageConstPtr& image_msg, const sensor_msgs::PointCloud2ConstPtr& pc_msg, const sensor_msgs::PointCloud2ConstPtr& result_msg) { 38 | sensor_msgs::ImagePtr eval_img_msg = nullptr; 39 | 40 | if (!initialized) { 41 | tracking_evaluator.set_start_time (std::chrono::steady_clock::now()); 42 | initialized = true; 43 | } 44 | 45 | double time_from_start; 46 | time_from_start = std::chrono::duration_cast(std::chrono::steady_clock::now() - tracking_evaluator.start_time()).count(); 47 | time_from_start = (time_from_start / 1000.0 + 3.0) * tracking_evaluator.rate(); // initialization usually takes 1.5 seconds 48 | std::cout << time_from_start << "; " << tracking_evaluator.exit_time() << std::endl; 49 | 50 | if (tracking_evaluator.exit_time() == -1) { 51 | if (callback_count >= tracking_evaluator.length() - 3) { 52 | std::cout << "Shutting down evaluator..." << std::endl; 53 | ros::shutdown(); 54 | } 55 | } 56 | else { 57 | if (time_from_start > tracking_evaluator.exit_time() || callback_count >= tracking_evaluator.length() - 3) { 58 | std::cout << "Shutting down evaluator..." << std::endl; 59 | ros::shutdown(); 60 | } 61 | } 62 | 63 | callback_count += 1; 64 | std::cout << "callback: " << callback_count << std::endl; 65 | 66 | Mat cur_image_orig = cv_bridge::toCvShare(image_msg, "bgr8")->image; 67 | 68 | // for visualizing results 69 | Mat eval_img; 70 | cur_image_orig.copyTo(eval_img); 71 | 72 | // process original pc 73 | pcl::PCLPointCloud2* cloud = new pcl::PCLPointCloud2; 74 | pcl_conversions::toPCL(*pc_msg, *cloud); 75 | pcl::PointCloud cloud_xyz; 76 | pcl::fromPCLPointCloud2(*cloud, cloud_xyz); 77 | 78 | // process result pc 79 | pcl::PCLPointCloud2* result_cloud = new pcl::PCLPointCloud2; 80 | pcl_conversions::toPCL(*result_msg, *result_cloud); 81 | pcl::PointCloud result_cloud_xyz; 82 | pcl::fromPCLPointCloud2(*result_cloud, result_cloud_xyz); 83 | MatrixXd Y_track = result_cloud_xyz.getMatrixXfMap().topRows(3).transpose().cast(); 84 | 85 | int top_left_x = -1; 86 | int top_left_y = -1; 87 | int bottom_right_x = -1; 88 | int bottom_right_y = -1; 89 | 90 | double cur_error = -1; 91 | if (time_from_start > tracking_evaluator.recording_start_time()) { 92 | 93 | if (bag_file != 3) { 94 | MatrixXd gt_nodes = tracking_evaluator.get_ground_truth_nodes(cur_image_orig, cloud_xyz); 95 | MatrixXd Y_true = gt_nodes.replicate(1, 1); 96 | // if not initialized 97 | if (head_node(0, 0) == 0.0 && head_node(0, 1) == 0.0 && head_node(0, 2) == 0.0) { 98 | // the one with greater x. this holds true for all 3 bag files 99 | if (Y_track(0, 0) > Y_track(Y_track.rows()-1, 0)) { 100 | head_node = Y_track.row(Y_track.rows()-1).replicate(1, 1); 101 | } 102 | else { 103 | head_node = Y_track.row(0).replicate(1, 1); 104 | } 105 | } 106 | Y_true = tracking_evaluator.sort_pts(gt_nodes, head_node); 107 | 108 | // update head node 109 | head_node = Y_true.row(0).replicate(1, 1); 110 | std::cout << "Y_true size: " << Y_true.rows() << "; Y_track size: " << Y_track.rows() << std::endl; 111 | 112 | if (time_from_start > tracking_evaluator.recording_start_time() + tracking_evaluator.wait_before_occlusion()) { 113 | if (bag_file == 0) { 114 | // simulate occlusion: occlude the first n nodes 115 | // strategy: first calculate the 3D boundary box based on point cloud, then project the four corners back to the image 116 | int num_of_occluded_nodes = static_cast(Y_track.rows() * tracking_evaluator.pct_occlusion() / 100.0); 117 | 118 | if (num_of_occluded_nodes != 0) { 119 | 120 | double min_x = Y_true(0, 0); 121 | double min_y = Y_true(0, 1); 122 | double min_z = Y_true(0, 2); 123 | 124 | double max_x = Y_true(0, 0); 125 | double max_y = Y_true(0, 1); 126 | double max_z = Y_true(0, 2); 127 | 128 | for (int i = 0; i < num_of_occluded_nodes; i ++) { 129 | if (Y_true(i, 0) < min_x) { 130 | min_x = Y_true(i, 0); 131 | } 132 | if (Y_true(i, 1) < min_y) { 133 | min_y = Y_true(i, 1); 134 | } 135 | if (Y_true(i, 2) < min_z) { 136 | min_z = Y_true(i, 2); 137 | } 138 | 139 | if (Y_true(i, 0) > max_x) { 140 | max_x = Y_true(i, 0); 141 | } 142 | if (Y_true(i, 1) > max_y) { 143 | max_y = Y_true(i, 1); 144 | } 145 | if (Y_true(i, 2) > max_z) { 146 | max_z = Y_true(i, 2); 147 | } 148 | } 149 | 150 | MatrixXd min_corner(1, 3); 151 | min_corner << min_x, min_y, min_z; 152 | MatrixXd max_corner(1, 3); 153 | max_corner << max_x, max_y, max_z; 154 | 155 | MatrixXd corners = MatrixXd::Zero(2, 3); 156 | corners.row(0) = min_corner.replicate(1, 1); 157 | corners.row(1) = max_corner.replicate(1, 1); 158 | 159 | // project to find occlusion block coorindate 160 | MatrixXd nodes_h = corners.replicate(1, 1); 161 | nodes_h.conservativeResize(nodes_h.rows(), nodes_h.cols()+1); 162 | nodes_h.col(nodes_h.cols()-1) = MatrixXd::Ones(nodes_h.rows(), 1); 163 | MatrixXd image_coords = (proj_matrix * nodes_h.transpose()).transpose(); 164 | 165 | int pix_coord_1_x = static_cast(image_coords(0, 0)/image_coords(0, 2)); 166 | int pix_coord_1_y = static_cast(image_coords(0, 1)/image_coords(0, 2)); 167 | int pix_coord_2_x = static_cast(image_coords(1, 0)/image_coords(1, 2)); 168 | int pix_coord_2_y = static_cast(image_coords(1, 1)/image_coords(1, 2)); 169 | 170 | int extra_border = 30; 171 | 172 | // best scenarios: min_corner and max_corner are the top left and bottom right corners 173 | if (pix_coord_1_x <= pix_coord_2_x && pix_coord_1_y <= pix_coord_2_y) { 174 | // cv::Point p1(pix_coord_1_x - extra_border, pix_coord_1_y - extra_border); 175 | // cv::Point p2(pix_coord_2_x + extra_border, pix_coord_2_y + extra_border); 176 | // cv::rectangle(occlusion_mask, p1, p2, cv::Scalar(0, 0, 0), -1); 177 | top_left_x = pix_coord_1_x - extra_border; 178 | if (top_left_x < 0) {top_left_x = 0;} 179 | top_left_y = pix_coord_1_y - extra_border; 180 | if (top_left_y < 0) {top_left_y = 0;} 181 | bottom_right_x = pix_coord_2_x + extra_border; 182 | if (bottom_right_x >= cur_image_orig.cols) {bottom_right_x = cur_image_orig.cols-1;} 183 | bottom_right_y = pix_coord_2_y + extra_border; 184 | if (bottom_right_y >= cur_image_orig.rows) {bottom_right_y = cur_image_orig.rows-1;} 185 | } 186 | // best scenarios: min_corner and max_corner are the top left and bottom right corners 187 | else if (pix_coord_2_x <= pix_coord_1_x && pix_coord_2_y <= pix_coord_1_y) { 188 | // cv::Point p1(pix_coord_2_x - extra_border, pix_coord_2_y - extra_border); 189 | // cv::Point p2(pix_coord_1_x + extra_border, pix_coord_1_y + extra_border); 190 | // cv::rectangle(occlusion_mask, p1, p2, cv::Scalar(0, 0, 0), -1); 191 | top_left_x = pix_coord_2_x - extra_border; 192 | if (top_left_x < 0) {top_left_x = 0;} 193 | top_left_y = pix_coord_2_y - extra_border; 194 | if (top_left_y < 0) {top_left_y = 0;} 195 | bottom_right_x = pix_coord_1_x + extra_border; 196 | if (bottom_right_x >= cur_image_orig.cols) {bottom_right_x = cur_image_orig.cols-1;} 197 | bottom_right_y = pix_coord_1_y + extra_border; 198 | if (bottom_right_y >= cur_image_orig.rows) {bottom_right_y = cur_image_orig.rows-1;} 199 | } 200 | // min_corner is top right, max_corner is bottom left 201 | else if (pix_coord_2_x <= pix_coord_1_x && pix_coord_1_y <= pix_coord_2_y) { 202 | // cv::Point p1(pix_coord_2_x - extra_border, pix_coord_1_y - extra_border); 203 | // cv::Point p2(pix_coord_1_x + extra_border, pix_coord_2_y + extra_border); 204 | // cv::rectangle(occlusion_mask, p1, p2, cv::Scalar(0, 0, 0), -1); 205 | top_left_x = pix_coord_2_x - extra_border; 206 | if (top_left_x < 0) {top_left_x = 0;} 207 | top_left_y = pix_coord_1_y - extra_border; 208 | if (top_left_y < 0) {top_left_y = 0;} 209 | bottom_right_x = pix_coord_1_x + extra_border; 210 | if (bottom_right_x >= cur_image_orig.cols) {bottom_right_x = cur_image_orig.cols-1;} 211 | bottom_right_y = pix_coord_2_y + extra_border; 212 | if (bottom_right_y >= cur_image_orig.rows) {bottom_right_y = cur_image_orig.rows-1;} 213 | } 214 | // max_corner is top right, min_corner is bottom left 215 | else { 216 | // cv::Point p1(pix_coord_1_x - extra_border, pix_coord_2_y - extra_border); 217 | // cv::Point p2(pix_coord_2_x + extra_border, pix_coord_1_y + extra_border); 218 | // cv::rectangle(occlusion_mask, p1, p2, cv::Scalar(0, 0, 0), -1); 219 | top_left_x = pix_coord_1_x - extra_border; 220 | if (top_left_x < 0) {top_left_x = 0;} 221 | top_left_y = pix_coord_2_y - extra_border; 222 | if (top_left_y < 0) {top_left_y = 0;} 223 | bottom_right_x = pix_coord_2_x + extra_border; 224 | if (bottom_right_x >= cur_image_orig.cols) {bottom_right_x = cur_image_orig.cols-1;} 225 | bottom_right_y = pix_coord_1_y + extra_border; 226 | if (bottom_right_y >= cur_image_orig.rows) {bottom_right_y = cur_image_orig.rows-1;} 227 | } 228 | 229 | std_msgs::Int32MultiArray corners_arr; 230 | corners_arr.data = {top_left_x, top_left_y, bottom_right_x, bottom_right_y}; 231 | corners_arr_pub.publish(corners_arr); 232 | } 233 | } 234 | 235 | else if (bag_file == 1) { 236 | top_left_x = 840; 237 | top_left_y = 408; 238 | bottom_right_x = 1191; 239 | bottom_right_y = 678; 240 | 241 | std_msgs::Int32MultiArray corners_arr; 242 | corners_arr.data = {top_left_x, top_left_y, bottom_right_x, bottom_right_y}; 243 | corners_arr_pub.publish(corners_arr); 244 | } 245 | 246 | else if (bag_file == 2) { 247 | top_left_x = 780; 248 | top_left_y = 120; 249 | bottom_right_x = 1050; 250 | bottom_right_y = 290; 251 | 252 | std_msgs::Int32MultiArray corners_arr; 253 | corners_arr.data = {top_left_x, top_left_y, bottom_right_x, bottom_right_y}; 254 | corners_arr_pub.publish(corners_arr); 255 | } 256 | 257 | else if (bag_file == 4) { 258 | top_left_x = 543; 259 | top_left_y = 276; 260 | bottom_right_x = 738; 261 | bottom_right_y = 383; 262 | 263 | std_msgs::Int32MultiArray corners_arr; 264 | corners_arr.data = {top_left_x, top_left_y, bottom_right_x, bottom_right_y}; 265 | corners_arr_pub.publish(corners_arr); 266 | } 267 | 268 | else if (bag_file == 5) { 269 | top_left_x = 300; 270 | top_left_y = 317; 271 | bottom_right_x = 698; 272 | bottom_right_y = 440; 273 | 274 | std_msgs::Int32MultiArray corners_arr; 275 | corners_arr.data = {top_left_x, top_left_y, bottom_right_x, bottom_right_y}; 276 | corners_arr_pub.publish(corners_arr); 277 | } 278 | 279 | else { 280 | throw std::invalid_argument("Invalid bag file ID!"); 281 | } 282 | } 283 | 284 | // compute error 285 | if (save_errors) { 286 | cur_error = tracking_evaluator.compute_and_save_error(Y_track, Y_true); 287 | } 288 | else { 289 | cur_error = tracking_evaluator.compute_error(Y_track, Y_true); 290 | } 291 | std::cout << "error = " << cur_error << std::endl; 292 | 293 | // optional pub and save result image 294 | if (time_from_start > tracking_evaluator.recording_start_time() + tracking_evaluator.wait_before_occlusion()) { 295 | cv::Point p1(top_left_x, top_left_y); 296 | cv::Point p2(bottom_right_x, bottom_right_y); 297 | cv::rectangle(eval_img, p1, p2, cv::Scalar(0, 0, 0), -1); 298 | eval_img = 0.5*eval_img + 0.5*cur_image_orig; 299 | 300 | if (bag_file == 4) { 301 | // cv::putText(tracking_img, "occlusion", cv::Point(occlusion_corner_j-190, occlusion_corner_i_2-5), cv::FONT_HERSHEY_DUPLEX, 1.2, cv::Scalar(0, 0, 240), 2); 302 | cv::putText(eval_img, "occlusion", cv::Point(top_left_x-190, bottom_right_y-5), cv::FONT_HERSHEY_DUPLEX, 1.2, cv::Scalar(0, 0, 240), 2); 303 | } 304 | else { 305 | cv::putText(eval_img, "occlusion", cv::Point(top_left_x, top_left_y-10), cv::FONT_HERSHEY_DUPLEX, 1.2, cv::Scalar(0, 0, 240), 2); 306 | } 307 | } 308 | else { 309 | cur_error = tracking_evaluator.compute_error(Y_track, Y_true); 310 | } 311 | } 312 | } 313 | 314 | // project tracking results onto the image 315 | MatrixXd Y_track_h = Y_track.replicate(1, 1); 316 | Y_track_h.conservativeResize(Y_track_h.rows(), Y_track_h.cols()+1); 317 | Y_track_h.col(Y_track_h.cols()-1) = MatrixXd::Ones(Y_track_h.rows(), 1); 318 | 319 | // project and pub image 320 | MatrixXd image_coords_Y = (proj_matrix * Y_track_h.transpose()).transpose(); 321 | // draw points 322 | for (int i = 0; i < image_coords_Y.rows(); i ++) { 323 | 324 | int row = static_cast(image_coords_Y(i, 0)/image_coords_Y(i, 2)); 325 | int col = static_cast(image_coords_Y(i, 1)/image_coords_Y(i, 2)); 326 | 327 | cv::Scalar point_color; 328 | cv::Scalar line_color; 329 | int line_width = 3; 330 | int point_radius = 7; 331 | // was 2 and 5 332 | 333 | if (row <= bottom_right_x && row >= top_left_x && col <= bottom_right_y && col >= top_left_y) { 334 | point_color = cv::Scalar(0, 0, 255); 335 | line_color = cv::Scalar(0, 0, 255); 336 | } 337 | else { 338 | point_color = cv::Scalar(0, 150, 255); 339 | line_color = cv::Scalar(0, 255, 0); 340 | } 341 | 342 | if (i != image_coords_Y.rows()-1) { 343 | cv::line(eval_img, cv::Point(row, col), 344 | cv::Point(static_cast(image_coords_Y(i+1, 0)/image_coords_Y(i+1, 2)), 345 | static_cast(image_coords_Y(i+1, 1)/image_coords_Y(i+1, 2))), 346 | line_color, line_width); 347 | } 348 | 349 | cv::circle(eval_img, cv::Point(row, col), point_radius, point_color, -1); 350 | } 351 | 352 | // if (cur_error != -1) { 353 | // std::string err = "Avg error per node: " + std::to_string(cur_error * 1000); 354 | // cv::putText(eval_img, err.substr(0, err.find(".")+3) + "mm", cv::Point(20, eval_img.rows - 20), cv::FONT_HERSHEY_DUPLEX, 1.2, cv::Scalar(0, 0, 0), 2); 355 | // } 356 | 357 | // save image 358 | if (save_images) { 359 | double diff = time_from_start - tracking_evaluator.recording_start_time(); 360 | double time_step = 0.5; 361 | if (bag_file == 0) { 362 | time_step = 1.0; 363 | } 364 | if ((int)(diff/time_step) == tracking_evaluator.image_counter() && fabs(diff-(tracking_evaluator.image_counter()*time_step)) <= 0.1) { 365 | std::string dir; 366 | // 0 -> statinary.bag; 1 -> with_gripper_perpendicular.bag; 2 -> with_gripper_parallel.bag 367 | if (bag_file == 0) { 368 | dir = save_location + "images/" + alg + "_" + std::to_string(trial) + "_" + std::to_string(pct_occlusion) + "_stationary_frame_" + std::to_string(tracking_evaluator.image_counter()) + ".png"; 369 | } 370 | else if (bag_file == 1) { 371 | dir = save_location + "images/" + alg + "_" + std::to_string(trial) + "_" + std::to_string(pct_occlusion) + "_perpendicular_motion_frame_" + std::to_string(tracking_evaluator.image_counter()) + ".png"; 372 | } 373 | else if (bag_file == 2) { 374 | dir = save_location + "images/" + alg + "_" + std::to_string(trial) + "_" + std::to_string(pct_occlusion) + "_parallel_motion_frame_" + std::to_string(tracking_evaluator.image_counter()) + ".png"; 375 | } 376 | else if (bag_file == 3) { 377 | dir = save_location + "images/" + alg + "_" + std::to_string(trial) + "_" + std::to_string(pct_occlusion) + "_self_occlusion_frame_" + std::to_string(tracking_evaluator.image_counter()) + ".png"; 378 | } 379 | else if (bag_file == 4) { 380 | dir = save_location + "images/" + alg + "_" + std::to_string(trial) + "_" + std::to_string(pct_occlusion) + "_short_rope_folding_frame_" + std::to_string(tracking_evaluator.image_counter()) + ".png"; 381 | } 382 | else if (bag_file == 5) { 383 | dir = save_location + "images/" + alg + "_" + std::to_string(trial) + "_" + std::to_string(pct_occlusion) + "_short_rope_stationary_frame_" + std::to_string(tracking_evaluator.image_counter()) + ".png"; 384 | } 385 | cv::imwrite(dir, eval_img); 386 | tracking_evaluator.increment_image_counter(); 387 | } 388 | } 389 | 390 | eval_img_msg = cv_bridge::CvImage(std_msgs::Header(), "bgr8", eval_img).toImageMsg(); 391 | return eval_img_msg; 392 | } 393 | 394 | int main(int argc, char **argv) { 395 | ros::init(argc, argv, "evaluation"); 396 | ros::NodeHandle nh; 397 | 398 | proj_matrix << 918.359130859375, 0.0, 645.8908081054688, 0.0, 399 | 0.0, 916.265869140625, 354.02392578125, 0.0, 400 | 0.0, 0.0, 1.0, 0.0; 401 | 402 | // load params 403 | nh.getParam("/evaluation/bag_file", bag_file); 404 | nh.getParam("/evaluation/trial", trial); 405 | nh.getParam("/evaluation/alg", alg); 406 | nh.getParam("/evaluation/bag_dir", bag_dir); 407 | nh.getParam("/evaluation/save_location", save_location); 408 | nh.getParam("/evaluation/pct_occlusion", pct_occlusion); 409 | nh.getParam("/evaluation/start_record_at", start_record_at); 410 | nh.getParam("/evaluation/exit_at", exit_at); 411 | nh.getParam("/evaluation/wait_before_occlusion", wait_before_occlusion); 412 | nh.getParam("/evaluation/bag_rate", bag_rate); 413 | nh.getParam("/evaluation/num_of_nodes", num_of_nodes); 414 | nh.getParam("/evaluation/save_images", save_images); 415 | nh.getParam("/evaluation/save_errors", save_errors); 416 | 417 | // get bag file length 418 | std::vector topics; 419 | topics.push_back("/camera/color/image_raw"); 420 | topics.push_back("/camera/aligned_depth_to_color/image_raw"); 421 | topics.push_back("/camera/depth/color/points"); 422 | topics.push_back("/tf"); 423 | topics.push_back("/tf_static"); 424 | 425 | rosbag::Bag bag(bag_dir, rosbag::bagmode::Read); 426 | rosbag::View view(bag, rosbag::TopicQuery(topics)); 427 | 428 | int rgb_count = 0; 429 | int depth_count = 0; 430 | int pc_count = 0; 431 | 432 | for (rosbag::MessageInstance const& msg: view) { 433 | if (msg.getTopic() == topics[0]) { 434 | rgb_count += 1; 435 | } 436 | if (msg.getTopic() == topics[1]) { 437 | depth_count += 1; 438 | } 439 | if (msg.getTopic() == topics[2]) { 440 | pc_count += 1; 441 | } 442 | } 443 | 444 | std::cout << "num of rgb images: " << rgb_count << std::endl; 445 | std::cout << "num of depth images: " << depth_count << std::endl; 446 | std::cout << "num of point cloud messages: " << pc_count << std::endl; 447 | 448 | // initialize evaluator 449 | tracking_evaluator = evaluator(rgb_count, trial, pct_occlusion, alg, bag_file, save_location, start_record_at, exit_at, wait_before_occlusion, bag_rate, num_of_nodes); 450 | 451 | image_transport::ImageTransport it(nh); 452 | image_transport::Publisher eval_img_pub = it.advertise("/eval_img", 10); 453 | corners_arr_pub = nh.advertise("/corners", 10); 454 | 455 | message_filters::Subscriber image_sub(nh, "/camera/color/image_raw", 10); 456 | message_filters::Subscriber pc_sub(nh, "/camera/depth/color/points", 10); 457 | message_filters::Subscriber result_sub(nh, "/" + alg + "/results_pc", 10); 458 | message_filters::TimeSynchronizer sync(image_sub, pc_sub, result_sub, 10); 459 | 460 | sync.registerCallback, 464 | const boost::shared_ptr, 465 | const boost::shared_ptr, 466 | const boost::shared_ptr, 467 | const boost::shared_ptr, 468 | const boost::shared_ptr)>> 469 | ( 470 | [&](const sensor_msgs::ImageConstPtr& img_msg, 471 | const sensor_msgs::PointCloud2ConstPtr& pc_msg, 472 | const sensor_msgs::PointCloud2ConstPtr& result_msg, 473 | const boost::shared_ptr var1, 474 | const boost::shared_ptr var2, 475 | const boost::shared_ptr var3, 476 | const boost::shared_ptr var4, 477 | const boost::shared_ptr var5, 478 | const boost::shared_ptr var6) 479 | { 480 | sensor_msgs::ImagePtr eval_img = Callback(img_msg, pc_msg, result_msg); 481 | eval_img_pub.publish(eval_img); 482 | } 483 | ); 484 | 485 | ros::spin(); 486 | } -------------------------------------------------------------------------------- /trackdlo/src/utils.cpp: -------------------------------------------------------------------------------- 1 | #include "../include/trackdlo.h" 2 | #include "../include/utils.h" 3 | 4 | using Eigen::MatrixXd; 5 | using Eigen::RowVectorXd; 6 | using cv::Mat; 7 | 8 | void signal_callback_handler(int signum) { 9 | // Terminate program 10 | exit(signum); 11 | } 12 | 13 | double pt2pt_dis_sq (MatrixXd pt1, MatrixXd pt2) { 14 | return (pt1 - pt2).rowwise().squaredNorm().sum(); 15 | } 16 | 17 | double pt2pt_dis (MatrixXd pt1, MatrixXd pt2) { 18 | return (pt1 - pt2).rowwise().norm().sum(); 19 | } 20 | 21 | void reg (MatrixXd pts, MatrixXd& Y, double& sigma2, int M, double mu, int max_iter) { 22 | // initial guess 23 | MatrixXd X = pts.replicate(1, 1); 24 | Y = MatrixXd::Zero(M, 3); 25 | for (int i = 0; i < M; i ++) { 26 | Y(i, 1) = 0.1 / static_cast(M) * static_cast(i); 27 | Y(i, 0) = 0; 28 | Y(i, 2) = 0; 29 | } 30 | 31 | int N = X.rows(); 32 | int D = 3; 33 | 34 | // diff_xy should be a (M * N) matrix 35 | MatrixXd diff_xy = MatrixXd::Zero(M, N); 36 | for (int i = 0; i < M; i ++) { 37 | for (int j = 0; j < N; j ++) { 38 | diff_xy(i, j) = (Y.row(i) - X.row(j)).squaredNorm(); 39 | } 40 | } 41 | 42 | // initialize sigma2 43 | sigma2 = diff_xy.sum() / static_cast(D * M * N); 44 | 45 | for (int it = 0; it < max_iter; it ++) { 46 | // update diff_xy 47 | for (int i = 0; i < M; i ++) { 48 | for (int j = 0; j < N; j ++) { 49 | diff_xy(i, j) = (Y.row(i) - X.row(j)).squaredNorm(); 50 | } 51 | } 52 | 53 | MatrixXd P = (-0.5 * diff_xy / sigma2).array().exp(); 54 | MatrixXd P_stored = P.replicate(1, 1); 55 | double c = pow((2 * M_PI * sigma2), static_cast(D)/2) * mu / (1 - mu) * static_cast(M)/N; 56 | P = P.array().rowwise() / (P.colwise().sum().array() + c); 57 | 58 | MatrixXd Pt1 = P.colwise().sum(); 59 | MatrixXd P1 = P.rowwise().sum(); 60 | double Np = P1.sum(); 61 | MatrixXd PX = P * X; 62 | 63 | MatrixXd P1_expanded = MatrixXd::Zero(M, D); 64 | P1_expanded.col(0) = P1; 65 | P1_expanded.col(1) = P1; 66 | P1_expanded.col(2) = P1; 67 | 68 | Y = PX.cwiseQuotient(P1_expanded); 69 | 70 | double numerator = 0; 71 | double denominator = 0; 72 | 73 | for (int m = 0; m < M; m ++) { 74 | for (int n = 0; n < N; n ++) { 75 | numerator += P(m, n)*diff_xy(m, n); 76 | denominator += P(m, n)*D; 77 | } 78 | } 79 | 80 | sigma2 = numerator / denominator; 81 | } 82 | } 83 | 84 | // link to original code: https://stackoverflow.com/a/46303314 85 | void remove_row(MatrixXd& matrix, unsigned int rowToRemove) { 86 | unsigned int numRows = matrix.rows()-1; 87 | unsigned int numCols = matrix.cols(); 88 | 89 | if( rowToRemove < numRows ) 90 | matrix.block(rowToRemove,0,numRows-rowToRemove,numCols) = matrix.bottomRows(numRows-rowToRemove); 91 | 92 | matrix.conservativeResize(numRows,numCols); 93 | } 94 | 95 | MatrixXd sort_pts (MatrixXd Y_0) { 96 | int N = Y_0.rows(); 97 | MatrixXd Y_0_sorted = MatrixXd::Zero(N, 3); 98 | std::vector Y_0_sorted_vec = {}; 99 | std::vector selected_node(N, false); 100 | selected_node[0] = true; 101 | int last_visited_b = 0; 102 | 103 | MatrixXd G = MatrixXd::Zero(N, N); 104 | for (int i = 0; i < N; i ++) { 105 | for (int j = 0; j < N; j ++) { 106 | G(i, j) = (Y_0.row(i) - Y_0.row(j)).squaredNorm(); 107 | } 108 | } 109 | 110 | int reverse = 0; 111 | int counter = 0; 112 | int reverse_on = 0; 113 | int insertion_counter = 0; 114 | 115 | while (counter < N-1) { 116 | double minimum = INFINITY; 117 | int a = 0; 118 | int b = 0; 119 | 120 | for (int m = 0; m < N; m ++) { 121 | if (selected_node[m] == true) { 122 | for (int n = 0; n < N; n ++) { 123 | if ((!selected_node[n]) && (G(m, n) != 0.0)) { 124 | if (minimum > G(m, n)) { 125 | minimum = G(m, n); 126 | a = m; 127 | b = n; 128 | } 129 | } 130 | } 131 | } 132 | } 133 | 134 | if (counter == 0) { 135 | Y_0_sorted_vec.push_back(Y_0.row(a)); 136 | Y_0_sorted_vec.push_back(Y_0.row(b)); 137 | } 138 | else { 139 | if (last_visited_b != a) { 140 | reverse += 1; 141 | reverse_on = a; 142 | insertion_counter = 1; 143 | } 144 | 145 | if (reverse % 2 == 1) { 146 | auto it = find(Y_0_sorted_vec.begin(), Y_0_sorted_vec.end(), Y_0.row(a)); 147 | Y_0_sorted_vec.insert(it, Y_0.row(b)); 148 | } 149 | else if (reverse != 0) { 150 | auto it = find(Y_0_sorted_vec.begin(), Y_0_sorted_vec.end(), Y_0.row(reverse_on)); 151 | Y_0_sorted_vec.insert(it + insertion_counter, Y_0.row(b)); 152 | insertion_counter += 1; 153 | } 154 | else { 155 | Y_0_sorted_vec.push_back(Y_0.row(b)); 156 | } 157 | } 158 | 159 | last_visited_b = b; 160 | selected_node[b] = true; 161 | counter += 1; 162 | } 163 | 164 | // copy to Y_0_sorted 165 | for (int i = 0; i < N; i ++) { 166 | Y_0_sorted.row(i) = Y_0_sorted_vec[i]; 167 | } 168 | 169 | return Y_0_sorted; 170 | } 171 | 172 | bool isBetween (MatrixXd x, MatrixXd a, MatrixXd b) { 173 | bool in_bound = true; 174 | 175 | for (int i = 0; i < 3; i ++) { 176 | if (!(a(0, i)-0.0001 <= x(0, i) && x(0, i) <= b(0, i)+0.0001) && 177 | !(b(0, i)-0.0001 <= x(0, i) && x(0, i) <= a(0, i)+0.0001)) { 178 | in_bound = false; 179 | } 180 | } 181 | 182 | return in_bound; 183 | } 184 | 185 | std::vector line_sphere_intersection (MatrixXd point_A, MatrixXd point_B, MatrixXd sphere_center, double radius) { 186 | std::vector intersections = {}; 187 | 188 | double a = pt2pt_dis_sq(point_A, point_B); 189 | double b = 2 * ((point_B(0, 0) - point_A(0, 0))*(point_A(0, 0) - sphere_center(0, 0)) + 190 | (point_B(0, 1) - point_A(0, 1))*(point_A(0, 1) - sphere_center(0, 1)) + 191 | (point_B(0, 2) - point_A(0, 2))*(point_A(0, 2) - sphere_center(0, 2))); 192 | double c = pt2pt_dis_sq(point_A, sphere_center) - pow(radius, 2); 193 | 194 | double delta = pow(b, 2) - 4*a*c; 195 | 196 | double d1 = (-b + sqrt(delta)) / (2*a); 197 | double d2 = (-b - sqrt(delta)) / (2*a); 198 | 199 | if (delta < 0) { 200 | // no solution 201 | return {}; 202 | } 203 | else if (delta > 0) { 204 | // two solutions 205 | // the first one 206 | double x1 = point_A(0, 0) + d1*(point_B(0, 0) - point_A(0, 0)); 207 | double y1 = point_A(0, 1) + d1*(point_B(0, 1) - point_A(0, 1)); 208 | double z1 = point_A(0, 2) + d1*(point_B(0, 2) - point_A(0, 2)); 209 | MatrixXd pt1(1, 3); 210 | pt1 << x1, y1, z1; 211 | 212 | // the second one 213 | double x2 = point_A(0, 0) + d2*(point_B(0, 0) - point_A(0, 0)); 214 | double y2 = point_A(0, 1) + d2*(point_B(0, 1) - point_A(0, 1)); 215 | double z2 = point_A(0, 2) + d2*(point_B(0, 2) - point_A(0, 2)); 216 | MatrixXd pt2(1, 3); 217 | pt2 << x2, y2, z2; 218 | 219 | if (isBetween(pt1, point_A, point_B)) { 220 | intersections.push_back(pt1); 221 | } 222 | if (isBetween(pt2, point_A, point_B)) { 223 | intersections.push_back(pt2); 224 | } 225 | } 226 | else { 227 | // one solution 228 | d1 = -b / (2*a); 229 | double x1 = point_A(0, 0) + d1*(point_B(0, 0) - point_A(0, 0)); 230 | double y1 = point_A(0, 1) + d1*(point_B(0, 1) - point_A(0, 1)); 231 | double z1 = point_A(0, 2) + d1*(point_B(0, 2) - point_A(0, 2)); 232 | MatrixXd pt1(1, 3); 233 | pt1 << x1, y1, z1; 234 | 235 | if (isBetween(pt1, point_A, point_B)) { 236 | intersections.push_back(pt1); 237 | } 238 | } 239 | 240 | return intersections; 241 | } 242 | 243 | // node color and object color are in rgba format and range from 0-1 244 | visualization_msgs::MarkerArray MatrixXd2MarkerArray (MatrixXd Y, 245 | std::string marker_frame, 246 | std::string marker_ns, 247 | std::vector node_color, 248 | std::vector line_color, 249 | double node_scale, 250 | double line_scale, 251 | std::vector visible_nodes, 252 | std::vector occluded_node_color, 253 | std::vector occluded_line_color) { // publish the results as a marker array 254 | 255 | visualization_msgs::MarkerArray results = visualization_msgs::MarkerArray(); 256 | 257 | bool last_node_visible = true; 258 | for (int i = 0; i < Y.rows(); i ++) { 259 | visualization_msgs::Marker cur_node_result = visualization_msgs::Marker(); 260 | 261 | // add header 262 | cur_node_result.header.frame_id = marker_frame; 263 | // cur_node_result.header.stamp = ros::Time::now(); 264 | cur_node_result.type = visualization_msgs::Marker::SPHERE; 265 | cur_node_result.action = visualization_msgs::Marker::ADD; 266 | cur_node_result.ns = marker_ns + "_node_" + std::to_string(i); 267 | cur_node_result.id = i; 268 | 269 | // add position 270 | cur_node_result.pose.position.x = Y(i, 0); 271 | cur_node_result.pose.position.y = Y(i, 1); 272 | cur_node_result.pose.position.z = Y(i, 2); 273 | 274 | // add orientation 275 | cur_node_result.pose.orientation.w = 1.0; 276 | cur_node_result.pose.orientation.x = 0.0; 277 | cur_node_result.pose.orientation.y = 0.0; 278 | cur_node_result.pose.orientation.z = 0.0; 279 | 280 | // set scale 281 | cur_node_result.scale.x = node_scale; 282 | cur_node_result.scale.y = node_scale; 283 | cur_node_result.scale.z = node_scale; 284 | 285 | // set color 286 | bool cur_node_visible; 287 | if (visible_nodes.size() != 0 && std::find(visible_nodes.begin(), visible_nodes.end(), i) == visible_nodes.end()) { 288 | cur_node_result.color.r = occluded_node_color[0]; 289 | cur_node_result.color.g = occluded_node_color[1]; 290 | cur_node_result.color.b = occluded_node_color[2]; 291 | cur_node_result.color.a = occluded_node_color[3]; 292 | cur_node_visible = false; 293 | } 294 | else { 295 | cur_node_result.color.r = node_color[0]; 296 | cur_node_result.color.g = node_color[1]; 297 | cur_node_result.color.b = node_color[2]; 298 | cur_node_result.color.a = node_color[3]; 299 | cur_node_visible = true; 300 | } 301 | 302 | results.markers.push_back(cur_node_result); 303 | 304 | // don't add line if at the first node 305 | if (i == 0) { 306 | continue; 307 | } 308 | 309 | visualization_msgs::Marker cur_line_result = visualization_msgs::Marker(); 310 | 311 | // add header 312 | cur_line_result.header.frame_id = marker_frame; 313 | cur_line_result.type = visualization_msgs::Marker::CYLINDER; 314 | cur_line_result.action = visualization_msgs::Marker::ADD; 315 | cur_line_result.ns = marker_ns + "_line_" + std::to_string(i); 316 | cur_line_result.id = i; 317 | 318 | // add position 319 | cur_line_result.pose.position.x = (Y(i, 0) + Y(i-1, 0)) / 2.0; 320 | cur_line_result.pose.position.y = (Y(i, 1) + Y(i-1, 1)) / 2.0; 321 | cur_line_result.pose.position.z = (Y(i, 2) + Y(i-1, 2)) / 2.0; 322 | 323 | // add orientation 324 | Eigen::Quaternionf q; 325 | Eigen::Vector3f vec1(0.0, 0.0, 1.0); 326 | Eigen::Vector3f vec2(Y(i, 0) - Y(i-1, 0), Y(i, 1) - Y(i-1, 1), Y(i, 2) - Y(i-1, 2)); 327 | q.setFromTwoVectors(vec1, vec2); 328 | 329 | cur_line_result.pose.orientation.w = q.w(); 330 | cur_line_result.pose.orientation.x = q.x(); 331 | cur_line_result.pose.orientation.y = q.y(); 332 | cur_line_result.pose.orientation.z = q.z(); 333 | 334 | // set scale 335 | cur_line_result.scale.x = line_scale; 336 | cur_line_result.scale.y = line_scale; 337 | cur_line_result.scale.z = pt2pt_dis(Y.row(i), Y.row(i-1)); 338 | 339 | // set color 340 | if (last_node_visible && cur_node_visible) { 341 | cur_line_result.color.r = line_color[0]; 342 | cur_line_result.color.g = line_color[1]; 343 | cur_line_result.color.b = line_color[2]; 344 | cur_line_result.color.a = line_color[3]; 345 | } 346 | else { 347 | cur_line_result.color.r = occluded_line_color[0]; 348 | cur_line_result.color.g = occluded_line_color[1]; 349 | cur_line_result.color.b = occluded_line_color[2]; 350 | cur_line_result.color.a = occluded_line_color[3]; 351 | } 352 | 353 | results.markers.push_back(cur_line_result); 354 | } 355 | 356 | return results; 357 | } 358 | 359 | // overload function 360 | visualization_msgs::MarkerArray MatrixXd2MarkerArray (std::vector Y, 361 | std::string marker_frame, 362 | std::string marker_ns, 363 | std::vector node_color, 364 | std::vector line_color, 365 | double node_scale, 366 | double line_scale, 367 | std::vector visible_nodes, 368 | std::vector occluded_node_color, 369 | std::vector occluded_line_color) { 370 | // publish the results as a marker array 371 | visualization_msgs::MarkerArray results = visualization_msgs::MarkerArray(); 372 | 373 | bool last_node_visible = true; 374 | for (int i = 0; i < Y.size(); i ++) { 375 | visualization_msgs::Marker cur_node_result = visualization_msgs::Marker(); 376 | 377 | int dim = Y[0].cols(); 378 | 379 | // add header 380 | cur_node_result.header.frame_id = marker_frame; 381 | // cur_node_result.header.stamp = ros::Time::now(); 382 | cur_node_result.type = visualization_msgs::Marker::SPHERE; 383 | cur_node_result.action = visualization_msgs::Marker::ADD; 384 | cur_node_result.ns = marker_ns + "_node_" + std::to_string(i); 385 | cur_node_result.id = i; 386 | 387 | // add position 388 | cur_node_result.pose.position.x = Y[i](0, dim-3); 389 | cur_node_result.pose.position.y = Y[i](0, dim-2); 390 | cur_node_result.pose.position.z = Y[i](0, dim-1); 391 | 392 | // add orientation 393 | cur_node_result.pose.orientation.w = 1.0; 394 | cur_node_result.pose.orientation.x = 0.0; 395 | cur_node_result.pose.orientation.y = 0.0; 396 | cur_node_result.pose.orientation.z = 0.0; 397 | 398 | // set scale 399 | cur_node_result.scale.x = 0.01; 400 | cur_node_result.scale.y = 0.01; 401 | cur_node_result.scale.z = 0.01; 402 | 403 | // set color 404 | bool cur_node_visible; 405 | if (visible_nodes.size() != 0 && std::find(visible_nodes.begin(), visible_nodes.end(), i) == visible_nodes.end()) { 406 | cur_node_result.color.r = occluded_node_color[0]; 407 | cur_node_result.color.g = occluded_node_color[1]; 408 | cur_node_result.color.b = occluded_node_color[2]; 409 | cur_node_result.color.a = occluded_node_color[3]; 410 | cur_node_visible = false; 411 | } 412 | else { 413 | cur_node_result.color.r = node_color[0]; 414 | cur_node_result.color.g = node_color[1]; 415 | cur_node_result.color.b = node_color[2]; 416 | cur_node_result.color.a = node_color[3]; 417 | cur_node_visible = true; 418 | } 419 | 420 | results.markers.push_back(cur_node_result); 421 | 422 | // don't add line if at the first node 423 | if (i == 0) { 424 | continue; 425 | } 426 | 427 | visualization_msgs::Marker cur_line_result = visualization_msgs::Marker(); 428 | 429 | // add header 430 | cur_line_result.header.frame_id = marker_frame; 431 | cur_line_result.type = visualization_msgs::Marker::CYLINDER; 432 | cur_line_result.action = visualization_msgs::Marker::ADD; 433 | cur_line_result.ns = marker_ns + "_line_" + std::to_string(i); 434 | cur_line_result.id = i; 435 | 436 | // add position 437 | cur_line_result.pose.position.x = (Y[i](0, dim-3) + Y[i-1](0, dim-3)) / 2.0; 438 | cur_line_result.pose.position.y = (Y[i](0, dim-2) + Y[i-1](0, dim-2)) / 2.0; 439 | cur_line_result.pose.position.z = (Y[i](0, dim-1) + Y[i-1](0, dim-1)) / 2.0; 440 | 441 | // add orientation 442 | Eigen::Quaternionf q; 443 | Eigen::Vector3f vec1(0.0, 0.0, 1.0); 444 | Eigen::Vector3f vec2(Y[i](0, dim-3) - Y[i-1](0, dim-3), Y[i](0, dim-2) - Y[i-1](0, dim-2), Y[i](0, dim-1) - Y[i-1](0, dim-1)); 445 | q.setFromTwoVectors(vec1, vec2); 446 | 447 | cur_line_result.pose.orientation.w = q.w(); 448 | cur_line_result.pose.orientation.x = q.x(); 449 | cur_line_result.pose.orientation.y = q.y(); 450 | cur_line_result.pose.orientation.z = q.z(); 451 | 452 | // set scale 453 | cur_line_result.scale.x = 0.005; 454 | cur_line_result.scale.y = 0.005; 455 | cur_line_result.scale.z = sqrt(pow(Y[i](0, dim-3) - Y[i-1](0, dim-3), 2) + pow(Y[i](0, dim-2) - Y[i-1](0, dim-2), 2) + pow(Y[i](0, dim-1) - Y[i-1](0, dim-1), 2)); 456 | 457 | // set color 458 | if (last_node_visible && cur_node_visible) { 459 | cur_line_result.color.r = line_color[0]; 460 | cur_line_result.color.g = line_color[1]; 461 | cur_line_result.color.b = line_color[2]; 462 | cur_line_result.color.a = line_color[3]; 463 | } 464 | else { 465 | cur_line_result.color.r = occluded_line_color[0]; 466 | cur_line_result.color.g = occluded_line_color[1]; 467 | cur_line_result.color.b = occluded_line_color[2]; 468 | cur_line_result.color.a = occluded_line_color[3]; 469 | } 470 | 471 | results.markers.push_back(cur_line_result); 472 | } 473 | 474 | return results; 475 | } 476 | 477 | MatrixXd cross_product (MatrixXd vec1, MatrixXd vec2) { 478 | MatrixXd ret = MatrixXd::Zero(1, 3); 479 | 480 | ret(0, 0) = vec1(0, 1)*vec2(0, 2) - vec1(0, 2)*vec2(0, 1); 481 | ret(0, 1) = -(vec1(0, 0)*vec2(0, 2) - vec1(0, 2)*vec2(0, 0)); 482 | ret(0, 2) = vec1(0, 0)*vec2(0, 1) - vec1(0, 1)*vec2(0, 0); 483 | 484 | return ret; 485 | } 486 | 487 | double dot_product (MatrixXd vec1, MatrixXd vec2) { 488 | return vec1(0, 0)*vec2(0, 0) + vec1(0, 1)*vec2(0, 1) + vec1(0, 2)*vec2(0, 2); 489 | } -------------------------------------------------------------------------------- /trackdlo/src/utils.py: -------------------------------------------------------------------------------- 1 | from skimage.morphology import skeletonize 2 | from scipy.optimize import linear_sum_assignment 3 | import cv2 4 | import numpy as np 5 | from PIL import Image, ImageFilter 6 | 7 | from visualization_msgs.msg import Marker 8 | from visualization_msgs.msg import MarkerArray 9 | from scipy.spatial.transform import Rotation as R 10 | 11 | def pt2pt_dis_sq(pt1, pt2): 12 | return np.sum(np.square(pt1 - pt2)) 13 | 14 | def pt2pt_dis(pt1, pt2): 15 | return np.linalg.norm(pt1-pt2) 16 | 17 | # from geeksforgeeks: https://www.geeksforgeeks.org/check-if-two-given-line-segments-intersect/ 18 | class Point_2D: 19 | def __init__(self, x, y): 20 | self.x = x 21 | self.y = y 22 | 23 | # from geeksforgeeks: https://www.geeksforgeeks.org/check-if-two-given-line-segments-intersect/ 24 | # Given three collinear points p, q, r, the function checks if 25 | # point q lies on line segment 'pr' 26 | def onSegment(p, q, r): 27 | if ( (q.x <= max(p.x, r.x)) and (q.x >= min(p.x, r.x)) and 28 | (q.y <= max(p.y, r.y)) and (q.y >= min(p.y, r.y))): 29 | return True 30 | return False 31 | 32 | # from geeksforgeeks: https://www.geeksforgeeks.org/check-if-two-given-line-segments-intersect/ 33 | def orientation(p, q, r): 34 | # to find the orientation of an ordered triplet (p,q,r) 35 | # function returns the following values: 36 | # 0 : Collinear points 37 | # 1 : Clockwise points 38 | # 2 : Counterclockwise 39 | 40 | # See https://www.geeksforgeeks.org/orientation-3-ordered-points/amp/ 41 | # for details of below formula. 42 | 43 | val = (np.float64(q.y - p.y) * (r.x - q.x)) - (np.float64(q.x - p.x) * (r.y - q.y)) 44 | if (val > 0): 45 | 46 | # Clockwise orientation 47 | return 1 48 | elif (val < 0): 49 | 50 | # Counterclockwise orientation 51 | return 2 52 | else: 53 | 54 | # Collinear orientation 55 | return 0 56 | 57 | # from geeksforgeeks: https://www.geeksforgeeks.org/check-if-two-given-line-segments-intersect/ 58 | # The main function that returns true if 59 | # the line segment 'p1q1' and 'p2q2' intersect. 60 | def doIntersect(p1,q1,p2,q2): 61 | 62 | # Find the 4 orientations required for 63 | # the general and special cases 64 | o1 = orientation(p1, q1, p2) 65 | o2 = orientation(p1, q1, q2) 66 | o3 = orientation(p2, q2, p1) 67 | o4 = orientation(p2, q2, q1) 68 | 69 | # General case 70 | if ((o1 != o2) and (o3 != o4)): 71 | return True 72 | 73 | # Special Cases 74 | 75 | # p1 , q1 and p2 are collinear and p2 lies on segment p1q1 76 | if ((o1 == 0) and onSegment(p1, p2, q1)): 77 | return True 78 | 79 | # p1 , q1 and q2 are collinear and q2 lies on segment p1q1 80 | if ((o2 == 0) and onSegment(p1, q2, q1)): 81 | return True 82 | 83 | # p2 , q2 and p1 are collinear and p1 lies on segment p2q2 84 | if ((o3 == 0) and onSegment(p2, p1, q2)): 85 | return True 86 | 87 | # p2 , q2 and q1 are collinear and q1 lies on segment p2q2 88 | if ((o4 == 0) and onSegment(p2, q1, q2)): 89 | return True 90 | 91 | # If none of the cases 92 | return False 93 | 94 | def build_rect (pt1, pt2, width): 95 | # arctan2 (y, x) 96 | line_angle = np.arctan2(pt2.y - pt1.y, pt2.x - pt1.x) 97 | angle1 = line_angle + np.pi/2 98 | angle2 = line_angle - np.pi/2 99 | rect_pt1 = Point_2D(pt1.x + width/2.0*np.cos(angle1), pt1.y + width/2.0*np.sin(angle1)) 100 | rect_pt2 = Point_2D(pt1.x + width/2.0*np.cos(angle2), pt1.y + width/2.0*np.sin(angle2)) 101 | rect_pt3 = Point_2D(pt2.x + width/2.0*np.cos(angle1), pt2.y + width/2.0*np.sin(angle1)) 102 | rect_pt4 = Point_2D(pt2.x + width/2.0*np.cos(angle2), pt2.y + width/2.0*np.sin(angle2)) 103 | 104 | return [rect_pt1, rect_pt2, rect_pt4, rect_pt3] 105 | 106 | def check_rect_overlap (rect1, rect2): 107 | overlap = False 108 | for i in range (-1, 3): 109 | for j in range (-1, 3): 110 | if doIntersect(rect1[i], rect1[i+1], rect2[j], rect2[j+1]): 111 | overlap = True 112 | return overlap 113 | return overlap 114 | 115 | # mode: 116 | # 0: start + start 117 | # 1: start + end 118 | # 2: end + start 119 | # 3: end + end 120 | def compute_cost (chain1, chain2, w_e, w_c, mode): 121 | chain1 = np.array(chain1) 122 | chain2 = np.array(chain2) 123 | 124 | # start + start 125 | if mode == 0: 126 | # treat the second chain as needing to be reversed 127 | cost_euclidean = np.linalg.norm(chain1[0] - chain2[0]) 128 | cost_curvature_1 = np.arccos(np.dot(chain1[0] - chain2[0], chain1[1] - chain1[0]) / (np.linalg.norm(chain1[0] - chain1[1]) * cost_euclidean)) 129 | cost_curvature_2 = np.arccos(np.dot(chain1[0] - chain2[0], chain2[0] - chain2[1]) / (np.linalg.norm(chain2[0] - chain2[1]) * cost_euclidean)) 130 | total_cost = w_e * cost_euclidean + w_c * (np.abs(cost_curvature_1) + np.abs(cost_curvature_2)) / 2.0 131 | # start + end 132 | elif mode == 1: 133 | cost_euclidean = np.linalg.norm(chain1[0] - chain2[-1]) 134 | cost_curvature_1 = np.arccos(np.dot(chain1[0] - chain2[-1], chain1[1] - chain1[0]) / (np.linalg.norm(chain1[0] - chain1[1]) * cost_euclidean)) 135 | cost_curvature_2 = np.arccos(np.dot(chain1[0] - chain2[-1], chain2[-1] - chain2[-2]) / (np.linalg.norm(chain2[-1] - chain2[-2]) * cost_euclidean)) 136 | total_cost = w_e * cost_euclidean + w_c * (np.abs(cost_curvature_1) + np.abs(cost_curvature_2)) / 2.0 137 | # end + start 138 | elif mode == 2: 139 | cost_euclidean = np.linalg.norm(chain1[-1] - chain2[0]) 140 | cost_curvature_1 = np.arccos(np.dot(chain2[0] - chain1[-1], chain1[-1] - chain1[-2]) / (np.linalg.norm(chain1[-1] - chain1[-2]) * cost_euclidean)) 141 | cost_curvature_2 = np.arccos(np.dot(chain2[0] - chain1[-1], chain2[1] - chain2[0]) / (np.linalg.norm(chain2[0] - chain2[1]) * cost_euclidean)) 142 | total_cost = w_e * cost_euclidean + w_c * (np.abs(cost_curvature_1) + np.abs(cost_curvature_2)) / 2.0 143 | # end + end 144 | else: 145 | # treat the second chain as needing to be reversed 146 | cost_euclidean = np.linalg.norm(chain1[-1] - chain2[-1]) 147 | cost_curvature_1 = np.arccos(np.dot(chain2[-1] - chain1[-1], chain1[-1] - chain1[-2]) / (np.linalg.norm(chain1[-1] - chain1[-2]) * cost_euclidean)) 148 | cost_curvature_2 = np.arccos(np.dot(chain2[-1] - chain1[-1], chain2[-2] - chain2[-1]) / (np.linalg.norm(chain2[-1] - chain2[-2]) * cost_euclidean)) 149 | total_cost = w_e * cost_euclidean + w_c * (np.abs(cost_curvature_1) + np.abs(cost_curvature_2)) / 2.0 150 | 151 | if total_cost is np.nan or cost_curvature_1 is np.nan or cost_curvature_2 is np.nan: 152 | print('total cost is nan!') 153 | print('chain1 =', chain1) 154 | print('chain2 =', chain2) 155 | print('euclidean cost = {}, w_e*cost = {}; curvature cost = {}, w_c*cost = {}'.format(cost_euclidean, w_e*cost_euclidean, (np.abs(cost_curvature_1) + np.abs(cost_curvature_2)) / 2.0, w_c * (np.abs(cost_curvature_1) + np.abs(cost_curvature_2)) / 2.0)) 156 | return total_cost 157 | 158 | # partial implementation of paper "Deformable One-Dimensional Object Detection for Routing and Manipulation" 159 | # paper link: https://ieeexplore.ieee.org/abstract/document/9697357 160 | def extract_connected_skeleton (visualize_process, mask, img_scale=10, seg_length=3, max_curvature=30): # note: mask is one channel 161 | 162 | # smooth image 163 | im = Image.fromarray(mask) 164 | smoothed_im = im.filter(ImageFilter.ModeFilter(size=15)) 165 | mask = np.array(smoothed_im) 166 | 167 | # resize if necessary for better skeletonization performance 168 | mask = cv2.resize(mask, (int(mask.shape[1]/img_scale), int(mask.shape[0]/img_scale))) 169 | 170 | if visualize_process: 171 | cv2.imshow('init frame', mask) 172 | while True: 173 | key = cv2.waitKey(10) 174 | if key == 27: # escape 175 | cv2.destroyAllWindows() 176 | break 177 | 178 | # skeletonization 179 | result = skeletonize(mask, method='zha') 180 | gray = cv2.cvtColor(result.copy(), cv2.COLOR_BGR2GRAY) 181 | gray[gray > 100] = 255 182 | print('Finished skeletonization. Traversing skeleton contours...') 183 | 184 | if visualize_process: 185 | cv2.imshow('after skeletonization', gray) 186 | while True: 187 | key = cv2.waitKey(10) 188 | if key == 27: # escape 189 | cv2.destroyAllWindows() 190 | break 191 | 192 | # extract contour 193 | contours, _ = cv2.findContours(gray, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[-2:] 194 | 195 | chains = [] 196 | 197 | # for each object segment 198 | for a, contour in enumerate(contours): 199 | c_area = cv2.contourArea(contour) 200 | mask = np.zeros(gray.shape, np.uint8) 201 | # mask = result.copy() 202 | 203 | last_segment_dir = None 204 | chain = [] 205 | cur_seg_start_point = None 206 | 207 | for i, coord in enumerate(contour): 208 | 209 | # special case: if reached the end of current contour, append chain if not empty 210 | if i == len(contour)-1: 211 | if len(chain) != 0: 212 | chains.append(chain) 213 | break 214 | 215 | # graph for visualization 216 | mask = cv2.line(mask, tuple(contour[i][0]), tuple(contour[i+1][0]), 255, 1) 217 | 218 | # if haven't initialized, perform initialization 219 | if cur_seg_start_point is None: 220 | cur_seg_start_point = contour[i][0].copy() 221 | 222 | # keep traversing until reach segment length 223 | if np.sqrt((contour[i][0][0] - cur_seg_start_point[0])**2 + (contour[i][0][1] - cur_seg_start_point[1])**2) <= seg_length: 224 | continue 225 | 226 | cur_seg_end_point = contour[i][0].copy() 227 | 228 | # record segment direction 229 | cur_segment_dir = [cur_seg_end_point[0] - cur_seg_start_point[0], cur_seg_end_point[1] - cur_seg_start_point[1]] 230 | if last_segment_dir is None: 231 | last_segment_dir = cur_segment_dir.copy() 232 | 233 | # if direction hasn't changed much, add current segment to chain 234 | elif np.dot(np.array(cur_segment_dir), np.array(last_segment_dir)) / \ 235 | (np.sqrt(cur_segment_dir[0]**2 + cur_segment_dir[1]**2) * np.sqrt(last_segment_dir[0]**2 + last_segment_dir[1]**2)) >= np.cos(max_curvature/180*np.pi): 236 | # if chain is empty, append both start and end point 237 | if len(chain) == 0: 238 | chain.append(cur_seg_start_point.tolist()) 239 | chain.append(cur_seg_end_point.tolist()) 240 | # if chain is not empty, only append end point 241 | else: 242 | chain.append(cur_seg_end_point.tolist()) 243 | 244 | # next start point <- end point 245 | cur_seg_start_point = cur_seg_end_point.copy() 246 | 247 | # update last segment direction 248 | last_segment_dir = cur_segment_dir.copy() 249 | 250 | # if direction changed, start a new chain 251 | else: 252 | # append current chain to chains if chain is not empty 253 | if len(chain) != 0: 254 | chains.append(chain) 255 | 256 | # reinitialize all variables 257 | last_segment_dir = None 258 | chain = [] 259 | cur_seg_start_point = None 260 | 261 | print('Finished contour traversal. Pruning extracted chains...') 262 | 263 | if visualize_process: 264 | mask = np.zeros((gray.shape[0], gray.shape[1], 3), np.uint8) 265 | for chain in chains: 266 | color = (int(np.random.random()*200)+55, int(np.random.random()*200)+55, int(np.random.random()*200)+55) 267 | for i in range (0, len(chain)-1): 268 | mask = cv2.line(mask, chain[i], chain[i+1], color, 1) 269 | cv2.imshow('added all chains frame', mask) 270 | while True: 271 | key = cv2.waitKey(10) 272 | if key == 27: # escape 273 | cv2.destroyAllWindows() 274 | break 275 | 276 | # another pruning method 277 | all_chain_length = [] 278 | line_seg_to_rect_dict = {} 279 | rect_width = 3 280 | for chain in chains: 281 | all_chain_length.append(np.sum(np.sqrt(np.sum(np.square(np.diff(np.array(chain), axis=0)), axis=1)))) 282 | for i in range (0, len(chain)-1): 283 | line_seg_to_rect_dict[(tuple(chain[i]), tuple(chain[i+1]))] = \ 284 | build_rect(Point_2D(chain[i][0], chain[i][1]), Point_2D(chain[i+1][0], chain[i+1][1]), rect_width) 285 | 286 | all_chain_length = np.array(all_chain_length) 287 | sorted_idx = np.argsort(all_chain_length.copy()) 288 | chains = np.asarray(chains, dtype=list) 289 | sorted_chains = chains[sorted_idx] 290 | 291 | pruned_chains = [] 292 | for i in range (0, len(chains)): 293 | leftover_chains = [] 294 | cur_chain = sorted_chains[-1] 295 | chains_to_check = [] 296 | 297 | for j in range (0, len(sorted_chains)-1): # -1 because the last one is cur_chain 298 | test_chain = sorted_chains[j] 299 | new_test_chain = [] 300 | for l in range (0, len(test_chain)-1): 301 | rect_test_seg = line_seg_to_rect_dict[(tuple(test_chain[l]), tuple(test_chain[l+1]))] 302 | # check against all segments in the current chain 303 | no_overlap = True 304 | for k in range (0, len(cur_chain)-1): 305 | rect_cur_seg = line_seg_to_rect_dict[(tuple(cur_chain[k]), tuple(cur_chain[k+1]))] 306 | if check_rect_overlap(rect_cur_seg, rect_test_seg): 307 | no_overlap = False 308 | break 309 | # only keep the segment in the test chain if it does not overlap with any segments in the current chain 310 | if no_overlap: 311 | if len(new_test_chain) == 0: 312 | new_test_chain.append(test_chain[l]) 313 | new_test_chain.append(test_chain[l+1]) 314 | else: 315 | new_test_chain.append(test_chain[l+1]) 316 | # finally, add trimmed test chain back into the pile for checking in the next loop 317 | leftover_chains.append(new_test_chain) 318 | 319 | # added current chain into pruned chains 320 | if len(cur_chain) != 0: 321 | pruned_chains.append(cur_chain) 322 | 323 | # reinitialize sorted chains 324 | all_chain_length = [] 325 | for chain in leftover_chains: 326 | if len(chain) != 0: 327 | all_chain_length.append(np.sum(np.sqrt(np.sum(np.square(np.diff(np.array(chain), axis=0)), axis=1)))) 328 | else: 329 | all_chain_length.append(0) 330 | 331 | all_chain_length = np.array(all_chain_length) 332 | sorted_idx = np.argsort(all_chain_length.copy()) 333 | leftover_chains = np.asarray(leftover_chains, dtype=list) 334 | sorted_chains = leftover_chains[sorted_idx] 335 | 336 | print('Finished pruning. Merging remaining chains...') 337 | 338 | if visualize_process: 339 | mask = np.zeros((gray.shape[0], gray.shape[1], 3), np.uint8) 340 | for chain in pruned_chains: 341 | color = (int(np.random.random()*200)+55, int(np.random.random()*200)+55, int(np.random.random()*200)+55) 342 | for i in range (0, len(chain)-1): 343 | mask = cv2.line(mask, chain[i], chain[i+1], color, 1) 344 | cv2.imshow("after pruning", mask) 345 | while True: 346 | key = cv2.waitKey(10) 347 | if key == 27: # escape 348 | cv2.destroyAllWindows() 349 | break 350 | 351 | if len(pruned_chains) == 1: 352 | return pruned_chains 353 | 354 | # total number of possible matches = number of ends * (number of ends - 2) / 2 = 2*len(chains) * (len(chains) - 1) 355 | # cost matrix size: num of tips + 2 356 | # cost matrix entry label format: tip1 start, tip1 end, tip2 start, tip2 end, ... 357 | matrix_size = 2*len(pruned_chains) + 2 358 | cost_matrix = np.zeros((matrix_size, matrix_size)) 359 | w_e = 0.001 360 | w_c = 1 361 | for i in range (0, len(pruned_chains)): 362 | for j in range (0, len(pruned_chains)): 363 | # two types of matches should be discouraged: 364 | # matching with itself and matching with the other tip on the same segment 365 | # tj_start tj_end 366 | # ti_start ... ... 367 | # ti_end ... ... 368 | if i == j: 369 | cost_matrix[2*i, 2*j] = 100000 370 | cost_matrix[2*i, 2*j+1] = 100000 371 | cost_matrix[2*i+1, 2*j] = 100000 372 | cost_matrix[2*i+1, 2*j+1] = 100000 373 | else: 374 | # ti_start to tj_start 375 | cost_matrix[2*i, 2*j] = compute_cost(pruned_chains[i], pruned_chains[j], w_e, w_c, 0) 376 | # ti_start to tj_end 377 | cost_matrix[2*i, 2*j+1] = compute_cost(pruned_chains[i], pruned_chains[j], w_e, w_c, 1) 378 | # ti_end to tj_start 379 | cost_matrix[2*i+1, 2*j] = compute_cost(pruned_chains[i], pruned_chains[j], w_e, w_c, 2) 380 | # ti_end to ti_end 381 | cost_matrix[2*i+1, 2*j+1] = compute_cost(pruned_chains[i], pruned_chains[j], w_e, w_c, 3) 382 | 383 | # cost for being the dlo's two ends 384 | cost_matrix[:, -1] = 1000 385 | cost_matrix[:, -2] = 1000 386 | cost_matrix[-1, :] = 1000 387 | cost_matrix[-2, :] = 1000 388 | # prevent matching with itself 389 | cost_matrix[matrix_size-2:matrix_size, matrix_size-2:matrix_size] = 100000 390 | 391 | row_idx, col_idx = linear_sum_assignment(cost_matrix) 392 | cur_idx = col_idx[row_idx[-1]] 393 | ordered_chains = [] 394 | 395 | mask = np.zeros((gray.shape[0], gray.shape[1], 3), np.uint8) 396 | while True: 397 | cur_chain_idx = int(cur_idx/2) 398 | cur_chain = pruned_chains[cur_chain_idx] 399 | 400 | if cur_idx % 2 == 1: 401 | cur_chain.reverse() 402 | ordered_chains.append(cur_chain) 403 | 404 | if cur_idx % 2 == 0: 405 | next_idx = col_idx[cur_idx+1] # find where this chain's end is connecting to 406 | else: 407 | next_idx = col_idx[cur_idx-1] # find where this chain's start is connecting to 408 | 409 | # if visualize_process: 410 | # mask = np.zeros((gray.shape[0], gray.shape[1], 3), np.uint8) 411 | # color = (int(np.random.random()*200)+55, int(np.random.random()*200)+55, int(np.random.random()*200)+55) 412 | # for j in range (0, len(cur_chain)-1): 413 | # mask = cv2.line(mask, cur_chain[j], cur_chain[j+1], color, 1) 414 | # cv2.imshow('in merging', mask) 415 | # while True: 416 | # key = cv2.waitKey(10) 417 | # if key == 27: # escape 418 | # cv2.destroyAllWindows() 419 | # break 420 | 421 | # if reached the other end, stop 422 | if next_idx == matrix_size-1 or next_idx == matrix_size-2: 423 | break 424 | cur_idx = next_idx 425 | 426 | print('Finished merging.') 427 | 428 | # visualization code for debug 429 | if visualize_process: 430 | mask = np.zeros((gray.shape[0], gray.shape[1], 3), np.uint8) 431 | for i in range (0, len(ordered_chains)): 432 | chain = ordered_chains[i] 433 | color = (int(np.random.random()*200)+55, int(np.random.random()*200)+55, int(np.random.random()*200)+55) 434 | for j in range (0, len(chain)-1): 435 | mask = cv2.line(mask, chain[j], chain[j+1], color, 1) 436 | 437 | cv2.imshow('after merging', mask) 438 | while True: 439 | key = cv2.waitKey(10) 440 | if key == 27: # escape 441 | cv2.destroyAllWindows() 442 | break 443 | 444 | if i == len(ordered_chains)-1: 445 | break 446 | 447 | pt1 = ordered_chains[i][-1] 448 | pt2 = ordered_chains[i+1][0] 449 | mask = cv2.line(mask, pt1, pt2, (255, 255, 255), 2) 450 | mask = cv2.circle(mask, pt1, 3, (255, 255, 255)) 451 | mask = cv2.circle(mask, pt2, 3, (255, 255, 255)) 452 | 453 | return ordered_chains 454 | 455 | # original post: https://stackoverflow.com/a/59204638 456 | def rotation_matrix_from_vectors(vec1, vec2): 457 | """ Find the rotation matrix that aligns vec1 to vec2 458 | :param vec1: A 3d "source" vector 459 | :param vec2: A 3d "destination" vector 460 | :return mat: A transform matrix (3x3) which when applied to vec1, aligns it with vec2. 461 | """ 462 | a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3) 463 | v = np.cross(a, b) 464 | c = np.dot(a, b) 465 | s = np.linalg.norm(v) 466 | kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) 467 | rotation_matrix = np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2)) 468 | return rotation_matrix 469 | 470 | def ndarray2MarkerArray (Y, marker_frame, node_color, line_color): 471 | results = MarkerArray() 472 | for i in range (0, len(Y)): 473 | cur_node_result = Marker() 474 | cur_node_result.header.frame_id = marker_frame 475 | cur_node_result.type = Marker.SPHERE 476 | cur_node_result.action = Marker.ADD 477 | cur_node_result.ns = "node_results" + str(i) 478 | cur_node_result.id = i 479 | 480 | cur_node_result.pose.position.x = Y[i, 0] 481 | cur_node_result.pose.position.y = Y[i, 1] 482 | cur_node_result.pose.position.z = Y[i, 2] 483 | cur_node_result.pose.orientation.w = 1.0 484 | cur_node_result.pose.orientation.x = 0.0 485 | cur_node_result.pose.orientation.y = 0.0 486 | cur_node_result.pose.orientation.z = 0.0 487 | 488 | cur_node_result.scale.x = 0.01 489 | cur_node_result.scale.y = 0.01 490 | cur_node_result.scale.z = 0.01 491 | cur_node_result.color.r = node_color[0] 492 | cur_node_result.color.g = node_color[1] 493 | cur_node_result.color.b = node_color[2] 494 | cur_node_result.color.a = node_color[3] 495 | 496 | results.markers.append(cur_node_result) 497 | 498 | if i == len(Y)-1: 499 | break 500 | 501 | cur_line_result = Marker() 502 | cur_line_result.header.frame_id = marker_frame 503 | cur_line_result.type = Marker.CYLINDER 504 | cur_line_result.action = Marker.ADD 505 | cur_line_result.ns = "line_results" + str(i) 506 | cur_line_result.id = i 507 | 508 | cur_line_result.pose.position.x = ((Y[i] + Y[i+1])/2)[0] 509 | cur_line_result.pose.position.y = ((Y[i] + Y[i+1])/2)[1] 510 | cur_line_result.pose.position.z = ((Y[i] + Y[i+1])/2)[2] 511 | 512 | rot_matrix = rotation_matrix_from_vectors(np.array([0, 0, 1]), (Y[i+1]-Y[i])/pt2pt_dis(Y[i+1], Y[i])) 513 | r = R.from_matrix(rot_matrix) 514 | x = r.as_quat()[0] 515 | y = r.as_quat()[1] 516 | z = r.as_quat()[2] 517 | w = r.as_quat()[3] 518 | 519 | cur_line_result.pose.orientation.w = w 520 | cur_line_result.pose.orientation.x = x 521 | cur_line_result.pose.orientation.y = y 522 | cur_line_result.pose.orientation.z = z 523 | cur_line_result.scale.x = 0.005 524 | cur_line_result.scale.y = 0.005 525 | cur_line_result.scale.z = pt2pt_dis(Y[i], Y[i+1]) 526 | cur_line_result.color.r = line_color[0] 527 | cur_line_result.color.g = line_color[1] 528 | cur_line_result.color.b = line_color[2] 529 | cur_line_result.color.a = line_color[3] 530 | 531 | results.markers.append(cur_line_result) 532 | 533 | return results -------------------------------------------------------------------------------- /utils/collect_pointcloud.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import rospy 4 | import ros_numpy 5 | from sensor_msgs.msg import PointCloud2, PointField, Image 6 | import sensor_msgs.point_cloud2 as pcl2 7 | import std_msgs.msg 8 | from sensor_msgs.msg import JointState 9 | 10 | import math 11 | import struct 12 | import time 13 | import cv2 14 | import numpy as np 15 | import pyrealsense2 as rs 16 | 17 | import time 18 | import tf 19 | import csv 20 | import pickle as pkl 21 | from scipy.spatial.transform import Rotation as R 22 | import os 23 | from os.path import dirname, abspath, join 24 | 25 | cur_pc = [] 26 | cur_image_arr = [] 27 | cur_result = [] 28 | cur_tracking_image_arr = [] 29 | 30 | # this gives pcl in the camera's frame 31 | def update_cur_pc(data): 32 | global cur_pc 33 | pc_arr = ros_numpy.point_cloud2.pointcloud2_to_array(data) 34 | cur_pc = ros_numpy.point_cloud2.get_xyz_points(pc_arr) 35 | 36 | def update_cur_result(data): 37 | global cur_result 38 | result_arr = ros_numpy.point_cloud2.pointcloud2_to_array(data) 39 | cur_result = ros_numpy.point_cloud2.get_xyz_points(result_arr) 40 | 41 | def update_img(data): 42 | global cur_image_arr 43 | cur_image = ros_numpy.numpify(data) 44 | cur_image_arr = cv2.cvtColor(cur_image, cv2.COLOR_BGR2RGB) 45 | 46 | def update_tracking_img(data): 47 | global cur_tracking_image_arr 48 | cur_tracking_image_arr = ros_numpy.numpify(data) 49 | # cur_tracking_image_arr = cv2.cvtColor(cur_tracking_image, cv2.COLOR_BGR2RGB) 50 | 51 | def record(main_dir, start=0, save_image=False, save_results=False): 52 | 53 | i = start 54 | 55 | while not rospy.is_shutdown(): 56 | sample_id = '' 57 | if len(str(i)) == 1: 58 | sample_id = '00' + str(i) 59 | else: 60 | sample_id = '0' + str(i) 61 | 62 | print("======================================================================") 63 | print("Press enter to collect and save point cloud and camera pose data") 64 | print("Press q + enter to exit the program") 65 | key_pressed = input("sample_id = " + sample_id + "\n") 66 | 67 | if key_pressed == 'q': 68 | print("Shutting down... \n") 69 | rospy.signal_shutdown('') 70 | else: 71 | rospy.Subscriber("/camera/color/image_raw", Image, update_img) 72 | rospy.Subscriber("/tracking_img", Image, update_tracking_img) 73 | rospy.Subscriber("/trackdlo_results_pc", PointCloud2, update_cur_result) 74 | rospy.Subscriber("/camera/depth/color/points", PointCloud2, update_cur_pc) 75 | 76 | if save_image: 77 | if len(cur_image_arr) == 0: 78 | print(" ") 79 | print("Could not capture image, please try again! \n") 80 | continue 81 | cv2.imwrite(main_dir + sample_id + "_rgb.png", cur_image_arr) 82 | 83 | if save_results: 84 | if len(cur_result) == 0: 85 | print(" ") 86 | print("Could not capture results, please try again! \n") 87 | continue 88 | 89 | f = open(main_dir + sample_id + "_results.json", "wb") 90 | pkl.dump(cur_result, f) 91 | f.close() 92 | 93 | if len(cur_tracking_image_arr) == 0: 94 | print(" ") 95 | print("Could not capture tracking image, please try again! \n") 96 | continue 97 | cv2.imwrite(main_dir + sample_id + "_result.png", cur_tracking_image_arr) 98 | 99 | # sometimes pointcloud can be empty 100 | if len(cur_pc) == 0: 101 | print(" ") 102 | print("Could not capture point cloud, please try again! \n") 103 | continue 104 | 105 | f = open(main_dir + sample_id + "_pc.json", "wb") 106 | pkl.dump(cur_pc, f) 107 | f.close() 108 | 109 | print("Data saved successfully! \n") 110 | i += 1 111 | 112 | if __name__ == '__main__': 113 | rospy.init_node('record_data', anonymous=True) 114 | main_dir = setting_path = join(dirname(dirname(abspath(__file__))), "data/") 115 | 116 | # try: 117 | os.listdir(main_dir) 118 | print("######################################################################") 119 | print("Collected data will be saved at the following directory:") 120 | print(main_dir) 121 | print("###################################################################### \n") 122 | 123 | record(main_dir, start=0, save_image=True, save_results=True) 124 | # except: 125 | # print("Invalid directory!") 126 | # rospy.signal_shutdown('') -------------------------------------------------------------------------------- /utils/color_picker.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import argparse 3 | import numpy as np 4 | 5 | def nothing(x): 6 | pass 7 | 8 | 9 | parser = argparse.ArgumentParser(description="Pick HSV color threshold") 10 | parser.add_argument("path", help="Indicate /full-or-relative/path/to/image_file.png") 11 | args = parser.parse_args() 12 | img_path = f"{args.path}" 13 | 14 | # Create a window 15 | cv2.namedWindow('image') 16 | 17 | # create trackbars for color change 18 | cv2.createTrackbar('HMin','image',0,179,nothing) # Hue is from 0-179 for Opencv 19 | cv2.createTrackbar('SMin','image',0,255,nothing) 20 | cv2.createTrackbar('VMin','image',0,255,nothing) 21 | cv2.createTrackbar('HMax','image',0,179,nothing) 22 | cv2.createTrackbar('SMax','image',0,255,nothing) 23 | cv2.createTrackbar('VMax','image',0,255,nothing) 24 | 25 | # Set default value for MAX HSV trackbars. 26 | cv2.setTrackbarPos('HMax', 'image', 179) 27 | cv2.setTrackbarPos('SMax', 'image', 255) 28 | cv2.setTrackbarPos('VMax', 'image', 255) 29 | 30 | # Initialize to check if HSV min/max value changes 31 | hMin = sMin = vMin = hMax = sMax = vMax = 0 32 | phMin = psMin = pvMin = phMax = psMax = pvMax = 0 33 | 34 | img = cv2.imread(img_path) 35 | img = cv2.resize(img, (640, 480)) 36 | output = img 37 | waitTime = 33 38 | 39 | while(1): 40 | 41 | # get current positions of all trackbars 42 | hMin = cv2.getTrackbarPos('HMin','image') 43 | sMin = cv2.getTrackbarPos('SMin','image') 44 | vMin = cv2.getTrackbarPos('VMin','image') 45 | 46 | hMax = cv2.getTrackbarPos('HMax','image') 47 | sMax = cv2.getTrackbarPos('SMax','image') 48 | vMax = cv2.getTrackbarPos('VMax','image') 49 | 50 | # Set minimum and max HSV values to display 51 | lower = np.array([hMin, sMin, vMin]) 52 | upper = np.array([hMax, sMax, vMax]) 53 | 54 | # Create HSV Image and threshold into a range. 55 | hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV) 56 | mask = cv2.inRange(hsv, lower, upper) 57 | output = cv2.bitwise_and(img,img, mask= mask) 58 | 59 | # Print if there is a change in HSV value 60 | if( (phMin != hMin) | (psMin != sMin) | (pvMin != vMin) | (phMax != hMax) | (psMax != sMax) | (pvMax != vMax) ): 61 | print("(hMin = %d , sMin = %d, vMin = %d), (hMax = %d , sMax = %d, vMax = %d)" % (hMin , sMin , vMin, hMax, sMax , vMax)) 62 | phMin = hMin 63 | psMin = sMin 64 | pvMin = vMin 65 | phMax = hMax 66 | psMax = sMax 67 | pvMax = vMax 68 | 69 | # Display output image 70 | cv2.imshow('image',output) 71 | 72 | # Wait longer to prevent freeze for videos. 73 | if cv2.waitKey(waitTime) & 0xFF == ord('q'): 74 | break 75 | 76 | cv2.destroyAllWindows() 77 | -------------------------------------------------------------------------------- /utils/mask.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import rospy 4 | import ros_numpy 5 | from sensor_msgs.msg import PointCloud2, PointField, Image 6 | import std_msgs.msg 7 | import cv2 8 | import numpy as np 9 | import message_filters 10 | 11 | def callback (rgb, pc): 12 | 13 | # process rgb image 14 | cur_image = ros_numpy.numpify(rgb) 15 | # cur_image = cv2.cvtColor(cur_image.copy(), cv2.COLOR_BGR2RGB) 16 | hsv_image = cv2.cvtColor(cur_image.copy(), cv2.COLOR_RGB2HSV) 17 | 18 | # color thresholding 19 | 20 | # rope blue 21 | # lower = (100, 120, 30) 22 | # upper = (130, 255, 255) 23 | # mask = cv2.inRange(hsv_image, lower, upper) 24 | 25 | # latex blue 26 | lower = (100, 230, 60) 27 | upper = (130, 255, 255) 28 | mask = cv2.inRange(hsv_image, lower, upper) 29 | 30 | # green 31 | # lower = (60, 130, 60) 32 | # upper = (95, 255, 255) 33 | # mask_green = cv2.inRange(hsv_image, lower, upper) 34 | 35 | # mask = cv2.bitwise_or(mask_blue, mask_green) 36 | mask = cv2.cvtColor(mask.copy(), cv2.COLOR_GRAY2BGR) 37 | 38 | # publish mask 39 | mask_img_msg = ros_numpy.msgify(Image, mask, 'rgb8') 40 | mask_img_pub.publish(mask_img_msg) 41 | 42 | 43 | if __name__=='__main__': 44 | rospy.init_node('test', anonymous=True) 45 | 46 | rgb_sub = message_filters.Subscriber('/camera/color/image_raw', Image) 47 | pc_sub = message_filters.Subscriber('/camera/depth/color/points', PointCloud2) 48 | 49 | # header 50 | header = std_msgs.msg.Header() 51 | header.stamp = rospy.Time.now() 52 | header.frame_id = 'camera_color_optical_frame' 53 | fields = [PointField('x', 0, PointField.FLOAT32, 1), 54 | PointField('y', 4, PointField.FLOAT32, 1), 55 | PointField('z', 8, PointField.FLOAT32, 1), 56 | PointField('rgba', 12, PointField.UINT32, 1)] 57 | pc_pub = rospy.Publisher ('/trackdlo/filtered_pointcloud', PointCloud2, queue_size=10) 58 | mask_img_pub = rospy.Publisher('/trackdlo/results_img', Image, queue_size=10) 59 | 60 | ts = message_filters.TimeSynchronizer([rgb_sub, pc_sub], 10) 61 | ts.registerCallback(callback) 62 | 63 | rospy.spin() -------------------------------------------------------------------------------- /utils/simulate_occlusion.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import cv2 4 | import numpy as np 5 | import rospy 6 | import ros_numpy 7 | from sensor_msgs.msg import Image 8 | 9 | class OcclusionSimulation: 10 | def __init__(self): 11 | 12 | self.rect = (0,0,0,0) 13 | self.startPoint = False 14 | self.endPoint = False 15 | self.start_moving = False 16 | self.rect_center = None 17 | self.offsets = None 18 | self.resting = False 19 | 20 | # update the mask each iteration 21 | self.mouse_mask = None 22 | 23 | self.rgb_sub = rospy.Subscriber('/camera/color/image_raw', Image, self.callback) 24 | self.occlusion_mask_img_pub = rospy.Publisher('/mask_with_occlusion', Image, queue_size=100) 25 | 26 | def callback(self,rgb): 27 | cur_image = ros_numpy.numpify(rgb) 28 | 29 | # convert color for opencv display 30 | cur_image = cv2.cvtColor(cur_image.copy(), cv2.COLOR_BGR2RGB) 31 | 32 | # resize for smaller window 33 | height, width, layers = cur_image.shape 34 | new_h = int(height / 1.5) 35 | new_w = int(width / 1.5) 36 | frame = cv2.resize(cur_image, (new_w, new_h)) 37 | 38 | # initialize mask if none 39 | if self.mouse_mask is None: 40 | self.mouse_mask = np.ones(frame.shape) 41 | 42 | # filter with mask 43 | frame = (frame * np.clip(self.mouse_mask, 0.5, 1)).astype('uint8') 44 | 45 | cv2.namedWindow('frame') 46 | cv2.setMouseCallback('frame', self.on_mouse) 47 | 48 | key = cv2.waitKey(10) 49 | 50 | # print(np.array(self.rect)*1.5) 51 | 52 | if key == 114: # r 53 | # reset everyhting 54 | frame = cv2.resize(cur_image, (new_w, new_h)) 55 | self.startPoint = False 56 | self.endPoint = False 57 | self.start_moving = False 58 | self.mouse_mask = np.ones(frame.shape) 59 | cv2.imshow('frame', frame) 60 | elif self.start_moving == True and self.resting == False: 61 | # first reset 62 | self.mouse_mask = np.ones(frame.shape) 63 | self.mouse_mask[self.rect[1]:self.rect[3], self.rect[0]:self.rect[2], :] = 0 64 | cv2.imshow('frame', frame) 65 | else: 66 | # drawing rectangle 67 | if self.startPoint == True and self.endPoint != True: 68 | cv2.rectangle(frame, (self.rect[0], self.rect[1]), (self.rect[2], self.rect[3]), (0, 0, 255), 2) 69 | if self.rect[0] < self.rect[2] and self.rect[1] < self.rect[3]: 70 | frame = cv2.putText(frame, 'occlusion', (self.rect[0], self.rect[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 240), 2) 71 | elif self.rect[0] < self.rect[2] and self.rect[1] > self.rect[3]: 72 | frame = cv2.putText(frame, 'occlusion', (self.rect[0], self.rect[3]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 240), 2) 73 | elif self.rect[0] > self.rect[2] and self.rect[1] < self.rect[3]: 74 | frame = cv2.putText(frame, 'occlusion', (self.rect[2], self.rect[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 240), 2) 75 | else: 76 | frame = cv2.putText(frame, 'occlusion', (self.rect[2], self.rect[3]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 240), 2) 77 | 78 | # if another rectangle is drawn, update mask 79 | if self.startPoint == True and self.endPoint == True: 80 | if self.rect[1] > self.rect[3]: 81 | self.rect[1], self.rect[3] = self.rect[3], self.rect[1] 82 | 83 | if self.rect[0] > self.rect[2]: 84 | self.rect[0], self.rect[2] = self.rect[2], self.rect[0] 85 | 86 | self.mouse_mask[self.rect[1]:self.rect[3], self.rect[0]:self.rect[2], :] = 0 87 | 88 | if not self.mouse_mask.all() == 1: 89 | frame = cv2.putText(frame, 'occlusion', (self.rect[0], self.rect[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 240), 2) 90 | 91 | cv2.imshow('frame', frame) 92 | 93 | # publish mask 94 | occlusion_mask = (self.mouse_mask*255).astype('uint8') 95 | 96 | # resize back for pub 97 | occlusion_mask = cv2.resize(occlusion_mask, (width, height)) 98 | 99 | occlusion_mask_img_msg = ros_numpy.msgify(Image, occlusion_mask, 'rgb8') 100 | self.occlusion_mask_img_pub.publish(occlusion_mask_img_msg) 101 | 102 | def on_mouse(self,event, x, y, flags, params): 103 | 104 | # get mouse click 105 | if event == cv2.EVENT_LBUTTONDOWN: 106 | 107 | if self.startPoint == True and self.endPoint == True: 108 | self.startPoint = False 109 | self.endPoint = False 110 | self.rect = [0, 0, 0, 0] 111 | 112 | if self.startPoint == False: 113 | self.rect = [x, y, x, y] 114 | self.startPoint = True 115 | elif self.endPoint == False: 116 | self.rect = [self.rect[0], self.rect[1], x, y] 117 | self.endPoint = True 118 | 119 | # draw rectangle when mouse hovering 120 | elif event == cv2.EVENT_MOUSEMOVE and self.startPoint == True and self.endPoint == False: 121 | self.rect = [self.rect[0], self.rect[1], x, y] 122 | 123 | elif event == cv2.EVENT_MBUTTONDOWN and self.start_moving == False and np.sum(self.mouse_mask[y, x]) == 0: 124 | self.start_moving = True 125 | # record rect center 126 | self.rect_center = (x, y) 127 | # offsets: left, up, right, down 128 | self.offsets = (self.rect[0]-self.rect_center[0], self.rect[1]-self.rect_center[1], self.rect[2]-self.rect_center[0], self.rect[3]-self.rect_center[1]) 129 | 130 | elif event == cv2.EVENT_MOUSEMOVE and self.start_moving == True: 131 | self.rect = [x+self.offsets[0], y+self.offsets[1], x+self.offsets[2], y+self.offsets[3]] 132 | self.resting = False 133 | 134 | elif event == cv2.EVENT_MBUTTONDOWN and self.start_moving == True: 135 | self.start_moving = False 136 | 137 | elif not event == cv2.EVENT_MOUSEMOVE and self.start_moving == True: 138 | self.resting = True 139 | 140 | if __name__=='__main__': 141 | rospy.init_node("simulated_occlusion") 142 | t = OcclusionSimulation() 143 | try: 144 | rospy.spin() 145 | except: 146 | print("Shutting down") -------------------------------------------------------------------------------- /utils/simulate_occlusion_eval.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import rospy 4 | import ros_numpy 5 | from sensor_msgs.msg import Image 6 | from std_msgs.msg import Int32MultiArray 7 | import cv2 8 | import numpy as np 9 | 10 | def callback (arr_msg): 11 | arr = arr_msg.data 12 | 13 | mouse_mask = np.ones((720, 1280, 3)) 14 | mouse_mask[arr[1]:arr[3], arr[0]:arr[2], :] = 0 15 | 16 | # publish mask 17 | occlusion_mask = (mouse_mask*255).astype('uint8') 18 | 19 | occlusion_mask_img_msg = ros_numpy.msgify(Image, occlusion_mask, 'rgb8') 20 | occlusion_mask_img_pub.publish(occlusion_mask_img_msg) 21 | 22 | 23 | if __name__=='__main__': 24 | rospy.init_node('simulate_occlusion_eval', anonymous=True) 25 | 26 | arr_sub = rospy.Subscriber('/corners', Int32MultiArray, callback) 27 | occlusion_mask_img_pub = rospy.Publisher('/mask_with_occlusion', Image, queue_size=10) 28 | 29 | rospy.spin() -------------------------------------------------------------------------------- /utils/tracking_result_img_from_pointcloud_topic.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import matplotlib.pyplot as plt 4 | import rospy 5 | import ros_numpy 6 | from sensor_msgs.msg import PointCloud2, PointField, Image 7 | import sensor_msgs.point_cloud2 as pcl2 8 | import std_msgs.msg 9 | 10 | import struct 11 | import time 12 | import cv2 13 | import numpy as np 14 | import math 15 | 16 | import time 17 | import pickle as pkl 18 | 19 | import message_filters 20 | import open3d as o3d 21 | from scipy import ndimage 22 | from scipy import interpolate 23 | 24 | cur_image = [] 25 | cur_image_arr = [] 26 | def update_rgb (data): 27 | global cur_image 28 | global cur_image_arr 29 | temp = ros_numpy.numpify(data) 30 | if len(cur_image_arr) <= 3: 31 | cur_image_arr.append(temp) 32 | cur_image = cur_image_arr[0] 33 | else: 34 | cur_image_arr.append(temp) 35 | cur_image = cur_image_arr[0] 36 | cur_image_arr.pop(0) 37 | 38 | bmask = [] 39 | mask = [] 40 | def update_mask (data): 41 | global bmask 42 | global mask 43 | mask = ros_numpy.numpify(data) 44 | bmask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY) 45 | 46 | def callback (pc): 47 | global cur_image 48 | global bmask 49 | global mask 50 | 51 | proj_matrix = np.array([[918.359130859375, 0.0, 645.8908081054688, 0.0], \ 52 | [ 0.0, 916.265869140625, 354.02392578125, 0.0], \ 53 | [ 0.0, 0.0, 1.0, 0.0]]) 54 | 55 | # process point cloud 56 | pc_data = ros_numpy.point_cloud2.pointcloud2_to_array(pc) 57 | result_pc = ros_numpy.point_cloud2.get_xyz_points(pc_data) 58 | nodes = result_pc.copy() 59 | 60 | # determined which nodes are occluded from mask information 61 | mask_dis_threshold = 10 62 | # projection 63 | nodes_h = np.hstack((nodes, np.ones((len(nodes), 1)))) 64 | # proj_matrix: 3*4; nodes_h.T: 4*M; result: 3*M 65 | image_coords = np.matmul(proj_matrix, nodes_h.T).T 66 | us = (image_coords[:, 0] / image_coords[:, 2]).astype(int) 67 | vs = (image_coords[:, 1] / image_coords[:, 2]).astype(int) 68 | 69 | # limit the range of calculated image coordinates 70 | us = np.where(us >= 1280, 1279, us) 71 | vs = np.where(vs >= 720, 719, vs) 72 | 73 | uvs = np.vstack((vs, us)).T 74 | uvs_t = tuple(map(tuple, uvs.T)) 75 | 76 | # invert bmask for distance transform 77 | bmask_transformed = ndimage.distance_transform_edt(255 - bmask) 78 | vis = bmask_transformed[uvs_t] 79 | 80 | tracking_img = cur_image.copy() 81 | for i in range (len(image_coords)): 82 | # draw circle 83 | uv = (us[i], vs[i]) 84 | if vis[i] < mask_dis_threshold: 85 | cv2.circle(tracking_img, uv, 5, (255, 150, 0), -1) 86 | else: 87 | cv2.circle(tracking_img, uv, 5, (255, 0, 0), -1) 88 | 89 | # draw line 90 | if i != len(image_coords)-1: 91 | if vis[i] < mask_dis_threshold: 92 | cv2.line(tracking_img, uv, (us[i+1], vs[i+1]), (0, 255, 0), 2) 93 | else: 94 | cv2.line(tracking_img, uv, (us[i+1], vs[i+1]), (255, 0, 0), 2) 95 | 96 | tracking_img_msg = ros_numpy.msgify(Image, tracking_img, 'rgb8') 97 | tracking_img_pub.publish(tracking_img_msg) 98 | 99 | if __name__=='__main__': 100 | rospy.init_node('cdcpd_image', anonymous=True) 101 | 102 | rospy.Subscriber('/mask', Image, update_mask) 103 | rospy.Subscriber('/camera/color/image_raw', Image, update_rgb) 104 | rospy.Subscriber('/cdcpd2_no_gripper_results_pc', PointCloud2, callback) 105 | tracking_img_pub = rospy.Publisher ('/tracking_img', Image, queue_size=10) 106 | 107 | rospy.spin() -------------------------------------------------------------------------------- /utils/tracking_test.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import rospy 4 | import ros_numpy 5 | from sensor_msgs.msg import PointCloud2, PointField, Image 6 | import sensor_msgs.point_cloud2 as pcl2 7 | import std_msgs.msg 8 | 9 | import struct 10 | import time 11 | import cv2 12 | import numpy as np 13 | import time 14 | 15 | import message_filters 16 | import open3d as o3d 17 | from scipy import ndimage 18 | 19 | from visualization_msgs.msg import Marker 20 | from visualization_msgs.msg import MarkerArray 21 | from scipy.spatial.transform import Rotation as R 22 | 23 | proj_matrix = np.array([[918.359130859375, 0.0, 645.8908081054688, 0.0], \ 24 | [ 0.0, 916.265869140625, 354.02392578125, 0.0], \ 25 | [ 0.0, 0.0, 1.0, 0.0]]) 26 | 27 | def pt2pt_dis_sq(pt1, pt2): 28 | return np.sum(np.square(pt1 - pt2)) 29 | 30 | def pt2pt_dis(pt1, pt2): 31 | return np.sqrt(np.sum(np.square(pt1 - pt2))) 32 | 33 | occlusion_mask_rgb = None 34 | def update_occlusion_mask(data): 35 | global occlusion_mask_rgb 36 | occlusion_mask_rgb = ros_numpy.numpify(data) 37 | 38 | # original post: https://stackoverflow.com/a/59204638 39 | def rotation_matrix_from_vectors(vec1, vec2): 40 | """ Find the rotation matrix that aligns vec1 to vec2 41 | :param vec1: A 3d "source" vector 42 | :param vec2: A 3d "destination" vector 43 | :return mat: A transform matrix (3x3) which when applied to vec1, aligns it with vec2. 44 | """ 45 | a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 / np.linalg.norm(vec2)).reshape(3) 46 | v = np.cross(a, b) 47 | c = np.dot(a, b) 48 | s = np.linalg.norm(v) 49 | kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) 50 | rotation_matrix = np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2)) 51 | return rotation_matrix 52 | 53 | def ndarray2MarkerArray (Y, marker_frame, node_color, line_color): 54 | results = MarkerArray() 55 | for i in range (0, len(Y)): 56 | cur_node_result = Marker() 57 | cur_node_result.header.frame_id = marker_frame 58 | cur_node_result.type = Marker.SPHERE 59 | cur_node_result.action = Marker.ADD 60 | cur_node_result.ns = "node_results" + str(i) 61 | cur_node_result.id = i 62 | 63 | cur_node_result.pose.position.x = Y[i, 0] 64 | cur_node_result.pose.position.y = Y[i, 1] 65 | cur_node_result.pose.position.z = Y[i, 2] 66 | cur_node_result.pose.orientation.w = 1.0 67 | cur_node_result.pose.orientation.x = 0.0 68 | cur_node_result.pose.orientation.y = 0.0 69 | cur_node_result.pose.orientation.z = 0.0 70 | 71 | cur_node_result.scale.x = 0.01 72 | cur_node_result.scale.y = 0.01 73 | cur_node_result.scale.z = 0.01 74 | cur_node_result.color.r = node_color[0] 75 | cur_node_result.color.g = node_color[1] 76 | cur_node_result.color.b = node_color[2] 77 | cur_node_result.color.a = node_color[3] 78 | 79 | results.markers.append(cur_node_result) 80 | 81 | if i == len(Y)-1: 82 | break 83 | 84 | cur_line_result = Marker() 85 | cur_line_result.header.frame_id = marker_frame 86 | cur_line_result.type = Marker.CYLINDER 87 | cur_line_result.action = Marker.ADD 88 | cur_line_result.ns = "line_results" + str(i) 89 | cur_line_result.id = i 90 | 91 | cur_line_result.pose.position.x = ((Y[i] + Y[i+1])/2)[0] 92 | cur_line_result.pose.position.y = ((Y[i] + Y[i+1])/2)[1] 93 | cur_line_result.pose.position.z = ((Y[i] + Y[i+1])/2)[2] 94 | 95 | rot_matrix = rotation_matrix_from_vectors(np.array([0, 0, 1]), (Y[i+1]-Y[i])/pt2pt_dis(Y[i+1], Y[i])) 96 | r = R.from_matrix(rot_matrix) 97 | x = r.as_quat()[0] 98 | y = r.as_quat()[1] 99 | z = r.as_quat()[2] 100 | w = r.as_quat()[3] 101 | 102 | cur_line_result.pose.orientation.w = w 103 | cur_line_result.pose.orientation.x = x 104 | cur_line_result.pose.orientation.y = y 105 | cur_line_result.pose.orientation.z = z 106 | cur_line_result.scale.x = 0.005 107 | cur_line_result.scale.y = 0.005 108 | cur_line_result.scale.z = pt2pt_dis(Y[i], Y[i+1]) 109 | cur_line_result.color.r = line_color[0] 110 | cur_line_result.color.g = line_color[1] 111 | cur_line_result.color.b = line_color[2] 112 | cur_line_result.color.a = line_color[3] 113 | 114 | results.markers.append(cur_line_result) 115 | 116 | return results 117 | 118 | def register(pts, M, mu=0, max_iter=50): 119 | 120 | # initial guess 121 | X = pts.copy() 122 | Y = np.vstack((np.arange(0, 0.1, (0.1/M)), np.zeros(M), np.zeros(M))).T 123 | if len(pts[0]) == 2: 124 | Y = np.vstack((np.arange(0, 0.1, (0.1/M)), np.zeros(M))).T 125 | s = 1 126 | N = len(pts) 127 | D = len(pts[0]) 128 | 129 | def get_estimates (Y, s): 130 | 131 | # construct the P matrix 132 | P = np.sum((X[None, :, :] - Y[:, None, :]) ** 2, axis=2) 133 | 134 | c = (2 * np.pi * s) ** (D / 2) 135 | c = c * mu / (1 - mu) 136 | c = c * M / N 137 | 138 | P = np.exp(-P / (2 * s)) 139 | den = np.sum(P, axis=0) 140 | den = np.tile(den, (M, 1)) 141 | den[den == 0] = np.finfo(float).eps 142 | den += c 143 | 144 | P = np.divide(P, den) # P is M*N 145 | Pt1 = np.sum(P, axis=0) # equivalent to summing from 0 to M (results in N terms) 146 | P1 = np.sum(P, axis=1) # equivalent to summing from 0 to N (results in M terms) 147 | Np = np.sum(P1) 148 | PX = np.matmul(P, X) 149 | 150 | # get new Y 151 | P1_expanded = np.full((D, M), P1).T 152 | new_Y = PX / P1_expanded 153 | 154 | # get new sigma2 155 | Y_N_arr = np.full((N, M, 3), Y) 156 | Y_N_arr = np.swapaxes(Y_N_arr, 0, 1) 157 | X_M_arr = np.full((M, N, 3), X) 158 | diff = Y_N_arr - X_M_arr 159 | diff = np.square(diff) 160 | diff = np.sum(diff, 2) 161 | new_s = np.sum(np.sum(P*diff, axis=1), axis=0) / (Np*D) 162 | 163 | return new_Y, new_s 164 | 165 | prev_Y, prev_s = Y, s 166 | new_Y, new_s = get_estimates(prev_Y, prev_s) 167 | 168 | for it in range (max_iter): 169 | prev_Y, prev_s = new_Y, new_s 170 | new_Y, new_s = get_estimates(prev_Y, prev_s) 171 | 172 | return new_Y, new_s 173 | 174 | def sort_pts (Y_0): 175 | diff = Y_0[:, None, :] - Y_0[None, :, :] 176 | diff = np.square(diff) 177 | diff = np.sum(diff, 2) 178 | 179 | N = len(diff) 180 | G = diff.copy() 181 | 182 | selected_node = np.zeros(N,).tolist() 183 | selected_node[0] = True 184 | Y_0_sorted = [] 185 | 186 | reverse = 0 187 | counter = 0 188 | reverse_on = 0 189 | insertion_counter = 0 190 | last_visited_b = 0 191 | while (counter < N - 1): 192 | 193 | minimum = 999999 194 | a = 0 195 | b = 0 196 | for m in range(N): 197 | if selected_node[m]: 198 | for n in range(N): 199 | if ((not selected_node[n]) and G[m][n]): 200 | # not in selected and there is an edge 201 | if minimum > G[m][n]: 202 | minimum = G[m][n] 203 | a = m 204 | b = n 205 | 206 | if len(Y_0_sorted) == 0: 207 | Y_0_sorted.append(Y_0[a].tolist()) 208 | Y_0_sorted.append(Y_0[b].tolist()) 209 | else: 210 | if last_visited_b != a: 211 | reverse += 1 212 | reverse_on = a 213 | insertion_counter = 0 214 | 215 | if reverse % 2 == 1: 216 | # switch direction 217 | Y_0_sorted.insert(Y_0_sorted.index(Y_0[a].tolist()), Y_0[b].tolist()) 218 | elif reverse != 0: 219 | Y_0_sorted.insert(Y_0_sorted.index(Y_0[reverse_on].tolist())+1+insertion_counter, Y_0[b].tolist()) 220 | insertion_counter += 1 221 | else: 222 | Y_0_sorted.append(Y_0[b].tolist()) 223 | 224 | last_visited_b = b 225 | selected_node[b] = True 226 | 227 | counter += 1 228 | 229 | return np.array(Y_0_sorted) 230 | 231 | # assuming Y is sorted 232 | # k -- going left for k indices, going right for k indices. a total of 2k neighbors. 233 | def get_nearest_indices (k, Y, idx): 234 | if idx - k < 0: 235 | # use more neighbors from the other side? 236 | indices_arr = np.append(np.arange(0, idx, 1), np.arange(idx+1, idx+k+1+np.abs(idx-k))) 237 | # indices_arr = np.append(np.arange(0, idx, 1), np.arange(idx+1, idx+k+1)) 238 | return indices_arr 239 | elif idx + k >= len(Y): 240 | last_index = len(Y) - 1 241 | # use more neighbots from the other side? 242 | indices_arr = np.append(np.arange(idx-k-(idx+k-last_index), idx, 1), np.arange(idx+1, last_index+1, 1)) 243 | # indices_arr = np.append(np.arange(idx-k, idx, 1), np.arange(idx+1, last_index+1, 1)) 244 | return indices_arr 245 | else: 246 | indices_arr = np.append(np.arange(idx-k, idx, 1), np.arange(idx+1, idx+k+1, 1)) 247 | return indices_arr 248 | 249 | def calc_LLE_weights (k, X): 250 | W = np.zeros((len(X), len(X))) 251 | for i in range (0, len(X)): 252 | indices = get_nearest_indices(int(k/2), X, i) 253 | xi, Xi = X[i], X[indices, :] 254 | component = np.full((len(Xi), len(xi)), xi).T - Xi.T 255 | Gi = np.matmul(component.T, component) 256 | # Gi might be singular when k is large 257 | try: 258 | Gi_inv = np.linalg.inv(Gi) 259 | except: 260 | epsilon = 0.00001 261 | Gi_inv = np.linalg.inv(Gi + epsilon*np.identity(len(Gi))) 262 | wi = np.matmul(Gi_inv, np.ones((len(Xi), 1))) / np.matmul(np.matmul(np.ones(len(Xi),), Gi_inv), np.ones((len(Xi), 1))) 263 | W[i, indices] = np.squeeze(wi.T) 264 | 265 | return W 266 | 267 | def indices_array(n): 268 | r = np.arange(n) 269 | out = np.empty((n,n,2),dtype=int) 270 | out[:,:,0] = r[:,None] 271 | out[:,:,1] = r 272 | return out 273 | 274 | def cpd_lle (X, Y_0, beta, alpha, gamma, mu, max_iter=50, tol=0.00001, include_lle=True, use_geodesic=False, use_prev_sigma2=False, sigma2_0=None): 275 | 276 | # define params 277 | M = len(Y_0) 278 | N = len(X) 279 | D = len(X[0]) 280 | 281 | # initialization 282 | # faster G calculation 283 | diff = Y_0[:, None, :] - Y_0[None, :, :] 284 | diff = np.square(diff) 285 | diff = np.sum(diff, 2) 286 | 287 | converted_node_dis = [] 288 | if not use_geodesic: 289 | # Gaussian Kernel 290 | G = np.exp(-diff / (2 * beta**2)) 291 | else: 292 | # compute the geodesic distances between nodes 293 | seg_dis = np.sqrt(np.sum(np.square(np.diff(Y_0, axis=0)), axis=1)) 294 | converted_node_coord = [] 295 | last_pt = 0 296 | converted_node_coord.append(last_pt) 297 | for i in range (1, M): 298 | last_pt += seg_dis[i-1] 299 | converted_node_coord.append(last_pt) 300 | converted_node_coord = np.array(converted_node_coord) 301 | converted_node_dis = np.abs(converted_node_coord[None, :] - converted_node_coord[:, None]) 302 | converted_node_dis_sq = np.square(converted_node_dis) 303 | 304 | # Gaussian Kernel 305 | G = np.exp(-converted_node_dis_sq / (2 * beta**2)) 306 | 307 | # temp 308 | # G[converted_node_dis > 0.07] = 0 309 | 310 | Y = Y_0.copy() 311 | 312 | # initialize sigma2 313 | if not use_prev_sigma2: 314 | (N, D) = X.shape 315 | (M, _) = Y.shape 316 | diff = X[None, :, :] - Y[:, None, :] 317 | err = diff ** 2 318 | sigma2 = np.sum(err) / (D * M * N) 319 | else: 320 | sigma2 = sigma2_0 321 | 322 | # get the LLE matrix 323 | L = calc_LLE_weights(6, Y_0) 324 | H = np.matmul((np.identity(M) - L).T, np.identity(M) - L) 325 | 326 | # loop until convergence or max_iter reached 327 | for it in range (0, max_iter): 328 | 329 | # ----- E step: compute posteriori probability matrix P ----- 330 | # faster P computation 331 | pts_dis_sq = np.sum((X[None, :, :] - Y[:, None, :]) ** 2, axis=2) 332 | c = (2 * np.pi * sigma2) ** (D / 2) 333 | c = c * mu / (1 - mu) 334 | c = c * M / N 335 | P = np.exp(-pts_dis_sq / (2 * sigma2)) 336 | den = np.sum(P, axis=0) 337 | den = np.tile(den, (M, 1)) 338 | den[den == 0] = np.finfo(float).eps 339 | den += c 340 | P = np.divide(P, den) 341 | 342 | max_p_nodes = np.argmax(P, axis=0) 343 | 344 | # if use geodesic, overwrite P 345 | # this section looks long, but it is simply replacing the Euclidean distances in P with geodesic distances 346 | if use_geodesic: 347 | potential_2nd_max_p_nodes_1 = max_p_nodes - 1 348 | potential_2nd_max_p_nodes_2 = max_p_nodes + 1 349 | potential_2nd_max_p_nodes_1 = np.where(potential_2nd_max_p_nodes_1 < 0, 1, potential_2nd_max_p_nodes_1) 350 | potential_2nd_max_p_nodes_2 = np.where(potential_2nd_max_p_nodes_2 > M-1, M-2, potential_2nd_max_p_nodes_2) 351 | potential_2nd_max_p_nodes_1_select = np.vstack((np.arange(0, N), potential_2nd_max_p_nodes_1)).T 352 | potential_2nd_max_p_nodes_2_select = np.vstack((np.arange(0, N), potential_2nd_max_p_nodes_2)).T 353 | potential_2nd_max_p_1 = P.T[tuple(map(tuple, potential_2nd_max_p_nodes_1_select.T))] 354 | potential_2nd_max_p_2 = P.T[tuple(map(tuple, potential_2nd_max_p_nodes_2_select.T))] 355 | next_max_p_nodes = np.where(potential_2nd_max_p_1 > potential_2nd_max_p_2, potential_2nd_max_p_nodes_1, potential_2nd_max_p_nodes_2) 356 | node_indices_diff = max_p_nodes - next_max_p_nodes 357 | max_node_smaller_index = np.arange(0, N)[node_indices_diff < 0] 358 | max_node_larger_index = np.arange(0, N)[node_indices_diff > 0] 359 | dis_to_max_p_nodes = np.sqrt(np.sum(np.square(Y[max_p_nodes]-X), axis=1)) 360 | dis_to_2nd_largest_p_nodes = np.sqrt(np.sum(np.square(Y[next_max_p_nodes]-X), axis=1)) 361 | converted_P = np.zeros((M, N)).T 362 | 363 | for idx in max_node_smaller_index: 364 | converted_P[idx, 0:max_p_nodes[idx]+1] = converted_node_dis[max_p_nodes[idx], 0:max_p_nodes[idx]+1] + dis_to_max_p_nodes[idx] 365 | converted_P[idx, next_max_p_nodes[idx]:M] = converted_node_dis[next_max_p_nodes[idx], next_max_p_nodes[idx]:M] + dis_to_2nd_largest_p_nodes[idx] 366 | 367 | for idx in max_node_larger_index: 368 | converted_P[idx, 0:next_max_p_nodes[idx]+1] = converted_node_dis[next_max_p_nodes[idx], 0:next_max_p_nodes[idx]+1] + dis_to_2nd_largest_p_nodes[idx] 369 | converted_P[idx, max_p_nodes[idx]:M] = converted_node_dis[max_p_nodes[idx], max_p_nodes[idx]:M] + dis_to_max_p_nodes[idx] 370 | 371 | converted_P = converted_P.T 372 | 373 | P = np.exp(-np.square(converted_P) / (2 * sigma2)) 374 | den = np.sum(P, axis=0) 375 | den = np.tile(den, (M, 1)) 376 | den[den == 0] = np.finfo(float).eps 377 | c = (2 * np.pi * sigma2) ** (D / 2) 378 | c = c * mu / (1 - mu) 379 | c = c * M / N 380 | den += c 381 | 382 | P = np.divide(P, den) 383 | 384 | Pt1 = np.sum(P, axis=0) 385 | P1 = np.sum(P, axis=1) 386 | Np = np.sum(P1) 387 | PX = np.matmul(P, X) 388 | 389 | # print(Pt1) 390 | 391 | # ----- M step: solve for new weights and variance ----- 392 | if include_lle: 393 | A_matrix = np.matmul(np.diag(P1), G) + alpha * sigma2 * np.identity(M) + sigma2 * gamma * np.matmul(H, G) 394 | B_matrix = PX - np.matmul(np.diag(P1) + sigma2*gamma*H, Y_0) 395 | else: 396 | A_matrix = np.matmul(np.diag(P1), G) + alpha * sigma2 * np.identity(M) 397 | B_matrix = PX - np.matmul(np.diag(P1), Y_0) 398 | 399 | # solve for W 400 | W = np.linalg.solve(A_matrix, B_matrix) 401 | 402 | T = Y_0 + np.matmul(G, W) 403 | trXtdPt1X = np.trace(np.matmul(np.matmul(X.T, np.diag(Pt1)), X)) 404 | trPXtT = np.trace(np.matmul(PX.T, T)) 405 | trTtdP1T = np.trace(np.matmul(np.matmul(T.T, np.diag(P1)), T)) 406 | 407 | # solve for sigma^2 408 | sigma2 = (trXtdPt1X - 2*trPXtT + trTtdP1T) / (Np * D) 409 | 410 | # update Y 411 | if pt2pt_dis_sq(Y, Y_0 + np.matmul(G, W)) < tol: 412 | # if converged, break loop 413 | Y = Y_0 + np.matmul(G, W) 414 | print("iteration until convergence:", it) 415 | break 416 | else: 417 | # keep going until max iteration is reached 418 | Y = Y_0 + np.matmul(G, W) 419 | 420 | if it == max_iter - 1: 421 | print("did not converge!") 422 | 423 | return Y, sigma2 424 | 425 | initialized = False 426 | use_eval_rope = True 427 | pub_tracking_img = True 428 | init_nodes = [] 429 | nodes = [] 430 | guide_nodes_Y_0 = [] 431 | sigma2 = 0 432 | guide_nodes_sigma2_0 = 0 433 | total_len = 0 434 | geodesic_coord = [] 435 | def callback (rgb, pc): 436 | global initialized 437 | global init_nodes, nodes, sigma2 438 | global total_len, geodesic_coord 439 | global guide_nodes_Y_0, guide_nodes_sigma2_0 440 | global params, read_params 441 | global occlusion_mask_rgb 442 | 443 | # log time 444 | cur_time_cb = time.time() 445 | print('----------') 446 | 447 | # process rgb image 448 | cur_image = ros_numpy.numpify(rgb) 449 | # cur_image = cv2.cvtColor(cur_image.copy(), cv2.COLOR_BGR2RGB) 450 | hsv_image = cv2.cvtColor(cur_image.copy(), cv2.COLOR_RGB2HSV) 451 | 452 | # process point cloud 453 | pc_data = ros_numpy.point_cloud2.pointcloud2_to_array(pc) 454 | cur_pc = ros_numpy.point_cloud2.get_xyz_points(pc_data) 455 | cur_pc = cur_pc.reshape((720, 1280, 3)) 456 | 457 | # process opencv mask 458 | if occlusion_mask_rgb is None: 459 | occlusion_mask_rgb = np.ones(cur_image.shape).astype('uint8')*255 460 | occlusion_mask = cv2.cvtColor(occlusion_mask_rgb.copy(), cv2.COLOR_RGB2GRAY) 461 | 462 | if not use_eval_rope: 463 | # color thresholding 464 | lower = (90, 90, 90) 465 | upper = (120, 255, 255) 466 | mask = cv2.inRange(hsv_image, lower, upper) 467 | else: 468 | # color thresholding 469 | # --- rope blue --- 470 | lower = (90, 60, 40) 471 | upper = (130, 255, 255) 472 | mask_dlo = cv2.inRange(hsv_image, lower, upper).astype('uint8') 473 | 474 | # --- tape red --- 475 | lower = (130, 60, 40) 476 | upper = (255, 255, 255) 477 | mask_red_1 = cv2.inRange(hsv_image, lower, upper).astype('uint8') 478 | lower = (0, 60, 40) 479 | upper = (10, 255, 255) 480 | mask_red_2 = cv2.inRange(hsv_image, lower, upper).astype('uint8') 481 | mask_marker = cv2.bitwise_or(mask_red_1.copy(), mask_red_2.copy()).astype('uint8') 482 | 483 | # combine masks 484 | mask = cv2.bitwise_or(mask_marker.copy(), mask_dlo.copy()) 485 | mask = cv2.bitwise_and(mask.copy(), occlusion_mask.copy()) 486 | 487 | bmask = mask.copy() # for checking visibility, max = 255 488 | mask = cv2.cvtColor(mask.copy(), cv2.COLOR_GRAY2BGR) 489 | 490 | # publish mask 491 | mask_img_msg = ros_numpy.msgify(Image, mask, 'rgb8') 492 | mask_img_pub.publish(mask_img_msg) 493 | 494 | mask = (mask/255).astype(int) 495 | 496 | filtered_pc = cur_pc*mask 497 | filtered_pc = filtered_pc[((filtered_pc[:, :, 0] != 0) | (filtered_pc[:, :, 1] != 0) | (filtered_pc[:, :, 2] != 0))] 498 | # filtered_pc = filtered_pc[filtered_pc[:, 2] < 0.705] 499 | filtered_pc = filtered_pc[filtered_pc[:, 2] > 0.58] 500 | 501 | # downsample with open3d 502 | pcd = o3d.geometry.PointCloud() 503 | pcd.points = o3d.utility.Vector3dVector(filtered_pc) 504 | downpcd = pcd.voxel_down_sample(voxel_size=0.005) 505 | filtered_pc = np.asarray(downpcd.points) 506 | 507 | rospy.loginfo("Downsampled point cloud size: " + str(len(filtered_pc))) 508 | 509 | # add color 510 | pc_rgba = struct.unpack('I', struct.pack('BBBB', 255, 40, 40, 255))[0] 511 | pc_rgba_arr = np.full((len(filtered_pc), 1), pc_rgba) 512 | filtered_pc_colored = np.hstack((filtered_pc, pc_rgba_arr)).astype('O') 513 | filtered_pc_colored[:, 3] = filtered_pc_colored[:, 3].astype(int) 514 | 515 | # filtered_pc = filtered_pc.reshape((len(filtered_pc)*len(filtered_pc[0]), 3)) 516 | header.stamp = rospy.Time.now() 517 | converted_points = pcl2.create_cloud(header, fields, filtered_pc_colored) 518 | pc_pub.publish(converted_points) 519 | 520 | rospy.logwarn('callback before initialized: ' + str((time.time() - cur_time_cb)*1000) + ' ms') 521 | 522 | # register nodes 523 | if not initialized: 524 | 525 | init_nodes, sigma2 = register(filtered_pc, 40, 0.05, max_iter=100) 526 | init_nodes = sort_pts(init_nodes) 527 | 528 | nodes = init_nodes.copy() 529 | 530 | # compute preset coord and total len. one time action 531 | seg_dis = np.sqrt(np.sum(np.square(np.diff(init_nodes, axis=0)), axis=1)) 532 | geodesic_coord = [] 533 | last_pt = 0 534 | geodesic_coord.append(last_pt) 535 | for i in range (1, len(init_nodes)): 536 | last_pt += seg_dis[i-1] 537 | geodesic_coord.append(last_pt) 538 | geodesic_coord = np.array(geodesic_coord) 539 | total_len = np.sum(np.sqrt(np.sum(np.square(np.diff(init_nodes, axis=0)), axis=1))) 540 | 541 | initialized = True 542 | # header.stamp = rospy.Time.now() 543 | # converted_init_nodes = pcl2.create_cloud(header, fields, init_nodes) 544 | # nodes_pub.publish(converted_init_nodes) 545 | 546 | # cpd 547 | if initialized: 548 | # determined which nodes are occluded from mask information 549 | mask_dis_threshold = 10 550 | # projection 551 | init_nodes_h = np.hstack((init_nodes, np.ones((len(init_nodes), 1)))) 552 | # proj_matrix: 3*4; nodes_h.T: 4*M; result: 3*M 553 | image_coords = np.matmul(proj_matrix, init_nodes_h.T).T 554 | us = (image_coords[:, 0] / image_coords[:, 2]).astype(int) 555 | vs = (image_coords[:, 1] / image_coords[:, 2]).astype(int) 556 | 557 | # temp 558 | us = np.where(us >= 1280, 1279, us) 559 | vs = np.where(vs >= 720, 719, vs) 560 | 561 | uvs = np.vstack((vs, us)).T 562 | uvs_t = tuple(map(tuple, uvs.T)) 563 | 564 | # invert bmask for distance transform 565 | bmask_transformed = ndimage.distance_transform_edt(255 - bmask) 566 | # bmask_transformed = bmask_transformed / np.amax(bmask_transformed) 567 | vis = bmask_transformed[uvs_t] 568 | # occluded_nodes = np.where(vis > mask_dis_threshold)[0] 569 | 570 | # log time 571 | cur_time = time.time() 572 | nodes, sigma2 = cpd_lle(filtered_pc, nodes, 0.7, 5, 1, 0.05, 50, 0.00001, True, False, False, sigma2) 573 | rospy.logwarn('tracking_step total: ' + str((time.time() - cur_time)*1000) + ' ms') 574 | 575 | init_nodes = nodes.copy() 576 | 577 | results = ndarray2MarkerArray(nodes, "camera_color_optical_frame", [255, 150, 0, 0.75], [0, 255, 0, 0.75]) 578 | results_pub.publish(results) 579 | 580 | if pub_tracking_img: 581 | # project and pub tracking image 582 | nodes_h = np.hstack((nodes, np.ones((len(nodes), 1)))) 583 | 584 | # proj_matrix: 3*4; nodes_h.T: 4*M; result: 3*M 585 | image_coords = np.matmul(proj_matrix, nodes_h.T).T 586 | us = (image_coords[:, 0] / image_coords[:, 2]).astype(int) 587 | vs = (image_coords[:, 1] / image_coords[:, 2]).astype(int) 588 | 589 | cur_image_masked = cv2.bitwise_and(cur_image, occlusion_mask_rgb) 590 | tracking_img = (cur_image*0.5 + cur_image_masked*0.5).astype(np.uint8) 591 | 592 | for i in range (len(image_coords)): 593 | # draw circle 594 | uv = (us[i], vs[i]) 595 | if vis[i] < mask_dis_threshold: 596 | cv2.circle(tracking_img, uv, 5, (255, 150, 0), -1) 597 | else: 598 | cv2.circle(tracking_img, uv, 5, (255, 0, 0), -1) 599 | 600 | # draw line 601 | if i != len(image_coords)-1: 602 | if vis[i] < mask_dis_threshold: 603 | cv2.line(tracking_img, uv, (us[i+1], vs[i+1]), (0, 255, 0), 2) 604 | else: 605 | cv2.line(tracking_img, uv, (us[i+1], vs[i+1]), (255, 0, 0), 2) 606 | 607 | tracking_img_msg = ros_numpy.msgify(Image, tracking_img, 'rgb8') 608 | tracking_img_pub.publish(tracking_img_msg) 609 | 610 | rospy.logwarn('callback total: ' + str((time.time() - cur_time_cb)*1000) + ' ms') 611 | 612 | if __name__=='__main__': 613 | 614 | rospy.init_node('tracking_test', anonymous=True) 615 | 616 | rgb_sub = message_filters.Subscriber('/camera/color/image_raw', Image) 617 | pc_sub = message_filters.Subscriber('/camera/depth/color/points', PointCloud2) 618 | 619 | # header 620 | header = std_msgs.msg.Header() 621 | header.stamp = rospy.Time.now() 622 | header.frame_id = 'camera_color_optical_frame' 623 | fields = [PointField('x', 0, PointField.FLOAT32, 1), 624 | PointField('y', 4, PointField.FLOAT32, 1), 625 | PointField('z', 8, PointField.FLOAT32, 1), 626 | PointField('rgba', 12, PointField.UINT32, 1)] 627 | pc_pub = rospy.Publisher ('/pts', PointCloud2, queue_size=10) 628 | results_pub = rospy.Publisher ('/results', MarkerArray, queue_size=10) 629 | tracking_img_pub = rospy.Publisher ('/tracking_img', Image, queue_size=10) 630 | mask_img_pub = rospy.Publisher('/mask', Image, queue_size=10) 631 | 632 | opencv_mask_sub = rospy.Subscriber('/mask_with_occlusion', Image, update_occlusion_mask) 633 | 634 | ts = message_filters.TimeSynchronizer([rgb_sub, pc_sub], 10) 635 | ts.registerCallback(callback) 636 | 637 | rospy.spin() --------------------------------------------------------------------------------