<commit_before>#include <ros/ros.h>
#include <cv_bridge/cv_bridge.h>
#include <geometry_msgs/Point.h>
#include <geometry_msgs/PointStamped.h>
#include <geometry_msgs/Pose.h>
#include <geometry_msgs/PoseArray.h>
#include <image_transport/image_transport.h>
#include <sensor_msgs/image_encodings.h>

#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/video/video.hpp>

class ShowImage {
public:
  ShowImage()
    : m_gamma(2048),
      RGB_WINDOW_NAME {"RGB image"},
      DEPTH_WINDOW_NAME {"Depth image"},
      BINARY_WINDOW_NAME {"Binary image"}
  {
    for(int i = 0; i < 2048; i++) {
      float v = i / 2048.0;
      v = std::pow(v, 3) * 6;
      m_gamma[i] = v * 6 * 256;
    }

    cv::namedWindow(RGB_WINDOW_NAME);
    cv::namedWindow(DEPTH_WINDOW_NAME);
    cv::namedWindow(BINARY_WINDOW_NAME);
  }

  ~ShowImage() {
    cv::destroyWindow(RGB_WINDOW_NAME);
    cv::destroyWindow(DEPTH_WINDOW_NAME);
    cv::destroyWindow(BINARY_WINDOW_NAME);
  }

  void showRGB(const cv::Mat& rgb) {
    imshow(RGB_WINDOW_NAME, rgb);
  }

  void showBinary(const cv::Mat& binary) {
    imshow(BINARY_WINDOW_NAME, binary);
  }

  void showDepth(const cv::Mat& depth) {
    cv::Mat buf = cv::Mat::zeros(depth.size(), CV_16UC1);
    cv::Mat output = cv::Mat::zeros(depth.size(), CV_8UC3);
    cv::normalize(depth, buf, 0, 2013, cv::NORM_MINMAX);

    for (int i = 0; i < depth.rows * depth.cols; i++) {
      int y = (int)(i / depth.cols);
      int x = (int)(i % depth.cols);
      int pval = m_gamma[buf.at<uint16_t>(i)];
      int lb = pval & 0xff;

      switch(pval >> 8) {
      case 0:
        output.at<cv::Vec3b>(y,x)[2] = 255;
        output.at<cv::Vec3b>(y,x)[1] = 255-lb;
        output.at<cv::Vec3b>(y,x)[0] = 255-lb;
        break;
      case 1:
        output.at<cv::Vec3b>(y,x)[2] = 255;
        output.at<cv::Vec3b>(y,x)[1] = lb;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      case 2:
        output.at<cv::Vec3b>(y,x)[2] = 255-lb;
        output.at<cv::Vec3b>(y,x)[1] = 255;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      case 3:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 255;
        output.at<cv::Vec3b>(y,x)[0] = lb;
        break;
      case 4:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 255-lb;
        output.at<cv::Vec3b>(y,x)[0] = 255;
        break;
      case 5:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 0;
        output.at<cv::Vec3b>(y,x)[0] = 255-lb;
        break;
      default:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 0;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      }
    }
    imshow(DEPTH_WINDOW_NAME, output);
  }

private:
  std::vector<uint16_t> m_gamma;
  const std::string RGB_WINDOW_NAME;
  const std::string DEPTH_WINDOW_NAME;
  const std::string BINARY_WINDOW_NAME;
};

class SearchTomato {
public:
  SearchTomato()
    : tomato_contours {},
      siObj {}
  {
  }

  void update(const cv::Mat& capture_rgb) {
    siObj.showRGB(capture_rgb);
    cv::Mat binary_mat = cv::Mat::zeros(capture_rgb.size(), CV_8UC1);
    imageProsessing(capture_rgb, binary_mat);
    siObj.showBinary(binary_mat);
    searchTomatoContours(binary_mat, tomato_contours);
  }

  bool searchTomato(const cv::Mat& capture_depth, geometry_msgs::Point& tomato_point) {
    siObj.showDepth(capture_depth);
    cv::Mat maskedDepth = cv::Mat::zeros(capture_depth.size(), CV_16UC1);
    cv::Mat mask = cv::Mat::zeros(capture_depth.size(), CV_8UC1);
    cv::drawContours(mask, tomato_contours, -1, cv::Scalar(255), CV_FILLED, 8);
    siObj.showBinary(mask);

    capture_depth.copyTo(maskedDepth, mask);
    return findClosePoint(maskedDepth, tomato_point);
  }

  bool searchTomato(const cv::Mat& capture_depth, geometry_msgs::PoseArray& tomato_poses) {
    tomato_poses.poses.clear();
    siObj.showDepth(capture_depth);
    cv::Mat maskedDepth = cv::Mat::zeros(capture_depth.size(), CV_16UC1);
    cv::Mat mask = cv::Mat::zeros(capture_depth.size(), CV_8UC1);
    for (std::size_t i {0}; i < tomato_contours.size(); ++i) {
      cv::drawContours(mask, tomato_contours, i, cv::Scalar(255), CV_FILLED, 8);
      siObj.showBinary(mask);

      capture_depth.copyTo(maskedDepth, mask);
      findClosePoint(maskedDepth, tomato_poses);
    }
    return tomato_poses.poses.empty() ? false : true;
  }

private:
  void imageProsessing(const cv::Mat& rgb, cv::Mat& binary) {
    cv::Mat blur, hsv;
    cv::GaussianBlur(rgb, blur, cv::Size(5, 5), 4.0, 4.0);
    cv::cvtColor(blur, hsv, CV_BGR2HSV);
    redFilter(hsv, binary);
    cv::dilate(binary, binary, cv::Mat(), cv::Point(-1, -1), 2);
    cv::erode(binary, binary, cv::Mat(), cv::Point(-1, -1), 4);
    cv::dilate(binary, binary, cv::Mat(), cv::Point(-1, -1), 1);
  }

  void searchTomatoContours(const cv::Mat& binary, std::vector<std::vector<cv::Point> >& tomato_contours) {
    std::vector<std::vector<cv::Point> > contours;
    std::vector<cv::Point> contour;
    cv::Rect bounding_box;
    float ratio_balance;

    cv::findContours(binary, contours, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_NONE);

    tomato_contours.clear();
    while (!contours.empty()) {
      contour = contours.back();
      contours.pop_back();

      bounding_box = cv::boundingRect(contour);
      ratio_balance = (float)bounding_box.width / (float)bounding_box.height;
      if (ratio_balance > 1.0f) ratio_balance = 1.0f / ratio_balance;

      // delete mismach. that is smaller or spier or empty
      if (bounding_box.area() >= 400 && bounding_box.area() < cv::contourArea(contour)*2 && ratio_balance > 0.4f)
        tomato_contours.push_back(contour);
    }
  }

  void redFilter(const cv::Mat& hsv, cv::Mat& binary) {
    int a, x, y;
    for(y = 0; y < hsv.rows; y++) {
      for(x = 0; x < hsv.cols; x++) {
        a = hsv.step*y+(x*3);
        if((hsv.data[a] <= 15 || hsv.data[a] >= 300) && hsv.data[a+1] >= 30 && hsv.data[a+2] >= 50)
          binary.at<unsigned char>(y,x) = 255;
        else
          binary.at<unsigned char>(y,x) = 0;
      }
    }
  }

  bool findClosePoint(const cv::Mat& maskedDepth, geometry_msgs::Point& tomato_point) {
    uint16_t depth;
    constexpr uint16_t OVER_RANGE = 3000;
    uint16_t close_depth = OVER_RANGE;

    for (int y = 0; y < maskedDepth.rows; y++) {
      const uint16_t* line_point = maskedDepth.ptr<uint16_t>(y);
      for (int x = 0; x < maskedDepth.cols; x++) {
        depth = line_point[x];
        if (512 < depth && depth <  close_depth) {
          tomato_point.x = depth; // for ros coordinate system
          tomato_point.y = x - maskedDepth.cols / 2; // for ros coordinate system
          tomato_point.z = -(y - maskedDepth.rows / 2); // for ros coordinate system
          close_depth = depth;
        }
      }
    }
    return (close_depth != OVER_RANGE ? true : false);
  }

  bool findClosePoint(const cv::Mat& maskedDepth, geometry_msgs::PoseArray& tomato_poses) {
    geometry_msgs::Point tomato_point {};
    if (findClosePoint(maskedDepth, tomato_point)) {
      geometry_msgs::Pose tomato_pose {};
      tomato_pose.position = tomato_point;
      tomato_pose.orientation.w = 1.0;
      tomato_poses.poses.push_back(tomato_pose);
      return true;
    }
    return false;
  }

  std::vector<std::vector<cv::Point> > tomato_contours;
  ShowImage siObj;
};

class ImageConverter {
private:
  ros::NodeHandle nh;
  ros::NodeHandle tomapo_nh;
  ros::Publisher poses_pub;
  geometry_msgs::PoseArray pub_msg;
  image_transport::ImageTransport it;
  image_transport::Subscriber rgb_sub;
  image_transport::Subscriber depth_sub;
  cv_bridge::CvImageConstPtr rgb_ptr;
  cv_bridge::CvImageConstPtr depth_ptr;
  SearchTomato stObj;
  const std::size_t height_;
  const double angle_x_per_piccell_;
  const double angle_y_per_piccell_;
  static constexpr double PI {3.141592653589793};

public:
  ImageConverter(const std::size_t width, const std::size_t height, const double angle_of_view_x, const double angle_of_view_y)
    : nh {},
      tomapo_nh {nh, "tomato_point"},
      poses_pub {tomapo_nh.advertise<geometry_msgs::PoseArray>("array", 1)},
      it {nh},
      rgb_sub {it.subscribe("color", 1, &ImageConverter::rgbCb, this)},
      depth_sub {it.subscribe("depth", 1, &ImageConverter::depthCb, this)},
      stObj {},
      height_ {height},
      angle_x_per_piccell_ {(angle_of_view_x * PI / 180.0) / width},
      angle_y_per_piccell_ {(angle_of_view_y * PI / 180.0) / height}
  {
    pub_msg.header.frame_id = "kinect";
  }

  /**
   * search tomato function.
   * get rgb image, and search tomato.
   * will tomato_point have data.
   *
   * @author Yusuke Doi
   */
  void rgbCb(const sensor_msgs::ImageConstPtr& msg) {
    // get rgb image on kinect
    try {
      rgb_ptr = cv_bridge::toCvShare(msg, "bgr8");
    } catch (cv_bridge::Exception& e) {
      ROS_ERROR("cv_bridge exception by rgb: %s", e.what());
    }

    stObj.update(rgb_ptr->image);

    cv::waitKey(100);
  }

  void depthCb(const sensor_msgs::ImageConstPtr& msg) {
    try {
      depth_ptr = cv_bridge::toCvShare(msg, sensor_msgs::image_encodings::TYPE_16UC1);
    } catch (cv_bridge::Exception& e) {
      ROS_ERROR("cv_bridge exception by depth: %s", e.what());
    }

    if (stObj.searchTomato(depth_ptr->image, pub_msg)) {
      pub_msg.header.stamp = ros::Time::now();
      for (auto tomato_pose& : pub_msg.poses) {
        tomato_pose.position.y = tomato_pose.position.x * tan(angle_x_per_piccell_ * tomato_pose.position.y) / 1000.0; // convert to m
        tomato_pose.position.z = tomato_pose.position.x * tan(angle_y_per_piccell_ * tomato_pose.position.z) / 1000.0; // convert to m
        tomato_pose.position.x = tomato_pose.position.x / 1000.0; // convert to m
      }
      poses_pub.publish(pub_msg);
    }

    cv::waitKey(100);
  }
};

