example of visualizing markers

marker_viz.py

#!/usr/bin/env python3

import sys

import rospy
import rostopic

import numpy as np

import message_filters
from std_msgs.msg import ColorRGBA, Header
from visualization_msgs.msg import Marker, MarkerArray
from vision_msgs.msg import ObjectHypothesis, Detection2D, Detection2DArray


labelMap = [
    "background",
    "aeroplane",
    "bicycle",
    "bird",
    "boat",
    "bottle",
    "bus",
    "car",
    "cat",
    "chair",
    "cow",
    "diningtable",
    "dog",
    "horse",
    "motorbike",
    "person",
    "pottedplant",
    "sheep",
    "sofa",
    "train",
    "tvmonitor",
]


def object_annotator(
    reid,
    roi,
    position,
    time: rospy.Time,
    frame_id: str = "oak_rgb_camera_optical_frame",
    offset=0.2,
    cube_size=0.2,
):
    cube = Marker()
    cube.header.frame_id = frame_id
    cube.header.stamp = time
    cube.id = int(reid.replace("_", "")[-3:] + "2")
    cube.ns = "detection_visualizer"
    cube.type = Marker.CUBE
    cube.action = Marker.DELETEALL
    cube.action = Marker.ADD

    cube.scale.x = cube_size
    cube.scale.y = cube_size
    cube.scale.z = cube_size

    cube.pose.position.x = position[0]
    cube.pose.position.y = position[1]
    cube.pose.position.z = position[2]

    cube.pose.orientation.x = 0
    cube.pose.orientation.y = 0
    cube.pose.orientation.z = 0
    cube.pose.orientation.w = 1

    cube.color = ColorRGBA(0.2, 0.8, 0.85, 0.4)
    cube.lifetime = rospy.Duration(0.1)

    return cube


class DetectionVisualizer:
    """
    This node publishes Markers for Object Detections.
    """

    def __init__(self, topic, rate=10, queue_len=10):
        rostopic._check_master()
        rospy.init_node("detection_visualizer", anonymous=True)
        rospy.loginfo("%s started" % rospy.get_name())

        det_sub = rospy.Subscriber(topic, Detection2DArray, self.callback)

        self.publisher = rospy.Publisher(
            "matched_detections", Detection2DArray, queue_size=queue_len
        )
        self._target_frame = rospy.get_param("~target_frame", "oak_rgb_camera_frame")

        self.rate = rospy.Rate(rate)
        self.viz_publisher = rospy.Publisher(
            "detection_markers", MarkerArray, queue_size=1
        )
        self.rate.sleep()

    def from_ros(self, tracklet):
        center, width, height = (
            tracklet.bbox.center,
            tracklet.bbox.size_x,
            tracklet.bbox.size_y,
        )
        x, y = center.x, center.y
        xmin, xmax = x - width // 2, x + width // 2
        ymin, ymax = y - height // 2, y + height // 2
        roi = (xmin, ymin, xmax, ymax)
        label = tracklet.results[0].id if tracklet.results else -1
        point = (
            tracklet.spatialCoordinate.x,
            tracklet.spatialCoordinate.y,
            tracklet.spatialCoordinate.z,
        )
        return point, roi, label

    def to_ros(self, det):
        detection = Detection2D()
        detection.results.append(ObjectHypothesis(id=det[2], score=1.0))
        return detection

    def callback(self, spatial_detections):
        trackedTracklets = []
        for tracklet in spatial_detections.detections:
            trackedTracklets.append(self.from_ros(tracklet))
        matchedDetection2DArray = Detection2DArray()
        markers = MarkerArray()
        for det in trackedTracklets:
            matchedDetection2DArray.detections.append(self.to_ros(det))
            box = object_annotator(
                labelMap[det[2]] + "_" + str(np.random.randint(100)).zfill(4),
                det[1],
                (det[0][0], det[0][1], det[0][2]),
                rospy.Time.now(),
            )
            markers.markers.append(box)
        matchedDetection2DArray.header = Header(
            stamp=rospy.Time.now(), frame_id=self._target_frame
        )
        self.publisher.publish(matchedDetection2DArray)
        self.viz_publisher.publish(markers)
        self.rate.sleep()

    def shutdown(self):
        rospy.loginfo("Shutting down node")
        rospy.signal_shutdown("Shutdown all threads")
        sys.exit(0)


if __name__ == "__main__":
    topic = "/rgb_publisher/color/mobilenet_detections"
    ros_node = DetectionVisualizer(topic=topic)
    rospy.on_shutdown(ros_node.shutdown)
    rospy.spin()