image_classification.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 <dirent.h>
#include <iomanip>
#include <random>
#include <string>
#include <vector>

#include <boost/date_time/posix_time/posix_time.hpp>
#include <cv_bridge/cv_bridge.h>
#include <object_msgs/ClassifyObject.h>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <ros/ros.h>

#define LINESPACING 30
#define MOVEWINDOW 1000
#define DEFAULT_PARALLEL_SIZE 1
#define DEFAULT_IMAGE_BASE_PATH "/opt/movidius/ncappzoo/data/images/"
#define DEFAULT_DEMO_MODE 0
#define DEFAULT_PARALLEL_FLAG 1

// UI Enhancement constants
#define TEXT_FONT cv::FONT_HERSHEY_SIMPLEX
#define TEXT_SCALE 0.7
#define TEXT_THICKNESS 2
#define BOX_THICKNESS 2
#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 drawTextWithBackground(cv::Mat& image, const std::string& text, cv::Point position, 
                           cv::Scalar textColor, cv::Scalar bgColor, int fontFace, double fontScale, int thickness) {
  int baseline = 0;
  cv::Size textSize = cv::getTextSize(text, fontFace, fontScale, thickness, &baseline);
  
  // Draw background rectangle
  cv::Point bgTopLeft(position.x - 5, position.y - textSize.height - 5);
  cv::Point bgBottomRight(position.x + textSize.width + 5, position.y + baseline + 5);
  cv::rectangle(image, bgTopLeft, bgBottomRight, bgColor, -1);
  
  // Draw text
  cv::putText(image, text, position, fontFace, fontScale, textColor, thickness);
}

std::vector<std::string> getImagePath(std::string image_dir)
{
  if (image_dir.back() != '/')
  {
    image_dir += "/";
  }

  std::vector<std::string> files;

  DIR* dir;
  struct dirent* ptr;

  if ((dir = opendir(image_dir.c_str())) == NULL)
  {
    std::cerr << "Open Dir error..." << std::endl;
    exit(1);
  }

  while ((ptr = readdir(dir)) != NULL)
  {
    if (strcmp(ptr->d_name, ".") == 0 || strcmp(ptr->d_name, "..") == 0)
      continue;
    else if (ptr->d_type == DT_REG)
      files.push_back(image_dir + ptr->d_name);
  }
  closedir(dir);

  return files;
}