int main(int argc, char** argv) {
  ros::init(argc, argv, "kinect_tomato_searcher_node");
  ImageConverter ic {512, 424, 70, 60};
  ros::spin();
  return 0;
}
<commit_msg>Fix typo<commit_after>#include <ros/ros.h>
#include <cv_bridge/cv_bridge.h>
#include <geometry_msgs/Point.h>
#include <geometry_msgs/PointStamped.h>
#include <geometry_msgs/Pose.h>
#include <geometry_msgs/PoseArray.h>
#include <image_transport/image_transport.h>
#include <sensor_msgs/image_encodings.h>

#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/video/video.hpp>

class ShowImage {
public:
  ShowImage()
    : m_gamma(2048),
      RGB_WINDOW_NAME {"RGB image"},
      DEPTH_WINDOW_NAME {"Depth image"},
      BINARY_WINDOW_NAME {"Binary image"}
  {
    for(int i = 0; i < 2048; i++) {
      float v = i / 2048.0;
      v = std::pow(v, 3) * 6;
      m_gamma[i] = v * 6 * 256;
    }

    cv::namedWindow(RGB_WINDOW_NAME);
    cv::namedWindow(DEPTH_WINDOW_NAME);
    cv::namedWindow(BINARY_WINDOW_NAME);
  }

  ~ShowImage() {
    cv::destroyWindow(RGB_WINDOW_NAME);
    cv::destroyWindow(DEPTH_WINDOW_NAME);
    cv::destroyWindow(BINARY_WINDOW_NAME);
  }

  void showRGB(const cv::Mat& rgb) {
    imshow(RGB_WINDOW_NAME, rgb);
  }

  void showBinary(const cv::Mat& binary) {
    imshow(BINARY_WINDOW_NAME, binary);
  }

  void showDepth(const cv::Mat& depth) {
    cv::Mat buf = cv::Mat::zeros(depth.size(), CV_16UC1);
    cv::Mat output = cv::Mat::zeros(depth.size(), CV_8UC3);
    cv::normalize(depth, buf, 0, 2013, cv::NORM_MINMAX);

    for (int i = 0; i < depth.rows * depth.cols; i++) {
      int y = (int)(i / depth.cols);
      int x = (int)(i % depth.cols);
      int pval = m_gamma[buf.at<uint16_t>(i)];
      int lb = pval & 0xff;

      switch(pval >> 8) {
      case 0:
        output.at<cv::Vec3b>(y,x)[2] = 255;
        output.at<cv::Vec3b>(y,x)[1] = 255-lb;
        output.at<cv::Vec3b>(y,x)[0] = 255-lb;
        break;
      case 1:
        output.at<cv::Vec3b>(y,x)[2] = 255;
        output.at<cv::Vec3b>(y,x)[1] = lb;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      case 2:
        output.at<cv::Vec3b>(y,x)[2] = 255-lb;
        output.at<cv::Vec3b>(y,x)[1] = 255;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      case 3:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 255;
        output.at<cv::Vec3b>(y,x)[0] = lb;
        break;
      case 4:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 255-lb;
        output.at<cv::Vec3b>(y,x)[0] = 255;
        break;
      case 5:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 0;
        output.at<cv::Vec3b>(y,x)[0] = 255-lb;
        break;
      default:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 0;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      }
    }
    imshow(DEPTH_WINDOW_NAME, output);
  }

private:
  std::vector<uint16_t> m_gamma;
  const std::string RGB_WINDOW_NAME;
  const std::string DEPTH_WINDOW_NAME;
  const std::string BINARY_WINDOW_NAME;
};

class SearchTomato {
public:
  SearchTomato()
    : tomato_contours {},
      siObj {}
  {
  }

  void update(const cv::Mat& capture_rgb) {
    siObj.showRGB(capture_rgb);
    cv::Mat binary_mat = cv::Mat::zeros(capture_rgb.size(), CV_8UC1);
    imageProsessing(capture_rgb, binary_mat);
    siObj.showBinary(binary_mat);
    searchTomatoContours(binary_mat, tomato_contours);
  }

  bool searchTomato(const cv::Mat& capture_depth, geometry_msgs::Point& tomato_point) {
    siObj.showDepth(capture_depth);
    cv::Mat maskedDepth = cv::Mat::zeros(capture_depth.size(), CV_16UC1);
    cv::Mat mask = cv::Mat::zeros(capture_depth.size(), CV_8UC1);
    cv::drawContours(mask, tomato_contours, -1, cv::Scalar(255), CV_FILLED, 8);
    siObj.showBinary(mask);

    capture_depth.copyTo(maskedDepth, mask);
    return findClosePoint(maskedDepth, tomato_point);
  }

  bool searchTomato(const cv::Mat& capture_depth, geometry_msgs::PoseArray& tomato_poses) {
    tomato_poses.poses.clear();
    siObj.showDepth(capture_depth);
    cv::Mat maskedDepth = cv::Mat::zeros(capture_depth.size(), CV_16UC1);
    cv::Mat mask = cv::Mat::zeros(capture_depth.size(), CV_8UC1);
    for (std::size_t i {0}; i < tomato_contours.size(); ++i) {
      cv::drawContours(mask, tomato_contours, i, cv::Scalar(255), CV_FILLED, 8);
      siObj.showBinary(mask);

      capture_depth.copyTo(maskedDepth, mask);
      findClosePoint(maskedDepth, tomato_poses);
    }
    return tomato_poses.poses.empty() ? false : true;
  }

private:
  void imageProsessing(const cv::Mat& rgb, cv::Mat& binary) {
    cv::Mat blur, hsv;
    cv::GaussianBlur(rgb, blur, cv::Size(5, 5), 4.0, 4.0);
    cv::cvtColor(blur, hsv, CV_BGR2HSV);
    redFilter(hsv, binary);
    cv::dilate(binary, binary, cv::Mat(), cv::Point(-1, -1), 2);
    cv::erode(binary, binary, cv::Mat(), cv::Point(-1, -1), 4);
    cv::dilate(binary, binary, cv::Mat(), cv::Point(-1, -1), 1);
  }

  void searchTomatoContours(const cv::Mat& binary, std::vector<std::vector<cv::Point> >& tomato_contours) {
    std::vector<std::vector<cv::Point> > contours;
    std::vector<cv::Point> contour;
    cv::Rect bounding_box;
    float ratio_balance;

    cv::findContours(binary, contours, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_NONE);

    tomato_contours.clear();
    while (!contours.empty()) {
      contour = contours.back();
      contours.pop_back();

      bounding_box = cv::boundingRect(contour);
      ratio_balance = (float)bounding_box.width / (float)bounding_box.height;
      if (ratio_balance > 1.0f) ratio_balance = 1.0f / ratio_balance;

      // delete mismach. that is smaller or spier or empty
      if (bounding_box.area() >= 400 && bounding_box.area() < cv::contourArea(contour)*2 && ratio_balance > 0.4f)
        tomato_contours.push_back(contour);
    }
  }

  void redFilter(const cv::Mat& hsv, cv::Mat& binary) {
    int a, x, y;
    for(y = 0; y < hsv.rows; y++) {
      for(x = 0; x < hsv.cols; x++) {
        a = hsv.step*y+(x*3);
        if((hsv.data[a] <= 15 || hsv.data[a] >= 300) && hsv.data[a+1] >= 30 && hsv.data[a+2] >= 50)
          binary.at<unsigned char>(y,x) = 255;
        else
          binary.at<unsigned char>(y,x) = 0;
      }
    }
  }

  bool findClosePoint(const cv::Mat& maskedDepth, geometry_msgs::Point& tomato_point) {
    uint16_t depth;
    constexpr uint16_t OVER_RANGE = 3000;
    uint16_t close_depth = OVER_RANGE;

    for (int y = 0; y < maskedDepth.rows; y++) {
      const uint16_t* line_point = maskedDepth.ptr<uint16_t>(y);
      for (int x = 0; x < maskedDepth.cols; x++) {
        depth = line_point[x];
        if (512 < depth && depth <  close_depth) {
          tomato_point.x = depth; // for ros coordinate system
          tomato_point.y = x - maskedDepth.cols / 2; // for ros coordinate system
          tomato_point.z = -(y - maskedDepth.rows / 2); // for ros coordinate system
          close_depth = depth;
        }
      }
    }
    return (close_depth != OVER_RANGE ? true : false);
  }

  bool findClosePoint(const cv::Mat& maskedDepth, geometry_msgs::PoseArray& tomato_poses) {
    geometry_msgs::Point tomato_point {};
    if (findClosePoint(maskedDepth, tomato_point)) {
      geometry_msgs::Pose tomato_pose {};
      tomato_pose.position = tomato_point;
      tomato_pose.orientation.w = 1.0;
      tomato_poses.poses.push_back(tomato_pose);
      return true;
    }
    return false;
  }

  std::vector<std::vector<cv::Point> > tomato_contours;
  ShowImage siObj;
};

class ImageConverter {
private:
  ros::NodeHandle nh;
  ros::NodeHandle tomapo_nh;
  ros::Publisher poses_pub;
  geometry_msgs::PoseArray pub_msg;
  image_transport::ImageTransport it;
  image_transport::Subscriber rgb_sub;
  image_transport::Subscriber depth_sub;
  cv_bridge::CvImageConstPtr rgb_ptr;
  cv_bridge::CvImageConstPtr depth_ptr;
  SearchTomato stObj;
  const std::size_t height_;
  const double angle_x_per_piccell_;
  const double angle_y_per_piccell_;
  static constexpr double PI {3.141592653589793};

public:
  ImageConverter(const std::size_t width, const std::size_t height, const double angle_of_view_x, const double angle_of_view_y)
    : nh {},
      tomapo_nh {nh, "tomato_point"},
      poses_pub {tomapo_nh.advertise<geometry_msgs::PoseArray>("array", 1)},
      it {nh},
      rgb_sub {it.subscribe("color", 1, &ImageConverter::rgbCb, this)},
      depth_sub {it.subscribe("depth", 1, &ImageConverter::depthCb, this)},
      stObj {},
      height_ {height},
      angle_x_per_piccell_ {(angle_of_view_x * PI / 180.0) / width},
      angle_y_per_piccell_ {(angle_of_view_y * PI / 180.0) / height}
  {
    pub_msg.header.frame_id = "kinect";
  }

  /**
   * search tomato function.
   * get rgb image, and search tomato.
   * will tomato_point have data.
   *
   * @author Yusuke Doi
   */
  void rgbCb(const sensor_msgs::ImageConstPtr& msg) {
    // get rgb image on kinect
    try {
      rgb_ptr = cv_bridge::toCvShare(msg, "bgr8");
    } catch (cv_bridge::Exception& e) {
      ROS_ERROR("cv_bridge exception by rgb: %s", e.what());
    }

    stObj.update(rgb_ptr->image);

    cv::waitKey(100);
  }

