talksik/mediapipe
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 612e50bb8db2ec3dc1c30049372d87a80c3848db
This commit is contained in:
MediaPipe Team
2020-08-30 19:52:55 -04:00
committed by chuoling
parent a7225b938a
commit c0124fb83c
248 changed files with 5225 additions and 1914 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/calculators/audio/BUILD
+1 -1
View File
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
licenses(["notice"]) # Apache 2.0
licenses(["notice"])
package(default_visibility = ["//visibility:private"])
mediapipe/calculators/audio/testdata/BUILD
Vendored
+1 -1
View File
@@ -13,7 +13,7 @@
# limitations under the License.
#
licenses(["notice"]) # Apache 2.0
licenses(["notice"])
filegroup(
name = "test_audios",
mediapipe/calculators/core/BUILD
+4 -1
View File
@@ -15,7 +15,7 @@
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
licenses(["notice"]) # Apache 2.0
licenses(["notice"])
package(default_visibility = ["//visibility:private"])
@@ -290,7 +290,9 @@ cc_library(
deps = [
":concatenate_vector_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/formats:classification_cc_proto",
"//mediapipe/framework/formats:landmark_cc_proto",
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"@org_tensorflow//tensorflow/lite:framework",
@@ -1119,6 +1121,7 @@ cc_library(
":constant_side_packet_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:collection_item_id",
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
],
mediapipe/calculators/core/concatenate_vector_calculator.cc
+13
View File
@@ -16,7 +16,9 @@
#include <vector>
#include "mediapipe/framework/formats/classification.pb.h"
#include "mediapipe/framework/formats/landmark.pb.h"
#include "mediapipe/framework/port/integral_types.h"
#include "tensorflow/lite/interpreter.h"
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
@@ -45,6 +47,9 @@ REGISTER_CALCULATOR(ConcatenateFloatVectorCalculator);
typedef ConcatenateVectorCalculator<int32> ConcatenateInt32VectorCalculator;
REGISTER_CALCULATOR(ConcatenateInt32VectorCalculator);
typedef ConcatenateVectorCalculator<uint64> ConcatenateUInt64VectorCalculator;
REGISTER_CALCULATOR(ConcatenateUInt64VectorCalculator);
// Example config:
// node {
// calculator: "ConcatenateTfLiteTensorVectorCalculator"
@@ -60,6 +65,14 @@ typedef ConcatenateVectorCalculator<::mediapipe::NormalizedLandmark>
ConcatenateLandmarkVectorCalculator;
REGISTER_CALCULATOR(ConcatenateLandmarkVectorCalculator);
typedef ConcatenateVectorCalculator<::mediapipe::NormalizedLandmarkList>
ConcatenateLandmarListVectorCalculator;
REGISTER_CALCULATOR(ConcatenateLandmarListVectorCalculator);
typedef ConcatenateVectorCalculator<mediapipe::ClassificationList>
ConcatenateClassificationListVectorCalculator;
REGISTER_CALCULATOR(ConcatenateClassificationListVectorCalculator);
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
typedef ConcatenateVectorCalculator<::tflite::gpu::gl::GlBuffer>
ConcatenateGlBufferVectorCalculator;
mediapipe/calculators/core/concatenate_vector_calculator.h
+46 -19
View File
@@ -15,6 +15,7 @@
#ifndef MEDIAPIPE_CALCULATORS_CORE_CONCATENATE_VECTOR_CALCULATOR_H_
#define MEDIAPIPE_CALCULATORS_CORE_CONCATENATE_VECTOR_CALCULATOR_H_
#include <string>
#include <type_traits>
#include <vector>
@@ -26,10 +27,10 @@
namespace mediapipe {
// Concatenates several std::vector<T> following stream index order. This class
// assumes that every input stream contains the vector<T> type. To use this
// class for a particular type T, regisiter a calculator using
// ConcatenateVectorCalculator<T>.
// Concatenates several objects of type T or std::vector<T> following stream
// index order. This class assumes that every input stream contains either T or
// vector<T> type. To use this class for a particular type T, regisiter a
// calculator using ConcatenateVectorCalculator<T>.
template <typename T>
class ConcatenateVectorCalculator : public CalculatorBase {
public:
@@ -38,7 +39,8 @@ class ConcatenateVectorCalculator : public CalculatorBase {
RET_CHECK(cc->Outputs().NumEntries() == 1);
for (int i = 0; i < cc->Inputs().NumEntries(); ++i) {
cc->Inputs().Index(i).Set<std::vector<T>>();
// Actual type T or vector<T> will be validated in Process().
cc->Inputs().Index(i).SetAny();
}
cc->Outputs().Index(0).Set<std::vector<T>>();
@@ -69,9 +71,19 @@ class ConcatenateVectorCalculator : public CalculatorBase {
CalculatorContext* cc) {
auto output = absl::make_unique<std::vector<U>>();
for (int i = 0; i < cc->Inputs().NumEntries(); ++i) {
if (cc->Inputs().Index(i).IsEmpty()) continue;
const std::vector<U>& input = cc->Inputs().Index(i).Get<std::vector<U>>();
output->insert(output->end(), input.begin(), input.end());
auto& input = cc->Inputs().Index(i);
if (input.IsEmpty()) continue;
if (input.Value().ValidateAsType<U>().ok()) {
const U& value = input.Get<U>();
output->push_back(value);
} else if (input.Value().ValidateAsType<std::vector<U>>().ok()) {
const std::vector<U>& value = input.Get<std::vector<U>>();
output->insert(output->end(), value.begin(), value.end());
} else {
return ::mediapipe::InvalidArgumentError("Invalid input stream type.");
}
}
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
return ::mediapipe::OkStatus();
@@ -88,17 +100,32 @@ class ConcatenateVectorCalculator : public CalculatorBase {
CalculatorContext* cc) {
auto output = absl::make_unique<std::vector<U>>();
for (int i = 0; i < cc->Inputs().NumEntries(); ++i) {
if (cc->Inputs().Index(i).IsEmpty()) continue;
::mediapipe::StatusOr<std::unique_ptr<std::vector<U>>> input_status =
cc->Inputs().Index(i).Value().Consume<std::vector<U>>();
if (input_status.ok()) {
std::unique_ptr<std::vector<U>> input_vector =
std::move(input_status).ValueOrDie();
output->insert(output->end(),
std::make_move_iterator(input_vector->begin()),
std::make_move_iterator(input_vector->end()));
auto& input = cc->Inputs().Index(i);
if (input.IsEmpty()) continue;
if (input.Value().ValidateAsType<U>().ok()) {
::mediapipe::StatusOr<std::unique_ptr<U>> value_status =
input.Value().Consume<U>();
if (value_status.ok()) {
std::unique_ptr<U> value = std::move(value_status).ValueOrDie();
output->push_back(std::move(*value));
} else {
return value_status.status();
}
} else if (input.Value().ValidateAsType<std::vector<U>>().ok()) {
::mediapipe::StatusOr<std::unique_ptr<std::vector<U>>> value_status =
input.Value().Consume<std::vector<U>>();
if (value_status.ok()) {
std::unique_ptr<std::vector<U>> value =
std::move(value_status).ValueOrDie();
output->insert(output->end(), std::make_move_iterator(value->begin()),
std::make_move_iterator(value->end()));
} else {
return value_status.status();
}
} else {
return input_status.status();
return ::mediapipe::InvalidArgumentError("Invalid input stream type.");
}
}
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
@@ -109,7 +136,7 @@ class ConcatenateVectorCalculator : public CalculatorBase {
::mediapipe::Status ConsumeAndConcatenateVectors(std::false_type,
CalculatorContext* cc) {
return ::mediapipe::InternalError(
"Cannot copy or move input vectors to concatenate them");
"Cannot copy or move inputs to concatenate them");
}
private:
mediapipe/calculators/core/concatenate_vector_calculator_test.cc
+149 -2
View File
@@ -30,11 +30,29 @@ namespace mediapipe {
typedef ConcatenateVectorCalculator<int> TestConcatenateIntVectorCalculator;
REGISTER_CALCULATOR(TestConcatenateIntVectorCalculator);
void AddInputVector(int index, const std::vector<int>& input, int64 timestamp,
CalculatorRunner* runner) {
runner->MutableInputs()->Index(index).packets.push_back(
MakePacket<std::vector<int>>(input).At(Timestamp(timestamp)));
}
void AddInputVectors(const std::vector<std::vector<int>>& inputs,
int64 timestamp, CalculatorRunner* runner) {
for (int i = 0; i < inputs.size(); ++i) {
runner->MutableInputs()->Index(i).packets.push_back(
MakePacket<std::vector<int>>(inputs[i]).At(Timestamp(timestamp)));
AddInputVector(i, inputs[i], timestamp, runner);
}
}
void AddInputItem(int index, int input, int64 timestamp,
CalculatorRunner* runner) {
runner->MutableInputs()->Index(index).packets.push_back(
MakePacket<int>(input).At(Timestamp(timestamp)));
}
void AddInputItems(const std::vector<int>& inputs, int64 timestamp,
CalculatorRunner* runner) {
for (int i = 0; i < inputs.size(); ++i) {
AddInputItem(i, inputs[i], timestamp, runner);
}
}
@@ -131,6 +149,135 @@ TEST(TestConcatenateIntVectorCalculatorTest, OneEmptyStreamNoOutput) {
EXPECT_EQ(0, outputs.size());
}
TEST(TestConcatenateIntVectorCalculatorTest, ItemsOneTimestamp) {
CalculatorRunner runner("TestConcatenateIntVectorCalculator",
/*options_string=*/"", /*num_inputs=*/3,
/*num_outputs=*/1, /*num_side_packets=*/0);
std::vector<int> inputs = {1, 2, 3};
AddInputItems(inputs, /*timestamp=*/1, &runner);
MP_ASSERT_OK(runner.Run());
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
EXPECT_EQ(1, outputs.size());
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
std::vector<int> expected_vector = {1, 2, 3};
EXPECT_EQ(expected_vector, outputs[0].Get<std::vector<int>>());
}
TEST(TestConcatenateIntVectorCalculatorTest, ItemsTwoInputsAtTwoTimestamps) {
CalculatorRunner runner("TestConcatenateIntVectorCalculator",
/*options_string=*/"", /*num_inputs=*/3,
/*num_outputs=*/1, /*num_side_packets=*/0);
{
std::vector<int> inputs = {1, 2, 3};
AddInputItems(inputs, /*timestamp=*/1, &runner);
}
{
std::vector<int> inputs = {4, 5, 6};
AddInputItems(inputs, /*timestamp=*/2, &runner);
}
MP_ASSERT_OK(runner.Run());
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
EXPECT_EQ(2, outputs.size());
{
EXPECT_EQ(3, outputs[0].Get<std::vector<int>>().size());
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
std::vector<int> expected_vector = {1, 2, 3};
EXPECT_EQ(expected_vector, outputs[0].Get<std::vector<int>>());
}
{
EXPECT_EQ(3, outputs[1].Get<std::vector<int>>().size());
EXPECT_EQ(Timestamp(2), outputs[1].Timestamp());
std::vector<int> expected_vector = {4, 5, 6};
EXPECT_EQ(expected_vector, outputs[1].Get<std::vector<int>>());
}
}
TEST(TestConcatenateIntVectorCalculatorTest, ItemsOneEmptyStreamStillOutput) {
CalculatorRunner runner("TestConcatenateIntVectorCalculator",
/*options_string=*/"", /*num_inputs=*/3,
/*num_outputs=*/1, /*num_side_packets=*/0);
// No third input item.
std::vector<int> inputs = {1, 2};
AddInputItems(inputs, /*timestamp=*/1, &runner);
MP_ASSERT_OK(runner.Run());
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
EXPECT_EQ(1, outputs.size());
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
std::vector<int> expected_vector = {1, 2};
EXPECT_EQ(expected_vector, outputs[0].Get<std::vector<int>>());
}
TEST(TestConcatenateIntVectorCalculatorTest, ItemsOneEmptyStreamNoOutput) {
CalculatorRunner runner("TestConcatenateIntVectorCalculator",
/*options_string=*/
"[mediapipe.ConcatenateVectorCalculatorOptions.ext]: "
"{only_emit_if_all_present: true}",
/*num_inputs=*/3,
/*num_outputs=*/1, /*num_side_packets=*/0);
// No third input item.
std::vector<int> inputs = {1, 2};
AddInputItems(inputs, /*timestamp=*/1, &runner);
MP_ASSERT_OK(runner.Run());
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
EXPECT_EQ(0, outputs.size());
}
TEST(TestConcatenateIntVectorCalculatorTest, MixedVectorsAndItems) {
CalculatorRunner runner("TestConcatenateIntVectorCalculator",
/*options_string=*/"", /*num_inputs=*/4,
/*num_outputs=*/1, /*num_side_packets=*/0);
std::vector<int> vector_0 = {1, 2};
std::vector<int> vector_1 = {3, 4, 5};
int item_0 = 6;
int item_1 = 7;
AddInputVector(/*index*/ 0, vector_0, /*timestamp=*/1, &runner);
AddInputVector(/*index*/ 1, vector_1, /*timestamp=*/1, &runner);
AddInputItem(/*index*/ 2, item_0, /*timestamp=*/1, &runner);
AddInputItem(/*index*/ 3, item_1, /*timestamp=*/1, &runner);
MP_ASSERT_OK(runner.Run());
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
EXPECT_EQ(1, outputs.size());
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
std::vector<int> expected_vector = {1, 2, 3, 4, 5, 6, 7};
EXPECT_EQ(expected_vector, outputs[0].Get<std::vector<int>>());
}
TEST(TestConcatenateIntVectorCalculatorTest, MixedVectorsAndItemsAnother) {
CalculatorRunner runner("TestConcatenateIntVectorCalculator",
/*options_string=*/"", /*num_inputs=*/4,
/*num_outputs=*/1, /*num_side_packets=*/0);
int item_0 = 1;
std::vector<int> vector_0 = {2, 3};
std::vector<int> vector_1 = {4, 5, 6};
int item_1 = 7;
AddInputItem(/*index*/ 0, item_0, /*timestamp=*/1, &runner);
AddInputVector(/*index*/ 1, vector_0, /*timestamp=*/1, &runner);
AddInputVector(/*index*/ 2, vector_1, /*timestamp=*/1, &runner);
AddInputItem(/*index*/ 3, item_1, /*timestamp=*/1, &runner);
MP_ASSERT_OK(runner.Run());
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
EXPECT_EQ(1, outputs.size());
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
std::vector<int> expected_vector = {1, 2, 3, 4, 5, 6, 7};
EXPECT_EQ(expected_vector, outputs[0].Get<std::vector<int>>());
}
void AddInputVectors(const std::vector<std::vector<float>>& inputs,
int64 timestamp, CalculatorRunner* runner) {
for (int i = 0; i < inputs.size(); ++i) {
mediapipe/calculators/core/constant_side_packet_calculator.cc
+5
View File
@@ -18,6 +18,7 @@
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/collection_item_id.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
@@ -71,6 +72,8 @@ class ConstantSidePacketCalculator : public CalculatorBase {
packet.Set<bool>();
} else if (packet_options.has_string_value()) {
packet.Set<std::string>();
} else if (packet_options.has_uint64_value()) {
packet.Set<uint64>();
} else {
return ::mediapipe::InvalidArgumentError(
"None of supported values were specified in options.");
@@ -95,6 +98,8 @@ class ConstantSidePacketCalculator : public CalculatorBase {
packet.Set(MakePacket<bool>(packet_options.bool_value()));
} else if (packet_options.has_string_value()) {
packet.Set(MakePacket<std::string>(packet_options.string_value()));
} else if (packet_options.has_uint64_value()) {
packet.Set(MakePacket<uint64>(packet_options.uint64_value()));
} else {
return ::mediapipe::InvalidArgumentError(
"None of supported values were specified in options.");
mediapipe/calculators/core/constant_side_packet_calculator.proto
+1
View File
@@ -29,6 +29,7 @@ message ConstantSidePacketCalculatorOptions {
float float_value = 2;
bool bool_value = 3;
string string_value = 4;
uint64 uint64_value = 5;
}
}
mediapipe/calculators/image/BUILD
+1 -1
View File
@@ -14,7 +14,7 @@
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
licenses(["notice"]) # Apache 2.0
licenses(["notice"])
package(default_visibility = ["//visibility:private"])
mediapipe/calculators/image/testdata/BUILD
Vendored
+1 -1
View File
@@ -13,7 +13,7 @@
# limitations under the License.
#
licenses(["notice"]) # Apache 2.0
licenses(["notice"])
filegroup(
name = "test_images",
mediapipe/calculators/internal/BUILD
+1 -1
View File
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
licenses(["notice"]) # Apache 2.0
licenses(["notice"])
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
mediapipe/calculators/tensorflow/BUILD
+7 -1
View File
@@ -15,7 +15,7 @@
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
licenses(["notice"]) # Apache 2.0
licenses(["notice"])
package(default_visibility = ["//visibility:private"])
@@ -427,6 +427,10 @@ cc_library(
deps = [
":tensorflow_session",
":tensorflow_inference_calculator_cc_proto",
"//mediapipe/framework:timestamp",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/memory",
"//mediapipe/framework:calculator_context",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/tool:status_util",
"@com_google_absl//absl/strings",
@@ -434,6 +438,8 @@ cc_library(
"//mediapipe/framework/deps:clock",
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:map_util",
"//mediapipe/framework:packet",
] + select({
"//conditions:default": [
"@org_tensorflow//tensorflow/core:framework",
mediapipe/calculators/tensorflow/lapped_tensor_buffer_calculator.cc
+1 -1
View File
@@ -93,7 +93,7 @@ REGISTER_CALCULATOR(LappedTensorBufferCalculator);
cc->Inputs().Index(0).Set<tf::Tensor>(
// tensorflow::Tensor stream.
);
RET_CHECK_EQ(cc->Inputs().NumEntries(), 1)
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1)
<< "Only one output stream is supported.";
if (cc->InputSidePackets().HasTag(kBufferSize)) {
mediapipe/calculators/tensorflow/tensorflow_inference_calculator.cc
+158 -57
View File
@@ -19,16 +19,22 @@
#include <unordered_set>
#include <vector>
#include "absl/base/thread_annotations.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_split.h"
#include "absl/synchronization/mutex.h"
#include "mediapipe/calculators/tensorflow/tensorflow_inference_calculator.pb.h"
#include "mediapipe/calculators/tensorflow/tensorflow_session.h"
#include "mediapipe/framework/calculator_context.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/deps/clock.h"
#include "mediapipe/framework/deps/monotonic_clock.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/map_util.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/status_macros.h"
#include "mediapipe/framework/timestamp.h"
#include "mediapipe/framework/tool/status_util.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
@@ -77,6 +83,17 @@ class SimpleSemaphore {
absl::Mutex mutex_;
absl::CondVar cond_;
};
class InferenceState {
public:
InferenceState() : input_tensor_batches_(), batch_timestamps_() {}
// A mapping between stream tags and the tensors we are collecting as a
// batch.
std::map<std::string, std::vector<tf::Tensor>> input_tensor_batches_;
// The timestamps that go into a batch.
std::vector<Timestamp> batch_timestamps_;
};
} // namespace
// This calculator performs inference on a trained TensorFlow model.
@@ -218,11 +235,16 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
}
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
const auto& options = cc->Options<TensorFlowInferenceCalculatorOptions>();
RET_CHECK(!cc->Inputs().GetTags().empty());
for (const std::string& tag : cc->Inputs().GetTags()) {
// The tensorflow::Tensor with the tag equal to the graph node. May
// have a TimeSeriesHeader if all present TimeSeriesHeaders match.
cc->Inputs().Tag(tag).Set<tf::Tensor>();
if (!options.batched_input()) {
cc->Inputs().Tag(tag).Set<tf::Tensor>();
} else {
cc->Inputs().Tag(tag).Set<std::vector<mediapipe::Packet>>();
}
}
RET_CHECK(!cc->Outputs().GetTags().empty());
for (const std::string& tag : cc->Outputs().GetTags()) {
@@ -242,6 +264,22 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
return ::mediapipe::OkStatus();
}
std::unique_ptr<InferenceState> CreateInferenceState(CalculatorContext* cc)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
std::unique_ptr<InferenceState> inference_state =
absl::make_unique<InferenceState>();
if (cc->InputSidePackets().HasTag("RECURRENT_INIT_TENSORS") &&
!cc->InputSidePackets().Tag("RECURRENT_INIT_TENSORS").IsEmpty()) {
std::map<std::string, tf::Tensor>* init_tensor_map;
init_tensor_map = GetFromUniquePtr<std::map<std::string, tf::Tensor>>(
cc->InputSidePackets().Tag("RECURRENT_INIT_TENSORS"));
for (const auto& p : *init_tensor_map) {
inference_state->input_tensor_batches_[p.first].emplace_back(p.second);
}
}
return inference_state;
}
::mediapipe::Status Open(CalculatorContext* cc) override {
options_ = cc->Options<TensorFlowInferenceCalculatorOptions>();
@@ -275,15 +313,6 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
recurrent_feed_tags_.insert(tags[0]);
recurrent_fetch_tags_to_feed_tags_[tags[1]] = tags[0];
}
if (cc->InputSidePackets().HasTag("RECURRENT_INIT_TENSORS") &&
!cc->InputSidePackets().Tag("RECURRENT_INIT_TENSORS").IsEmpty()) {
std::map<std::string, tf::Tensor>* init_tensor_map;
init_tensor_map = GetFromUniquePtr<std::map<std::string, tf::Tensor>>(
cc->InputSidePackets().Tag("RECURRENT_INIT_TENSORS"));
for (const auto& p : *init_tensor_map) {
input_tensor_batches_[p.first].emplace_back(p.second);
}
}
// Check that all tags are present in this signature bound to tensors.
for (const std::string& tag : cc->Inputs().GetTags()) {
@@ -297,9 +326,15 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
<< options_.signature_name();
}
if (options_.batch_size() == 1) {
{
absl::WriterMutexLock l(&mutex_);
inference_state_ = std::unique_ptr<InferenceState>();
}
if (options_.batch_size() == 1 || options_.batched_input()) {
cc->SetOffset(0);
}
return ::mediapipe::OkStatus();
}
@@ -316,6 +351,24 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
return ::mediapipe::OkStatus();
}
::mediapipe::Status AggregateTensorPacket(
const std::string& tag_name, const Packet& packet,
std::map<Timestamp, std::map<std::string, tf::Tensor>>*
input_tensors_by_tag_by_timestamp,
InferenceState* inference_state) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_) {
tf::Tensor input_tensor(packet.Get<tf::Tensor>());
RET_CHECK_OK(AddBatchDimension(&input_tensor));
if (::mediapipe::ContainsKey(recurrent_feed_tags_, tag_name)) {
// If we receive an input on a recurrent tag, override the state.
// It's OK to override the global state because there is just one
// input stream allowed for recurrent tensors.
inference_state_->input_tensor_batches_[tag_name].clear();
}
(*input_tensors_by_tag_by_timestamp)[packet.Timestamp()].insert(
std::make_pair(tag_name, input_tensor));
return ::mediapipe::OkStatus();
}
// Removes the batch dimension of the output tensor if specified in the
// calculator options.
::mediapipe::Status RemoveBatchDimension(tf::Tensor* output_tensor) {
@@ -331,48 +384,85 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
}
::mediapipe::Status Process(CalculatorContext* cc) override {
std::map<std::string, tf::Tensor> input_tensors_by_tag;
for (const std::string& tag_as_node_name : cc->Inputs().GetTags()) {
if (cc->Inputs().Tag(tag_as_node_name).IsEmpty()) {
// Recurrent tensors can be empty.
if (!::mediapipe::ContainsKey(recurrent_feed_tags_, tag_as_node_name)) {
if (options_.skip_on_missing_features()) {
return ::mediapipe::OkStatus();
} else {
return ::mediapipe::InvalidArgumentError(absl::StrCat(
"Tag ", tag_as_node_name,
" not present at timestamp: ", cc->InputTimestamp().Value()));
std::unique_ptr<InferenceState> inference_state_to_process;
{
absl::WriterMutexLock l(&mutex_);
if (inference_state_ == nullptr) {
inference_state_ = CreateInferenceState(cc);
}
std::map<Timestamp, std::map<std::string, tf::Tensor>>
input_tensors_by_tag_by_timestamp;
for (const std::string& tag_as_node_name : cc->Inputs().GetTags()) {
if (cc->Inputs().Tag(tag_as_node_name).IsEmpty()) {
// Recurrent tensors can be empty.
if (!::mediapipe::ContainsKey(recurrent_feed_tags_,
tag_as_node_name)) {
if (options_.skip_on_missing_features()) {
return ::mediapipe::OkStatus();
} else {
return ::mediapipe::InvalidArgumentError(absl::StrCat(
"Tag ", tag_as_node_name,
" not present at timestamp: ", cc->InputTimestamp().Value()));
}
}
} else if (options_.batched_input()) {
const auto& tensor_packets =
cc->Inputs().Tag(tag_as_node_name).Get<std::vector<Packet>>();
if (tensor_packets.size() > options_.batch_size()) {
return ::mediapipe::InvalidArgumentError(absl::StrCat(
"Batch for tag ", tag_as_node_name,
" has more packets than batch capacity. batch_size: ",
options_.batch_size(), " packets: ", tensor_packets.size()));
}
for (const auto& packet : tensor_packets) {
RET_CHECK_OK(AggregateTensorPacket(
tag_as_node_name, packet, &input_tensors_by_tag_by_timestamp,
inference_state_.get()));
}
} else {
RET_CHECK_OK(AggregateTensorPacket(
tag_as_node_name, cc->Inputs().Tag(tag_as_node_name).Value(),
&input_tensors_by_tag_by_timestamp, inference_state_.get()));
}
} else {
tf::Tensor input_tensor(
cc->Inputs().Tag(tag_as_node_name).Get<tf::Tensor>());
RET_CHECK_OK(AddBatchDimension(&input_tensor));
if (::mediapipe::ContainsKey(recurrent_feed_tags_, tag_as_node_name)) {
// If we receive an input on a recurrent tag, override the state.
// It's OK to override the global state because there is just one
// input stream allowed for recurrent tensors.
input_tensor_batches_[tag_as_node_name].clear();
}
for (const auto& timestamp_and_input_tensors_by_tag :
input_tensors_by_tag_by_timestamp) {
inference_state_->batch_timestamps_.emplace_back(
timestamp_and_input_tensors_by_tag.first);
for (const auto& input_tensor_and_tag :
timestamp_and_input_tensors_by_tag.second) {
inference_state_->input_tensor_batches_[input_tensor_and_tag.first]
.emplace_back(input_tensor_and_tag.second);
}
input_tensors_by_tag.insert(
std::make_pair(tag_as_node_name, input_tensor));
}
if (inference_state_->batch_timestamps_.size() == options_.batch_size() ||
options_.batched_input()) {
inference_state_to_process = std::move(inference_state_);
inference_state_ = std::unique_ptr<InferenceState>();
}
}
batch_timestamps_.emplace_back(cc->InputTimestamp());
for (const auto& input_tensor_and_tag : input_tensors_by_tag) {
input_tensor_batches_[input_tensor_and_tag.first].emplace_back(
input_tensor_and_tag.second);
if (inference_state_to_process) {
MP_RETURN_IF_ERROR(
OutputBatch(cc, std::move(inference_state_to_process)));
}
if (batch_timestamps_.size() == options_.batch_size()) {
MP_RETURN_IF_ERROR(OutputBatch(cc));
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status Close(CalculatorContext* cc) override {
if (!batch_timestamps_.empty()) {
MP_RETURN_IF_ERROR(OutputBatch(cc));
std::unique_ptr<InferenceState> inference_state_to_process = nullptr;
{
absl::WriterMutexLock l(&mutex_);
if (cc->GraphStatus().ok() && inference_state_ != nullptr &&
!inference_state_->batch_timestamps_.empty()) {
inference_state_to_process = std::move(inference_state_);
inference_state_ = std::unique_ptr<InferenceState>();
}
}
if (inference_state_to_process) {
MP_RETURN_IF_ERROR(
OutputBatch(cc, std::move(inference_state_to_process)));
}
return ::mediapipe::OkStatus();
}
@@ -385,10 +475,12 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
// memory buffer. Therefore, copies are cheap and should not cause the memory
// buffer to fall out of scope. In contrast, concat is only used where
// necessary.
::mediapipe::Status OutputBatch(CalculatorContext* cc) {
::mediapipe::Status OutputBatch(
CalculatorContext* cc, std::unique_ptr<InferenceState> inference_state) {
const int64 start_time = absl::ToUnixMicros(clock_->TimeNow());
std::vector<std::pair<mediapipe::ProtoString, tf::Tensor>> input_tensors;
for (auto& keyed_tensors : input_tensor_batches_) {
for (auto& keyed_tensors : inference_state->input_tensor_batches_) {
if (options_.batch_size() == 1) {
// Short circuit to avoid the cost of deep copying tensors in concat.
if (!keyed_tensors.second.empty()) {
@@ -404,7 +496,8 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
} else {
// Pad by replicating the first tens or, then ignore the values.
keyed_tensors.second.resize(options_.batch_size());
std::fill(keyed_tensors.second.begin() + batch_timestamps_.size(),
std::fill(keyed_tensors.second.begin() +
inference_state->batch_timestamps_.size(),
keyed_tensors.second.end(), keyed_tensors.second[0]);
tf::Tensor concated;
const tf::Status concat_status =
@@ -414,7 +507,7 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
concated);
}
}
input_tensor_batches_.clear();
inference_state->input_tensor_batches_.clear();
std::vector<mediapipe::ProtoString> output_tensor_names;
std::vector<std::string> output_name_in_signature;
for (const std::string& tag : cc->Outputs().GetTags()) {
@@ -466,9 +559,11 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
int pos = std::find(output_name_in_signature.begin(),
output_name_in_signature.end(), tag_pair.first) -
output_name_in_signature.begin();
input_tensor_batches_[tag_pair.second].emplace_back(outputs[pos]);
inference_state->input_tensor_batches_[tag_pair.second].emplace_back(
outputs[pos]);
}
absl::WriterMutexLock l(&mutex_);
// Set that we want to split on each index of the 0th dimension.
std::vector<tf::int64> split_vector(options_.batch_size(), 1);
for (int i = 0; i < output_tensor_names.size(); ++i) {
@@ -478,7 +573,8 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
RET_CHECK_OK(RemoveBatchDimension(&output_tensor));
cc->Outputs()
.Tag(output_name_in_signature[i])
.Add(new tf::Tensor(output_tensor), batch_timestamps_[0]);
.Add(new tf::Tensor(output_tensor),
inference_state->batch_timestamps_[0]);
}
} else {
std::vector<tf::Tensor> split_tensors;
@@ -486,22 +582,30 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
tf::tensor::Split(outputs[i], split_vector, &split_tensors);
CHECK(split_status.ok()) << split_status.ToString();
// Loop over timestamps so that we don't copy the padding.
for (int j = 0; j < batch_timestamps_.size(); ++j) {
for (int j = 0; j < inference_state->batch_timestamps_.size(); ++j) {
tf::Tensor output_tensor(split_tensors[j]);
RET_CHECK_OK(RemoveBatchDimension(&output_tensor));
cc->Outputs()
.Tag(output_name_in_signature[i])
.Add(new tf::Tensor(output_tensor), batch_timestamps_[j]);
.Add(new tf::Tensor(output_tensor),
inference_state->batch_timestamps_[j]);
}
}
}
// Get end time and report.
const int64 end_time = absl::ToUnixMicros(clock_->TimeNow());
cc->GetCounter(kTotalUsecsCounterSuffix)
->IncrementBy(end_time - start_time);
cc->GetCounter(kTotalProcessedTimestampsCounterSuffix)
->IncrementBy(batch_timestamps_.size());
batch_timestamps_.clear();
->IncrementBy(inference_state->batch_timestamps_.size());
// Make sure we hold on to the recursive state.
if (!options_.recurrent_tag_pair().empty()) {
inference_state_ = std::move(inference_state);
inference_state_->batch_timestamps_.clear();
}
return ::mediapipe::OkStatus();
}
@@ -514,11 +618,8 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
// A mapping between stream tags and the tensor names they are bound to.
std::map<std::string, std::string> tag_to_tensor_map_;
// A mapping between stream tags and the tensors we are collecting as a batch.
std::map<std::string, std::vector<tf::Tensor>> input_tensor_batches_;
// The timestamps that go into a batch.
std::vector<Timestamp> batch_timestamps_;
absl::Mutex mutex_;
std::unique_ptr<InferenceState> inference_state_ ABSL_GUARDED_BY(mutex_);
// The options for the calculator.
TensorFlowInferenceCalculatorOptions options_;
mediapipe/calculators/tensorflow/tensorflow_inference_calculator.proto
+9
View File
@@ -76,4 +76,13 @@ message TensorFlowInferenceCalculatorOptions {
// only works in the local process, not "globally" across multiple processes
// or replicas (if any). Default to 0, i.e. no limit.
optional int32 max_concurrent_session_runs = 6 [default = 0];
// If turned on, the Calculator expects a vector of batched packages as input.
// This will make sure that you can turn on max_in_flight for batch_size
// greater than 1. Otherwise it results in problems of none-monotonically
// increasing timestamps.
// Use BatchSequentialCalculator to create the batches. The batch_size
// should agree for both calculators. All the data in a batch is processed
// together. The BatchSequentialCalculator can't run with max_in_flight.
optional bool batched_input = 7;
}
mediapipe/calculators/tensorflow/tensorflow_inference_calculator_test.cc
+233 -4
View File
@@ -89,17 +89,31 @@ class TensorflowInferenceCalculatorTest : public ::testing::Test {
output_side_packets.Tag("SESSION");
}
// Create tensor from Vector and add as a Packet to the provided tag as input.
void AddVectorToInputsAsTensor(const std::vector<int32>& input,
const std::string& tag, int64 time) {
Packet CreateTensorPacket(const std::vector<int32>& input, int64 time) {
tf::TensorShape tensor_shape;
tensor_shape.AddDim(input.size());
auto tensor = absl::make_unique<tf::Tensor>(tf::DT_INT32, tensor_shape);
for (int i = 0; i < input.size(); ++i) {
tensor->vec<int32>()(i) = input[i];
}
return Adopt(tensor.release()).At(Timestamp(time));
}
// Create tensor from Vector and add as a Packet to the provided tag as input.
void AddVectorToInputsAsTensor(const std::vector<int32>& input,
const std::string& tag, int64 time) {
runner_->MutableInputs()->Tag(tag).packets.push_back(
Adopt(tensor.release()).At(Timestamp(time)));
CreateTensorPacket(input, time));
}
// Create tensor from Vector and add as a Packet to the provided tag as input.
void AddVectorToInputsAsPacket(const std::vector<Packet>& packets,
const std::string& tag) {
CHECK(!packets.empty())
<< "Please specify at least some data in the packet";
auto packets_ptr = absl::make_unique<std::vector<Packet>>(packets);
runner_->MutableInputs()->Tag(tag).packets.push_back(
Adopt(packets_ptr.release()).At(packets.begin()->Timestamp()));
}
std::unique_ptr<CalculatorRunner> runner_;
@@ -183,6 +197,45 @@ TEST_F(TensorflowInferenceCalculatorTest, GetComputed) {
EXPECT_THAT(run_status.ToString(), testing::HasSubstr("Tag B"));
}
TEST_F(TensorflowInferenceCalculatorTest, GetComputed_MaxInFlight) {
CalculatorGraphConfig::Node config;
config.set_calculator("TensorFlowInferenceCalculator");
config.add_input_stream("A:tensor_a");
config.add_input_stream("B:tensor_b");
config.add_output_stream("MULTIPLIED:tensor_o1");
config.add_input_side_packet("SESSION:session");
config.set_max_in_flight(2);
CalculatorOptions options;
options.MutableExtension(TensorFlowInferenceCalculatorOptions::ext)
->set_batch_size(1);
options.MutableExtension(TensorFlowInferenceCalculatorOptions::ext)
->set_add_batch_dim_to_tensors(false);
*config.mutable_options() = options;
runner_ = absl::make_unique<CalculatorRunner>(config);
AddSessionInputSidePacket();
AddVectorToInputsAsTensor({2, 2, 2}, "A", 0);
AddVectorToInputsAsTensor({3, 4, 5}, "B", 0);
MP_ASSERT_OK(runner_->Run());
const std::vector<Packet>& output_packets_mult =
runner_->Outputs().Tag("MULTIPLIED").packets;
ASSERT_EQ(1, output_packets_mult.size());
const tf::Tensor& tensor_mult = output_packets_mult[0].Get<tf::Tensor>();
tf::TensorShape expected_shape({3});
auto expected_tensor = tf::test::AsTensor<int32>({6, 8, 10}, expected_shape);
tf::test::ExpectTensorEqual<int32>(expected_tensor, tensor_mult);
// Add only one of the two expected tensors at the next timestamp, expect
// useful failure message.
AddVectorToInputsAsTensor({1, 2, 3}, "A", 1);
auto run_status = runner_->Run();
ASSERT_FALSE(run_status.ok());
EXPECT_THAT(run_status.ToString(),
testing::HasSubstr("TensorFlowInferenceCalculator"));
EXPECT_THAT(run_status.ToString(), testing::HasSubstr("Tag B"));
}
TEST_F(TensorflowInferenceCalculatorTest, BadTag) {
CalculatorGraphConfig::Node config;
config.set_calculator("TensorFlowInferenceCalculator");
@@ -235,6 +288,86 @@ TEST_F(TensorflowInferenceCalculatorTest, GetMultiBatchComputed) {
->Get());
}
TEST_F(TensorflowInferenceCalculatorTest, GetMultiBatchComputed_MaxInFlight) {
CalculatorGraphConfig::Node config;
config.set_calculator("TensorFlowInferenceCalculator");
config.add_input_stream("A:tensor_a");
config.add_input_stream("B:tensor_b");
config.add_output_stream("MULTIPLIED:tensor_o1");
config.add_input_side_packet("SESSION:session");
config.set_max_in_flight(2);
CalculatorOptions options;
options.MutableExtension(TensorFlowInferenceCalculatorOptions::ext)
->set_batch_size(1);
*config.mutable_options() = options;
runner_ = absl::make_unique<CalculatorRunner>(config);
AddSessionInputSidePacket();
AddVectorToInputsAsTensor({2, 2, 2}, "A", 0);
AddVectorToInputsAsTensor({3, 4, 5}, "B", 0);
AddVectorToInputsAsTensor({3, 3, 3}, "A", 1);
AddVectorToInputsAsTensor({3, 4, 5}, "B", 1);
MP_ASSERT_OK(runner_->Run());
const std::vector<Packet>& output_packets_mult =
runner_->Outputs().Tag("MULTIPLIED").packets;
ASSERT_EQ(2, output_packets_mult.size());
const tf::Tensor& tensor_mult = output_packets_mult[0].Get<tf::Tensor>();
auto expected_tensor = tf::test::AsTensor<int32>({6, 8, 10});
tf::test::ExpectTensorEqual<int32>(tensor_mult, expected_tensor);
const tf::Tensor& tensor_mult1 = output_packets_mult[1].Get<tf::Tensor>();
auto expected_tensor1 = tf::test::AsTensor<int32>({9, 12, 15});
tf::test::ExpectTensorEqual<int32>(tensor_mult1, expected_tensor1);
EXPECT_EQ(2, runner_
->GetCounter(
"TensorFlowInferenceCalculator-TotalProcessedTimestamps")
->Get());
}
TEST_F(TensorflowInferenceCalculatorTest,
GetMultiBatchComputed_MoreThanMaxInFlight) {
CalculatorGraphConfig::Node config;
config.set_calculator("TensorFlowInferenceCalculator");
config.add_input_stream("A:tensor_a");
config.add_input_stream("B:tensor_b");
config.add_output_stream("MULTIPLIED:tensor_o1");
config.add_input_side_packet("SESSION:session");
config.set_max_in_flight(2);
CalculatorOptions options;
options.MutableExtension(TensorFlowInferenceCalculatorOptions::ext)
->set_batch_size(1);
*config.mutable_options() = options;
runner_ = absl::make_unique<CalculatorRunner>(config);
AddSessionInputSidePacket();
AddVectorToInputsAsTensor({2, 2, 2}, "A", 0);
AddVectorToInputsAsTensor({3, 4, 5}, "B", 0);
AddVectorToInputsAsTensor({3, 3, 3}, "A", 1);
AddVectorToInputsAsTensor({3, 4, 5}, "B", 1);
AddVectorToInputsAsTensor({4, 4, 4}, "A", 2);
AddVectorToInputsAsTensor({3, 4, 5}, "B", 2);
MP_ASSERT_OK(runner_->Run());
const std::vector<Packet>& output_packets_mult =
runner_->Outputs().Tag("MULTIPLIED").packets;
ASSERT_EQ(3, output_packets_mult.size());
const tf::Tensor& tensor_mult = output_packets_mult[0].Get<tf::Tensor>();
auto expected_tensor = tf::test::AsTensor<int32>({6, 8, 10});
tf::test::ExpectTensorEqual<int32>(tensor_mult, expected_tensor);
const tf::Tensor& tensor_mult1 = output_packets_mult[1].Get<tf::Tensor>();
auto expected_tensor1 = tf::test::AsTensor<int32>({9, 12, 15});
tf::test::ExpectTensorEqual<int32>(tensor_mult1, expected_tensor1);
const tf::Tensor& tensor_mult2 = output_packets_mult[2].Get<tf::Tensor>();
auto expected_tensor2 = tf::test::AsTensor<int32>({12, 16, 20});
tf::test::ExpectTensorEqual<int32>(tensor_mult2, expected_tensor2);
EXPECT_EQ(3, runner_
->GetCounter(
"TensorFlowInferenceCalculator-TotalProcessedTimestamps")
->Get());
}
TEST_F(TensorflowInferenceCalculatorTest, GetSingleBatchComputed) {
CalculatorGraphConfig::Node config;
config.set_calculator("TensorFlowInferenceCalculator");
@@ -311,6 +444,66 @@ TEST_F(TensorflowInferenceCalculatorTest, GetCloseBatchComputed) {
->Get());
}
TEST_F(TensorflowInferenceCalculatorTest, GetBatchComputed_MaxInFlight) {
CalculatorGraphConfig::Node config;
config.set_calculator("TensorFlowInferenceCalculator");
config.add_input_stream("A:tensor_a");
config.add_input_stream("B:tensor_b");
config.add_output_stream("MULTIPLIED:tensor_o1");
config.add_input_side_packet("SESSION:session");
config.set_max_in_flight(2);
CalculatorOptions options;
options.MutableExtension(TensorFlowInferenceCalculatorOptions::ext)
->set_batch_size(2);
options.MutableExtension(TensorFlowInferenceCalculatorOptions::ext)
->set_add_batch_dim_to_tensors(true);
options.MutableExtension(TensorFlowInferenceCalculatorOptions::ext)
->set_batched_input(true);
*config.mutable_options() = options;
runner_ = absl::make_unique<CalculatorRunner>(config);
AddSessionInputSidePacket();
AddVectorToInputsAsPacket(
{CreateTensorPacket({2, 2, 2}, 0), CreateTensorPacket({3, 3, 3}, 1)},
"A");
AddVectorToInputsAsPacket(
{CreateTensorPacket({3, 4, 5}, 0), CreateTensorPacket({3, 4, 5}, 1)},
"B");
AddVectorToInputsAsPacket(
{CreateTensorPacket({4, 4, 4}, 2), CreateTensorPacket({5, 5, 5}, 3)},
"A");
AddVectorToInputsAsPacket(
{CreateTensorPacket({3, 4, 5}, 2), CreateTensorPacket({3, 4, 5}, 3)},
"B");
AddVectorToInputsAsPacket({CreateTensorPacket({6, 6, 6}, 4)}, "A");
AddVectorToInputsAsPacket({CreateTensorPacket({3, 4, 5}, 4)}, "B");
MP_ASSERT_OK(runner_->Run());
const std::vector<Packet>& output_packets_mult =
runner_->Outputs().Tag("MULTIPLIED").packets;
ASSERT_EQ(5, output_packets_mult.size());
const tf::Tensor& tensor_mult = output_packets_mult[0].Get<tf::Tensor>();
auto expected_tensor = tf::test::AsTensor<int32>({6, 8, 10});
tf::test::ExpectTensorEqual<int32>(tensor_mult, expected_tensor);
const tf::Tensor& tensor_mult1 = output_packets_mult[1].Get<tf::Tensor>();
auto expected_tensor1 = tf::test::AsTensor<int32>({9, 12, 15});
tf::test::ExpectTensorEqual<int32>(tensor_mult1, expected_tensor1);
const tf::Tensor& tensor_mult2 = output_packets_mult[2].Get<tf::Tensor>();
auto expected_tensor2 = tf::test::AsTensor<int32>({12, 16, 20});
tf::test::ExpectTensorEqual<int32>(tensor_mult2, expected_tensor2);
const tf::Tensor& tensor_mult3 = output_packets_mult[3].Get<tf::Tensor>();
auto expected_tensor3 = tf::test::AsTensor<int32>({15, 20, 25});
tf::test::ExpectTensorEqual<int32>(tensor_mult3, expected_tensor3);
const tf::Tensor& tensor_mult4 = output_packets_mult[4].Get<tf::Tensor>();
auto expected_tensor4 = tf::test::AsTensor<int32>({18, 24, 30});
tf::test::ExpectTensorEqual<int32>(tensor_mult4, expected_tensor4);
EXPECT_EQ(5, runner_
->GetCounter(
"TensorFlowInferenceCalculator-TotalProcessedTimestamps")
->Get());
}
TEST_F(TensorflowInferenceCalculatorTest, TestRecurrentStates) {
CalculatorGraphConfig::Node config;
config.set_calculator("TensorFlowInferenceCalculator");
@@ -509,4 +702,40 @@ TEST_F(TensorflowInferenceCalculatorTest,
->Get());
}
TEST_F(TensorflowInferenceCalculatorTest, BatchedInputTooBigBatch) {
CalculatorGraphConfig::Node config;
config.set_calculator("TensorFlowInferenceCalculator");
config.add_input_stream("A:tensor_a");
config.add_input_stream("B:tensor_b");
config.add_output_stream("MULTIPLIED:tensor_o1");
config.add_input_side_packet("SESSION:session");
config.set_max_in_flight(2);
CalculatorOptions options;
options.MutableExtension(TensorFlowInferenceCalculatorOptions::ext)
->set_batch_size(2);
options.MutableExtension(TensorFlowInferenceCalculatorOptions::ext)
->set_add_batch_dim_to_tensors(true);
options.MutableExtension(TensorFlowInferenceCalculatorOptions::ext)
->set_batched_input(true);
*config.mutable_options() = options;
runner_ = absl::make_unique<CalculatorRunner>(config);
AddSessionInputSidePacket();
AddVectorToInputsAsPacket(
{CreateTensorPacket({2, 2, 2}, 0), CreateTensorPacket({3, 3, 3}, 1),
CreateTensorPacket({4, 4, 4}, 2)},
"A");
AddVectorToInputsAsPacket(
{CreateTensorPacket({3, 4, 5}, 0), CreateTensorPacket({3, 4, 5}, 1),
CreateTensorPacket({3, 4, 5}, 2)},
"B");
auto status = runner_->Run();
ASSERT_FALSE(status.ok());
EXPECT_THAT(
status.message(),
::testing::HasSubstr(
"has more packets than batch capacity. batch_size: 2 packets: 3"));
}
} // namespace mediapipe
mediapipe/calculators/tensorflow/unpack_media_sequence_calculator.cc
+9 -2
View File
@@ -29,6 +29,7 @@ namespace mediapipe {
// Streams:
const char kBBoxTag[] = "BBOX";
const char kImageTag[] = "IMAGE";
const char kKeypointsTag[] = "KEYPOINTS";
const char kFloatFeaturePrefixTag[] = "FLOAT_FEATURE_";
const char kForwardFlowImageTag[] = "FORWARD_FLOW_ENCODED";
@@ -150,7 +151,6 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
<< "or" << kAudioDecoderOptions;
}
// Optional streams.
if (cc->Outputs().HasTag(kForwardFlowImageTag)) {
cc->Outputs().Tag(kForwardFlowImageTag).Set<std::string>();
}
@@ -244,6 +244,10 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
const auto& sequence = cc->InputSidePackets()
.Tag(kSequenceExampleTag)
.Get<tensorflow::SequenceExample>();
if (cc->Outputs().HasTag(kKeypointsTag)) {
keypoint_names_ = absl::StrSplit(options.keypoint_names(), ',');
default_keypoint_location_ = options.default_keypoint_location();
}
if (cc->OutputSidePackets().HasTag(kDataPath)) {
std::string root_directory = "";
if (cc->InputSidePackets().HasTag(kDatasetRootDirTag)) {
@@ -357,7 +361,6 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
end_timestamp =
timestamps_[last_timestamp_key_][current_timestamp_index_ + 1];
}
for (const auto& map_kv : timestamps_) {
for (int i = 0; i < map_kv.second.size(); ++i) {
if (map_kv.second[i] >= start_timestamp &&
@@ -454,6 +457,10 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
int current_timestamp_index_;
// Store the very first timestamp, so we output everything on the first frame.
int64 first_timestamp_seen_;
// List of keypoint names.
std::vector<std::string> keypoint_names_;
// Default keypoint location when missing.
float default_keypoint_location_;
};
REGISTER_CALCULATOR(UnpackMediaSequenceCalculator);
} // namespace mediapipe
mediapipe/calculators/tflite/BUILD
+2 -1
View File
@@ -16,7 +16,7 @@
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
load("@bazel_skylib//lib:selects.bzl", "selects")
licenses(["notice"]) # Apache 2.0
licenses(["notice"])
package(default_visibility = ["//visibility:private"])
@@ -257,6 +257,7 @@ cc_library(
}) + select({
"//conditions:default": [],
"//mediapipe:android": [
"//mediapipe/util/android/file/base",
"@org_tensorflow//tensorflow/lite/delegates/nnapi:nnapi_delegate",
],
}) + select({
mediapipe/calculators/tflite/tflite_inference_calculator.cc
+57
View File
@@ -33,6 +33,12 @@
#include "tensorflow/lite/kernels/register.h"
#include "tensorflow/lite/model.h"
#if defined(MEDIAPIPE_ANDROID)
#include "mediapipe/util/android/file/base/file.h"
#include "mediapipe/util/android/file/base/filesystem.h"
#include "mediapipe/util/android/file/base/helpers.h"
#endif // ANDROID
#if MEDIAPIPE_TFLITE_GL_INFERENCE
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gpu_buffer.h"
@@ -219,6 +225,8 @@ class TfLiteInferenceCalculator : public CalculatorBase {
::mediapipe::Status Close(CalculatorContext* cc) override;
private:
::mediapipe::Status ReadKernelsFromFile();
::mediapipe::Status WriteKernelsToFile();
::mediapipe::Status LoadModel(CalculatorContext* cc);
::mediapipe::StatusOr<Packet> GetModelAsPacket(const CalculatorContext& cc);
::mediapipe::Status LoadDelegate(CalculatorContext* cc);
@@ -273,6 +281,9 @@ class TfLiteInferenceCalculator : public CalculatorBase {
bool use_quantized_tensors_ = false;
bool use_advanced_gpu_api_ = false;
bool use_kernel_caching_ = false;
std::string cached_kernel_filename_;
};
REGISTER_CALCULATOR(TfLiteInferenceCalculator);
@@ -354,6 +365,17 @@ bool ShouldUseGpu(CC* cc) {
options.has_delegate() &&
options.delegate().has_gpu() &&
options.delegate().gpu().use_advanced_gpu_api();
use_kernel_caching_ =
use_advanced_gpu_api_ && options.delegate().gpu().use_kernel_caching();
if (use_kernel_caching_) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE && defined(MEDIAPIPE_ANDROID)
cached_kernel_filename_ =
"/sdcard/" + mediapipe::File::Basename(options.model_path()) + ".ker";
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE && MEDIAPIPE_ANDROID
}
if (use_advanced_gpu_api_ && !gpu_input_) {
LOG(WARNING) << "Cannot use advanced GPU APIs, input must be GPU buffers."
"Falling back to the default TFLite API.";
@@ -423,7 +445,23 @@ bool ShouldUseGpu(CC* cc) {
});
}
::mediapipe::Status TfLiteInferenceCalculator::WriteKernelsToFile() {
#if MEDIAPIPE_TFLITE_GL_INFERENCE && defined(MEDIAPIPE_ANDROID)
if (use_kernel_caching_) {
// Save kernel file.
auto kernel_cache = absl::make_unique<std::vector<uint8_t>>(
tflite_gpu_runner_->GetSerializedBinaryCache());
std::string cache_str(kernel_cache->begin(), kernel_cache->end());
MP_RETURN_IF_ERROR(
mediapipe::file::SetContents(cached_kernel_filename_, cache_str));
}
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE && MEDIAPIPE_ANDROID
return ::mediapipe::OkStatus();
}
::mediapipe::Status TfLiteInferenceCalculator::Close(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(WriteKernelsToFile());
return RunInContextIfNeeded([this]() -> ::mediapipe::Status {
if (delegate_) {
interpreter_ = nullptr;
@@ -635,6 +673,22 @@ bool ShouldUseGpu(CC* cc) {
return ::mediapipe::OkStatus();
}
::mediapipe::Status TfLiteInferenceCalculator::ReadKernelsFromFile() {
#if MEDIAPIPE_TFLITE_GL_INFERENCE && defined(MEDIAPIPE_ANDROID)
if (use_kernel_caching_) {
// Load pre-compiled kernel file.
if (mediapipe::File::Exists(cached_kernel_filename_)) {
std::string cache_str;
MP_RETURN_IF_ERROR(
mediapipe::file::GetContents(cached_kernel_filename_, &cache_str));
std::vector<uint8_t> cache_vec(cache_str.begin(), cache_str.end());
tflite_gpu_runner_->SetSerializedBinaryCache(std::move(cache_vec));
}
}
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE && MEDIAPIPE_ANDROID
return ::mediapipe::OkStatus();
}
::mediapipe::Status TfLiteInferenceCalculator::InitTFLiteGPURunner(
CalculatorContext* cc) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE
@@ -692,6 +746,9 @@ bool ShouldUseGpu(CC* cc) {
::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_[i]->elements, &gpu_data_out_[i]->buffer));
}
MP_RETURN_IF_ERROR(ReadKernelsFromFile());
MP_RETURN_IF_ERROR(tflite_gpu_runner_->Build());
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
mediapipe/calculators/tflite/tflite_inference_calculator.proto
+4
View File
@@ -48,6 +48,10 @@ message TfLiteInferenceCalculatorOptions {
// example:
// delegate: { gpu { use_advanced_gpu_api: true } }
optional bool use_advanced_gpu_api = 1 [default = false];
// Load pre-compiled serialized binary cache to accelerate init process.
// Only available for OpenCL delegate on Android.
optional bool use_kernel_caching = 2 [default = false];
}
// Android only.
message Nnapi {}
mediapipe/calculators/util/BUILD
+2 -1
View File
@@ -14,7 +14,7 @@
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
licenses(["notice"]) # Apache 2.0
licenses(["notice"])
package(default_visibility = ["//visibility:public"])
@@ -783,6 +783,7 @@ mediapipe_cc_proto_library(
cc_library(
name = "landmarks_to_render_data_calculator",
srcs = ["landmarks_to_render_data_calculator.cc"],
hdrs = ["landmarks_to_render_data_calculator.h"],
visibility = ["//visibility:public"],
deps = [
":landmarks_to_render_data_calculator_cc_proto",
mediapipe/calculators/util/annotation_overlay_calculator.cc
-2
View File
@@ -389,8 +389,6 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
// Upload render target to GPU.
{
glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
glBindTexture(GL_TEXTURE_2D, image_mat_tex_);
glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, width_canvas_, height_canvas_,
GL_RGB, GL_UNSIGNED_BYTE, overlay_image);
mediapipe/calculators/util/landmarks_to_render_data_calculator.cc
+2 -41
View File
@@ -11,6 +11,7 @@
// 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/calculators/util/landmarks_to_render_data_calculator.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
@@ -34,8 +35,6 @@ constexpr char kRenderDataTag[] = "RENDER_DATA";
constexpr char kLandmarkLabel[] = "KEYPOINT";
constexpr int kMaxLandmarkThickness = 18;
using ::mediapipe::RenderAnnotation_Point;
inline void SetColor(RenderAnnotation* annotation, const Color& color) {
annotation->mutable_color()->set_r(color.r());
annotation->mutable_color()->set_g(color.g());
@@ -162,45 +161,6 @@ RenderAnnotation* AddPointRenderData(const Color& landmark_color,
} // namespace
// A calculator that converts Landmark proto to RenderData proto for
// visualization. The input should be LandmarkList proto. It is also possible
// to specify the connections between landmarks.
//
// Example config:
// node {
// calculator: "LandmarksToRenderDataCalculator"
// input_stream: "NORM_LANDMARKS:landmarks"
// output_stream: "RENDER_DATA:render_data"
// options {
// [LandmarksToRenderDataCalculatorOptions.ext] {
// landmark_connections: [0, 1, 1, 2]
// landmark_color { r: 0 g: 255 b: 0 }
// connection_color { r: 0 g: 255 b: 0 }
// thickness: 4.0
// }
// }
// }
class LandmarksToRenderDataCalculator : public CalculatorBase {
public:
LandmarksToRenderDataCalculator() {}
~LandmarksToRenderDataCalculator() override {}
LandmarksToRenderDataCalculator(const LandmarksToRenderDataCalculator&) =
delete;
LandmarksToRenderDataCalculator& operator=(
const LandmarksToRenderDataCalculator&) = delete;
static ::mediapipe::Status GetContract(CalculatorContract* cc);
::mediapipe::Status Open(CalculatorContext* cc) override;
::mediapipe::Status Process(CalculatorContext* cc) override;
private:
LandmarksToRenderDataCalculatorOptions options_;
std::vector<int> landmark_connections_;
};
REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
::mediapipe::Status LandmarksToRenderDataCalculator::GetContract(
CalculatorContract* cc) {
RET_CHECK(cc->Inputs().HasTag(kLandmarksTag) ||
@@ -354,4 +314,5 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
return ::mediapipe::OkStatus();
}
REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
} // namespace mediapipe
mediapipe/calculators/util/landmarks_to_render_data_calculator.h
+69
View File
@@ -0,0 +1,69 @@
// Copyright 2020 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.
#ifndef MEDIAPIPE_CALCULATORS_UTIL_LANDMARKS_TO_RENDER_DATA_CALCULATOR_H_
#define MEDIAPIPE_CALCULATORS_UTIL_LANDMARKS_TO_RENDER_DATA_CALCULATOR_H_
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "mediapipe/calculators/util/landmarks_to_render_data_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_options.pb.h"
#include "mediapipe/framework/formats/landmark.pb.h"
#include "mediapipe/framework/formats/location_data.pb.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/util/color.pb.h"
#include "mediapipe/util/render_data.pb.h"
namespace mediapipe {
// A calculator that converts Landmark proto to RenderData proto for
// visualization. The input should be LandmarkList proto. It is also possible
// to specify the connections between landmarks.
//
// Example config:
// node {
// calculator: "LandmarksToRenderDataCalculator"
// input_stream: "NORM_LANDMARKS:landmarks"
// output_stream: "RENDER_DATA:render_data"
// options {
// [LandmarksToRenderDataCalculatorOptions.ext] {
// landmark_connections: [0, 1, 1, 2]
// landmark_color { r: 0 g: 255 b: 0 }
// connection_color { r: 0 g: 255 b: 0 }
// thickness: 4.0
// }
// }
// }
class LandmarksToRenderDataCalculator : public CalculatorBase {
public:
LandmarksToRenderDataCalculator() {}
~LandmarksToRenderDataCalculator() override {}
LandmarksToRenderDataCalculator(const LandmarksToRenderDataCalculator&) =
delete;
LandmarksToRenderDataCalculator& operator=(
const LandmarksToRenderDataCalculator&) = delete;
static ::mediapipe::Status GetContract(CalculatorContract* cc);
::mediapipe::Status Open(CalculatorContext* cc) override;
::mediapipe::Status Process(CalculatorContext* cc) override;
protected:
::mediapipe::LandmarksToRenderDataCalculatorOptions options_;
std::vector<int> landmark_connections_;
};
} // namespace mediapipe
#endif // MEDIAPIPE_CALCULATORS_UTIL_LANDMARKS_TO_RENDER_DATA_CALCULATOR_H_
mediapipe/calculators/video/BUILD
+1 -1
View File
@@ -19,7 +19,7 @@ load(
"mediapipe_binary_graph",
)
licenses(["notice"]) # Apache 2.0
licenses(["notice"])
package(default_visibility = ["//visibility:private"])
mediapipe/calculators/video/tool/BUILD
+1 -1
View File
@@ -15,7 +15,7 @@
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
licenses(["notice"]) # Apache 2.0
licenses(["notice"])
package(default_visibility = ["//mediapipe/calculators/video:__subpackages__"])