diff --git a/inference.py b/inference.py index 9bb6463..8519073 100644 --- a/inference.py +++ b/inference.py @@ -175,7 +175,30 @@ class YoloProcessing: return box_centers, box_scales, objness, class_pred - def postprocessing(self, endnodes): + def process_to_picture(self, endnodes, data): + logits = self.postprocessing(endnodes) + self.visualize_image(logits, data) + + + def visualize_image(self, logits, data): + labels = data.get_labels("data/daria_labels.json") + image = visualize_boxes_and_labels_on_image_array( + data.dataset[0], + logits['detection_boxes'].numpy()[0], + logits['detection_classes'][0], + logits['detection_scores'].numpy()[0], + labels, + use_normalized_coordinates=True, + max_boxes_to_draw=100, + min_score_thresh=.5, + agnostic_mode=False, + line_thickness=4) + + Image.fromarray(np.uint8(image)).save('/home/maintenance/test.png') + print("Successfully saved image") + + + def postprocessing(self, endnodes, count): """ endnodes is a list of 3 output tensors: endnodes[0] - stride 32 of input @@ -235,6 +258,7 @@ class YoloProcessing: max_total_size=100) + # adding offset to the class prediction and cast to integer def translate_coco_2017_to_2014(nmsed_classes): return np.vectorize(COCO_17_14.get)(nmsed_classes).astype(np.int32) @@ -242,6 +266,8 @@ class YoloProcessing: nmsed_classes = tf.cast(tf.add(nmsed_classes, labels_offset), tf.int16) nmsed_classes = translate_coco_2017_to_2014(nmsed_classes) + print(count) + return {'detection_boxes': nmsed_boxes, 'detection_scores': nmsed_scores, 'detection_classes': nmsed_classes, @@ -336,53 +362,43 @@ def test_async_yolo5(): fps = 0 now = time.time() - for i in range(1000): + for i in range(100): fps += 1 if now + 1 < time.time(): - print(fps) fps = 0 now = time.time() hailo.hailo_input(data.dataset) out = None while(out == None): - time.sleep(0.01) + time.sleep(0.0001) out = hailo.hailo_output() + + Thread(target=processor.postprocessing, args=[out, i]).start() hailo.stop_hailo_thread() - logits = processor.postprocessing(out) - - - labels = data.get_labels("data/daria_labels.json") - image = visualize_boxes_and_labels_on_image_array( - data.dataset[0], - logits['detection_boxes'].numpy()[0], - logits['detection_classes'][0], - logits['detection_scores'].numpy()[0], - labels, - use_normalized_coordinates=True, - max_boxes_to_draw=100, - min_score_thresh=.5, - agnostic_mode=False, - line_thickness=4) - - Image.fromarray(np.uint8(image)).save('/home/maintenance/test.png') - print("Successfully saved image") - - def test_process_yolo5(): imageMeta = ImageMeta(640, 640, 3) - processor = YoloProcessing(imageMeta, classes=3) + processor = YoloProcessing(imageMeta, classes=4) data = DataHandler('./data', imageMeta) data.load_data(processor.preproc) hailo = HailoHandler('hef/yolov5m_daria.hef') - out = hailo.run_hailo(data.dataset) - logits = processor.postprocessing(out) + now = time.time() + fps = 0 + for i in range(100): + fps += 1 + if now + 1 < time.time(): + print(fps) + fps = 0 + now = time.time() + + out = hailo.run_hailo(data.dataset) + logits = processor.postprocessing(out) labels = data.get_labels("data/daria_labels.json") @@ -398,7 +414,6 @@ def test_process_yolo5(): agnostic_mode=False, line_thickness=4) - Image.fromarray(np.uint8(image)).save('/home/maintenance/test.png') print("Successfully saved image") if __name__ == "__main__": diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..ccc4f07 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +vision_msgs diff --git a/ros_inference.py b/ros_inference.py index bbd10b2..5593414 100644 --- a/ros_inference.py +++ b/ros_inference.py @@ -2,9 +2,11 @@ import json import os import io import time +import copy from PIL import Image from threading import Thread +from multiprocessing import Process import ipdb @@ -24,6 +26,7 @@ from tensorflow.image import combined_non_max_suppression import rclpy from rclpy.node import Node +from std_msgs.msg import String from sensor_msgs.msg import Image as ImageMsg from vision_msgs.msg import Detection2DArray, Detection2D, BoundingBox2D, ObjectHypothesisWithPose from geometry_msgs.msg import Pose2D @@ -190,16 +193,10 @@ class YoloProcessing: box_scales = (raw_box_scales * 2) ** 2 * anchors_for_stride # dim [N, HxW, 3, 2] return box_centers, box_scales, objness, class_pred - - def process_to_picture(self, endnodes, data): - logits = self.postprocessing(endnodes) - self.visualize_image(logits, data) - - - def visualize_image(self, logits, data): - labels = data.get_labels("data/daria_labels.json") + def visualize_image(self, logits, image): + labels = get_labels("data/daria_labels.json") image = visualize_boxes_and_labels_on_image_array( - data.dataset[0], + image, logits['detection_boxes'].numpy()[0], logits['detection_classes'][0], logits['detection_scores'].numpy()[0], @@ -211,7 +208,7 @@ class YoloProcessing: line_thickness=4) Image.fromarray(np.uint8(image)).save('/home/maintenance/test.png') - print("Successfully saved image") + Image.fromarray(np.uint8(image)).show() def postprocessing(self, endnodes): @@ -292,7 +289,6 @@ class YoloProcessing: class HailoHandler: def __init__(self, hef_path='hef/yolov5m.hef'): target = PcieDevice() - self.hef = HEF(hef_path) # Configure network groups @@ -326,6 +322,7 @@ class HailoHandler: self.hailo_async = True self.hailo_block = False self.input_data = None + self._infer_results = None self.hailo_thread = Thread(target=self._hailo_async) self.hailo_thread.start() @@ -339,26 +336,25 @@ class HailoHandler: def _hailo_async_loop(self, infer_pipeline): while self.hailo_async: if(not self.hailo_block and type(self.input_data) != type(None)): - self.infer_results = None + self._infer_results = None self.hailo_block = True infer_results = infer_pipeline.infer(self.input_data) - self.infer_results = [infer_results[i.name] for i in self.output_vstream_infos] + self._infer_results = [infer_results[i.name] for i in self.output_vstream_infos] self.input_data = None self.hailo_block = False def hailo_input(self, input_data): while self.hailo_block: - time.sleep(0.01) + time.sleep(0.001) self.hailo_block = True self.input_data = input_data self.input_data = {self.input_vstream_info.name: input_data} - self.infer_results = None self.hailo_block = False def hailo_output(self): while self.hailo_block: - time.sleep(0.01) - return self.infer_results + time.sleep(0.001) + return self._infer_results def stop_hailo_thread(self): @@ -368,49 +364,98 @@ class HailoHandler: class HailoNode(Node): def __init__(self): + + self._ros_init() + self._metadata_init() + self._object_init() + self._thread_init() + + def __del__(self): + self.hailo.stop_hailo_thread() + self._thread_run = False + self._post_process.join() + + def _ros_init(self): super().__init__('hailo_image_subscriber') - self.sub = self.create_subscription(ImageMsg, '/camera/color/image_raw', self.image_callback, 10) + self.sub = self.create_subscription(ImageMsg, '/r3_cam_left_0', self._image_callback, 10) self.pub = self.create_publisher(Detection2DArray, '/hailo_bounding_boxes', 10) + self.pub_ping = self.create_publisher(String, '/ping', 1) - self.bridge = CvBridge() - - # metadata init + def _metadata_init(self): + # TODO into yaml file classes = 3 self.image_meta = ImageMeta(640, 640, 3) self.processor = YoloProcessing(self.image_meta, classes) + self.hailo_hef = 'hef/yolov5m_daria.hef' - # hailo init - self.hailo = HailoHandler('hef/yolov5m_daria.hef') + def _object_init(self): + self.hailo = HailoHandler(self.hailo_hef) + self.bridge = CvBridge() + + def _thread_init(self): + self._thread_run = True + self._new_input = False + self.yolo_image = None self.hailo.start_hailo_thread() - def image_callback(self, data): - img = self.convert(data) - self.image_infer(img) + self.detections = None + self.detections_new = False + self.detections_mutex = False + self._post_process = Thread(target=self._thread_postprocessing).start() + self.publish_thread = Thread(target=self._thread_publish).start() - def image_infer(self, image): + def _image_callback(self, ros_image): + image = self._convert_ros_to_pil(ros_image) + self.yolo_image = self._preprocess(image) + self.image_infer(self.yolo_image) + self._new_input = True + + def _preprocess(self, image): image = self.processor.preproc(image) - dataset = self.dataset_from_image(image) - self.hailo.hailo_input(dataset) + return self._dataset_from_image(image) - out = None - while(out == None): - time.sleep(0.0001) - out = self.hailo.hailo_output() + def image_infer(self, data): + self.hailo.hailo_input(data) - Thread(target=self._thread_postprocessing, args=[out]).start() + def _thread_postprocessing(self): + while self._thread_run: + output = None + while(output == None or not self._new_input): + time.sleep(0.001) + output = self.hailo.hailo_output() + + + self._new_input = False + now = time.time() + self.detections_mutex = True + self.detections = self.processor.postprocessing(output) + self.detections_new = True + self.detections_mutex = False + print("postprocessing time: ", time.time() - now) + self.processor.visualize_image(self.detections, self.yolo_image[0]) + def _thread_publish(self): + while self._thread_run: + while self.detections_mutex or not self.detections_new: + time.sleep(0.001) + self._publish_detection(self.detections) + self.detections_new = False - def _thread_postprocessing(self, out): - logits = self.processor.postprocessing(out) + def _publish_ping(self, msg="ping"): + s = String() + s.data = msg + self.pub_ping.publish(s) + + def _publish_detection(self, detections): labels = get_labels("data/daria_labels.json") detection_array = Detection2DArray() - for bb in range(len(logits['detection_boxes'].numpy()[0])): - boxes = logits['detection_boxes'].numpy()[0][bb] - classes = logits['detection_classes'][0][bb] - scores = logits['detection_scores'].numpy()[0][bb] + for bb in range(len(detections['detection_boxes'].numpy()[0])): + boxes = detections['detection_boxes'].numpy()[0][bb] + classes = detections['detection_classes'][0][bb] + scores = detections['detection_scores'].numpy()[0][bb] if(scores > 0.01): bb = BoundingBox2D() bb.center = Pose2D() @@ -424,8 +469,7 @@ class HailoNode(Node): self.pub.publish(detection_array) - # convert ros image to PIL image - def convert(self, ros_image): + def _convert_ros_to_pil(self, ros_image): try: img = self.bridge.imgmsg_to_cv2(ros_image, "rgb8") image = Image.fromarray(img) @@ -433,7 +477,7 @@ class HailoNode(Node): print(e) return image - def dataset_from_image(self, image): + def _dataset_from_image(self, image): dataset = np.zeros((1, self.image_meta.image_height, self.image_meta.image_width, self.image_meta.channels), @@ -441,77 +485,12 @@ class HailoNode(Node): dataset[0, :, :, :] = np.array(image) return dataset -def test_async_yolo5(): - imageMeta = ImageMeta(640, 640, 3) - processor = YoloProcessing(imageMeta, classes=3) - data = DataHandler('./data', imageMeta) - data.load_data(processor.preproc) - - hailo = HailoHandler('hef/yolov5m_daria.hef') - hailo.start_hailo_thread() - - fps = 0 - now = time.time() - for i in range(100): - fps += 1 - if now + 1 < time.time(): - fps = 0 - now = time.time() - - hailo.hailo_input(data.dataset) - out = None - while(out == None): - time.sleep(0.0001) - out = hailo.hailo_output() - - Thread(target=processor.postprocessing, args=[out]).start() - - hailo.stop_hailo_thread() - - -def test_process_yolo5(): - - imageMeta = ImageMeta(640, 640, 3) - processor = YoloProcessing(imageMeta, classes=4) - data = DataHandler('./data', imageMeta) - data.load_data(processor.preproc) - - hailo = HailoHandler('hef/yolov5m_daria.hef') - - now = time.time() - fps = 0 - for i in range(100): - fps += 1 - if now + 1 < time.time(): - print(fps) - fps = 0 - now = time.time() - - out = hailo.run_hailo(data.dataset) - logits = processor.postprocessing(out) - - - labels = data.get_labels("data/daria_labels.json") - image = visualize_boxes_and_labels_on_image_array( - data.dataset[0], - logits['detection_boxes'].numpy()[0], - logits['detection_classes'][0], - logits['detection_scores'].numpy()[0], - labels, - use_normalized_coordinates=True, - max_boxes_to_draw=100, - min_score_thresh=.5, - agnostic_mode=False, - line_thickness=4) - - print("Successfully saved image") def main(args=None): rclpy.init(args=args) - hailo_node = HailoNode() - rclpy.spin(hailo_node) + if __name__ == "__main__": main()