  void depthCb(const sensor_msgs::ImageConstPtr& msg) {
    try {
      depth_ptr = cv_bridge::toCvShare(msg, sensor_msgs::image_encodings::TYPE_16UC1);
    } catch (cv_bridge::Exception& e) {
      ROS_ERROR("cv_bridge exception by depth: %s", e.what());
    }

    if (stObj.searchTomato(depth_ptr->image, pub_msg)) {
      pub_msg.header.stamp = ros::Time::now();
      for (auto& tomato_pose : pub_msg.poses) {
        tomato_pose.position.y = tomato_pose.position.x * tan(angle_x_per_piccell_ * tomato_pose.position.y) / 1000.0; // convert to m
        tomato_pose.position.z = tomato_pose.position.x * tan(angle_y_per_piccell_ * tomato_pose.position.z) / 1000.0; // convert to m
        tomato_pose.position.x = tomato_pose.position.x / 1000.0; // convert to m
      }
      poses_pub.publish(pub_msg);
    }

    cv::waitKey(100);
  }
};

int main(int argc, char** argv) {
  ros::init(argc, argv, "kinect_tomato_searcher_node");
  ImageConverter ic {512, 424, 70, 60};
  ros::spin();
  return 0;
}
<|endoftext|>


<commit_before>#include <ros/ros.h>
#include <cv_bridge/cv_bridge.h>
#include <geometry_msgs/Point.h>
#include <geometry_msgs/PointStamped.h>
#include <geometry_msgs/Pose.h>
#include <geometry_msgs/PoseArray.h>
#include <image_transport/image_transport.h>
#include <sensor_msgs/image_encodings.h>

#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/video/video.hpp>

class ShowImage {
public:
  ShowImage()
    : m_gamma(2048),
      RGB_WINDOW_NAME {"RGB image"},
      DEPTH_WINDOW_NAME {"Depth image"},
      BINARY_WINDOW_NAME {"Binary image"}
  {
    for(int i = 0; i < 2048; i++) {
      float v = i / 2048.0;
      v = std::pow(v, 3) * 6;
      m_gamma[i] = v * 6 * 256;
    }

    cv::namedWindow(RGB_WINDOW_NAME);
    cv::namedWindow(DEPTH_WINDOW_NAME);
    cv::namedWindow(BINARY_WINDOW_NAME);
  }

  ~ShowImage() {
    cv::destroyWindow(RGB_WINDOW_NAME);
    cv::destroyWindow(DEPTH_WINDOW_NAME);
    cv::destroyWindow(BINARY_WINDOW_NAME);
  }

  void showRGB(const cv::Mat& rgb) {
    imshow(RGB_WINDOW_NAME, rgb);
  }

  void showBinary(const cv::Mat& binary) {
    imshow(BINARY_WINDOW_NAME, binary);
  }

  void showDepth(const cv::Mat& depth) {
    cv::Mat buf = cv::Mat::zeros(depth.size(), CV_16UC1);
    cv::Mat output = cv::Mat::zeros(depth.size(), CV_8UC3);
    cv::normalize(depth, buf, 0, 2013, cv::NORM_MINMAX);

    for (int i = 0; i < depth.rows * depth.cols; i++) {
      int y = (int)(i / depth.cols);
      int x = (int)(i % depth.cols);
      int pval = m_gamma[buf.at<uint16_t>(i)];
      int lb = pval & 0xff;

      switch(pval >> 8) {
      case 0:
        output.at<cv::Vec3b>(y,x)[2] = 255;
        output.at<cv::Vec3b>(y,x)[1] = 255-lb;
        output.at<cv::Vec3b>(y,x)[0] = 255-lb;
        break;
      case 1:
        output.at<cv::Vec3b>(y,x)[2] = 255;
        output.at<cv::Vec3b>(y,x)[1] = lb;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      case 2:
        output.at<cv::Vec3b>(y,x)[2] = 255-lb;
        output.at<cv::Vec3b>(y,x)[1] = 255;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      case 3:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 255;
        output.at<cv::Vec3b>(y,x)[0] = lb;
        break;
      case 4:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 255-lb;
        output.at<cv::Vec3b>(y,x)[0] = 255;
        break;
      case 5:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 0;
        output.at<cv::Vec3b>(y,x)[0] = 255-lb;
        break;
      default:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 0;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      }
    }
    imshow(DEPTH_WINDOW_NAME, output);
  }

private:
  std::vector<uint16_t> m_gamma;
  const std::string RGB_WINDOW_NAME;
  const std::string DEPTH_WINDOW_NAME;
  const std::string BINARY_WINDOW_NAME;
};

class SearchTomato {
public:
  SearchTomato()
    : tomato_contours {},
      siObj {}
  {
  }

  void update(const cv::Mat& capture_rgb) {
    siObj.showRGB(capture_rgb);
    cv::Mat binary_mat = cv::Mat::zeros(capture_rgb.size(), CV_8UC1);
    imageProsessing(capture_rgb, binary_mat);
    siObj.showBinary(binary_mat);
    searchTomatoContours(binary_mat, tomato_contours);
  }

  bool searchTomato(const cv::Mat& capture_depth, geometry_msgs::Point& tomato_point) {
    siObj.showDepth(capture_depth);
    cv::Mat maskedDepth = cv::Mat::zeros(capture_depth.size(), CV_16UC1);
    cv::Mat mask = cv::Mat::zeros(capture_depth.size(), CV_8UC1);
    cv::drawContours(mask, tomato_contours, -1, cv::Scalar(255), CV_FILLED, 8);
    siObj.showBinary(mask);

    capture_depth.copyTo(maskedDepth, mask);
    return findClosePoint(maskedDepth, tomato_point);
  }

  bool searchTomato(const cv::Mat& capture_depth, geometry_msgs::PoseArray& tomato_poses) {
    tomato_poses.poses.clear();
    siObj.showDepth(capture_depth);
    cv::Mat maskedDepth = cv::Mat::zeros(capture_depth.size(), CV_16UC1);
    cv::Mat mask = cv::Mat::zeros(capture_depth.size(), CV_8UC1);
    for (std::size_t i {0}; i < tomato_contours.size(); ++i) {
      cv::drawContours(mask, tomato_contours, i, cv::Scalar(255), CV_FILLED, 8);
      siObj.showBinary(mask);

      capture_depth.copyTo(maskedDepth, mask);
      findClosePoint(maskedDepth, tomato_poses);
    }
    return tomato_poses.poses.empty() ? false : true;
  }

private:
  void imageProsessing(const cv::Mat& rgb, cv::Mat& binary) {
    cv::Mat blur, hsv;
    cv::GaussianBlur(rgb, blur, cv::Size(5, 5), 4.0, 4.0);
    cv::cvtColor(blur, hsv, CV_BGR2HSV);
    redFilter(hsv, binary);
    cv::dilate(binary, binary, cv::Mat(), cv::Point(-1, -1), 2);
    cv::erode(binary, binary, cv::Mat(), cv::Point(-1, -1), 4);
    cv::dilate(binary, binary, cv::Mat(), cv::Point(-1, -1), 1);
  }

  void searchTomatoContours(const cv::Mat& binary, std::vector<std::vector<cv::Point> >& tomato_contours) {
    std::vector<std::vector<cv::Point> > contours;
    std::vector<cv::Point> contour;
    cv::Rect bounding_box;
    float ratio_balance;

    cv::findContours(binary, contours, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_NONE);

    tomato_contours.clear();
    while (!contours.empty()) {
      contour = contours.back();
      contours.pop_back();

      bounding_box = cv::boundingRect(contour);
      ratio_balance = (float)bounding_box.width / (float)bounding_box.height;
      if (ratio_balance > 1.0f) ratio_balance = 1.0f / ratio_balance;

      // delete mismach. that is smaller or spier or empty
      if (bounding_box.area() >= 400 && bounding_box.area() < cv::contourArea(contour)*2 && ratio_balance > 0.4f)
        tomato_contours.push_back(contour);
    }
  }

  void redFilter(const cv::Mat& hsv, cv::Mat& binary) {
    int a, x, y;
    for(y = 0; y < hsv.rows; y++) {
      for(x = 0; x < hsv.cols; x++) {
        a = hsv.step*y+(x*3);
        if((hsv.data[a] <= 15 || hsv.data[a] >= 300) && hsv.data[a+1] >= 30 && hsv.data[a+2] >= 50)
          binary.at<unsigned char>(y,x) = 255;
        else
          binary.at<unsigned char>(y,x) = 0;
      }
    }
  }

  bool findClosePoint(const cv::Mat& maskedDepth, geometry_msgs::Point& tomato_point) {
    uint16_t depth;
    constexpr uint16_t OVER_RANGE = 3000;
    uint16_t close_depth = OVER_RANGE;

    for (int y = 0; y < maskedDepth.rows; y++) {
      const uint16_t* line_point = maskedDepth.ptr<uint16_t>(y);
      for (int x = 0; x < maskedDepth.cols; x++) {
        depth = line_point[x];
        if (512 < depth && depth <  close_depth) {
          tomato_point.x = depth; // for ros coordinate system
          tomato_point.y = x - maskedDepth.cols / 2; // for ros coordinate system
          tomato_point.z = -(y - maskedDepth.rows / 2); // for ros coordinate system
          close_depth = depth;
        }
      }
    }
    return (close_depth != OVER_RANGE ? true : false);
  }

  bool findClosePoint(const cv::Mat& maskedDepth, geometry_msgs::PoseArray& tomato_poses) {
    geometry_msgs::Point tomato_point {};
    if (findClosePoint(maskedDepth, tomato_point)) {
      geometry_msgs::Pose tomato_pose {};
      tomato_pose.position = tomato_point;
      tomato_pose.orientation.w = 1.0;
      tomato_poses.poses.push_back(tomato_pose);
      return true;
    }
    return false;
  }

  std::vector<std::vector<cv::Point> > tomato_contours;
  ShowImage siObj;
};

class ImageConverter {
private:
  ros::NodeHandle nh;
  ros::NodeHandle tomapo_nh;
  ros::Publisher poses_pub;
  geometry_msgs::PoseArray pub_msg;
  image_transport::ImageTransport it;
  image_transport::Subscriber rgb_sub;
  image_transport::Subscriber depth_sub;
  cv_bridge::CvImageConstPtr rgb_ptr;
  cv_bridge::CvImageConstPtr depth_ptr;
  SearchTomato stObj;
  const std::size_t height_;
  const double angle_x_per_piccell_;
  const double angle_y_per_piccell_;
  static constexpr double PI {3.141592653589793};

public:
  ImageConverter(const std::size_t width, const std::size_t height, const double angle_of_view_x, const double angle_of_view_y)
    : nh {},
      tomapo_nh {nh, "tomato_point"},
      poses_pub {tomapo_nh.advertise<geometry_msgs::PoseArray>("raw", 1)},
      it {nh},
      rgb_sub {it.subscribe("rgb", 1, &ImageConverter::rgbCb, this)},
      depth_sub {it.subscribe("depth", 1, &ImageConverter::depthCb, this)},
      stObj {},
      height_ {height},
      angle_x_per_piccell_ {(angle_of_view_x * PI / 180.0) / width},
      angle_y_per_piccell_ {(angle_of_view_y * PI / 180.0) / height}
  {
  }

