225 lines
7.9 KiB
C++
225 lines
7.9 KiB
C++
// Copyright 2019 The MediaPipe Authors.
|
|
//
|
|
// 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 "mediapipe/util/tracking/tracking_visualization_utilities.h"
|
|
|
|
#include "absl/strings/str_format.h"
|
|
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
|
|
#include "mediapipe/util/tracking/box_tracker.h"
|
|
#include "mediapipe/util/tracking/tracking.h"
|
|
|
|
namespace mediapipe {
|
|
|
|
void RenderState(const MotionBoxState& box_state, bool print_stats,
|
|
cv::Mat* frame) {
|
|
#ifndef NO_RENDERING
|
|
CHECK(frame != nullptr);
|
|
|
|
const int frame_width = frame->cols;
|
|
const int frame_height = frame->rows;
|
|
|
|
const Vector2_f top_left(box_state.pos_x() * frame_width,
|
|
box_state.pos_y() * frame_height);
|
|
|
|
cv::rectangle(*frame,
|
|
cv::Point(box_state.inlier_center_x() * frame_width - 2,
|
|
box_state.inlier_center_y() * frame_height - 2),
|
|
cv::Point(box_state.inlier_center_x() * frame_width + 2,
|
|
box_state.inlier_center_y() * frame_height + 2),
|
|
cv::Scalar(255, 255, 0, 255), 2);
|
|
|
|
cv::rectangle(
|
|
*frame,
|
|
cv::Point(
|
|
(box_state.inlier_center_x() - box_state.inlier_width() * 0.5) *
|
|
frame_width,
|
|
(box_state.inlier_center_y() - box_state.inlier_height() * 0.5) *
|
|
frame_height),
|
|
cv::Point(
|
|
(box_state.inlier_center_x() + box_state.inlier_width() * 0.5) *
|
|
frame_width,
|
|
(box_state.inlier_center_y() + box_state.inlier_height() * 0.5) *
|
|
frame_height),
|
|
cv::Scalar(0, 0, 255, 255), 1);
|
|
|
|
std::vector<Vector2_f> inlier_locs;
|
|
std::vector<Vector2_f> outlier_locs;
|
|
MotionBoxInlierLocations(box_state, &inlier_locs);
|
|
MotionBoxOutlierLocations(box_state, &outlier_locs);
|
|
float scale_x = 1.0f;
|
|
float scale_y = 1.0f;
|
|
ScaleFromAspect(frame_width * 1.0f / frame_height, true, &scale_x, &scale_y);
|
|
|
|
for (const Vector2_f& loc : inlier_locs) {
|
|
cv::circle(*frame,
|
|
cv::Point(loc.x() * scale_x * frame_width,
|
|
loc.y() * scale_y * frame_height),
|
|
4.0, cv::Scalar(0, 255, 0, 128), 1);
|
|
}
|
|
for (const Vector2_f& loc : outlier_locs) {
|
|
cv::circle(*frame,
|
|
cv::Point(loc.x() * scale_x * frame_width,
|
|
loc.y() * scale_y * frame_height),
|
|
4.0, cv::Scalar(255, 0, 0, 128), 1);
|
|
}
|
|
|
|
if (!print_stats) {
|
|
return;
|
|
}
|
|
|
|
std::vector<std::string> stats;
|
|
stats.push_back(
|
|
absl::StrFormat("Motion: %.4f, %.4f", box_state.dx(), box_state.dy()));
|
|
stats.push_back(
|
|
absl::StrFormat("KinEnergy: %.4f", box_state.kinetic_energy()));
|
|
stats.push_back(
|
|
absl::StrFormat("Disparity: %.2f", box_state.motion_disparity()));
|
|
stats.push_back(absl::StrFormat("Discrimination: %.2f",
|
|
box_state.background_discrimination()));
|
|
stats.push_back(
|
|
absl::StrFormat("InlierRatio: %2.2f", box_state.inlier_ratio()));
|
|
stats.push_back(
|
|
absl::StrFormat("InlierNum: %3d", box_state.inlier_ids_size()));
|
|
|
|
stats.push_back(absl::StrFormat("Prior: %.2f", box_state.prior_weight()));
|
|
stats.push_back(absl::StrFormat("TrackingConfidence: %.2f",
|
|
box_state.tracking_confidence()));
|
|
|
|
for (int k = 0; k < stats.size(); ++k) {
|
|
// Display some stats below the box.
|
|
cv::putText(*frame, stats[k],
|
|
cv::Point(top_left.x(), top_left.y() + (k + 1) * 12),
|
|
cv::FONT_HERSHEY_PLAIN, 1.0, cv::Scalar(0, 0, 0, 255));
|
|
|
|
cv::putText(*frame, stats[k],
|
|
cv::Point(top_left.x() + 1, top_left.y() + (k + 1) * 12 + 1),
|
|
cv::FONT_HERSHEY_PLAIN, 1.0, cv::Scalar(255, 255, 255, 255));
|
|
}
|
|
|
|
// Visualize locking state via heuristically chosen thresholds. Locking
|
|
// state is purely for visualization/illustration and has no further
|
|
// meaning.
|
|
std::string lock_text;
|
|
cv::Scalar lock_color;
|
|
if (box_state.motion_disparity() > 0.8) {
|
|
lock_text = "Lock lost";
|
|
lock_color = cv::Scalar(255, 0, 0, 255);
|
|
} else if (box_state.motion_disparity() > 0.4) {
|
|
lock_text = "Aquiring lock";
|
|
lock_color = cv::Scalar(255, 255, 0, 255);
|
|
} else if (box_state.motion_disparity() > 0.1) {
|
|
lock_text = "Locked";
|
|
lock_color = cv::Scalar(0, 255, 0, 255);
|
|
}
|
|
|
|
cv::putText(*frame, lock_text, cv::Point(top_left.x() + 1, top_left.y() - 4),
|
|
cv::FONT_HERSHEY_PLAIN, 0.8, cv::Scalar(255, 255, 255, 255));
|
|
|
|
cv::putText(*frame, lock_text, cv::Point(top_left.x(), top_left.y() - 5),
|
|
cv::FONT_HERSHEY_PLAIN, 0.8, lock_color);
|
|
#else
|
|
LOG(FATAL) << "Code stripped out because of NO_RENDERING";
|
|
#endif
|
|
}
|
|
|
|
void RenderInternalState(const MotionBoxInternalState& internal,
|
|
cv::Mat* frame) {
|
|
#ifndef NO_RENDERING
|
|
CHECK(frame != nullptr);
|
|
|
|
const int num_vectors = internal.pos_x_size();
|
|
|
|
// Determine max inlier score for visualization.
|
|
float max_score = 0;
|
|
for (int k = 0; k < num_vectors; ++k) {
|
|
max_score = std::max(max_score, internal.inlier_score(k));
|
|
}
|
|
|
|
float alpha_scale = 1.0f;
|
|
if (max_score > 0) {
|
|
alpha_scale = 1.0f / max_score;
|
|
}
|
|
|
|
const int frame_width = frame->cols;
|
|
const int frame_height = frame->rows;
|
|
for (int k = 0; k < num_vectors; ++k) {
|
|
const MotionVector v = MotionVector::FromInternalState(internal, k);
|
|
cv::Point p1(v.pos.x() * frame_width, v.pos.y() * frame_height);
|
|
Vector2_f match = v.pos + v.object;
|
|
|
|
cv::Point p2(match.x() * frame_width, match.y() * frame_height);
|
|
const float alpha = internal.inlier_score(k) * alpha_scale;
|
|
|
|
cv::Scalar color_scaled(0 * alpha + (1.0f - alpha) * 255,
|
|
255 * alpha + (1.0f - alpha) * 0,
|
|
0 * alpha + (1.0f - alpha) * 0);
|
|
|
|
cv::line(*frame, p1, p2, color_scaled, 1, cv::LINE_AA);
|
|
cv::circle(*frame, p1, 2.0, color_scaled, 1);
|
|
}
|
|
#else
|
|
LOG(FATAL) << "Code stripped out because of NO_RENDERING";
|
|
#endif
|
|
}
|
|
|
|
void RenderTrackingData(const TrackingData& data, cv::Mat* mat,
|
|
bool antialiasing) {
|
|
#ifndef NO_RENDERING
|
|
CHECK(mat != nullptr);
|
|
|
|
MotionVectorFrame mvf;
|
|
MotionVectorFrameFromTrackingData(data, &mvf);
|
|
const float aspect_ratio = mvf.aspect_ratio;
|
|
float scale_x, scale_y;
|
|
ScaleFromAspect(aspect_ratio, true, &scale_x, &scale_y);
|
|
scale_x *= mat->cols;
|
|
scale_y *= mat->rows;
|
|
|
|
for (const auto& motion_vector : mvf.motion_vectors) {
|
|
Vector2_f pos = motion_vector.Location();
|
|
Vector2_f match = motion_vector.MatchLocation();
|
|
|
|
cv::Point p1(pos.x() * scale_x, pos.y() * scale_y);
|
|
cv::Point p2(match.x() * scale_x, match.y() * scale_y);
|
|
|
|
// iOS cannot display a width 1 antialiased line, so we provide
|
|
// an option for 8-connected drawing instead.
|
|
cv::line(*mat, p1, p2, cv::Scalar(0, 255, 0, 255), 1,
|
|
antialiasing ? cv::LINE_AA : 8);
|
|
}
|
|
#else
|
|
LOG(FATAL) << "Code stripped out because of NO_RENDERING";
|
|
#endif
|
|
}
|
|
|
|
void RenderBox(const TimedBoxProto& box_proto, cv::Mat* mat) {
|
|
#ifndef NO_RENDERING
|
|
CHECK(mat != nullptr);
|
|
|
|
TimedBox box = TimedBox::FromProto(box_proto);
|
|
std::array<Vector2_f, 4> corners = box.Corners(mat->cols, mat->rows);
|
|
|
|
for (int k = 0; k < 4; ++k) {
|
|
cv::line(*mat, cv::Point2f(corners[k].x(), corners[k].y()),
|
|
cv::Point2f(corners[(k + 1) % 4].x(), corners[(k + 1) % 4].y()),
|
|
cv::Scalar(255, 0, 0, 255), // Red.
|
|
4);
|
|
}
|
|
#else
|
|
LOG(FATAL) << "Code stripped out because of NO_RENDERING";
|
|
#endif
|
|
}
|
|
|
|
} // namespace mediapipe
|