postprocess.cc 6.8 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
//   Copyright (c) 2021 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.

#include <sstream>
// for setprecision
#include <chrono>
18
#include <iomanip>
19
#include <iostream>
20 21 22 23 24 25
#include "include/postprocess.h"

namespace PaddleDetection {

cv::Scalar GetColor(int idx) {
  idx = idx * 3;
26 27
  cv::Scalar color =
      cv::Scalar((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255);
28 29 30 31
  return color;
}

cv::Mat VisualizeTrackResult(const cv::Mat& img,
32 33 34
                             const MOTResult& results,
                             const float fps,
                             const int frame_id) {
35 36 37
  cv::Mat vis_img = img.clone();
  int im_h = img.rows;
  int im_w = img.cols;
38
  float text_scale = std::max(1, static_cast<int>(im_w / 1600.));
39
  float text_thickness = 2.;
40
  float line_thickness = std::max(1, static_cast<int>(im_w / 500.));
41 42 43

  std::ostringstream oss;
  oss << std::setiosflags(std::ios::fixed) << std::setprecision(4);
44 45
  oss << "frame: " << frame_id << " ";
  oss << "fps: " << fps << " ";
46 47 48 49 50
  oss << "num: " << results.size();
  std::string text = oss.str();

  cv::Point origin;
  origin.x = 0;
51 52 53 54 55 56 57 58
  origin.y = static_cast<int>(15 * text_scale);
  cv::putText(vis_img,
              text,
              origin,
              cv::FONT_HERSHEY_PLAIN,
              text_scale,
              (0, 0, 255),
              2);
59 60 61 62

  for (int i = 0; i < results.size(); ++i) {
    const int obj_id = results[i].ids;
    const float score = results[i].score;
63

64 65 66 67
    cv::Scalar color = GetColor(obj_id);

    cv::Point pt1 = cv::Point(results[i].rects.left, results[i].rects.top);
    cv::Point pt2 = cv::Point(results[i].rects.right, results[i].rects.bottom);
68 69 70 71
    cv::Point id_pt =
        cv::Point(results[i].rects.left, results[i].rects.top + 10);
    cv::Point score_pt =
        cv::Point(results[i].rects.left, results[i].rects.top - 10);
72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
    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,
                cv::Scalar(0, 255, 255),
                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,
                cv::Scalar(0, 255, 255),
                text_thickness);
  }
  return vis_img;
}

103 104
void FlowStatistic(const MOTResult& results,
                   const int frame_id,
105 106 107
                   const int secs_interval,
                   const bool do_entrance_counting,
                   const int video_fps,
108
                   const Rect entrance,
109
                   std::set<int>* id_set,
110 111
                   std::set<int>* interval_id_set,
                   std::vector<int>* in_id_list,
112 113 114
                   std::vector<int>* out_id_list,
                   std::map<int, std::vector<float>>* prev_center,
                   std::vector<std::string>* records) {
115 116
  if (frame_id == 0) interval_id_set->clear();

117
  if (do_entrance_counting) {
118
    // Count in and out number:
119
    // Use horizontal center line as the entrance just for simplification.
120
    // If a person located in the above the horizontal center line
121 122
    // at the previous frame and is in the below the line at the current frame,
    // the in number is increased by one.
123 124 125
    // If a person was in the below the horizontal center line
    // at the previous frame and locates in the below the line at the current
    // frame,
126
    // the out number is increased by one.
127
    // TODO(qianhui): if the entrance is not the horizontal center line,
128
    // the counting method should be optimized.
129 130

    float entrance_y = entrance.top;
131 132 133 134
    for (const auto& result : results) {
      float center_x = (result.rects.left + result.rects.right) / 2;
      float center_y = (result.rects.top + result.rects.bottom) / 2;
      int ids = result.ids;
135
      std::map<int, std::vector<float>>::iterator iter;
136 137 138 139 140 141 142 143 144 145 146 147
      iter = prev_center->find(ids);
      if (iter != prev_center->end()) {
        if (iter->second[1] <= entrance_y && center_y > entrance_y) {
          in_id_list->push_back(ids);
        }
        if (iter->second[1] >= entrance_y && center_y < entrance_y) {
          out_id_list->push_back(ids);
        }
        (*prev_center)[ids][0] = center_x;
        (*prev_center)[ids][1] = center_y;
      } else {
        prev_center->insert(
148
            std::pair<int, std::vector<float>>(ids, {center_x, center_y}));
149 150 151 152 153 154 155 156 157 158 159
      }
    }
  }

  // Count totol number, number at a manual-setting interval
  for (const auto& result : results) {
    id_set->insert(result.ids);
    interval_id_set->insert(result.ids);
  }

  std::ostringstream os;
160
  os << "Frame id: " << frame_id << ", Total count: " << id_set->size();
161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177
  if (do_entrance_counting) {
    os << ", In count: " << in_id_list->size()
       << ", Out count: " << out_id_list->size();
  }

  // Reset counting at the interval beginning
  int curr_interval_count = -1;
  if (frame_id % video_fps == 0 && frame_id / video_fps % secs_interval == 0) {
    curr_interval_count = interval_id_set->size();
    os << ", Count during " << secs_interval
       << " secs: " << curr_interval_count;
    interval_id_set->clear();
  }
  os << "\n";
  std::string record = os.str();
  records->push_back(record);
  LOG(INFO) << record;
178 179
}

180 181 182
void SaveMOTResult(const MOTResult& results,
                   const int frame_id,
                   std::vector<std::string>* records) {
W
wangguanzhong 已提交
183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199
  // result format: frame_id, track_id, x1, y1, w, h
  std::string record;
  for (int i = 0; i < results.size(); ++i) {
    MOTTrack mot_track = results[i];
    int ids = mot_track.ids;
    float score = mot_track.score;
    Rect rects = mot_track.rects;
    float x1 = rects.left;
    float y1 = rects.top;
    float x2 = rects.right;
    float y2 = rects.bottom;
    float w = x2 - x1;
    float h = y2 - y1;
    if (w == 0 || h == 0) {
      continue;
    }
    std::ostringstream os;
200 201
    os << frame_id << " " << ids << "" << x1 << " " << y1 << " " << w << " "
       << h << "\n";
W
wangguanzhong 已提交
202
    record = os.str();
203
    records->push_back(record);
W
wangguanzhong 已提交
204
  }
205 206
}

207
}  // namespace PaddleDetection