208 lines
6.8 KiB
C++
208 lines
6.8 KiB
C++
// 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>
|
|
#include <iomanip>
|
|
#include <iostream>
|
|
#include "include/postprocess.h"
|
|
|
|
namespace PaddleDetection {
|
|
|
|
cv::Scalar GetColor(int idx) {
|
|
idx = idx * 3;
|
|
cv::Scalar color =
|
|
cv::Scalar((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255);
|
|
return color;
|
|
}
|
|
|
|
cv::Mat VisualizeTrackResult(const cv::Mat& img,
|
|
const MOTResult& results,
|
|
const float fps,
|
|
const int frame_id) {
|
|
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.));
|
|
|
|
std::ostringstream oss;
|
|
oss << std::setiosflags(std::ios::fixed) << std::setprecision(4);
|
|
oss << "frame: " << frame_id << " ";
|
|
oss << "fps: " << fps << " ";
|
|
oss << "num: " << results.size();
|
|
std::string text = oss.str();
|
|
|
|
cv::Point origin;
|
|
origin.x = 0;
|
|
origin.y = static_cast<int>(15 * text_scale);
|
|
cv::putText(vis_img,
|
|
text,
|
|
origin,
|
|
cv::FONT_HERSHEY_PLAIN,
|
|
text_scale,
|
|
(0, 0, 255),
|
|
2);
|
|
|
|
for (int i = 0; i < results.size(); ++i) {
|
|
const int obj_id = results[i].ids;
|
|
const float score = results[i].score;
|
|
|
|
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);
|
|
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);
|
|
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;
|
|
}
|
|
|
|
void FlowStatistic(const MOTResult& results,
|
|
const int frame_id,
|
|
const int secs_interval,
|
|
const bool do_entrance_counting,
|
|
const int video_fps,
|
|
const Rect entrance,
|
|
std::set<int>* id_set,
|
|
std::set<int>* interval_id_set,
|
|
std::vector<int>* in_id_list,
|
|
std::vector<int>* out_id_list,
|
|
std::map<int, std::vector<float>>* prev_center,
|
|
std::vector<std::string>* records) {
|
|
if (frame_id == 0) interval_id_set->clear();
|
|
|
|
if (do_entrance_counting) {
|
|
// Count in and out number:
|
|
// Use horizontal center line as the entrance just for simplification.
|
|
// If a person located in the above the horizontal center line
|
|
// at the previous frame and is in the below the line at the current frame,
|
|
// the in number is increased by one.
|
|
// 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,
|
|
// the out number is increased by one.
|
|
// TODO(qianhui): if the entrance is not the horizontal center line,
|
|
// the counting method should be optimized.
|
|
|
|
float entrance_y = entrance.top;
|
|
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;
|
|
std::map<int, std::vector<float>>::iterator iter;
|
|
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(
|
|
std::pair<int, std::vector<float>>(ids, {center_x, center_y}));
|
|
}
|
|
}
|
|
}
|
|
|
|
// 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;
|
|
os << "Frame id: " << frame_id << ", Total count: " << id_set->size();
|
|
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;
|
|
}
|
|
|
|
void SaveMOTResult(const MOTResult& results,
|
|
const int frame_id,
|
|
std::vector<std::string>* records) {
|
|
// 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;
|
|
os << frame_id << " " << ids << " " << x1 << " " << y1 << " " << w << " "
|
|
<< h << "\n";
|
|
record = os.str();
|
|
records->push_back(record);
|
|
}
|
|
}
|
|
|
|
} // namespace PaddleDetection
|