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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 8e1da4611d93ccb7d9674713157d43be0348d98f
This commit is contained in:
MediaPipe Team
2021-07-27 22:36:23 -04:00
committed by chuoling
parent 50c92c6623
commit b899d17f18
79 changed files with 1808 additions and 946 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/calculator_graph_stopping_test.cc
+12 -8
View File
@@ -44,28 +44,32 @@ using mediapipe::MakePacket;
using mediapipe::OutputStreamShardSet;
using mediapipe::Timestamp;
namespace proto_ns = mediapipe::proto_ns;
constexpr char kEventTag[] = "EVENT";
constexpr char kOutTag[] = "OUT";
using mediapipe::CalculatorGraph;
using mediapipe::Packet;
class InfiniteSequenceCalculator : public mediapipe::CalculatorBase {
public:
static absl::Status GetContract(mediapipe::CalculatorContract* cc) {
cc->Outputs().Tag("OUT").Set<int>();
cc->Outputs().Tag("EVENT").Set<int>();
cc->Outputs().Tag(kOutTag).Set<int>();
cc->Outputs().Tag(kEventTag).Set<int>();
return absl::OkStatus();
}
absl::Status Open(CalculatorContext* cc) override {
cc->Outputs().Tag("EVENT").AddPacket(MakePacket<int>(1).At(Timestamp(1)));
cc->Outputs().Tag(kEventTag).AddPacket(MakePacket<int>(1).At(Timestamp(1)));
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) override {
cc->Outputs().Tag("OUT").AddPacket(
cc->Outputs().Tag(kOutTag).AddPacket(
MakePacket<int>(count_).At(Timestamp(count_)));
count_++;
return absl::OkStatus();
}
absl::Status Close(CalculatorContext* cc) override {
cc->Outputs().Tag("EVENT").AddPacket(MakePacket<int>(2).At(Timestamp(2)));
cc->Outputs().Tag(kEventTag).AddPacket(MakePacket<int>(2).At(Timestamp(2)));
return absl::OkStatus();
}
@@ -81,11 +85,11 @@ class StoppingPassThroughCalculator : public mediapipe::CalculatorBase {
cc->Inputs().Get("", i).SetAny();
cc->Outputs().Get("", i).SetSameAs(&cc->Inputs().Get("", i));
}
cc->Outputs().Tag("EVENT").Set<int>();
cc->Outputs().Tag(kEventTag).Set<int>();
return absl::OkStatus();
}
absl::Status Open(CalculatorContext* cc) override {
cc->Outputs().Tag("EVENT").AddPacket(MakePacket<int>(1).At(Timestamp(1)));
cc->Outputs().Tag(kEventTag).AddPacket(MakePacket<int>(1).At(Timestamp(1)));
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) override {
@@ -98,7 +102,7 @@ class StoppingPassThroughCalculator : public mediapipe::CalculatorBase {
: mediapipe::tool::StatusStop();
}
absl::Status Close(CalculatorContext* cc) override {
cc->Outputs().Tag("EVENT").AddPacket(MakePacket<int>(2).At(Timestamp(2)));
cc->Outputs().Tag(kEventTag).AddPacket(MakePacket<int>(2).At(Timestamp(2)));
return absl::OkStatus();
}
mediapipe/framework/calculator_graph_test.cc
+31 -21
View File
@@ -65,6 +65,16 @@ namespace mediapipe {
namespace {
constexpr char kCounter2Tag[] = "COUNTER2";
constexpr char kCounter1Tag[] = "COUNTER1";
constexpr char kExtraTag[] = "EXTRA";
constexpr char kWaitSemTag[] = "WAIT_SEM";
constexpr char kPostSemTag[] = "POST_SEM";
constexpr char kErrorOnOpenTag[] = "ERROR_ON_OPEN";
constexpr char kOutputTag[] = "OUTPUT";
constexpr char kInputTag[] = "INPUT";
constexpr char kSelectTag[] = "SELECT";
using testing::ElementsAre;
using testing::HasSubstr;
@@ -125,8 +135,8 @@ class DemuxTimedCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc) {
RET_CHECK_EQ(cc->Inputs().NumEntries(), 2);
cc->Inputs().Tag("SELECT").Set<int>();
PacketType* data_input = &cc->Inputs().Tag("INPUT");
cc->Inputs().Tag(kSelectTag).Set<int>();
PacketType* data_input = &cc->Inputs().Tag(kInputTag);
data_input->SetAny();
for (CollectionItemId id = cc->Outputs().BeginId("OUTPUT");
id < cc->Outputs().EndId("OUTPUT"); ++id) {
@@ -182,7 +192,7 @@ REGISTER_CALCULATOR(DemuxTimedCalculator);
class MuxTimedCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Tag("SELECT").Set<int>();
cc->Inputs().Tag(kSelectTag).Set<int>();
CollectionItemId data_input_id = cc->Inputs().BeginId("INPUT");
PacketType* data_input0 = &cc->Inputs().Get(data_input_id);
data_input0->SetAny();
@@ -191,7 +201,7 @@ class MuxTimedCalculator : public CalculatorBase {
cc->Inputs().Get(data_input_id).SetSameAs(data_input0);
}
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1);
cc->Outputs().Tag("OUTPUT").SetSameAs(data_input0);
cc->Outputs().Tag(kOutputTag).SetSameAs(data_input0);
cc->SetTimestampOffset(TimestampDiff(0));
return absl::OkStatus();
}
@@ -598,12 +608,12 @@ class ErrorOnOpenCalculator : public CalculatorBase {
static absl::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).SetAny();
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
cc->InputSidePackets().Tag("ERROR_ON_OPEN").Set<bool>();
cc->InputSidePackets().Tag(kErrorOnOpenTag).Set<bool>();
return absl::OkStatus();
}
absl::Status Open(CalculatorContext* cc) final {
if (cc->InputSidePackets().Tag("ERROR_ON_OPEN").Get<bool>()) {
if (cc->InputSidePackets().Tag(kErrorOnOpenTag).Get<bool>()) {
return absl::NotFoundError("expected error");
}
return absl::OkStatus();
@@ -920,8 +930,8 @@ class SemaphoreCalculator : public CalculatorBase {
static absl::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).SetAny();
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
cc->InputSidePackets().Tag("POST_SEM").Set<Semaphore*>();
cc->InputSidePackets().Tag("WAIT_SEM").Set<Semaphore*>();
cc->InputSidePackets().Tag(kPostSemTag).Set<Semaphore*>();
cc->InputSidePackets().Tag(kWaitSemTag).Set<Semaphore*>();
cc->SetTimestampOffset(TimestampDiff(0));
return absl::OkStatus();
}
@@ -929,8 +939,8 @@ class SemaphoreCalculator : public CalculatorBase {
absl::Status Open(CalculatorContext* cc) override { return absl::OkStatus(); }
absl::Status Process(CalculatorContext* cc) override {
cc->InputSidePackets().Tag("POST_SEM").Get<Semaphore*>()->Release(1);
cc->InputSidePackets().Tag("WAIT_SEM").Get<Semaphore*>()->Acquire(1);
cc->InputSidePackets().Tag(kPostSemTag).Get<Semaphore*>()->Release(1);
cc->InputSidePackets().Tag(kWaitSemTag).Get<Semaphore*>()->Acquire(1);
cc->Outputs().Index(0).AddPacket(cc->Inputs().Index(0).Value());
return absl::OkStatus();
}
@@ -1177,9 +1187,9 @@ class IncrementingStatusHandler : public StatusHandler {
static absl::Status FillExpectations(
const MediaPipeOptions& extendable_options,
PacketTypeSet* input_side_packets) {
input_side_packets->Tag("EXTRA").SetAny().Optional();
input_side_packets->Tag("COUNTER1").Set<std::unique_ptr<int>>();
input_side_packets->Tag("COUNTER2").Set<std::unique_ptr<int>>();
input_side_packets->Tag(kExtraTag).SetAny().Optional();
input_side_packets->Tag(kCounter1Tag).Set<std::unique_ptr<int>>();
input_side_packets->Tag(kCounter2Tag).Set<std::unique_ptr<int>>();
return absl::OkStatus();
}
@@ -1187,7 +1197,7 @@ class IncrementingStatusHandler : public StatusHandler {
const MediaPipeOptions& extendable_options,
const PacketSet& input_side_packets, //
const absl::Status& pre_run_status) {
int* counter = GetFromUniquePtr<int>(input_side_packets.Tag("COUNTER1"));
int* counter = GetFromUniquePtr<int>(input_side_packets.Tag(kCounter1Tag));
(*counter)++;
return pre_run_status_result_;
}
@@ -1195,7 +1205,7 @@ class IncrementingStatusHandler : public StatusHandler {
static absl::Status HandleStatus(const MediaPipeOptions& extendable_options,
const PacketSet& input_side_packets, //
const absl::Status& run_status) {
int* counter = GetFromUniquePtr<int>(input_side_packets.Tag("COUNTER2"));
int* counter = GetFromUniquePtr<int>(input_side_packets.Tag(kCounter2Tag));
(*counter)++;
return post_run_status_result_;
}
@@ -2228,20 +2238,20 @@ class DemuxUntimedCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc) {
RET_CHECK_EQ(cc->Inputs().NumEntries(), 2);
cc->Inputs().Tag("INPUT").SetAny();
cc->Inputs().Tag("SELECT").Set<int>();
cc->Inputs().Tag(kInputTag).SetAny();
cc->Inputs().Tag(kSelectTag).Set<int>();
for (CollectionItemId id = cc->Outputs().BeginId("OUTPUT");
id < cc->Outputs().EndId("OUTPUT"); ++id) {
cc->Outputs().Get(id).SetSameAs(&cc->Inputs().Tag("INPUT"));
cc->Outputs().Get(id).SetSameAs(&cc->Inputs().Tag(kInputTag));
}
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) final {
int index = cc->Inputs().Tag("SELECT").Get<int>();
if (!cc->Inputs().Tag("INPUT").IsEmpty()) {
int index = cc->Inputs().Tag(kSelectTag).Get<int>();
if (!cc->Inputs().Tag(kInputTag).IsEmpty()) {
cc->Outputs()
.Get("OUTPUT", index)
.AddPacket(cc->Inputs().Tag("INPUT").Value());
.AddPacket(cc->Inputs().Tag(kInputTag).Value());
} else {
cc->Outputs()
.Get("OUTPUT", index)
mediapipe/framework/calculator_runner_test.cc
+11 -6
View File
@@ -32,6 +32,11 @@
namespace mediapipe {
namespace {
constexpr char kTag[] = "";
constexpr char kBTag[] = "B";
constexpr char kATag[] = "A";
constexpr char kSideOutputTag[] = "SIDE_OUTPUT";
// Inputs: 2 streams with ints. Headers are strings.
// Input side packets: 1.
// Outputs: 3 streams with ints. #0 and #1 will contain the negated values from
@@ -48,7 +53,7 @@ class CalculatorRunnerTestCalculator : public CalculatorBase {
cc->Outputs().Index(2).SetSameAs(&cc->InputSidePackets().Index(0));
cc->InputSidePackets().Index(0).SetAny();
cc->OutputSidePackets()
.Tag("SIDE_OUTPUT")
.Tag(kSideOutputTag)
.SetSameAs(&cc->InputSidePackets().Index(0));
return absl::OkStatus();
}
@@ -64,7 +69,7 @@ class CalculatorRunnerTestCalculator : public CalculatorBase {
Adopt(new std::string(absl::StrCat(input_header_string, i))));
}
cc->OutputSidePackets()
.Tag("SIDE_OUTPUT")
.Tag(kSideOutputTag)
.Set(cc->InputSidePackets().Index(0));
return absl::OkStatus();
}
@@ -152,7 +157,7 @@ TEST(CalculatorRunner, RunsCalculator) {
Adopt(new int(input_side_packet_content));
MP_ASSERT_OK(runner.Run());
EXPECT_EQ(input_side_packet_content,
runner.OutputSidePackets().Tag("SIDE_OUTPUT").Get<int>());
runner.OutputSidePackets().Tag(kSideOutputTag).Get<int>());
const auto& outputs = runner.Outputs();
ASSERT_EQ(3, outputs.NumEntries());
@@ -209,9 +214,9 @@ TEST(CalculatorRunner, MultiTagTestCalculatorOk) {
const auto& outputs = runner.Outputs();
ASSERT_EQ(3, outputs.NumEntries());
for (int ts = 0; ts < 5; ++ts) {
const std::vector<Packet>& a_packets = outputs.Tag("A").packets;
const std::vector<Packet>& b_packets = outputs.Tag("B").packets;
const std::vector<Packet>& c_packets = outputs.Tag("").packets;
const std::vector<Packet>& a_packets = outputs.Tag(kATag).packets;
const std::vector<Packet>& b_packets = outputs.Tag(kBTag).packets;
const std::vector<Packet>& c_packets = outputs.Tag(kTag).packets;
EXPECT_EQ(Timestamp(ts), a_packets[ts].Timestamp());
EXPECT_EQ(Timestamp(ts), b_packets[ts].Timestamp());
EXPECT_EQ(Timestamp(ts), c_packets[ts].Timestamp());
mediapipe/framework/collection_test.cc
+10 -6
View File
@@ -24,6 +24,10 @@
namespace mediapipe {
namespace {
constexpr char kTag2Tag[] = "TAG_2";
constexpr char kTag0Tag[] = "TAG_0";
constexpr char kTag1Tag[] = "TAG_1";
TEST(CollectionTest, BasicByIndex) {
tool::TagAndNameInfo info;
info.names.push_back("name_1");
@@ -55,14 +59,14 @@ TEST(CollectionTest, BasicByTag) {
info.names.push_back("name_2");
info.tags.push_back("TAG_2");
internal::Collection<int> collection(info);
collection.Tag("TAG_1") = 101;
collection.Tag("TAG_0") = 100;
collection.Tag("TAG_2") = 102;
collection.Tag(kTag1Tag) = 101;
collection.Tag(kTag0Tag) = 100;
collection.Tag(kTag2Tag) = 102;
// Test the stored values.
EXPECT_EQ(100, collection.Tag("TAG_0"));
EXPECT_EQ(101, collection.Tag("TAG_1"));
EXPECT_EQ(102, collection.Tag("TAG_2"));
EXPECT_EQ(100, collection.Tag(kTag0Tag));
EXPECT_EQ(101, collection.Tag(kTag1Tag));
EXPECT_EQ(102, collection.Tag(kTag2Tag));
// Test access using a range based for.
int i = 0;
for (int num : collection) {
mediapipe/framework/formats/tensor.cc
+54 -4
View File
@@ -134,6 +134,21 @@ void Tensor::AllocateMtlBuffer(id<MTLDevice> device) const {
#endif // MEDIAPIPE_METAL_ENABLED
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
bool Tensor::NeedsHalfFloatRenderTarget() const {
static bool has_color_buffer_float =
gl_context_->HasGlExtension("WEBGL_color_buffer_float") ||
gl_context_->HasGlExtension("EXT_color_buffer_float");
if (!has_color_buffer_float) {
static bool has_color_buffer_half_float =
gl_context_->HasGlExtension("EXT_color_buffer_half_float");
LOG_IF(FATAL, !has_color_buffer_half_float)
<< "EXT_color_buffer_half_float or WEBGL_color_buffer_float "
<< "required on web to use MP tensor";
return true;
}
return false;
}
Tensor::OpenGlTexture2dView Tensor::GetOpenGlTexture2dReadView() const {
LOG_IF(FATAL, valid_ == kValidNone)
<< "Tensor must be written prior to read from.";
@@ -164,8 +179,24 @@ Tensor::OpenGlTexture2dView Tensor::GetOpenGlTexture2dReadView() const {
// Set alignment for the proper value (default) to avoid address sanitizer
// error "out of boundary reading".
glPixelStorei(GL_UNPACK_ALIGNMENT, 4);
glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, texture_width_, texture_height_,
GL_RGBA, GL_FLOAT, temp_buffer.get());
#ifdef __EMSCRIPTEN__
// Under WebGL1, format must match in order to use glTexSubImage2D, so if we
// have a half-float texture, then uploading from GL_FLOAT here would fail.
// We change the texture's data type to float here to accommodate.
// Furthermore, for a full-image replacement operation, glTexImage2D is
// expected to be more performant than glTexSubImage2D. Note that for WebGL2
// we cannot use glTexImage2D, because we allocate using glTexStorage2D in
// that case, which is incompatible.
if (gl_context_->GetGlVersion() == mediapipe::GlVersion::kGLES2) {
glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, texture_width_, texture_height_,
0, GL_RGBA, GL_FLOAT, temp_buffer.get());
texture_is_half_float_ = false;
} else
#endif // __EMSCRIPTEN__
{
glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, texture_width_, texture_height_,
GL_RGBA, GL_FLOAT, temp_buffer.get());
}
glBindTexture(GL_TEXTURE_2D, 0);
valid_ |= kValidOpenGlTexture2d;
}
@@ -175,6 +206,16 @@ Tensor::OpenGlTexture2dView Tensor::GetOpenGlTexture2dReadView() const {
Tensor::OpenGlTexture2dView Tensor::GetOpenGlTexture2dWriteView() const {
auto lock = absl::make_unique<absl::MutexLock>(&view_mutex_);
AllocateOpenGlTexture2d();
#ifdef __EMSCRIPTEN__
// On web, we may have to change type from float to half-float
if (!texture_is_half_float_ && NeedsHalfFloatRenderTarget()) {
glBindTexture(GL_TEXTURE_2D, opengl_texture2d_);
glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, texture_width_, texture_height_, 0,
GL_RGBA, GL_HALF_FLOAT_OES, 0 /* data */);
glBindTexture(GL_TEXTURE_2D, 0);
texture_is_half_float_ = true;
}
#endif
valid_ = kValidOpenGlTexture2d;
return {opengl_texture2d_, std::move(lock)};
}
@@ -255,8 +296,18 @@ void Tensor::AllocateOpenGlTexture2d() const {
<< "with GLES 2.0";
// Allocate the image data; note that it's no longer RGBA32F, so will be
// lower precision.
auto type = GL_FLOAT;
// On web, we might need to change type to half-float (e.g. for iOS-
// Safari) in order to have a valid framebuffer. See b/194442743 for more
// details.
#ifdef __EMSCRIPTEN__
if (NeedsHalfFloatRenderTarget()) {
type = GL_HALF_FLOAT_OES;
texture_is_half_float_ = true;
}
#endif // __EMSCRIPTEN
glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, texture_width_, texture_height_,
0, GL_RGBA, GL_FLOAT, 0 /* data */);
0, GL_RGBA, type, 0 /* data */);
}
glBindTexture(GL_TEXTURE_2D, 0);
glGenFramebuffers(1, &frame_buffer_);
@@ -443,7 +494,6 @@ Tensor::CpuReadView Tensor::GetCpuReadView() const {
glPixelStorei(GL_PACK_ALIGNMENT, 4);
glReadPixels(0, 0, texture_width_, texture_height_, GL_RGBA, GL_FLOAT,
buffer);
uint8_t* dest_buffer = reinterpret_cast<uint8_t*>(cpu_buffer_);
const int actual_depth_size =
BhwcDepthFromShape(shape_) * element_size();
mediapipe/framework/formats/tensor.h
+4
View File
@@ -266,11 +266,15 @@ class Tensor {
mutable GLuint frame_buffer_ = GL_INVALID_INDEX;
mutable int texture_width_;
mutable int texture_height_;
#ifdef __EMSCRIPTEN__
mutable bool texture_is_half_float_ = false;
#endif // __EMSCRIPTEN__
void AllocateOpenGlTexture2d() const;
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
mutable GLuint opengl_buffer_ = GL_INVALID_INDEX;
void AllocateOpenGlBuffer() const;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
bool NeedsHalfFloatRenderTarget() const;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
};
mediapipe/framework/graph_validation_test.cc
+14 -8
View File
@@ -31,6 +31,11 @@ namespace mediapipe {
namespace {
constexpr char kOutputTag[] = "OUTPUT";
constexpr char kEnableTag[] = "ENABLE";
constexpr char kSelectTag[] = "SELECT";
constexpr char kSideinputTag[] = "SIDEINPUT";
// Shows validation success for a graph and a subgraph.
TEST(GraphValidationTest, InitializeGraphFromProtos) {
auto config_1 = ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
@@ -323,20 +328,21 @@ TEST(GraphValidationTest, OptionalSubgraphStreamsMismatched) {
class OptionalSideInputTestCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc) {
cc->InputSidePackets().Tag("SIDEINPUT").Set<std::string>().Optional();
cc->Inputs().Tag("SELECT").Set<int>().Optional();
cc->Inputs().Tag("ENABLE").Set<bool>().Optional();
cc->Outputs().Tag("OUTPUT").Set<std::string>();
cc->InputSidePackets().Tag(kSideinputTag).Set<std::string>().Optional();
cc->Inputs().Tag(kSelectTag).Set<int>().Optional();
cc->Inputs().Tag(kEnableTag).Set<bool>().Optional();
cc->Outputs().Tag(kOutputTag).Set<std::string>();
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) final {
std::string value("default");
if (cc->InputSidePackets().HasTag("SIDEINPUT")) {
value = cc->InputSidePackets().Tag("SIDEINPUT").Get<std::string>();
if (cc->InputSidePackets().HasTag(kSideinputTag)) {
value = cc->InputSidePackets().Tag(kSideinputTag).Get<std::string>();
}
cc->Outputs().Tag("OUTPUT").Add(new std::string(value),
cc->InputTimestamp());
cc->Outputs()
.Tag(kOutputTag)
.Add(new std::string(value), cc->InputTimestamp());
return absl::OkStatus();
}
};
mediapipe/framework/packet_registration_test.cc
+7 -4
View File
@@ -26,17 +26,20 @@ namespace {
namespace test_ns {
constexpr char kOutTag[] = "OUT";
constexpr char kInTag[] = "IN";
class TestSinkCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Tag("IN").Set<mediapipe::InputOnlyProto>();
cc->Outputs().Tag("OUT").Set<int>();
cc->Inputs().Tag(kInTag).Set<mediapipe::InputOnlyProto>();
cc->Outputs().Tag(kOutTag).Set<int>();
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) override {
int x = cc->Inputs().Tag("IN").Get<mediapipe::InputOnlyProto>().x();
cc->Outputs().Tag("OUT").AddPacket(
int x = cc->Inputs().Tag(kInTag).Get<mediapipe::InputOnlyProto>().x();
cc->Outputs().Tag(kOutTag).AddPacket(
MakePacket<int>(x).At(cc->InputTimestamp()));
return absl::OkStatus();
}
mediapipe/framework/test_calculators.cc
+41 -27
View File
@@ -34,6 +34,19 @@
namespace mediapipe {
constexpr char kOutTag[] = "OUT";
constexpr char kClockTag[] = "CLOCK";
constexpr char kSleepMicrosTag[] = "SLEEP_MICROS";
constexpr char kCloseTag[] = "CLOSE";
constexpr char kProcessTag[] = "PROCESS";
constexpr char kOpenTag[] = "OPEN";
constexpr char kTag[] = "";
constexpr char kMeanTag[] = "MEAN";
constexpr char kDataTag[] = "DATA";
constexpr char kPairTag[] = "PAIR";
constexpr char kLowTag[] = "LOW";
constexpr char kHighTag[] = "HIGH";
using RandomEngine = std::mt19937_64;
// A Calculator that outputs twice the value of its input packet (an int).
@@ -95,9 +108,9 @@ class TaggedIntSplitterPacketGenerator : public PacketGenerator {
PacketTypeSet* input_side_packets, //
PacketTypeSet* output_side_packets) {
input_side_packets->Index(0).Set<uint64>();
output_side_packets->Tag("HIGH").Set<uint32>();
output_side_packets->Tag("LOW").Set<uint32>();
output_side_packets->Tag("PAIR").Set<std::pair<uint32, uint32>>();
output_side_packets->Tag(kHighTag).Set<uint32>();
output_side_packets->Tag(kLowTag).Set<uint32>();
output_side_packets->Tag(kPairTag).Set<std::pair<uint32, uint32>>();
return absl::OkStatus();
}
@@ -108,9 +121,9 @@ class TaggedIntSplitterPacketGenerator : public PacketGenerator {
uint64 value = input_side_packets.Index(0).Get<uint64>();
uint32 high = value >> 32;
uint32 low = value & 0xFFFFFFFF;
output_side_packets->Tag("HIGH") = Adopt(new uint32(high));
output_side_packets->Tag("LOW") = Adopt(new uint32(low));
output_side_packets->Tag("PAIR") =
output_side_packets->Tag(kHighTag) = Adopt(new uint32(high));
output_side_packets->Tag(kLowTag) = Adopt(new uint32(low));
output_side_packets->Tag(kPairTag) =
Adopt(new std::pair<uint32, uint32>(high, low));
return absl::OkStatus();
}
@@ -221,8 +234,8 @@ class StdDevCalculator : public CalculatorBase {
StdDevCalculator() {}
static absl::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Tag("DATA").Set<int>();
cc->Inputs().Tag("MEAN").Set<double>();
cc->Inputs().Tag(kDataTag).Set<int>();
cc->Inputs().Tag(kMeanTag).Set<double>();
cc->Outputs().Index(0).Set<int>();
return absl::OkStatus();
}
@@ -234,15 +247,15 @@ class StdDevCalculator : public CalculatorBase {
absl::Status Process(CalculatorContext* cc) final {
if (cc->InputTimestamp() == Timestamp::PreStream()) {
RET_CHECK(cc->Inputs().Tag("DATA").Value().IsEmpty());
RET_CHECK(!cc->Inputs().Tag("MEAN").Value().IsEmpty());
mean_ = cc->Inputs().Tag("MEAN").Get<double>();
RET_CHECK(cc->Inputs().Tag(kDataTag).Value().IsEmpty());
RET_CHECK(!cc->Inputs().Tag(kMeanTag).Value().IsEmpty());
mean_ = cc->Inputs().Tag(kMeanTag).Get<double>();
initialized_ = true;
} else {
RET_CHECK(initialized_);
RET_CHECK(!cc->Inputs().Tag("DATA").Value().IsEmpty());
RET_CHECK(cc->Inputs().Tag("MEAN").Value().IsEmpty());
double diff = cc->Inputs().Tag("DATA").Get<int>() - mean_;
RET_CHECK(!cc->Inputs().Tag(kDataTag).Value().IsEmpty());
RET_CHECK(cc->Inputs().Tag(kMeanTag).Value().IsEmpty());
double diff = cc->Inputs().Tag(kDataTag).Get<int>() - mean_;
cummulative_variance_ += diff * diff;
++count_;
}
@@ -564,8 +577,8 @@ class LambdaCalculator : public CalculatorBase {
id < cc->Outputs().EndId(); ++id) {
cc->Outputs().Get(id).SetAny();
}
if (cc->InputSidePackets().HasTag("") > 0) {
cc->InputSidePackets().Tag("").Set<ProcessFunction>();
if (cc->InputSidePackets().HasTag(kTag) > 0) {
cc->InputSidePackets().Tag(kTag).Set<ProcessFunction>();
}
for (const std::string& tag : {"OPEN", "PROCESS", "CLOSE"}) {
if (cc->InputSidePackets().HasTag(tag)) {
@@ -576,24 +589,24 @@ class LambdaCalculator : public CalculatorBase {
}
absl::Status Open(CalculatorContext* cc) final {
if (cc->InputSidePackets().HasTag("OPEN")) {
if (cc->InputSidePackets().HasTag(kOpenTag)) {
return GetContextFn(cc, "OPEN")(cc);
}
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) final {
if (cc->InputSidePackets().HasTag("PROCESS")) {
if (cc->InputSidePackets().HasTag(kProcessTag)) {
return GetContextFn(cc, "PROCESS")(cc);
}
if (cc->InputSidePackets().HasTag("") > 0) {
if (cc->InputSidePackets().HasTag(kTag) > 0) {
return GetProcessFn(cc, "")(cc->Inputs(), &cc->Outputs());
}
return absl::OkStatus();
}
absl::Status Close(CalculatorContext* cc) final {
if (cc->InputSidePackets().HasTag("CLOSE")) {
if (cc->InputSidePackets().HasTag(kCloseTag)) {
return GetContextFn(cc, "CLOSE")(cc);
}
return absl::OkStatus();
@@ -645,17 +658,18 @@ class PassThroughWithSleepCalculator : public CalculatorBase {
static absl::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).Set<int>();
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
cc->InputSidePackets().Tag("SLEEP_MICROS").Set<int>();
cc->InputSidePackets().Tag("CLOCK").Set<std::shared_ptr<Clock>>();
cc->InputSidePackets().Tag(kSleepMicrosTag).Set<int>();
cc->InputSidePackets().Tag(kClockTag).Set<std::shared_ptr<Clock>>();
return absl::OkStatus();
}
absl::Status Open(CalculatorContext* cc) final {
cc->SetOffset(TimestampDiff(0));
sleep_micros_ = cc->InputSidePackets().Tag("SLEEP_MICROS").Get<int>();
sleep_micros_ = cc->InputSidePackets().Tag(kSleepMicrosTag).Get<int>();
if (sleep_micros_ < 0) {
return absl::InternalError("SLEEP_MICROS should be >= 0");
}
clock_ = cc->InputSidePackets().Tag("CLOCK").Get<std::shared_ptr<Clock>>();
clock_ =
cc->InputSidePackets().Tag(kClockTag).Get<std::shared_ptr<Clock>>();
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) final {
@@ -678,8 +692,8 @@ class MultiplyIntCalculator : public CalculatorBase {
cc->Inputs().Index(0).Set<int>();
cc->Inputs().Index(1).SetSameAs(&cc->Inputs().Index(0));
// cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
RET_CHECK(cc->Outputs().HasTag("OUT"));
cc->Outputs().Tag("OUT").SetSameAs(&cc->Inputs().Index(0));
RET_CHECK(cc->Outputs().HasTag(kOutTag));
cc->Outputs().Tag(kOutTag).SetSameAs(&cc->Inputs().Index(0));
return absl::OkStatus();
}
absl::Status Open(CalculatorContext* cc) final {
@@ -689,7 +703,7 @@ class MultiplyIntCalculator : public CalculatorBase {
absl::Status Process(CalculatorContext* cc) final {
int x = cc->Inputs().Index(0).Value().Get<int>();
int y = cc->Inputs().Index(1).Value().Get<int>();
cc->Outputs().Tag("OUT").Add(new int(x * y), cc->InputTimestamp());
cc->Outputs().Tag(kOutTag).Add(new int(x * y), cc->InputTimestamp());
return absl::OkStatus();
}
};