FastDeploy/fastdeploy/vision/visualize/mot.cc

// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// 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.

#ifdef ENABLE_VISION_VISUALIZE
#include "fastdeploy/vision/visualize/visualize.h"
#include <iomanip>

namespace fastdeploy {
namespace vision {

cv::Scalar GetMOTBoxColor(int idx) {
  idx = idx * 3;
  cv::Scalar color = cv::Scalar((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255);
  return color;
}

cv::Mat VisMOT(const cv::Mat &img, const MOTResult &results,
               float score_threshold, tracking::TrailRecorder* recorder) {
  cv::Mat vis_img = img.clone();
  int im_h = img.rows;
  int im_w = img.cols;
  float text_scale = std::max(1, static_cast<int>(im_w / 1600.));
  float text_thickness = 2.;
  float line_thickness = std::max(1, static_cast<int>(im_w / 500.));
  for (int i = 0; i < results.boxes.size(); ++i) {
    if (results.scores[i] < score_threshold) {
        continue;
    }
    const int obj_id = results.ids[i];
    const float score = results.scores[i];
    cv::Scalar color = GetMOTBoxColor(obj_id);
    if (recorder != nullptr){
      int id = results.ids[i];
      auto iter = recorder->records.find(id);
      if (iter != recorder->records.end()) {
        for (int j = 0; j < iter->second.size(); j++) {
            cv::Point center(iter->second[j][0], iter->second[j][1]);
            cv::circle(vis_img, center, text_thickness, color);
        }
      }
    }
    cv::Point pt1 = cv::Point(results.boxes[i][0], results.boxes[i][1]);
    cv::Point pt2 = cv::Point(results.boxes[i][2], results.boxes[i][3]);
    cv::Point id_pt =
            cv::Point(results.boxes[i][0], results.boxes[i][1] + 10);
    cv::Point score_pt =
            cv::Point(results.boxes[i][0], results.boxes[i][1] - 10);
    cv::rectangle(vis_img, pt1, pt2, color, line_thickness);
    std::ostringstream idoss;
    idoss << std::setiosflags(std::ios::fixed) << std::setprecision(4);
    idoss << obj_id;
    std::string id_text = idoss.str();

    cv::putText(vis_img,
                id_text,
                id_pt,
                cv::FONT_HERSHEY_PLAIN,
                text_scale,
                color,
                text_thickness);

    std::ostringstream soss;
    soss << std::setiosflags(std::ios::fixed) << std::setprecision(2);
    soss << score;
    std::string score_text = soss.str();

    cv::putText(vis_img,
                score_text,
                score_pt,
                cv::FONT_HERSHEY_PLAIN,
                text_scale,
                color,
                text_thickness);
  }
  return vis_img;
}
}// namespace vision
} //namespace fastdepoly
#endif