  /**
   * search tomato function.
   * get rgb image, and search tomato.
   * will tomato_point have data.
   *
   * @author Yusuke Doi
   */
  void rgbCb(const sensor_msgs::ImageConstPtr& msg) {
    // get rgb image on kinect
    try {
      rgb_ptr = cv_bridge::toCvShare(msg, "bgr8");
    } catch (cv_bridge::Exception& e) {
      ROS_ERROR("cv_bridge exception by rgb: %s", e.what());
    }

    stObj.update(rgb_ptr->image);

    cv::waitKey(100);
  }

  void depthCb(const sensor_msgs::ImageConstPtr& msg) {
    try {
      depth_ptr = cv_bridge::toCvShare(msg, sensor_msgs::image_encodings::TYPE_16UC1);
    } catch (cv_bridge::Exception& e) {
      ROS_ERROR("cv_bridge exception by depth: %s", e.what());
    }

    if (stObj.searchTomato(depth_ptr->image, pub_msg)) {
      for (auto tomato_pose : pub_msg.poses) {
        tomato_pose.position.y = tomato_pose.position.x * tan(angle_x_per_piccell_ * tomato_pose.position.y) / 1000.0; // convert to m
        tomato_pose.position.z = tomato_pose.position.x * tan(angle_y_per_piccell_ * tomato_pose.position.z) / 1000.0; // convert to m
        tomato_pose.position.x = tomato_pose.position.x / 1000.0; // convert to m
      }
      poses_pub.publish(pub_msg);
    }

    cv::waitKey(100);
  }
};

int main(int argc, char** argv) {
  ros::init(argc, argv, "kinect_tomato_searcher_node");
  ImageConverter ic {512, 424, 70, 60};
  ros::spin();
  return 0;
}
<commit_msg>Update topic name for PoseArray<commit_after>#include <ros/ros.h>
#include <cv_bridge/cv_bridge.h>
#include <geometry_msgs/Point.h>
#include <geometry_msgs/PointStamped.h>
#include <geometry_msgs/Pose.h>
#include <geometry_msgs/PoseArray.h>
#include <image_transport/image_transport.h>
#include <sensor_msgs/image_encodings.h>

#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/video/video.hpp>

class ShowImage {
public:
  ShowImage()
    : m_gamma(2048),
      RGB_WINDOW_NAME {"RGB image"},
      DEPTH_WINDOW_NAME {"Depth image"},
      BINARY_WINDOW_NAME {"Binary image"}
  {
    for(int i = 0; i < 2048; i++) {
      float v = i / 2048.0;
      v = std::pow(v, 3) * 6;
      m_gamma[i] = v * 6 * 256;
    }

    cv::namedWindow(RGB_WINDOW_NAME);
    cv::namedWindow(DEPTH_WINDOW_NAME);
    cv::namedWindow(BINARY_WINDOW_NAME);
  }

  ~ShowImage() {
    cv::destroyWindow(RGB_WINDOW_NAME);
    cv::destroyWindow(DEPTH_WINDOW_NAME);
    cv::destroyWindow(BINARY_WINDOW_NAME);
  }

  void showRGB(const cv::Mat& rgb) {
    imshow(RGB_WINDOW_NAME, rgb);
  }

  void showBinary(const cv::Mat& binary) {
    imshow(BINARY_WINDOW_NAME, binary);
  }

  void showDepth(const cv::Mat& depth) {
    cv::Mat buf = cv::Mat::zeros(depth.size(), CV_16UC1);
    cv::Mat output = cv::Mat::zeros(depth.size(), CV_8UC3);
    cv::normalize(depth, buf, 0, 2013, cv::NORM_MINMAX);

    for (int i = 0; i < depth.rows * depth.cols; i++) {
      int y = (int)(i / depth.cols);
      int x = (int)(i % depth.cols);
      int pval = m_gamma[buf.at<uint16_t>(i)];
      int lb = pval & 0xff;

      switch(pval >> 8) {
      case 0:
        output.at<cv::Vec3b>(y,x)[2] = 255;
        output.at<cv::Vec3b>(y,x)[1] = 255-lb;
        output.at<cv::Vec3b>(y,x)[0] = 255-lb;
        break;
      case 1:
        output.at<cv::Vec3b>(y,x)[2] = 255;
        output.at<cv::Vec3b>(y,x)[1] = lb;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      case 2:
        output.at<cv::Vec3b>(y,x)[2] = 255-lb;
        output.at<cv::Vec3b>(y,x)[1] = 255;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      case 3:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 255;
        output.at<cv::Vec3b>(y,x)[0] = lb;
        break;
      case 4:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 255-lb;
        output.at<cv::Vec3b>(y,x)[0] = 255;
        break;
      case 5:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 0;
        output.at<cv::Vec3b>(y,x)[0] = 255-lb;
        break;
      default:
        output.at<cv::Vec3b>(y,x)[2] = 0;
        output.at<cv::Vec3b>(y,x)[1] = 0;
        output.at<cv::Vec3b>(y,x)[0] = 0;
        break;
      }
    }
    imshow(DEPTH_WINDOW_NAME, output);
  }

private:
  std::vector<uint16_t> m_gamma;
  const std::string RGB_WINDOW_NAME;
  const std::string DEPTH_WINDOW_NAME;
  const std::string BINARY_WINDOW_NAME;
};

class SearchTomato {
public:
  SearchTomato()
    : tomato_contours {},
      siObj {}
  {
  }

  void update(const cv::Mat& capture_rgb) {
    siObj.showRGB(capture_rgb);
    cv::Mat binary_mat = cv::Mat::zeros(capture_rgb.size(), CV_8UC1);
    imageProsessing(capture_rgb, binary_mat);
    siObj.showBinary(binary_mat);
    searchTomatoContours(binary_mat, tomato_contours);
  }

  bool searchTomato(const cv::Mat& capture_depth, geometry_msgs::Point& tomato_point) {
    siObj.showDepth(capture_depth);
    cv::Mat maskedDepth = cv::Mat::zeros(capture_depth.size(), CV_16UC1);
    cv::Mat mask = cv::Mat::zeros(capture_depth.size(), CV_8UC1);
    cv::drawContours(mask, tomato_contours, -1, cv::Scalar(255), CV_FILLED, 8);
    siObj.showBinary(mask);

    capture_depth.copyTo(maskedDepth, mask);
    return findClosePoint(maskedDepth, tomato_point);
  }

  bool searchTomato(const cv::Mat& capture_depth, geometry_msgs::PoseArray& tomato_poses) {
    tomato_poses.poses.clear();
    siObj.showDepth(capture_depth);
    cv::Mat maskedDepth = cv::Mat::zeros(capture_depth.size(), CV_16UC1);
    cv::Mat mask = cv::Mat::zeros(capture_depth.size(), CV_8UC1);
    for (std::size_t i {0}; i < tomato_contours.size(); ++i) {
      cv::drawContours(mask, tomato_contours, i, cv::Scalar(255), CV_FILLED, 8);
      siObj.showBinary(mask);

      capture_depth.copyTo(maskedDepth, mask);
      findClosePoint(maskedDepth, tomato_poses);
    }
    return tomato_poses.poses.empty() ? false : true;
  }

private:
  void imageProsessing(const cv::Mat& rgb, cv::Mat& binary) {
    cv::Mat blur, hsv;
    cv::GaussianBlur(rgb, blur, cv::Size(5, 5), 4.0, 4.0);
    cv::cvtColor(blur, hsv, CV_BGR2HSV);
    redFilter(hsv, binary);
    cv::dilate(binary, binary, cv::Mat(), cv::Point(-1, -1), 2);
    cv::erode(binary, binary, cv::Mat(), cv::Point(-1, -1), 4);
    cv::dilate(binary, binary, cv::Mat(), cv::Point(-1, -1), 1);
  }

  void searchTomatoContours(const cv::Mat& binary, std::vector<std::vector<cv::Point> >& tomato_contours) {
    std::vector<std::vector<cv::Point> > contours;
    std::vector<cv::Point> contour;
    cv::Rect bounding_box;
    float ratio_balance;

    cv::findContours(binary, contours, CV_RETR_EXTERNAL, CV_CHAIN_APPROX_NONE);

    tomato_contours.clear();
    while (!contours.empty()) {
      contour = contours.back();
      contours.pop_back();

      bounding_box = cv::boundingRect(contour);
      ratio_balance = (float)bounding_box.width / (float)bounding_box.height;
      if (ratio_balance > 1.0f) ratio_balance = 1.0f / ratio_balance;

      // delete mismach. that is smaller or spier or empty
      if (bounding_box.area() >= 400 && bounding_box.area() < cv::contourArea(contour)*2 && ratio_balance > 0.4f)
        tomato_contours.push_back(contour);
    }
  }

  void redFilter(const cv::Mat& hsv, cv::Mat& binary) {
    int a, x, y;
    for(y = 0; y < hsv.rows; y++) {
      for(x = 0; x < hsv.cols; x++) {
        a = hsv.step*y+(x*3);
        if((hsv.data[a] <= 15 || hsv.data[a] >= 300) && hsv.data[a+1] >= 30 && hsv.data[a+2] >= 50)
          binary.at<unsigned char>(y,x) = 255;
        else
          binary.at<unsigned char>(y,x) = 0;
      }
    }
  }

  bool findClosePoint(const cv::Mat& maskedDepth, geometry_msgs::Point& tomato_point) {
    uint16_t depth;
    constexpr uint16_t OVER_RANGE = 3000;
    uint16_t close_depth = OVER_RANGE;

    for (int y = 0; y < maskedDepth.rows; y++) {
      const uint16_t* line_point = maskedDepth.ptr<uint16_t>(y);
      for (int x = 0; x < maskedDepth.cols; x++) {
        depth = line_point[x];
        if (512 < depth && depth <  close_depth) {
          tomato_point.x = depth; // for ros coordinate system
          tomato_point.y = x - maskedDepth.cols / 2; // for ros coordinate system
          tomato_point.z = -(y - maskedDepth.rows / 2); // for ros coordinate system
          close_depth = depth;
        }
      }
    }
    return (close_depth != OVER_RANGE ? true : false);
  }

  bool findClosePoint(const cv::Mat& maskedDepth, geometry_msgs::PoseArray& tomato_poses) {
    geometry_msgs::Point tomato_point {};
    if (findClosePoint(maskedDepth, tomato_point)) {
      geometry_msgs::Pose tomato_pose {};
      tomato_pose.position = tomato_point;
      tomato_pose.orientation.w = 1.0;
      tomato_poses.poses.push_back(tomato_pose);
      return true;
    }
    return false;
  }

  std::vector<std::vector<cv::Point> > tomato_contours;
  ShowImage siObj;
};

class ImageConverter {
private:
  ros::NodeHandle nh;
  ros::NodeHandle tomapo_nh;
  ros::Publisher poses_pub;
  geometry_msgs::PoseArray pub_msg;
  image_transport::ImageTransport it;
  image_transport::Subscriber rgb_sub;
  image_transport::Subscriber depth_sub;
  cv_bridge::CvImageConstPtr rgb_ptr;
  cv_bridge::CvImageConstPtr depth_ptr;
  SearchTomato stObj;
  const std::size_t height_;
  const double angle_x_per_piccell_;
  const double angle_y_per_piccell_;
  static constexpr double PI {3.141592653589793};

public:
  ImageConverter(const std::size_t width, const std::size_t height, const double angle_of_view_x, const double angle_of_view_y)
    : nh {},
      tomapo_nh {nh, "tomato_point"},
      poses_pub {tomapo_nh.advertise<geometry_msgs::PoseArray>("array", 1)},
      it {nh},
      rgb_sub {it.subscribe("rgb", 1, &ImageConverter::rgbCb, this)},
      depth_sub {it.subscribe("depth", 1, &ImageConverter::depthCb, this)},
      stObj {},
      height_ {height},
      angle_x_per_piccell_ {(angle_of_view_x * PI / 180.0) / width},
      angle_y_per_piccell_ {(angle_of_view_y * PI / 180.0) / height}
  {
  }

