From f22b253d7adf1062e9802348be2d3a8c8241d98a Mon Sep 17 00:00:00 2001 From: tmralmeida Date: Wed, 27 May 2020 20:20:18 +0100 Subject: [PATCH] visualization_bboxes --- CMakeLists.txt | 1 + src/node_detectnet.cpp | 450 ++++++++++++++++++++++------------------- 2 files changed, 244 insertions(+), 207 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 25fdddd..67ed998 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -9,6 +9,7 @@ find_package(catkin REQUIRED COMPONENTS sensor_msgs vision_msgs std_msgs + cv_bridge ) find_package(jetson-utils REQUIRED) diff --git a/src/node_detectnet.cpp b/src/node_detectnet.cpp index 933a6ec..26c7bff 100644 --- a/src/node_detectnet.cpp +++ b/src/node_detectnet.cpp @@ -23,8 +23,10 @@ #include #include +#include #include #include +#include #include #include @@ -35,226 +37,260 @@ // globals -detectNet* net = NULL; -imageConverter* cvt = NULL; +detectNet *net = NULL; +imageConverter *cvt = NULL; -ros::Publisher* detection_pub = NULL; +ros::Publisher *detection_pub = NULL; +ros::Publisher *image_pub = NULL; vision_msgs::VisionInfo info_msg; // callback triggered when a new subscriber connected to vision_info topic -void info_connect( const ros::SingleSubscriberPublisher& pub ) -{ - ROS_INFO("new subscriber '%s' connected to vision_info topic '%s', sending VisionInfo msg", pub.getSubscriberName().c_str(), pub.getTopic().c_str()); - pub.publish(info_msg); +void info_connect(const ros::SingleSubscriberPublisher &pub) { + ROS_INFO("new subscriber '%s' connected to vision_info topic '%s', sending VisionInfo msg", + pub.getSubscriberName().c_str(), pub.getTopic().c_str()); + pub.publish(info_msg); } // input image subscriber callback -void img_callback( const sensor_msgs::ImageConstPtr& input ) -{ - // convert the image to reside on GPU - if( !cvt || !cvt->Convert(input) ) - { - ROS_INFO("failed to convert %ux%u %s image", input->width, input->height, input->encoding.c_str()); - return; - } - - // classify the image - detectNet::Detection* detections = NULL; - - const int numDetections = net->Detect(cvt->ImageGPU(), cvt->GetWidth(), cvt->GetHeight(), &detections, detectNet::OVERLAY_NONE); - - // verify success - if( numDetections < 0 ) - { - ROS_ERROR("failed to run object detection on %ux%u image", input->width, input->height); - return; - } - - // if objects were detected, send out message - if( numDetections > 0 ) - { - ROS_INFO("detected %i objects in %ux%u image", numDetections, input->width, input->height); +void img_callback(const sensor_msgs::ImageConstPtr &input) { + // convert the image to reside on GPU + if (!cvt || !cvt->Convert(input)) { + ROS_INFO("failed to convert %ux%u %s image", input->width, input->height, input->encoding.c_str()); + return; + } + + // classify the image + detectNet::Detection *detections = NULL; + + const int numDetections = net->Detect(cvt->ImageGPU(), cvt->GetWidth(), cvt->GetHeight(), &detections, + detectNet::OVERLAY_NONE); + + // verify success + if (numDetections < 0) { + ROS_ERROR("failed to run object detection on %ux%u image", input->width, input->height); + return; + } + + // if objects were detected, send out message + vision_msgs::Detection2DArray msg; + msg.header.stamp = ros::Time::now(); + if (numDetections > 0) { + ROS_INFO("detected %i objects in %ux%u image", numDetections, input->width, input->height); + + // create a detection for each bounding box + for (int n = 0; n < numDetections; n++) { + detectNet::Detection *det = detections + n; + + printf("object %i class #%u (%s) confidence=%f\n", n, det->ClassID, net->GetClassDesc(det->ClassID), + det->Confidence); + printf("object %i bounding box (%f, %f) (%f, %f) w=%f h=%f\n", n, det->Left, det->Top, det->Right, + det->Bottom, det->Width(), det->Height()); + + // create a detection sub-message + vision_msgs::Detection2D detMsg; + detMsg.header.stamp = ros::Time::now(); + + detMsg.bbox.size_x = det->Width(); + detMsg.bbox.size_y = det->Height(); + + float cx, cy; + det->Center(&cx, &cy); + + detMsg.bbox.center.x = cx; + detMsg.bbox.center.y = cy; + + detMsg.bbox.center.theta = 0.0f; // TODO optionally output object image + + // create classification hypothesis + vision_msgs::ObjectHypothesisWithPose hyp; + + hyp.id = det->ClassID; + hyp.score = det->Confidence; + + detMsg.results.push_back(hyp); + msg.detections.push_back(detMsg); + } + + // publish the detection message + detection_pub->publish(msg); + } + sensor_msgs::Image output; + output.data = input->data; + output.step = input->step; + output.width = input->width; + output.header = input->header; + output.height = input->height; + output.encoding = input->encoding; + output.is_bigendian = input->is_bigendian; + + cv_bridge::CvImagePtr cv_ptr; + try { + cv_ptr = cv_bridge::toCvCopy(output, sensor_msgs::image_encodings::BGR8); + } + catch (cv_bridge::Exception &e) { + ROS_ERROR("cv_bridge exception: %s", e.what()); + return; + } + for (const auto &detection : msg.detections) { + const auto bbox = detection.bbox; + cv::Point upper_left(bbox.center.x - bbox.size_x / 2, bbox.center.y - bbox.size_y / 2); + cv::Point class_pst(25 + bbox.center.x - bbox.size_x / 2, bbox.center.y - bbox.size_y / 2); + cv::Point lower_right(bbox.center.x + bbox.size_x / 2, bbox.center.y + bbox.size_y / 2); + + std::stringstream textStream_confidence; + std::stringstream textStream_class; + float* color; - // create a detection for each bounding box - vision_msgs::Detection2DArray msg; - - for( int n=0; n < numDetections; n++ ) - { - detectNet::Detection* det = detections + n; - - printf("object %i class #%u (%s) confidence=%f\n", n, det->ClassID, net->GetClassDesc(det->ClassID), det->Confidence); - printf("object %i bounding box (%f, %f) (%f, %f) w=%f h=%f\n", n, det->Left, det->Top, det->Right, det->Bottom, det->Width(), det->Height()); - - // create a detection sub-message - vision_msgs::Detection2D detMsg; - - detMsg.bbox.size_x = det->Width(); - detMsg.bbox.size_y = det->Height(); - - float cx, cy; - det->Center(&cx, &cy); - - detMsg.bbox.center.x = cx; - detMsg.bbox.center.y = cy; - - detMsg.bbox.center.theta = 0.0f; // TODO optionally output object image - - // create classification hypothesis - vision_msgs::ObjectHypothesisWithPose hyp; - - hyp.id = det->ClassID; - hyp.score = det->Confidence; - - detMsg.results.push_back(hyp); - msg.detections.push_back(detMsg); - } - - // publish the detection message - detection_pub->publish(msg); - } + for (const auto hyp : detection.results) { + textStream_confidence << std::setprecision(2) << hyp.score; + textStream_class << net->GetClassDesc(hyp.id); + color = net->GetClassColor(hyp.id); + } + cv::rectangle(cv_ptr->image, upper_left, lower_right, cv::Scalar(*color, *color, *color), 2); + cv::putText(cv_ptr->image, textStream_confidence.str(), upper_left, cv::FONT_HERSHEY_PLAIN, + 0.6,cv::Scalar(*color, *color, *color), 1); + cv::putText(cv_ptr->image, textStream_class.str(), class_pst, cv::FONT_HERSHEY_PLAIN, + 0.6, cv::Scalar(*color, *color, *color), 1); + } + image_pub->publish(cv_ptr->toImageMsg()); } // node main loop -int main(int argc, char **argv) -{ - ros::init(argc, argv, "detectnet"); - - ros::NodeHandle nh; - ros::NodeHandle private_nh("~"); - - /* - * retrieve parameters - */ - std::string class_labels_path; - std::string prototxt_path; - std::string model_path; - std::string model_name; - - bool use_model_name = false; - - // determine if custom model paths were specified - if( !private_nh.getParam("prototxt_path", prototxt_path) || - !private_nh.getParam("model_path", model_path) ) - { - // without custom model, use one of the built-in pretrained models - private_nh.param("model_name", model_name, "ssd-mobilenet-v2"); - use_model_name = true; - } - - // set mean pixel and threshold defaults - float mean_pixel = 0.0f; - float threshold = 0.5f; - - private_nh.param("mean_pixel_value", mean_pixel, mean_pixel); - private_nh.param("threshold", threshold, threshold); - - - /* - * load object detection network - */ - if( use_model_name ) - { - // determine which built-in model was requested - detectNet::NetworkType model = detectNet::NetworkTypeFromStr(model_name.c_str()); - - if( model == detectNet::CUSTOM ) - { - ROS_ERROR("invalid built-in pretrained model name '%s', defaulting to pednet", model_name.c_str()); - model = detectNet::PEDNET; - } - - // create network using the built-in model - net = detectNet::Create(model); - } - else - { - // get the class labels path (optional) - private_nh.getParam("class_labels_path", class_labels_path); - - // create network using custom model paths - net = detectNet::Create(prototxt_path.c_str(), model_path.c_str(), mean_pixel, class_labels_path.c_str(), threshold); - } - - if( !net ) - { - ROS_ERROR("failed to load detectNet model"); - return 0; - } - - - /* - * create the class labels parameter vector - */ - std::hash model_hasher; // hash the model path to avoid collisions on the param server - std::string model_hash_str = std::string(net->GetModelPath()) + std::string(net->GetClassPath()); - const size_t model_hash = model_hasher(model_hash_str); - - ROS_INFO("model hash => %zu", model_hash); - ROS_INFO("hash string => %s", model_hash_str.c_str()); - - // obtain the list of class descriptions - std::vector class_descriptions; - const uint32_t num_classes = net->GetNumClasses(); - - for( uint32_t n=0; n < num_classes; n++ ) - class_descriptions.push_back(net->GetClassDesc(n)); - - // create the key on the param server - std::string class_key = std::string("class_labels_") + std::to_string(model_hash); - private_nh.setParam(class_key, class_descriptions); - - // populate the vision info msg - std::string node_namespace = private_nh.getNamespace(); - ROS_INFO("node namespace => %s", node_namespace.c_str()); - - info_msg.database_location = node_namespace + std::string("/") + class_key; - info_msg.database_version = 0; - info_msg.method = net->GetModelPath(); - - ROS_INFO("class labels => %s", info_msg.database_location.c_str()); - - - /* - * create an image converter object - */ - cvt = new imageConverter(); - - if( !cvt ) - { - ROS_ERROR("failed to create imageConverter object"); - return 0; - } - - - /* - * advertise publisher topics - */ - ros::Publisher pub = private_nh.advertise("detections", 25); - detection_pub = &pub; // we need to publish from the subscriber callback - - // the vision info topic only publishes upon a new connection - ros::Publisher info_pub = private_nh.advertise("vision_info", 1, (ros::SubscriberStatusCallback)info_connect); - - - /* - * subscribe to image topic - */ - //image_transport::ImageTransport it(nh); // BUG - stack smashing on TX2? - //image_transport::Subscriber img_sub = it.subscribe("image", 1, img_callback); - ros::Subscriber img_sub = private_nh.subscribe("image_in", 5, img_callback); - - - /* - * wait for messages - */ - ROS_INFO("detectnet node initialized, waiting for messages"); - - ros::spin(); - - return 0; -} +int main(int argc, char **argv) { + ros::init(argc, argv, "detectnet"); + + ros::NodeHandle nh; + ros::NodeHandle private_nh("~"); + + /* + * retrieve parameters + */ + std::string class_labels_path; + std::string prototxt_path; + std::string model_path; + std::string model_name; + + bool use_model_name = false; + + // determine if custom model paths were specified + if (!private_nh.getParam("prototxt_path", prototxt_path) || + !private_nh.getParam("model_path", model_path)) { + // without custom model, use one of the built-in pretrained models + private_nh.param("model_name", model_name, "ssd-mobilenet-v2"); + use_model_name = true; + } + + // set mean pixel and threshold defaults + float mean_pixel = 0.0f; + float threshold = 0.5f; + + private_nh.param("mean_pixel_value", mean_pixel, mean_pixel); + private_nh.param("threshold", threshold, threshold); + + + /* + * load object detection network + */ + if (use_model_name) { + // determine which built-in model was requested + detectNet::NetworkType model = detectNet::NetworkTypeFromStr(model_name.c_str()); + + if (model == detectNet::CUSTOM) { + ROS_ERROR("invalid built-in pretrained model name '%s', defaulting to pednet", model_name.c_str()); + model = detectNet::PEDNET; + } + + // create network using the built-in model + net = detectNet::Create(model); + } else { + // get the class labels path (optional) + private_nh.getParam("class_labels_path", class_labels_path); + + // create network using custom model paths + net = detectNet::Create(prototxt_path.c_str(), model_path.c_str(), mean_pixel, class_labels_path.c_str(), + threshold); + } + + if (!net) { + ROS_ERROR("failed to load detectNet model"); + return 0; + } + + + /* + * create the class labels parameter vector + */ + std::hash model_hasher; // hash the model path to avoid collisions on the param server + std::string model_hash_str = std::string(net->GetModelPath()) + std::string(net->GetClassPath()); + const size_t model_hash = model_hasher(model_hash_str); + + ROS_INFO("model hash => %zu", model_hash); + ROS_INFO("hash string => %s", model_hash_str.c_str()); + + // obtain the list of class descriptions + std::vector class_descriptions; + const uint32_t num_classes = net->GetNumClasses(); + + for (uint32_t n = 0; n < num_classes; n++) + class_descriptions.push_back(net->GetClassDesc(n)); + + // create the key on the param server + std::string class_key = std::string("class_labels_") + std::to_string(model_hash); + private_nh.setParam(class_key, class_descriptions); + + // populate the vision info msg + std::string node_namespace = private_nh.getNamespace(); + ROS_INFO("node namespace => %s", node_namespace.c_str()); + + info_msg.database_location = node_namespace + std::string("/") + class_key; + info_msg.database_version = 0; + info_msg.method = net->GetModelPath(); + + ROS_INFO("class labels => %s", info_msg.database_location.c_str()); + + + /* + * create an image converter object + */ + cvt = new imageConverter(); + + if (!cvt) { + ROS_ERROR("failed to create imageConverter object"); + return 0; + } + + + /* + * advertise publisher topics + */ + ros::Publisher pub = private_nh.advertise("detections", 25); + detection_pub = &pub; // we need to publish from the subscriber callback + ros::Publisher img_publisher = private_nh.advertise("/camera/image_detections", 1); + image_pub = &img_publisher; + + // the vision info topic only publishes upon a new connection + ros::Publisher info_pub = private_nh.advertise("vision_info", 1, + (ros::SubscriberStatusCallback) info_connect); + + + /* + * subscribe to image topictoImageMsg + */ + //image_transport::ImageTransport it(nh); // BUG - stack smashing on TX2? + //image_transport::Subscriber img_sub = it.subscribe("image", 1, img_callback); + ros::Subscriber img_sub = private_nh.subscribe("image_in", 5, img_callback); + + + /* + * wait for messages + */ + ROS_INFO("detectnet node initialized, waiting for messages"); + + ros::spin(); + return 0; +} \ No newline at end of file