|
| 1 | +# Copyright (c) 2021-2022, NVIDIA CORPORATION. All rights reserved. |
| 2 | +# |
| 3 | +# NVIDIA CORPORATION and its licensors retain all intellectual property |
| 4 | +# and proprietary rights in and to this software, related documentation |
| 5 | +# and any modifications thereto. Any use, reproduction, disclosure or |
| 6 | +# distribution of this software and related documentation without an express |
| 7 | +# license agreement from NVIDIA CORPORATION is strictly prohibited. |
| 8 | + |
| 9 | +""" |
| 10 | +Basic Proof-Of-Life test for the Isaac ROS CenterPose Node. |
| 11 | +
|
| 12 | +This test checks that an image can be encoder into a tensor, run through |
| 13 | +Triton using a CenterPose model, and the inference decoded into a series of poses. |
| 14 | +""" |
| 15 | +import os |
| 16 | +import pathlib |
| 17 | +import time |
| 18 | + |
| 19 | +from ament_index_python.packages import get_package_share_directory |
| 20 | +from isaac_ros_test import IsaacROSBaseTest, JSONConversion |
| 21 | +from launch_ros.actions import ComposableNodeContainer, Node |
| 22 | +from launch_ros.descriptions import ComposableNode |
| 23 | + |
| 24 | +import pytest |
| 25 | +import rclpy |
| 26 | + |
| 27 | +from sensor_msgs.msg import Image |
| 28 | +from visualization_msgs.msg import MarkerArray |
| 29 | + |
| 30 | +MODEL_GENERATION_TIMEOUT_SEC = 300 |
| 31 | +INIT_WAIT_SEC = 10 |
| 32 | +MODEL_PATH = '/tmp/centerpose_trt_engine.plan' |
| 33 | + |
| 34 | + |
| 35 | +@pytest.mark.rostest |
| 36 | +def generate_test_description(): |
| 37 | + launch_dir_path = pathlib.Path(os.path.dirname(os.path.realpath(__file__))) |
| 38 | + model_dir_path = launch_dir_path / 'models' |
| 39 | + model_file_name = 'shoe_resnet_140.onnx' |
| 40 | + config = os.path.join( |
| 41 | + get_package_share_directory('isaac_ros_centerpose'), |
| 42 | + 'config', |
| 43 | + 'decoder_params_test.yaml' |
| 44 | + ) |
| 45 | + |
| 46 | + centerpose_encoder_node = ComposableNode( |
| 47 | + name='centerpose_encoder', |
| 48 | + package='isaac_ros_dnn_encoders', |
| 49 | + plugin='nvidia::isaac_ros::dnn_inference::DnnImageEncoderNode', |
| 50 | + namespace=IsaacROSCenterPosePOLTest.generate_namespace(), |
| 51 | + parameters=[{ |
| 52 | + 'network_image_width': 512, |
| 53 | + 'network_image_height': 512, |
| 54 | + 'image_mean': [0.408, 0.447, 0.47], |
| 55 | + 'image_stddev': [0.289, 0.274, 0.278] |
| 56 | + }], |
| 57 | + remappings=[('encoded_tensor', 'tensor_pub')]) |
| 58 | + |
| 59 | + centerpose_inference_node = ComposableNode( |
| 60 | + name='centerpose_inference', |
| 61 | + package='isaac_ros_tensor_rt', |
| 62 | + plugin='nvidia::isaac_ros::dnn_inference::TensorRTNode', |
| 63 | + namespace=IsaacROSCenterPosePOLTest.generate_namespace(), |
| 64 | + parameters=[{ |
| 65 | + 'model_file_path': str(model_dir_path / model_file_name), |
| 66 | + 'engine_file_path': MODEL_PATH, |
| 67 | + 'input_tensor_names': ['input_tensor'], |
| 68 | + 'input_binding_names': ['input'], |
| 69 | + 'input_tensor_formats': ['nitros_tensor_list_nchw_rgb_f32'], |
| 70 | + 'output_tensor_names': ['hm', 'wh', 'hps', 'reg', 'hm_hp', 'hp_offset', 'scale'], |
| 71 | + 'output_binding_names': ['hm', 'wh', 'hps', 'reg', 'hm_hp', 'hp_offset', 'scale'], |
| 72 | + 'output_tensor_formats': ['nitros_tensor_list_nhwc_rgb_f32'], |
| 73 | + 'verbose': False, |
| 74 | + 'force_engine_update': False |
| 75 | + }]) |
| 76 | + |
| 77 | + centerpose_decoder_node = Node( |
| 78 | + name='centerpose_decoder_node', |
| 79 | + package='isaac_ros_centerpose', |
| 80 | + namespace=IsaacROSCenterPosePOLTest.generate_namespace(), |
| 81 | + executable='CenterPoseDecoder', |
| 82 | + parameters=[config], |
| 83 | + output='screen' |
| 84 | + ) |
| 85 | + |
| 86 | + rclcpp_container = ComposableNodeContainer( |
| 87 | + name='rclcpp_container', |
| 88 | + namespace='', |
| 89 | + package='rclcpp_components', |
| 90 | + executable='component_container_mt', |
| 91 | + composable_node_descriptions=[ |
| 92 | + centerpose_encoder_node, centerpose_inference_node], |
| 93 | + output='screen', |
| 94 | + ) |
| 95 | + |
| 96 | + return IsaacROSCenterPosePOLTest.generate_test_description([rclcpp_container, |
| 97 | + centerpose_decoder_node]) |
| 98 | + |
| 99 | + |
| 100 | +class IsaacROSCenterPosePOLTest(IsaacROSBaseTest): |
| 101 | + filepath = pathlib.Path(os.path.dirname(__file__)) |
| 102 | + |
| 103 | + def test_pol(self): |
| 104 | + |
| 105 | + self.node._logger.info( |
| 106 | + f'Generating model (timeout={MODEL_GENERATION_TIMEOUT_SEC}s)') |
| 107 | + start_time = time.time() |
| 108 | + wait_cycles = 1 |
| 109 | + while not os.path.isfile(MODEL_PATH): |
| 110 | + time_diff = time.time() - start_time |
| 111 | + if time_diff > MODEL_GENERATION_TIMEOUT_SEC: |
| 112 | + self.fail('Model generation timed out') |
| 113 | + if time_diff > wait_cycles*10: |
| 114 | + self.node._logger.info( |
| 115 | + f'Waiting for model generation to finish... ({int(time_diff)}s passed)') |
| 116 | + wait_cycles += 1 |
| 117 | + time.sleep(1) |
| 118 | + |
| 119 | + # Wait for TensorRT engine |
| 120 | + time.sleep(INIT_WAIT_SEC) |
| 121 | + |
| 122 | + self.node._logger.info( |
| 123 | + f'Model generation was finished (took {(time.time() - start_time)}s)') |
| 124 | + |
| 125 | + received_messages = {} |
| 126 | + |
| 127 | + self.generate_namespace_lookup(['image', 'object_poses']) |
| 128 | + |
| 129 | + image_pub = self.node.create_publisher( |
| 130 | + Image, self.namespaces['image'], self.DEFAULT_QOS) |
| 131 | + |
| 132 | + # The current CenterPose decoder outputs MarkerArray |
| 133 | + subs = self.create_logging_subscribers( |
| 134 | + [('object_poses', MarkerArray)], received_messages) |
| 135 | + |
| 136 | + try: |
| 137 | + json_file = self.filepath / 'test_cases/gray_shoe/image.json' |
| 138 | + image = JSONConversion.load_image_from_json(json_file) |
| 139 | + |
| 140 | + TIMEOUT = 60 |
| 141 | + end_time = time.time() + TIMEOUT |
| 142 | + done = False |
| 143 | + |
| 144 | + while time.time() < end_time: |
| 145 | + image_pub.publish(image) |
| 146 | + rclpy.spin_once(self.node, timeout_sec=(0.1)) |
| 147 | + if 'object_poses' in received_messages: |
| 148 | + done = True |
| 149 | + break |
| 150 | + self.assertTrue(done, 'Timeout. Appropriate output not received') |
| 151 | + self.node._logger.info('A message was successfully received') |
| 152 | + |
| 153 | + finally: |
| 154 | + self.node.destroy_subscription(subs) |
| 155 | + self.node.destroy_publisher(image_pub) |
0 commit comments