Robotics StackExchange | Archived questions

Gazebo Camera Images Fed Into Custom Model in .onnx File Format

Hi all,

Below I have template code I intend to use to incorporate my custom model that's in .onnx format. Currently have 'getresult' receive images from Gazebo robot camera, but then I want to be able to feed that image into a custom model. If you are familiar with jetson-inference, my thought process was to do something similar to the 'detectnet.py' file. Load model ('net = jetson.inference.detectNet(somepretrained_network, threshold'), and the calling 'net.Detect(img)' to get results on the image given by the Gazebo camera. I'm taking ideas to any approach, all is helpful. Thank you very much!

#/usr/bin/env python

import cv2
import numpy as np
import jetson.utils
from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError

def show_image(img):
    """
        Displays images

    Parameters
    +--------------------+
    img: The image to be displayed. Should be of type 'ndarray'
"""
cv2.imshow('image', img)
cv2.waitKey(1)

def get_result(img):
    """
        Given an image from 'situation_awareness' ROS node report back results such as
        bounding box coordinates of a given light source w/in Gazebo

    Parameters
    +--------------------+
    img: 'Image' message type that was forwarded from 'situational_awareness' node for
            processing

    Return value(s)
    +--------------------+
    data: Python dictionary that stores model results in key value pairs. Keys:
            {"class_id", "confidence", "predicted_class", "left", "top", "right",
            "bottom", "width", "height", "area", "center"}
"""
# Initialize return value
data = None
# Initialize CvBridge class
bridge = CvBridge()
# TODO: Import exported custom object detection model
# model = import(custom_model)
# Do necessary preprocessing on 'img' parameter to feed into model
try:
    # Convert ROS Image class to OpenCV2 image
    cv2_img = bridge.imgmsg_to_cv2(img, desired_encoding="bgr8")
    # Convert OpenCV2 image (numpy ndarray) to cudaImage
    cuda_mem = jetson.utils.cudaFromNumpy(cv2_img)
except CvBridgeError as e:
    print(e)
# Copy above image
drawImg = cv2_img
# Display the image
show_image(drawImg)
# TODO: Produce results by passing processed image into the model
# detections = model(cuda_mem, img.width, img.height)
# TODO: Parse results from predictions/detections
# if (len(detections) > 0):
    # data = {}
    # detection = detections[0]
    # data['class_id'] = detection.ClassID
    # data['confidence'] = detection.Confidence
    # data['predicted_class'] = model.GetClassDesc(class_id)
    # data['left'] = detection.Left
    # data['top'] = detection.Top
    # data['right'] = detection.Right
    # data['bottom'] = detection.Bottom
    # data['width'] = detection.Width
    # data['height'] = detection.Height
    # data['area'] = detection.Area
    # data['center_x'] = detection.Center[0]
    # data['center_y'] = detection.Center[1]
# TODO: Return results in a data structure of some sort for later 
# processing by sender
return data

Asked by aalex on 2021-08-04 21:58:54 UTC

Comments

Answers