  /**
   * search tomato function.
   * get rgb image, and search tomato.
   * will tomato_point have data.
   *
   * @author Yusuke Doi
   */
  void rgbCb(const sensor_msgs::ImageConstPtr& msg) {
    // get rgb image on kinect
    try {
      rgb_ptr = cv_bridge::toCvShare(msg, "bgr8");
    } catch (cv_bridge::Exception& e) {
      ROS_ERROR("cv_bridge exception by rgb: %s", e.what());
    }

    stObj.update(rgb_ptr->image);

    cv::waitKey(100);
  }

  void depthCb(const sensor_msgs::ImageConstPtr& msg) {
    try {
      depth_ptr = cv_bridge::toCvShare(msg, sensor_msgs::image_encodings::TYPE_16UC1);
    } catch (cv_bridge::Exception& e) {
      ROS_ERROR("cv_bridge exception by depth: %s", e.what());
    }

    if (stObj.searchTomato(depth_ptr->image, pub_msg)) {
      for (auto tomato_pose : pub_msg.poses) {
        tomato_pose.position.y = tomato_pose.position.x * tan(angle_x_per_piccell_ * tomato_pose.position.y) / 1000.0; // convert to m
        tomato_pose.position.z = tomato_pose.position.x * tan(angle_y_per_piccell_ * tomato_pose.position.z) / 1000.0; // convert to m
        tomato_pose.position.x = tomato_pose.position.x / 1000.0; // convert to m
      }
      poses_pub.publish(pub_msg);
    }

    cv::waitKey(100);
  }
};

int main(int argc, char** argv) {
  ros::init(argc, argv, "kinect_tomato_searcher_node");
  ImageConverter ic {512, 424, 70, 60};
  ros::spin();
  return 0;
}
<|endoftext|>


<commit_before>
#include "AliHBTReaderITSv1.h"
#include "AliHBTEvent.h"
#include "AliHBTRun.h"
#include "AliHBTParticle.h"
#include "AliHBTParticleCut.h"

#include <iostream.h>

#include <TROOT.h>
#include <TFile.h>
#include <TTree.h>
#include <TBranch.h>
#include <TObjArray.h>
#include <TParticle.h>
#include <TString.h>
#include <TObjString.h>

#include <AliRun.h>
#include <AliMagF.h>
#include <AliKalmanTrack.h>
#include <AliITSIOTrack.h>

ClassImp(AliHBTReaderITSv1)
/********************************************************************/

AliHBTReaderITSv1::
AliHBTReaderITSv1(const Char_t* tracksfilename,const Char_t* galicefilename):
                 fITSTracksFileName(tracksfilename),fGAliceFileName(galicefilename)
 {
     fParticles = new AliHBTRun();
     fTracks    = new AliHBTRun();
     fIsRead = kFALSE;
 }
/********************************************************************/

AliHBTReaderITSv1::
AliHBTReaderITSv1(TObjArray* dirs, const Char_t* tracksfilename,const Char_t* galicefilename):
                 AliHBTReader(dirs),
                 fITSTracksFileName(tracksfilename),fGAliceFileName(galicefilename)
 {
   fParticles = new AliHBTRun();
   fTracks    = new AliHBTRun();
   fIsRead    = kFALSE;
 }
/********************************************************************/

AliHBTReaderITSv1::~AliHBTReaderITSv1()
{
   delete fParticles;
   delete fTracks;
}
/********************************************************************/

AliHBTEvent* AliHBTReaderITSv1::GetParticleEvent(Int_t n)
 {
 //returns Nth event with simulated particles
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetParticleEvent","Error in reading");
       return 0x0;
     }

   return fParticles->GetEvent(n);
 }
/********************************************************************/

AliHBTEvent* AliHBTReaderITSv1::GetTrackEvent(Int_t n)
 {
 //returns Nth event with reconstructed tracks
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetTrackEvent","Error in reading");
       return 0x0;
     }
   return fTracks->GetEvent(n);
 }
/********************************************************************/

Int_t AliHBTReaderITSv1::GetNumberOfPartEvents()
 {
 //returns number of events of particles
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetNumberOfPartEvents","Error in reading");
       return 0;
     }
   return fParticles->GetNumberOfEvents();
 }

/********************************************************************/
Int_t AliHBTReaderITSv1::GetNumberOfTrackEvents()
 {
 //returns number of events of tracks
  if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetNumberOfTrackEvents","Error in reading");
       return 0;
     }
  return fTracks->GetNumberOfEvents();
 }
/********************************************************************/

Int_t AliHBTReaderITSv1::Read(AliHBTRun* particles, AliHBTRun *tracks)
{
 cout<<"AliHBTReaderITSv1::Read()"<<endl;
 Int_t Nevents = 0;
 AliITSIOTrack *iotrack=new AliITSIOTrack;
 Int_t currentdir = 0;
 Int_t Ndirs;
 Int_t totalNevents = 0;

 if (fDirs)
  {
    Ndirs = fDirs->GetEntries();
  }
 else
  {
    Ndirs = 0;
  }

 do //do while is good even if
  {
   TFile* gAliceFile = OpenGAliceFile(currentdir);
   if(gAliceFile == 0x0)
    {
       Error("Read","Can not open the file with gAlice");
       delete iotrack;
       return 1;
    }
   TFile *file = OpenTrackFile(currentdir);
   if(file == 0x0)
    {
       Error("Read","Can not open the file with ITS tracks V1");
       delete iotrack;
       return 2;
    }

   if (gAlice->TreeE())//check if tree E exists
     {
      Nevents = (Int_t)gAlice->TreeE()->GetEntries();//if yes get number of events in gAlice
      cout<<"________________________________________________________\n";
      cout<<"Found "<<Nevents<<" event(s) in directory "<<GetDirName(currentdir)<<endl;
      cout<<"Setting Magnetic Field. Factor is "<<gAlice->Field()->Factor()<<endl;
      AliKalmanTrack::SetConvConst(100/0.299792458/0.2/gAlice->Field()->Factor());
     }
    else
     {//if not return an error
       Error("Read","Can not find Header tree (TreeE) in gAlice");
       delete iotrack;
       return 4;
     }

   Int_t naccepted = 0;
   char tname[30];

   for (Int_t currentEvent = 0; currentEvent < Nevents; currentEvent++)
    {
      cout<<"Reading Event "<<currentEvent;

      sprintf(tname,"TreeT%d",currentEvent);
      file->cd();
      TTree *tracktree=(TTree*)file->Get(tname);
      TBranch *tbranch=tracktree->GetBranch("ITStracks");
      tbranch->SetAddress(&iotrack);

      gAliceFile->cd();
      gAlice->GetEvent(currentEvent);
      gAlice->Particles();

      Int_t nentr=(Int_t)tracktree->GetEntries();

      cout<<".  Found "<<nentr<<" tracks.";
      fflush(0);

      for (Int_t i=0; i<nentr; i++)
       {

        tracktree->GetEvent(i);
        if(!iotrack) continue;
        Int_t label = iotrack->GetLabel();
        if (label < 0)
         {
           continue;
           delete iotrack;
         }
        TParticle *p = (TParticle*)gAlice->Particle(label);
        if(!p)
         {
           Warning("Read","Can not get particle with label &d",label);
           continue;
         }
        if(Pass(p->GetPdgCode())) continue; //check if we are intersted with particles of this type
                                           //if not take next partilce

        AliHBTParticle* part = new AliHBTParticle(*p);
        if(Pass(part)) { delete part; continue;}//check if meets all criteria of any of our cuts
                                                //if it does not delete it and take next good track

        Double_t px=iotrack->GetPx();
        Double_t py=iotrack->GetPy();
        Double_t pz=iotrack->GetPz();
        Double_t mass = p->GetMass();
        Double_t tEtot = TMath::Sqrt(px*px + py*py + pz*pz + mass*mass);//total energy of the track

        Double_t x= iotrack->GetX();
        Double_t y= iotrack->GetY();
        Double_t z= iotrack->GetZ();

        AliHBTParticle* track = new AliHBTParticle(p->GetPdgCode(), px, py , pz, tEtot, x, y, z, 0.);
        if(Pass(track)) { delete  track;continue;}//check if meets all criteria of any of our cuts
                                                  //if it does not delete it and take next good track

        particles->AddParticle(totalNevents,part);//put track and particle on the run
        tracks->AddParticle(totalNevents,track);
        naccepted++;
       }//end loop over tracks in the event

       totalNevents++;
       cout<<"  Accepted "<<naccepted<<" tracks"<<endl;
     }//end of loop over events in current directory

    gAliceFile->Close();
    delete gAliceFile;
    gAliceFile = 0;

    file->Close();
    delete file;
    file = 0;
    currentdir++;
   }while(currentdir < Ndirs);//end of loop over directories specified in fDirs Obj Array

  delete iotrack;
  fIsRead = kTRUE;
  return 0;

 }
/********************************************************************/

TFile* AliHBTReaderITSv1::OpenTrackFile(Int_t ndir)
{
//opens files to be read for given directoru nomber in fDirs Array
   const TString& dirname = GetDirName(ndir);
   if (dirname == "")
    {
      Error("OpenGAliceFile","Can not get directory name");
      return 0x0;
    }
   TString filename = dirname + "/" + fITSTracksFileName;
   TFile *file = (TFile*)gROOT->GetListOfFiles()->FindObject(filename.Data());
   if (!file) file = new TFile(filename.Data());

   if (!file)
    {
      Error("Read","Can not open file %s",filename.Data());
      return 0x0;
    }
   if (!file->IsOpen())
    {
      Error("Read","Can not open file %s",filename.Data());
      return 0x0;
    }

   return file;
}

/********************************************************************/
TFile* AliHBTReaderITSv1::OpenGAliceFile(Int_t ndir)
{
  const TString& dirname = GetDirName(ndir);
   if (dirname == "")
    {
      Error("OpenGAliceFile","Can not get directory name");
      return 0x0;
    }

  TString filename = dirname + "/" + fGAliceFileName;

  TFile* gAliceFile = TFile::Open(filename.Data());
  if ( gAliceFile== 0x0)
   {
     Error("OpenFiles","Can't open file named %s",filename.Data());
     return 0x0;
   }
  if (!gAliceFile->IsOpen())
   {
     Error("OpenFiles","Can't open file named %s",filename.Data());
     return 0x0;
   }

  if (!(gAlice=(AliRun*)gAliceFile->Get("gAlice")))
   {
     Error("OpenFiles","gAlice have not been found on %s !\n",filename.Data());
     gAliceFile->Close();
     delete gAliceFile;
     return 0x0;
   }

  return gAliceFile;
}

/********************************************************************/
/********************************************************************/

<commit_msg>Opening file with tracks changed<commit_after>
#include "AliHBTReaderITSv1.h"
#include "AliHBTEvent.h"
#include "AliHBTRun.h"
#include "AliHBTParticle.h"
#include "AliHBTParticleCut.h"

