Module not found: onnxruntime, despite virtual environement and dependencies

edit

Hello, I've been trying to have a node run a classification model using onnxruntime. The problem is that whatever I do I still get a ModuleNotFoundError: No module named 'onnxruntime'.

To install the python package I followed the instructions on this page, using a virtual environment. I use pyenv to make sure I'm on python3.8 and virtualenv to make the virtual environment. And then I installed every dependency including onnxruntime-gpu.

I added a dependency in my package.xml : <exec_depend>python3-onnxruntime-gpu-pip</exec_depend> which seems to be correct since I get this:

rosdep install -i --from-path src --rosdistro galactic -y
#All required rosdeps installed successfully

For reference here is my code:

import rclpy
import cv2
import onnxruntime as ort
import numpy as np

from rclpy.node import Node
from sensor_msgs.msg import Image, RegionOfInterest
from smart_cart_msgs.msg import DetectedObjectList, DetectedObject
from cv_bridge import CvBridge


def classifier_dummy(detected_objects, ort_sess):
    preds = []
    for obj in detected_objects:
      image = obj.image_crop
      outputs = ort_sess.run(None, {'input': image.numpy()})
      preds.append(outputs[0].argmax(axis=1)+1)
    return DetectedObjectList(preds)

class classifierNode(Node):
  def __init__(self, ort_sess):
    super().__init__('classifier')
    self.subscription = self.create_subscription(
      DetectedObjectList, 
      'object_detection', 
      self.listener_callback, 
      10)
    self.subscription
    self.br = CvBridge()
    self.publisher_ = self.create_publisher(
        DetectedObjectList, 
        'object_classification', 
        10
    )
    self.ort_sess = ort_sess

  def listener_callback(self, data):
    self.get_logger().info('Classifier receiving video frame')
    self.get_logger().info('Performing classification')
    ret_value = classifier_dummy(data, self.ort_sess)
    self.get_logger().info('Sending classification results')
    self.publisher_.publish(ret_value)

def main(args=None):
  rclpy.init(args=args)
  ort_sess = ort.InferenceSession("weights.onnx")
  classifier_node = classifierNode(ort_sess)
  rclpy.spin(classifier_node)

  classifier_node.destroy_node()

  rclpy.shutdown()

if __name__ == '__main__':
  main()

So classifierNode is subscribed to the topic object_detection and will receive a list of cropped image containing detected objects and will perform inference using the model run with onnxruntme. And then sends the results to the next topic object_classification.

.

My package.xml:

  <exec_depend>python3-onnxruntime-gpu-pip</exec_depend>
  <exec_depend>rclpy</exec_depend>
  <exec_depend>image_transport</exec_depend>
  <exec_depend>cv_bridge</exec_depend>
  <exec_depend>sensor_msgs</exec_depend>
  <exec_depend>std_msgs</exec_depend>
  <exec_depend>smart_cart_msgs</exec_depend>

  <build_depend>rosidl_default_generators</build_depend>
  <exec_depend>rosidl_default_runtime</exec_depend>
  <member_of_group>rosidl_interface_packages</member_of_group>

  <test_depend>ament_copyright</test_depend>
  <test_depend>ament_flake8</test_depend>
  <test_depend>ament_pep257</test_depend>
  <test_depend>python3-pytest</test_depend>

  <export>
    <build_type>ament_python</build_type>
  </export>

(env | grep ROS)

ROS_DISTRO=galactic
ROS_LOCALHOST_ONLY=0
ROS_PYTHON_VERSION=3
ROS_VERSION=2

pip list --local

Package           Version
----------------- --------
catkin-pkg        0.5.2
docutils          0.18.1
empy              3.3.4
flatbuffers       2.0
lark              1.1.2
numpy             1.22.4
onnx              1.11.0
onnxruntime-gpu   1.11.1
opencv-python     4.6.0.66
pip               22.0.4
protobuf          4.21.1
pyparsing         3.0.9
python-dateutil   2.8.2
setuptools        62.1.0
six               1.16.0
typing_extensions 4.2.0
wheel             0.37.1