diff --git a/CMakeLists.txt b/CMakeLists.txt index ea637846..bb264cd1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -81,7 +81,7 @@ ENDIF(RTABMAP_GUI OR rviz_FOUND) ## Uncomment this if the package has a setup.py. This macro ensures ## modules and global scripts declared therein get installed ## See http://ros.org/doc/api/catkin/html/user_guide/setup_dot_py.html -# catkin_python_setup() +catkin_python_setup() ####################################### ## Declare ROS messages and services ## @@ -517,6 +517,7 @@ catkin_install_python(PROGRAMS scripts/point_to_tf.py scripts/transform_to_tf.py scripts/yaml_to_camera_info.py + scripts/netvlad_tf_ros.py scripts/wifi_signal_pub.py DESTINATION ${CATKIN_PACKAGE_BIN_DESTINATION} ) diff --git a/python/rtabmap_ros/__init__.py b/python/rtabmap_ros/__init__.py new file mode 100644 index 00000000..139597f9 --- /dev/null +++ b/python/rtabmap_ros/__init__.py @@ -0,0 +1,2 @@ + + diff --git a/python/rtabmap_ros/compression.py b/python/rtabmap_ros/compression.py new file mode 100644 index 00000000..140f7f6f --- /dev/null +++ b/python/rtabmap_ros/compression.py @@ -0,0 +1,34 @@ + +import zlib +import struct +import numpy as np + +def compress(data): + assert data.ndim == 1 or data.ndim == 2 + + dim1 = 1 + if data.ndim == 1: + dim1 = 1 + dim2 = len(data) + else: + dim1 = data.shape[0] + dim2 = data.shape[1] + + numpy_type_to_cvtype = {'uint8': 0, 'int8': 1, 'uint16': 2, + 'int16': 3, 'int32': 4, 'float32': 5, + 'float64': 6} + + compressed_data = bytearray(zlib.compress(data.tobytes())) + compressed_data.extend(struct.pack("iii", dim1, dim2, numpy_type_to_cvtype[data.dtype.name])) + + return compressed_data + +def uncompress(bytes): + cvtype_to_numpy_type = {0: 'uint8', 1: 'int8', 2: 'uint16', + 3: 'int16', 4: 'int32', 5: 'float32', + 6: 'float64'} + out = zlib.decompress(bytes[:len(bytes)-3*4]) + rows, cols, datatype = struct.unpack_from("iii", bytes, offset=len(bytes)-3*4) + data = np.frombuffer(out, dtype=cvtype_to_numpy_type[datatype]) + return data.reshape((rows, cols)) + diff --git a/scripts/netvlad_tf_ros.py b/scripts/netvlad_tf_ros.py new file mode 100755 index 00000000..1a079b73 --- /dev/null +++ b/scripts/netvlad_tf_ros.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python + +# Using netvlad tensorflow-v1 implementation from https://github.com/uzh-rpg/netvlad_tf_open/ +# For ROS melodic, follow the following instructions to rebuild cv_bridge with Python3 +# https://medium.com/@beta_b0t/how-to-setup-ros-with-python-3-44a69ca36674 +# On Jetpack 4.4 (18.04 and OpenCV4), use vision_opencv's noetic branch. In cv_bridge/CMakeLists.txt, +# apply this patch: +# -find_package(Boost REQUIRED python37) +# +find_package(Boost REQUIRED python3) + +from __future__ import print_function + +import roslib +import sys +import rospy +import cv2 +import numpy as np +import tensorflow as tf +import time + +import netvlad_tf.net_from_mat as nfm +import netvlad_tf.nets as nets + +from std_msgs.msg import String +from sensor_msgs.msg import Image +from cv_bridge import CvBridge, CvBridgeError +from rtabmap_ros import compression as cp +from rtabmap_ros.msg import GlobalDescriptor + +class netvlad_ros: + + def __init__(self): + + self.dim = rospy.get_param('~dim', 4096) + self.scale = rospy.get_param('~scale', 1.0) + rospy.loginfo("Parameter dim=%d", self.dim) + rospy.loginfo("Parameter scale=%d", self.scale) + + tf.reset_default_graph() + + self.image_batch = tf.placeholder( + dtype=tf.float32, shape=[None, None, None, 3]) + + self.net_out = nets.vgg16NetvladPca(self.image_batch) + self.saver = tf.train.Saver() + + self.sess = tf.Session() + self.saver.restore(self.sess, nets.defaultCheckpoint()) + + self.pub = rospy.Publisher('netvlad_descriptor', GlobalDescriptor, queue_size=1) + + self.bridge = CvBridge() + self.image_sub = rospy.Subscriber("image",Image,self.callback, queue_size=1) + + def callback(self,data): + start = time.time() + try: + cv_image = self.bridge.imgmsg_to_cv2(data, "rgb8") + except CvBridgeError as e: + print(e) + + if self.scale != 1.0: + width = int(cv_image.shape[1] * self.scale) + height = int(cv_image.shape[0] * self.scale) + cv_image = cv2.resize(cv_image, (width, height), interpolation = cv2.INTER_AREA) + + batch = np.expand_dims(cv_image, axis=0) + result = self.sess.run(self.net_out, feed_dict={self.image_batch: batch}) + result = result[:,:self.dim] + + descriptor = GlobalDescriptor() + descriptor.type = 0 + descriptor.header = data.header + descriptor.data = cp.compress(result) + self.pub.publish(descriptor) + end = time.time() + rospy.loginfo("Extracting descriptor (img=%dx%d, dim=%d): %fs", cv_image.shape[1], cv_image.shape[0], self.dim, end-start) + +def main(args): + rospy.init_node('netvlad', anonymous=True) + n = netvlad_ros() + try: + rospy.spin() + except KeyboardInterrupt: + print("Shutting down") + +if __name__ == '__main__': + main(sys.argv) + diff --git a/setup.py b/setup.py new file mode 100644 index 00000000..7279001b --- /dev/null +++ b/setup.py @@ -0,0 +1,10 @@ +#!/usr/bin/env python +from setuptools import setup +from catkin_pkg.python_setup import generate_distutils_setup + +d = generate_distutils_setup() + +d['packages'] = ['rtabmap_ros'] +d['package_dir'] = {'' : 'python'} + +setup(**d)