#include <iostream.h>

#include <TROOT.h>
#include <TFile.h>
#include <TTree.h>
#include <TBranch.h>
#include <TObjArray.h>
#include <TParticle.h>
#include <TString.h>
#include <TObjString.h>

#include <AliRun.h>
#include <AliMagF.h>
#include <AliKalmanTrack.h>
#include <AliITSIOTrack.h>

ClassImp(AliHBTReaderITSv1)
/********************************************************************/

AliHBTReaderITSv1::
AliHBTReaderITSv1(const Char_t* tracksfilename,const Char_t* galicefilename):
                 fITSTracksFileName(tracksfilename),fGAliceFileName(galicefilename)
 {
     fParticles = new AliHBTRun();
     fTracks    = new AliHBTRun();
     fIsRead = kFALSE;
 }
/********************************************************************/

AliHBTReaderITSv1::
AliHBTReaderITSv1(TObjArray* dirs, const Char_t* tracksfilename,const Char_t* galicefilename):
                 AliHBTReader(dirs),
                 fITSTracksFileName(tracksfilename),fGAliceFileName(galicefilename)
 {
   fParticles = new AliHBTRun();
   fTracks    = new AliHBTRun();
   fIsRead    = kFALSE;
 }
/********************************************************************/

AliHBTReaderITSv1::~AliHBTReaderITSv1()
{
   delete fParticles;
   delete fTracks;
}
/********************************************************************/

AliHBTEvent* AliHBTReaderITSv1::GetParticleEvent(Int_t n)
 {
 //returns Nth event with simulated particles
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetParticleEvent","Error in reading");
       return 0x0;
     }

   return fParticles->GetEvent(n);
 }
/********************************************************************/

AliHBTEvent* AliHBTReaderITSv1::GetTrackEvent(Int_t n)
 {
 //returns Nth event with reconstructed tracks
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetTrackEvent","Error in reading");
       return 0x0;
     }
   return fTracks->GetEvent(n);
 }
/********************************************************************/

Int_t AliHBTReaderITSv1::GetNumberOfPartEvents()
 {
 //returns number of events of particles
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetNumberOfPartEvents","Error in reading");
       return 0;
     }
   return fParticles->GetNumberOfEvents();
 }

/********************************************************************/
Int_t AliHBTReaderITSv1::GetNumberOfTrackEvents()
 {
 //returns number of events of tracks
  if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetNumberOfTrackEvents","Error in reading");
       return 0;
     }
  return fTracks->GetNumberOfEvents();
 }
/********************************************************************/

Int_t AliHBTReaderITSv1::Read(AliHBTRun* particles, AliHBTRun *tracks)
{
 cout<<"AliHBTReaderITSv1::Read()"<<endl;
 Int_t Nevents = 0;
 AliITSIOTrack *iotrack=new AliITSIOTrack;
 Int_t currentdir = 0;
 Int_t Ndirs;
 Int_t totalNevents = 0;

 if (fDirs)
  {
    Ndirs = fDirs->GetEntries();
  }
 else
  {
    Ndirs = 0;
  }

 do //do while is good even if
  {
   TFile* gAliceFile = OpenGAliceFile(currentdir);
   if(gAliceFile == 0x0)
    {
       Error("Read","Can not open the file with gAlice");
       delete iotrack;
       return 1;
    }
   TFile *file = OpenTrackFile(currentdir);
   if(file == 0x0)
    {
       Error("Read","Can not open the file with ITS tracks V1");
       delete iotrack;
       return 2;
    }

   if (gAlice->TreeE())//check if tree E exists
     {
      Nevents = (Int_t)gAlice->TreeE()->GetEntries();//if yes get number of events in gAlice
      cout<<"________________________________________________________\n";
      cout<<"Found "<<Nevents<<" event(s) in directory "<<GetDirName(currentdir)<<endl;
      cout<<"Setting Magnetic Field. Factor is "<<gAlice->Field()->Factor()<<endl;
      AliKalmanTrack::SetConvConst(100/0.299792458/0.2/gAlice->Field()->Factor());
     }
    else
     {//if not return an error
       Error("Read","Can not find Header tree (TreeE) in gAlice");
       delete iotrack;
       return 4;
     }

   Int_t naccepted = 0;
   char tname[30];

   for (Int_t currentEvent = 0; currentEvent < Nevents; currentEvent++)
    {
      cout<<"Reading Event "<<currentEvent;

      sprintf(tname,"TreeT%d",currentEvent);
      file->cd();
      TTree *tracktree=(TTree*)file->Get(tname);
      TBranch *tbranch=tracktree->GetBranch("ITStracks");
      tbranch->SetAddress(&iotrack);

      gAliceFile->cd();
      gAlice->GetEvent(currentEvent);
      gAlice->Particles();

      Int_t nentr=(Int_t)tracktree->GetEntries();

      cout<<".  Found "<<nentr<<" tracks.";
      fflush(0);

      for (Int_t i=0; i<nentr; i++)
       {

        tracktree->GetEvent(i);
        if(!iotrack) continue;
        Int_t label = iotrack->GetLabel();
        if (label < 0)
         {
           continue;
           delete iotrack;
         }
        TParticle *p = (TParticle*)gAlice->Particle(label);
        if(!p)
         {
           Warning("Read","Can not get particle with label &d",label);
           continue;
         }
        if(Pass(p->GetPdgCode())) continue; //check if we are intersted with particles of this type
                                           //if not take next partilce

        AliHBTParticle* part = new AliHBTParticle(*p);
        if(Pass(part)) { delete part; continue;}//check if meets all criteria of any of our cuts
                                                //if it does not delete it and take next good track

        Double_t px=iotrack->GetPx();
        Double_t py=iotrack->GetPy();
        Double_t pz=iotrack->GetPz();
        Double_t mass = p->GetMass();
        Double_t tEtot = TMath::Sqrt(px*px + py*py + pz*pz + mass*mass);//total energy of the track

        Double_t x= iotrack->GetX();
        Double_t y= iotrack->GetY();
        Double_t z= iotrack->GetZ();

        AliHBTParticle* track = new AliHBTParticle(p->GetPdgCode(), px, py , pz, tEtot, x, y, z, 0.);
        if(Pass(track)) { delete  track;continue;}//check if meets all criteria of any of our cuts
                                                  //if it does not delete it and take next good track

        particles->AddParticle(totalNevents,part);//put track and particle on the run
        tracks->AddParticle(totalNevents,track);
        naccepted++;
       }//end loop over tracks in the event

       totalNevents++;
       cout<<"  Accepted "<<naccepted<<" tracks"<<endl;
     }//end of loop over events in current directory

    gAliceFile->Close();
    delete gAliceFile;
    gAliceFile = 0;

    file->Close();
    delete file;
    file = 0;
    currentdir++;
   }while(currentdir < Ndirs);//end of loop over directories specified in fDirs Obj Array

  delete iotrack;
  fIsRead = kTRUE;
  return 0;

 }
/********************************************************************/

TFile* AliHBTReaderITSv1::OpenTrackFile(Int_t ndir)
{
//opens files to be read for given directoru nomber in fDirs Array
   const TString& dirname = GetDirName(ndir);
   if (dirname == "")
    {
      Error("OpenGAliceFile","Can not get directory name");
      return 0x0;
    }
   TString filename = dirname + "/" + fITSTracksFileName;

   TFile *file = TFile::Open(filename.Data());
   if (!file)
    {
      Error("Read","Can not open file %s",filename.Data());
      return 0x0;
    }
   if (!file->IsOpen())
    {
      Error("Read","Can not open file %s",filename.Data());
      return 0x0;
    }

   return file;
}

/********************************************************************/
TFile* AliHBTReaderITSv1::OpenGAliceFile(Int_t ndir)
{
  const TString& dirname = GetDirName(ndir);
   if (dirname == "")
    {
      Error("OpenGAliceFile","Can not get directory name");
      return 0x0;
    }

  TString filename = dirname + "/" + fGAliceFileName;

  TFile* gAliceFile = TFile::Open(filename.Data());
  if ( gAliceFile== 0x0)
   {
     Error("OpenFiles","Can't open file named %s",filename.Data());
     return 0x0;
   }
  if (!gAliceFile->IsOpen())
   {
     Error("OpenFiles","Can't open file named %s",filename.Data());
     return 0x0;
   }

  if (!(gAlice=(AliRun*)gAliceFile->Get("gAlice")))
   {
     Error("OpenFiles","gAlice have not been found on %s !\n",filename.Data());
     gAliceFile->Close();
     delete gAliceFile;
     return 0x0;
   }

  return gAliceFile;
}

/********************************************************************/
/********************************************************************/

<|endoftext|>


<commit_before>
#include "AliHBTReaderITSv2.h"

#include <iostream.h>
#include <fstream.h>
#include <TString.h>
#include <TObjString.h>
#include <TTree.h>
#include <TFile.h>
#include <TParticle.h>

#include <AliRun.h>
#include <AliMagF.h>
#include <AliITStrackV2.h>
//#include <AliITSParam.h>
#include <AliITStrackerV2.h>
#include <AliITSgeom.h>

#include "AliHBTRun.h"
#include "AliHBTEvent.h"
#include "AliHBTParticle.h"
#include "AliHBTParticleCut.h"

ClassImp(AliHBTReaderITSv2)

AliHBTReaderITSv2::
 AliHBTReaderITSv2(const Char_t* trackfilename, const Char_t* clusterfilename,
	const Char_t* galicefilename)
                  :fTrackFileName(trackfilename),fClusterFileName(clusterfilename),
                    fGAliceFileName(galicefilename)
{
  //constructor,
  //Defaults:
  //  trackfilename = "AliITStracksV2.root"
  //  clusterfilename = "AliITSclustersV2.root"
  //  galicefilename = "galice.root"
  fParticles = new AliHBTRun();
  fTracks    = new AliHBTRun();
  fIsRead = kFALSE;
}
/********************************************************************/

AliHBTReaderITSv2::
 AliHBTReaderITSv2(TObjArray* dirs, const Char_t* trackfilename,
                   const Char_t* clusterfilename, const Char_t* galicefilename)
                  : AliHBTReader(dirs),
                    fTrackFileName(trackfilename),fClusterFileName(clusterfilename),
                    fGAliceFileName(galicefilename)
{
  //constructor,
  //Defaults:
  //  trackfilename = "AliITStracksV2.root"
  //  clusterfilename = "AliITSclustersV2.root"
  //  galicefilename = "galice.root"

  fParticles = new AliHBTRun();
  fTracks    = new AliHBTRun();
  fIsRead = kFALSE;
}
/********************************************************************/

AliHBTReaderITSv2::~AliHBTReaderITSv2()
 {
   if (fParticles) delete fParticles;
   if (fTracks) delete fTracks;
 }
/********************************************************************/
/********************************************************************/

AliHBTEvent* AliHBTReaderITSv2::GetParticleEvent(Int_t n)
 {
 //returns Nth event with simulated particles
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetParticleEvent","Error in reading");
       return 0x0;
     }

   return fParticles->GetEvent(n);
 }
/********************************************************************/

AliHBTEvent* AliHBTReaderITSv2::GetTrackEvent(Int_t n)
 {
 //returns Nth event with reconstructed tracks
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetTrackEvent","Error in reading");
       return 0x0;
     }
   return fTracks->GetEvent(n);
 }
