dnn_detect.cpp

Go to the documentation of this file.

/*
 * Copyright (c) 2017, Ubiquity Robotics
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 * The views and conclusions contained in the software and documentation are
 * those of the authors and should not be interpreted as representing official
 * policies, either expressed or implied, of the FreeBSD Project.
 *
 */

#include <assert.h>
#include <sys/time.h>
#include <unistd.h>

#include <ros/ros.h>
#include <image_transport/image_transport.h>
#include <cv_bridge/cv_bridge.h>

#include "dnn_detect/DetectedObject.h"
#include "dnn_detect/DetectedObjectArray.h"
#include "dnn_detect/Detect.h"

#include <opencv2/highgui.hpp>
#include <opencv2/dnn.hpp>
#include <opencv2/calib3d.hpp>

#include <list>
#include <string>
#include <boost/algorithm/string.hpp>
#include <boost/format.hpp>

#include <thread>
#include <mutex>
#include <condition_variable>

using namespace std;
using namespace cv;

std::condition_variable cond;
std::mutex mutx;

class DnnNode {
  private:
    ros::Publisher results_pub;

    image_transport::ImageTransport it;
    image_transport::Subscriber img_sub;

    // if set, we publish the images that contain objects
    bool publish_images;

    int frame_num;
    float min_confidence;
    int im_size;
    int rotate_flag;
    float scale_factor;
    float mean_val;
    std::vector<std::string> class_names;

    image_transport::Publisher image_pub;

    cv::dnn::Net net;
    cv::Mat resized_image;
    cv::Mat rotated_image;

    bool single_shot;
    volatile bool triggered;
    volatile bool processed;

    dnn_detect::DetectedObjectArray results;

    ros::ServiceServer detect_srv;

    bool trigger_callback(dnn_detect::Detect::Request &req,
                          dnn_detect::Detect::Response &res);

    void image_callback(const sensor_msgs::ImageConstPtr &msg);

  public:
    DnnNode(ros::NodeHandle &nh);
};

bool DnnNode::  trigger_callback(dnn_detect::Detect::Request &req,
                                 dnn_detect::Detect::Response &res)
{
    ROS_INFO("Got service request");
    triggered = true;

    std::unique_lock<std::mutex> lock(mutx);

    while (!processed) {
      cond.wait(lock);
    }
    res.result = results;
    processed = false;
    return true;
}


