rospy CLIP Node

This simple CLIP embedding computing ROS node is provided as a utility for running the tutorial. It has two dependencies outside of ROS that can be installed via:

pip3 install opencv-python git+https://github.com/openai/CLIP.git

In order to be able to launch it, ensure to make it runnable with chmod +x clip_rospy.py

The complete file below.

#!/usr/bin/python3
import rospy
import sys
import cv2
from std_msgs.msg import String, Float64MultiArray
from sensor_msgs.msg import Image as ImageMsg
from cv_bridge import CvBridge, CvBridgeError
import time
import clip
from PIL import Image

DEFAULT_CLIP_MODEL = "ViT-L/14"


class ClipROSNode:
    def __init__(self):
        self.bridge = CvBridge()

        self.image_sub = rospy.Subscriber("/zed2/camera/left/image_raw_throttled", ImageMsg, self.callback)
        self.emb_pub = rospy.Publisher("/clip/embeddings", Float64MultiArray, queue_size=10)

        self.device = "cuda"
        self.embedding_model, self.embedding_preprocess = clip.load(DEFAULT_CLIP_MODEL, device=self.device)

        self.msg = Float64MultiArray()

    def callback(self, data):
        t1 = time.time()
        try:
            cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8")
        except CvBridgeError as e:
            print(e)

        t6 = time.time()
        cv2_img = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
        pil_img = Image.fromarray(cv2_img)
        t7 = time.time()

        image = self.embedding_preprocess(pil_img).unsqueeze(0).to(self.device)
        embeddings = self.embedding_model.encode_image(image).detach().cpu().numpy()[0]
        t8 = time.time()
        self.msg.data = embeddings.tolist()
        self.emb_pub.publish(self.msg)

        t9 = time.time()
        rospy.loginfo_throttle(1, "Clip Embeddings Timings: {} {} {} {}".format(t9 - t8, t8 - t7, t7 - t6, t6 - t1))


def main(args):
    ic = ClipROSNode()
    rospy.init_node('ClipROSNode', anonymous=True)
    try:
        rospy.spin()
    except KeyboardInterrupt:
        print("Shutting down")


if __name__ == '__main__':
    main(sys.argv)