/********************************************************************/

Int_t AliHBTReaderITSv2::GetNumberOfPartEvents()
 {
 //returns number of events of particles
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetNumberOfPartEvents","Error in reading");
       return 0;
     }
   return fParticles->GetNumberOfEvents();
 }

/********************************************************************/
Int_t AliHBTReaderITSv2::GetNumberOfTrackEvents()
 {
 //returns number of events of tracks
  if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetNumberOfTrackEvents","Error in reading");
       return 0;
     }
  return fTracks->GetNumberOfEvents();
 }

/********************************************************************/
/********************************************************************/
Int_t AliHBTReaderITSv2::Read(AliHBTRun* particles, AliHBTRun *tracks)
{
 Int_t Nevents = 0; //number of events found in given directory
 Int_t Ndirs; //number of the directories to be read
 Int_t Ntracks; //number of tracks in current event
 Int_t currentdir = 0; //number of events in the current directory
 Int_t totalNevents = 0; //total number of events read from all directories up to now
 register Int_t i; //iterator

 TFile *aTracksFile;//file with tracks
 TFile *aClustersFile;//file with clusters
 TFile *aGAliceFile;//file name with galice

 AliITStrackerV2 *tracker; // ITS tracker - used for cooking labels
 TTree *tracktree; // tree for tracks

 Double_t xk;
 Double_t par[5]; //Kalman track parameters
 Float_t phi, lam, pt;//angles and transverse momentum
 Int_t label; //label of the current track

 char tname[100]; //buffer for tree name
 AliITStrackV2 *iotrack= new AliITStrackV2(); //buffer track for reading data from tree

 if (!particles) //check if an object is instatiated
  {
    Error("Read"," particles object must instatiated before passing it to the reader");
  }
 if (!tracks)  //check if an object is instatiated
  {
    Error("Read"," tracks object must instatiated before passing it to the reader");
  }
 particles->Reset();//clear runs == delete all old events
 tracks->Reset();

 if (fDirs) //if array with directories is supplied by user
  {
    Ndirs = fDirs->GetEntries(); //get the number if directories
  }
 else
  {
    Ndirs = 0; //if the array is not supplied read only from current directory
  }

// cout<<"Found "<<Ndirs<<" directory entries"<<endl;

 do //do while is good even if Ndirs==0 (than read from current directory)
   {
    if( (i=OpenFiles(aTracksFile,aClustersFile,aGAliceFile,currentdir)) )
     {
       Error("Read","Exiting due to problems with opening files. Errorcode %d",i);
       delete iotrack;
       return i;
     }

    if (gAlice->TreeE())//check if tree E exists
     {
      Nevents = (Int_t)gAlice->TreeE()->GetEntries();//if yes get number of events in gAlice
      cout<<"________________________________________________________\n";
      cout<<"Found "<<Nevents<<" event(s) in directory "<<GetDirName(currentdir)<<endl;
      cout<<"Setting Magnetic Field. Factor is "<<gAlice->Field()->Factor()<<endl;
      AliKalmanTrack::SetConvConst(100/0.299792458/0.2/gAlice->Field()->Factor());
     }
    else
     {//if not return an error
       Error("Read","Can not find Header tree (TreeE) in gAlice");
       delete iotrack;
       return 1;
     }

    AliITSgeom *geom=(AliITSgeom*)aClustersFile->Get("AliITSgeom");
    if (!geom)
     {
       Error("Read","Can't get the ITS geometry!");
       delete iotrack;
       return 2;
     }

    for(Int_t currentEvent =0; currentEvent<Nevents;currentEvent++)//loop over all events
     {
       cout<<"Reading Event "<<currentEvent<<endl;

       aGAliceFile->cd();
       gAlice->GetEvent(currentEvent);
       TParticle * part = gAlice->Particle(0);
       Double_t orz=part->Vz();

       aClustersFile->cd();
       tracker = new AliITStrackerV2(geom,currentEvent,orz);

       sprintf(tname,"TreeT_ITS_%d",currentEvent);

       tracktree=(TTree*)aTracksFile->Get(tname);
       if (!tracktree)
         {
           Error("Read","Can't get a tree with ITS tracks");
           delete iotrack;
           delete tracker;
           return 4;
         }
       TBranch *tbranch=tracktree->GetBranch("tracks");

       Ntracks=(Int_t)tracktree->GetEntries();

       Int_t accepted = 0;
       Int_t tpcfault = 0;
       Int_t itsfault = 0;
       for (i=0; i<Ntracks; i++) //loop over all tpc tracks
        {
          if(i%100 == 0)cout<<"all: "<<i<<"   accepted: "<<accepted<<"   tpc faults: "<<tpcfault<<"   its faults: "<<itsfault<<"\r";

          tbranch->SetAddress(&iotrack);
          tracktree->GetEvent(i);

          label=iotrack->GetLabel();
          if (label < 0)
           {
             tpcfault++;
             continue;
           }
          tracker->CookLabel(iotrack,0.);
          Int_t itsLabel=iotrack->GetLabel();
          if (itsLabel != label)
           {
            itsfault++;
            continue;
           }

          TParticle *p = (TParticle*)gAlice->Particle(label);
          if(Pass(p->GetPdgCode())) continue; //check if we are intersted with particles of this type
                                              //if not take next partilce

          AliHBTParticle* part = new AliHBTParticle(*p);
          if(Pass(part)) { delete part; continue;}//check if meets all criteria of any of our cuts
                                                  //if it does not delete it and take next good track

          iotrack->Propagate(iotrack->GetAlpha(),3.,0.1/65.19*1.848,0.1*1.848);
          iotrack->PropagateToVertex();

          iotrack->GetExternalParameters(xk,par);     //get properties of the track
          phi=TMath::ASin(par[2]) + iotrack->GetAlpha();
          if (phi<-TMath::Pi()) phi+=2*TMath::Pi();
          if (phi>=TMath::Pi()) phi-=2*TMath::Pi();
          lam=par[3];
          pt=1.0/TMath::Abs(par[4]);

          Double_t tpx = pt * TMath::Cos(phi); //track x coordinate of momentum
          Double_t tpy = pt * TMath::Sin(phi); //track y coordinate of momentum
          Double_t tpz = pt * lam; //track z coordinate of momentum

          Double_t mass = p->GetMass();
          Double_t tEtot = TMath::Sqrt( tpx*tpx + tpy*tpy + tpz*tpz + mass*mass);//total energy of the track

          AliHBTParticle* track = new AliHBTParticle(p->GetPdgCode(), tpx, tpy , tpz, tEtot, 0., 0., 0., 0.);
          if(Pass(track))//check if meets all criteria of any of our cuts
                         //if it does not delete it and take next good track
           {
            delete track;
            delete part;
            continue;
           }
          particles->AddParticle(totalNevents,part);//put track and particle on the run
          tracks->AddParticle(totalNevents,track);
          accepted++;
        }//end of loop over tracks in the event

       aTracksFile->Delete(tname);
       aTracksFile->Delete("tracks");
       delete tracker;

       totalNevents++;
       CloseFiles(aTracksFile,aClustersFile,aGAliceFile);
       cout<<"all: "<<i<<"   accepted: "<<accepted<<"   tpc faults: "<<tpcfault<<"   its faults: "<<itsfault<<endl;

     }//end of loop over events in current directory
     currentdir++;
   }while(currentdir < Ndirs);//end of loop over directories specified in fDirs Obj Array

 delete iotrack;
 fIsRead = kTRUE;
 return 0;
}

/********************************************************************/
Int_t AliHBTReaderITSv2::OpenFiles
(TFile*& aTracksFile, TFile*& aClustersFile, TFile*& agAliceFile,Int_t event)
{
 //opens all the files

   const TString& dirname = GetDirName(event);
   if (dirname == "")
    {
      Error("OpenFiles","Can not get directory name");
      return 4;
    }

   TString filename = dirname +"/"+ fTrackFileName;
   aTracksFile = TFile::Open(filename.Data());
   if ( aTracksFile  == 0x0 )
     {
       Error("OpenFiles","Can't open file with tacks named %s",filename.Data());
       return 1;
     }
   if (!aTracksFile->IsOpen())
     {
       Error("OpenFiles","Can't open file with tacks named %s",filename.Data());
       return 1;
     }

   filename = dirname +"/"+ fClusterFileName;
   aClustersFile = TFile::Open(filename.Data());
   if ( aClustersFile == 0x0 )
    {
      Error("OpenFiles","Can't open file with TPC clusters named %s",filename.Data());
      return 2;
    }
   if (!aClustersFile->IsOpen())
    {
      Error("OpenFiles","Can't open file with TPC clusters named %s",filename.Data());
      return 2;
    }

   filename = dirname +"/"+ fGAliceFileName;
   agAliceFile = TFile::Open(filename.Data());
   if ( agAliceFile== 0x0)
    {
      Error("OpenFiles","Can't open file with TPC clusters named %s",filename.Data());
      return 3;
    }
   if (!agAliceFile->IsOpen())
    {
      Error("OpenFiles","Can't open file with TPC clusters named %s",filename.Data());
      return 3;
    }

   if (!(gAlice=(AliRun*)agAliceFile->Get("gAlice")))
    {
      Error("OpenFiles","gAlice have not been found on %s !\n",filename.Data());
      return 5;
    }

   return 0;
}
/********************************************************************/

/********************************************************************/

void AliHBTReaderITSv2::CloseFiles(TFile*& tracksFile, TFile*& clustersFile, TFile*& gAliceFile)
{
  //closes the files
  tracksFile->Close();
  delete tracksFile;
  tracksFile = 0x0;

  clustersFile->Close();
  delete clustersFile;
  clustersFile = 0x0;

  delete gAlice;
  gAlice = 0;

  if (gAliceFile)
   {
     gAliceFile->Close();
     delete gAliceFile;
     gAliceFile = 0x0;
   }
}

/********************************************************************/

<commit_msg>Bug correction<commit_after>
#include "AliHBTReaderITSv2.h"

#include <iostream.h>
#include <fstream.h>
#include <TString.h>
#include <TObjString.h>
#include <TTree.h>
#include <TFile.h>
#include <TParticle.h>

#include <AliRun.h>
#include <AliMagF.h>
#include <AliITStrackV2.h>
//#include <AliITSParam.h>
#include <AliITStrackerV2.h>
#include <AliITSgeom.h>

#include "AliHBTRun.h"
#include "AliHBTEvent.h"
#include "AliHBTParticle.h"
#include "AliHBTParticleCut.h"

ClassImp(AliHBTReaderITSv2)

AliHBTReaderITSv2::
 AliHBTReaderITSv2(const Char_t* trackfilename, const Char_t* clusterfilename,
	const Char_t* galicefilename)
                  :fTrackFileName(trackfilename),fClusterFileName(clusterfilename),
                    fGAliceFileName(galicefilename)
{
  //constructor,
  //Defaults:
  //  trackfilename = "AliITStracksV2.root"
  //  clusterfilename = "AliITSclustersV2.root"
  //  galicefilename = "galice.root"
  fParticles = new AliHBTRun();
  fTracks    = new AliHBTRun();
  fIsRead = kFALSE;
}
/********************************************************************/

