<ROS1> Jetson-inference를 사용한 실시간 객체 검출

DDOKKON·2024년 8월 24일
#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy
import rosnode
import cv2
import jetson_utils
from std_msgs.msg import String
from jetson_inference import detectNet
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import time

model_path = 'data_real'
gstreamer_pipeline = (
    "nvarguscamerasrc sensor-id=0 ! video/x-raw(memory:NVMM), width=(int)1920, height=(int)1080, framerate=(fraction)30/1 ! "
    "nvvidconv flip-method=2 ! video/x-raw, width=(int)640, height=(int)480, format=(string)BGRx ! videoconvert ! "
    "video/x-raw, format=(string)BGR ! appsink"
)

cap = cv2.VideoCapture(gstreamer_pipeline)
if not cap.isOpened():
    print("Failed to open CSI camera (csi://0)")
    exit()

net = detectNet(argv=[
    '--model=/home/farm/catkin_ws/src/ros_vision/src/models/'+model_path+'/ssd-mobilenet.onnx',
    '--labels=/home/farm/catkin_ws/src/ros_vision/src/models/'+model_path+'/labels.txt',
    '--input-blob=input_0',
    '--output-cvg=scores',
    '--output-bbox=boxes'
])

rospy.init_node('ros_vision_detectnet')
data_pub = rospy.Publisher('inference', String, queue_size=10)
data_list = ["", 0]
data = ""
current_spent_time = 0
previous_spent_time = 0

opencv_publisher = rospy.Publisher('opencv_topic', Image, queue_size=10)
inference_publisher = rospy.Publisher('inference_topic', Image, queue_size=10)
bridge = CvBridge()

while not rospy.is_shutdown():
    current_spent_time = time.time()
    _, frame = cap.read()

    opencv_topic = bridge.cv2_to_imgmsg(frame, "bgr8")
    opencv_publisher.publish(opencv_topic)
    
    img_cuda = jetson_utils.cudaFromNumpy(frame)
    detections = net.Detect(img_cuda, overlay="box,labels,conf")

    for detection in detections:
        confidence = detection.Confidence
        left = int(detection.Left)
        top = int(detection.Top)
        right = int(detection.Right)
        bottom = int(detection.Bottom)
        
        cv2.rectangle(frame, (left, top), (right, bottom), (0, 255, 0), 2)
        
        label = net.GetClassDesc(detection.ClassID)
        text = "%s (%.1f%%)" % (label, confidence * 100)
        
        cv2.putText(frame, text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        data_list[0] = label
        data_list[1] = confidence * 100
        data = ','.join(map(str, data_list))
        
        if confidence > 0.65:
            if 'stop' in data and current_spent_time - previous_spent_time > 20 and confidence > 0.997:
                data_pub.publish(data)
                rospy.loginfo(data)  #240814 modified
                rospy.loginfo('정지 표지판을 감지했다는 토픽을 보냈습니다.')
                previous_spent_time = current_spent_time

            elif 'human' in data and current_spent_time - previous_spent_time > 8:
                data_pub.publish(data)
                rospy.loginfo(data)  #240814 modified
                rospy.loginfo('보행자를 감지했다는 토픽을 보냈습니다.')
                previous_spent_time = current_spent_time

    inference_topic = bridge.cv2_to_imgmsg(frame, "bgr8")
    inference_publisher.publish(inference_topic)
    
    if cv2.waitKey(1) == ord('q'):
        break

try:
    rosnode.kill_nodes(['ros_detectnet_subscriber'])
except:
    rospy.loginfo("Couldn't kill 'ros_detectnet_subscriber'")

try:
    rosnode.kill_nodes(['ros_detectnet'])
except:
    rospy.loginfo("Couldn't kill 'ros_detectnet'")

cv2.destroyAllWindows()

토픽으로 보내기 전에 넘파이이미지배열에서 ROS이미지메시지로 변형시켜줘야함. (cv2_to_imgmsg)

ROS에서 토픽으로 받은 이미지를 opencv로 띄워주거나 처리 하려면 당연히 다시 넘파이 배열로 변환 먼저 해야함.

profile
School of Electronic Engineering. KNU

0개의 댓글