mirror of
https://github.com/introlab/rtabmap_ros.git
synced 2026-10-03 16:27:46 +08:00
90 lines
2.6 KiB
Python
Executable File
90 lines
2.6 KiB
Python
Executable File
#!/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)
|
|
|