AliHBTReaderITSv2::
 AliHBTReaderITSv2(TObjArray* dirs, const Char_t* trackfilename,
                   const Char_t* clusterfilename, const Char_t* galicefilename)
                  : AliHBTReader(dirs),
                    fTrackFileName(trackfilename),fClusterFileName(clusterfilename),
                    fGAliceFileName(galicefilename)
{
  //constructor,
  //Defaults:
  //  trackfilename = "AliITStracksV2.root"
  //  clusterfilename = "AliITSclustersV2.root"
  //  galicefilename = "galice.root"

  fParticles = new AliHBTRun();
  fTracks    = new AliHBTRun();
  fIsRead = kFALSE;
}
/********************************************************************/

AliHBTReaderITSv2::~AliHBTReaderITSv2()
 {
   if (fParticles) delete fParticles;
   if (fTracks) delete fTracks;
 }
/********************************************************************/
/********************************************************************/

AliHBTEvent* AliHBTReaderITSv2::GetParticleEvent(Int_t n)
 {
 //returns Nth event with simulated particles
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetParticleEvent","Error in reading");
       return 0x0;
     }

   return fParticles->GetEvent(n);
 }
/********************************************************************/

AliHBTEvent* AliHBTReaderITSv2::GetTrackEvent(Int_t n)
 {
 //returns Nth event with reconstructed tracks
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetTrackEvent","Error in reading");
       return 0x0;
     }
   return fTracks->GetEvent(n);
 }
/********************************************************************/

Int_t AliHBTReaderITSv2::GetNumberOfPartEvents()
 {
 //returns number of events of particles
   if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetNumberOfPartEvents","Error in reading");
       return 0;
     }
   return fParticles->GetNumberOfEvents();
 }

/********************************************************************/
Int_t AliHBTReaderITSv2::GetNumberOfTrackEvents()
 {
 //returns number of events of tracks
  if (!fIsRead)
    if(Read(fParticles,fTracks))
     {
       Error("GetNumberOfTrackEvents","Error in reading");
       return 0;
     }
  return fTracks->GetNumberOfEvents();
 }

/********************************************************************/
/********************************************************************/
Int_t AliHBTReaderITSv2::Read(AliHBTRun* particles, AliHBTRun *tracks)
{
 Int_t Nevents = 0; //number of events found in given directory
 Int_t Ndirs; //number of the directories to be read
 Int_t Ntracks; //number of tracks in current event
 Int_t currentdir = 0; //number of events in the current directory
 Int_t totalNevents = 0; //total number of events read from all directories up to now
 register Int_t i; //iterator

 TFile *aTracksFile;//file with tracks
 TFile *aClustersFile;//file with clusters
 TFile *aGAliceFile;//file name with galice

 AliITStrackerV2 *tracker; // ITS tracker - used for cooking labels
 TTree *tracktree; // tree for tracks

 Double_t xk;
 Double_t par[5]; //Kalman track parameters
 Float_t phi, lam, pt;//angles and transverse momentum
 Int_t label; //label of the current track

 char tname[100]; //buffer for tree name
 AliITStrackV2 *iotrack= new AliITStrackV2(); //buffer track for reading data from tree

 if (!particles) //check if an object is instatiated
  {
    Error("Read"," particles object must instatiated before passing it to the reader");
  }
 if (!tracks)  //check if an object is instatiated
  {
    Error("Read"," tracks object must instatiated before passing it to the reader");
  }
 particles->Reset();//clear runs == delete all old events
 tracks->Reset();

 if (fDirs) //if array with directories is supplied by user
  {
    Ndirs = fDirs->GetEntries(); //get the number if directories
  }
 else
  {
    Ndirs = 0; //if the array is not supplied read only from current directory
  }

// cout<<"Found "<<Ndirs<<" directory entries"<<endl;

 do //do while is good even if Ndirs==0 (than read from current directory)
   {
    if( (i=OpenFiles(aTracksFile,aClustersFile,aGAliceFile,currentdir)) )
     {
       Error("Read","Exiting due to problems with opening files. Errorcode %d",i);
       delete iotrack;
       return i;
     }

    if (gAlice->TreeE())//check if tree E exists
     {
      Nevents = (Int_t)gAlice->TreeE()->GetEntries();//if yes get number of events in gAlice
      cout<<"________________________________________________________\n";
      cout<<"Found "<<Nevents<<" event(s) in directory "<<GetDirName(currentdir)<<endl;
      cout<<"Setting Magnetic Field. Factor is "<<gAlice->Field()->Factor()<<endl;
      AliKalmanTrack::SetConvConst(100/0.299792458/0.2/gAlice->Field()->Factor());
     }
    else
     {//if not return an error
       Error("Read","Can not find Header tree (TreeE) in gAlice");
       delete iotrack;
       return 1;
     }

    AliITSgeom *geom=(AliITSgeom*)aClustersFile->Get("AliITSgeom");
    if (!geom)
     {
       Error("Read","Can't get the ITS geometry!");
       delete iotrack;
       return 2;
     }

    for(Int_t currentEvent =0; currentEvent<Nevents;currentEvent++)//loop over all events
     {
       cout<<"Reading Event "<<currentEvent<<endl;

       aGAliceFile->cd();
       gAlice->GetEvent(currentEvent);
       TParticle * part = gAlice->Particle(0);
       Double_t orz=part->Vz();

       aClustersFile->cd();
       tracker = new AliITStrackerV2(geom,currentEvent,orz);

       sprintf(tname,"TreeT_ITS_%d",currentEvent);

       tracktree=(TTree*)aTracksFile->Get(tname);
       if (!tracktree)
         {
           Error("Read","Can't get a tree with ITS tracks");
           delete iotrack;
           delete tracker;
           return 4;
         }
       TBranch *tbranch=tracktree->GetBranch("tracks");

       Ntracks=(Int_t)tracktree->GetEntries();

       Int_t accepted = 0;
       Int_t tpcfault = 0;
       Int_t itsfault = 0;
       for (i=0; i<Ntracks; i++) //loop over all tpc tracks
        {
          if(i%100 == 0)cout<<"all: "<<i<<"   accepted: "<<accepted<<"   tpc faults: "<<tpcfault<<"   its faults: "<<itsfault<<"\r";

          tbranch->SetAddress(&iotrack);
          tracktree->GetEvent(i);

          label=iotrack->GetLabel();
          if (label < 0)
           {
             tpcfault++;
             continue;
           }
          tracker->CookLabel(iotrack,0.);
          Int_t itsLabel=iotrack->GetLabel();
          if (itsLabel != label)
           {
            itsfault++;
            continue;
           }

          TParticle *p = (TParticle*)gAlice->Particle(label);
          if(Pass(p->GetPdgCode())) continue; //check if we are intersted with particles of this type
                                              //if not take next partilce

          AliHBTParticle* part = new AliHBTParticle(*p);
          if(Pass(part)) { delete part; continue;}//check if meets all criteria of any of our cuts
                                                  //if it does not delete it and take next good track

          iotrack->Propagate(iotrack->GetAlpha(),3.,0.1/65.19*1.848,0.1*1.848);
          iotrack->PropagateToVertex();

          iotrack->GetExternalParameters(xk,par);     //get properties of the track
          phi=TMath::ASin(par[2]) + iotrack->GetAlpha();
          if (phi<-TMath::Pi()) phi+=2*TMath::Pi();
          if (phi>=TMath::Pi()) phi-=2*TMath::Pi();
          lam=par[3];
          pt=1.0/TMath::Abs(par[4]);

          Double_t tpx = pt * TMath::Cos(phi); //track x coordinate of momentum
          Double_t tpy = pt * TMath::Sin(phi); //track y coordinate of momentum
          Double_t tpz = pt * lam; //track z coordinate of momentum

          Double_t mass = p->GetMass();
          Double_t tEtot = TMath::Sqrt( tpx*tpx + tpy*tpy + tpz*tpz + mass*mass);//total energy of the track

          AliHBTParticle* track = new AliHBTParticle(p->GetPdgCode(), tpx, tpy , tpz, tEtot, 0., 0., 0., 0.);
          if(Pass(track))//check if meets all criteria of any of our cuts
                         //if it does not delete it and take next good track
           {
            delete track;
            delete part;
            continue;
           }
          particles->AddParticle(totalNevents,part);//put track and particle on the run
          tracks->AddParticle(totalNevents,track);
          accepted++;
        }//end of loop over tracks in the event

       aTracksFile->Delete(tname);
       aTracksFile->Delete("tracks");
       delete tracker;

       totalNevents++;
       cout<<"all: "<<i<<"   accepted: "<<accepted<<"   tpc faults: "<<tpcfault<<"   its faults: "<<itsfault<<endl;

     }//end of loop over events in current directory
     currentdir++;
     CloseFiles(aTracksFile,aClustersFile,aGAliceFile);

   }while(currentdir < Ndirs);//end of loop over directories specified in fDirs Obj Array

 delete iotrack;
 fIsRead = kTRUE;
 return 0;
}

/********************************************************************/
Int_t AliHBTReaderITSv2::OpenFiles
(TFile*& aTracksFile, TFile*& aClustersFile, TFile*& agAliceFile,Int_t event)
{
 //opens all the files

   const TString& dirname = GetDirName(event);
   if (dirname == "")
    {
      Error("OpenFiles","Can not get directory name");
      return 4;
    }

   TString filename = dirname +"/"+ fTrackFileName;
   aTracksFile = TFile::Open(filename.Data());
   if ( aTracksFile  == 0x0 )
     {
       Error("OpenFiles","Can't open file with tacks named %s",filename.Data());
       return 1;
     }
   if (!aTracksFile->IsOpen())
     {
       Error("OpenFiles","Can't open file with tacks named %s",filename.Data());
       return 1;
     }

   filename = dirname +"/"+ fClusterFileName;
   aClustersFile = TFile::Open(filename.Data());
   if ( aClustersFile == 0x0 )
    {
      Error("OpenFiles","Can't open file with TPC clusters named %s",filename.Data());
      return 2;
    }
   if (!aClustersFile->IsOpen())
    {
      Error("OpenFiles","Can't open file with TPC clusters named %s",filename.Data());
      return 2;
    }

   filename = dirname +"/"+ fGAliceFileName;
   agAliceFile = TFile::Open(filename.Data());
   if ( agAliceFile== 0x0)
    {
      Error("OpenFiles","Can't open file with TPC clusters named %s",filename.Data());
      return 3;
    }
   if (!agAliceFile->IsOpen())
    {
      Error("OpenFiles","Can't open file with TPC clusters named %s",filename.Data());
      return 3;
    }

   if (!(gAlice=(AliRun*)agAliceFile->Get("gAlice")))
    {
      Error("OpenFiles","gAlice have not been found on %s !\n",filename.Data());
      return 5;
    }

   return 0;
}
/********************************************************************/

/********************************************************************/

void AliHBTReaderITSv2::CloseFiles(TFile*& tracksFile, TFile*& clustersFile, TFile*& gAliceFile)
{
  //closes the files
  tracksFile->Close();
  delete tracksFile;
  tracksFile = 0x0;

  clustersFile->Close();
  delete clustersFile;
  clustersFile = 0x0;

  delete gAlice;
  gAlice = 0;

  if (gAliceFile)
   {
     gAliceFile->Close();
     delete gAliceFile;
     gAliceFile = 0x0;
   }
}

/********************************************************************/

<|endoftext|>