Segnet ROS - arieldo/MorBot GitHub Wiki

#!/usr/bin/env python

import sys
import rospy
import cv2
import numpy as np

import jetson.utils
from jetson_inference import segNet
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image ,PointCloud
from geometry_msgs.msg import Point32
from std_msgs.msg import Int32MultiArray

bridge = CvBridge()
rospy.init_node('segnet_estimation_node', anonymous=True)

# Fetch desired publishing frequency from ROS parameter server (defaulting to 10 Hz if not set)
pub_frequency = rospy.get_param('~pub_frequency', 30)
rate = rospy.Rate(pub_frequency)  # Set the rate for the loop

# Publishers
pub_overlay = rospy.Publisher("/segnet_estimations/overlay", Image, queue_size=10)
pub_positions = rospy.Publisher("/segnet_estimations/positions", PointCloud, queue_size=10)
size_pub = rospy.Publisher('/image_size', Int32MultiArray, queue_size=10)

# Loading the segNet model
net = segNet("fcn-resnet18-sun")

def image_callback(msg):
    global net

    try:
        cv2_img = bridge.imgmsg_to_cv2(msg, "bgr8")
        height, width,channels = cv2_img.shape
    except CvBridgeError as e:
        print(e)
        return

    cv2_img = cv2.cvtColor(cv2_img, cv2.COLOR_RGBA2RGB)

    cuda_img = jetson.utils.cudaFromNumpy(cv2_img)
    
    net.Process(cuda_img)

    # Get class mask (this will contain class IDs for each pixel)
    class_mask_cuda = jetson.utils.cudaAllocMapped(width=cuda_img.width, height=cuda_img.height, format='gray8')
    net.Mask(class_mask_cuda, class_mask_cuda.width, class_mask_cuda.height)

    class_mask_np = jetson.utils.cudaToNumpy(class_mask_cuda)

    # Here we filter out everything but the desired class. 
    # Let's say the class ID of the desired class is 2 (see below)
    desired_class_id = 2
    class_mask_np[class_mask_np != desired_class_id] = 0
    class_mask_np[class_mask_np == desired_class_id] = 255  # This sets the desired class pixels to white

    # Overlay this mask with original image. This step might need further adjustments based on your need.
    overlayed_img = cv2.bitwise_and(cv2_img, cv2_img, mask=class_mask_np)

    overlay_img_msg = bridge.cv2_to_imgmsg(overlayed_img, encoding="rgb8")
    pub_overlay.publish(overlay_img_msg)

    ys, xs = np.where(class_mask_np[:,:,0] == 255)
    
    centroid_x = xs.mean() if len(xs) else None
    centroid_y = ys.mean() if len(ys) else None


# Populate the PointCloud message only if there's a valid centroid
    if centroid_x is not None and centroid_y is not None:
        msg = PointCloud()
        msg.header.stamp = rospy.Time.now()
        msg.header.frame_id = "camera_frame" # or whatever your frame is
        msg.points = [Point32(centroid_x, centroid_y, 0)]

    # Publish the message
        pub_positions.publish(msg)

    # Publish input img size
    size_msg = Int32MultiArray()
    size_msg.data = [width, height]
    size_pub.publish(size_msg)
    
    rate.sleep()  # This makes sure the loop runs at the desired frequency

def main(args):
    rospy.Subscriber("/camera/image_raw", Image, image_callback,queue_size=10)
    try:
        rospy.spin()
    except KeyboardInterrupt:
        print("Shutting down ROS segNet estimation module")

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