cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
2044e45b1a
mediapipe/mediapipe2/util/tracking/box_tracker.cc
AK391 2044e45b1a examples
2021-06-10 23:01:19 +00:00

1063 lines
34 KiB
C++
Raw Blame History

// 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/box_tracker.h"
#include <sys/stat.h>
#include <fstream>
#include <limits>
#include "absl/strings/str_cat.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/util/tracking/measure_time.h"
#include "mediapipe/util/tracking/tracking.pb.h"
namespace mediapipe {
// Time within which close checkpoints are removed.
static constexpr int kSnapMs = 1000;
static constexpr int kInitCheckpoint = -1;
void MotionBoxStateQuadToVertices(const MotionBoxState::Quad& quad,
std::vector<Vector2_f>* vertices) {
CHECK_EQ(TimedBox::kNumQuadVertices * 2, quad.vertices_size());
CHECK(vertices != nullptr);
vertices->clear();
for (int i = 0; i < TimedBox::kNumQuadVertices; ++i) {
vertices->push_back(
Vector2_f(quad.vertices(i * 2), quad.vertices(i * 2 + 1)));
}
}
void VerticesToMotionBoxStateQuad(const std::vector<Vector2_f>& vertices,
MotionBoxState::Quad* quad) {
CHECK_EQ(TimedBox::kNumQuadVertices, vertices.size());
CHECK(quad != nullptr);
for (const Vector2_f& vertex : vertices) {
quad->add_vertices(vertex.x());
quad->add_vertices(vertex.y());
}
}
void MotionBoxStateFromTimedBox(const TimedBox& box, MotionBoxState* state) {
CHECK(state);
state->set_pos_x(box.left);
state->set_pos_y(box.top);
state->set_width(box.right - box.left);
state->set_height(box.bottom - box.top);
state->set_rotation(box.rotation);
state->set_request_grouping(box.request_grouping);
if (box.quad_vertices.size() == TimedBox::kNumQuadVertices) {
VerticesToMotionBoxStateQuad(box.quad_vertices, state->mutable_quad());
if (box.aspect_ratio > 0.0f) {
state->set_aspect_ratio(box.aspect_ratio);
}
// set pos_x and pos_y to be the top-left vertex x and y coordinates
// set width and height to be max - min of x and y.
float min_x = std::numeric_limits<float>::max();
float max_x = std::numeric_limits<float>::lowest();
float min_y = std::numeric_limits<float>::max();
float max_y = std::numeric_limits<float>::lowest();
for (const auto& vertex : box.quad_vertices) {
min_x = std::min(min_x, vertex.x());
max_x = std::max(max_x, vertex.x());
min_y = std::min(min_y, vertex.y());
max_y = std::max(max_y, vertex.y());
}
state->set_pos_x(min_x);
state->set_pos_y(min_y);
state->set_width(max_x - min_x);
state->set_height(max_y - min_y);
}
}
void TimedBoxFromMotionBoxState(const MotionBoxState& state, TimedBox* box) {
CHECK(box);
const float scale_dx = state.width() * (state.scale() - 1.0f) * 0.5f;
const float scale_dy = state.height() * (state.scale() - 1.0f) * 0.5f;
box->left = state.pos_x() - scale_dx;
box->top = state.pos_y() - scale_dy;
box->right = state.pos_x() + state.width() + scale_dx;
box->bottom = state.pos_y() + state.height() + scale_dy;
box->rotation = state.rotation();
box->confidence = state.tracking_confidence();
box->request_grouping = state.request_grouping();
if (state.has_quad()) {
MotionBoxStateQuadToVertices(state.quad(), &(box->quad_vertices));
if (state.has_aspect_ratio()) {
box->aspect_ratio = state.aspect_ratio();
}
}
}
namespace {
TimedBox BlendTimedBoxes(const TimedBox& lhs, const TimedBox& rhs,
int64 time_msec) {
CHECK_LT(lhs.time_msec, rhs.time_msec);
const double alpha =
(time_msec - lhs.time_msec) * 1.0 / (rhs.time_msec - lhs.time_msec);
return TimedBox::Blend(lhs, rhs, alpha);
}
} // namespace.
TimedBox TimedBox::Blend(const TimedBox& lhs, const TimedBox& rhs, double alpha,
double beta) {
// Due to large timestamps alpha beta should be in double.
TimedBox result;
result.top = alpha * lhs.top + beta * rhs.top;
result.left = alpha * lhs.left + beta * rhs.left;
result.bottom = alpha * lhs.bottom + beta * rhs.bottom;
result.right = alpha * lhs.right + beta * rhs.right;
result.rotation = alpha * lhs.rotation + beta * rhs.rotation;
result.time_msec = std::round(alpha * lhs.time_msec + beta * rhs.time_msec);
result.confidence = alpha * lhs.confidence + beta * rhs.confidence;
if (lhs.quad_vertices.size() == kNumQuadVertices &&
rhs.quad_vertices.size() == kNumQuadVertices) {
result.quad_vertices.clear();
for (int i = 0; i < lhs.quad_vertices.size(); ++i) {
result.quad_vertices.push_back(alpha * lhs.quad_vertices[i] +
beta * rhs.quad_vertices[i]);
}
// Since alpha and beta are not necessarily sum to 1, aspect ratio can not
// be derived with alpha and beta. Here we are simply averaging the
// aspect_ratio as the blended box aspect ratio.
if (lhs.aspect_ratio > 0 && rhs.aspect_ratio > 0) {
result.aspect_ratio = 0.5f * lhs.aspect_ratio + 0.5f * rhs.aspect_ratio;
}
}
return result;
}
TimedBox TimedBox::Blend(const TimedBox& lhs, const TimedBox& rhs,
double alpha) {
return Blend(lhs, rhs, 1.0 - alpha, alpha);
}
std::array<Vector2_f, 4> TimedBox::Corners(float width, float height) const {
if (quad_vertices.size() == kNumQuadVertices) {
std::array<Vector2_f, 4> corners{{
Vector2_f(quad_vertices[0].x() * width, quad_vertices[0].y() * height),
Vector2_f(quad_vertices[1].x() * width, quad_vertices[1].y() * height),
Vector2_f(quad_vertices[2].x() * width, quad_vertices[2].y() * height),
Vector2_f(quad_vertices[3].x() * width, quad_vertices[3].y() * height),
}};
return corners;
} else {
// Rotate 4 corner w.r.t. center.
const Vector2_f center(0.5f * (left + right) * width,
0.5f * (top + bottom) * height);
const std::array<Vector2_f, 4> corners{{
Vector2_f(left * width, top * height),
Vector2_f(left * width, bottom * height),
Vector2_f(right * width, bottom * height),
Vector2_f(right * width, top * height),
}};
const float cos_a = std::cos(rotation);
const float sin_a = std::sin(rotation);
std::array<Vector2_f, 4> transformed_corners;
for (int k = 0; k < 4; ++k) {
// Scale and rotate w.r.t. center.
const Vector2_f rad = corners[k] - center;
const Vector2_f rot_rad(cos_a * rad.x() - sin_a * rad.y(),
sin_a * rad.x() + cos_a * rad.y());
transformed_corners[k] = center + rot_rad;
}
return transformed_corners;
}
}
TimedBox TimedBox::FromProto(const TimedBoxProto& proto) {
TimedBox box;
box.top = proto.top();
box.left = proto.left();
box.bottom = proto.bottom();
box.right = proto.right();
box.rotation = proto.rotation();
box.time_msec = proto.time_msec();
box.request_grouping = proto.request_grouping();
if (proto.has_quad() &&
proto.quad().vertices_size() == kNumQuadVertices * 2) {
MotionBoxStateQuadToVertices(proto.quad(), &(box.quad_vertices));
if (proto.has_aspect_ratio()) {
box.aspect_ratio = proto.aspect_ratio();
}
}
return box;
}
TimedBoxProto TimedBox::ToProto() const {
TimedBoxProto proto;
proto.set_top(top);
proto.set_left(left);
proto.set_bottom(bottom);
proto.set_right(right);
proto.set_rotation(rotation);
proto.set_time_msec(time_msec);
proto.set_confidence(confidence);
proto.set_request_grouping(request_grouping);
if (quad_vertices.size() == kNumQuadVertices) {
VerticesToMotionBoxStateQuad(quad_vertices, proto.mutable_quad());
if (aspect_ratio > 0.0f) {
proto.set_aspect_ratio(aspect_ratio);
}
}
return proto;
}
BoxTracker::BoxTracker(const std::string& cache_dir,
const BoxTrackerOptions& options)
: options_(options), cache_dir_(cache_dir) {
tracking_workers_.reset(new ThreadPool(options_.num_tracking_workers()));
tracking_workers_->StartWorkers();
}
BoxTracker::BoxTracker(
const std::vector<const TrackingDataChunk*>& tracking_data, bool copy_data,
const BoxTrackerOptions& options)
: BoxTracker("", options) {
AddTrackingDataChunks(tracking_data, copy_data);
}
void BoxTracker::AddTrackingDataChunk(const TrackingDataChunk* chunk,
bool copy_data) {
CHECK_GT(chunk->item_size(), 0) << "Empty chunk.";
int64 chunk_time_msec = chunk->item(0).timestamp_usec() / 1000;
int chunk_idx = ChunkIdxFromTime(chunk_time_msec);
CHECK_GE(chunk_idx, tracking_data_.size()) << "Chunk is out of order.";
if (chunk_idx > tracking_data_.size()) {
LOG(INFO) << "Resize tracking_data_ to " << chunk_idx;
tracking_data_.resize(chunk_idx);
}
if (copy_data) {
tracking_data_buffer_.emplace_back(new TrackingDataChunk(*chunk));
tracking_data_.push_back(tracking_data_buffer_.back().get());
} else {
tracking_data_.emplace_back(chunk);
}
}
void BoxTracker::AddTrackingDataChunks(
const std::vector<const TrackingDataChunk*>& tracking_data,
bool copy_data) {
for (const auto item : tracking_data) {
AddTrackingDataChunk(item, copy_data);
}
}
void BoxTracker::NewBoxTrack(const TimedBox& initial_pos, int id,
int64 min_msec, int64 max_msec) {
VLOG(1) << "New box track: " << id << " : " << initial_pos.ToString()
<< " from " << min_msec << " to " << max_msec;
// Mark initialization with checkpoint -1.
absl::MutexLock lock(&status_mutex_);
if (canceling_) {
LOG(WARNING) << "Box Tracker is in cancel state. Refusing request.";
return;
}
++track_status_[id][kInitCheckpoint].tracks_ongoing;
auto operation = [this, initial_pos, id, min_msec, max_msec]() {
this->NewBoxTrackAsync(initial_pos, id, min_msec, max_msec);
};
tracking_workers_->Schedule(operation);
}
std::pair<int64, int64> BoxTracker::TrackInterval(int id) {
absl::MutexLock lock(&path_mutex_);
const Path& path = paths_[id];
if (path.empty()) {
return std::make_pair(-1, -1);
}
auto first_interval = path.begin()->second;
auto last_interval = path.rbegin()->second;
return std::make_pair(first_interval.front().time_msec,
last_interval.back().time_msec);
}
void BoxTracker::NewBoxTrackAsync(const TimedBox& initial_pos, int id,
int64 min_msec, int64 max_msec) {
VLOG(1) << "Async track for id: " << id << " from " << min_msec << " to "
<< max_msec;
// Determine start position and track forward and backward.
int chunk_idx = ChunkIdxFromTime(initial_pos.time_msec);
VLOG(1) << "Starting at chunk " << chunk_idx;
AugmentedChunkPtr tracking_chunk(ReadChunk(id, kInitCheckpoint, chunk_idx));
if (!tracking_chunk.first) {
absl::MutexLock lock(&status_mutex_);
--track_status_[id][kInitCheckpoint].tracks_ongoing;
LOG(ERROR) << "Could not read tracking chunk from file: " << chunk_idx
<< " for start position: " << initial_pos.ToString();
return;
}
// Grab ownership here, to avoid any memory leaks due to early return.
std::unique_ptr<const TrackingDataChunk> chunk_owned;
if (tracking_chunk.second) {
chunk_owned.reset(tracking_chunk.first);
}
const int start_frame =
ClosestFrameIndex(initial_pos.time_msec, *tracking_chunk.first);
VLOG(1) << "Local start frame: " << start_frame;
// Update starting position to coincide with a frame.
TimedBox start_pos = initial_pos;
start_pos.time_msec =
tracking_chunk.first->item(start_frame).timestamp_usec() / 1000;
VLOG(1) << "Request at " << initial_pos.time_msec << " revised to "
<< start_pos.time_msec;
const int checkpoint = start_pos.time_msec;
// TODO:
// Compute min and max for tracking based on existing check points.
if (!WaitToScheduleId(id)) {
// Could not schedule, id already being canceled.
return;
}
// If another checkpoint is close by, cancel that one.
VLOG(1) << "Removing close checkpoints";
absl::MutexLock lock(&status_mutex_);
RemoveCloseCheckpoints(id, checkpoint);
VLOG(1) << "Cancel existing tracks";
CancelTracking(id, checkpoint);
// Remove checkpoint results (to be replaced with current one).
ClearCheckpoint(id, checkpoint);
MotionBoxState start_state;
MotionBoxStateFromTimedBox(start_pos, &start_state);
VLOG(1) << "Adding initial result";
AddBoxResult(start_pos, id, checkpoint, start_state);
// Perform forward and backward tracking and add to current PathSegment.
// Track forward.
track_status_[id][checkpoint].tracks_ongoing += 2;
VLOG(1) << "Starting tracking workers ... ";
AugmentedChunkPtr forward_chunk = tracking_chunk;
AugmentedChunkPtr backward_chunk = tracking_chunk;
if (tracking_chunk.second) { // We have ownership, need a copy here.
forward_chunk = std::make_pair(new TrackingDataChunk(*chunk_owned), true);
backward_chunk = std::make_pair(chunk_owned.release(), true);
}
auto forward_operation = [this, forward_chunk, start_state, start_frame,
chunk_idx, id, checkpoint, min_msec, max_msec]() {
this->TrackingImpl(TrackingImplArgs(forward_chunk, start_state, start_frame,
chunk_idx, id, checkpoint, true, true,
min_msec, max_msec));
};
tracking_workers_->Schedule(forward_operation);
// Track backward.
auto backward_operation = [this, backward_chunk, start_state, start_frame,
chunk_idx, id, checkpoint, min_msec, max_msec]() {
this->TrackingImpl(TrackingImplArgs(backward_chunk, start_state,
start_frame, chunk_idx, id, checkpoint,
false, true, min_msec, max_msec));
};
tracking_workers_->Schedule(backward_operation);
DoneSchedulingId(id);
// Tell a waiting request that we are done scheduling.
status_condvar_.SignalAll();
VLOG(1) << "Scheduling done for " << id;
}
void BoxTracker::RemoveCloseCheckpoints(int id, int checkpoint) {
if (track_status_[id].empty()) {
return;
}
auto pos = track_status_[id].lower_bound(checkpoint);
// Test current and previous location (if possible).
int num_turns = 1;
if (pos != track_status_[id].begin()) {
--pos;
++num_turns;
}
for (int k = 0; k < num_turns; ++k, ++pos) {
if (pos != track_status_[id].end()) {
const int check_pos = pos->first;
// Ignore marker init checkpoint from track_status_.
if (check_pos > kInitCheckpoint &&
std::abs(check_pos - checkpoint) < kSnapMs) {
CancelTracking(id, check_pos);
ClearCheckpoint(id, check_pos);
}
} else {
break;
}
}
}
bool BoxTracker::WaitToScheduleId(int id) {
VLOG(1) << "Wait to schedule id: " << id;
absl::MutexLock lock(&status_mutex_);
while (new_box_track_[id]) {
// Box tracking is currently ongoing for this id.
if (track_status_[id][kInitCheckpoint].canceled) {
// Canceled, remove myself from ongoing tracks.
--track_status_[id][kInitCheckpoint].tracks_ongoing;
status_condvar_.SignalAll();
return false;
}
// Only one request can be processing in the section till end of the
// function NewBoxTrackAsync at a time.
status_condvar_.Wait(&status_mutex_);
}
// We got canceled already, don't proceed.
if (track_status_[id][kInitCheckpoint].canceled) {
--track_status_[id][kInitCheckpoint].tracks_ongoing;
status_condvar_.SignalAll();
return false;
}
// Signal we are about to schedule new tracking.
new_box_track_[id] = true;
VLOG(1) << "Ready to schedule id: " << id;
return true;
}
void BoxTracker::DoneSchedulingId(int id) {
new_box_track_[id] = false;
--track_status_[id][kInitCheckpoint].tracks_ongoing;
}
void BoxTracker::CancelTracking(int id, int checkpoint) {
// Wait for ongoing requests to terminate.
while (track_status_[id][checkpoint].tracks_ongoing != 0) {
// Cancel all ongoing requests.
track_status_[id][checkpoint].canceled = true;
status_condvar_.Wait(&status_mutex_);
}
track_status_[id][checkpoint].canceled = false;
}
bool BoxTracker::GetTimedPosition(int id, int64 time_msec, TimedBox* result,
std::vector<MotionBoxState>* states) {
CHECK(result);
MotionBoxState* lhs_box_state = nullptr;
MotionBoxState* rhs_box_state = nullptr;
if (states) {
CHECK(options_.record_path_states())
<< "Requesting corresponding tracking states requires option "
<< "record_path_states to be set";
states->resize(1);
lhs_box_state = rhs_box_state = &states->at(0);
}
VLOG(1) << "Obtaining result at " << time_msec;
absl::MutexLock lock(&path_mutex_);
const Path& path = paths_[id];
if (path.size() < 1) {
LOG(ERROR) << "Empty path!";
return false;
}
// Find corresponding checkpoint.
auto check_pos = path.lower_bound(time_msec);
if (check_pos == path.begin()) {
VLOG(1) << "To left";
// We are to the left of the earliest checkpoint.
return TimedBoxAtTime(check_pos->second, time_msec, result, lhs_box_state);
}
if (check_pos == path.end()) {
VLOG(1) << "To right";
--check_pos;
// We are to the right of the lastest checkpoint.
return TimedBoxAtTime(check_pos->second, time_msec, result, rhs_box_state);
}
VLOG(1) << "Blending ...";
// We are inbetween checkpoints, get result for each, then blend.
const PathSegment& rhs = check_pos->second;
const int check_rhs = check_pos->first;
--check_pos;
const PathSegment& lhs = check_pos->second;
const int check_lhs = check_pos->first;
TimedBox lhs_box;
TimedBox rhs_box;
if (states) {
states->resize(2);
lhs_box_state = &states->at(0);
rhs_box_state = &states->at(1);
}
if (!TimedBoxAtTime(lhs, time_msec, &lhs_box, lhs_box_state)) {
return false;
}
if (!TimedBoxAtTime(rhs, time_msec, &rhs_box, rhs_box_state)) {
return false;
}
VLOG(1) << "Blending: " << lhs_box.ToString() << " and "
<< rhs_box.ToString();
const double alpha = (time_msec - check_lhs) * 1.0 / (check_rhs - check_lhs);
*result = TimedBox::Blend(lhs_box, rhs_box, alpha);
return true;
}
bool BoxTracker::IsTrackingOngoingForId(int id) {
absl::MutexLock lock(&status_mutex_);
for (const auto& item : track_status_[id]) {
if (item.second.tracks_ongoing > 0) {
return true;
}
}
return false;
}
bool BoxTracker::IsTrackingOngoing() {
absl::MutexLock lock(&status_mutex_);
return IsTrackingOngoingMutexHeld();
}
bool BoxTracker::IsTrackingOngoingMutexHeld() {
for (const auto& id : track_status_) {
for (const auto& item : id.second) {
if (item.second.tracks_ongoing > 0) {
return true;
}
}
}
return false;
}
BoxTracker::AugmentedChunkPtr BoxTracker::ReadChunk(int id, int checkpoint,
int chunk_idx) {
VLOG(1) << __FUNCTION__ << " id=" << id << " chunk_idx=" << chunk_idx;
if (cache_dir_.empty() && !tracking_data_.empty()) {
if (chunk_idx < tracking_data_.size()) {
return std::make_pair(tracking_data_[chunk_idx], false);
} else {
LOG(ERROR) << "chunk_idx >= tracking_data_.size()";
return std::make_pair(nullptr, false);
}
} else {
std::unique_ptr<TrackingDataChunk> chunk_data(
ReadChunkFromCache(id, checkpoint, chunk_idx));
return std::make_pair(chunk_data.release(), true);
}
}
std::unique_ptr<TrackingDataChunk> BoxTracker::ReadChunkFromCache(
int id, int checkpoint, int chunk_idx) {
VLOG(1) << __FUNCTION__ << " id=" << id << " chunk_idx=" << chunk_idx;
auto format_runtime =
absl::ParsedFormat<'d'>::New(options_.cache_file_format());
std::string chunk_file;
if (format_runtime) {
chunk_file = cache_dir_ + "/" + absl::StrFormat(*format_runtime, chunk_idx);
} else {
LOG(ERROR) << "chache_file_format wrong. fall back to chunk_%04d.";
chunk_file = cache_dir_ + "/" + absl::StrFormat("chunk_%04d", chunk_idx);
}
VLOG(1) << "Reading chunk from cache: " << chunk_file;
std::unique_ptr<TrackingDataChunk> chunk_data(new TrackingDataChunk());
struct stat tmp;
if (stat(chunk_file.c_str(), &tmp)) {
if (!WaitForChunkFile(id, checkpoint, chunk_file)) {
return nullptr;
}
}
VLOG(1) << "File exists, reading ...";
std::ifstream in(chunk_file, std::ios::in | std::ios::binary);
if (!in) {
LOG(ERROR) << "Could not read chunk file: " << chunk_file;
return nullptr;
}
std::string data;
in.seekg(0, std::ios::end);
data.resize(in.tellg());
in.seekg(0, std::ios::beg);
in.read(&data[0], data.size());
in.close();
chunk_data->ParseFromString(data);
VLOG(1) << "Read success";
return chunk_data;
}
bool BoxTracker::WaitForChunkFile(int id, int checkpoint,
const std::string& chunk_file) {
VLOG(1) << "Chunk no exists, waiting for file: " << chunk_file;
const int timeout_msec = options_.read_chunk_timeout_msec();
// Exponential backoff sleep till file exists.
int wait_time_msec = 20;
VLOG(1) << "In wait for chunk ...: " << chunk_file;
// Maximum wait time.
const int kMaxWaitPeriod = 5000;
int total_wait_msec = 0;
bool file_exists = false;
while (!file_exists && total_wait_msec < timeout_msec) {
// Check if we got canceled.
{
absl::MutexLock lock(&status_mutex_);
if (track_status_[id][checkpoint].canceled) {
return false;
}
}
absl::SleepFor(absl::Milliseconds(wait_time_msec));
total_wait_msec += wait_time_msec;
struct stat tmp;
file_exists = stat(chunk_file.c_str(), &tmp) == 0;
if (file_exists) {
VLOG(1) << "Sucessfully waited on " << chunk_file << " for "
<< total_wait_msec;
break;
}
if (wait_time_msec < kMaxWaitPeriod) {
wait_time_msec *= 1.5;
}
}
return file_exists;
}
int BoxTracker::ClosestFrameIndex(int64 msec,
const TrackingDataChunk& chunk) const {
CHECK_GT(chunk.item_size(), 0);
typedef TrackingDataChunk::Item Item;
Item item_to_find;
item_to_find.set_timestamp_usec(msec * 1000);
int pos =
std::lower_bound(chunk.item().begin(), chunk.item().end(), item_to_find,
[](const Item& lhs, const Item& rhs) -> bool {
return lhs.timestamp_usec() < rhs.timestamp_usec();
}) -
chunk.item().begin();
// Skip end.
if (pos == chunk.item_size()) {
return pos - 1;
} else if (pos == 0) {
// Nothing smaller exists.
return 0;
}
// Determine closest timestamp.
const int64 lhs_diff = msec - chunk.item(pos - 1).timestamp_usec() / 1000;
const int64 rhs_diff = chunk.item(pos).timestamp_usec() / 1000 - msec;
if (std::min(lhs_diff, rhs_diff) >= 67) {
LOG(ERROR) << "No frame found within 67ms, probably using wrong chunk.";
}
if (lhs_diff < rhs_diff) {
return pos - 1;
} else {
return pos;
}
}
void BoxTracker::AddBoxResult(const TimedBox& box, int id, int checkpoint,
const MotionBoxState& state) {
absl::MutexLock lock(&path_mutex_);
PathSegment& segment = paths_[id][checkpoint];
auto insert_pos = std::lower_bound(segment.begin(), segment.end(), box);
const bool store_state = options_.record_path_states();
// Don't overwrite an existing box.
if (insert_pos == segment.end() || insert_pos->time_msec != box.time_msec) {
segment.insert(insert_pos,
InternalTimedBox(
box, store_state ? new MotionBoxState(state) : nullptr));
}
}
void BoxTracker::ClearCheckpoint(int id, int checkpoint) {
absl::MutexLock lock(&path_mutex_);
PathSegment& segment = paths_[id][checkpoint];
segment.clear();
}
void BoxTracker::TrackingImpl(const TrackingImplArgs& a) {
TrackStepOptions track_step_options = options_.track_step_options();
ChangeTrackingDegreesBasedOnStartPos(a.start_state, &track_step_options);
MotionBox motion_box(track_step_options);
const int chunk_data_size = a.chunk_data->item_size();
CHECK_GE(a.start_frame, 0);
CHECK_LT(a.start_frame, chunk_data_size);
VLOG(1) << " a.start_frame = " << a.start_frame << " @"
<< a.chunk_data->item(a.start_frame).timestamp_usec() << " with "
<< chunk_data_size << " items";
motion_box.ResetAtFrame(a.start_frame, a.start_state);
auto cleanup_func = [&a, this]() -> void {
if (a.first_call) {
// Signal we are done processing in this direction.
absl::MutexLock lock(&status_mutex_);
--track_status_[a.id][a.checkpoint].tracks_ongoing;
status_condvar_.SignalAll();
}
};
if (a.forward) {
// TrackingData at frame f, contains tracking information from
// frame f to f - 1. Get information at frame f + 1 and invert:
// Tracking from f to f + 1.
for (int f = a.start_frame; f + 1 < chunk_data_size; ++f) {
// Note: we use / 1000 instead of * 1000 to avoid overflow.
if (a.chunk_data->item(f + 1).timestamp_usec() / 1000 > a.max_msec) {
VLOG(2) << "Reached maximum tracking timestamp @" << a.max_msec;
break;
}
VLOG(1) << "Track forward from " << f;
MotionVectorFrame mvf;
MotionVectorFrameFromTrackingData(
a.chunk_data->item(f + 1).tracking_data(), &mvf);
const int track_duration_ms =
TrackingDataDurationMs(a.chunk_data->item(f + 1));
if (track_duration_ms > 0) {
mvf.duration_ms = track_duration_ms;
}
// If this is the first frame in a chunk, there might be an unobserved
// chunk boundary at the first frame.
if (f == 0 && a.chunk_data->item(0).tracking_data().frame_flags() &
TrackingData::FLAG_CHUNK_BOUNDARY) {
mvf.is_chunk_boundary = true;
}
MotionVectorFrame mvf_inverted;
InvertMotionVectorFrame(mvf, &mvf_inverted);
const bool forward_tracking = true;
if (!motion_box.TrackStep(f, mvf_inverted, forward_tracking)) {
VLOG(1) << "Failed forward at frame: " << f;
break;
} else {
// Test if current request is canceled.
{
absl::MutexLock lock(&status_mutex_);
if (track_status_[a.id][a.checkpoint].canceled) {
--track_status_[a.id][a.checkpoint].tracks_ongoing;
status_condvar_.SignalAll();
return;
}
}
TimedBox result;
const MotionBoxState& result_state = motion_box.StateAtFrame(f + 1);
TimedBoxFromMotionBoxState(result_state, &result);
result.time_msec = a.chunk_data->item(f + 1).timestamp_usec() / 1000;
AddBoxResult(result, a.id, a.checkpoint, result_state);
}
if (f + 2 == chunk_data_size && !a.chunk_data->last_chunk()) {
// Last frame, successful track, continue;
AugmentedChunkPtr next_chunk(
ReadChunk(a.id, a.checkpoint, a.chunk_idx + 1));
if (next_chunk.first != nullptr) {
TrackingImplArgs next_args(next_chunk, motion_box.StateAtFrame(f + 1),
0, a.chunk_idx + 1, a.id, a.checkpoint,
a.forward, false, a.min_msec, a.max_msec);
TrackingImpl(next_args);
} else {
cleanup_func();
LOG(ERROR) << "Can't read expected chunk file!";
}
}
}
} else {
// Backward tracking.
// Don't attempt to track from the very first frame backwards.
const int first_frame = a.chunk_data->first_chunk() ? 1 : 0;
for (int f = a.start_frame; f >= first_frame; --f) {
if (a.chunk_data->item(f).timestamp_usec() / 1000 < a.min_msec) {
VLOG(2) << "Reached minimum tracking timestamp @" << a.min_msec;
break;
}
VLOG(1) << "Track backward from " << f;
MotionVectorFrame mvf;
MotionVectorFrameFromTrackingData(a.chunk_data->item(f).tracking_data(),
&mvf);
const int64 track_duration_ms =
TrackingDataDurationMs(a.chunk_data->item(f));
if (track_duration_ms > 0) {
mvf.duration_ms = track_duration_ms;
}
const bool forward_tracking = false;
if (!motion_box.TrackStep(f, mvf, forward_tracking)) {
VLOG(1) << "Failed backward at frame: " << f;
break;
} else {
// Test if current request is canceled.
{
absl::MutexLock lock(&status_mutex_);
if (track_status_[a.id][a.checkpoint].canceled) {
--track_status_[a.id][a.checkpoint].tracks_ongoing;
status_condvar_.SignalAll();
return;
}
}
TimedBox result;
const MotionBoxState& result_state = motion_box.StateAtFrame(f - 1);
TimedBoxFromMotionBoxState(result_state, &result);
result.time_msec = a.chunk_data->item(f).prev_timestamp_usec() / 1000;
AddBoxResult(result, a.id, a.checkpoint, result_state);
}
if (f == first_frame && !a.chunk_data->first_chunk()) {
VLOG(1) << "Read next chunk: " << f << "==" << first_frame << " in "
<< a.chunk_idx;
// First frame, successful track, continue.
AugmentedChunkPtr prev_chunk(
ReadChunk(a.id, a.checkpoint, a.chunk_idx - 1));
if (prev_chunk.first != nullptr) {
const int last_frame = prev_chunk.first->item_size() - 1;
TrackingImplArgs prev_args(prev_chunk, motion_box.StateAtFrame(f - 1),
last_frame, a.chunk_idx - 1, a.id,
a.checkpoint, a.forward, false, a.min_msec,
a.max_msec);
TrackingImpl(prev_args);
} else {
cleanup_func();
LOG(ERROR) << "Can't read expected chunk file! " << a.chunk_idx - 1
<< " while tracking @"
<< a.chunk_data->item(f).timestamp_usec() / 1000
<< " with cutoff " << a.min_msec;
return;
}
}
}
}
cleanup_func();
}
bool TimedBoxAtTime(const PathSegment& segment, int64 time_msec, TimedBox* box,
MotionBoxState* state) {
CHECK(box);
if (segment.empty()) {
return false;
}
TimedBox to_find;
to_find.time_msec = time_msec;
auto pos = std::lower_bound(segment.begin(), segment.end(), to_find);
if (pos != segment.end() && pos->time_msec == time_msec) {
*box = *pos;
if (state) {
*state = *pos->state;
}
return true;
}
constexpr int kMaxDiff = 67;
if (pos == segment.begin()) {
if (pos->time_msec - time_msec < kMaxDiff) {
*box = *pos;
if (state && pos->state) {
*state = *pos->state;
}
return true;
} else {
return false;
}
}
if (pos == segment.end()) {
if (time_msec - pos[-1].time_msec < kMaxDiff) {
*box = pos[-1];
if (state && pos[-1].state) {
*state = *pos[-1].state;
}
return true;
} else {
return false;
}
}
// Interpolation necessary.
*box = BlendTimedBoxes(pos[-1], pos[0], time_msec);
if (state) {
// Grab closest state.
if (std::abs(pos[-1].time_msec - time_msec) <
std::abs(pos[0].time_msec - time_msec)) {
if (pos[-1].state) {
*state = *pos[-1].state;
}
} else {
if (pos[0].state) {
*state = *pos[0].state;
}
}
}
return true;
}
void BoxTracker::ResumeTracking() {
absl::MutexLock lock(&status_mutex_);
canceling_ = false;
}
void BoxTracker::CancelAllOngoingTracks() {
// Get a list of items to be canceled (id, checkpoint)
absl::MutexLock lock(&status_mutex_);
canceling_ = true;
std::vector<std::pair<int, int>> to_be_canceled;
for (auto& id : track_status_) {
for (auto& checkpoint : id.second) {
if (checkpoint.second.tracks_ongoing > 0) {
checkpoint.second.canceled = true;
to_be_canceled.push_back(std::make_pair(id.first, checkpoint.first));
}
}
}
// Wait for ongoing requests to terminate.
auto on_going_test = [&to_be_canceled, this]() -> bool {
status_mutex_.AssertHeld();
for (const auto& item : to_be_canceled) {
if (track_status_[item.first][item.second].tracks_ongoing > 0) {
return true;
}
}
return false;
};
while (on_going_test()) {
status_condvar_.Wait(&status_mutex_);
}
// Indicate we are done canceling.
for (const auto& item : to_be_canceled) {
track_status_[item.first][item.second].canceled = false;
}
}
bool BoxTracker::WaitForAllOngoingTracks(int timeout_us) {
MEASURE_TIME << "Tracking time ...";
absl::MutexLock lock(&status_mutex_);
// Infinite wait for timeout <= 0.
absl::Duration timeout = timeout_us > 0 ? absl::Microseconds(timeout_us)
: absl::InfiniteDuration();
while (timeout > absl::ZeroDuration() && IsTrackingOngoingMutexHeld()) {
absl::Time start_wait = absl::Now();
status_condvar_.WaitWithTimeout(&status_mutex_, timeout);
absl::Duration elapsed = absl::Now() - start_wait;
timeout -= elapsed;
}
return !IsTrackingOngoingMutexHeld();
}
bool BoxTracker::GetTrackingData(int id, int64 request_time_msec,
TrackingData* tracking_data,
int* tracking_data_msec) {
CHECK(tracking_data);
int chunk_idx = ChunkIdxFromTime(request_time_msec);
AugmentedChunkPtr tracking_chunk(ReadChunk(id, kInitCheckpoint, chunk_idx));
if (!tracking_chunk.first) {
absl::MutexLock lock(&status_mutex_);
--track_status_[id][kInitCheckpoint].tracks_ongoing;
LOG(ERROR) << "Could not read tracking chunk from file.";
return false;
}
std::unique_ptr<const TrackingDataChunk> owned_chunk;
if (tracking_chunk.second) {
owned_chunk.reset(tracking_chunk.first);
}
const int closest_frame =
ClosestFrameIndex(request_time_msec, *tracking_chunk.first);
*tracking_data = tracking_chunk.first->item(closest_frame).tracking_data();
if (tracking_data_msec) {
*tracking_data_msec =
tracking_chunk.first->item(closest_frame).timestamp_usec() / 1000;
}
return true;
}
} // namespace mediapipe
Reference in New Issue View Git Blame Copy Permalink