150 lines
4.2 KiB
C++
150 lines
4.2 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/streaming_buffer.h"
|
|
|
|
#include "absl/strings/str_cat.h"
|
|
|
|
namespace mediapipe {
|
|
|
|
StreamingBuffer::StreamingBuffer(
|
|
const std::vector<TaggedType>& data_configuration, int overlap)
|
|
: overlap_(overlap) {
|
|
CHECK_GE(overlap, 0);
|
|
for (auto& item : data_configuration) {
|
|
CHECK(data_config_.find(item.first) == data_config_.end())
|
|
<< "Tag " << item.first << " already exists";
|
|
data_config_[item.first] = item.second;
|
|
// Init deque.
|
|
data_[item.first].clear();
|
|
}
|
|
}
|
|
|
|
bool StreamingBuffer::HasTag(const std::string& tag) const {
|
|
return data_config_.find(tag) != data_config_.end();
|
|
}
|
|
|
|
bool StreamingBuffer::HasTags(const std::vector<std::string>& tags) const {
|
|
for (const auto& tag : tags) {
|
|
if (!HasTag(tag)) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
int StreamingBuffer::BufferSize(const std::string& tag) const {
|
|
CHECK(HasTag(tag));
|
|
return data_.find(tag)->second.size();
|
|
}
|
|
|
|
int StreamingBuffer::MaxBufferSize() const {
|
|
int max_buffer = 0;
|
|
for (const auto& elem : data_) {
|
|
max_buffer = std::max(max_buffer, BufferSize(elem.first));
|
|
}
|
|
return max_buffer;
|
|
}
|
|
|
|
bool StreamingBuffer::HaveEqualSize(
|
|
const std::vector<std::string>& tags) const {
|
|
if (tags.size() < 2) {
|
|
return true;
|
|
}
|
|
int first_size = BufferSize(tags[0]);
|
|
for (int k = 1; k < tags.size(); ++k) {
|
|
if (BufferSize(tags[1]) != first_size) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
std::vector<std::string> StreamingBuffer::AllTags() const {
|
|
std::vector<std::string> all_tags;
|
|
for (auto& item : data_config_) {
|
|
all_tags.push_back(item.first);
|
|
}
|
|
return all_tags;
|
|
}
|
|
|
|
bool StreamingBuffer::TruncateBuffer(bool flush) {
|
|
// Only truncate if sufficient elements have been buffered.
|
|
const int elems_to_clear =
|
|
std::max(0, MaxBufferSize() - (flush ? 0 : overlap_));
|
|
|
|
if (elems_to_clear == 0) {
|
|
return true;
|
|
}
|
|
|
|
bool is_consistent = true;
|
|
for (auto& item : data_) {
|
|
auto& buffer = item.second;
|
|
const int buffer_elems_to_clear =
|
|
std::min<int>(elems_to_clear, buffer.size());
|
|
if (buffer_elems_to_clear < elems_to_clear) {
|
|
LOG(WARNING) << "For tag " << item.first << " got "
|
|
<< elems_to_clear - buffer_elems_to_clear
|
|
<< "fewer elements than buffer can hold.";
|
|
is_consistent = false;
|
|
}
|
|
buffer.erase(buffer.begin(), buffer.begin() + buffer_elems_to_clear);
|
|
}
|
|
|
|
first_frame_index_ += elems_to_clear;
|
|
|
|
const int remaining_elems = flush ? 0 : overlap_;
|
|
for (const auto& item : data_) {
|
|
const auto& buffer = item.second;
|
|
if (buffer.size() != remaining_elems) {
|
|
LOG(WARNING) << "After trunctation, for tag " << item.first << "got "
|
|
<< buffer.size() << " elements, "
|
|
<< "expected " << remaining_elems;
|
|
is_consistent = false;
|
|
}
|
|
}
|
|
|
|
return is_consistent;
|
|
}
|
|
|
|
void StreamingBuffer::DiscardDatum(const std::string& tag, int num_frames) {
|
|
CHECK(HasTag(tag));
|
|
auto& queue = data_[tag];
|
|
if (queue.empty()) {
|
|
return;
|
|
}
|
|
queue.erase(queue.begin(),
|
|
queue.begin() + std::min<int>(queue.size(), num_frames));
|
|
}
|
|
|
|
void StreamingBuffer::DiscardDatumFromEnd(const std::string& tag,
|
|
int num_frames) {
|
|
CHECK(HasTag(tag));
|
|
auto& queue = data_[tag];
|
|
if (queue.empty()) {
|
|
return;
|
|
}
|
|
queue.erase(queue.end() - std::min<int>(queue.size(), num_frames),
|
|
queue.end());
|
|
}
|
|
|
|
void StreamingBuffer::DiscardData(const std::vector<std::string>& tags,
|
|
int num_frames) {
|
|
for (const std::string& tag : tags) {
|
|
DiscardDatum(tag, num_frames);
|
|
}
|
|
}
|
|
|
|
} // namespace mediapipe
|