stream_detection.cpp

/*
 * Copyright (c) 2017 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <cv_bridge/cv_bridge.h>
#include <iomanip>
#include <message_filters/subscriber.h>
#include <message_filters/time_synchronizer.h>
#include <object_msgs/ObjectInBox.h>
#include <object_msgs/ObjectsInBoxes.h>
#include <opencv2/opencv.hpp>
#include <ros/ros.h>

#define LINESPACING 20
#define DEFAULT_PARALLEL_FLAG 1
#define MOVEWINDOW 1000

// UI Enhancement constants
#define TEXT_FONT cv::FONT_HERSHEY_SIMPLEX
#define TEXT_SCALE 0.6
#define TEXT_THICKNESS 2
#define BOX_THICKNESS 3
#define SHADOW_OFFSET 2

// UI Helper Functions
cv::Scalar getConfidenceColor(double confidence) {
  if (confidence >= 0.8) return cv::Scalar(0, 255, 0);      // Green for high confidence
  else if (confidence >= 0.6) return cv::Scalar(0, 255, 255); // Yellow for medium confidence
  else return cv::Scalar(0, 165, 255);                      // Orange for low confidence
}

void drawTextWithShadow(cv::Mat& image, const std::string& text, cv::Point position, 
                       cv::Scalar color, int fontFace, double fontScale, int thickness) {
  // Draw shadow
  cv::putText(image, text, cv::Point(position.x + SHADOW_OFFSET, position.y + SHADOW_OFFSET), 
              fontFace, fontScale, cv::Scalar(0, 0, 0), thickness + 1);
  // Draw main text
  cv::putText(image, text, position, fontFace, fontScale, color, thickness);
}

void drawDetectionBox(cv::Mat& image, const std::string& label, cv::Point topLeft, cv::Point bottomRight, 
                     cv::Scalar boxColor, cv::Scalar textColor, int fontFace, double fontScale, int thickness) {
  // Draw main bounding box
  cv::rectangle(image, topLeft, bottomRight, boxColor, BOX_THICKNESS, cv::LINE_8, 0);
  
  // Draw label background
  int baseline = 0;
  cv::Size textSize = cv::getTextSize(label, fontFace, fontScale, thickness, &baseline);
  cv::Point labelTopLeft(topLeft.x, topLeft.y - textSize.height - 5);
  cv::Point labelBottomRight(topLeft.x + textSize.width + 10, topLeft.y);
  cv::rectangle(image, labelTopLeft, labelBottomRight, boxColor, -1);
  
  // Draw label text with shadow
  drawTextWithShadow(image, label, cv::Point(topLeft.x + 5, topLeft.y - 5), 
                    textColor, fontFace, fontScale, thickness);
}

int getFPS()
{
  static int fps = 0;
  static boost::posix_time::ptime duration_start = boost::posix_time::microsec_clock::local_time();
  static int frame_cnt = 0;

  frame_cnt++;

  boost::posix_time::ptime current = boost::posix_time::microsec_clock::local_time();
  boost::posix_time::time_duration msdiff = current - duration_start;

  if (msdiff.total_milliseconds() > 1000)
  {
    fps = frame_cnt;
    frame_cnt = 0;
    duration_start = current;
  }

  return fps;
}

int parallel_flag;
void syncCb(const sensor_msgs::ImageConstPtr& img, const object_msgs::ObjectsInBoxes::ConstPtr& objs_in_boxes)
{
  cv::Mat cvImage = cv_bridge::toCvShare(img, "bgr8")->image;
  int width = img->width;
  int height = img->height;

  // Draw header
  std::stringstream header_ss;
  header_ss << "F.Project 2022 - Live Object Detection Stream";
  drawTextWithShadow(cvImage, header_ss.str(), cv::Point(LINESPACING, LINESPACING), 
                    cv::Scalar(255, 255, 255), TEXT_FONT, TEXT_SCALE + 0.3, TEXT_THICKNESS + 1);

  for (auto obj : objs_in_boxes->objects_vector)
  {
    std::stringstream ss;
    ss << obj.object.object_name << ": " << std::fixed << std::setprecision(1) << obj.object.probability * 100 << '%';

    int xmin = obj.roi.x_offset;
    int ymin = obj.roi.y_offset;
    int w = obj.roi.width;
    int h = obj.roi.height;

    int xmax = ((xmin + w) < width) ? (xmin + w) : width;
    int ymax = ((ymin + h) < height) ? (ymin + h) : height;

    cv::Point left_top = cv::Point(xmin, ymin);
    cv::Point right_bottom = cv::Point(xmax, ymax);
    
    // Get confidence-based colors
    cv::Scalar boxColor = getConfidenceColor(obj.object.probability);
    cv::Scalar textColor = cv::Scalar(255, 255, 255); // White text for better contrast
    
    // Draw enhanced detection box
    drawDetectionBox(cvImage, ss.str(), left_top, right_bottom, 
                    boxColor, textColor, TEXT_FONT, TEXT_SCALE, TEXT_THICKNESS);
  }

  // Draw FPS with background
  std::stringstream fps_ss;
  int fps = getFPS();
  fps_ss << "FPS: " << fps;
  drawTextWithShadow(cvImage, fps_ss.str(), cv::Point(LINESPACING, LINESPACING + 30), 
                    cv::Scalar(0, 255, 0), TEXT_FONT, TEXT_SCALE, TEXT_THICKNESS);

  // Draw help text
  std::stringstream help_ss;
  help_ss << "Press ESC, Enter, Space, or 'q' to exit";
  drawTextWithShadow(cvImage, help_ss.str(), cv::Point(LINESPACING, cvImage.rows - 30), 
                    cv::Scalar(200, 200, 200), TEXT_FONT, TEXT_SCALE - 0.2, TEXT_THICKNESS - 1);

  if (parallel_flag == 0)
  {
    cv::imshow("F.Project 2022 - Live Detection (Single Device)", cvImage);
  }
  else
  {
    cv::namedWindow("F.Project 2022 - Live Detection (Multiple Devices)");
    cv::moveWindow("F.Project 2022 - Live Detection (Multiple Devices)", MOVEWINDOW, 0);
    cv::imshow("F.Project 2022 - Live Detection (Multiple Devices)", cvImage);
  }

  int key = cv::waitKey(5);
  if (key == 13 || key == 27 || key == 32 || key == 113)
  {
    ros::shutdown();
  }
}

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

  parallel_flag = DEFAULT_PARALLEL_FLAG;
  if (!nh.getParam("parallel_flag", parallel_flag))
  {
    ROS_WARN("param parallel_flag not set, use default");
  }
  ROS_INFO_STREAM("use parallel_flag = " << parallel_flag);

  if (parallel_flag == 0)
  {
    message_filters::Subscriber<sensor_msgs::Image> camSub(nh, "/camera/color/image_raw", 1);
    message_filters::Subscriber<object_msgs::ObjectsInBoxes> objSub(nh, "/movidius_ncs_nodelet/detected_objects_single",
                                                                    1);
    message_filters::TimeSynchronizer<sensor_msgs::Image, object_msgs::ObjectsInBoxes> sync(camSub, objSub, 60);
    sync.registerCallback(boost::bind(&syncCb, _1, _2));
    ros::spin();
  }
  else
  {
    message_filters::Subscriber<sensor_msgs::Image> camSub(nh, "/camera/color/image_raw", 1);
    message_filters::Subscriber<object_msgs::ObjectsInBoxes> objSub(
        nh, "/movidius_ncs_nodelet/detected_objects_multiple", 1);
    message_filters::TimeSynchronizer<sensor_msgs::Image, object_msgs::ObjectsInBoxes> sync(camSub, objSub, 60);
    sync.registerCallback(boost::bind(&syncCb, _1, _2));
    ros::spin();
  }

  return 0;
}