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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 7e1d382a1788ebd8412c5626581b4c4cf2fe75ea
This commit is contained in:
MediaPipe Team
2021-11-16 14:32:04 -05:00
committed by chuoling
parent f4e7f6cc48
commit cf101e62a9
39 changed files with 1102 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/api2/builder.h
+122 -80
View File
@@ -75,6 +75,9 @@ class TagIndexMap {
std::map<std::string, std::vector<std::unique_ptr<T>>> map_;
};
class Graph;
class NodeBase;
// These structs are used internally to store information about the endpoints
// of a connection.
struct SourceBase;
@@ -109,7 +112,7 @@ class MultiPort : public Single {
// These classes wrap references to the underlying source/destination
// endpoints, adding type information and the user-visible API.
template <bool AllowMultiple, bool IsSide, typename T = internal::Generic>
template <bool IsSide, typename T = internal::Generic>
class DestinationImpl {
public:
using Base = DestinationBase;
@@ -121,13 +124,12 @@ class DestinationImpl {
};
template <bool IsSide, typename T>
class DestinationImpl<true, IsSide, T>
: public MultiPort<DestinationImpl<false, IsSide, T>> {
class MultiDestinationImpl : public MultiPort<DestinationImpl<IsSide, T>> {
public:
using MultiPort<DestinationImpl<false, IsSide, T>>::MultiPort;
using MultiPort<DestinationImpl<IsSide, T>>::MultiPort;
};
template <bool AllowMultiple, bool IsSide, typename T = internal::Generic>
template <bool IsSide, typename T = internal::Generic>
class SourceImpl {
public:
using Base = SourceBase;
@@ -135,9 +137,9 @@ class SourceImpl {
// Src is used as the return type of fluent methods below. Since these are
// single-port methods, it is desirable to always decay to a reference to the
// single-port superclass, even if they are called on a multiport.
using Src = SourceImpl<false, IsSide, T>;
using Src = SourceImpl<IsSide, T>;
template <typename U>
using Dst = DestinationImpl<false, IsSide, U>;
using Dst = DestinationImpl<IsSide, U>;
// clang-format off
template <typename U>
@@ -173,10 +175,9 @@ class SourceImpl {
};
template <bool IsSide, typename T>
class SourceImpl<true, IsSide, T>
: public MultiPort<SourceImpl<false, IsSide, T>> {
class MultiSourceImpl : public MultiPort<SourceImpl<IsSide, T>> {
public:
using MultiPort<SourceImpl<false, IsSide, T>>::MultiPort;
using MultiPort<SourceImpl<IsSide, T>>::MultiPort;
};
// A source and a destination correspond to an output/input stream on a node,
@@ -185,14 +186,23 @@ class SourceImpl<true, IsSide, T>
// For graph inputs/outputs, however, the inputs are sources, and the outputs
// are destinations. This is because graph ports are connected "from inside"
// when building the graph.
template <bool AllowMultiple = false, typename T = internal::Generic>
using Source = SourceImpl<AllowMultiple, false, T>;
template <bool AllowMultiple = false, typename T = internal::Generic>
using SideSource = SourceImpl<AllowMultiple, true, T>;
template <bool AllowMultiple = false, typename T = internal::Generic>
using Destination = DestinationImpl<AllowMultiple, false, T>;
template <bool AllowMultiple = false, typename T = internal::Generic>
using SideDestination = DestinationImpl<AllowMultiple, true, T>;
template <typename T = internal::Generic>
using Source = SourceImpl<false, T>;
template <typename T = internal::Generic>
using MultiSource = MultiSourceImpl<false, T>;
template <typename T = internal::Generic>
using SideSource = SourceImpl<true, T>;
template <typename T = internal::Generic>
using MultiSideSource = MultiSourceImpl<true, T>;
template <typename T = internal::Generic>
using Destination = DestinationImpl<false, T>;
template <typename T = internal::Generic>
using SideDestination = DestinationImpl<true, T>;
template <typename T = internal::Generic>
using MultiDestination = MultiDestinationImpl<false, T>;
template <typename T = internal::Generic>
using MultiSideDestination = MultiDestinationImpl<true, T>;
class NodeBase {
public:
@@ -202,45 +212,67 @@ 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).
Source<true> Out(const std::string& tag) {
return Source<true>(&out_streams_[tag]);
MultiSource<> Out(const std::string& tag) {
return MultiSource<>(&out_streams_[tag]);
}
Destination<true> In(const std::string& tag) {
return Destination<true>(&in_streams_[tag]);
MultiDestination<> In(const std::string& tag) {
return MultiDestination<>(&in_streams_[tag]);
}
SideSource<true> SideOut(const std::string& tag) {
return SideSource<true>(&out_sides_[tag]);
MultiSideSource<> SideOut(const std::string& tag) {
return MultiSideSource<>(&out_sides_[tag]);
}
SideDestination<true> SideIn(const std::string& tag) {
return SideDestination<true>(&in_sides_[tag]);
MultiSideDestination<> SideIn(const std::string& tag) {
return MultiSideDestination<>(&in_sides_[tag]);
}
template <typename B, typename T, bool kIsOptional, bool kIsMultiple>
auto operator[](const PortCommon<B, T, kIsOptional, kIsMultiple>& port) {
using PayloadT =
typename PortCommon<B, T, kIsOptional, kIsMultiple>::PayloadT;
if constexpr (std::is_same_v<B, OutputBase>) {
return Source<kIsMultiple, T>(&out_streams_[port.Tag()]);
auto* base = &out_streams_[port.Tag()];
if constexpr (kIsMultiple) {
return MultiSource<PayloadT>(base);
} else {
return Source<PayloadT>(base);
}
} else if constexpr (std::is_same_v<B, InputBase>) {
return Destination<kIsMultiple, T>(&in_streams_[port.Tag()]);
auto* base = &in_streams_[port.Tag()];
if constexpr (kIsMultiple) {
return MultiDestination<PayloadT>(base);
} else {
return Destination<PayloadT>(base);
}
} else if constexpr (std::is_same_v<B, SideOutputBase>) {
return SideSource<kIsMultiple, T>(&out_sides_[port.Tag()]);
auto* base = &out_sides_[port.Tag()];
if constexpr (kIsMultiple) {
return MultiSideSource<PayloadT>(base);
} else {
return SideSource<PayloadT>(base);
}
} else if constexpr (std::is_same_v<B, SideInputBase>) {
return SideDestination<kIsMultiple, T>(&in_sides_[port.Tag()]);
auto* base = &in_sides_[port.Tag()];
if constexpr (kIsMultiple) {
return MultiSideDestination<PayloadT>(base);
} else {
return SideDestination<PayloadT>(base);
}
} else {
static_assert(dependent_false<B>::value, "Type not supported.");
}
}
// Convenience methods for accessing purely index-based ports.
Source<false> Out(int index) { return Out("")[index]; }
Source<> Out(int index) { return Out("")[index]; }
Destination<false> In(int index) { return In("")[index]; }
Destination<> In(int index) { return In("")[index]; }
SideSource<false> SideOut(int index) { return SideOut("")[index]; }
SideSource<> SideOut(int index) { return SideOut("")[index]; }
SideDestination<false> SideIn(int index) { return SideIn("")[index]; }
SideDestination<> SideIn(int index) { return SideIn("")[index]; }
template <typename T>
T& GetOptions() {
@@ -277,11 +309,6 @@ class Node<internal::Generic> : public NodeBase {
using GenericNode = Node<internal::Generic>;
template <template <bool, class> class BP, class Port, class TagIndexMapT>
auto MakeBuilderPort(const Port& port, TagIndexMapT& streams) {
return BP<Port::kMultiple, typename Port::PayloadT>(&streams[port.Tag()]);
}
template <class Calc>
class Node : public NodeBase {
public:
@@ -298,25 +325,25 @@ class Node : public NodeBase {
template <class Tag>
auto Out(Tag tag) {
constexpr auto& port = Calc::Contract::TaggedOutputs::get(tag);
return MakeBuilderPort<Source>(port, out_streams_);
return NodeBase::operator[](port);
}
template <class Tag>
auto In(Tag tag) {
constexpr auto& port = Calc::Contract::TaggedInputs::get(tag);
return MakeBuilderPort<Destination>(port, in_streams_);
return NodeBase::operator[](port);
}
template <class Tag>
auto SideOut(Tag tag) {
constexpr auto& port = Calc::Contract::TaggedSideOutputs::get(tag);
return MakeBuilderPort<SideSource>(port, out_sides_);
return NodeBase::operator[](port);
}
template <class Tag>
auto SideIn(Tag tag) {
constexpr auto& port = Calc::Contract::TaggedSideInputs::get(tag);
return MakeBuilderPort<SideDestination>(port, in_sides_);
return NodeBase::operator[](port);
}
// We could allow using the non-checked versions with typed nodes too, but
@@ -332,17 +359,17 @@ class PacketGenerator {
public:
PacketGenerator(std::string type) : type_(std::move(type)) {}
SideSource<true> SideOut(const std::string& tag) {
return SideSource<true>(&out_sides_[tag]);
MultiSideSource<> SideOut(const std::string& tag) {
return MultiSideSource<>(&out_sides_[tag]);
}
SideDestination<true> SideIn(const std::string& tag) {
return SideDestination<true>(&in_sides_[tag]);
MultiSideDestination<> SideIn(const std::string& tag) {
return MultiSideDestination<>(&in_sides_[tag]);
}
// Convenience methods for accessing purely index-based ports.
SideSource<false> SideOut(int index) { return SideOut("")[index]; }
SideDestination<false> SideIn(int index) { return SideIn("")[index]; }
SideSource<> SideOut(int index) { return SideOut("")[index]; }
SideDestination<> SideIn(int index) { return SideIn("")[index]; }
template <typename T>
T& GetOptions() {
@@ -402,70 +429,85 @@ class Graph {
}
// Graph ports, non-typed.
Source<true> In(const std::string& graph_input) {
MultiSource<> In(const std::string& graph_input) {
return graph_boundary_.Out(graph_input);
}
Destination<true> Out(const std::string& graph_output) {
MultiDestination<> Out(const std::string& graph_output) {
return graph_boundary_.In(graph_output);
}
SideSource<true> SideIn(const std::string& graph_input) {
MultiSideSource<> SideIn(const std::string& graph_input) {
return graph_boundary_.SideOut(graph_input);
}
SideDestination<true> SideOut(const std::string& graph_output) {
MultiSideDestination<> SideOut(const std::string& graph_output) {
return graph_boundary_.SideIn(graph_output);
}
// Convenience methods for accessing purely index-based ports.
Source<false> In(int index) { return In("")[0]; }
Source<> In(int index) { return In("")[index]; }
Destination<false> Out(int index) { return Out("")[0]; }
Destination<> Out(int index) { return Out("")[index]; }
SideSource<false> SideIn(int index) { return SideIn("")[0]; }
SideSource<> SideIn(int index) { return SideIn("")[index]; }
SideDestination<false> SideOut(int index) { return SideOut("")[0]; }
SideDestination<> SideOut(int index) { return SideOut("")[index]; }
// Graph ports, typed.
// TODO: make graph_boundary_ a typed node!
template <class PortT, class Payload = typename PortT::PayloadT,
class Src = Source<PortT::kMultiple, Payload>>
Src In(const PortT& graph_input) {
return Src(&graph_boundary_.out_streams_[graph_input.Tag()]);
template <class PortT, class Payload = typename PortT::PayloadT>
auto In(const PortT& graph_input) {
return (*this)[graph_input];
}
template <class PortT, class Payload = typename PortT::PayloadT,
class Dst = Destination<PortT::kMultiple, Payload>>
Dst Out(const PortT& graph_output) {
return Dst(&graph_boundary_.in_streams_[graph_output.Tag()]);
template <class PortT, class Payload = typename PortT::PayloadT>
auto Out(const PortT& graph_output) {
return (*this)[graph_output];
}
template <class PortT, class Payload = typename PortT::PayloadT,
class Src = SideSource<PortT::kMultiple, Payload>>
Src SideIn(const PortT& graph_input) {
return Src(&graph_boundary_.out_sides_[graph_input.Tag()]);
template <class PortT, class Payload = typename PortT::PayloadT>
auto SideIn(const PortT& graph_input) {
return (*this)[graph_input];
}
template <class PortT, class Payload = typename PortT::PayloadT,
class Dst = SideDestination<PortT::kMultiple, Payload>>
Dst SideOut(const PortT& graph_output) {
return Dst(&graph_boundary_.in_sides_[graph_output.Tag()]);
template <class PortT, class Payload = typename PortT::PayloadT>
auto SideOut(const PortT& graph_output) {
return (*this)[graph_output];
}
template <typename B, typename T, bool kIsOptional, bool kIsMultiple>
auto operator[](const PortCommon<B, T, kIsOptional, kIsMultiple>& port) {
using PayloadT =
typename PortCommon<B, T, kIsOptional, kIsMultiple>::PayloadT;
if constexpr (std::is_same_v<B, OutputBase>) {
return Destination<kIsMultiple, T>(
&graph_boundary_.in_streams_[port.Tag()]);
auto* base = &graph_boundary_.in_streams_[port.Tag()];
if constexpr (kIsMultiple) {
return MultiDestination<PayloadT>(base);
} else {
return Destination<PayloadT>(base);
}
} else if constexpr (std::is_same_v<B, InputBase>) {
return Source<kIsMultiple, T>(&graph_boundary_.out_streams_[port.Tag()]);
auto* base = &graph_boundary_.out_streams_[port.Tag()];
if constexpr (kIsMultiple) {
return MultiSource<PayloadT>(base);
} else {
return Source<PayloadT>(base);
}
} else if constexpr (std::is_same_v<B, SideOutputBase>) {
return SideDestination<kIsMultiple, T>(
&graph_boundary_.in_sides_[port.Tag()]);
auto* base = &graph_boundary_.in_sides_[port.Tag()];
if constexpr (kIsMultiple) {
return MultiSideDestination<PayloadT>(base);
} else {
return SideDestination<PayloadT>(base);
}
} else if constexpr (std::is_same_v<B, SideInputBase>) {
return SideSource<kIsMultiple, T>(
&graph_boundary_.out_sides_[port.Tag()]);
auto* base = &graph_boundary_.out_sides_[port.Tag()];
if constexpr (kIsMultiple) {
return MultiSideSource<PayloadT>(base);
} else {
return SideSource<PayloadT>(base);
}
} else {
static_assert(dependent_false<B>::value, "Type not supported.");
}
mediapipe/framework/api2/builder_test.cc
+42 -15
View File
@@ -50,21 +50,21 @@ TEST(BuilderTest, BuildGraph) {
TEST(BuilderTest, CopyableSource) {
builder::Graph graph;
builder::Source<false, int> a = graph[Input<int>("A")];
builder::Source<int> a = graph[Input<int>("A")];
a.SetName("a");
builder::Source<false, int> b = graph[Input<int>("B")];
builder::Source<int> b = graph[Input<int>("B")];
b.SetName("b");
builder::SideSource<false, float> side_a = graph[SideInput<float>("SIDE_A")];
builder::SideSource<float> side_a = graph[SideInput<float>("SIDE_A")];
side_a.SetName("side_a");
builder::SideSource<false, float> side_b = graph[SideInput<float>("SIDE_B")];
builder::SideSource<float> side_b = graph[SideInput<float>("SIDE_B")];
side_b.SetName("side_b");
builder::Destination<false, int> out = graph[Output<int>("OUT")];
builder::SideDestination<false, float> side_out =
builder::Destination<int> out = graph[Output<int>("OUT")];
builder::SideDestination<float> side_out =
graph[SideOutput<float>("SIDE_OUT")];
builder::Source<false, int> input = a;
builder::Source<int> input = a;
input = b;
builder::SideSource<false, float> side_input = side_b;
builder::SideSource<float> side_input = side_b;
side_input = side_a;
input >> out;
@@ -85,27 +85,26 @@ TEST(BuilderTest, CopyableSource) {
TEST(BuilderTest, BuildGraphWithFunctions) {
builder::Graph graph;
builder::Source<false, int> base = graph[Input<int>("IN")];
builder::Source<int> base = graph[Input<int>("IN")];
base.SetName("base");
builder::SideSource<false, float> side = graph[SideInput<float>("SIDE")];
builder::SideSource<float> side = graph[SideInput<float>("SIDE")];
side.SetName("side");
auto foo_fn = [](builder::Source<false, int> base,
builder::SideSource<false, float> side,
auto foo_fn = [](builder::Source<int> base, builder::SideSource<float> side,
builder::Graph& graph) {
auto& foo = graph.AddNode("Foo");
base >> foo[Input<int>("BASE")];
side >> foo[SideInput<float>("SIDE")];
return foo[Output<double>("OUT")];
};
builder::Source<false, double> foo_out = foo_fn(base, side, graph);
builder::Source<double> foo_out = foo_fn(base, side, graph);
auto bar_fn = [](builder::Source<false, double> in, builder::Graph& graph) {
auto bar_fn = [](builder::Source<double> in, builder::Graph& graph) {
auto& bar = graph.AddNode("Bar");
in >> bar[Input<double>("IN")];
return bar[Output<double>("OUT")];
};
builder::Source<false, double> bar_out = bar_fn(foo_out, graph);
builder::Source<double> bar_out = bar_fn(foo_out, graph);
bar_out.SetName("out");
bar_out >> graph[Output<double>("OUT")];
@@ -298,6 +297,34 @@ TEST(BuilderTest, EmptyTag) {
EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
}
TEST(BuilderTest, GraphIndexes) {
builder::Graph graph;
auto& foo = graph.AddNode("Foo");
graph.In(0).SetName("a") >> foo.In("")[0];
graph.In(1).SetName("c") >> foo.In("")[2];
graph.In(2).SetName("b") >> foo.In("")[1];
foo.Out("")[0].SetName("x") >> graph.Out(1);
foo.Out("")[1].SetName("y") >> graph.Out(0);
CalculatorGraphConfig expected =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
input_stream: "a"
input_stream: "c"
input_stream: "b"
output_stream: "y"
output_stream: "x"
node {
calculator: "Foo"
input_stream: "a"
input_stream: "b"
input_stream: "c"
output_stream: "x"
output_stream: "y"
}
)pb");
EXPECT_THAT(graph.GetConfig(), EqualsProto(expected));
}
} // namespace test
} // namespace api2
} // namespace mediapipe
mediapipe/framework/formats/tensor_internal.h
+28
View File
@@ -0,0 +1,28 @@
#ifndef MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_INTERNAL_H_
#define MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_INTERNAL_H_
#include <cstdint>
namespace mediapipe {
// Generates unique view id at compile-time using FILE and LINE.
#define TENSOR_UNIQUE_VIEW_ID() \
static constexpr uint64_t kId = tensor_internal::FnvHash64( \
__FILE__, tensor_internal::FnvHash64(TENSOR_INT_TO_STRING(__LINE__)))
namespace tensor_internal {
#define TENSOR_INT_TO_STRING2(x) #x
#define TENSOR_INT_TO_STRING(x) TENSOR_INT_TO_STRING2(x)
// Compile-time hash function
// https://en.wikipedia.org/wiki/Fowler%E2%80%93Noll%E2%80%93Vo_hash_function
constexpr uint64_t kFnvPrime = 0x00000100000001B3;
constexpr uint64_t kFnvOffsetBias = 0xcbf29ce484222325;
constexpr uint64_t FnvHash64(const char* str, uint64_t hash = kFnvOffsetBias) {
return (str[0] == 0) ? hash : FnvHash64(str + 1, (hash ^ str[0]) * kFnvPrime);
}
} // namespace tensor_internal
} // namespace mediapipe
#endif // MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_INTERNAL_H_
mediapipe/framework/packet.cc
+17
View File
@@ -21,6 +21,7 @@
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/status_builder.h"
#include "mediapipe/framework/tool/type_util.h"
namespace mediapipe {
namespace packet_internal {
@@ -105,6 +106,22 @@ std::string Packet::DebugString() const {
return result;
}
absl::Status Packet::ValidateAsType(const tool::TypeInfo& type_info) const {
if (ABSL_PREDICT_FALSE(IsEmpty())) {
return absl::InternalError(
absl::StrCat("Expected a Packet of type: ",
MediaPipeTypeStringOrDemangled(type_info),
", but received an empty Packet."));
}
bool holder_is_right_type = holder_->GetTypeId() == type_info.hash_code();
if (ABSL_PREDICT_FALSE(!holder_is_right_type)) {
return absl::InvalidArgumentError(absl::StrCat(
"The Packet stores \"", holder_->DebugTypeName(), "\", but \"",
MediaPipeTypeStringOrDemangled(type_info), "\" was requested."));
}
return absl::OkStatus();
}
absl::Status Packet::ValidateAsProtoMessageLite() const {
if (ABSL_PREDICT_FALSE(IsEmpty())) {
return absl::InternalError("Packet is empty.");
mediapipe/framework/packet.h
+5 -16
View File
@@ -179,7 +179,9 @@ class Packet {
// Returns an error if the packet does not contain data of type T.
template <typename T>
absl::Status ValidateAsType() const;
absl::Status ValidateAsType() const {
return ValidateAsType(tool::TypeId<T>());
}
// Returns an error if the packet is not an instance of
// a protocol buffer message.
@@ -218,6 +220,8 @@ class Packet {
friend std::shared_ptr<packet_internal::HolderBase>
packet_internal::GetHolderShared(Packet&& packet);
absl::Status ValidateAsType(const tool::TypeInfo& type_info) const;
std::shared_ptr<packet_internal::HolderBase> holder_;
class Timestamp timestamp_;
};
@@ -770,21 +774,6 @@ inline const T& Packet::Get() const {
return holder->data();
}
template <typename T>
absl::Status Packet::ValidateAsType() const {
if (ABSL_PREDICT_FALSE(IsEmpty())) {
return absl::InternalError(absl::StrCat(
"Expected a Packet of type: ", MediaPipeTypeStringOrDemangled<T>(),
", but received an empty Packet."));
}
if (ABSL_PREDICT_FALSE(holder_->As<T>() == nullptr)) {
return absl::InvalidArgumentError(absl::StrCat(
"The Packet stores \"", holder_->DebugTypeName(), "\", but \"",
MediaPipeTypeStringOrDemangled<T>(), "\" was requested."));
}
return absl::OkStatus();
}
inline Timestamp Packet::Timestamp() const { return timestamp_; }
template <typename T>
mediapipe/framework/packet_type.h
+5
View File
@@ -84,6 +84,11 @@ class PacketType {
// Returns true iff this and other are consistent, meaning they do
// not expect different types. IsAny() is consistent with anything.
// IsNone() is only consistent with IsNone() and IsAny().
// Note: this is definied as a symmetric relationship, but within the
// framework, it is consistently invoked as:
// input_port_type.IsConsistentWith(connected_output_port_type)
// TODO: consider making this explicitly directional, and
// sharing some logic with the packet validation check.
bool IsConsistentWith(const PacketType& other) const;
// Returns OK if the packet contains an object of the appropriate type.
mediapipe/framework/type_map.h
+10 -4
View File
@@ -373,16 +373,22 @@ inline const std::string* MediaPipeTypeString() {
return MediaPipeTypeStringFromTypeId(tool::GetTypeHash<T>());
}
template <typename T>
const std::string MediaPipeTypeStringOrDemangled() {
const std::string* type_string = MediaPipeTypeString<T>();
inline std::string MediaPipeTypeStringOrDemangled(
const tool::TypeInfo& type_info) {
const std::string* type_string =
MediaPipeTypeStringFromTypeId(type_info.hash_code());
if (type_string) {
return *type_string;
} else {
return mediapipe::Demangle(tool::TypeId<T>().name());
return mediapipe::Demangle(type_info.name());
}
}
template <typename T>
std::string MediaPipeTypeStringOrDemangled() {
return MediaPipeTypeStringOrDemangled(tool::TypeId<T>());
}
// Returns type hash id of type identified by type_string or NULL if not
// registered.
inline const size_t* MediaPipeTypeId(const std::string& type_string) {