int main(int argc, char** argv)
{
  ros::init(argc, argv, "movidius_ncs_example_classification");
  ros::NodeHandle n("~");
  ros::ServiceClient client;

  std::string image_base_path = DEFAULT_IMAGE_BASE_PATH;
  if (!n.getParam("image_base_path", image_base_path))
  {
    ROS_WARN("param image_base_path not set, use default");
  }
  ROS_INFO_STREAM("use image_base_path = " << image_base_path);
  std::vector<std::string> image_paths = getImagePath(image_base_path);

  int parallel_size = DEFAULT_PARALLEL_SIZE;
  if (!n.getParam("parallel_size", parallel_size))
  {
    ROS_WARN("param parallel_size not set, use default");
  }
  ROS_INFO_STREAM("use parallel_size = " << parallel_size);

  int demo_mode = DEFAULT_DEMO_MODE;
  if (!n.getParam("demo_mode", demo_mode))
  {
    ROS_WARN("param demo_mode not set, use default");
  }
  ROS_INFO_STREAM("use demo_mode = " << demo_mode);

  int parallel_flag = DEFAULT_PARALLEL_FLAG;
  if (!n.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)
  {
    client = n.serviceClient<object_msgs::ClassifyObject>("/movidius_ncs_image_single/classify_object");
  }
  else
  {
    client = n.serviceClient<object_msgs::ClassifyObject>("/movidius_ncs_image_multiple/classify_object");
  }

  if (demo_mode == 0)
  {
    object_msgs::ClassifyObject srv;
    srv.request.image_paths = image_paths;
    boost::posix_time::ptime start = boost::posix_time::microsec_clock::local_time();
    if (!client.call(srv))
    {
      ROS_ERROR("failed to call service ClassifyObject");
      return 1;
    }
    boost::posix_time::ptime end = boost::posix_time::microsec_clock::local_time();
    boost::posix_time::time_duration msdiff = end - start;

    for (unsigned int i = 0; i < srv.response.objects.size(); i++)
    {
      cv_bridge::CvImage cv_image;
      cv_image.image = cv::imread(image_paths[i]);
      cv_image.encoding = "bgr8";
      int cnt = 0;

      ROS_INFO("Classification result for image No.%u:", i + 1);
      
      // Draw header
      std::stringstream header_ss;
      header_ss << "Image " << (i + 1) << " Classification Results";
      drawTextWithBackground(cv_image.image, header_ss.str(), cv::Point(LINESPACING, LINESPACING), 
                            cv::Scalar(255, 255, 255), cv::Scalar(50, 50, 50), TEXT_FONT, TEXT_SCALE + 0.2, TEXT_THICKNESS);
      
      for (unsigned int j = 0; j < srv.response.objects[i].objects_vector.size(); j++)
      {
        std::stringstream ss;
        ss << srv.response.objects[i].objects_vector[j].object_name << ": "
           << std::fixed << std::setprecision(1) << srv.response.objects[i].objects_vector[j].probability * 100 << "%";

        ROS_INFO("%d: object: %s\nprobability: %lf%%", j, srv.response.objects[i].objects_vector[j].object_name.c_str(),
                 srv.response.objects[i].objects_vector[j].probability * 100);

        // Get confidence-based color
        cv::Scalar textColor = getConfidenceColor(srv.response.objects[i].objects_vector[j].probability);
        
        // Draw text with shadow for better visibility
        drawTextWithShadow(cv_image.image, ss.str(), cv::Point(LINESPACING, LINESPACING * (++cnt + 1)), 
                          textColor, TEXT_FONT, TEXT_SCALE, TEXT_THICKNESS);
      }

      cv::imshow("F.Project 2022 - Image Classification", cv_image.image);
      cv::waitKey(0);
    }
    ROS_INFO("inference %lu images during %ld ms", srv.response.objects.size(), msdiff.total_milliseconds());
  }
  else
  {
    while (1)
    {
      object_msgs::ClassifyObject srv;
      std::vector<int> random_index_list;
      for (int i = 0; i < parallel_size; i++)
      {
        std::random_device rd;
        std::default_random_engine engine(rd());
        std::uniform_int_distribution<> dis(0, image_paths.size() - 1);
        auto dice = std::bind(dis, engine);
        int random_index = dice();
        random_index_list.push_back(random_index);
        srv.request.image_paths.push_back(image_paths[random_index]);
      }

      if (!client.call(srv))
      {
        ROS_ERROR("failed to call service ClassifyObject");
        exit(1);
      }

      for (unsigned int i = 0; i < srv.response.objects.size(); i++)
      {
        cv_bridge::CvImage cv_image;
        cv_image.image = cv::imread(image_paths[random_index_list[i]]);
        cv_image.encoding = "bgr8";
        int cnt = 0;

        // Draw header
        std::stringstream header_ss;
        header_ss << "Real-time Classification Results";
        drawTextWithBackground(cv_image.image, header_ss.str(), cv::Point(LINESPACING, LINESPACING), 
                              cv::Scalar(255, 255, 255), cv::Scalar(50, 50, 50), TEXT_FONT, TEXT_SCALE + 0.2, TEXT_THICKNESS);
        
        for (unsigned int j = 0; j < srv.response.objects[i].objects_vector.size(); j++)
        {
          std::stringstream ss;
          ss << srv.response.objects[i].objects_vector[j].object_name << ": "
             << std::fixed << std::setprecision(1) << srv.response.objects[i].objects_vector[j].probability * 100 << "%";

          // Get confidence-based color
          cv::Scalar textColor = getConfidenceColor(srv.response.objects[i].objects_vector[j].probability);
          
          // Draw text with shadow for better visibility
          drawTextWithShadow(cv_image.image, ss.str(), cv::Point(LINESPACING, LINESPACING * (++cnt + 1)), 
                            textColor, TEXT_FONT, TEXT_SCALE, TEXT_THICKNESS);
        }

        if (parallel_flag == 0)
        {
          cv::imshow("F.Project 2022 - Real-time Classification (Single Device)", cv_image.image);
          cv::waitKey(20);
        }
        else
        {
          cv::namedWindow("F.Project 2022 - Real-time Classification (Multiple Devices)");
          cv::moveWindow("F.Project 2022 - Real-time Classification (Multiple Devices)", MOVEWINDOW, 0);
          cv::imshow("F.Project 2022 - Real-time Classification (Multiple Devices)", cv_image.image);
          cv::waitKey(20);
        }
      }
    }
  }

  return 0;
}