1+ #!/usr/bin/env python3
2+ # -*- coding:utf-8 -*-
3+ # Copyright (c) Megvii, Inc. and its affiliates.
4+
5+ # ROS2 rclpy -- Ar-Ray-code 2021
6+ import rclpy
7+
8+ import argparse
9+ import logging as log
10+ import os
11+ import sys
12+
13+ import cv2
14+ import numpy as np
15+
16+ from openvino .inference_engine import IECore
17+
18+ from yolox .data .data_augment import preproc as preprocess
19+ from yolox .data .datasets import COCO_CLASSES
20+ from yolox .utils import multiclass_nms , demo_postprocess , vis
21+
22+ # ROS2 =====================================
23+ import rclpy
24+ from rclpy .node import Node
25+
26+ from std_msgs .msg import Header
27+ from cv_bridge import CvBridge
28+ from sensor_msgs .msg import Image
29+
30+ from bboxes_ex_msgs .msg import BoundingBoxes
31+ from bboxes_ex_msgs .msg import BoundingBox
32+
33+ # from darkself.net_ros_msgs.msg import BoundingBoxes
34+ # from darkself.net_ros_msgs.msg import BoundingBox
35+
36+ class yolox_ros (Node ):
37+ def __init__ (self ) -> None :
38+
39+ # ROS2 init
40+ super ().__init__ ('yolox_ros' )
41+
42+ self .setting_yolox_exp ()
43+
44+ if (self .imshow_isshow ):
45+ cv2 .namedWindow ("YOLOX" )
46+
47+ self .bridge = CvBridge ()
48+
49+ self .pub = self .create_publisher (BoundingBoxes ,"yolox/bounding_boxes" , 10 )
50+ self .pub_image = self .create_publisher (Image ,"yolox/image_raw" , 10 )
51+ self .sub = self .create_subscription (Image ,"image_raw" ,self .imageflow_callback , 10 )
52+
53+ def setting_yolox_exp (self ) -> None :
54+ # set environment variables for distributed training
55+
56+ # ==============================================================
57+
58+ WEIGHTS_PATH = '/home/ubuntu/ros2_ws/install/yolox_ros_py/share/yolox_ros_py/yolox_s.xml'
59+
60+ self .declare_parameter ('imshow_isshow' ,True )
61+
62+ self .declare_parameter ('model_path' , WEIGHTS_PATH )
63+ self .declare_parameter ('conf' , 0.3 )
64+ self .declare_parameter ('device' , "CPU" )
65+
66+ self .declare_parameter ('image_size/width' , 640 )
67+ self .declare_parameter ('image_size/height' , 480 )
68+
69+
70+ # =============================================================
71+ self .imshow_isshow = self .get_parameter ('imshow_isshow' ).value
72+
73+ model_path = self .get_parameter ('model_path' ).value
74+ self .conf = self .get_parameter ('conf' ).value
75+ device = self .get_parameter ('device' ).value
76+
77+ self .input_width = self .get_parameter ('image_size/width' ).value
78+ self .input_height = self .get_parameter ('image_size/height' ).value
79+
80+ # ==============================================================
81+
82+ print ('Creating Inference Engine' )
83+ ie = IECore ()
84+ print (f'Reading the self.network: { model_path } )
85+ # (.xml and .bin files) or (.onnx file)
86+
87+ self .net = ie .read_network (model = model_path )
88+ print ('Configuring input and output blobs' )
89+ # Get names of input and output blobs
90+ self .input_blob = next (iter (self .net .input_info ))
91+ self .out_blob = next (iter (self .net .outputs ))
92+
93+ # Set input and output precision manually
94+ self .net .input_info [self .input_blob ].precision = 'FP32'
95+ self .net .outputs [self .out_blob ].precision = 'FP16'
96+
97+ # Get a number of classes recognized by a model
98+ num_of_classes = max (self .net .outputs [self .out_blob ].shape )
99+
100+ print ('Loading the model to the plugin' )
101+ self .exec_net = ie .load_network (network = self .net , device_name = device )
102+
103+
104+ def yolox2bboxes_msgs (self , bboxes , scores , cls , cls_names , img_header :Header ):
105+ bboxes_msg = BoundingBoxes ()
106+ bboxes_msg .header = img_header
107+ i = 0
108+ for bbox in bboxes :
109+ one_box = BoundingBox ()
110+ one_box .xmin = int (bbox [0 ])
111+ one_box .ymin = int (bbox [1 ])
112+ one_box .xmax = int (bbox [2 ])
113+ one_box .ymax = int (bbox [3 ])
114+ one_box .probability = float (scores [i ])
115+ one_box .class_id = str (cls_names [int (cls [i ])])
116+ bboxes_msg .bounding_boxes .append (one_box )
117+ i = i + 1
118+
119+ return bboxes_msg
120+
121+ def imageflow_callback (self ,msg :Image ) -> None :
122+ try :
123+ # fps start
124+ start_time = cv2 .getTickCount ()
125+ bboxes = BoundingBoxes ()
126+ img_rgb = self .bridge .imgmsg_to_cv2 (msg ,"bgr8" )
127+
128+ origin_img = img_rgb
129+ _ , _ , h , w = self .net .input_info [self .input_blob ].input_data .shape
130+ mean = (0.485 , 0.456 , 0.406 )
131+ std = (0.229 , 0.224 , 0.225 )
132+ image , ratio = preprocess (origin_img , (h , w ), mean , std )
133+
134+ res = self .exec_net .infer (inputs = {self .input_blob : image })
135+ res = res [self .out_blob ]
136+
137+ # Predictions is result
138+ predictions = demo_postprocess (res , (h , w ), p6 = False )[0 ]
139+
140+ boxes = predictions [:, :4 ]
141+ scores = predictions [:, 4 , None ] * predictions [:, 5 :]
142+
143+ boxes_xyxy = np .ones_like (boxes )
144+ boxes_xyxy [:, 0 ] = boxes [:, 0 ] - boxes [:, 2 ]/ 2.
145+ boxes_xyxy [:, 1 ] = boxes [:, 1 ] - boxes [:, 3 ]/ 2.
146+ boxes_xyxy [:, 2 ] = boxes [:, 0 ] + boxes [:, 2 ]/ 2.
147+ boxes_xyxy [:, 3 ] = boxes [:, 1 ] + boxes [:, 3 ]/ 2.
148+ boxes_xyxy /= ratio
149+ dets = multiclass_nms (boxes_xyxy , scores , nms_thr = 0.45 , score_thr = 0.1 )
150+
151+ if dets is not None :
152+ final_boxes = dets [:, :4 ]
153+ final_scores , final_cls_inds = dets [:, 4 ], dets [:, 5 ]
154+ origin_img = vis (origin_img , final_boxes , final_scores , final_cls_inds ,
155+ conf = self .conf , class_names = COCO_CLASSES )
156+ # ==============================================================
157+ end_time = cv2 .getTickCount ()
158+ time_took = (end_time - start_time ) / cv2 .getTickFrequency ()
159+
160+ # rclpy log FPS
161+ self .get_logger ().info (f'FPS: { 1 / time_took } )
162+ bboxes = self .yolox2bboxes_msgs (dets [:, :4 ], final_scores , final_cls_inds , COCO_CLASSES , msg .header )
163+ self .get_logger ().info (f'bboxes: { bboxes } )
164+ try :
165+ bboxes = self .yolox2bboxes_msgs (dets [:, :4 ], final_scores , final_cls_inds , COCO_CLASSES , msg .header )
166+
167+ if (self .imshow_isshow ):
168+ cv2 .imshow ("YOLOX" ,origin_img )
169+ cv2 .waitKey (1 )
170+
171+ except :
172+ self .get_logger ().info ('No object detected' )
173+ if (self .imshow_isshow ):
174+ cv2 .imshow ("YOLOX" ,origin_img )
175+ cv2 .waitKey (1 )
176+
177+ self .pub .publish (bboxes )
178+ self .pub_image .publish (self .bridge .cv2_to_imgmsg (img_rgb ,"bgr8" ))
179+
180+ except :
181+ pass
182+
183+ # origin_img is output
184+
185+ def ros_main (args = None ) -> None :
186+ rclpy .init (args = args )
187+
188+ yolox_ros_class = yolox_ros ()
189+ rclpy .spin (yolox_ros_class )
190+
191+ yolox_ros_class .destroy_node ()
192+ rclpy .shutdown ()
193+ cv2 .destroyAllWindows ()
194+
195+ if __name__ == "__main__" :
196+ ros_main ()
0 commit comments