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

Internal change

Browse Source 
PiperOrigin-RevId: 477538515
This commit is contained in:
MediaPipe Team
2022-09-28 21:32:36 +00:00
committed by Sebastian Schmidt
parent 6cdc6443b6
commit f8af41b1eb
236 changed files with 12865 additions and 1923 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/api2/BUILD
+2
View File
@@ -21,7 +21,9 @@ cc_library(
":port",
"//mediapipe/framework:calculator_base",
"//mediapipe/framework:calculator_contract",
"@com_google_absl//absl/container:btree",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
],
)
mediapipe/framework/api2/builder.h
+24 -12
View File
@@ -4,7 +4,9 @@
#include <string>
#include <type_traits>
#include "absl/container/btree_map.h"
#include "absl/container/flat_hash_map.h"
#include "absl/strings/string_view.h"
#include "mediapipe/framework/api2/const_str.h"
#include "mediapipe/framework/api2/contract.h"
#include "mediapipe/framework/api2/node.h"
@@ -46,7 +48,7 @@ struct TagIndexLocation {
template <typename T>
class TagIndexMap {
public:
std::vector<std::unique_ptr<T>>& operator[](const std::string& tag) {
std::vector<std::unique_ptr<T>>& operator[](absl::string_view tag) {
return map_[tag];
}
@@ -72,7 +74,7 @@ class TagIndexMap {
// Note: entries are held by a unique_ptr to ensure pointers remain valid.
// Should use absl::flat_hash_map but ordering keys for now.
std::map<std::string, std::vector<std::unique_ptr<T>>> map_;
absl::btree_map<std::string, std::vector<std::unique_ptr<T>>> map_;
};
class Graph;
@@ -169,6 +171,16 @@ class SourceImpl {
return AddTarget(dest);
}
template <typename U>
struct AllowCast
: public std::integral_constant<bool, std::is_same_v<T, AnyType> &&
!std::is_same_v<T, U>> {};
template <typename U, std::enable_if_t<AllowCast<U>{}, int> = 0>
SourceImpl<IsSide, U> Cast() {
return SourceImpl<IsSide, U>(base_);
}
private:
// Never null.
SourceBase* base_;
@@ -212,19 +224,19 @@ class NodeBase {
// of its entries by index. However, for nodes without visible contracts we
// can't know whether a tag is indexable or not, so we would need the
// multi-port to also be usable as a port directly (representing index 0).
MultiSource<> Out(const std::string& tag) {
MultiSource<> Out(absl::string_view tag) {
return MultiSource<>(&out_streams_[tag]);
}
MultiDestination<> In(const std::string& tag) {
MultiDestination<> In(absl::string_view tag) {
return MultiDestination<>(&in_streams_[tag]);
}
MultiSideSource<> SideOut(const std::string& tag) {
MultiSideSource<> SideOut(absl::string_view tag) {
return MultiSideSource<>(&out_sides_[tag]);
}
MultiSideDestination<> SideIn(const std::string& tag) {
MultiSideDestination<> SideIn(absl::string_view tag) {
return MultiSideDestination<>(&in_sides_[tag]);
}
@@ -359,11 +371,11 @@ class PacketGenerator {
public:
PacketGenerator(std::string type) : type_(std::move(type)) {}
MultiSideSource<> SideOut(const std::string& tag) {
MultiSideSource<> SideOut(absl::string_view tag) {
return MultiSideSource<>(&out_sides_[tag]);
}
MultiSideDestination<> SideIn(const std::string& tag) {
MultiSideDestination<> SideIn(absl::string_view tag) {
return MultiSideDestination<>(&in_sides_[tag]);
}
@@ -452,19 +464,19 @@ class Graph {
}
// Graph ports, non-typed.
MultiSource<> In(const std::string& graph_input) {
MultiSource<> In(absl::string_view graph_input) {
return graph_boundary_.Out(graph_input);
}
MultiDestination<> Out(const std::string& graph_output) {
MultiDestination<> Out(absl::string_view graph_output) {
return graph_boundary_.In(graph_output);
}
MultiSideSource<> SideIn(const std::string& graph_input) {
MultiSideSource<> SideIn(absl::string_view graph_input) {
return graph_boundary_.SideOut(graph_input);
}
MultiSideDestination<> SideOut(const std::string& graph_output) {
MultiSideDestination<> SideOut(absl::string_view graph_output) {
return graph_boundary_.SideIn(graph_output);
}
mediapipe/framework/api2/builder_test.cc
+77 -16
View File
@@ -2,6 +2,7 @@
#include <functional>
#include "absl/strings/string_view.h"
#include "absl/strings/substitute.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/api2/packet.h"
@@ -296,6 +297,32 @@ TEST(BuilderTest, EmptyTag) {
EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
}
TEST(BuilderTest, StringLikeTags) {
const char kA[] = "A";
const std::string kB = "B";
constexpr absl::string_view kC = "C";
builder::Graph graph;
auto& foo = graph.AddNode("Foo");
graph.In(kA).SetName("a") >> foo.In(kA);
graph.In(kB).SetName("b") >> foo.In(kB);
foo.Out(kC).SetName("c") >> graph.Out(kC);
CalculatorGraphConfig expected =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
input_stream: "A:a"
input_stream: "B:b"
output_stream: "C:c"
node {
calculator: "Foo"
input_stream: "A:a"
input_stream: "B:b"
output_stream: "C:c"
}
)pb");
EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
}
TEST(BuilderTest, GraphIndexes) {
builder::Graph graph;
auto& foo = graph.AddNode("Foo");
@@ -326,57 +353,91 @@ TEST(BuilderTest, GraphIndexes) {
class AnyAndSameTypeCalculator : public NodeIntf {
public:
static constexpr Input<AnyType> kAnyTypeInput{"INPUT"};
static constexpr Output<AnyType> kAnyTypeOutput{"ANY_OUTPUT"};
static constexpr Output<SameType<kAnyTypeInput>> kSameTypeOutput{
static constexpr Input<AnyType>::Optional kAnyTypeInput{"INPUT"};
static constexpr Output<AnyType>::Optional kAnyTypeOutput{"ANY_OUTPUT"};
static constexpr Output<SameType<kAnyTypeInput>>::Optional kSameTypeOutput{
"SAME_OUTPUT"};
static constexpr Output<SameType<kSameTypeOutput>> kRecursiveSameTypeOutput{
"RECURSIVE_SAME_OUTPUT"};
static constexpr Input<int> kIntInput{"INT_INPUT"};
static constexpr Input<int>::Optional kIntInput{"INT_INPUT"};
// `SameType` usage for this output is only for testing purposes.
//
// `SameType` is designed to work with inputs of `AnyType` and, normally, you
// would not use `Output<SameType<kIntInput>>` in a real calculator. You
// should write `Output<int>` instead, since the type is known.
static constexpr Output<SameType<kIntInput>> kSameIntOutput{
static constexpr Output<SameType<kIntInput>>::Optional kSameIntOutput{
"SAME_INT_OUTPUT"};
static constexpr Output<SameType<kSameIntOutput>> kRecursiveSameIntOutput{
"RECURSIVE_SAME_INT_OUTPUT"};
MEDIAPIPE_NODE_INTERFACE(AnyTypeCalculator, kAnyTypeInput, kAnyTypeOutput,
kSameTypeOutput);
MEDIAPIPE_NODE_INTERFACE(AnyAndSameTypeCalculator, kAnyTypeInput,
kAnyTypeOutput, kSameTypeOutput);
};
TEST(BuilderTest, AnyAndSameTypeHandledProperly) {
builder::Graph graph;
builder::Source<internal::Generic> any_input =
graph[Input<AnyType>{"GRAPH_ANY_INPUT"}];
builder::Source<AnyType> any_input = graph[Input<AnyType>{"GRAPH_ANY_INPUT"}];
builder::Source<int> int_input = graph[Input<int>{"GRAPH_INT_INPUT"}];
auto& node = graph.AddNode("AnyAndSameTypeCalculator");
any_input >> node[AnyAndSameTypeCalculator::kAnyTypeInput];
int_input >> node[AnyAndSameTypeCalculator::kIntInput];
builder::Source<internal::Generic> any_type_output =
builder::Source<AnyType> any_type_output =
node[AnyAndSameTypeCalculator::kAnyTypeOutput];
any_type_output.SetName("any_type_output");
builder::Source<internal::Generic> same_type_output =
builder::Source<AnyType> same_type_output =
node[AnyAndSameTypeCalculator::kSameTypeOutput];
same_type_output.SetName("same_type_output");
builder::Source<internal::Generic> same_int_output =
builder::Source<AnyType> recursive_same_type_output =
node[AnyAndSameTypeCalculator::kRecursiveSameTypeOutput];
recursive_same_type_output.SetName("recursive_same_type_output");
builder::Source<int> same_int_output =
node[AnyAndSameTypeCalculator::kSameIntOutput];
same_int_output.SetName("same_int_output");
builder::Source<int> recursive_same_int_type_output =
node[AnyAndSameTypeCalculator::kRecursiveSameIntOutput];
recursive_same_int_type_output.SetName("recursive_same_int_type_output");
CalculatorGraphConfig expected = mediapipe::ParseTextProtoOrDie<
CalculatorGraphConfig>(R"pb(
node {
calculator: "AnyAndSameTypeCalculator"
input_stream: "INPUT:__stream_0"
input_stream: "INT_INPUT:__stream_1"
output_stream: "ANY_OUTPUT:any_type_output"
output_stream: "RECURSIVE_SAME_INT_OUTPUT:recursive_same_int_type_output"
output_stream: "RECURSIVE_SAME_OUTPUT:recursive_same_type_output"
output_stream: "SAME_INT_OUTPUT:same_int_output"
output_stream: "SAME_OUTPUT:same_type_output"
}
input_stream: "GRAPH_ANY_INPUT:__stream_0"
input_stream: "GRAPH_INT_INPUT:__stream_1"
)pb");
EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
}
TEST(BuilderTest, AnyTypeCanBeCast) {
builder::Graph graph;
builder::Source<std::string> any_input =
graph.In("GRAPH_ANY_INPUT").Cast<std::string>();
auto& node = graph.AddNode("AnyAndSameTypeCalculator");
any_input >> node[AnyAndSameTypeCalculator::kAnyTypeInput];
builder::Source<double> any_type_output =
node[AnyAndSameTypeCalculator::kAnyTypeOutput].Cast<double>();
any_type_output.SetName("any_type_output");
CalculatorGraphConfig expected =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
node {
calculator: "AnyAndSameTypeCalculator"
input_stream: "INPUT:__stream_0"
input_stream: "INT_INPUT:__stream_1"
output_stream: "ANY_OUTPUT:any_type_output"
output_stream: "SAME_INT_OUTPUT:same_int_output"
output_stream: "SAME_OUTPUT:same_type_output"
}
input_stream: "GRAPH_ANY_INPUT:__stream_0"
input_stream: "GRAPH_INT_INPUT:__stream_1"
)pb");
EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
}
mediapipe/framework/api2/packet.h
+1 -3
View File
@@ -27,9 +27,7 @@ using HolderBase = mediapipe::packet_internal::HolderBase;
template <typename T>
class Packet;
struct DynamicType {};
struct AnyType : public DynamicType {
struct AnyType {
AnyType() = delete;
};
mediapipe/framework/api2/port.h
+21 -16
View File
@@ -73,14 +73,12 @@ class SideOutputBase : public PortBase {
};
struct NoneType {
private:
NoneType() = delete;
};
template <auto& P>
class SameType : public DynamicType {
public:
static constexpr const decltype(P)& kPort = P;
template <auto& kP>
struct SameType {
static constexpr const decltype(kP)& kPort = kP;
};
class PacketTypeAccess;
@@ -137,21 +135,28 @@ struct IsOneOf : std::false_type {};
template <class... T>
struct IsOneOf<OneOf<T...>> : std::true_type {};
template <typename T, typename std::enable_if<
!std::is_base_of<DynamicType, T>{} && !IsOneOf<T>{},
int>::type = 0>
template <class T>
struct IsSameType : std::false_type {};
template <class P, P& kP>
struct IsSameType<SameType<kP>> : std::true_type {};
template <typename T,
typename std::enable_if<!std::is_same<T, AnyType>{} &&
!IsOneOf<T>{} && !IsSameType<T>{},
int>::type = 0>
inline void SetType(CalculatorContract* cc, PacketType& pt) {
pt.Set<T>();
}
template <typename T, typename std::enable_if<std::is_base_of<DynamicType, T>{},
int>::type = 0>
template <typename T, typename std::enable_if<IsSameType<T>{}, int>::type = 0>
inline void SetType(CalculatorContract* cc, PacketType& pt) {
pt.SetSameAs(&internal::GetCollection(cc, T::kPort).Tag(T::kPort.Tag()));
}
template <>
inline void SetType<AnyType>(CalculatorContract* cc, PacketType& pt) {
template <typename T,
typename std::enable_if<std::is_same<T, AnyType>{}, int>::type = 0>
inline void SetType(CalculatorContract* cc, PacketType& pt) {
pt.SetAny();
}
@@ -289,15 +294,15 @@ struct SideBase<InputBase> {
};
// TODO: maybe return a PacketBase instead of a Packet<internal::Generic>?
template <typename T, class = void>
template <typename T, typename = void>
struct ActualPayloadType {
using type = T;
};
template <typename T>
struct ActualPayloadType<
T, std::enable_if_t<std::is_base_of<DynamicType, T>{}, void>> {
using type = internal::Generic;
struct ActualPayloadType<T, std::enable_if_t<IsSameType<T>{}, void>> {
using type = typename ActualPayloadType<
typename std::decay_t<decltype(T::kPort)>::value_t>::type;
};
} // namespace internal