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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: b2062656e5b3d33264e28ed0cbca31c4b93fe1bf
This commit is contained in:
MediaPipe Team
2020-07-29 21:18:36 -04:00
committed by chuoling
parent e9fbe868e5
commit bdfdaef305
132 changed files with 2567 additions and 819 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/calculators/core/BUILD
+15 -1
View File
@@ -450,6 +450,21 @@ cc_library(
alwayslink = 1,
)
cc_test(
name = "mux_calculator_test",
srcs = ["mux_calculator_test.cc"],
deps = [
":mux_calculator",
":round_robin_demux_calculator",
":split_vector_calculator",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:status",
],
)
cc_library(
name = "packet_cloner_calculator",
srcs = ["packet_cloner_calculator.cc"],
@@ -947,7 +962,6 @@ cc_test(
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:status",
],
)
mediapipe/calculators/core/gate_calculator.cc
+64 -15
View File
@@ -56,12 +56,19 @@ std::string ToString(GateState state) {
// disallowing the corresponding packets in other input streams. The behavior
// can be inverted with a calculator option.
//
// ALLOW or DISALLOW can also be specified as an input side packet. The rules
// for evaluation remain the same as above.
//
// ALLOW/DISALLOW inputs must be specified either using input stream or
// via input side packet but not both.
//
// Intended to be used with the default input stream handler, which synchronizes
// all data input streams with the ALLOW/DISALLOW control input stream.
//
// Example config:
// node {
// calculator: "GateCalculator"
// input_side_packet: "ALLOW:allow" or "DISALLOW:disallow"
// input_stream: "input_stream0"
// input_stream: "input_stream1"
// input_stream: "input_streamN"
@@ -75,10 +82,40 @@ class GateCalculator : public CalculatorBase {
public:
GateCalculator() {}
static ::mediapipe::Status CheckAndInitAllowDisallowInputs(
CalculatorContract* cc) {
bool input_via_side_packet = cc->InputSidePackets().HasTag("ALLOW") ||
cc->InputSidePackets().HasTag("DISALLOW");
bool input_via_stream =
cc->Inputs().HasTag("ALLOW") || cc->Inputs().HasTag("DISALLOW");
// Only one of input_side_packet or input_stream may specify ALLOW/DISALLOW
// input.
RET_CHECK(input_via_side_packet ^ input_via_stream);
if (input_via_side_packet) {
RET_CHECK(cc->InputSidePackets().HasTag("ALLOW") ^
cc->InputSidePackets().HasTag("DISALLOW"));
if (cc->InputSidePackets().HasTag("ALLOW")) {
cc->InputSidePackets().Tag("ALLOW").Set<bool>();
} else {
cc->InputSidePackets().Tag("DISALLOW").Set<bool>();
}
} else {
RET_CHECK(cc->Inputs().HasTag("ALLOW") ^ cc->Inputs().HasTag("DISALLOW"));
if (cc->Inputs().HasTag("ALLOW")) {
cc->Inputs().Tag("ALLOW").Set<bool>();
} else {
cc->Inputs().Tag("DISALLOW").Set<bool>();
}
}
return ::mediapipe::OkStatus();
}
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
// Assume that input streams do not have a tag and that gating signal is
// tagged either ALLOW or DISALLOW.
RET_CHECK(cc->Inputs().HasTag("ALLOW") ^ cc->Inputs().HasTag("DISALLOW"));
RET_CHECK_OK(CheckAndInitAllowDisallowInputs(cc));
const int num_data_streams = cc->Inputs().NumEntries("");
RET_CHECK_GE(num_data_streams, 1);
RET_CHECK_EQ(cc->Outputs().NumEntries(""), num_data_streams)
@@ -88,11 +125,6 @@ class GateCalculator : public CalculatorBase {
cc->Inputs().Get("", i).SetAny();
cc->Outputs().Get("", i).SetSameAs(&cc->Inputs().Get("", i));
}
if (cc->Inputs().HasTag("ALLOW")) {
cc->Inputs().Tag("ALLOW").Set<bool>();
} else {
cc->Inputs().Tag("DISALLOW").Set<bool>();
}
if (cc->Outputs().HasTag("STATE_CHANGE")) {
cc->Outputs().Tag("STATE_CHANGE").Set<bool>();
@@ -102,6 +134,17 @@ class GateCalculator : public CalculatorBase {
}
::mediapipe::Status Open(CalculatorContext* cc) final {
use_side_packet_for_allow_disallow_ = false;
if (cc->InputSidePackets().HasTag("ALLOW")) {
use_side_packet_for_allow_disallow_ = true;
allow_by_side_packet_decision_ =
cc->InputSidePackets().Tag("ALLOW").Get<bool>();
} else if (cc->InputSidePackets().HasTag("DISALLOW")) {
use_side_packet_for_allow_disallow_ = true;
allow_by_side_packet_decision_ =
!cc->InputSidePackets().Tag("DISALLOW").Get<bool>();
}
cc->SetOffset(TimestampDiff(0));
num_data_streams_ = cc->Inputs().NumEntries("");
last_gate_state_ = GATE_UNINITIALIZED;
@@ -115,14 +158,18 @@ class GateCalculator : public CalculatorBase {
::mediapipe::Status Process(CalculatorContext* cc) final {
bool allow = empty_packets_as_allow_;
if (cc->Inputs().HasTag("ALLOW") && !cc->Inputs().Tag("ALLOW").IsEmpty()) {
allow = cc->Inputs().Tag("ALLOW").Get<bool>();
if (use_side_packet_for_allow_disallow_) {
allow = allow_by_side_packet_decision_;
} else {
if (cc->Inputs().HasTag("ALLOW") &&
!cc->Inputs().Tag("ALLOW").IsEmpty()) {
allow = cc->Inputs().Tag("ALLOW").Get<bool>();
}
if (cc->Inputs().HasTag("DISALLOW") &&
!cc->Inputs().Tag("DISALLOW").IsEmpty()) {
allow = !cc->Inputs().Tag("DISALLOW").Get<bool>();
}
}
if (cc->Inputs().HasTag("DISALLOW") &&
!cc->Inputs().Tag("DISALLOW").IsEmpty()) {
allow = !cc->Inputs().Tag("DISALLOW").Get<bool>();
}
const GateState new_gate_state = allow ? GATE_ALLOW : GATE_DISALLOW;
if (cc->Outputs().HasTag("STATE_CHANGE")) {
@@ -157,6 +204,8 @@ class GateCalculator : public CalculatorBase {
GateState last_gate_state_ = GATE_UNINITIALIZED;
int num_data_streams_;
bool empty_packets_as_allow_;
bool use_side_packet_for_allow_disallow_;
bool allow_by_side_packet_decision_;
};
REGISTER_CALCULATOR(GateCalculator);
mediapipe/calculators/core/gate_calculator_test.cc
+145 -1
View File
@@ -24,6 +24,21 @@ namespace {
class GateCalculatorTest : public ::testing::Test {
protected:
// Helper to run a graph and return status.
static ::mediapipe::Status RunGraph(const std::string& proto) {
auto runner = absl::make_unique<CalculatorRunner>(
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(proto));
return runner->Run();
}
// Use this when ALLOW/DISALLOW input is provided as a side packet.
void RunTimeStep(int64 timestamp, bool stream_payload) {
runner_->MutableInputs()->Get("", 0).packets.push_back(
MakePacket<bool>(stream_payload).At(Timestamp(timestamp)));
MP_ASSERT_OK(runner_->Run()) << "Calculator execution failed.";
}
// Use this when ALLOW/DISALLOW input is provided as an input stream.
void RunTimeStep(int64 timestamp, const std::string& control_tag,
bool control) {
runner_->MutableInputs()->Get("", 0).packets.push_back(
@@ -31,7 +46,6 @@ class GateCalculatorTest : public ::testing::Test {
runner_->MutableInputs()
->Tag(control_tag)
.packets.push_back(MakePacket<bool>(control).At(Timestamp(timestamp)));
MP_ASSERT_OK(runner_->Run()) << "Calculator execution failed.";
}
@@ -46,6 +60,136 @@ class GateCalculatorTest : public ::testing::Test {
std::unique_ptr<CalculatorRunner> runner_;
};
TEST_F(GateCalculatorTest, InvalidInputs) {
EXPECT_TRUE(absl::IsInternal(GateCalculatorTest::RunGraph(R"(
calculator: "GateCalculator"
input_stream: "test_input"
input_stream: "ALLOW:gating_stream"
input_stream: "DISALLOW:gating_stream"
output_stream: "test_output"
)")));
EXPECT_TRUE(absl::IsInternal(GateCalculatorTest::RunGraph(R"(
calculator: "GateCalculator"
input_stream: "test_input"
input_side_packet: "ALLOW:gating_stream"
input_side_packet: "DISALLOW:gating_stream"
output_stream: "test_output"
)")));
EXPECT_TRUE(absl::IsInternal(GateCalculatorTest::RunGraph(R"(
calculator: "GateCalculator"
input_stream: "test_input"
input_stream: "ALLOW:gating_stream"
input_side_packet: "ALLOW:gating_stream"
output_stream: "test_output"
)")));
EXPECT_TRUE(absl::IsInternal(GateCalculatorTest::RunGraph(R"(
calculator: "GateCalculator"
input_stream: "test_input"
input_stream: "DISALLOW:gating_stream"
input_side_packet: "DISALLOW:gating_stream"
output_stream: "test_output"
)")));
EXPECT_TRUE(absl::IsInternal(GateCalculatorTest::RunGraph(R"(
calculator: "GateCalculator"
input_stream: "test_input"
input_stream: "ALLOW:gating_stream"
input_side_packet: "DISALLOW:gating_stream"
output_stream: "test_output"
)")));
EXPECT_TRUE(absl::IsInternal(GateCalculatorTest::RunGraph(R"(
calculator: "GateCalculator"
input_stream: "test_input"
input_stream: "DISALLOW:gating_stream"
input_side_packet: "ALLOW:gating_stream"
output_stream: "test_output"
)")));
}
TEST_F(GateCalculatorTest, AllowByALLOWSidePacketSetToTrue) {
SetRunner(R"(
calculator: "GateCalculator"
input_side_packet: "ALLOW:gating_stream"
input_stream: "test_input"
output_stream: "test_output"
)");
runner()->MutableSidePackets()->Tag("ALLOW") = Adopt(new bool(true));
constexpr int64 kTimestampValue0 = 42;
RunTimeStep(kTimestampValue0, true);
constexpr int64 kTimestampValue1 = 43;
RunTimeStep(kTimestampValue1, false);
const std::vector<Packet>& output = runner()->Outputs().Get("", 0).packets;
ASSERT_EQ(2, output.size());
EXPECT_EQ(kTimestampValue0, output[0].Timestamp().Value());
EXPECT_EQ(kTimestampValue1, output[1].Timestamp().Value());
EXPECT_EQ(true, output[0].Get<bool>());
EXPECT_EQ(false, output[1].Get<bool>());
}
TEST_F(GateCalculatorTest, AllowByDisallowSidePacketSetToFalse) {
SetRunner(R"(
calculator: "GateCalculator"
input_side_packet: "DISALLOW:gating_stream"
input_stream: "test_input"
output_stream: "test_output"
)");
runner()->MutableSidePackets()->Tag("DISALLOW") = Adopt(new bool(false));
constexpr int64 kTimestampValue0 = 42;
RunTimeStep(kTimestampValue0, true);
constexpr int64 kTimestampValue1 = 43;
RunTimeStep(kTimestampValue1, false);
const std::vector<Packet>& output = runner()->Outputs().Get("", 0).packets;
ASSERT_EQ(2, output.size());
EXPECT_EQ(kTimestampValue0, output[0].Timestamp().Value());
EXPECT_EQ(kTimestampValue1, output[1].Timestamp().Value());
EXPECT_EQ(true, output[0].Get<bool>());
EXPECT_EQ(false, output[1].Get<bool>());
}
TEST_F(GateCalculatorTest, DisallowByALLOWSidePacketSetToFalse) {
SetRunner(R"(
calculator: "GateCalculator"
input_side_packet: "ALLOW:gating_stream"
input_stream: "test_input"
output_stream: "test_output"
)");
runner()->MutableSidePackets()->Tag("ALLOW") = Adopt(new bool(false));
constexpr int64 kTimestampValue0 = 42;
RunTimeStep(kTimestampValue0, true);
constexpr int64 kTimestampValue1 = 43;
RunTimeStep(kTimestampValue1, false);
const std::vector<Packet>& output = runner()->Outputs().Get("", 0).packets;
ASSERT_EQ(0, output.size());
}
TEST_F(GateCalculatorTest, DisallowByDISALLOWSidePacketSetToTrue) {
SetRunner(R"(
calculator: "GateCalculator"
input_side_packet: "DISALLOW:gating_stream"
input_stream: "test_input"
output_stream: "test_output"
)");
runner()->MutableSidePackets()->Tag("DISALLOW") = Adopt(new bool(true));
constexpr int64 kTimestampValue0 = 42;
RunTimeStep(kTimestampValue0, true);
constexpr int64 kTimestampValue1 = 43;
RunTimeStep(kTimestampValue1, false);
const std::vector<Packet>& output = runner()->Outputs().Get("", 0).packets;
ASSERT_EQ(0, output.size());
}
TEST_F(GateCalculatorTest, Allow) {
SetRunner(R"(
calculator: "GateCalculator"
mediapipe/calculators/core/immediate_mux_calculator.cc
+7
View File
@@ -37,6 +37,10 @@ namespace mediapipe {
// the RoundRobinDemuxCalculator. Therefore, packets from different
// input streams are normally not expected to have the same timestamp.
//
// NOTE: this calculator can drop packets non-deterministically, depending on
// how fast the input streams are fed. In most cases, MuxCalculator should be
// preferred. In particular, dropping packets can interfere with rate limiting
// mechanisms.
class ImmediateMuxCalculator : public CalculatorBase {
public:
// This calculator combines any set of input streams into a single
@@ -76,6 +80,9 @@ REGISTER_CALCULATOR(ImmediateMuxCalculator);
if (!packet.IsEmpty()) {
if (packet.Timestamp() >= cc->Outputs().Index(0).NextTimestampBound()) {
cc->Outputs().Index(0).AddPacket(packet);
} else {
LOG_FIRST_N(WARNING, 5)
<< "Dropping a packet with timestamp " << packet.Timestamp();
}
if (cc->Outputs().NumEntries() >= 2) {
Timestamp output_timestamp = std::max(
mediapipe/calculators/core/mux_calculator.cc
+41 -10
View File
@@ -17,28 +17,49 @@
namespace mediapipe {
namespace {
constexpr char kSelectTag[] = "SELECT";
constexpr char kInputTag[] = "INPUT";
} // namespace
// A Calculator that selects an input stream from "INPUT:0", "INPUT:1", ...,
// using the integer value (0, 1, ...) in the packet on the "SELECT" input
// using the integer value (0, 1, ...) in the packet on the kSelectTag input
// stream, and passes the packet on the selected input stream to the "OUTPUT"
// output stream.
// The kSelectTag input can also be passed in as an input side packet, instead
// of as an input stream. Either of input stream or input side packet must be
// specified but not both.
//
// Note that this calculator defaults to use MuxInputStreamHandler, which is
// required for this calculator.
// required for this calculator. However, it can be overridden to work with
// other InputStreamHandlers. Check out the unit tests on for an example usage
// with DefaultInputStreamHandler.
class MuxCalculator : public CalculatorBase {
public:
static ::mediapipe::Status CheckAndInitAllowDisallowInputs(
CalculatorContract* cc) {
RET_CHECK(cc->Inputs().HasTag(kSelectTag) ^
cc->InputSidePackets().HasTag(kSelectTag));
if (cc->Inputs().HasTag(kSelectTag)) {
cc->Inputs().Tag(kSelectTag).Set<int>();
} else {
cc->InputSidePackets().Tag(kSelectTag).Set<int>();
}
return ::mediapipe::OkStatus();
}
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Tag("SELECT").Set<int>();
CollectionItemId data_input_id = cc->Inputs().BeginId("INPUT");
RET_CHECK_OK(CheckAndInitAllowDisallowInputs(cc));
CollectionItemId data_input_id = cc->Inputs().BeginId(kInputTag);
PacketType* data_input0 = &cc->Inputs().Get(data_input_id);
data_input0->SetAny();
++data_input_id;
for (; data_input_id < cc->Inputs().EndId("INPUT"); ++data_input_id) {
for (; data_input_id < cc->Inputs().EndId(kInputTag); ++data_input_id) {
cc->Inputs().Get(data_input_id).SetSameAs(data_input0);
}
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1);
cc->Outputs().Tag("OUTPUT").SetSameAs(data_input0);
// Assign this calculator's default InputStreamHandler.
cc->SetInputStreamHandler("MuxInputStreamHandler");
MediaPipeOptions options;
cc->SetInputStreamHandlerOptions(options);
@@ -47,16 +68,24 @@ class MuxCalculator : public CalculatorBase {
}
::mediapipe::Status Open(CalculatorContext* cc) final {
select_input_ = cc->Inputs().GetId("SELECT", 0);
data_input_base_ = cc->Inputs().GetId("INPUT", 0);
num_data_inputs_ = cc->Inputs().NumEntries("INPUT");
use_side_packet_select_ = false;
if (cc->InputSidePackets().HasTag(kSelectTag)) {
use_side_packet_select_ = true;
selected_index_ = cc->InputSidePackets().Tag(kSelectTag).Get<int>();
} else {
select_input_ = cc->Inputs().GetId(kSelectTag, 0);
}
data_input_base_ = cc->Inputs().GetId(kInputTag, 0);
num_data_inputs_ = cc->Inputs().NumEntries(kInputTag);
output_ = cc->Outputs().GetId("OUTPUT", 0);
cc->SetOffset(TimestampDiff(0));
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
int select = cc->Inputs().Get(select_input_).Get<int>();
int select = use_side_packet_select_
? selected_index_
: cc->Inputs().Get(select_input_).Get<int>();
RET_CHECK(0 <= select && select < num_data_inputs_);
if (!cc->Inputs().Get(data_input_base_ + select).IsEmpty()) {
cc->Outputs().Get(output_).AddPacket(
@@ -70,6 +99,8 @@ class MuxCalculator : public CalculatorBase {
CollectionItemId data_input_base_;
int num_data_inputs_ = 0;
CollectionItemId output_;
bool use_side_packet_select_;
int selected_index_;
};
REGISTER_CALCULATOR(MuxCalculator);
mediapipe/calculators/core/mux_calculator_test.cc
+237
View File
@@ -0,0 +1,237 @@
// 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/calculators/core/split_vector_calculator.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_runner.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/status_matchers.h"
namespace mediapipe {
typedef SplitVectorCalculator<int, false> SplitIntVectorCalculator;
REGISTER_CALCULATOR(SplitIntVectorCalculator);
namespace {
// Graph with default input stream handler, and the input selection is driven
// by an input stream. All MuxCalculator inputs are present at each timestamp.
constexpr char kTestGraphConfig1[] = R"proto(
input_stream: "input"
output_stream: "test_output"
node {
calculator: "SplitIntVectorCalculator"
input_stream: "input"
output_stream: "stream0"
output_stream: "stream1"
output_stream: "stream2"
output_stream: "input_select"
options {
[mediapipe.SplitVectorCalculatorOptions.ext] {
ranges: { begin: 0 end: 1 }
ranges: { begin: 1 end: 2 }
ranges: { begin: 2 end: 3 }
ranges: { begin: 3 end: 4 }
element_only: true
}
}
}
node {
calculator: "MuxCalculator"
input_stream: "INPUT:0:stream0"
input_stream: "INPUT:1:stream1"
input_stream: "INPUT:2:stream2"
input_stream: "SELECT:input_select"
output_stream: "OUTPUT:test_output"
input_stream_handler { input_stream_handler: "DefaultInputStreamHandler" }
}
)proto";
// Graph with default input stream handler, and the input selection is driven
// by an input side packet. All MuxCalculator inputs are present at each
// timestamp.
constexpr char kTestGraphConfig2[] = R"proto(
input_side_packet: "input_selector"
input_stream: "input"
output_stream: "test_output"
node {
calculator: "SplitIntVectorCalculator"
input_stream: "input"
output_stream: "stream0"
output_stream: "stream1"
output_stream: "stream2"
options {
[mediapipe.SplitVectorCalculatorOptions.ext] {
ranges: { begin: 0 end: 1 }
ranges: { begin: 1 end: 2 }
ranges: { begin: 2 end: 3 }
element_only: true
}
}
}
node {
calculator: "MuxCalculator"
input_stream: "INPUT:0:stream0"
input_stream: "INPUT:1:stream1"
input_stream: "INPUT:2:stream2"
input_side_packet: "SELECT:input_selector"
output_stream: "OUTPUT:test_output"
input_stream_handler { input_stream_handler: "DefaultInputStreamHandler" }
}
)proto";
// Graph with mux input stream handler, and the input selection is driven
// by an input stream. Only one MuxCalculator input is present at each
// timestamp.
constexpr char kTestGraphConfig3[] = R"proto(
input_stream: "input"
output_stream: "test_output"
node {
calculator: "RoundRobinDemuxCalculator"
input_stream: "input"
output_stream: "OUTPUT:0:stream0"
output_stream: "OUTPUT:1:stream1"
output_stream: "OUTPUT:2:stream2"
output_stream: "SELECT:input_select"
}
node {
calculator: "MuxCalculator"
input_stream: "INPUT:0:stream0"
input_stream: "INPUT:1:stream1"
input_stream: "INPUT:2:stream2"
input_stream: "SELECT:input_select"
output_stream: "OUTPUT:test_output"
}
)proto";
constexpr char kOutputName[] = "test_output";
constexpr char kInputName[] = "input";
constexpr char kInputSelector[] = "input_selector";
// Helper to run a graph with the given inputs and generate outputs, asserting
// each step along the way.
// Inputs:
// graph_config_proto - graph config protobuf
// extra_side_packets - input side packets name to value map
// input_stream_name - name of the input
void RunGraph(const std::string& graph_config_proto,
const std::map<std::string, Packet>& extra_side_packets,
const std::string& input_stream_name, int num_input_packets,
std::function<Packet(int)> input_fn,
const std::string& output_stream_name,
std::function<::mediapipe::Status(const Packet&)> output_fn) {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
graph_config_proto);
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.ObserveOutputStream(output_stream_name, output_fn));
MP_ASSERT_OK(graph.StartRun(extra_side_packets));
for (int i = 0; i < num_input_packets; ++i) {
MP_ASSERT_OK(graph.AddPacketToInputStream(input_stream_name, input_fn(i)));
}
MP_ASSERT_OK(graph.CloseAllInputStreams());
MP_ASSERT_OK(graph.WaitUntilDone());
}
TEST(MuxCalculatorTest, InputStreamSelector_DefaultInputStreamHandler) {
// Input and handling.
std::vector<std::vector<int>> input_packets = {
{1, 1, 2, 1}, {3, 5, 8, 2}, {13, 21, 34, 0},
{55, 89, 144, 2}, {233, 377, 610, 0}, {987, 1597, 2584, 1},
{4181, 6765, 10946, 2},
};
int packet_time_stamp = 22;
// This function will return the i-th input packet.
auto input_fn = [&packet_time_stamp, &input_packets](int i) -> Packet {
return MakePacket<std::vector<int>>(input_packets[i])
.At(Timestamp(packet_time_stamp++));
};
// Output and handling.
std::vector<int> output;
// This function collects the output from the packet.
auto output_fn = [&output](const Packet& p) -> ::mediapipe::Status {
output.push_back(p.Get<int>());
return ::mediapipe::OkStatus();
};
RunGraph(kTestGraphConfig1, {}, kInputName, input_packets.size(), input_fn,
kOutputName, output_fn);
EXPECT_THAT(output, testing::ElementsAre(1, 8, 13, 144, 233, 1597, 10946));
}
TEST(MuxCalculatorTest, InputSidePacketSelector_DefaultInputStreamHandler) {
// Input and handling.
std::vector<std::vector<int>> input_packets = {
{1, 1, 2}, {3, 5, 8}, {13, 21, 34}, {55, 89, 144},
{233, 377, 610}, {987, 1597, 2584}, {4181, 6765, 10946},
};
int packet_time_stamp = 22;
// This function will return the i-th input packet.
auto input_fn = [&packet_time_stamp, &input_packets](int i) -> Packet {
return MakePacket<std::vector<int>>(input_packets[i])
.At(Timestamp(packet_time_stamp++));
};
// Output and handling.
std::vector<int> output;
// This function collects the output from the packet.
auto output_fn = [&output](const Packet& p) -> ::mediapipe::Status {
output.push_back(p.Get<int>());
return ::mediapipe::OkStatus();
};
RunGraph(kTestGraphConfig2, {{kInputSelector, MakePacket<int>(0)}},
kInputName, input_packets.size(), input_fn, kOutputName, output_fn);
EXPECT_THAT(output, testing::ElementsAre(1, 3, 13, 55, 233, 987, 4181));
output.clear();
RunGraph(kTestGraphConfig2, {{kInputSelector, MakePacket<int>(1)}},
kInputName, input_packets.size(), input_fn, kOutputName, output_fn);
EXPECT_THAT(output, testing::ElementsAre(1, 5, 21, 89, 377, 1597, 6765));
output.clear();
RunGraph(kTestGraphConfig2, {{kInputSelector, MakePacket<int>(2)}},
kInputName, input_packets.size(), input_fn, kOutputName, output_fn);
EXPECT_THAT(output, testing::ElementsAre(2, 8, 34, 144, 610, 2584, 10946));
}
TEST(MuxCalculatorTest, InputStreamSelector_MuxInputStreamHandler) {
// Input and handling.
std::vector<int> input_packets = {1, 1, 2, 3, 5, 8, 13,
21, 34, 55, 89, 144, 233, 377,
610, 987, 1597, 2584, 4181, 6765, 10946};
int packet_time_stamp = 22;
// This function will return the i-th input packet.
auto input_fn = [&packet_time_stamp, &input_packets](int i) -> Packet {
return MakePacket<int>(input_packets[i]).At(Timestamp(packet_time_stamp++));
};
// Output and handling.
std::vector<int> output;
// This function collects the output from the packet.
auto output_fn = [&output](const Packet& p) -> ::mediapipe::Status {
output.push_back(p.Get<int>());
return ::mediapipe::OkStatus();
};
RunGraph(kTestGraphConfig3, {}, kInputName, input_packets.size(), input_fn,
kOutputName, output_fn);
EXPECT_EQ(output, input_packets);
}
} // namespace
} // namespace mediapipe
mediapipe/calculators/core/previous_loopback_calculator.cc
+9 -3
View File
@@ -128,11 +128,17 @@ class PreviousLoopbackCalculator : public CalculatorBase {
loop_packets_.pop_front();
main_packet_specs_.pop_front();
}
// We can close PREV_LOOP output stream as soon as we processed last
// possible MAIN packet. That can happen in two cases:
// a) Non-empty MAIN packet has been received with Timestamp::Max()
// b) Empty MAIN packet has been received with Timestamp::Max() indicating
// MAIN is done.
if (main_spec.timestamp == Timestamp::Done().PreviousAllowedInStream()) {
prev_loop.Close();
}
}
if (main_packet_specs_.empty() && cc->Inputs().Get(main_id_).IsDone()) {
prev_loop.Close();
}
return ::mediapipe::OkStatus();
}
mediapipe/calculators/core/previous_loopback_calculator_test.cc
+98
View File
@@ -228,6 +228,104 @@ TEST(PreviousLoopbackCalculator, ClosesCorrectly) {
MP_EXPECT_OK(graph_.WaitUntilDone());
}
TEST(PreviousLoopbackCalculator, ProcessesMaxTimestamp) {
std::vector<Packet> out_and_previous_packets;
CalculatorGraphConfig graph_config =
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: 'in'
node {
calculator: 'PreviousLoopbackCalculator'
input_stream: 'MAIN:in'
input_stream: 'LOOP:out'
input_stream_info: { tag_index: 'LOOP' back_edge: true }
output_stream: 'PREV_LOOP:previous'
}
node {
calculator: 'PassThroughCalculator'
input_stream: 'in'
input_stream: 'previous'
output_stream: 'out'
output_stream: 'previous2'
}
node {
calculator: 'MakePairCalculator'
input_stream: 'out'
input_stream: 'previous'
output_stream: 'out_and_previous'
}
)");
tool::AddVectorSink("out_and_previous", &graph_config,
&out_and_previous_packets);
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(graph_config, {}));
MP_ASSERT_OK(graph.StartRun({}));
MP_EXPECT_OK(graph.AddPacketToInputStream(
"in", MakePacket<int>(1).At(Timestamp::Max())));
MP_EXPECT_OK(graph.WaitUntilIdle());
EXPECT_THAT(out_and_previous_packets,
ElementsAre(PairPacket(Timestamp::Max(),
Pair(IntPacket(1), EmptyPacket()))));
MP_EXPECT_OK(graph.CloseAllInputStreams());
MP_EXPECT_OK(graph.WaitUntilIdle());
MP_EXPECT_OK(graph.WaitUntilDone());
}
TEST(PreviousLoopbackCalculator, ProcessesMaxTimestampNonEmptyPrevious) {
std::vector<Packet> out_and_previous_packets;
CalculatorGraphConfig graph_config =
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: 'in'
node {
calculator: 'PreviousLoopbackCalculator'
input_stream: 'MAIN:in'
input_stream: 'LOOP:out'
input_stream_info: { tag_index: 'LOOP' back_edge: true }
output_stream: 'PREV_LOOP:previous'
}
node {
calculator: 'PassThroughCalculator'
input_stream: 'in'
input_stream: 'previous'
output_stream: 'out'
output_stream: 'previous2'
}
node {
calculator: 'MakePairCalculator'
input_stream: 'out'
input_stream: 'previous'
output_stream: 'out_and_previous'
}
)");
tool::AddVectorSink("out_and_previous", &graph_config,
&out_and_previous_packets);
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(graph_config, {}));
MP_ASSERT_OK(graph.StartRun({}));
MP_EXPECT_OK(graph.AddPacketToInputStream(
"in", MakePacket<int>(1).At(Timestamp::Min())));
MP_EXPECT_OK(graph.AddPacketToInputStream(
"in", MakePacket<int>(2).At(Timestamp::Max())));
MP_EXPECT_OK(graph.WaitUntilIdle());
EXPECT_THAT(
out_and_previous_packets,
ElementsAre(
PairPacket(Timestamp::Min(), Pair(IntPacket(1), EmptyPacket())),
PairPacket(Timestamp::Max(), Pair(IntPacket(2), IntPacket(1)))));
MP_EXPECT_OK(graph.CloseAllInputStreams());
MP_EXPECT_OK(graph.WaitUntilIdle());
MP_EXPECT_OK(graph.WaitUntilDone());
}
// Demonstrates that downstream calculators won't be blocked by
// always-empty-LOOP-stream.
TEST(PreviousLoopbackCalculator, EmptyLoopForever) {
mediapipe/calculators/core/side_packet_to_stream_calculator.cc
+41 -8
View File
@@ -34,6 +34,8 @@ constexpr char kTagAtPostStream[] = "AT_POSTSTREAM";
constexpr char kTagAtZero[] = "AT_ZERO";
constexpr char kTagAtTick[] = "AT_TICK";
constexpr char kTagTick[] = "TICK";
constexpr char kTagAtTimestamp[] = "AT_TIMESTAMP";
constexpr char kTagSideInputTimestamp[] = "TIMESTAMP";
static std::map<std::string, Timestamp>* kTimestampMap = []() {
auto* res = new std::map<std::string, Timestamp>();
@@ -41,6 +43,7 @@ static std::map<std::string, Timestamp>* kTimestampMap = []() {
res->emplace(kTagAtPostStream, Timestamp::PostStream());
res->emplace(kTagAtZero, Timestamp(0));
res->emplace(kTagAtTick, Timestamp::Unset());
res->emplace(kTagAtTimestamp, Timestamp::Unset());
return res;
}();
@@ -56,9 +59,10 @@ std::string GetOutputTag(const CC& cc) {
// timestamp, depending on the tag used to define output stream(s). (One tag can
// be used only.)
//
// Valid tags are AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO and AT_TICK and
// corresponding timestamps are Timestamp::PreStream(), Timestamp::PostStream(),
// Timestamp(0) and timestamp of a packet received in TICK input.
// Valid tags are AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO, AT_TICK, AT_TIMESTAMP
// and corresponding timestamps are Timestamp::PreStream(),
// Timestamp::PostStream(), Timestamp(0), timestamp of a packet received in TICK
// input, and timestamp received from a side input.
//
// Examples:
// node {
@@ -73,6 +77,13 @@ std::string GetOutputTag(const CC& cc) {
// input_side_packet: "side_packet"
// output_stream: "AT_TICK:packet"
// }
//
// node {
// calculator: "SidePacketToStreamCalculator"
// input_side_packet: "TIMESTAMP:timestamp"
// input_side_packet: "side_packet"
// output_stream: "AT_TIMESTAMP:packet"
// }
class SidePacketToStreamCalculator : public CalculatorBase {
public:
SidePacketToStreamCalculator() = default;
@@ -93,16 +104,29 @@ REGISTER_CALCULATOR(SidePacketToStreamCalculator);
CalculatorContract* cc) {
const auto& tags = cc->Outputs().GetTags();
RET_CHECK(tags.size() == 1 && kTimestampMap->count(*tags.begin()) == 1)
<< "Only one of AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO and AT_TICK tags is "
"allowed and required to specify output stream(s).";
<< "Only one of AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO, AT_TICK and "
"AT_TIMESTAMP tags is allowed and required to specify output "
"stream(s).";
RET_CHECK(
(cc->Outputs().HasTag(kTagAtTick) && cc->Inputs().HasTag(kTagTick)) ||
(!cc->Outputs().HasTag(kTagAtTick) && !cc->Inputs().HasTag(kTagTick)))
<< "Either both of TICK and AT_TICK should be used or none of them.";
RET_CHECK((cc->Outputs().HasTag(kTagAtTimestamp) &&
cc->InputSidePackets().HasTag(kTagSideInputTimestamp)) ||
(!cc->Outputs().HasTag(kTagAtTimestamp) &&
!cc->InputSidePackets().HasTag(kTagSideInputTimestamp)))
<< "Either both TIMESTAMP and AT_TIMESTAMP should be used or none of "
"them.";
const std::string output_tag = GetOutputTag(*cc);
const int num_entries = cc->Outputs().NumEntries(output_tag);
RET_CHECK_EQ(num_entries, cc->InputSidePackets().NumEntries())
<< "Same number of input side packets and output streams is required.";
if (cc->Outputs().HasTag(kTagAtTimestamp)) {
RET_CHECK_EQ(num_entries + 1, cc->InputSidePackets().NumEntries())
<< "For AT_TIMESTAMP tag, 2 input side packets are required.";
cc->InputSidePackets().Tag(kTagSideInputTimestamp).Set<int64>();
} else {
RET_CHECK_EQ(num_entries, cc->InputSidePackets().NumEntries())
<< "Same number of input side packets and output streams is required.";
}
for (int i = 0; i < num_entries; ++i) {
cc->InputSidePackets().Index(i).SetAny();
cc->Outputs()
@@ -147,13 +171,22 @@ REGISTER_CALCULATOR(SidePacketToStreamCalculator);
}
::mediapipe::Status SidePacketToStreamCalculator::Close(CalculatorContext* cc) {
if (!cc->Outputs().HasTag(kTagAtTick)) {
if (!cc->Outputs().HasTag(kTagAtTick) &&
!cc->Outputs().HasTag(kTagAtTimestamp)) {
const auto& timestamp = kTimestampMap->at(output_tag_);
for (int i = 0; i < cc->Outputs().NumEntries(output_tag_); ++i) {
cc->Outputs()
.Get(output_tag_, i)
.AddPacket(cc->InputSidePackets().Index(i).At(timestamp));
}
} else if (cc->Outputs().HasTag(kTagAtTimestamp)) {
int64 timestamp =
cc->InputSidePackets().Tag(kTagSideInputTimestamp).Get<int64>();
for (int i = 0; i < cc->Outputs().NumEntries(output_tag_); ++i) {
cc->Outputs()
.Get(output_tag_, i)
.AddPacket(cc->InputSidePackets().Index(i).At(Timestamp(timestamp)));
}
}
return ::mediapipe::OkStatus();
}
mediapipe/calculators/core/side_packet_to_stream_calculator_test.cc
+101 -4
View File
@@ -51,6 +51,27 @@ TEST(SidePacketToStreamCalculator, WrongConfig_MissingTick) {
"Either both of TICK and AT_TICK should be used or none of them.");
}
TEST(SidePacketToStreamCalculator, WrongConfig_MissingTimestampSideInput) {
CalculatorGraphConfig graph_config =
ParseTextProtoOrDie<CalculatorGraphConfig>(
R"(
input_stream: "timestamp"
input_side_packet: "side_packet"
output_stream: "packet"
node {
calculator: "SidePacketToStreamCalculator"
input_side_packet: "side_packet"
output_stream: "AT_TIMESTAMP:packet"
}
)");
CalculatorGraph graph;
auto status = graph.Initialize(graph_config);
EXPECT_FALSE(status.ok());
EXPECT_PRED2(
absl::StrContains, status.message(),
"Either both TIMESTAMP and AT_TIMESTAMP should be used or none of them.");
}
TEST(SidePacketToStreamCalculator, WrongConfig_NonExistentTag) {
CalculatorGraphConfig graph_config =
ParseTextProtoOrDie<CalculatorGraphConfig>(
@@ -68,8 +89,9 @@ TEST(SidePacketToStreamCalculator, WrongConfig_NonExistentTag) {
auto status = graph.Initialize(graph_config);
EXPECT_FALSE(status.ok());
EXPECT_PRED2(absl::StrContains, status.message(),
"Only one of AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO and AT_TICK "
"tags is allowed and required to specify output stream(s).");
"Only one of AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO, AT_TICK and "
"AT_TIMESTAMP tags is allowed and required to specify output "
"stream(s).");
}
TEST(SidePacketToStreamCalculator, WrongConfig_MixedTags) {
@@ -91,8 +113,9 @@ TEST(SidePacketToStreamCalculator, WrongConfig_MixedTags) {
auto status = graph.Initialize(graph_config);
EXPECT_FALSE(status.ok());
EXPECT_PRED2(absl::StrContains, status.message(),
"Only one of AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO and AT_TICK "
"tags is allowed and required to specify output stream(s).");
"Only one of AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO, AT_TICK and "
"AT_TIMESTAMP tags is allowed and required to specify output "
"stream(s).");
}
TEST(SidePacketToStreamCalculator, WrongConfig_NotEnoughSidePackets) {
@@ -271,5 +294,79 @@ TEST(SidePacketToStreamCalculator, AtTick_MultipleSidePackets) {
tick_and_verify(/*at_timestamp=*/1025);
}
TEST(SidePacketToStreamCalculator, AtTimestamp) {
CalculatorGraphConfig graph_config =
ParseTextProtoOrDie<CalculatorGraphConfig>(
R"(
input_side_packet: "timestamp"
input_side_packet: "side_packet"
output_stream: "packet"
node {
calculator: "SidePacketToStreamCalculator"
input_side_packet: "TIMESTAMP:timestamp"
input_side_packet: "side_packet"
output_stream: "AT_TIMESTAMP:packet"
}
)");
std::vector<Packet> output_packets;
tool::AddVectorSink("packet", &graph_config, &output_packets);
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(graph_config));
const int expected_value = 20;
const int64 expected_timestamp = 5;
MP_ASSERT_OK(
graph.StartRun({{"side_packet", MakePacket<int>(expected_value)},
{"timestamp", MakePacket<int64>(expected_timestamp)}}));
MP_ASSERT_OK(graph.WaitUntilDone());
ASSERT_FALSE(output_packets.empty());
EXPECT_EQ(Timestamp(expected_timestamp), output_packets.back().Timestamp());
EXPECT_EQ(expected_value, output_packets.back().Get<int>());
}
TEST(SidePacketToStreamCalculator, AtTimestamp_MultipleOutputs) {
CalculatorGraphConfig graph_config =
ParseTextProtoOrDie<CalculatorGraphConfig>(
R"(
input_side_packet: "timestamp"
input_side_packet: "side_packet0"
input_side_packet: "side_packet1"
output_stream: "packet"
node {
calculator: "SidePacketToStreamCalculator"
input_side_packet: "TIMESTAMP:timestamp"
input_side_packet: "side_packet0"
input_side_packet: "side_packet1"
output_stream: "AT_TIMESTAMP:0:packet0"
output_stream: "AT_TIMESTAMP:1:packet1"
}
)");
std::vector<Packet> output_packets0;
tool::AddVectorSink("packet0", &graph_config, &output_packets0);
std::vector<Packet> output_packets1;
tool::AddVectorSink("packet1", &graph_config, &output_packets1);
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(graph_config));
const int expected_value0 = 20;
const int expected_value1 = 15;
const int64 expected_timestamp = 5;
MP_ASSERT_OK(
graph.StartRun({{"side_packet0", MakePacket<int>(expected_value0)},
{"side_packet1", MakePacket<int>(expected_value1)},
{"timestamp", MakePacket<int64>(expected_timestamp)}}));
MP_ASSERT_OK(graph.WaitUntilDone());
ASSERT_FALSE(output_packets0.empty());
EXPECT_EQ(Timestamp(expected_timestamp), output_packets0.back().Timestamp());
EXPECT_EQ(expected_value0, output_packets0.back().Get<int>());
ASSERT_FALSE(output_packets1.empty());
EXPECT_EQ(Timestamp(expected_timestamp), output_packets1.back().Timestamp());
EXPECT_EQ(expected_value1, output_packets1.back().Get<int>());
}
} // namespace
} // namespace mediapipe
mediapipe/calculators/image/image_transformation_calculator.cc
-4
View File
@@ -449,19 +449,15 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
switch (rotation_) {
case mediapipe::RotationMode_Mode_UNKNOWN:
case mediapipe::RotationMode_Mode_ROTATION_0:
LOG(ERROR) << "Not rotating image.";
rotated_mat = input_mat;
break;
case mediapipe::RotationMode_Mode_ROTATION_90:
LOG(ERROR) << "Rotating image by 90 degrees ccw.";
cv::rotate(input_mat, rotated_mat, cv::ROTATE_90_COUNTERCLOCKWISE);
break;
case mediapipe::RotationMode_Mode_ROTATION_180:
LOG(ERROR) << "Rotating image by 180 degrees.";
cv::rotate(input_mat, rotated_mat, cv::ROTATE_180);
break;
case mediapipe::RotationMode_Mode_ROTATION_270:
LOG(ERROR) << "Rotating image by 90 degrees cw.";
cv::rotate(input_mat, rotated_mat, cv::ROTATE_90_CLOCKWISE);
break;
}
mediapipe/calculators/tensorflow/BUILD
+14 -18
View File
@@ -57,22 +57,6 @@ proto_library(
deps = ["//mediapipe/framework:calculator_proto"],
)
proto_library(
name = "tensorflow_session_from_saved_model_generator_proto",
srcs = ["tensorflow_session_from_saved_model_generator.proto"],
visibility = ["//visibility:public"],
deps = ["//mediapipe/framework:packet_generator_proto"],
)
proto_library(
name = "tensorflow_session_from_saved_model_calculator_proto",
srcs = ["tensorflow_session_from_saved_model_calculator.proto"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_proto",
],
)
proto_library(
name = "tensor_squeeze_dimensions_calculator_proto",
srcs = ["tensor_squeeze_dimensions_calculator.proto"],
@@ -212,7 +196,10 @@ mediapipe_cc_proto_library(
mediapipe_cc_proto_library(
name = "tensorflow_session_from_saved_model_generator_cc_proto",
srcs = ["tensorflow_session_from_saved_model_generator.proto"],
cc_deps = ["//mediapipe/framework:packet_generator_cc_proto"],
cc_deps = [
"//mediapipe/framework:packet_generator_cc_proto",
"@org_tensorflow//tensorflow/core:protos_all_cc",
],
visibility = ["//visibility:public"],
deps = [":tensorflow_session_from_saved_model_generator_proto"],
)
@@ -220,7 +207,10 @@ mediapipe_cc_proto_library(
mediapipe_cc_proto_library(
name = "tensorflow_session_from_saved_model_calculator_cc_proto",
srcs = ["tensorflow_session_from_saved_model_calculator.proto"],
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
cc_deps = [
"//mediapipe/framework:calculator_cc_proto",
"@org_tensorflow//tensorflow/core:protos_all_cc",
],
visibility = ["//visibility:public"],
deps = [":tensorflow_session_from_saved_model_calculator_proto"],
)
@@ -488,6 +478,8 @@ cc_library(
"//mediapipe/calculators/tensorflow:tensorflow_session_from_frozen_graph_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/tool:status_util",
"//mediapipe/framework/deps:clock",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:ret_check",
] + select({
@@ -518,6 +510,8 @@ cc_library(
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/tool:status_util",
"//mediapipe/framework/port:status",
"//mediapipe/framework/deps:clock",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:ret_check",
] + select({
"//conditions:default": [
@@ -929,6 +923,7 @@ cc_test(
"@com_google_absl//absl/strings",
"@org_tensorflow//tensorflow/core:all_kernels",
"@org_tensorflow//tensorflow/core:direct_session",
"@org_tensorflow//tensorflow/core:protos_all_cc",
],
)
@@ -954,6 +949,7 @@ cc_test(
"@com_google_absl//absl/strings",
"@org_tensorflow//tensorflow/core:all_kernels",
"@org_tensorflow//tensorflow/core:direct_session",
"@org_tensorflow//tensorflow/core:protos_all_cc",
],
)
mediapipe/calculators/tensorflow/tensorflow_session_from_frozen_graph_calculator.cc
+28
View File
@@ -26,9 +26,14 @@
#include "mediapipe/calculators/tensorflow/tensorflow_session.h"
#include "mediapipe/calculators/tensorflow/tensorflow_session_from_frozen_graph_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/deps/clock.h"
#include "mediapipe/framework/deps/monotonic_clock.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/tool/status_util.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/public/session_options.h"
#if defined(MEDIAPIPE_MOBILE)
@@ -41,6 +46,17 @@ namespace mediapipe {
namespace tf = ::tensorflow;
namespace {
// Updates the graph nodes to use the device as specified by device_id.
void SetPreferredDevice(tf::GraphDef* graph_def, absl::string_view device_id) {
for (auto& node : *graph_def->mutable_node()) {
if (node.device().empty()) {
node.set_device(device_id);
}
}
}
} // namespace
class TensorFlowSessionFromFrozenGraphCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
@@ -77,6 +93,9 @@ class TensorFlowSessionFromFrozenGraphCalculator : public CalculatorBase {
}
::mediapipe::Status Open(CalculatorContext* cc) override {
auto clock = std::unique_ptr<mediapipe::Clock>(
mediapipe::MonotonicClock::CreateSynchronizedMonotonicClock());
const uint64 start_time = absl::ToUnixMicros(clock->TimeNow());
const auto& options =
cc->Options<TensorFlowSessionFromFrozenGraphCalculatorOptions>();
// Output bundle packet.
@@ -108,6 +127,12 @@ class TensorFlowSessionFromFrozenGraphCalculator : public CalculatorBase {
tensorflow::GraphDef graph_def;
RET_CHECK(graph_def.ParseFromString(graph_def_serialized));
// Update the graph nodes to use the preferred device, if set.
if (!options.preferred_device_id().empty()) {
SetPreferredDevice(&graph_def, options.preferred_device_id());
}
const tf::Status tf_status = session->session->Create(graph_def);
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.ToString();
@@ -123,6 +148,9 @@ class TensorFlowSessionFromFrozenGraphCalculator : public CalculatorBase {
}
cc->OutputSidePackets().Tag("SESSION").Set(Adopt(session.release()));
const uint64 end_time = absl::ToUnixMicros(clock->TimeNow());
LOG(INFO) << "Loaded frozen model in: " << end_time - start_time
<< " microseconds.";
return ::mediapipe::OkStatus();
}
mediapipe/calculators/tensorflow/tensorflow_session_from_frozen_graph_calculator.proto
+8
View File
@@ -69,4 +69,12 @@ message TensorFlowSessionFromFrozenGraphCalculatorOptions {
// Graph nodes to run to initialize the model. Any output of these ops is
// ignored.
repeated string initialization_op_names = 4;
// The id of the device you would prefer to execute the graph nodes on.
// If set, all graph nodes without a previously specified device, will be set
// to run on preferred_device_id. Example values include:
// ["/device:GPU:0","/device:CPU:0", ...]
// NOTE: If config.allow_soft_placement = false, and the device is not found,
// an error will be thrown.
optional string preferred_device_id = 5;
}
mediapipe/calculators/tensorflow/tensorflow_session_from_frozen_graph_calculator_test.cc
+1
View File
@@ -66,6 +66,7 @@ class TensorFlowSessionFromFrozenGraphCalculatorTest : public ::testing::Test {
(*calculator_options_->mutable_tag_to_tensor_names())["B"] = "b:0";
calculator_options_->mutable_config()->set_intra_op_parallelism_threads(1);
calculator_options_->mutable_config()->set_inter_op_parallelism_threads(2);
calculator_options_->set_preferred_device_id("/device:CPU:0");
}
void VerifySignatureMap(const TensorFlowSession& session) {
mediapipe/calculators/tensorflow/tensorflow_session_from_frozen_graph_generator.cc
+28
View File
@@ -27,16 +27,32 @@
#include "mediapipe/calculators/tensorflow/tensorflow_session.h"
#include "mediapipe/calculators/tensorflow/tensorflow_session_from_frozen_graph_generator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/deps/clock.h"
#include "mediapipe/framework/deps/monotonic_clock.h"
#include "mediapipe/framework/port/file_helpers.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/tool/status_util.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/public/session_options.h"
namespace mediapipe {
namespace tf = ::tensorflow;
namespace {
// Updates the graph nodes to use the device as specified by device_id.
void SetPreferredDevice(tf::GraphDef* graph_def, absl::string_view device_id) {
for (auto& node : *graph_def->mutable_node()) {
if (node.device().empty()) {
node.set_device(device_id);
}
}
}
} // namespace
class TensorFlowSessionFromFrozenGraphGenerator : public PacketGenerator {
public:
static ::mediapipe::Status FillExpectations(
@@ -77,6 +93,9 @@ class TensorFlowSessionFromFrozenGraphGenerator : public PacketGenerator {
static ::mediapipe::Status Generate(
const PacketGeneratorOptions& packet_generator_options,
const PacketSet& input_side_packets, PacketSet* output_side_packets) {
auto clock = std::unique_ptr<mediapipe::Clock>(
mediapipe::MonotonicClock::CreateSynchronizedMonotonicClock());
const uint64 start_time = absl::ToUnixMicros(clock->TimeNow());
const TensorFlowSessionFromFrozenGraphGeneratorOptions& options =
packet_generator_options.GetExtension(
TensorFlowSessionFromFrozenGraphGeneratorOptions::ext);
@@ -108,6 +127,12 @@ class TensorFlowSessionFromFrozenGraphGenerator : public PacketGenerator {
tensorflow::GraphDef graph_def;
RET_CHECK(graph_def.ParseFromString(graph_def_serialized));
// Update the graph nodes to use the preferred device, if set.
if (!options.preferred_device_id().empty()) {
SetPreferredDevice(&graph_def, options.preferred_device_id());
}
const tf::Status tf_status = session->session->Create(graph_def);
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.ToString();
@@ -123,6 +148,9 @@ class TensorFlowSessionFromFrozenGraphGenerator : public PacketGenerator {
}
output_side_packets->Tag("SESSION") = Adopt(session.release());
const uint64 end_time = absl::ToUnixMicros(clock->TimeNow());
LOG(INFO) << "Loaded frozen model in: " << end_time - start_time
<< " microseconds.";
return ::mediapipe::OkStatus();
}
};
mediapipe/calculators/tensorflow/tensorflow_session_from_frozen_graph_generator.proto
+8
View File
@@ -69,4 +69,12 @@ message TensorFlowSessionFromFrozenGraphGeneratorOptions {
// Graph nodes to run to initialize the model. Any output of these ops is
// ignored.
repeated string initialization_op_names = 4;
// The id of the device you would prefer to execute the graph nodes on.
// If set, all graph nodes without a previously specified device, will be set
// to run on preferred_device_id. Example values include:
// ["/device:GPU:0","/device:CPU:0", ...]
// NOTE: If config.allow_soft_placement = false, and the device is not found,
// an error will be thrown.
optional string preferred_device_id = 5;
}
mediapipe/calculators/tensorflow/tensorflow_session_from_frozen_graph_generator_test.cc
+1
View File
@@ -66,6 +66,7 @@ class TensorFlowSessionFromFrozenGraphGeneratorTest : public ::testing::Test {
(*generator_options_->mutable_tag_to_tensor_names())["B"] = "b:0";
generator_options_->mutable_config()->set_intra_op_parallelism_threads(1);
generator_options_->mutable_config()->set_inter_op_parallelism_threads(2);
generator_options_->set_preferred_device_id("/device:CPU:0");
}
void VerifySignatureMap(PacketSet* output_side_packets) {
mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_calculator.cc
+1 -1
View File
@@ -134,8 +134,8 @@ class TensorFlowSessionFromSavedModelCalculator : public CalculatorBase {
}
tensorflow::RunOptions run_options;
// In the future, could construct session options from the options proto.
tensorflow::SessionOptions session_options;
session_options.config = options.session_config();
auto saved_model = absl::make_unique<tensorflow::SavedModelBundle>();
::tensorflow::Status status = tensorflow::LoadSavedModel(
session_options, run_options, path, tags_set, saved_model.get());
mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_calculator.proto
+4
View File
@@ -17,6 +17,7 @@ syntax = "proto2";
package mediapipe;
import "mediapipe/framework/calculator.proto";
import "tensorflow/core/protobuf/config.proto";
message TensorFlowSessionFromSavedModelCalculatorOptions {
extend mediapipe.CalculatorOptions {
@@ -55,4 +56,7 @@ message TensorFlowSessionFromSavedModelCalculatorOptions {
// If no tag is specified, then use "serve" as the default. Note that in order
// to use TPU accelerator hardware, the tag "tpu" needs to be specified.
repeated string saved_model_tag = 6;
// Tensorflow session config options.
optional tensorflow.ConfigProto session_config = 7;
}
mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_calculator_test.cc
+27
View File
@@ -26,6 +26,7 @@
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/framework/tool/tag_map_helper.h"
#include "mediapipe/framework/tool/validate_type.h"
#include "tensorflow/core/framework/device_attributes.pb.h"
namespace mediapipe {
@@ -204,5 +205,31 @@ TEST_F(TensorFlowSessionFromSavedModelCalculatorTest,
ASSERT_NE(session.session, nullptr);
}
TEST_F(TensorFlowSessionFromSavedModelCalculatorTest,
ConfiguresSessionGivenConfig) {
options_->set_saved_model_path(
std::string(file::SplitPath(GetSavedModelDir()).first));
options_->set_load_latest_model(true);
options_->mutable_session_config()->mutable_device_count()->insert(
{"CPU", 10});
CalculatorRunner runner(absl::Substitute(R"(
calculator: "TensorFlowSessionFromSavedModelCalculator"
output_side_packet: "SESSION:tf_model"
options {
[mediapipe.TensorFlowSessionFromSavedModelCalculatorOptions.ext]: {
$0
}
})",
options_->DebugString()));
MP_ASSERT_OK(runner.Run());
const TensorFlowSession& session =
runner.OutputSidePackets().Tag("SESSION").Get<TensorFlowSession>();
// Session must be set.
ASSERT_NE(session.session, nullptr);
std::vector<tensorflow::DeviceAttributes> devices;
ASSERT_EQ(session.session->ListDevices(&devices), tensorflow::Status::OK());
EXPECT_THAT(devices.size(), 10);
}
} // namespace
} // namespace mediapipe
mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_generator.cc
+1 -1
View File
@@ -129,8 +129,8 @@ class TensorFlowSessionFromSavedModelGenerator : public PacketGenerator {
}
tensorflow::RunOptions run_options;
// In the future, could construct session options from the options proto.
tensorflow::SessionOptions session_options;
session_options.config = options.session_config();
auto saved_model = absl::make_unique<tensorflow::SavedModelBundle>();
::tensorflow::Status status = tensorflow::LoadSavedModel(
session_options, run_options, path, tags_set, saved_model.get());
mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_generator.proto
+4
View File
@@ -17,6 +17,7 @@ syntax = "proto2";
package mediapipe;
import "mediapipe/framework/packet_generator.proto";
import "tensorflow/core/protobuf/config.proto";
message TensorFlowSessionFromSavedModelGeneratorOptions {
extend mediapipe.PacketGeneratorOptions {
@@ -55,4 +56,7 @@ message TensorFlowSessionFromSavedModelGeneratorOptions {
// If no tag is specified, then use "serve" as the default. Note that in order
// to use TPU accelerator hardware, the tag "tpu" needs to be specified.
repeated string saved_model_tag = 6;
// Tensorflow session config options.
optional tensorflow.ConfigProto session_config = 9;
}
mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_generator_test.cc
+25
View File
@@ -25,6 +25,7 @@
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/framework/tool/tag_map_helper.h"
#include "mediapipe/framework/tool/validate_type.h"
#include "tensorflow/core/framework/device_attributes.pb.h"
namespace mediapipe {
@@ -196,5 +197,29 @@ TEST_F(TensorFlowSessionFromSavedModelGeneratorTest,
ASSERT_NE(session.session, nullptr);
}
TEST_F(TensorFlowSessionFromSavedModelGeneratorTest,
ConfiguresSessionGivenConfig) {
generator_options_->set_saved_model_path(
std::string(file::SplitPath(GetSavedModelDir()).first));
generator_options_->set_load_latest_model(true);
generator_options_->mutable_session_config()->mutable_device_count()->insert(
{"CPU", 10});
PacketSet input_side_packets(tool::CreateTagMap({}).ValueOrDie());
PacketSet output_side_packets(
tool::CreateTagMap({"SESSION:session"}).ValueOrDie());
::mediapipe::Status run_status = tool::RunGenerateAndValidateTypes(
"TensorFlowSessionFromSavedModelGenerator", extendable_options_,
input_side_packets, &output_side_packets);
MP_EXPECT_OK(run_status) << run_status.message();
const TensorFlowSession& session =
output_side_packets.Tag("SESSION").Get<TensorFlowSession>();
// Session must be set.
ASSERT_NE(session.session, nullptr);
std::vector<tensorflow::DeviceAttributes> devices;
ASSERT_EQ(session.session->ListDevices(&devices), tensorflow::Status::OK());
EXPECT_THAT(devices.size(), 10);
}
} // namespace
} // namespace mediapipe
mediapipe/calculators/tensorflow/vector_float_to_tensor_calculator.cc
+4 -4
View File
@@ -91,11 +91,11 @@ REGISTER_CALCULATOR(VectorFloatToTensorCalculator);
cc->Inputs().Index(0).Value().Get<std::vector<std::vector<float>>>();
const int32 rows = input.size();
CHECK_GE(rows, 1);
RET_CHECK_GE(rows, 1);
const int32 cols = input[0].size();
CHECK_GE(cols, 1);
RET_CHECK_GE(cols, 1);
for (int i = 1; i < rows; ++i) {
CHECK_EQ(input[i].size(), cols);
RET_CHECK_EQ(input[i].size(), cols);
}
if (options_.transpose()) {
tensor_shape = tf::TensorShape({cols, rows});
@@ -116,7 +116,7 @@ REGISTER_CALCULATOR(VectorFloatToTensorCalculator);
} else if (options_.input_size() == INPUT_1D) {
const std::vector<float>& input =
cc->Inputs().Index(0).Value().Get<std::vector<float>>();
CHECK_GE(input.size(), 1);
RET_CHECK_GE(input.size(), 1);
const int32 length = input.size();
tensor_shape = tf::TensorShape({length});
auto output = ::absl::make_unique<tf::Tensor>(tf::DT_FLOAT, tensor_shape);
mediapipe/calculators/tflite/BUILD
+1 -12
View File
@@ -196,13 +196,6 @@ cc_test(
],
)
cc_library(
name = "util",
hdrs = ["util.h"],
visibility = ["//visibility:public"],
alwayslink = 1,
)
selects.config_setting_group(
name = "gpu_inference_disabled",
match_any = [
@@ -229,7 +222,6 @@ cc_library(
}),
visibility = ["//visibility:public"],
deps = [
":util",
":tflite_inference_calculator_cc_proto",
"@com_google_absl//absl/memory",
"//mediapipe/framework:calculator_framework",
@@ -295,7 +287,6 @@ cc_library(
visibility = ["//visibility:public"],
deps = [
"//mediapipe/util/tflite:config",
":util",
":tflite_converter_calculator_cc_proto",
"//mediapipe/util:resource_util",
"//mediapipe/framework:calculator_framework",
@@ -334,7 +325,6 @@ cc_library(
srcs = ["tflite_model_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
":util",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:packet",
"//mediapipe/framework/port:ret_check",
@@ -348,7 +338,6 @@ cc_library(
srcs = ["tflite_tensors_to_segmentation_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
":util",
":tflite_tensors_to_segmentation_calculator_cc_proto",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
@@ -418,7 +407,6 @@ cc_library(
visibility = ["//visibility:public"],
deps = [
"//mediapipe/util/tflite:config",
":util",
":tflite_tensors_to_detections_calculator_cc_proto",
"//mediapipe/framework/formats:detection_cc_proto",
"@com_google_absl//absl/strings:str_format",
@@ -551,6 +539,7 @@ cc_test(
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/tool:validate_type",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@org_tensorflow//tensorflow/lite:framework",
],
)
mediapipe/calculators/tflite/tflite_converter_calculator.cc
+78 -52
View File
@@ -16,7 +16,6 @@
#include <vector>
#include "mediapipe/calculators/tflite/tflite_converter_calculator.pb.h"
#include "mediapipe/calculators/tflite/util.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/matrix.h"
@@ -146,8 +145,7 @@ class TfLiteConverterCalculator : public CalculatorBase {
::mediapipe::Status LoadOptions(CalculatorContext* cc);
template <class T>
::mediapipe::Status NormalizeImage(const ImageFrame& image_frame,
bool zero_center, bool flip_vertically,
float* tensor_ptr);
bool flip_vertically, float* tensor_ptr);
::mediapipe::Status CopyMatrixToTensor(const Matrix& matrix,
float* tensor_ptr);
::mediapipe::Status ProcessCPU(CalculatorContext* cc);
@@ -165,10 +163,7 @@ class TfLiteConverterCalculator : public CalculatorBase {
bool initialized_ = false;
bool use_gpu_ = false;
bool zero_center_ = true; // normalize range to [-1,1] | otherwise [0,1]
bool use_custom_normalization_ = false;
float custom_div_ = -1.0f;
float custom_sub_ = -1.0f;
absl::optional<std::pair<float, float>> output_range_;
bool flip_vertically_ = false;
bool row_major_matrix_ = false;
bool use_quantized_tensors_ = false;
@@ -362,11 +357,11 @@ bool ShouldUseGpu(CC* cc) {
float* tensor_buffer = tensor->data.f;
RET_CHECK(tensor_buffer);
if (image_frame.ByteDepth() == 1) {
MP_RETURN_IF_ERROR(NormalizeImage<uint8>(
image_frame, zero_center_, flip_vertically_, tensor_buffer));
MP_RETURN_IF_ERROR(NormalizeImage<uint8>(image_frame, flip_vertically_,
tensor_buffer));
} else if (image_frame.ByteDepth() == 4) {
MP_RETURN_IF_ERROR(NormalizeImage<float>(
image_frame, zero_center_, flip_vertically_, tensor_buffer));
MP_RETURN_IF_ERROR(NormalizeImage<float>(image_frame, flip_vertically_,
tensor_buffer));
} else {
return ::mediapipe::InternalError(
"Only byte-based (8 bit) and float (32 bit) images supported.");
@@ -427,11 +422,11 @@ bool ShouldUseGpu(CC* cc) {
auto src = gpu_helper_.CreateSourceTexture(input);
glActiveTexture(GL_TEXTURE0 + 0);
glBindTexture(GL_TEXTURE_2D, src.name());
RET_CHECK_CALL(gpu_data_out_->buffer.BindToIndex(1));
MP_RETURN_IF_ERROR(gpu_data_out_->buffer.BindToIndex(1));
const tflite::gpu::uint3 workgroups = {
NumGroups(input.width(), kWorkgroupSize),
NumGroups(input.height(), kWorkgroupSize), 1};
RET_CHECK_CALL(gpu_data_out_->program.Dispatch(workgroups));
MP_RETURN_IF_ERROR(gpu_data_out_->program.Dispatch(workgroups));
glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
glBindTexture(GL_TEXTURE_2D, 0);
src.Release();
@@ -445,9 +440,9 @@ bool ShouldUseGpu(CC* cc) {
output_tensors->resize(1);
{
GpuTensor& tensor = output_tensors->at(0);
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
MP_RETURN_IF_ERROR(CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_->elements, &tensor));
RET_CHECK_CALL(CopyBuffer(gpu_data_out_->buffer, tensor));
MP_RETURN_IF_ERROR(CopyBuffer(gpu_data_out_->buffer, tensor));
}
return ::mediapipe::OkStatus();
}));
@@ -521,7 +516,7 @@ bool ShouldUseGpu(CC* cc) {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &include_alpha, &input, &single_channel]() -> ::mediapipe::Status {
// Device memory.
RET_CHECK_CALL(
MP_RETURN_IF_ERROR(
::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_->elements, &gpu_data_out_->buffer));
@@ -544,7 +539,13 @@ bool ShouldUseGpu(CC* cc) {
$6 // alpha channel
})",
/*$0=*/kWorkgroupSize, /*$1=*/input.width(), /*$2=*/input.height(),
/*$3=*/zero_center_ ? "pixel = (pixel - 0.5) * 2.0;" : "",
/*$3=*/
output_range_.has_value()
? absl::Substitute(
"pixel = pixel * float($0) + float($1);",
(output_range_->second - output_range_->first),
output_range_->first)
: "",
/*$4=*/flip_vertically_ ? "(width_height.y - 1 - gid.y)" : "gid.y",
/*$5=*/
single_channel
@@ -555,10 +556,10 @@ bool ShouldUseGpu(CC* cc) {
include_alpha ? "output_data.elements[linear_index + 3] = pixel.w;"
: "",
/*$7=*/max_num_channels_);
RET_CHECK_CALL(GlShader::CompileShader(GL_COMPUTE_SHADER, shader_source,
&gpu_data_out_->shader));
RET_CHECK_CALL(GlProgram::CreateWithShader(gpu_data_out_->shader,
&gpu_data_out_->program));
MP_RETURN_IF_ERROR(GlShader::CompileShader(
GL_COMPUTE_SHADER, shader_source, &gpu_data_out_->shader));
MP_RETURN_IF_ERROR(GlProgram::CreateWithShader(
gpu_data_out_->shader, &gpu_data_out_->program));
return ::mediapipe::OkStatus();
}));
@@ -599,7 +600,12 @@ bool ShouldUseGpu(CC* cc) {
)",
/*$0=*/include_alpha ? "float4" : "float3",
/*$1=*/include_alpha ? "rgba" : "rgb",
/*$2=*/zero_center_ ? "pixel = (pixel - 0.5) * 2.0;" : "",
/*$2=*/
output_range_.has_value()
? absl::Substitute("pixel = pixel * float($0) + float($1);",
(output_range_->second - output_range_->first),
output_range_->first)
: "",
/*$3=*/flip_vertically_ ? "(in_tex.get_height() - 1 - gid.y)" : "gid.y",
/*$4=*/include_alpha ? 4 : 3,
/*$5=*/include_alpha ? "out_buf[linear_index + 3] = pixel.w;" : "");
@@ -630,13 +636,27 @@ bool ShouldUseGpu(CC* cc) {
const auto& options =
cc->Options<::mediapipe::TfLiteConverterCalculatorOptions>();
// Get data normalization mode.
zero_center_ = options.zero_center();
// if zero_center, set output float range to match [-1, 1] as specified in
// calculator proto.
if (options.zero_center()) {
output_range_.emplace(std::pair<float, float>(-1.0, 1.0));
}
// Custom output_tensor_float_range values.
// If the float range is specified in pb text, use the specified values
// instead.
if (options.has_output_tensor_float_range()) {
output_range_.emplace(options.output_tensor_float_range().min(),
options.output_tensor_float_range().max());
CHECK_GT(output_range_->second, output_range_->first);
}
// Custom div and sub values.
use_custom_normalization_ = options.use_custom_normalization();
custom_div_ = options.custom_div();
custom_sub_ = options.custom_sub();
if (options.use_custom_normalization()) {
output_range_.emplace(std::pair<float, float>(
-options.custom_sub(),
-options.custom_sub() + 255.0 / options.custom_div()));
}
// Get y-flip mode.
flip_vertically_ = options.flip_vertically();
@@ -664,40 +684,46 @@ bool ShouldUseGpu(CC* cc) {
template <class T>
::mediapipe::Status TfLiteConverterCalculator::NormalizeImage(
const ImageFrame& image_frame, bool zero_center, bool flip_vertically,
float* tensor_ptr) {
const ImageFrame& image_frame, bool flip_vertically, float* tensor_ptr) {
const int height = image_frame.Height();
const int width = image_frame.Width();
const int channels = image_frame.NumberOfChannels();
const int channels_preserved = std::min(channels, max_num_channels_);
const int channels_ignored = channels - channels_preserved;
float div, sub;
if (output_range_.has_value()) {
// If the output float range is set and we are not using custom
// normalization, normalize the pixel values from [0, 255] to the specified
// output range.
RET_CHECK_NE(output_range_->first, output_range_->second);
const float scale = (output_range_->second - output_range_->first) / 255.0f;
const float bias = output_range_->first;
if (use_custom_normalization_) {
RET_CHECK_GT(custom_div_, 0.0f);
RET_CHECK_GE(custom_sub_, 0.0f);
div = custom_div_;
sub = custom_sub_;
} else if (zero_center) {
// [-1,1]
div = 127.5f;
sub = 1.0f;
for (int i = 0; i < height; ++i) {
const T* image_ptr = reinterpret_cast<const T*>(
image_frame.PixelData() +
(flip_vertically ? height - 1 - i : i) * image_frame.WidthStep());
for (int j = 0; j < width; ++j) {
for (int c = 0; c < channels_preserved; ++c) {
*tensor_ptr++ = *image_ptr++ * scale + bias;
}
image_ptr += channels_ignored;
}
}
} else {
// [0,1]
div = 255.0f;
sub = 0.0f;
}
for (int i = 0; i < height; ++i) {
const T* image_ptr = reinterpret_cast<const T*>(
image_frame.PixelData() +
(flip_vertically ? height - 1 - i : i) * image_frame.WidthStep());
for (int j = 0; j < width; ++j) {
for (int c = 0; c < channels_preserved; ++c) {
*tensor_ptr++ = *image_ptr++ / div - sub;
// [0,1], scale only (bias == 0)
// Verified that there are no precision issues with 1.0f / 255.0f expression
const float scale = 1.0f / 255.0f;
for (int i = 0; i < height; ++i) {
const T* image_ptr = reinterpret_cast<const T*>(
image_frame.PixelData() +
(flip_vertically ? height - 1 - i : i) * image_frame.WidthStep());
for (int j = 0; j < width; ++j) {
for (int c = 0; c < channels_preserved; ++c) {
*tensor_ptr++ = *image_ptr++ * scale;
}
image_ptr += channels_ignored;
}
image_ptr += channels_ignored;
}
}
mediapipe/calculators/tflite/tflite_converter_calculator.proto
+10
View File
@@ -56,4 +56,14 @@ message TfLiteConverterCalculatorOptions {
// Quantization option (CPU only).
// When true, output kTfLiteUInt8 tensor instead of kTfLiteFloat32.
optional bool use_quantized_tensors = 5 [default = false];
// Normalization option.
// Setting normalization_range results in the values normalized to
// the range [output_tensor_float_range.min, output_tensor_float_range.max].
optional TensorFloatRange output_tensor_float_range = 9;
message TensorFloatRange {
optional float min = 1;
optional float max = 2;
}
}
mediapipe/calculators/tflite/tflite_converter_calculator_test.cc
+68 -1
View File
@@ -16,6 +16,7 @@
#include <vector>
#include "absl/memory/memory.h"
#include "absl/strings/substitute.h"
#include "mediapipe/calculators/tflite/tflite_converter_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_runner.h"
@@ -40,6 +41,7 @@ constexpr char kTransposeOptionsString[] =
} // namespace
using RandomEngine = std::mt19937_64;
using testing::Eq;
const uint32 kSeed = 1234;
const int kNumSizes = 8;
const int sizes[kNumSizes][2] = {{1, 1}, {12, 1}, {1, 9}, {2, 2},
@@ -232,7 +234,6 @@ TEST_F(TfLiteConverterCalculatorTest, CustomDivAndSub) {
// Wait until the calculator done processing.
MP_ASSERT_OK(graph.WaitUntilIdle());
EXPECT_EQ(1, output_packets.size());
// Get and process results.
const std::vector<TfLiteTensor>& tensor_vec =
@@ -249,4 +250,70 @@ TEST_F(TfLiteConverterCalculatorTest, CustomDivAndSub) {
MP_ASSERT_OK(graph.WaitUntilDone());
}
TEST_F(TfLiteConverterCalculatorTest, SetOutputRange) {
std::vector<std::pair<float, float>> range_values = {
std::make_pair(0.0, 1.0), std::make_pair(-1.0, 1.0),
std::make_pair(-0.5, 0.5)};
for (std::pair<float, float> range : range_values) {
CalculatorGraph graph;
CalculatorGraphConfig graph_config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
absl::Substitute(R"(
input_stream: "input_image"
node {
calculator: "TfLiteConverterCalculator"
input_stream: "IMAGE:input_image"
output_stream: "TENSORS:tensor"
options {
[mediapipe.TfLiteConverterCalculatorOptions.ext] {
output_tensor_float_range {
min: $0
max: $1
}
}
}
}
)",
/*$0=*/range.first,
/*$1=*/range.second));
std::vector<Packet> output_packets;
tool::AddVectorSink("tensor", &graph_config, &output_packets);
// Run the graph.
MP_ASSERT_OK(graph.Initialize(graph_config));
MP_ASSERT_OK(graph.StartRun({}));
auto input_image = absl::make_unique<ImageFrame>(ImageFormat::GRAY8, 1, 1);
cv::Mat mat = ::mediapipe::formats::MatView(input_image.get());
mat.at<uint8>(0, 0) = 200;
MP_ASSERT_OK(graph.AddPacketToInputStream(
"input_image", Adopt(input_image.release()).At(Timestamp(0))));
// Wait until the calculator finishes processing.
MP_ASSERT_OK(graph.WaitUntilIdle());
EXPECT_THAT(output_packets.size(), Eq(1));
// Get and process results.
const std::vector<TfLiteTensor>& tensor_vec =
output_packets[0].Get<std::vector<TfLiteTensor>>();
EXPECT_THAT(tensor_vec.size(), Eq(1));
const TfLiteTensor* tensor = &tensor_vec[0];
// Calculate the expected normalized value:
float normalized_value =
range.first + (200 * (range.second - range.first)) / 255.0;
EXPECT_THAT(tensor->type, Eq(kTfLiteFloat32));
EXPECT_THAT(normalized_value,
testing::FloatNear(*tensor->data.f,
2.0f * std::abs(*tensor->data.f) *
std::numeric_limits<float>::epsilon()));
// Fully close graph at end, otherwise calculator+tensors are destroyed
// after calling WaitUntilDone().
MP_ASSERT_OK(graph.CloseInputStream("input_image"));
MP_ASSERT_OK(graph.WaitUntilDone());
}
}
} // namespace mediapipe
mediapipe/calculators/tflite/tflite_inference_calculator.cc
+13 -12
View File
@@ -19,7 +19,6 @@
#include "absl/memory/memory.h"
#include "mediapipe/calculators/tflite/tflite_inference_calculator.pb.h"
#include "mediapipe/calculators/tflite/util.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/util/tflite/config.h"
@@ -496,7 +495,7 @@ bool ShouldUseGpu(CC* cc) {
output_tensors_gpu->resize(gpu_data_out_.size());
for (int i = 0; i < gpu_data_out_.size(); ++i) {
GpuTensor& tensor = output_tensors_gpu->at(i);
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
MP_RETURN_IF_ERROR(CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_[i]->elements, &tensor));
MP_RETURN_IF_ERROR(
tflite_gpu_runner_->BindSSBOToOutputTensor(tensor.id(), i));
@@ -518,7 +517,7 @@ bool ShouldUseGpu(CC* cc) {
// Explicit copy input.
gpu_data_in_.resize(input_tensors.size());
for (int i = 0; i < input_tensors.size(); ++i) {
RET_CHECK_CALL(CopyBuffer(input_tensors[i], gpu_data_in_[i]->buffer));
MP_RETURN_IF_ERROR(CopyBuffer(input_tensors[i], gpu_data_in_[i]->buffer));
}
#elif MEDIAPIPE_TFLITE_METAL_INFERENCE
const auto& input_tensors =
@@ -582,7 +581,7 @@ bool ShouldUseGpu(CC* cc) {
for (int i = 0; i < tensor_indexes.size(); ++i) {
TfLiteTensor* tensor = interpreter_->tensor(tensor_indexes[i]);
std::vector<float> gpu_data(tensor->bytes / sizeof(float));
RET_CHECK_CALL(gpu_data_out_[i]->buffer.Read(
MP_RETURN_IF_ERROR(gpu_data_out_[i]->buffer.Read(
absl::MakeSpan(tensor->data.f, tensor->bytes)));
output_tensors_cpu->emplace_back(*tensor);
}
@@ -599,9 +598,9 @@ bool ShouldUseGpu(CC* cc) {
for (int i = 0; i < gpu_data_out_.size(); ++i) {
GpuTensor& tensor = output_tensors_gpu->at(i);
// Allocate output tensor.
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
MP_RETURN_IF_ERROR(CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_[i]->elements, &tensor));
RET_CHECK_CALL(CopyBuffer(gpu_data_out_[i]->buffer, tensor));
MP_RETURN_IF_ERROR(CopyBuffer(gpu_data_out_[i]->buffer, tensor));
}
cc->Outputs()
.Tag(kTensorsGpuTag)
@@ -655,7 +654,8 @@ bool ShouldUseGpu(CC* cc) {
options.priority3 = tflite::gpu::InferencePriority::AUTO;
options.usage = tflite::gpu::InferenceUsage::SUSTAINED_SPEED;
tflite_gpu_runner_ = std::make_unique<tflite::gpu::TFLiteGPURunner>(options);
RET_CHECK_CALL(tflite_gpu_runner_->InitializeWithModel(model, op_resolver));
MP_RETURN_IF_ERROR(
tflite_gpu_runner_->InitializeWithModel(model, op_resolver));
// Allocate interpreter memory for cpu output.
if (!gpu_output_) {
@@ -688,10 +688,11 @@ bool ShouldUseGpu(CC* cc) {
ASSIGN_OR_RETURN(gpu_data_out_[i]->elements,
tflite_gpu_runner_->GetOutputElements(i));
// Create and bind input buffer.
RET_CHECK_CALL(::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_[i]->elements, &gpu_data_out_[i]->buffer));
MP_RETURN_IF_ERROR(
::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_[i]->elements, &gpu_data_out_[i]->buffer));
}
RET_CHECK_CALL(tflite_gpu_runner_->Build());
MP_RETURN_IF_ERROR(tflite_gpu_runner_->Build());
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
return ::mediapipe::OkStatus();
@@ -841,7 +842,7 @@ bool ShouldUseGpu(CC* cc) {
gpu_data_in_[i]->elements *= tensor->dims->data[d];
}
// Create and bind input buffer.
RET_CHECK_CALL(
MP_RETURN_IF_ERROR(
::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
gpu_data_in_[i]->elements, &gpu_data_in_[i]->buffer));
RET_CHECK_EQ(TfLiteGpuDelegateBindBufferToTensor(
@@ -866,7 +867,7 @@ bool ShouldUseGpu(CC* cc) {
// Create and bind output buffers.
interpreter_->SetAllowBufferHandleOutput(true);
for (int i = 0; i < gpu_data_out_.size(); ++i) {
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
MP_RETURN_IF_ERROR(CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_[i]->elements, &gpu_data_out_[i]->buffer));
RET_CHECK_EQ(TfLiteGpuDelegateBindBufferToTensor(
delegate_.get(), gpu_data_out_[i]->buffer.id(),
mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.cc
+26 -24
View File
@@ -18,7 +18,6 @@
#include "absl/strings/str_format.h"
#include "absl/types/span.h"
#include "mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.pb.h"
#include "mediapipe/calculators/tflite/util.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/deps/file_path.h"
#include "mediapipe/framework/formats/detection.pb.h"
@@ -404,8 +403,10 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
&output_detections]()
-> ::mediapipe::Status {
// Copy inputs.
RET_CHECK_CALL(CopyBuffer(input_tensors[0], gpu_data_->raw_boxes_buffer));
RET_CHECK_CALL(CopyBuffer(input_tensors[1], gpu_data_->raw_scores_buffer));
MP_RETURN_IF_ERROR(
CopyBuffer(input_tensors[0], gpu_data_->raw_boxes_buffer));
MP_RETURN_IF_ERROR(
CopyBuffer(input_tensors[1], gpu_data_->raw_scores_buffer));
if (!anchors_init_) {
if (side_packet_anchors_) {
CHECK(!cc->InputSidePackets().Tag("ANCHORS").IsEmpty());
@@ -413,11 +414,11 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
cc->InputSidePackets().Tag("ANCHORS").Get<std::vector<Anchor>>();
std::vector<float> raw_anchors(num_boxes_ * kNumCoordsPerBox);
ConvertAnchorsToRawValues(anchors, num_boxes_, raw_anchors.data());
RET_CHECK_CALL(gpu_data_->raw_anchors_buffer.Write<float>(
MP_RETURN_IF_ERROR(gpu_data_->raw_anchors_buffer.Write<float>(
absl::MakeSpan(raw_anchors)));
} else {
CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
RET_CHECK_CALL(
MP_RETURN_IF_ERROR(
CopyBuffer(input_tensors[2], gpu_data_->raw_anchors_buffer));
}
anchors_init_ = true;
@@ -425,23 +426,24 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
// Run shaders.
// Decode boxes.
RET_CHECK_CALL(gpu_data_->decoded_boxes_buffer.BindToIndex(0));
RET_CHECK_CALL(gpu_data_->raw_boxes_buffer.BindToIndex(1));
RET_CHECK_CALL(gpu_data_->raw_anchors_buffer.BindToIndex(2));
MP_RETURN_IF_ERROR(gpu_data_->decoded_boxes_buffer.BindToIndex(0));
MP_RETURN_IF_ERROR(gpu_data_->raw_boxes_buffer.BindToIndex(1));
MP_RETURN_IF_ERROR(gpu_data_->raw_anchors_buffer.BindToIndex(2));
const tflite::gpu::uint3 decode_workgroups = {num_boxes_, 1, 1};
RET_CHECK_CALL(gpu_data_->decode_program.Dispatch(decode_workgroups));
MP_RETURN_IF_ERROR(gpu_data_->decode_program.Dispatch(decode_workgroups));
// Score boxes.
RET_CHECK_CALL(gpu_data_->scored_boxes_buffer.BindToIndex(0));
RET_CHECK_CALL(gpu_data_->raw_scores_buffer.BindToIndex(1));
MP_RETURN_IF_ERROR(gpu_data_->scored_boxes_buffer.BindToIndex(0));
MP_RETURN_IF_ERROR(gpu_data_->raw_scores_buffer.BindToIndex(1));
const tflite::gpu::uint3 score_workgroups = {num_boxes_, 1, 1};
RET_CHECK_CALL(gpu_data_->score_program.Dispatch(score_workgroups));
MP_RETURN_IF_ERROR(gpu_data_->score_program.Dispatch(score_workgroups));
// Copy decoded boxes from GPU to CPU.
std::vector<float> boxes(num_boxes_ * num_coords_);
RET_CHECK_CALL(gpu_data_->decoded_boxes_buffer.Read(absl::MakeSpan(boxes)));
MP_RETURN_IF_ERROR(
gpu_data_->decoded_boxes_buffer.Read(absl::MakeSpan(boxes)));
std::vector<float> score_class_id_pairs(num_boxes_ * 2);
RET_CHECK_CALL(gpu_data_->scored_boxes_buffer.Read(
MP_RETURN_IF_ERROR(gpu_data_->scored_boxes_buffer.Read(
absl::MakeSpan(score_class_id_pairs)));
// TODO: b/138851969. Is it possible to output a float vector
@@ -802,20 +804,20 @@ void main() {
// Shader program
GlShader decode_shader;
RET_CHECK_CALL(
MP_RETURN_IF_ERROR(
GlShader::CompileShader(GL_COMPUTE_SHADER, decode_src, &decode_shader));
RET_CHECK_CALL(GpuProgram::CreateWithShader(decode_shader,
&gpu_data_->decode_program));
MP_RETURN_IF_ERROR(GpuProgram::CreateWithShader(
decode_shader, &gpu_data_->decode_program));
// Outputs
size_t decoded_boxes_length = num_boxes_ * num_coords_;
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
MP_RETURN_IF_ERROR(CreateReadWriteShaderStorageBuffer<float>(
decoded_boxes_length, &gpu_data_->decoded_boxes_buffer));
// Inputs
size_t raw_boxes_length = num_boxes_ * num_coords_;
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
MP_RETURN_IF_ERROR(CreateReadWriteShaderStorageBuffer<float>(
raw_boxes_length, &gpu_data_->raw_boxes_buffer));
size_t raw_anchors_length = num_boxes_ * kNumCoordsPerBox;
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
MP_RETURN_IF_ERROR(CreateReadWriteShaderStorageBuffer<float>(
raw_anchors_length, &gpu_data_->raw_anchors_buffer));
// Parameters
glUseProgram(gpu_data_->decode_program.id());
@@ -896,17 +898,17 @@ void main() {
// Shader program
GlShader score_shader;
RET_CHECK_CALL(
MP_RETURN_IF_ERROR(
GlShader::CompileShader(GL_COMPUTE_SHADER, score_src, &score_shader));
RET_CHECK_CALL(
MP_RETURN_IF_ERROR(
GpuProgram::CreateWithShader(score_shader, &gpu_data_->score_program));
// Outputs
size_t scored_boxes_length = num_boxes_ * 2; // score, class
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
MP_RETURN_IF_ERROR(CreateReadWriteShaderStorageBuffer<float>(
scored_boxes_length, &gpu_data_->scored_boxes_buffer));
// Inputs
size_t raw_scores_length = num_boxes_ * num_classes_;
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
MP_RETURN_IF_ERROR(CreateReadWriteShaderStorageBuffer<float>(
raw_scores_length, &gpu_data_->raw_scores_buffer));
return ::mediapipe::OkStatus();
mediapipe/calculators/tflite/tflite_tensors_to_segmentation_calculator.cc
+12 -13
View File
@@ -17,7 +17,6 @@
#include "absl/strings/str_format.h"
#include "absl/types/span.h"
#include "mediapipe/calculators/tflite/tflite_tensors_to_segmentation_calculator.pb.h"
#include "mediapipe/calculators/tflite/util.h"
#include "mediapipe/framework/calculator_context.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_frame.h"
@@ -400,7 +399,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
// Create initial working mask texture.
::tflite::gpu::gl::GlTexture small_mask_texture;
RET_CHECK_CALL(CreateReadWriteRgbaImageTexture(
MP_RETURN_IF_ERROR(CreateReadWriteRgbaImageTexture(
tflite::gpu::DataType::UINT8, // GL_RGBA8
{tensor_width_, tensor_height_}, &small_mask_texture));
@@ -410,7 +409,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
: mediapipe::GlTexture();
// Copy input tensor.
RET_CHECK_CALL(CopyBuffer(input_tensors[0], *tensor_buffer_));
MP_RETURN_IF_ERROR(CopyBuffer(input_tensors[0], *tensor_buffer_));
// Run shader, process mask tensor.
// Run softmax over tensor output and blend with previous mask.
@@ -418,18 +417,18 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
const int output_index = 0;
glBindImageTexture(output_index, small_mask_texture.id(), 0, GL_FALSE, 0,
GL_WRITE_ONLY, GL_RGBA8);
RET_CHECK_CALL(tensor_buffer_->BindToIndex(2));
MP_RETURN_IF_ERROR(tensor_buffer_->BindToIndex(2));
const tflite::gpu::uint3 workgroups = {
NumGroups(tensor_width_, kWorkgroupSize),
NumGroups(tensor_height_, kWorkgroupSize), 1};
if (!has_prev_mask) {
RET_CHECK_CALL(mask_program_no_prev_->Dispatch(workgroups));
MP_RETURN_IF_ERROR(mask_program_no_prev_->Dispatch(workgroups));
} else {
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D, input_mask_texture.name());
RET_CHECK_CALL(mask_program_with_prev_->Dispatch(workgroups));
MP_RETURN_IF_ERROR(mask_program_with_prev_->Dispatch(workgroups));
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D, 0);
}
@@ -622,22 +621,22 @@ void main() {
// Shader programs.
GlShader shader_without_previous;
RET_CHECK_CALL(GlShader::CompileShader(
MP_RETURN_IF_ERROR(GlShader::CompileShader(
GL_COMPUTE_SHADER, shader_src_no_previous, &shader_without_previous));
mask_program_no_prev_ = absl::make_unique<GlProgram>();
RET_CHECK_CALL(GlProgram::CreateWithShader(shader_without_previous,
mask_program_no_prev_.get()));
MP_RETURN_IF_ERROR(GlProgram::CreateWithShader(
shader_without_previous, mask_program_no_prev_.get()));
GlShader shader_with_previous;
RET_CHECK_CALL(GlShader::CompileShader(
MP_RETURN_IF_ERROR(GlShader::CompileShader(
GL_COMPUTE_SHADER, shader_src_with_previous, &shader_with_previous));
mask_program_with_prev_ = absl::make_unique<GlProgram>();
RET_CHECK_CALL(GlProgram::CreateWithShader(shader_with_previous,
mask_program_with_prev_.get()));
MP_RETURN_IF_ERROR(GlProgram::CreateWithShader(
shader_with_previous, mask_program_with_prev_.get()));
// Buffer storage for input tensor.
size_t tensor_length = tensor_width_ * tensor_height_ * tensor_channels_;
tensor_buffer_ = absl::make_unique<GlBuffer>();
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
MP_RETURN_IF_ERROR(CreateReadWriteShaderStorageBuffer<float>(
tensor_length, tensor_buffer_.get()));
// Parameters.
mediapipe/calculators/tflite/util.h
-25
View File
@@ -1,25 +0,0 @@
// 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.
#ifndef MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
#define MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
#define RET_CHECK_CALL(call) \
do { \
const auto status = (call); \
if (ABSL_PREDICT_FALSE(!status.ok())) \
return ::mediapipe::InternalError(status.message()); \
} while (0);
#endif // MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
mediapipe/calculators/util/BUILD
+2
View File
@@ -700,6 +700,8 @@ mediapipe_cc_proto_library(
deps = [":rect_to_render_data_calculator_proto"],
)
# TODO: What is that one for?
mediapipe_cc_proto_library(
name = "detections_to_render_data_calculator_cc_proto",
srcs = ["detections_to_render_data_calculator.proto"],
mediapipe/calculators/util/annotation_overlay_calculator.cc
+24 -12
View File
@@ -160,6 +160,8 @@ class AnnotationOverlayCalculator : public CalculatorBase {
GLuint image_mat_tex_ = 0; // Overlay drawing image for GPU.
int width_ = 0;
int height_ = 0;
int width_gpu_ = 0; // Size of overlay drawing texture.
int height_gpu_ = 0;
#endif // MEDIAPIPE_DISABLE_GPU
};
REGISTER_CALCULATOR(AnnotationOverlayCalculator);
@@ -389,7 +391,7 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
glBindTexture(GL_TEXTURE_2D, image_mat_tex_);
glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, width_, height_, GL_RGB,
glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, width_gpu_, height_gpu_, GL_RGB,
GL_UNSIGNED_BYTE, overlay_image);
glBindTexture(GL_TEXTURE_2D, 0);
}
@@ -492,12 +494,12 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
if (format != mediapipe::ImageFormat::SRGBA &&
format != mediapipe::ImageFormat::SRGB)
RET_CHECK_FAIL() << "Unsupported GPU input format: " << format;
image_mat = absl::make_unique<cv::Mat>(height_, width_, CV_8UC3);
image_mat = absl::make_unique<cv::Mat>(height_gpu_, width_gpu_, CV_8UC3);
memset(image_mat->data, kAnnotationBackgroundColor,
height_ * width_ * image_mat->elemSize());
height_gpu_ * width_gpu_ * image_mat->elemSize());
} else {
image_mat = absl::make_unique<cv::Mat>(
options_.canvas_height_px(), options_.canvas_width_px(), CV_8UC3,
height_gpu_, width_gpu_, CV_8UC3,
cv::Scalar(options_.canvas_color().r(), options_.canvas_color().g(),
options_.canvas_color().b()));
}
@@ -632,18 +634,28 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
kAnnotationBackgroundColor / 255.0,
kAnnotationBackgroundColor / 255.0);
// Init texture for opencv rendered frame.
const auto& input_frame =
cc->Inputs().Tag(kInputFrameTagGpu).Get<mediapipe::GpuBuffer>();
// Ensure GPU texture is divisible by 4. See b/138751944 for more info.
width_ =
RoundUp(input_frame.width(), ImageFrame::kGlDefaultAlignmentBoundary);
height_ =
RoundUp(input_frame.height(), ImageFrame::kGlDefaultAlignmentBoundary);
const float alignment = ImageFrame::kGlDefaultAlignmentBoundary;
const float scale_factor = options_.gpu_scale_factor();
if (image_frame_available_) {
const auto& input_frame =
cc->Inputs().Tag(kInputFrameTagGpu).Get<mediapipe::GpuBuffer>();
width_ = RoundUp(input_frame.width(), alignment);
height_ = RoundUp(input_frame.height(), alignment);
} else {
width_ = RoundUp(options_.canvas_width_px(), alignment);
height_ = RoundUp(options_.canvas_height_px(), alignment);
}
width_gpu_ = RoundUp(width_ * scale_factor, alignment);
height_gpu_ = RoundUp(height_ * scale_factor, alignment);
// Init texture for opencv rendered frame.
{
glGenTextures(1, &image_mat_tex_);
glBindTexture(GL_TEXTURE_2D, image_mat_tex_);
glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB8, width_, height_, 0, GL_RGB,
// TODO
// OpenCV only renders to RGB images, not RGBA. Ideally this should be RGBA.
glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB8, width_gpu_, height_gpu_, 0, GL_RGB,
GL_UNSIGNED_BYTE, nullptr);
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
mediapipe/calculators/util/annotation_overlay_calculator.proto
+8
View File
@@ -45,4 +45,12 @@ message AnnotationOverlayCalculatorOptions {
// origin. (Historically, OpenGL uses bottom left origin, but most MediaPipe
// examples expect textures to have top-left origin.)
optional bool gpu_uses_top_left_origin = 6 [default = true];
// Scale factor for intermediate image for GPU rendering.
// This can be used to speed up annotation by drawing the annotation on an
// intermediate image with a reduced scale, e.g. 0.5 (of the input image width
// and height), before resizing and overlaying it on top of the input image.
// Should only be used if *all* render data uses normalized coordinates
// (or absolute coordinates are updated to scale accordingly).
optional float gpu_scale_factor = 7 [default = 1.0];
}
mediapipe/calculators/util/detections_to_render_data_calculator.cc
+4 -3
View File
@@ -235,9 +235,10 @@ void DetectionsToRenderDataCalculator::AddLabels(
const Detection& detection,
const DetectionsToRenderDataCalculatorOptions& options,
float text_line_height, RenderData* render_data) {
CHECK(detection.label().empty() || detection.label_id().empty())
<< "Either std::string or integer labels must be used for detection "
"but not both at the same time.";
CHECK(detection.label().empty() || detection.label_id().empty() ||
detection.label_size() == detection.label_id_size())
<< "String or integer labels should be of same size. Or only one of them "
"is present.";
const auto num_labels =
std::max(detection.label_size(), detection.label_id_size());
CHECK_EQ(detection.score_size(), num_labels)
mediapipe/calculators/video/BUILD
+1
View File
@@ -316,6 +316,7 @@ cc_library(
"//mediapipe/util/tracking",
"//mediapipe/util/tracking:box_tracker",
"//mediapipe/util/tracking:tracking_visualization_utilities",
"@com_google_absl//absl/container:node_hash_set",
"@com_google_absl//absl/strings",
],
alwayslink = 1,
mediapipe/calculators/video/box_tracker_calculator.cc
+3 -2
View File
@@ -18,6 +18,7 @@
#include <unordered_map>
#include <unordered_set>
#include "absl/container/node_hash_set.h"
#include "absl/strings/numbers.h"
#include "mediapipe/calculators/video/box_tracker_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
@@ -193,12 +194,12 @@ class BoxTrackerCalculator : public CalculatorBase {
TimedBoxProtoList initial_pos_;
// Keeps tracks boxes that have already been initialized.
std::unordered_set<int> initialized_ids_;
absl::node_hash_set<int> initialized_ids_;
// Non empty for batch mode tracking.
std::string cache_dir_;
// Ids to be tracked in batch_mode.
std::unordered_set<int> batch_track_ids_;
absl::node_hash_set<int> batch_track_ids_;
int frame_num_ = 0;