void DnnNode::image_callback(const sensor_msgs::ImageConstPtr & msg)
{
    if (single_shot && !triggered) {
        return;
    }
    triggered = false;

    ROS_INFO("Got image %d", msg->header.seq);
    frame_num++;

    cv_bridge::CvImagePtr cv_ptr;

    try {
        cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);

        int w = cv_ptr->image.cols;
        int h = cv_ptr->image.rows;

        if (rotate_flag >= 0) {
          cv::rotate(cv_ptr->image, rotated_image, rotate_flag);
          rotated_image.copyTo(cv_ptr->image);
        }

        cv::resize(cv_ptr->image, resized_image, cvSize(im_size, im_size));
        cv::Mat blob = cv::dnn::blobFromImage(resized_image, scale_factor,
          cvSize(im_size, im_size), mean_val, false);

        net.setInput(blob, "data");
        cv::Mat objs = net.forward("detection_out");

        cv::Mat detectionMat(objs.size[2], objs.size[3], CV_32F,
                             objs.ptr<float>());

        std::unique_lock<std::mutex> lock(mutx);
        results.header.frame_id = msg->header.frame_id;
        results.objects.clear();

        for(int i = 0; i < detectionMat.rows; i++) {

            float confidence = detectionMat.at<float>(i, 2);
            if (confidence > min_confidence) {
                int object_class = (int)(detectionMat.at<float>(i, 1));

                int x_min = static_cast<int>(detectionMat.at<float>(i, 3) * w);
                int y_min = static_cast<int>(detectionMat.at<float>(i, 4) * h);
                int x_max = static_cast<int>(detectionMat.at<float>(i, 5) * w);
                int y_max = static_cast<int>(detectionMat.at<float>(i, 6) * h);

                std::string class_name;
                if (object_class >= class_names.size()) {
                     class_name = "unknown";
                     ROS_ERROR("Object class %d out of range of class names",
                               object_class);
                }
                else {
                     class_name = class_names[object_class];
                }
                std::string label = str(boost::format{"%1% %2%"} %
                                        class_name % confidence);

                ROS_INFO("%s", label.c_str());
                dnn_detect::DetectedObject obj;
                obj.class_name = class_name;
                obj.confidence = confidence;
                obj.x_min = x_min;
                obj.x_max = x_max;
                obj.y_min = y_min;
                obj.y_max = y_max;
                results.objects.push_back(obj);

                Rect object(x_min, y_min, x_max-x_min, y_max-y_min);

                rectangle(cv_ptr->image, object, Scalar(0, 255, 0));
                int baseline=0;
                cv::Size text_size = cv::getTextSize(label,
                                     FONT_HERSHEY_SIMPLEX, 0.75, 2, &baseline);
                putText(cv_ptr->image, label, Point(x_min, y_min-text_size.height),
                        FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 255, 0));
            }
        }

        results_pub.publish(results);

        image_pub.publish(cv_ptr->toImageMsg());

    }
    catch(cv_bridge::Exception & e) {
        ROS_ERROR("cv_bridge exception: %s", e.what());
    }
    catch(cv::Exception & e) {
        ROS_ERROR("cv exception: %s", e.what());
    }
    ROS_DEBUG("Notifying condition variable");
    processed = true;
    cond.notify_all();
}

DnnNode::DnnNode(ros::NodeHandle & nh) : it(nh)
{
    frame_num = 0;

    std::string dir;
    std::string proto_net_file;
    std::string caffe_model_file;
    std::string classes("background,"
       "aeroplane,bicycle,bird,boat,bottle,bus,car,cat,chair,"
       "cow,diningtable,dog,horse,motorbike,person,pottedplant,"
       "sheep,sofa,train,tvmonitor");

    nh.param<bool>("single_shot", single_shot, false);

    nh.param<bool>("publish_images", publish_images, false);
    nh.param<string>("data_dir", dir, "");
    nh.param<string>("protonet_file", proto_net_file,
                     "MobileNetSSD_deploy.prototxt.txt");
    nh.param<string>("caffe_model_file", caffe_model_file,
                     "MobileNetSSD_deploy.caffemodel");
    nh.param<float>("min_confidence", min_confidence, 0.2);
    nh.param<int>("im_size", im_size, 300);
    nh.param<int>("rotate_flag", rotate_flag, -1);
    nh.param<float>("scale_factor", scale_factor, 0.007843f);
    nh.param<float>("mean_val", mean_val, 127.5f);
    nh.param<std::string>("class_names", classes, classes);

    boost::split(class_names, classes, boost::is_any_of(","));
    ROS_INFO("Read %d class names", (int)class_names.size());

    try {
        net = cv::dnn::readNetFromCaffe(dir + "/" + proto_net_file,
                                        dir + "/" + caffe_model_file);
    }
    catch(cv::Exception & e) {
        ROS_ERROR("cv exception: %s", e.what());
        exit(1);
    }

    triggered = false;

    detect_srv = nh.advertiseService("detect", &DnnNode::trigger_callback, this);

    results_pub =
        nh.advertise<dnn_detect::DetectedObjectArray>("/dnn_objects", 20);

    image_pub = it.advertise("/dnn_images", 1);

    img_sub = it.subscribe("/camera", 1,
                           &DnnNode::image_callback, this);

    ROS_INFO("DNN detection ready");
}

int main(int argc, char ** argv) {
    ros::init(argc, argv, "dnn_detect");
    ros::NodeHandle nh("~");

    DnnNode node = DnnNode(nh);
    ros::MultiThreadedSpinner spinner(2);
    spinner.spin();

    return 0;
}