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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 6e5aa035cd1f6a9333962df5d3ab97a05bd5744e
This commit is contained in:
MediaPipe Team
2022-06-28 12:11:05 +00:00
committed by Sebastian Schmidt
parent 4a20e9909d
commit c688862570
144 changed files with 5772 additions and 2118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/BUILD
+6 -1
View File
@@ -115,7 +115,10 @@ mediapipe_proto_library(
name = "packet_test_proto",
testonly = 1,
srcs = ["packet_test.proto"],
visibility = ["//mediapipe/framework:__subpackages__"],
visibility = [
":mediapipe_internal",
"//mediapipe/framework:__subpackages__",
],
)
mediapipe_proto_library(
@@ -973,6 +976,7 @@ cc_library(
],
}),
visibility = [
"//fitbit/research/sensing/mobisense:__subpackages__",
"//mediapipe/calculators:__subpackages__",
"//mediapipe/framework:__subpackages__",
"//mediapipe/framework/port:__pkg__",
@@ -1427,6 +1431,7 @@ cc_test(
"//mediapipe/framework/stream_handler:timestamp_align_input_stream_handler",
"//mediapipe/framework/tool:sink",
"//mediapipe/framework/tool:status_util",
"//mediapipe/gpu:graph_support",
"@com_google_absl//absl/container:fixed_array",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
mediapipe/framework/api2/BUILD
+1
View File
@@ -149,6 +149,7 @@ cc_library(
"//mediapipe/framework:calculator_contract",
"//mediapipe/framework:output_side_packet",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/tool:type_util",
"@com_google_absl//absl/strings",
],
)
mediapipe/framework/api2/port.h
+8 -15
View File
@@ -27,41 +27,34 @@
#include "mediapipe/framework/calculator_contract.h"
#include "mediapipe/framework/output_side_packet.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/tool/type_util.h"
namespace mediapipe {
namespace api2 {
// typeid is not constexpr, but a pointer to this is.
template <typename T>
size_t get_type_hash() {
return typeid(T).hash_code();
}
using type_id_fptr = size_t (*)();
// This is a base class for various types of port. It is not meant to be used
// directly by node code.
class PortBase {
public:
constexpr PortBase(std::size_t tag_size, const char* tag,
type_id_fptr get_type_id, bool optional, bool multiple)
constexpr PortBase(std::size_t tag_size, const char* tag, TypeId type_id,
bool optional, bool multiple)
: tag_(tag_size, tag),
optional_(optional),
multiple_(multiple),
type_id_getter_(get_type_id) {}
type_id_(type_id) {}
bool IsOptional() const { return optional_; }
bool IsMultiple() const { return multiple_; }
const char* Tag() const { return tag_.data(); }
size_t type_id() const { return type_id_getter_(); }
TypeId type_id() const { return type_id_; }
const const_str tag_;
const bool optional_;
const bool multiple_;
protected:
type_id_fptr type_id_getter_;
TypeId type_id_;
};
// These four base classes are used to distinguish between ports of different
@@ -340,7 +333,7 @@ class PortCommon : public Base {
template <std::size_t N>
explicit constexpr PortCommon(const char (&tag)[N])
: Base(N, tag, &get_type_hash<ValueT>, IsOptionalV, IsMultipleV) {}
: Base(N, tag, kTypeId<ValueT>, IsOptionalV, IsMultipleV) {}
using PayloadT = ActualPayloadT<ValueT>;
@@ -428,7 +421,7 @@ class SideFallbackT : public Base {
template <std::size_t N>
explicit constexpr SideFallbackT(const char (&tag)[N])
: Base(N, tag, &get_type_hash<ValueT>, IsOptionalV, IsMultipleV),
: Base(N, tag, kTypeId<ValueT>, IsOptionalV, IsMultipleV),
stream_port(tag),
side_port(tag) {}
mediapipe/framework/api2/port_test.cc
+1 -1
View File
@@ -8,7 +8,7 @@ namespace {
TEST(PortTest, IntInput) {
static constexpr auto port = Input<int>("FOO");
EXPECT_EQ(port.type_id(), typeid(int).hash_code());
EXPECT_EQ(port.type_id(), kTypeId<int>);
}
TEST(PortTest, OptionalInput) {
mediapipe/framework/calculator_contract.h
+1 -1
View File
@@ -59,7 +59,7 @@ class CalculatorContract {
const CalculatorOptions& Options() const { return node_config_->options(); }
// Returns the name given to this node.
const std::string& GetNodeName() { return node_name_; }
const std::string& GetNodeName() const { return node_name_; }
// Returns the options given to this calculator. Template argument T must
// be the type of the protobuf extension message or the protobuf::Any
mediapipe/framework/calculator_graph.cc
+11 -9
View File
@@ -120,10 +120,10 @@ CalculatorGraph::CalculatorGraph()
counter_factory_ = absl::make_unique<BasicCounterFactory>();
}
CalculatorGraph::CalculatorGraph(const CalculatorGraphConfig& config)
CalculatorGraph::CalculatorGraph(CalculatorGraphConfig config)
: CalculatorGraph() {
counter_factory_ = absl::make_unique<BasicCounterFactory>();
MEDIAPIPE_CHECK_OK(Initialize(config));
MEDIAPIPE_CHECK_OK(Initialize(std::move(config)));
}
// Defining the destructor here lets us use incomplete types in the header;
@@ -429,18 +429,17 @@ absl::Status CalculatorGraph::Initialize(
return absl::OkStatus();
}
absl::Status CalculatorGraph::Initialize(
const CalculatorGraphConfig& input_config) {
return Initialize(input_config, {});
absl::Status CalculatorGraph::Initialize(CalculatorGraphConfig input_config) {
return Initialize(std::move(input_config), {});
}
absl::Status CalculatorGraph::Initialize(
const CalculatorGraphConfig& input_config,
CalculatorGraphConfig input_config,
const std::map<std::string, Packet>& side_packets) {
auto validated_graph = absl::make_unique<ValidatedGraphConfig>();
MP_RETURN_IF_ERROR(validated_graph->Initialize(
input_config, /*graph_registry=*/nullptr, /*graph_options=*/nullptr,
&service_manager_));
std::move(input_config), /*graph_registry=*/nullptr,
/*graph_options=*/nullptr, &service_manager_));
return Initialize(std::move(validated_graph), side_packets);
}
@@ -675,6 +674,7 @@ absl::Status CalculatorGraph::PrepareForRun(
#endif // !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(PrepareServices());
#if !MEDIAPIPE_DISABLE_GPU
// TODO: should we do this on each run, or only once?
MP_RETURN_IF_ERROR(PrepareGpu());
additional_side_packets = MaybeCreateLegacyGpuSidePacket(legacy_sp);
#endif // !MEDIAPIPE_DISABLE_GPU
@@ -1251,7 +1251,9 @@ void CalculatorGraph::Resume() { scheduler_.Resume(); }
absl::Status CalculatorGraph::SetExecutorInternal(
const std::string& name, std::shared_ptr<Executor> executor) {
if (!executors_.emplace(name, executor).second) {
auto [it, inserted] = executors_.emplace(name, executor);
if (!inserted) {
if (it->second == executor) return absl::OkStatus();
return mediapipe::AlreadyExistsErrorBuilder(MEDIAPIPE_LOC)
<< "SetExecutor must be called only once for the executor \"" << name
<< "\"";
mediapipe/framework/calculator_graph.h
+3 -4
View File
@@ -119,17 +119,17 @@ class CalculatorGraph {
// Initializes the graph from its proto description (using Initialize())
// and crashes if something goes wrong.
explicit CalculatorGraph(const CalculatorGraphConfig& config);
explicit CalculatorGraph(CalculatorGraphConfig config);
virtual ~CalculatorGraph();
// Initializes the graph from a its proto description.
// side_packets that are provided at this stage are common across all Run()
// invocations and could be used to execute PacketGenerators immediately.
absl::Status Initialize(const CalculatorGraphConfig& config,
absl::Status Initialize(CalculatorGraphConfig config,
const std::map<std::string, Packet>& side_packets);
// Convenience version which does not take side packets.
absl::Status Initialize(const CalculatorGraphConfig& config);
absl::Status Initialize(CalculatorGraphConfig config);
// Initializes the CalculatorGraph from the specified graph and subgraph
// configs. Template graph and subgraph configs can be specified through
@@ -272,7 +272,6 @@ class CalculatorGraph {
absl::Status CloseInputStream(const std::string& stream_name);
// Closes all the graph input streams.
// TODO: deprecate this function in favor of CloseAllPacketSources.
absl::Status CloseAllInputStreams();
// Closes all the graph input streams and source calculator nodes.
mediapipe/framework/calculator_graph_test.cc
+21
View File
@@ -60,6 +60,7 @@
#include "mediapipe/framework/tool/sink.h"
#include "mediapipe/framework/tool/status_util.h"
#include "mediapipe/framework/type_map.h"
#include "mediapipe/gpu/graph_support.h"
namespace mediapipe {
@@ -2059,6 +2060,26 @@ TEST(CalculatorGraph, HandlersRun) {
input_side_packets.at("unavailable_input_counter2")));
}
TEST(CalculatorGraph, CalculatorGraphConfigCopyElision) {
CalculatorGraph graph;
CalculatorGraphConfig config =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
input_stream: 'in'
node {
calculator: 'PassThroughCalculator'
input_stream: 'in'
output_stream: 'out'
}
)pb");
// config is consumed and never copied, which avoid copying data.
MP_ASSERT_OK(graph.Initialize(std::move(config)));
MP_EXPECT_OK(graph.StartRun({}));
MP_EXPECT_OK(
graph.AddPacketToInputStream("in", MakePacket<int>(1).At(Timestamp(1))));
MP_EXPECT_OK(graph.CloseInputStream("in"));
MP_EXPECT_OK(graph.WaitUntilDone());
}
// Test that calling SetOffset() in Calculator::Process() results in the
// absl::StatusCode::kFailedPrecondition error.
TEST(CalculatorGraph, SetOffsetInProcess) {
mediapipe/framework/calculator_profile.proto
+1 -4
View File
@@ -11,10 +11,6 @@
// 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.
//
// Forked from mediapipe/framework/calculator_profile.proto.
// The forked proto must remain identical to the original proto and should be
// ONLY used by mediapipe open source project.
syntax = "proto2";
@@ -24,6 +20,7 @@ import "mediapipe/framework/calculator.proto";
option java_package = "com.google.mediapipe.proto";
option java_outer_classname = "CalculatorProfileProto";
option objc_class_prefix = "MediaPipe";
// Stores the profiling information.
//
mediapipe/framework/deps/BUILD
+4 -2
View File
@@ -88,7 +88,10 @@ cc_library(
testonly = True,
hdrs = ["message_matchers.h"],
# Use this library through "mediapipe/framework/port:gtest_main".
visibility = ["//mediapipe/framework/port:__pkg__"],
visibility = [
"//mediapipe/framework/port:__pkg__",
"//third_party/visionai/algorithms/tracking:__pkg__",
],
deps = [
"//mediapipe/framework/port:core_proto",
"@com_google_googletest//:gtest",
@@ -137,7 +140,6 @@ cc_library(
hdrs = ["image_resizer.h"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:opencv_imgproc",
],
)
mediapipe/framework/deps/image_resizer.h
-1
View File
@@ -15,7 +15,6 @@
#ifndef MEDIAPIPE_DEPS_IMAGE_RESIZER_H_
#define MEDIAPIPE_DEPS_IMAGE_RESIZER_H_
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
namespace mediapipe {
mediapipe/framework/encode_binary_proto.bzl
+14 -1
View File
@@ -140,9 +140,22 @@ _encode_binary_proto = rule(
)
def encode_binary_proto(name, input, message_type, deps, **kwargs):
if type(input) == type("string"):
input_label = input
textproto_srcs = [input]
elif type(input) == type(dict()):
# We cannot accept a select, as macros are unable to manipulate selects.
input_label = select(input)
srcs_dict = dict()
for k, v in input.items():
srcs_dict[k] = [v]
textproto_srcs = select(srcs_dict)
else:
fail("input should be a string or a dict, got %s" % input)
_encode_binary_proto(
name = name,
input = input,
input = input_label,
message_type = message_type,
deps = deps,
**kwargs
mediapipe/framework/formats/BUILD
+4
View File
@@ -448,6 +448,7 @@ cc_library(
srcs =
[
"tensor.cc",
"tensor_ahwb.cc",
],
hdrs = ["tensor.h"],
copts = select({
@@ -463,6 +464,9 @@ cc_library(
"-framework MetalKit",
],
"//conditions:default": [],
"//mediapipe:android": [
"-landroid",
],
}),
visibility = ["//visibility:public"],
deps = [
mediapipe/framework/formats/body_rig.proto
+3 -1
View File
@@ -19,7 +19,9 @@ package mediapipe;
// Joint of a 3D human model (e.g. elbow, knee, wrist). Contains 3D rotation of
// the joint and its visibility.
message Joint {
// Joint rotation in 6D contineous representation.
// Joint rotation in 6D contineous representation ordered as
// [a1, b1, a2, b2, a3, b3].
//
// Such representation is more sutable for NN model training and can be
// converted to quaternions and Euler angles if needed. Details can be found
// in https://arxiv.org/abs/1812.07035.
mediapipe/framework/formats/tensor.cc
+99 -37
View File
@@ -20,6 +20,9 @@
#include "absl/synchronization/mutex.h"
#include "mediapipe/framework/port.h"
#include "mediapipe/framework/port/logging.h"
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#include "mediapipe/gpu/gl_base.h"
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#if MEDIAPIPE_METAL_ENABLED
#include <mach/mach_init.h>
@@ -319,28 +322,41 @@ void Tensor::AllocateOpenGlTexture2d() const {
Tensor::OpenGlBufferView Tensor::GetOpenGlBufferReadView() const {
LOG_IF(FATAL, valid_ == kValidNone)
<< "Tensor must be written prior to read from.";
LOG_IF(FATAL, !(valid_ & (kValidCpu | kValidOpenGlBuffer)))
<< "Tensor conversion between different GPU resources is not supported "
"yet.";
LOG_IF(FATAL, !(valid_ & (kValidCpu |
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
kValidAHardwareBuffer |
#endif // MEDIAPIPE_TENSOR_USE_AHWB
kValidOpenGlBuffer)))
<< "Tensor conversion between different GPU resources is not supported.";
auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
AllocateOpenGlBuffer();
if (!(valid_ & kValidOpenGlBuffer)) {
glBindBuffer(GL_SHADER_STORAGE_BUFFER, opengl_buffer_);
void* ptr =
glMapBufferRange(GL_SHADER_STORAGE_BUFFER, 0, bytes(),
GL_MAP_INVALIDATE_BUFFER_BIT | GL_MAP_WRITE_BIT);
std::memcpy(ptr, cpu_buffer_, bytes());
glUnmapBuffer(GL_SHADER_STORAGE_BUFFER);
// If the call succeds then AHWB -> SSBO are synchronized so any usage of
// the SSBO is correct after this call.
if (!InsertAhwbToSsboFence()) {
glBindBuffer(GL_SHADER_STORAGE_BUFFER, opengl_buffer_);
void* ptr =
glMapBufferRange(GL_SHADER_STORAGE_BUFFER, 0, bytes(),
GL_MAP_INVALIDATE_BUFFER_BIT | GL_MAP_WRITE_BIT);
std::memcpy(ptr, cpu_buffer_, bytes());
glUnmapBuffer(GL_SHADER_STORAGE_BUFFER);
}
valid_ |= kValidOpenGlBuffer;
}
return {opengl_buffer_, std::move(lock)};
return {opengl_buffer_, std::move(lock),
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
&ssbo_read_
#else
nullptr
#endif // MEDIAPIPE_TENSOR_USE_AHWB
};
}
Tensor::OpenGlBufferView Tensor::GetOpenGlBufferWriteView() const {
auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
AllocateOpenGlBuffer();
valid_ = kValidOpenGlBuffer;
return {opengl_buffer_, std::move(lock)};
return {opengl_buffer_, std::move(lock), nullptr};
}
void Tensor::AllocateOpenGlBuffer() const {
@@ -349,7 +365,10 @@ void Tensor::AllocateOpenGlBuffer() const {
LOG_IF(FATAL, !gl_context_) << "GlContext is not bound to the thread.";
glGenBuffers(1, &opengl_buffer_);
glBindBuffer(GL_SHADER_STORAGE_BUFFER, opengl_buffer_);
glBufferData(GL_SHADER_STORAGE_BUFFER, bytes(), NULL, GL_STREAM_COPY);
if (!AllocateAhwbMapToSsbo()) {
glBufferData(GL_SHADER_STORAGE_BUFFER, bytes(), NULL, GL_STREAM_COPY);
}
glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
}
}
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
@@ -377,6 +396,8 @@ void Tensor::Move(Tensor* src) {
src->metal_buffer_ = nil;
#endif // MEDIAPIPE_METAL_ENABLED
MoveAhwbStuff(src);
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
gl_context_ = std::move(src->gl_context_);
frame_buffer_ = src->frame_buffer_;
@@ -395,27 +416,31 @@ void Tensor::Move(Tensor* src) {
Tensor::Tensor(ElementType element_type, const Shape& shape)
: element_type_(element_type), shape_(shape) {}
#if MEDIAPIPE_METAL_ENABLED
void Tensor::Invalidate() {
absl::MutexLock lock(&view_mutex_);
// If memory is allocated and not owned by the metal buffer.
// TODO: Re-design cpu buffer memory management.
if (cpu_buffer_ && !metal_buffer_) {
DeallocateVirtualMemory(cpu_buffer_, AlignToPageSize(bytes()));
}
metal_buffer_ = nil;
cpu_buffer_ = nullptr;
}
#else
void Tensor::Invalidate() {
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
GLuint cleanup_gl_tex = GL_INVALID_INDEX;
GLuint cleanup_gl_fb = GL_INVALID_INDEX;
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
GLuint cleanup_gl_buf = GL_INVALID_INDEX;
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
{
absl::MutexLock lock(&view_mutex_);
#if MEDIAPIPE_METAL_ENABLED
// If memory is allocated and not owned by the metal buffer.
// TODO: Re-design cpu buffer memory management.
if (cpu_buffer_ && !metal_buffer_) {
DeallocateVirtualMemory(cpu_buffer_, AlignToPageSize(bytes()));
}
metal_buffer_ = nil;
#else
if (cpu_buffer_) {
free(cpu_buffer_);
}
#endif // MEDIAPIPE_METAL_ENABLED
cpu_buffer_ = nullptr;
ReleaseAhwbStuff();
// Don't need to wait for the resource to be deleted bacause if will be
// released on last reference deletion inside the OpenGL driver.
@@ -429,28 +454,44 @@ void Tensor::Invalidate() {
}
// Do not hold the view mutex while invoking GlContext::RunWithoutWaiting,
// since that method may acquire the context's own lock.
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
if (cleanup_gl_tex != GL_INVALID_INDEX || cleanup_gl_fb != GL_INVALID_INDEX ||
cleanup_gl_buf != GL_INVALID_INDEX)
gl_context_->RunWithoutWaiting([cleanup_gl_tex, cleanup_gl_fb
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
,
cleanup_gl_buf
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
]() {
if (cleanup_gl_tex != GL_INVALID_INDEX || cleanup_gl_fb != GL_INVALID_INDEX ||
cleanup_gl_buf != GL_INVALID_INDEX) {
gl_context_->RunWithoutWaiting(
[cleanup_gl_tex, cleanup_gl_fb, cleanup_gl_buf]() {
glDeleteTextures(1, &cleanup_gl_tex);
glDeleteFramebuffers(1, &cleanup_gl_fb);
glDeleteBuffers(1, &cleanup_gl_buf);
});
}
#elif MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
if (cleanup_gl_tex != GL_INVALID_INDEX || cleanup_gl_fb != GL_INVALID_INDEX) {
gl_context_->RunWithoutWaiting([cleanup_gl_tex, cleanup_gl_fb]() {
glDeleteTextures(1, &cleanup_gl_tex);
glDeleteFramebuffers(1, &cleanup_gl_fb);
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
glDeleteBuffers(1, &cleanup_gl_buf);
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
});
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
}
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
if (cpu_buffer_) {
free(cpu_buffer_);
}
cpu_buffer_ = nullptr;
}
#endif // MEDIAPIPE_METAL_ENABLED
Tensor::CpuReadView Tensor::GetCpuReadView() const {
auto lock = absl::make_unique<absl::MutexLock>(&view_mutex_);
LOG_IF(FATAL, valid_ == kValidNone)
<< "Tensor must be written prior to read from.";
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
void* ptr = MapAhwbToCpuRead();
if (ptr) {
valid_ |= kValidCpu;
return {ptr, ahwb_, nullptr, std::move(lock)};
}
#endif // MEDIAPIPE_TENSOR_USE_AHWB
AllocateCpuBuffer();
if (!(valid_ & kValidCpu)) {
// GPU-to-CPU synchronization and read-back.
@@ -512,18 +553,33 @@ Tensor::CpuReadView Tensor::GetCpuReadView() const {
#endif // MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
valid_ |= kValidCpu;
}
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
return {cpu_buffer_, nullptr, nullptr, std::move(lock)};
#else
return {cpu_buffer_, std::move(lock)};
#endif // MEDIAPIPE_TENSOR_USE_AHWB
}
Tensor::CpuWriteView Tensor::GetCpuWriteView() const {
auto lock = absl::make_unique<absl::MutexLock>(&view_mutex_);
AllocateCpuBuffer();
valid_ = kValidCpu;
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
void* ptr = MapAhwbToCpuWrite();
if (ptr) {
return {ptr, ahwb_, &fence_fd_, std::move(lock)};
}
return {cpu_buffer_, nullptr, nullptr, std::move(lock)};
#else
return {cpu_buffer_, std::move(lock)};
#endif // MEDIAPIPE_TENSOR_USE_AHWB
}
void Tensor::AllocateCpuBuffer() const {
if (!cpu_buffer_) {
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
if (AllocateAHardwareBuffer()) return;
#endif // MEDIAPIPE_TENSOR_USE_AHWB
#if MEDIAPIPE_METAL_ENABLED
cpu_buffer_ = AllocateVirtualMemory(bytes());
#else
@@ -532,4 +588,10 @@ void Tensor::AllocateCpuBuffer() const {
}
}
void Tensor::SetPreferredStorageType(StorageType type) {
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
use_ahwb_ = type == StorageType::kAhwb;
#endif // MEDIAPIPE_TENSOR_USE_AHWB
}
} // namespace mediapipe
mediapipe/framework/formats/tensor.h
+144 -11
View File
@@ -30,6 +30,16 @@
#import <Metal/Metal.h>
#endif // MEDIAPIPE_METAL_ENABLED
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
#if __ANDROID_API__ < 26
#error MEDIAPIPE_TENSOR_USE_AHWB requires NDK version 26 or higher to be specified.
#endif // __ANDROID_API__ < 26
#include <android/hardware_buffer.h>
#include "third_party/GL/gl/include/EGL/egl.h"
#include "third_party/GL/gl/include/EGL/eglext.h"
#endif // MEDIAPIPE_TENSOR_USE_AHWB
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
#include "mediapipe/gpu/gl_base.h"
#include "mediapipe/gpu/gl_context.h"
@@ -108,14 +118,37 @@ class Tensor {
return static_cast<typename std::tuple_element<
std::is_const<T>::value, std::tuple<P*, const P*> >::type>(buffer_);
}
CpuView(CpuView&& src) : View(std::move(src)), buffer_(src.buffer_) {
src.buffer_ = nullptr;
CpuView(CpuView&& src) : View(std::move(src)) {
buffer_ = std::exchange(src.buffer_, nullptr);
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
ahwb_ = std::exchange(src.ahwb_, nullptr);
fence_fd_ = std::exchange(src.fence_fd_, nullptr);
#endif // MEDIAPIPE_TENSOR_USE_AHWB
}
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
~CpuView() {
if (ahwb_) {
auto error = AHardwareBuffer_unlock(ahwb_, fence_fd_);
CHECK(error == 0) << "AHardwareBuffer_unlock " << error;
}
}
#endif // MEDIAPIPE_TENSOR_USE_AHWB
protected:
friend class Tensor;
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
CpuView(T* buffer, AHardwareBuffer* ahwb, int* fence_fd,
std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)),
buffer_(buffer),
fence_fd_(fence_fd),
ahwb_(ahwb) {}
AHardwareBuffer* ahwb_;
int* fence_fd_;
#else
CpuView(T* buffer, std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)), buffer_(buffer) {}
#endif // MEDIAPIPE_TENSOR_USE_AHWB
T* buffer_;
};
using CpuReadView = CpuView<const void>;
@@ -150,6 +183,60 @@ class Tensor {
MtlBufferView GetMtlBufferWriteView(id<MTLDevice> device) const;
#endif // MEDIAPIPE_METAL_ENABLED
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
class AHardwareBufferView : public View {
public:
AHardwareBuffer* handle() const { return handle_; }
AHardwareBufferView(AHardwareBufferView&& src) : View(std::move(src)) {
handle_ = std::exchange(src.handle_, nullptr);
file_descriptor_ = src.file_descriptor_;
fence_fd_ = std::exchange(src.fence_fd_, nullptr);
ahwb_written_ = std::exchange(src.ahwb_written_, nullptr);
release_callback_ = std::exchange(src.release_callback_, nullptr);
}
int file_descriptor() const { return file_descriptor_; }
void SetReadingFinishedFunc(std::function<bool()>&& func) {
CHECK(ahwb_written_)
<< "AHWB write view can't accept 'reading finished callback'";
*ahwb_written_ = std::move(func);
}
void SetWritingFinishedFD(int fd) {
CHECK(fence_fd_)
<< "AHWB read view can't accept 'writing finished file descriptor'";
*fence_fd_ = fd;
}
// The function is called when the tensor is released.
void SetReleaseCallback(std::function<void()> callback) {
*release_callback_ = std::move(callback);
}
protected:
friend class Tensor;
AHardwareBufferView(AHardwareBuffer* handle, int file_descriptor,
int* fence_fd, std::function<bool()>* ahwb_written,
std::function<void()>* release_callback,
std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)),
handle_(handle),
file_descriptor_(file_descriptor),
fence_fd_(fence_fd),
ahwb_written_(ahwb_written),
release_callback_(release_callback) {}
AHardwareBuffer* handle_;
int file_descriptor_;
// The view sets some Tensor's fields. The view is released prior to tensor.
int* fence_fd_;
std::function<bool()>* ahwb_written_;
std::function<void()>* release_callback_;
};
AHardwareBufferView GetAHardwareBufferReadView() const;
// size_alignment is an optional argument to tell the API to allocate
// a buffer that is padded to multiples of size_alignment bytes.
// size_alignment must be power of 2, i.e. 2, 4, 8, 16, 64, etc.
// If size_alignment is 0, then the buffer will not be padded.
AHardwareBufferView GetAHardwareBufferWriteView(int size_alignment = 0) const;
#endif // MEDIAPIPE_TENSOR_USE_AHWB
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
// TODO: Use GlTextureView instead.
// Only float32 textures are supported with 1/2/3/4 depths.
@@ -188,16 +275,23 @@ class Tensor {
class OpenGlBufferView : public View {
public:
GLuint name() const { return name_; }
OpenGlBufferView(OpenGlBufferView&& src)
: View(std::move(src)), name_(src.name_) {
src.name_ = GL_INVALID_INDEX;
OpenGlBufferView(OpenGlBufferView&& src) : View(std::move(src)) {
name_ = std::exchange(src.name_, GL_INVALID_INDEX);
ssbo_read_ = std::exchange(src.ssbo_read_, nullptr);
}
~OpenGlBufferView() {
if (ssbo_read_) {
*ssbo_read_ = glFenceSync(GL_SYNC_GPU_COMMANDS_COMPLETE, 0);
}
}
protected:
friend class Tensor;
OpenGlBufferView(GLuint name, std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)), name_(name) {}
OpenGlBufferView(GLuint name, std::unique_ptr<absl::MutexLock>&& lock,
GLsync* ssbo_read)
: View(std::move(lock)), name_(name), ssbo_read_(ssbo_read) {}
GLuint name_;
GLsync* ssbo_read_;
};
// A valid OpenGL context must be bound to the calling thread due to possible
// GPU resource allocation.
@@ -223,16 +317,26 @@ class Tensor {
}
int bytes() const { return shape_.num_elements() * element_size(); }
bool ready_on_cpu() const { return valid_ & kValidCpu; }
bool ready_on_cpu() const {
return valid_ & (kValidAHardwareBuffer | kValidCpu);
}
bool ready_on_gpu() const {
return valid_ &
(kValidMetalBuffer | kValidOpenGlBuffer | kValidOpenGlTexture2d);
return valid_ & (kValidMetalBuffer | kValidOpenGlBuffer |
kValidAHardwareBuffer | kValidOpenGlTexture2d);
}
bool ready_as_metal_buffer() const { return valid_ & kValidMetalBuffer; }
bool ready_as_opengl_buffer() const { return valid_ & kValidOpenGlBuffer; }
bool ready_as_opengl_buffer() const {
return valid_ & (kValidAHardwareBuffer | kValidOpenGlBuffer);
}
bool ready_as_opengl_texture_2d() const {
return valid_ & kValidOpenGlTexture2d;
}
// Sets the type of underlying resource that is going to be allocated.
enum class StorageType {
kDefault,
kAhwb,
};
static void SetPreferredStorageType(StorageType type);
private:
void Move(Tensor*);
@@ -248,6 +352,7 @@ class Tensor {
kValidMetalBuffer = 1 << 1,
kValidOpenGlBuffer = 1 << 2,
kValidOpenGlTexture2d = 1 << 3,
kValidAHardwareBuffer = 1 << 5,
};
// A list of resource which are currently allocated and synchronized between
// each-other: valid_ = kValidCpu | kValidMetalBuffer;
@@ -264,6 +369,34 @@ class Tensor {
void AllocateMtlBuffer(id<MTLDevice> device) const;
#endif // MEDIAPIPE_METAL_ENABLED
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
mutable AHardwareBuffer* ahwb_ = nullptr;
// Signals when GPU finished writing into SSBO so AHWB can be used then. Or
// signals when writing into AHWB has been finished so GPU can read from SSBO.
// Sync and FD are bound together.
mutable EGLSyncKHR fence_sync_ = EGL_NO_SYNC_KHR;
// This FD signals when the writing into the SSBO has been finished.
mutable int ssbo_written_ = -1;
// An externally set FD that is wrapped with the EGL sync then to synchronize
// AHWB -> OpenGL SSBO.
mutable int fence_fd_ = -1;
// Reading from SSBO has been finished so SSBO can be released.
mutable GLsync ssbo_read_ = 0;
// An externally set function that signals when it is safe to release AHWB.
mutable std::function<bool()> ahwb_written_;
mutable std::function<void()> release_callback_;
bool AllocateAHardwareBuffer(int size_alignment = 0) const;
void CreateEglSyncAndFd() const;
// Use Ahwb for other views: OpenGL / CPU buffer.
static inline bool use_ahwb_ = false;
#endif // MEDIAPIPE_TENSOR_USE_AHWB
bool AllocateAhwbMapToSsbo() const;
bool InsertAhwbToSsboFence() const;
void MoveAhwbStuff(Tensor* src);
void ReleaseAhwbStuff();
void* MapAhwbToCpuRead() const;
void* MapAhwbToCpuWrite() const;
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
mutable std::shared_ptr<mediapipe::GlContext> gl_context_;
mutable GLuint opengl_texture2d_ = GL_INVALID_INDEX;
mediapipe/framework/formats/tensor_ahwb.cc
+382
View File
@@ -0,0 +1,382 @@
#include <cstdint>
#include <utility>
#include "mediapipe/framework/formats/tensor.h"
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
#include "absl/synchronization/mutex.h"
#include "mediapipe/framework/port.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/gpu/gl_base.h"
#include "third_party/GL/gl/include/EGL/egl.h"
#include "third_party/GL/gl/include/EGL/eglext.h"
#endif // MEDIAPIPE_TENSOR_USE_AHWB
namespace mediapipe {
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
namespace {
PFNGLBUFFERSTORAGEEXTERNALEXTPROC glBufferStorageExternalEXT;
PFNEGLGETNATIVECLIENTBUFFERANDROIDPROC eglGetNativeClientBufferANDROID;
PFNEGLDUPNATIVEFENCEFDANDROIDPROC eglDupNativeFenceFDANDROID;
PFNEGLCREATESYNCKHRPROC eglCreateSyncKHR;
PFNEGLWAITSYNCKHRPROC eglWaitSyncKHR;
PFNEGLCLIENTWAITSYNCKHRPROC eglClientWaitSyncKHR;
PFNEGLDESTROYSYNCKHRPROC eglDestroySyncKHR;
bool IsGlSupported() {
static const bool extensions_allowed = [] {
eglGetNativeClientBufferANDROID =
reinterpret_cast<PFNEGLGETNATIVECLIENTBUFFERANDROIDPROC>(
eglGetProcAddress("eglGetNativeClientBufferANDROID"));
glBufferStorageExternalEXT =
reinterpret_cast<PFNGLBUFFERSTORAGEEXTERNALEXTPROC>(
eglGetProcAddress("glBufferStorageExternalEXT"));
eglDupNativeFenceFDANDROID =
reinterpret_cast<PFNEGLDUPNATIVEFENCEFDANDROIDPROC>(
eglGetProcAddress("eglDupNativeFenceFDANDROID"));
eglCreateSyncKHR = reinterpret_cast<PFNEGLCREATESYNCKHRPROC>(
eglGetProcAddress("eglCreateSyncKHR"));
eglWaitSyncKHR = reinterpret_cast<PFNEGLWAITSYNCKHRPROC>(
eglGetProcAddress("eglWaitSyncKHR"));
eglClientWaitSyncKHR = reinterpret_cast<PFNEGLCLIENTWAITSYNCKHRPROC>(
eglGetProcAddress("eglClientWaitSyncKHR"));
eglDestroySyncKHR = reinterpret_cast<PFNEGLDESTROYSYNCKHRPROC>(
eglGetProcAddress("eglDestroySyncKHR"));
return eglClientWaitSyncKHR && eglWaitSyncKHR &&
eglGetNativeClientBufferANDROID && glBufferStorageExternalEXT &&
eglCreateSyncKHR && eglDupNativeFenceFDANDROID && eglDestroySyncKHR;
}();
return extensions_allowed;
}
absl::Status MapAHardwareBufferToGlBuffer(AHardwareBuffer* handle, size_t size,
GLuint name) {
if (!IsGlSupported()) {
return absl::UnknownError(
"No GL extension functions found to bind AHardwareBuffer and "
"OpenGL buffer");
}
EGLClientBuffer native_buffer = eglGetNativeClientBufferANDROID(handle);
if (!native_buffer) {
return absl::UnknownError("Can't get native buffer");
}
glBufferStorageExternalEXT(GL_SHADER_STORAGE_BUFFER, 0, size, native_buffer,
GL_MAP_READ_BIT | GL_MAP_WRITE_BIT |
GL_MAP_COHERENT_BIT_EXT |
GL_MAP_PERSISTENT_BIT_EXT);
if (glGetError() == GL_NO_ERROR) {
return absl::OkStatus();
} else {
return absl::InternalError("Error in glBufferStorageExternalEXT");
}
}
static inline int AlignedToPowerOf2(int value, int alignment) {
// alignment must be a power of 2
return ((value - 1) | (alignment - 1)) + 1;
}
// This class keeps tensor's resources while the tensor is in use on GPU or TPU
// but is already released on CPU. When a regular OpenGL buffer is bound to the
// GPU queue for execution and released on client side then the buffer is still
// not released because is being used by GPU. OpenGL driver keeps traking of
// that. When OpenGL buffer is build on top of AHWB then the traking is done
// with the DeleyedRelease which, actually, keeps record of all AHWBs allocated
// and releases each of them if already used. EGL/GL fences are used to check
// the status of a buffer.
class DelayedReleaser {
public:
// Non-copyable
DelayedReleaser(const DelayedReleaser&) = delete;
DelayedReleaser& operator=(const DelayedReleaser&) = delete;
// Non-movable
DelayedReleaser(DelayedReleaser&&) = delete;
DelayedReleaser& operator=(DelayedReleaser&&) = delete;
static void Add(AHardwareBuffer* ahwb, GLuint opengl_buffer,
EGLSyncKHR ssbo_sync, GLsync ssbo_read,
std::function<bool()>&& ahwb_written,
std::shared_ptr<mediapipe::GlContext> gl_context,
std::function<void()>&& callback) {
static absl::Mutex mutex;
absl::MutexLock lock(&mutex);
// Using `new` to access a non-public constructor.
to_release_.emplace_back(absl::WrapUnique(new DelayedReleaser(
ahwb, opengl_buffer, ssbo_sync, ssbo_read, std::move(ahwb_written),
gl_context, std::move(callback))));
for (auto it = to_release_.begin(); it != to_release_.end();) {
if ((*it)->IsSignaled()) {
it = to_release_.erase(it);
} else {
++it;
}
}
}
~DelayedReleaser() {
AHardwareBuffer_release(ahwb_);
if (release_callback_) release_callback_();
}
bool IsSignaled() {
CHECK(!(ssbo_read_ && ahwb_written_))
<< "ssbo_read_ and ahwb_written_ cannot both be set";
if (ahwb_written_) {
if (!ahwb_written_()) return false;
gl_context_->Run([this]() {
if (fence_sync_ != EGL_NO_SYNC_KHR && IsGlSupported()) {
auto egl_display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
if (egl_display != EGL_NO_DISPLAY) {
eglDestroySyncKHR(egl_display, fence_sync_);
}
fence_sync_ = EGL_NO_SYNC_KHR;
}
glDeleteBuffers(1, &opengl_buffer_);
opengl_buffer_ = GL_INVALID_INDEX;
});
return true;
}
gl_context_->Run([this]() {
if (ssbo_read_ != 0) {
GLenum status = glClientWaitSync(ssbo_read_, 0,
/* timeout ns = */ 0);
if (status != GL_CONDITION_SATISFIED && status != GL_ALREADY_SIGNALED) {
return;
}
glDeleteSync(ssbo_read_);
ssbo_read_ = 0;
// Don't wait on ssbo_sync because it is ahead of ssbo_read_sync.
if (fence_sync_ != EGL_NO_SYNC_KHR && IsGlSupported()) {
auto egl_display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
if (egl_display != EGL_NO_DISPLAY) {
eglDestroySyncKHR(egl_display, fence_sync_);
}
}
fence_sync_ = EGL_NO_SYNC_KHR;
glDeleteBuffers(1, &opengl_buffer_);
opengl_buffer_ = GL_INVALID_INDEX;
}
});
return opengl_buffer_ == GL_INVALID_INDEX;
}
protected:
AHardwareBuffer* ahwb_;
GLuint opengl_buffer_;
// TODO: use wrapper instead.
EGLSyncKHR fence_sync_;
// TODO: use wrapper instead.
GLsync ssbo_read_;
std::function<bool()> ahwb_written_;
std::shared_ptr<mediapipe::GlContext> gl_context_;
std::function<void()> release_callback_;
static inline std::deque<std::unique_ptr<DelayedReleaser>> to_release_;
DelayedReleaser(AHardwareBuffer* ahwb, GLuint opengl_buffer,
EGLSyncKHR fence_sync, GLsync ssbo_read,
std::function<bool()>&& ahwb_written,
std::shared_ptr<mediapipe::GlContext> gl_context,
std::function<void()>&& callback)
: ahwb_(ahwb),
opengl_buffer_(opengl_buffer),
fence_sync_(fence_sync),
ssbo_read_(ssbo_read),
ahwb_written_(std::move(ahwb_written)),
gl_context_(gl_context),
release_callback_(std::move(callback)) {}
};
} // namespace
Tensor::AHardwareBufferView Tensor::GetAHardwareBufferReadView() const {
auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
CHECK(valid_ != kValidNone) << "Tensor must be written prior to read from.";
CHECK(!(valid_ & kValidOpenGlTexture2d))
<< "Tensor conversion between OpenGL texture and AHardwareBuffer is not "
"supported.";
CHECK(ahwb_ || !(valid_ & kValidOpenGlBuffer))
<< "Interoperability bettween OpenGL buffer and AHardwareBuffer is not "
"supported on targe system.";
CHECK(AllocateAHardwareBuffer())
<< "AHardwareBuffer is not supported on the target system.";
valid_ |= kValidAHardwareBuffer;
if (valid_ & kValidOpenGlBuffer) CreateEglSyncAndFd();
return {ahwb_,
ssbo_written_,
&fence_fd_, // The FD is created for SSBO -> AHWB synchronization.
&ahwb_written_, // Filled by SetReadingFinishedFunc.
&release_callback_,
std::move(lock)};
}
void Tensor::CreateEglSyncAndFd() const {
gl_context_->Run([this]() {
if (IsGlSupported()) {
auto egl_display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
if (egl_display != EGL_NO_DISPLAY) {
fence_sync_ = eglCreateSyncKHR(egl_display,
EGL_SYNC_NATIVE_FENCE_ANDROID, nullptr);
if (fence_sync_ != EGL_NO_SYNC_KHR) {
ssbo_written_ = eglDupNativeFenceFDANDROID(egl_display, fence_sync_);
if (ssbo_written_ == -1) {
eglDestroySyncKHR(egl_display, fence_sync_);
fence_sync_ = EGL_NO_SYNC_KHR;
}
}
}
}
// Can't use Sync object.
if (fence_sync_ == EGL_NO_SYNC_KHR) glFinish();
});
}
Tensor::AHardwareBufferView Tensor::GetAHardwareBufferWriteView(
int size_alignment) const {
auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
CHECK(AllocateAHardwareBuffer(size_alignment))
<< "AHardwareBuffer is not supported on the target system.";
valid_ = kValidAHardwareBuffer;
return {ahwb_,
/*ssbo_written=*/-1,
&fence_fd_, // For SetWritingFinishedFD.
/*ahwb_written=*/nullptr, // The lifetime is managed by SSBO.
&release_callback_,
std::move(lock)};
}
bool Tensor::AllocateAHardwareBuffer(int size_alignment) const {
if (!use_ahwb_) return false;
if (ahwb_ == nullptr) {
AHardwareBuffer_Desc desc = {};
if (size_alignment == 0) {
desc.width = bytes();
} else {
// We expect allocations to be page-aligned, implicitly satisfying any
// requirements from Edge TPU. No need to add a check for this,
// since Edge TPU will check for us.
desc.width = AlignedToPowerOf2(bytes(), size_alignment);
}
desc.height = 1;
desc.layers = 1;
desc.format = AHARDWAREBUFFER_FORMAT_BLOB;
desc.usage = AHARDWAREBUFFER_USAGE_CPU_WRITE_OFTEN |
AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN |
AHARDWAREBUFFER_USAGE_GPU_DATA_BUFFER;
return AHardwareBuffer_allocate(&desc, &ahwb_) == 0;
}
return true;
}
bool Tensor::AllocateAhwbMapToSsbo() const {
if (AllocateAHardwareBuffer()) {
if (MapAHardwareBufferToGlBuffer(ahwb_, bytes(), opengl_buffer_).ok()) {
glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
return true;
}
// Unable to make OpenGL <-> AHWB binding. Use regular SSBO instead.
AHardwareBuffer_release(ahwb_);
ahwb_ = nullptr;
}
return false;
}
// SSBO is created on top of AHWB. A fence is inserted into the GPU queue before
// the GPU task that is going to read from the SSBO. When the writing into AHWB
// is finished then the GPU reads from the SSBO.
bool Tensor::InsertAhwbToSsboFence() const {
if (!ahwb_) return false;
if (fence_fd_ != -1) {
// Can't wait for FD to be signaled on GPU.
// TODO: wait on CPU instead.
if (!IsGlSupported()) return true;
// Server-side fence.
auto egl_display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
if (egl_display == EGL_NO_DISPLAY) return true;
EGLint sync_attribs[] = {EGL_SYNC_NATIVE_FENCE_FD_ANDROID,
(EGLint)fence_fd_, EGL_NONE};
fence_sync_ = eglCreateSyncKHR(egl_display, EGL_SYNC_NATIVE_FENCE_ANDROID,
sync_attribs);
if (fence_sync_ != EGL_NO_SYNC_KHR) {
eglWaitSyncKHR(egl_display, fence_sync_, 0);
}
}
return true;
}
void Tensor::MoveAhwbStuff(Tensor* src) {
ahwb_ = std::exchange(src->ahwb_, nullptr);
fence_sync_ = std::exchange(src->fence_sync_, EGL_NO_SYNC_KHR);
ssbo_read_ = std::exchange(src->ssbo_read_, static_cast<GLsync>(0));
ssbo_written_ = std::exchange(src->ssbo_written_, -1);
fence_fd_ = std::exchange(src->fence_fd_, -1);
ahwb_written_ = std::move(src->ahwb_written_);
release_callback_ = std::move(src->release_callback_);
}
void Tensor::ReleaseAhwbStuff() {
if (fence_fd_ != -1) {
close(fence_fd_);
fence_fd_ = -1;
}
if (ahwb_) {
if (ssbo_read_ != 0 || fence_sync_ != EGL_NO_SYNC_KHR) {
if (ssbo_written_ != -1) close(ssbo_written_);
DelayedReleaser::Add(ahwb_, opengl_buffer_, fence_sync_, ssbo_read_,
std::move(ahwb_written_), gl_context_,
std::move(release_callback_));
opengl_buffer_ = GL_INVALID_INDEX;
} else {
AHardwareBuffer_release(ahwb_);
}
}
}
void* Tensor::MapAhwbToCpuRead() const {
if (ahwb_) {
if (!(valid_ & kValidCpu) && (valid_ & kValidOpenGlBuffer) &&
ssbo_written_ == -1) {
// EGLSync is failed. Use another synchronization method.
// TODO: Use tflite::gpu::GlBufferSync and GlActiveSync.
glFinish();
}
void* ptr;
auto error =
AHardwareBuffer_lock(ahwb_, AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN,
ssbo_written_, nullptr, &ptr);
CHECK(error == 0) << "AHardwareBuffer_lock " << error;
close(ssbo_written_);
ssbo_written_ = -1;
return ptr;
}
return nullptr;
}
void* Tensor::MapAhwbToCpuWrite() const {
if (ahwb_) {
// TODO: If previously acquired view is GPU write view then need to
// be sure that writing is finished. That's a warning: two consequent write
// views should be interleaved with read view.
void* ptr;
auto error = AHardwareBuffer_lock(
ahwb_, AHARDWAREBUFFER_USAGE_CPU_WRITE_OFTEN, -1, nullptr, &ptr);
CHECK(error == 0) << "AHardwareBuffer_lock " << error;
return ptr;
}
return nullptr;
}
#else // MEDIAPIPE_TENSOR_USE_AHWB
bool Tensor::AllocateAhwbMapToSsbo() const { return false; }
bool Tensor::InsertAhwbToSsboFence() const { return false; }
void Tensor::MoveAhwbStuff(Tensor* src) {}
void Tensor::ReleaseAhwbStuff() {}
void* Tensor::MapAhwbToCpuRead() const { return nullptr; }
void* Tensor::MapAhwbToCpuWrite() const { return nullptr; }
#endif // MEDIAPIPE_TENSOR_USE_AHWB
} // namespace mediapipe
mediapipe/framework/formats/tensor_ahwb_test.cc
+59
View File
@@ -0,0 +1,59 @@
#include "mediapipe/gpu/gpu_test_base.h"
#include "testing/base/public/gmock.h"
#include "testing/base/public/gunit.h"
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
#include <android/hardware_buffer.h>
#include "mediapipe/framework/formats/tensor.h"
namespace mediapipe {
#if !MEDIAPIPE_DISABLE_GPU
class TensorAhwbTest : public mediapipe::GpuTestBase {
public:
};
TEST_F(TensorAhwbTest, TestCpuThenAHWB) {
Tensor tensor(Tensor::ElementType::kFloat32, Tensor::Shape{1});
{
auto ptr = tensor.GetCpuWriteView().buffer<float>();
EXPECT_NE(ptr, nullptr);
}
{
auto ahwb = tensor.GetAHardwareBufferReadView().handle();
EXPECT_NE(ahwb, nullptr);
}
}
TEST_F(TensorAhwbTest, TestAHWBThenCpu) {
Tensor tensor(Tensor::ElementType::kFloat32, Tensor::Shape{1});
{
auto ahwb = tensor.GetAHardwareBufferWriteView().handle();
EXPECT_NE(ahwb, nullptr);
}
{
auto ptr = tensor.GetCpuReadView().buffer<float>();
EXPECT_NE(ptr, nullptr);
}
}
TEST_F(TensorAhwbTest, TestCpuThenGl) {
RunInGlContext([] {
Tensor tensor(Tensor::ElementType::kFloat32, Tensor::Shape{1});
{
auto ptr = tensor.GetCpuWriteView().buffer<float>();
EXPECT_NE(ptr, nullptr);
}
{
auto ssbo = tensor.GetOpenGlBufferReadView().name();
EXPECT_GT(ssbo, 0);
}
});
}
} // namespace mediapipe
#endif // !MEDIAPIPE_DISABLE_GPU
#endif // MEDIAPIPE_TENSOR_USE_AHWB
mediapipe/framework/formats/time_series_header.proto
+2
View File
@@ -21,6 +21,8 @@ syntax = "proto2";
package mediapipe;
option objc_class_prefix = "MediaPipe";
// Header for a uniformly sampled time series stream. Each Packet in
// the stream is a Matrix, and each column is a (vector-valued) sample of
// the series, i.e. each column corresponds to a distinct sample in time.
mediapipe/framework/input_stream_manager.cc
+2
View File
@@ -204,6 +204,8 @@ absl::Status InputStreamManager::SetNextTimestampBound(const Timestamp bound,
// untimed scheduling policies.
if (bound > next_timestamp_bound_) {
next_timestamp_bound_ = bound;
VLOG(3) << "Next timestamp bound for input " << name_ << " is "
<< next_timestamp_bound_;
if (queue_.empty()) {
// If the queue was not empty then a change to the next_timestamp_bound_
// is not detectable by the consumer.
mediapipe/framework/output_stream_manager.cc
+2
View File
@@ -168,6 +168,8 @@ void OutputStreamManager::PropagateUpdatesToMirrors(
if (next_timestamp_bound != Timestamp::Unset()) {
absl::MutexLock lock(&stream_mutex_);
next_timestamp_bound_ = next_timestamp_bound;
VLOG(3) << "Next timestamp bound for output " << output_stream_spec_.name
<< " is " << next_timestamp_bound_;
}
}
std::list<Packet>* packets_to_propagate = output_stream_shard->OutputQueue();
mediapipe/framework/packet.cc
+6 -8
View File
@@ -106,19 +106,17 @@ std::string Packet::DebugString() const {
return result;
}
absl::Status Packet::ValidateAsType(const tool::TypeInfo& type_info) const {
absl::Status Packet::ValidateAsType(TypeId type_id) const {
if (ABSL_PREDICT_FALSE(IsEmpty())) {
return absl::InternalError(
absl::StrCat("Expected a Packet of type: ",
MediaPipeTypeStringOrDemangled(type_info),
", but received an empty Packet."));
return absl::InternalError(absl::StrCat(
"Expected a Packet of type: ", MediaPipeTypeStringOrDemangled(type_id),
", but received an empty Packet."));
}
bool holder_is_right_type =
holder_->GetTypeInfo().hash_code() == type_info.hash_code();
bool holder_is_right_type = holder_->GetTypeId() == type_id;
if (ABSL_PREDICT_FALSE(!holder_is_right_type)) {
return absl::InvalidArgumentError(absl::StrCat(
"The Packet stores \"", holder_->DebugTypeName(), "\", but \"",
MediaPipeTypeStringOrDemangled(type_info), "\" was requested."));
MediaPipeTypeStringOrDemangled(type_id), "\" was requested."));
}
return absl::OkStatus();
}
mediapipe/framework/packet.h
+10 -17
View File
@@ -21,7 +21,6 @@
#include <memory>
#include <string>
#include <type_traits>
#include <typeinfo>
#include "absl/base/macros.h"
#include "absl/memory/memory.h"
@@ -69,7 +68,7 @@ absl::StatusOr<Packet> PacketFromDynamicProto(const std::string& type_name,
// The preferred method of creating a Packet is with MakePacket<T>().
// The Packet typically owns the object that it contains, but
// PointToForeign allows a Packet to be constructed which does not
// own it's data.
// own its data.
//
// This class is thread compatible.
class Packet {
@@ -180,7 +179,7 @@ class Packet {
// Returns an error if the packet does not contain data of type T.
template <typename T>
absl::Status ValidateAsType() const {
return ValidateAsType(tool::TypeInfo::Get<T>());
return ValidateAsType(kTypeId<T>);
}
// Returns an error if the packet is not an instance of
@@ -189,11 +188,7 @@ class Packet {
// Get the type id for the underlying type stored in the Packet.
// Crashes if IsEmpty() == true.
size_t GetTypeId() const { return GetTypeInfo().hash_code(); }
// Get the type info for the underlying type stored in the Packet.
// Crashes if IsEmpty() == true.
const tool::TypeInfo& GetTypeInfo() const;
TypeId GetTypeId() const;
// Returns the timestamp.
class Timestamp Timestamp() const;
@@ -225,7 +220,7 @@ class Packet {
packet_internal::GetHolderShared(Packet&& packet);
friend class PacketType;
absl::Status ValidateAsType(const tool::TypeInfo& type_info) const;
absl::Status ValidateAsType(TypeId type_id) const;
std::shared_ptr<packet_internal::HolderBase> holder_;
class Timestamp timestamp_;
@@ -369,7 +364,7 @@ class HolderBase {
virtual ~HolderBase();
template <typename T>
bool PayloadIsOfType() const {
return GetTypeInfo().hash_code() == tool::GetTypeHash<T>();
return GetTypeId() == kTypeId<T>;
}
// Returns a printable string identifying the type stored in the holder.
virtual const std::string DebugTypeName() const = 0;
@@ -377,7 +372,7 @@ class HolderBase {
// empty string.
virtual const std::string RegisteredTypeName() const = 0;
// Get the type id of the underlying data type.
virtual const tool::TypeInfo& GetTypeInfo() const = 0;
virtual TypeId GetTypeId() const = 0;
// Downcasts this to Holder<T>. Returns nullptr if deserialization
// failed or if the requested type is not what is stored.
template <typename T>
@@ -428,7 +423,7 @@ StatusOr<std::vector<const proto_ns::MessageLite*>>
ConvertToVectorOfProtoMessageLitePtrs(const T* data,
/*is_proto_vector=*/std::false_type) {
return absl::InvalidArgumentError(absl::StrCat(
"The Packet stores \"", tool::TypeInfo::Get<T>().name(), "\"",
"The Packet stores \"", kTypeId<T>.name(), "\"",
"which is not convertible to vector<proto_ns::MessageLite*>."));
}
@@ -510,9 +505,7 @@ class Holder : public HolderBase {
HolderSupport<T>::EnsureStaticInit();
return *ptr_;
}
const tool::TypeInfo& GetTypeInfo() const final {
return tool::TypeInfo::Get<T>();
}
TypeId GetTypeId() const final { return kTypeId<T>; }
// Releases the underlying data pointer and transfers the ownership to a
// unique pointer.
// This method is dangerous and is only used by Packet::Consume() if the
@@ -748,9 +741,9 @@ inline Packet& Packet::operator=(Packet&& packet) {
inline bool Packet::IsEmpty() const { return holder_ == nullptr; }
inline const tool::TypeInfo& Packet::GetTypeInfo() const {
inline TypeId Packet::GetTypeId() const {
CHECK(holder_);
return holder_->GetTypeInfo();
return holder_->GetTypeId();
}
template <typename T>
mediapipe/framework/packet_test.proto
+2
View File
@@ -18,6 +18,8 @@ syntax = "proto2";
package mediapipe;
option objc_class_prefix = "MediaPipe";
message PacketTestProto {
// Tests that the tags used to encode the timestamp do not interfere with
// proto tags.
mediapipe/framework/packet_type.cc
+22 -26
View File
@@ -127,13 +127,13 @@ bool PacketType::IsOneOf() const {
}
bool PacketType::IsExactType() const {
return absl::holds_alternative<const tool::TypeInfo*>(type_spec_);
return absl::holds_alternative<TypeId>(type_spec_);
}
const std::string* PacketType::RegisteredTypeName() const {
if (auto* same_as = SameAsPtr()) return same_as->RegisteredTypeName();
if (auto* type_info = absl::get_if<const tool::TypeInfo*>(&type_spec_))
return MediaPipeTypeStringFromTypeId((**type_info).hash_code());
if (auto* type_id = absl::get_if<TypeId>(&type_spec_))
return MediaPipeTypeStringFromTypeId(*type_id);
if (auto* multi_type = absl::get_if<MultiType>(&type_spec_))
return multi_type->registered_type_name;
return nullptr;
@@ -141,8 +141,8 @@ const std::string* PacketType::RegisteredTypeName() const {
namespace internal {
struct TypeInfoFormatter {
void operator()(std::string* out, const tool::TypeInfo& t) const {
struct TypeIdFormatter {
void operator()(std::string* out, TypeId t) const {
absl::StrAppend(out, MediaPipeTypeStringOrDemangled(t));
}
};
@@ -167,12 +167,9 @@ explicit QuoteFormatter(Formatter f) -> QuoteFormatter<Formatter>;
} // namespace internal
std::string PacketType::TypeNameForOneOf(TypeInfoSpan types) {
std::string PacketType::TypeNameForOneOf(TypeIdSpan types) {
return absl::StrCat(
"OneOf<",
absl::StrJoin(types, ", ",
absl::DereferenceFormatter(internal::TypeInfoFormatter())),
">");
"OneOf<", absl::StrJoin(types, ", ", internal::TypeIdFormatter()), ">");
}
std::string PacketType::DebugTypeName() const {
@@ -185,8 +182,8 @@ std::string PacketType::DebugTypeName() const {
if (auto* special = absl::get_if<SpecialType>(&type_spec_)) {
return special->name_;
}
if (auto* type_info = absl::get_if<const tool::TypeInfo*>(&type_spec_)) {
return MediaPipeTypeStringOrDemangled(**type_info);
if (auto* type_id = absl::get_if<TypeId>(&type_spec_)) {
return MediaPipeTypeStringOrDemangled(*type_id);
}
if (auto* multi_type = absl::get_if<MultiType>(&type_spec_)) {
return TypeNameForOneOf(multi_type->types);
@@ -194,11 +191,11 @@ std::string PacketType::DebugTypeName() const {
return "[Undefined Type]";
}
static bool HaveCommonType(absl::Span<const tool::TypeInfo* const> types1,
absl::Span<const tool::TypeInfo* const> types2) {
static bool HaveCommonType(absl::Span<const TypeId> types1,
absl::Span<const TypeId> types2) {
for (const auto& first : types1) {
for (const auto& second : types2) {
if (first->hash_code() == second->hash_code()) {
if (first == second) {
return true;
}
}
@@ -216,35 +213,34 @@ absl::Status PacketType::Validate(const Packet& packet) const {
// in SetSameAs().
return GetSameAs()->Validate(packet);
}
if (auto* type_info = absl::get_if<const tool::TypeInfo*>(&type_spec_)) {
return packet.ValidateAsType(**type_info);
if (auto* type_id = absl::get_if<TypeId>(&type_spec_)) {
return packet.ValidateAsType(*type_id);
}
if (packet.IsEmpty()) {
return mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
<< "Empty packets are not allowed for type: " << DebugTypeName();
}
if (auto* multi_type = absl::get_if<MultiType>(&type_spec_)) {
auto* packet_type = &packet.GetTypeInfo();
auto packet_type = packet.GetTypeId();
if (HaveCommonType(multi_type->types, absl::MakeSpan(&packet_type, 1))) {
return absl::OkStatus();
} else {
return absl::InvalidArgumentError(absl::StrCat(
"The Packet stores \"", packet.DebugTypeName(), "\", but one of ",
absl::StrJoin(multi_type->types, ", ",
absl::DereferenceFormatter(internal::QuoteFormatter(
internal::TypeInfoFormatter()))),
internal::QuoteFormatter(internal::TypeIdFormatter())),
" was requested."));
}
}
if (auto* special = absl::get_if<SpecialType>(&type_spec_)) {
return special->accept_fn_(&packet.GetTypeInfo());
return special->accept_fn_(packet.GetTypeId());
}
return absl::OkStatus();
}
PacketType::TypeInfoSpan PacketType::GetTypeSpan(const TypeSpec& type_spec) {
if (auto* type_info = absl::get_if<const tool::TypeInfo*>(&type_spec))
return absl::MakeSpan(type_info, 1);
PacketType::TypeIdSpan PacketType::GetTypeSpan(const TypeSpec& type_spec) {
if (auto* type_id = absl::get_if<TypeId>(&type_spec))
return absl::MakeSpan(type_id, 1);
if (auto* multi_type = absl::get_if<MultiType>(&type_spec))
return multi_type->types;
return {};
@@ -254,8 +250,8 @@ bool PacketType::IsConsistentWith(const PacketType& other) const {
const PacketType* type1 = GetSameAs();
const PacketType* type2 = other.GetSameAs();
TypeInfoSpan types1 = GetTypeSpan(type1->type_spec_);
TypeInfoSpan types2 = GetTypeSpan(type2->type_spec_);
TypeIdSpan types1 = GetTypeSpan(type1->type_spec_);
TypeIdSpan types2 = GetTypeSpan(type2->type_spec_);
if (!types1.empty() && !types2.empty()) {
return HaveCommonType(types1, types2);
}
mediapipe/framework/packet_type.h
+8 -9
View File
@@ -121,15 +121,15 @@ class PacketType {
// We don't do union-find optimizations in order to avoid a mutex.
const PacketType* other;
};
using TypeInfoSpan = absl::Span<const tool::TypeInfo* const>;
using TypeIdSpan = absl::Span<const TypeId>;
struct MultiType {
TypeInfoSpan types;
TypeIdSpan types;
// TODO: refactor RegisteredTypeName, remove.
const std::string* registered_type_name;
};
struct SpecialType;
using TypeSpec = absl::variant<absl::monostate, const tool::TypeInfo*,
MultiType, SameAs, SpecialType>;
using TypeSpec =
absl::variant<absl::monostate, TypeId, MultiType, SameAs, SpecialType>;
typedef absl::Status (*AcceptsTypeFn)(const TypeSpec& type);
struct SpecialType {
std::string name_;
@@ -140,8 +140,8 @@ class PacketType {
static absl::Status AcceptNone(const TypeSpec& type);
const PacketType* SameAsPtr() const;
static TypeInfoSpan GetTypeSpan(const TypeSpec& type_spec);
static std::string TypeNameForOneOf(TypeInfoSpan types);
static TypeIdSpan GetTypeSpan(const TypeSpec& type_spec);
static std::string TypeNameForOneOf(TypeIdSpan types);
TypeSpec type_spec_;
@@ -259,14 +259,13 @@ absl::Status ValidatePacketTypeSet(const PacketTypeSet& packet_type_set);
template <typename T>
PacketType& PacketType::Set() {
type_spec_ = &tool::TypeInfo::Get<T>();
type_spec_ = kTypeId<T>;
return *this;
}
template <typename... T>
PacketType& PacketType::SetOneOf() {
static const NoDestructor<std::vector<const tool::TypeInfo*>> types{
{&tool::TypeInfo::Get<T>()...}};
static const NoDestructor<std::vector<TypeId>> types{{kTypeId<T>...}};
static const NoDestructor<std::string> name{TypeNameForOneOf(*types)};
type_spec_ = MultiType{*types, &*name};
return *this;
mediapipe/framework/profiler/graph_profiler.cc
+3 -3
View File
@@ -43,7 +43,7 @@ const int kDefaultLogFileCount = 2;
const char kDefaultLogFilePrefix[] = "mediapipe_trace_";
// The number of recent timestamps tracked for each input stream.
const int kPacketInfoRecentCount = 100;
const int kPacketInfoRecentCount = 400;
std::string PacketIdToString(const PacketId& packet_id) {
return absl::Substitute("stream_name: $0, timestamp_usec: $1",
@@ -507,8 +507,8 @@ int64 GraphProfiler::AddInputStreamTimeSamples(
// This is a condition rather than a failure CHECK because
// under certain conditions the consumer calculator's Process()
// can start before the producer calculator's Process() is finished.
LOG_EVERY_N(WARNING, 100) << "Expected packet info is missing for: "
<< PacketIdToString(packet_id);
LOG_FIRST_N(WARNING, 10) << "Expected packet info is missing for: "
<< PacketIdToString(packet_id);
continue;
}
AddTimeSample(
mediapipe/framework/subgraph.h
+11 -6
View File
@@ -36,7 +36,7 @@ class SubgraphContext {
public:
SubgraphContext() : SubgraphContext(nullptr, nullptr) {}
// @node and/or @service_manager can be nullptr.
SubgraphContext(const CalculatorGraphConfig::Node* node,
SubgraphContext(CalculatorGraphConfig::Node* node,
const GraphServiceManager* service_manager)
: default_node_(node ? absl::nullopt
: absl::optional<CalculatorGraphConfig::Node>(
@@ -48,14 +48,19 @@ class SubgraphContext {
: absl::optional<GraphServiceManager>(GraphServiceManager())),
service_manager_(service_manager ? *service_manager
: default_service_manager_.value()),
options_map_(std::move(tool::OptionsMap().Initialize(original_node_))) {
}
options_map_(
std::move(tool::MutableOptionsMap().Initialize(original_node_))) {}
template <typename T>
const T& Options() {
return options_map_.Get<T>();
}
template <typename T>
T* MutableOptions() {
return options_map_.GetMutable<T>();
}
const CalculatorGraphConfig::Node& OriginalNode() const {
return original_node_;
}
@@ -67,16 +72,16 @@ class SubgraphContext {
private:
// Populated if node is not provided during construction.
const absl::optional<CalculatorGraphConfig::Node> default_node_;
absl::optional<CalculatorGraphConfig::Node> default_node_;
const CalculatorGraphConfig::Node& original_node_;
CalculatorGraphConfig::Node& original_node_;
// Populated if service manager is not provided during construction.
const absl::optional<GraphServiceManager> default_service_manager_;
const GraphServiceManager& service_manager_;
tool::OptionsMap options_map_;
tool::MutableOptionsMap options_map_;
};
// Instances of this class are responsible for providing a subgraph config.
mediapipe/framework/test_calculators.proto
+2
View File
@@ -22,6 +22,8 @@ package mediapipe;
import "mediapipe/framework/calculator.proto";
option objc_class_prefix = "MediaPipe";
message RandomMatrixCalculatorOptions {
extend CalculatorOptions {
optional RandomMatrixCalculatorOptions ext = 52056136;
mediapipe/framework/tool/BUILD
+8 -1
View File
@@ -198,6 +198,7 @@ cc_library(
":name_util",
":options_registry",
":proto_util_lite",
":type_util",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:packet",
"//mediapipe/framework:packet_type",
@@ -277,9 +278,12 @@ cc_library(
hdrs = ["options_registry.h"],
visibility = ["//visibility:public"],
deps = [
":field_data_cc_proto",
":proto_util_lite",
"//mediapipe/framework/deps:registration",
"//mediapipe/framework/port:advanced_proto",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:ret_check",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/synchronization",
],
@@ -334,6 +338,7 @@ cc_library(
hdrs = ["proto_util_lite.h"],
visibility = ["//visibility:public"],
deps = [
":field_data_cc_proto",
"//mediapipe/framework:type_map",
"//mediapipe/framework/port:advanced_proto_lite",
"//mediapipe/framework/port:integral_types",
@@ -518,9 +523,11 @@ cc_library(
cc_library(
name = "type_util",
hdrs = ["type_util.h"],
visibility = ["//mediapipe/framework:mediapipe_internal"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:demangle",
"//mediapipe/framework:port",
"@com_google_absl//absl/base:core_headers",
],
)
mediapipe/framework/tool/calculator_graph_template.proto
+1
View File
@@ -3,6 +3,7 @@ syntax = "proto2";
package mediapipe;
import "mediapipe/framework/calculator.proto";
import "mediapipe/framework/calculator_options.proto";
import "mediapipe/framework/deps/proto_descriptor.proto";
option java_package = "com.google.mediapipe.proto";
mediapipe/framework/tool/options_field_util.cc
+215 -304
View File
@@ -7,6 +7,7 @@
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "mediapipe/framework/packet.h"
@@ -18,6 +19,7 @@
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/tool/name_util.h"
#include "mediapipe/framework/tool/proto_util_lite.h"
#include "mediapipe/framework/tool/type_util.h"
namespace mediapipe {
namespace tool {
@@ -41,165 +43,39 @@ FieldType AsFieldType(proto_ns::FieldDescriptorProto::Type type) {
return static_cast<FieldType>(type);
}
absl::Status WriteValue(const FieldData& value, FieldType field_type,
std::string* field_bytes) {
StringOutputStream sos(field_bytes);
CodedOutputStream out(&sos);
switch (field_type) {
case WireFormatLite::TYPE_INT32:
WireFormatLite::WriteInt32NoTag(value.int32_value(), &out);
break;
case WireFormatLite::TYPE_SINT32:
WireFormatLite::WriteSInt32NoTag(value.int32_value(), &out);
break;
case WireFormatLite::TYPE_INT64:
WireFormatLite::WriteInt64NoTag(value.int64_value(), &out);
break;
case WireFormatLite::TYPE_SINT64:
WireFormatLite::WriteSInt64NoTag(value.int64_value(), &out);
break;
case WireFormatLite::TYPE_UINT32:
WireFormatLite::WriteUInt32NoTag(value.uint32_value(), &out);
break;
case WireFormatLite::TYPE_UINT64:
WireFormatLite::WriteUInt64NoTag(value.uint64_value(), &out);
break;
case WireFormatLite::TYPE_DOUBLE:
WireFormatLite::WriteDoubleNoTag(value.uint64_value(), &out);
break;
case WireFormatLite::TYPE_FLOAT:
WireFormatLite::WriteFloatNoTag(value.float_value(), &out);
break;
case WireFormatLite::TYPE_BOOL:
WireFormatLite::WriteBoolNoTag(value.bool_value(), &out);
break;
case WireFormatLite::TYPE_ENUM:
WireFormatLite::WriteEnumNoTag(value.enum_value(), &out);
break;
case WireFormatLite::TYPE_STRING:
out.WriteString(value.string_value());
break;
case WireFormatLite::TYPE_MESSAGE:
out.WriteString(value.message_value().value());
break;
default:
return absl::UnimplementedError(
absl::StrCat("Cannot write type: ", field_type));
}
return absl::OkStatus();
}
// Serializes a packet value.
absl::Status WriteField(const FieldData& packet, const FieldDescriptor* field,
std::string* result) {
FieldType field_type = AsFieldType(field->type());
return WriteValue(packet, field_type, result);
}
template <typename ValueT, FieldType kFieldType>
static ValueT ReadValue(absl::string_view field_bytes, absl::Status* status) {
ArrayInputStream ais(field_bytes.data(), field_bytes.size());
CodedInputStream input(&ais);
ValueT result;
if (!WireFormatLite::ReadPrimitive<ValueT, kFieldType>(&input, &result)) {
status->Update(mediapipe::InvalidArgumentError(absl::StrCat(
"Bad serialized value: ", MediaPipeTypeStringOrDemangled<ValueT>(),
".")));
}
return result;
}
absl::Status ReadValue(absl::string_view field_bytes, FieldType field_type,
absl::string_view message_type, FieldData* result) {
absl::Status status;
result->Clear();
switch (field_type) {
case WireFormatLite::TYPE_INT32:
result->set_int32_value(
ReadValue<int32, WireFormatLite::TYPE_INT32>(field_bytes, &status));
break;
case WireFormatLite::TYPE_SINT32:
result->set_int32_value(
ReadValue<int32, WireFormatLite::TYPE_SINT32>(field_bytes, &status));
break;
case WireFormatLite::TYPE_INT64:
result->set_int64_value(
ReadValue<int64, WireFormatLite::TYPE_INT64>(field_bytes, &status));
break;
case WireFormatLite::TYPE_SINT64:
result->set_int64_value(
ReadValue<int64, WireFormatLite::TYPE_SINT64>(field_bytes, &status));
break;
case WireFormatLite::TYPE_UINT32:
result->set_uint32_value(
ReadValue<uint32, WireFormatLite::TYPE_UINT32>(field_bytes, &status));
break;
case WireFormatLite::TYPE_UINT64:
result->set_uint64_value(
ReadValue<uint32, WireFormatLite::TYPE_UINT32>(field_bytes, &status));
break;
case WireFormatLite::TYPE_DOUBLE:
result->set_double_value(
ReadValue<double, WireFormatLite::TYPE_DOUBLE>(field_bytes, &status));
break;
case WireFormatLite::TYPE_FLOAT:
result->set_float_value(
ReadValue<float, WireFormatLite::TYPE_FLOAT>(field_bytes, &status));
break;
case WireFormatLite::TYPE_BOOL:
result->set_bool_value(
ReadValue<bool, WireFormatLite::TYPE_BOOL>(field_bytes, &status));
break;
case WireFormatLite::TYPE_ENUM:
result->set_enum_value(
ReadValue<int32, WireFormatLite::TYPE_ENUM>(field_bytes, &status));
break;
case WireFormatLite::TYPE_STRING:
result->set_string_value(std::string(field_bytes));
break;
case WireFormatLite::TYPE_MESSAGE:
result->mutable_message_value()->set_value(std::string(field_bytes));
result->mutable_message_value()->set_type_url(TypeUrl(message_type));
break;
default:
status = absl::UnimplementedError(
absl::StrCat("Cannot read type: ", field_type));
break;
}
return status;
return ProtoUtilLite::WriteValue(packet, field->type(), result);
}
// Deserializes a packet from a protobuf field.
absl::Status ReadField(absl::string_view bytes, const FieldDescriptor* field,
absl::Status ReadField(absl::string_view bytes, const FieldDescriptor& field,
FieldData* result) {
RET_CHECK_NE(field, nullptr);
FieldType field_type = AsFieldType(field->type());
std::string message_type = (field_type == WireFormatLite::TYPE_MESSAGE)
? field->message_type()->full_name()
std::string message_type = (field.type() == WireFormatLite::TYPE_MESSAGE)
? field.message_type()->full_name()
: "";
return ReadValue(bytes, field_type, message_type, result);
return ProtoUtilLite::ReadValue(bytes, field.type(), message_type, result);
}
// Reads all values from a repeated field.
absl::Status GetFieldValues(const FieldData& message_data,
const FieldDescriptor& field,
std::vector<FieldData>* result) {
absl::StatusOr<std::vector<FieldData>> GetFieldValues(
const FieldData& message_data, const FieldDescriptor& field) {
std::vector<FieldData> result;
const std::string& message_bytes = message_data.message_value().value();
FieldType field_type = AsFieldType(field.type());
ProtoUtilLite proto_util;
ProtoUtilLite::ProtoPath proto_path = {{field.number(), 0}};
int count;
MP_RETURN_IF_ERROR(
proto_util.GetFieldCount(message_bytes, proto_path, field_type, &count));
MP_RETURN_IF_ERROR(ProtoUtilLite::GetFieldCount(message_bytes, proto_path,
field.type(), &count));
std::vector<std::string> field_values;
MP_RETURN_IF_ERROR(proto_util.GetFieldRange(message_bytes, proto_path, count,
field_type, &field_values));
for (int i = 0; i < count; ++i) {
MP_RETURN_IF_ERROR(ProtoUtilLite::GetFieldRange(
message_bytes, proto_path, count, field.type(), &field_values));
for (int i = 0; i < field_values.size(); ++i) {
FieldData r;
MP_RETURN_IF_ERROR(ReadField(field_values[i], &field, &r));
result->push_back(std::move(r));
MP_RETURN_IF_ERROR(ReadField(field_values[i], field, &r));
result.push_back(std::move(r));
}
return absl::OkStatus();
return result;
}
// Reads one value from a field.
@@ -207,42 +83,70 @@ absl::Status GetFieldValue(const FieldData& message_data,
const FieldPathEntry& entry, FieldData* result) {
RET_CHECK_NE(entry.field, nullptr);
const std::string& message_bytes = message_data.message_value().value();
FieldType field_type = AsFieldType(entry.field->type());
ProtoUtilLite proto_util;
ProtoUtilLite::ProtoPath proto_path = {{entry.field->number(), entry.index}};
FieldType field_type = entry.field->type();
int index = std::max(0, entry.index);
ProtoUtilLite::ProtoPath proto_path = {{entry.field->number(), index}};
std::vector<std::string> field_values;
MP_RETURN_IF_ERROR(proto_util.GetFieldRange(message_bytes, proto_path, 1,
field_type, &field_values));
MP_RETURN_IF_ERROR(ReadField(field_values[0], entry.field, result));
MP_RETURN_IF_ERROR(ProtoUtilLite::GetFieldRange(message_bytes, proto_path, 1,
field_type, &field_values));
MP_RETURN_IF_ERROR(ReadField(field_values[0], *entry.field, result));
return absl::OkStatus();
}
// Writes one value to a field.
absl::Status SetFieldValue(const FieldPathEntry& entry, const FieldData& value,
FieldData* result) {
std::vector<FieldData> field_values;
ProtoUtilLite proto_util;
FieldType field_type = AsFieldType(entry.field->type());
ProtoUtilLite::ProtoPath proto_path = {{entry.field->number(), entry.index}};
std::string* message_bytes = result->mutable_message_value()->mutable_value();
absl::Status SetFieldValue(FieldData& result, const FieldPathEntry& entry,
const FieldData& value) {
int index = std::max(0, entry.index);
ProtoUtilLite::ProtoPath proto_path = {{entry.field->number(), index}};
std::string* message_bytes = result.mutable_message_value()->mutable_value();
int field_count;
MP_RETURN_IF_ERROR(proto_util.GetFieldCount(*message_bytes, proto_path,
field_type, &field_count));
if (entry.index > field_count) {
MP_RETURN_IF_ERROR(ProtoUtilLite::GetFieldCount(
*message_bytes, proto_path, entry.field->type(), &field_count));
if (index > field_count) {
return absl::OutOfRangeError(
absl::StrCat("Option field index out of range: ", entry.index));
absl::StrCat("Option field index out of range: ", index));
}
int replace_length = entry.index < field_count ? 1 : 0;
int replace_length = index < field_count ? 1 : 0;
std::string field_value;
MP_RETURN_IF_ERROR(WriteField(value, entry.field, &field_value));
MP_RETURN_IF_ERROR(proto_util.ReplaceFieldRange(
message_bytes, proto_path, replace_length, field_type, {field_value}));
MP_RETURN_IF_ERROR(ProtoUtilLite::ReplaceFieldRange(
message_bytes, proto_path, replace_length, entry.field->type(),
{field_value}));
return absl::OkStatus();
}
// Writes several values to a repeated field.
// The specified |values| replace the specified |entry| index,
// or if no index is specified all field values are replaced.
absl::Status SetFieldValues(FieldData& result, const FieldPathEntry& entry,
const std::vector<FieldData>& values) {
if (entry.field == nullptr) {
return absl::InvalidArgumentError("Field not found.");
}
FieldType field_type = entry.field->type();
ProtoUtilLite::ProtoPath proto_path = {{entry.field->number(), 0}};
std::string* message_bytes = result.mutable_message_value()->mutable_value();
int field_count;
MP_RETURN_IF_ERROR(ProtoUtilLite::GetFieldCount(*message_bytes, proto_path,
field_type, &field_count));
int replace_start = 0, replace_length = field_count;
if (entry.index > -1) {
replace_start = entry.index;
replace_length = 1;
}
std::vector<std::string> field_values(values.size());
for (int i = 0; i < values.size(); ++i) {
MP_RETURN_IF_ERROR(WriteField(values[i], entry.field, &field_values[i]));
}
proto_path = {{entry.field->number(), replace_start}};
MP_RETURN_IF_ERROR(ProtoUtilLite::ReplaceFieldRange(
message_bytes, proto_path, replace_length, field_type, field_values));
return absl::OkStatus();
}
// Returns true for a field of type "google.protobuf.Any".
bool IsProtobufAny(const FieldDescriptor* field) {
return AsFieldType(field->type()) == FieldType::TYPE_MESSAGE &&
return field->type() == FieldType::TYPE_MESSAGE &&
field->message_type()->full_name() == kGoogleProtobufAny;
}
@@ -275,9 +179,7 @@ StatusOr<int> FindExtensionIndex(const FieldData& message_data,
}
std::string& extension_type = entry->extension_type;
std::vector<FieldData> field_values;
RET_CHECK_NE(entry->field, nullptr);
MP_RETURN_IF_ERROR(
GetFieldValues(message_data, *entry->field, &field_values));
ASSIGN_OR_RETURN(field_values, GetFieldValues(message_data, *entry->field));
for (int i = 0; i < field_values.size(); ++i) {
FieldData extension = ParseProtobufAny(field_values[i]);
if (extension_type == "*" ||
@@ -290,9 +192,9 @@ StatusOr<int> FindExtensionIndex(const FieldData& message_data,
// Returns true if the value of a field is available.
bool HasField(const FieldPath& field_path, const FieldData& message_data) {
FieldData value;
return GetField(field_path, message_data, &value).ok() &&
value.value_case() != mediapipe::FieldData::VALUE_NOT_SET;
auto value = GetField(message_data, field_path);
return value.ok() &&
value->value_case() != mediapipe::FieldData::VALUE_NOT_SET;
}
// Returns the extension field containing the specified extension-type.
@@ -330,43 +232,24 @@ void SetOptionsMessage(
*options_any->mutable_value() = node_options.message_value().value();
}
// Returns the count of values in a repeated field.
int FieldCount(const FieldData& message_data, const FieldDescriptor* field) {
const std::string& message_bytes = message_data.message_value().value();
FieldType field_type = AsFieldType(field->type());
ProtoUtilLite proto_util;
ProtoUtilLite::ProtoPath proto_path = {{field->number(), 0}};
int count;
if (proto_util.GetFieldCount(message_bytes, proto_path, field_type, &count)
.ok()) {
return count;
}
return 0;
}
} // anonymous namespace
// Deserializes a packet containing a MessageLite value.
absl::Status ReadMessage(const std::string& value, const std::string& type_name,
Packet* result) {
auto packet = packet_internal::PacketFromDynamicProto(type_name, value);
if (packet.ok()) {
*result = *packet;
}
return packet.status();
absl::StatusOr<Packet> ReadMessage(const std::string& value,
const std::string& type_name) {
return packet_internal::PacketFromDynamicProto(type_name, value);
}
// Merge two options FieldData values.
absl::Status MergeMessages(const FieldData& base, const FieldData& over,
FieldData* result) {
absl::StatusOr<FieldData> MergeMessages(const FieldData& base,
const FieldData& over) {
FieldData result;
absl::Status status;
if (over.value_case() == FieldData::VALUE_NOT_SET) {
*result = base;
return status;
return base;
}
if (base.value_case() == FieldData::VALUE_NOT_SET) {
*result = over;
return status;
return over;
}
if (over.value_case() != base.value_case()) {
return absl::InvalidArgumentError(absl::StrCat(
@@ -382,10 +265,9 @@ absl::Status MergeMessages(const FieldData& base, const FieldData& over,
absl::Cord merged_value;
merged_value.Append(base.message_value().value());
merged_value.Append(over.message_value().value());
result->mutable_message_value()->set_type_url(
base.message_value().type_url());
result->mutable_message_value()->set_value(std::string(merged_value));
return status;
result.mutable_message_value()->set_type_url(base.message_value().type_url());
result.mutable_message_value()->set_value(std::string(merged_value));
return result;
}
// Returns either the extension field or the repeated protobuf.Any field index
@@ -439,51 +321,48 @@ FieldPath GetExtensionPath(const std::string& parent_type,
}
// Returns the requested options protobuf for a graph node.
absl::Status GetNodeOptions(const FieldData& message_data,
const std::string& extension_type,
FieldData* result) {
absl::StatusOr<FieldData> GetNodeOptions(const FieldData& message_data,
const std::string& extension_type) {
constexpr char kOptionsName[] = "options";
constexpr char kNodeOptionsName[] = "node_options";
std::string parent_type = options_field_util::ParseTypeUrl(
std::string(message_data.message_value().type_url()));
FieldPath path;
Status status;
absl::Status status;
path = GetExtensionPath(parent_type, extension_type, kOptionsName, false);
status = GetField(path, message_data, result);
if (status.ok()) {
return status;
auto result = GetField(message_data, path);
if (result.ok()) {
return result;
}
path = GetExtensionPath(parent_type, extension_type, kNodeOptionsName, true);
status = GetField(path, message_data, result);
return status;
return GetField(message_data, path);
}
// Returns the requested options protobuf for a graph.
absl::Status GetGraphOptions(const FieldData& message_data,
const std::string& extension_type,
FieldData* result) {
absl::StatusOr<FieldData> GetGraphOptions(const FieldData& message_data,
const std::string& extension_type) {
constexpr char kOptionsName[] = "options";
constexpr char kGraphOptionsName[] = "graph_options";
std::string parent_type = options_field_util::ParseTypeUrl(
std::string(message_data.message_value().type_url()));
FieldPath path;
Status status;
absl::Status status;
path = GetExtensionPath(parent_type, extension_type, kOptionsName, false);
status = GetField(path, message_data, result);
if (status.ok()) {
return status;
auto result = GetField(message_data, path);
if (result.ok()) {
return result;
}
path = GetExtensionPath(parent_type, extension_type, kGraphOptionsName, true);
status = GetField(path, message_data, result);
return status;
return GetField(message_data, path);
}
// Reads a FieldData value from a protobuf field.
absl::Status GetField(const FieldPath& field_path,
const FieldData& message_data, FieldData* result) {
// Reads the FieldData values from a protobuf field.
absl::StatusOr<std::vector<FieldData>> GetFieldValues(
const FieldData& message_data, const FieldPath& field_path) {
std::vector<FieldData> results;
if (field_path.empty()) {
*result->mutable_message_value() = message_data.message_value();
return absl::OkStatus();
results.push_back(message_data);
return results;
}
FieldPathEntry head = field_path.front();
FieldPath tail = field_path;
@@ -491,65 +370,101 @@ absl::Status GetField(const FieldPath& field_path,
if (!head.extension_type.empty()) {
MP_RETURN_IF_ERROR(FindExtension(message_data, &head));
}
if (tail.empty() && FieldCount(message_data, head.field) == 0) {
return absl::OkStatus();
}
MP_RETURN_IF_ERROR(GetFieldValue(message_data, head, result));
RET_CHECK_NE(head.field, nullptr);
ASSIGN_OR_RETURN(results, GetFieldValues(message_data, *head.field));
if (IsProtobufAny(head.field)) {
*result = ParseProtobufAny(*result);
for (int i = 0; i < results.size(); ++i) {
results[i] = ParseProtobufAny(results[i]);
}
}
int index = tail.empty() ? head.index : std::max(0, head.index);
if ((int)results.size() <= index) {
return absl::OutOfRangeError(absl::StrCat(
"Missing feild value: ", head.field ? head.field->name() : "#",
" at index: ", index));
}
if (!tail.empty()) {
FieldData child = *result;
MP_RETURN_IF_ERROR(GetField(tail, child, result));
FieldData child = results.at(index);
ASSIGN_OR_RETURN(results, GetFieldValues(child, tail));
} else if (index > -1) {
FieldData child = results.at(index);
results.clear();
results.push_back(child);
}
return absl::OkStatus();
return results;
}
// Writes a FieldData value into protobuf field.
absl::Status SetField(const FieldPath& field_path, const FieldData& value,
FieldData* message_data) {
// Reads a FieldData value from a protobuf field.
absl::StatusOr<FieldData> GetField(const FieldData& message_data,
const FieldPath& field_path) {
std::vector<FieldData> results;
ASSIGN_OR_RETURN(results, GetFieldValues(message_data, field_path));
if (results.empty()) {
FieldPathEntry tail = field_path.back();
return absl::OutOfRangeError(absl::StrCat(
"Missing feild value: ", tail.field ? tail.field->name() : "##",
" at index: ", tail.index));
}
return results[0];
}
// Writes FieldData values into protobuf field.
absl::Status SetFieldValues(FieldData& message_data,
const FieldPath& field_path,
const std::vector<FieldData>& values) {
if (field_path.empty()) {
*message_data->mutable_message_value() = value.message_value();
if (values.empty()) {
return absl::InvalidArgumentError("Missing feild value.");
}
message_data = values[0];
return absl::OkStatus();
}
FieldPathEntry head = field_path.front();
FieldPath tail = field_path;
tail.erase(tail.begin());
if (!head.extension_type.empty()) {
MP_RETURN_IF_ERROR(FindExtension(*message_data, &head));
MP_RETURN_IF_ERROR(FindExtension(message_data, &head));
}
if (tail.empty()) {
MP_RETURN_IF_ERROR(SetFieldValue(head, value, message_data));
} else {
FieldData child;
MP_RETURN_IF_ERROR(GetFieldValue(*message_data, head, &child));
MP_RETURN_IF_ERROR(SetField(tail, value, &child));
if (IsProtobufAny(head.field)) {
child = SerializeProtobufAny(child);
}
MP_RETURN_IF_ERROR(SetFieldValue(head, child, message_data));
MP_RETURN_IF_ERROR(SetFieldValues(message_data, head, values));
return absl::OkStatus();
}
FieldData child;
MP_RETURN_IF_ERROR(GetFieldValue(message_data, head, &child));
MP_RETURN_IF_ERROR(SetFieldValues(child, tail, values));
if (IsProtobufAny(head.field)) {
child = SerializeProtobufAny(child);
}
MP_RETURN_IF_ERROR(SetFieldValue(message_data, head, child));
return absl::OkStatus();
}
// Merges a packet value into nested protobuf Message.
absl::Status MergeField(const FieldPath& field_path, const FieldData& value,
FieldData* message_data) {
// Writes a FieldData value into protobuf field.
absl::Status SetField(FieldData& message_data, const FieldPath& field_path,
const FieldData& value) {
return SetFieldValues(message_data, field_path, {value});
}
// Merges FieldData values into nested protobuf Message.
// For each new field index, any previous value is merged with the new value.
absl::Status MergeFieldValues(FieldData& message_data,
const FieldPath& field_path,
const std::vector<FieldData>& values) {
absl::Status status;
FieldType field_type = field_path.empty()
? FieldType::TYPE_MESSAGE
: AsFieldType(field_path.back().field->type());
std::string message_type =
(value.has_message_value())
? ParseTypeUrl(std::string(value.message_value().type_url()))
: "";
FieldData v = value;
FieldType field_type = field_path.empty() ? FieldType::TYPE_MESSAGE
: field_path.back().field->type();
std::vector<FieldData> results = values;
std::vector<FieldData> prevs;
ASSIGN_OR_RETURN(prevs, GetFieldValues(message_data, field_path));
if (field_type == FieldType::TYPE_MESSAGE) {
FieldData b;
status.Update(GetField(field_path, *message_data, &b));
status.Update(MergeMessages(b, v, &v));
for (int i = 0; i < std::min(values.size(), prevs.size()); ++i) {
FieldData& v = results[i];
FieldData& b = prevs[i];
ASSIGN_OR_RETURN(v, MergeMessages(b, v));
}
}
status.Update(SetField(field_path, v, message_data));
status.Update(SetFieldValues(message_data, field_path, results));
return status;
}
@@ -576,34 +491,35 @@ struct ProtoEnum {
int32 value;
};
absl::Status AsPacket(const FieldData& data, Packet* result) {
absl::StatusOr<Packet> AsPacket(const FieldData& data) {
Packet result;
switch (data.value_case()) {
case FieldData::ValueCase::kInt32Value:
*result = MakePacket<int32>(data.int32_value());
result = MakePacket<int32>(data.int32_value());
break;
case FieldData::ValueCase::kInt64Value:
*result = MakePacket<int64>(data.int64_value());
result = MakePacket<int64>(data.int64_value());
break;
case FieldData::ValueCase::kUint32Value:
*result = MakePacket<uint32>(data.uint32_value());
result = MakePacket<uint32>(data.uint32_value());
break;
case FieldData::ValueCase::kUint64Value:
*result = MakePacket<uint64>(data.uint64_value());
result = MakePacket<uint64>(data.uint64_value());
break;
case FieldData::ValueCase::kDoubleValue:
*result = MakePacket<double>(data.double_value());
result = MakePacket<double>(data.double_value());
break;
case FieldData::ValueCase::kFloatValue:
*result = MakePacket<float>(data.float_value());
result = MakePacket<float>(data.float_value());
break;
case FieldData::ValueCase::kBoolValue:
*result = MakePacket<bool>(data.bool_value());
result = MakePacket<bool>(data.bool_value());
break;
case FieldData::ValueCase::kEnumValue:
*result = MakePacket<ProtoEnum>(data.enum_value());
result = MakePacket<ProtoEnum>(data.enum_value());
break;
case FieldData::ValueCase::kStringValue:
*result = MakePacket<std::string>(data.string_value());
result = MakePacket<std::string>(data.string_value());
break;
case FieldData::ValueCase::kMessageValue: {
auto r = packet_internal::PacketFromDynamicProto(
@@ -612,32 +528,33 @@ absl::Status AsPacket(const FieldData& data, Packet* result) {
if (!r.ok()) {
return r.status();
}
*result = r.value();
result = r.value();
break;
}
case FieldData::VALUE_NOT_SET:
*result = Packet();
result = Packet();
}
return absl::OkStatus();
return result;
}
absl::Status AsFieldData(Packet packet, FieldData* result) {
static const auto* kTypeIds = new std::map<size_t, int32>{
{tool::GetTypeHash<int32>(), WireFormatLite::CPPTYPE_INT32},
{tool::GetTypeHash<int64>(), WireFormatLite::CPPTYPE_INT64},
{tool::GetTypeHash<uint32>(), WireFormatLite::CPPTYPE_UINT32},
{tool::GetTypeHash<uint64>(), WireFormatLite::CPPTYPE_UINT64},
{tool::GetTypeHash<double>(), WireFormatLite::CPPTYPE_DOUBLE},
{tool::GetTypeHash<float>(), WireFormatLite::CPPTYPE_FLOAT},
{tool::GetTypeHash<bool>(), WireFormatLite::CPPTYPE_BOOL},
{tool::GetTypeHash<ProtoEnum>(), WireFormatLite::CPPTYPE_ENUM},
{tool::GetTypeHash<std::string>(), WireFormatLite::CPPTYPE_STRING},
absl::StatusOr<FieldData> AsFieldData(Packet packet) {
static const auto* kTypeIds = new std::map<TypeId, int32>{
{kTypeId<int32>, WireFormatLite::CPPTYPE_INT32},
{kTypeId<int64>, WireFormatLite::CPPTYPE_INT64},
{kTypeId<uint32>, WireFormatLite::CPPTYPE_UINT32},
{kTypeId<uint64>, WireFormatLite::CPPTYPE_UINT64},
{kTypeId<double>, WireFormatLite::CPPTYPE_DOUBLE},
{kTypeId<float>, WireFormatLite::CPPTYPE_FLOAT},
{kTypeId<bool>, WireFormatLite::CPPTYPE_BOOL},
{kTypeId<ProtoEnum>, WireFormatLite::CPPTYPE_ENUM},
{kTypeId<std::string>, WireFormatLite::CPPTYPE_STRING},
};
FieldData result;
if (packet.ValidateAsProtoMessageLite().ok()) {
result->mutable_message_value()->set_value(
result.mutable_message_value()->set_value(
packet.GetProtoMessageLite().SerializeAsString());
result->mutable_message_value()->set_type_url(
result.mutable_message_value()->set_type_url(
TypeUrl(packet.GetProtoMessageLite().GetTypeName()));
return absl::OkStatus();
}
@@ -649,48 +566,42 @@ absl::Status AsFieldData(Packet packet, FieldData* result) {
switch (kTypeIds->at(packet.GetTypeId())) {
case WireFormatLite::CPPTYPE_INT32:
result->set_int32_value(packet.Get<int32>());
result.set_int32_value(packet.Get<int32>());
break;
case WireFormatLite::CPPTYPE_INT64:
result->set_int64_value(packet.Get<int64>());
result.set_int64_value(packet.Get<int64>());
break;
case WireFormatLite::CPPTYPE_UINT32:
result->set_uint32_value(packet.Get<uint32>());
result.set_uint32_value(packet.Get<uint32>());
break;
case WireFormatLite::CPPTYPE_UINT64:
result->set_uint64_value(packet.Get<uint64>());
result.set_uint64_value(packet.Get<uint64>());
break;
case WireFormatLite::CPPTYPE_DOUBLE:
result->set_double_value(packet.Get<double>());
result.set_double_value(packet.Get<double>());
break;
case WireFormatLite::CPPTYPE_FLOAT:
result->set_float_value(packet.Get<float>());
result.set_float_value(packet.Get<float>());
break;
case WireFormatLite::CPPTYPE_BOOL:
result->set_bool_value(packet.Get<bool>());
result.set_bool_value(packet.Get<bool>());
break;
case WireFormatLite::CPPTYPE_ENUM:
result->set_enum_value(packet.Get<ProtoEnum>().value);
result.set_enum_value(packet.Get<ProtoEnum>().value);
break;
case WireFormatLite::CPPTYPE_STRING:
result->set_string_value(packet.Get<std::string>());
result.set_string_value(packet.Get<std::string>());
break;
}
return absl::OkStatus();
return result;
}
std::string TypeUrl(absl::string_view type_name) {
constexpr std::string_view kTypeUrlPrefix = "type.googleapis.com/";
return absl::StrCat(std::string(kTypeUrlPrefix), std::string(type_name));
return ProtoUtilLite::TypeUrl(type_name);
}
std::string ParseTypeUrl(absl::string_view type_url) {
constexpr std::string_view kTypeUrlPrefix = "type.googleapis.com/";
if (std::string(type_url).rfind(kTypeUrlPrefix, 0) == 0) {
return std::string(
type_url.substr(kTypeUrlPrefix.length(), std::string::npos));
}
return std::string(type_url);
return ProtoUtilLite::ParseTypeUrl(type_url);
}
} // namespace options_field_util
mediapipe/framework/tool/options_field_util.h
+25 -17
View File
@@ -34,30 +34,38 @@ absl::Status SetField(const FieldPath& field_path, const FieldData& value,
FieldData* message_data);
// Reads a field value from a protobuf field.
absl::Status GetField(const FieldPath& field_path,
const FieldData& message_data, FieldData* result);
absl::StatusOr<FieldData> GetField(const FieldData& message_data,
const FieldPath& field_path);
// Merges a field value into nested protobuf Message.
absl::Status MergeField(const FieldPath& field_path, const FieldData& value,
FieldData* message_data);
// Reads one or all FieldData values from a protobuf field.
absl::StatusOr<std::vector<FieldData>> GetFieldValues(
const FieldData& message_data, const FieldPath& field_path);
// Writes FieldData values into a protobuf field.
absl::Status SetFieldValues(FieldData& message_data,
const FieldPath& field_path,
const std::vector<FieldData>& values);
// Merges FieldData values into a protobuf field.
absl::Status MergeFieldValues(FieldData& message_data,
const FieldPath& field_path,
const std::vector<FieldData>& values);
// Deserializes a packet containing a MessageLite value.
absl::Status ReadMessage(const std::string& value, const std::string& type_name,
Packet* result);
absl::StatusOr<Packet> ReadMessage(const std::string& value,
const std::string& type_name);
// Merge two options protobuf field values.
absl::Status MergeMessages(const FieldData& base, const FieldData& over,
FieldData* result);
absl::StatusOr<FieldData> MergeMessages(const FieldData& base,
const FieldData& over);
// Returns the requested options protobuf for a graph.
absl::Status GetNodeOptions(const FieldData& message_data,
const std::string& extension_type,
FieldData* result);
absl::StatusOr<FieldData> GetNodeOptions(const FieldData& message_data,
const std::string& extension_type);
// Returns the requested options protobuf for a graph node.
absl::Status GetGraphOptions(const FieldData& message_data,
const std::string& extension_type,
FieldData* result);
absl::StatusOr<FieldData> GetGraphOptions(const FieldData& message_data,
const std::string& extension_type);
// Sets the node_options field in a Node, and clears the options field.
void SetOptionsMessage(const FieldData& node_options,
@@ -67,10 +75,10 @@ void SetOptionsMessage(const FieldData& node_options,
FieldData AsFieldData(const proto_ns::MessageLite& message);
// Constructs a Packet for a FieldData proto.
absl::Status AsPacket(const FieldData& data, Packet* result);
absl::StatusOr<Packet> AsPacket(const FieldData& data);
// Constructs a FieldData proto for a Packet.
absl::Status AsFieldData(Packet packet, FieldData* result);
absl::StatusOr<FieldData> AsFieldData(Packet packet);
// Returns the protobuf type-url for a protobuf type-name.
std::string TypeUrl(absl::string_view type_name);
mediapipe/framework/tool/options_lib_template.cc
+7 -6
View File
@@ -25,11 +25,12 @@ constexpr char kDescriptorContents[] =
#include "{{DESCRIPTOR_INC_FILE_PATH}}"
; // NOLINT(whitespace/semicolon)
mediapipe::proto_ns::FileDescriptorSet ParseFileDescriptorSet(
const std::string& pb) {
mediapipe::proto_ns::FileDescriptorSet files;
files.ParseFromString(pb);
return files;
mediapipe::FieldData ReadFileDescriptorSet(const std::string& pb) {
mediapipe::FieldData result;
*result.mutable_message_value()->mutable_type_url() =
"proto2.FileDescriptorSet";
*result.mutable_message_value()->mutable_value() = pb;
return result;
}
} // namespace
@@ -39,6 +40,6 @@ namespace mediapipe {
template <>
const RegistrationToken tool::OptionsRegistry::registration_token<
MP_OPTION_TYPE_NS::MP_OPTION_TYPE_NAME> =
tool::OptionsRegistry::Register(ParseFileDescriptorSet(
tool::OptionsRegistry::Register(ReadFileDescriptorSet(
std::string(kDescriptorContents, sizeof(kDescriptorContents) - 1)));
} // namespace mediapipe
mediapipe/framework/tool/options_map.h
+66 -9
View File
@@ -30,15 +30,26 @@ struct IsExtension {
template <class T,
typename std::enable_if<IsExtension<T>::value, int>::type = 0>
void GetExtension(const CalculatorOptions& options, T* result) {
T* GetExtension(CalculatorOptions& options) {
if (options.HasExtension(T::ext)) {
*result = options.GetExtension(T::ext);
return options.MutableExtension(T::ext);
}
return nullptr;
}
template <class T,
typename std::enable_if<!IsExtension<T>::value, int>::type = 0>
void GetExtension(const CalculatorOptions& options, T* result) {}
T* GetExtension(const CalculatorOptions& options) {
return nullptr;
}
template <class T>
void GetExtension(const CalculatorOptions& options, T* result) {
T* r = GetExtension<T>(*const_cast<CalculatorOptions*>(&options));
if (r) {
*result = *r;
}
}
template <class T>
void GetNodeOptions(const CalculatorGraphConfig::Node& node_config, T* result) {
@@ -53,23 +64,39 @@ void GetNodeOptions(const CalculatorGraphConfig::Node& node_config, T* result) {
#endif
}
template <class T>
void SetNodeOptions(CalculatorGraphConfig::Node& node_config, const T& value) {
#if defined(MEDIAPIPE_PROTO_LITE) && defined(MEDIAPIPE_PROTO_THIRD_PARTY)
// protobuf::Any is unavailable with third_party/protobuf:protobuf-lite.
#else
for (mediapipe::protobuf::Any& options :
*node_config.mutable_node_options()) {
if (options.Is<T>()) {
options.PackFrom(value);
return;
}
}
node_config.add_node_options()->PackFrom(value);
#endif
}
// A map from object type to object.
class TypeMap {
public:
template <class T>
bool Has() const {
return content_.count(TypeInfo::Get<T>()) > 0;
return content_.count(kTypeId<T>) > 0;
}
template <class T>
T* Get() const {
if (!Has<T>()) {
content_[TypeInfo::Get<T>()] = std::make_shared<T>();
content_[kTypeId<T>] = std::make_shared<T>();
}
return static_cast<T*>(content_[TypeInfo::Get<T>()].get());
return static_cast<T*>(content_[kTypeId<T>].get());
}
private:
mutable std::map<TypeIndex, std::shared_ptr<void>> content_;
mutable std::map<TypeId, std::shared_ptr<void>> content_;
};
// Extracts the options message of a specified type from a
@@ -77,7 +104,7 @@ class TypeMap {
class OptionsMap {
public:
OptionsMap& Initialize(const CalculatorGraphConfig::Node& node_config) {
node_config_ = &node_config;
node_config_ = const_cast<CalculatorGraphConfig::Node*>(&node_config);
return *this;
}
@@ -97,10 +124,40 @@ class OptionsMap {
return *result;
}
const CalculatorGraphConfig::Node* node_config_;
CalculatorGraphConfig::Node* node_config_;
TypeMap options_;
};
class MutableOptionsMap : public OptionsMap {
public:
MutableOptionsMap& Initialize(CalculatorGraphConfig::Node& node_config) {
node_config_ = &node_config;
return *this;
}
template <class T>
void Set(const T& value) const {
*options_.Get<T>() = value;
if (node_config_->has_options()) {
*GetExtension<T>(*node_config_->mutable_options()) = value;
} else {
SetNodeOptions(*node_config_, value);
}
}
template <class T>
T* GetMutable() const {
if (options_.Has<T>()) {
return options_.Get<T>();
}
if (node_config_->has_options()) {
return GetExtension<T>(*node_config_->mutable_options());
}
T* result = options_.Get<T>();
GetNodeOptions(*node_config_, result);
return result;
}
};
} // namespace tool
} // namespace mediapipe
mediapipe/framework/tool/options_registry.cc
+140 -25
View File
@@ -1,6 +1,11 @@
#include "mediapipe/framework/tool/options_registry.h"
#include <string>
#include <vector>
#include "absl/synchronization/mutex.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/tool/proto_util_lite.h"
namespace mediapipe {
namespace tool {
@@ -9,37 +14,135 @@ namespace {
// Returns a canonical message type name, with any leading "." removed.
std::string CanonicalTypeName(const std::string& type_name) {
return (type_name.rfind('.', 0) == 0) ? type_name.substr(1) : type_name;
return (absl::StartsWith(type_name, ".")) ? type_name.substr(1) : type_name;
}
// Returns the values from a protobuf field as typed FieldData.
absl::StatusOr<std::vector<FieldData>> GetFieldValues(
const FieldData& message_data, std::string field_name) {
std::string type_name =
ProtoUtilLite::ParseTypeUrl(message_data.message_value().type_url());
const Descriptor* descriptor =
OptionsRegistry::GetProtobufDescriptor(type_name);
RET_CHECK_NE(descriptor, nullptr);
const FieldDescriptor* field = descriptor->FindFieldByName(field_name);
if (field == nullptr) {
return std::vector<FieldData>();
}
ProtoUtilLite::ProtoPath proto_path = {{field->number(), 0}};
ProtoUtilLite::FieldValue mesage_bytes = message_data.message_value().value();
int count;
MP_RETURN_IF_ERROR(ProtoUtilLite::GetFieldCount(mesage_bytes, proto_path,
field->type(), &count));
std::vector<std::string> field_values;
MP_RETURN_IF_ERROR(ProtoUtilLite::GetFieldRange(
mesage_bytes, proto_path, count, field->type(), &field_values));
std::vector<FieldData> result;
for (int i = 0; i < field_values.size(); ++i) {
FieldData r;
std::string message_type =
field->message_type() ? field->message_type()->full_name() : "";
MP_RETURN_IF_ERROR(ProtoUtilLite::ReadValue(field_values[i], field->type(),
message_type, &r));
result.push_back(std::move(r));
}
return result;
}
// Returns a single value from a protobuf string field.
std::string GetFieldString(const FieldData& message_data,
std::string field_name) {
auto values = GetFieldValues(message_data, field_name);
if (!values->empty()) {
return values->front().string_value();
}
return "";
}
// Registers the descriptors for the descriptor protobufs. These four
// descriptors are required to deserialize descriptors for other protobufs.
// This implementation avoids a code size problem introduced by
// proto_ns::DescriptorProto.
void RegisterDescriptorProtos(
absl::flat_hash_map<std::string, Descriptor>& result) {
std::vector<Descriptor> descriptors = {
{"proto2.FileDescriptorSet",
{
{"file", 1, FieldType::TYPE_MESSAGE, "proto2.FileDescriptorProto"},
}},
{"proto2.FileDescriptorProto",
{
{"package", 2, FieldType::TYPE_STRING, ""},
{"message_type", 4, FieldType::TYPE_MESSAGE,
"proto2.DescriptorProto"},
}},
{"proto2.DescriptorProto",
{
{"name", 1, FieldType::TYPE_STRING, ""},
{"field", 2, FieldType::TYPE_MESSAGE, "proto2.FieldDescriptorProto"},
{"extension", 6, FieldType::TYPE_MESSAGE,
"proto2.FieldDescriptorProto"},
{"nested_type", 3, FieldType::TYPE_MESSAGE,
"proto2.DescriptorProto"},
}},
{"proto2.FieldDescriptorProto",
{
{"name", 1, FieldType::TYPE_STRING, ""},
{"number", 3, FieldType::TYPE_INT32, ""},
{"type", 5, FieldType::TYPE_ENUM, ""},
{"type_name", 6, FieldType::TYPE_STRING, ""},
{"extendee", 2, FieldType::TYPE_STRING, ""},
}},
};
for (const auto& descriptor : descriptors) {
result[descriptor.full_name()] = descriptor;
}
}
} // namespace
RegistrationToken OptionsRegistry::Register(
const proto_ns::FileDescriptorSet& files) {
absl::MutexLock lock(&mutex());
for (auto& file : files.file()) {
for (auto& message_type : file.message_type()) {
Register(message_type, file.package());
const FieldData& file_descriptor_set) {
auto files = GetFieldValues(file_descriptor_set, "file");
for (auto& file : *files) {
std::string package_name = GetFieldString(file, "package");
auto message_types = GetFieldValues(file, "message_type");
for (auto& message_type : *message_types) {
Register(message_type, package_name);
}
}
return RegistrationToken([]() {});
}
void OptionsRegistry::Register(const proto_ns::DescriptorProto& message_type,
void OptionsRegistry::Register(const FieldData& message_type,
const std::string& parent_name) {
auto full_name = absl::StrCat(parent_name, ".", message_type.name());
descriptors()[full_name] = Descriptor(message_type, full_name);
for (auto& nested : message_type.nested_type()) {
std::string name = GetFieldString(message_type, "name");
std::string full_name = absl::StrCat(parent_name, ".", name);
Descriptor descriptor(full_name, message_type);
{
absl::MutexLock lock(&mutex());
descriptors()[full_name] = descriptor;
}
auto nested_types = GetFieldValues(message_type, "nested_type");
for (auto& nested : *nested_types) {
Register(nested, full_name);
}
for (auto& extension : message_type.extension()) {
extensions()[CanonicalTypeName(extension.extendee())].push_back(
FieldDescriptor(extension));
auto exts = GetFieldValues(message_type, "extension");
for (auto& extension : *exts) {
FieldDescriptor field(extension);
std::string extendee = GetFieldString(extension, "extendee");
{
absl::MutexLock lock(&mutex());
extensions()[CanonicalTypeName(extendee)].push_back(field);
}
}
}
const Descriptor* OptionsRegistry::GetProtobufDescriptor(
const std::string& type_name) {
if (descriptors().count("proto2.DescriptorProto") == 0) {
RegisterDescriptorProtos(descriptors());
}
absl::ReaderMutexLock lock(&mutex());
auto it = descriptors().find(CanonicalTypeName(type_name));
return (it == descriptors().end()) ? nullptr : &it->second;
@@ -73,11 +176,21 @@ absl::Mutex& OptionsRegistry::mutex() {
return *mutex;
}
Descriptor::Descriptor(const proto_ns::DescriptorProto& proto,
const std::string& full_name)
Descriptor::Descriptor(const std::string& full_name,
const FieldData& descriptor_proto)
: full_name_(full_name) {
for (auto& field : proto.field()) {
fields_[field.name()] = FieldDescriptor(field);
auto fields = GetFieldValues(descriptor_proto, "field");
for (const auto& field : *fields) {
FieldDescriptor f(field);
fields_[f.name()] = f;
}
}
Descriptor::Descriptor(const std::string& full_name,
const std::vector<FieldDescriptor>& fields)
: full_name_(full_name) {
for (const auto& field : fields) {
fields_[field.name()] = field;
}
}
@@ -89,20 +202,22 @@ const FieldDescriptor* Descriptor::FindFieldByName(
return (it != fields_.end()) ? &it->second : nullptr;
}
FieldDescriptor::FieldDescriptor(const proto_ns::FieldDescriptorProto& proto) {
name_ = proto.name();
message_type_ = CanonicalTypeName(proto.type_name());
type_ = proto.type();
number_ = proto.number();
FieldDescriptor::FieldDescriptor(const FieldData& field_proto) {
name_ = GetFieldString(field_proto, "name");
number_ = GetFieldValues(field_proto, "number")->front().int32_value();
type_ = (FieldType)GetFieldValues(field_proto, "type")->front().enum_value();
message_type_ = CanonicalTypeName(GetFieldString(field_proto, "type_name"));
}
FieldDescriptor::FieldDescriptor(std::string name, int number, FieldType type,
std::string message_type)
: name_(name), number_(number), type_(type), message_type_(message_type) {}
const std::string& FieldDescriptor::name() const { return name_; }
int FieldDescriptor::number() const { return number_; }
proto_ns::FieldDescriptorProto::Type FieldDescriptor::type() const {
return type_;
}
FieldType FieldDescriptor::type() const { return type_; }
const Descriptor* FieldDescriptor::message_type() const {
return OptionsRegistry::GetProtobufDescriptor(message_type_);
mediapipe/framework/tool/options_registry.h
+20 -12
View File
@@ -1,15 +1,20 @@
#ifndef MEDIAPIPE_FRAMEWORK_TOOL_OPTIONS_REGISTRY_H_
#define MEDIAPIPE_FRAMEWORK_TOOL_OPTIONS_REGISTRY_H_
#include <string>
#include "absl/container/flat_hash_map.h"
#include "mediapipe/framework/deps/registration.h"
#include "mediapipe/framework/port/advanced_proto_inc.h"
#include "mediapipe/framework/tool/field_data.pb.h"
namespace mediapipe {
namespace tool {
class Descriptor;
class FieldDescriptor;
using FieldType = mediapipe::proto_ns::internal::WireFormatLite::FieldType;
using mediapipe::FieldData;
// A static registry that stores descriptors for protobufs used in MediaPipe
// calculator options. Lite-proto builds do not normally include descriptors.
@@ -17,8 +22,8 @@ class FieldDescriptor;
// referenced and specified separately within CalculatorGraphConfigs.
class OptionsRegistry {
public:
// Registers the protobuf descriptors for a MessageLite.
static RegistrationToken Register(const proto_ns::FileDescriptorSet& files);
// Registers the protobuf descriptors for a FileDescriptorSet.
static RegistrationToken Register(const FieldData& file_descriptor_set);
// Finds the descriptor for a protobuf.
static const Descriptor* GetProtobufDescriptor(const std::string& type_name);
@@ -28,8 +33,8 @@ class OptionsRegistry {
std::vector<const FieldDescriptor*>* result);
private:
// Registers protobuf descriptors a MessageLite and nested types.
static void Register(const proto_ns::DescriptorProto& message_type,
// Registers protobuf descriptors for a message type and nested types.
static void Register(const FieldData& message_type,
const std::string& parent_name);
static absl::flat_hash_map<std::string, Descriptor>& descriptors();
@@ -46,9 +51,10 @@ class OptionsRegistry {
// avoids a code size problem introduced by proto_ns::FieldDescriptor.
class Descriptor {
public:
Descriptor() {}
Descriptor(const proto_ns::DescriptorProto& proto,
const std::string& full_name);
Descriptor() = default;
Descriptor(const std::string& full_name, const FieldData& descriptor_proto);
Descriptor(const std::string& full_name,
const std::vector<FieldDescriptor>& fields);
const std::string& full_name() const;
const FieldDescriptor* FindFieldByName(const std::string& name) const;
@@ -61,18 +67,20 @@ class Descriptor {
// avoids a code size problem introduced by proto_ns::FieldDescriptor.
class FieldDescriptor {
public:
FieldDescriptor() {}
FieldDescriptor(const proto_ns::FieldDescriptorProto& proto);
FieldDescriptor() = default;
FieldDescriptor(const FieldData& field_proto);
FieldDescriptor(std::string name, int number, FieldType type,
std::string message_type);
const std::string& name() const;
int number() const;
proto_ns::FieldDescriptorProto::Type type() const;
FieldType type() const;
const Descriptor* message_type() const;
private:
std::string name_;
std::string message_type_;
proto_ns::FieldDescriptorProto::Type type_;
int number_;
FieldType type_;
std::string message_type_;
};
} // namespace tool
mediapipe/framework/tool/options_syntax_util.cc
+2 -3
View File
@@ -91,8 +91,7 @@ class OptionsSyntaxUtil::OptionsSyntaxHelper {
int index;
if (absl::SimpleAtoi(option_name, &index)) {
result.back().index = index;
}
if (!ExtensionType(option_name).empty()) {
} else if (!ExtensionType(option_name).empty()) {
std::string extension_type = std::string(ExtensionType(option_name));
result.push_back({nullptr, 0, extension_type});
descriptor = OptionsRegistry::GetProtobufDescriptor(extension_type);
@@ -102,7 +101,7 @@ class OptionsSyntaxUtil::OptionsSyntaxHelper {
}
auto field = descriptor->FindFieldByName(std::string(option_name));
descriptor = field ? field->message_type() : nullptr;
result.push_back({std::move(field), 0});
result.push_back({std::move(field), -1});
}
}
return result;
mediapipe/framework/tool/options_util.cc
+24 -15
View File
@@ -26,10 +26,9 @@ namespace mediapipe {
namespace tool {
using options_field_util::FieldPath;
using options_field_util::GetField;
using options_field_util::GetGraphOptions;
using options_field_util::GetNodeOptions;
using options_field_util::MergeField;
using options_field_util::MergeFieldValues;
using options_field_util::MergeMessages;
// Returns the type for the root options message if specified.
@@ -56,10 +55,19 @@ std::string MessageType(FieldData message) {
std::string(message.message_value().type_url()));
}
// Assigns the value from a StatusOr if avialable.
#define ASSIGN_IF_OK(lhs, rexpr) \
{ \
auto statusor = (rexpr); \
if (statusor.ok()) { \
lhs = statusor.value(); \
} \
}
// Copy literal options from graph_options to node_options.
absl::Status CopyLiteralOptions(CalculatorGraphConfig::Node parent_node,
CalculatorGraphConfig* config) {
Status status;
absl::Status status;
FieldData graph_data = options_field_util::AsFieldData(*config);
FieldData parent_data = options_field_util::AsFieldData(parent_node);
@@ -75,25 +83,26 @@ absl::Status CopyLiteralOptions(CalculatorGraphConfig::Node parent_node,
std::string node_tag = syntax_util.OptionFieldsTag(tag_and_name[0]);
std::string node_extension_type = ExtensionType(node_tag);
FieldData graph_options;
GetGraphOptions(graph_data, graph_extension_type, &graph_options)
.IgnoreError();
ASSIGN_IF_OK(graph_options,
GetGraphOptions(graph_data, graph_extension_type));
FieldData parent_options;
GetNodeOptions(parent_data, graph_extension_type, &parent_options)
.IgnoreError();
status.Update(
MergeMessages(graph_options, parent_options, &graph_options));
ASSIGN_IF_OK(parent_options,
GetNodeOptions(parent_data, graph_extension_type));
ASSIGN_OR_RETURN(graph_options,
MergeMessages(graph_options, parent_options));
FieldData node_options;
status.Update(
GetNodeOptions(node_data, node_extension_type, &node_options));
ASSIGN_OR_RETURN(node_options,
GetNodeOptions(node_data, node_extension_type));
if (!node_options.has_message_value() ||
!graph_options.has_message_value()) {
continue;
}
FieldPath graph_path = GetPath(graph_tag, MessageType(graph_options));
FieldPath node_path = GetPath(node_tag, MessageType(node_options));
FieldData packet_data;
status.Update(GetField(graph_path, graph_options, &packet_data));
status.Update(MergeField(node_path, packet_data, &node_options));
std::vector<FieldData> packet_data;
ASSIGN_OR_RETURN(packet_data, GetFieldValues(graph_options, graph_path));
MP_RETURN_IF_ERROR(
MergeFieldValues(node_options, node_path, packet_data));
options_field_util::SetOptionsMessage(node_options, &node);
}
node.clear_option_value();
@@ -105,7 +114,7 @@ absl::Status CopyLiteralOptions(CalculatorGraphConfig::Node parent_node,
absl::Status DefineGraphOptions(const CalculatorGraphConfig::Node& parent_node,
CalculatorGraphConfig* config) {
MP_RETURN_IF_ERROR(CopyLiteralOptions(parent_node, config));
return mediapipe::OkStatus();
return absl::OkStatus();
}
} // namespace tool
mediapipe/framework/tool/options_util_test.cc
+353 -19
View File
@@ -13,8 +13,10 @@
// limitations under the License.
#include <memory>
#include <sstream>
#include <vector>
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/gmock.h"
@@ -30,23 +32,27 @@
namespace mediapipe {
namespace {
using ::mediapipe::proto_ns::FieldDescriptorProto;
using FieldType = ::mediapipe::proto_ns::FieldDescriptorProto::Type;
using ::testing::HasSubstr;
// Assigns the value from a StatusOr if avialable.
#define ASSERT_AND_ASSIGN(lhs, rexpr) \
{ \
auto statusor = (rexpr); \
MP_ASSERT_OK(statusor); \
lhs = statusor.value(); \
}
// A test Calculator using DeclareOptions and DefineOptions.
class NightLightCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc) {
return mediapipe::OkStatus();
return absl::OkStatus();
}
absl::Status Open(CalculatorContext* cc) final {
return mediapipe::OkStatus();
}
absl::Status Open(CalculatorContext* cc) final { return absl::OkStatus(); }
absl::Status Process(CalculatorContext* cc) final {
return mediapipe::OkStatus();
}
absl::Status Process(CalculatorContext* cc) final { return absl::OkStatus(); }
private:
NightLightCalculatorOptions options_;
@@ -124,7 +130,7 @@ TEST_F(OptionsUtilTest, CopyLiteralOptions) {
CalculatorGraph graph;
graph_config.set_num_threads(4);
MP_EXPECT_OK(graph.Initialize({subgraph_config, graph_config}, {}, {}));
MP_ASSERT_OK(graph.Initialize({subgraph_config, graph_config}, {}, {}));
CalculatorGraphConfig expanded_config = graph.Config();
expanded_config.clear_executor();
@@ -236,8 +242,8 @@ TEST_F(OptionsUtilTest, FindOptionsMessage) {
tool::options_field_util::FieldPath field_path =
syntax_util.OptionFieldPath(split[1], descriptor);
EXPECT_EQ(field_path.size(), 2);
EXPECT_TRUE(Equals(field_path[0], "sub_options", 0, ""));
EXPECT_TRUE(Equals(field_path[1], "num_lights", 0, ""));
EXPECT_TRUE(Equals(field_path[0], "sub_options", -1, ""));
EXPECT_TRUE(Equals(field_path[1], "num_lights", -1, ""));
{
// NightLightCalculatorOptions in Node.options.
@@ -252,11 +258,11 @@ TEST_F(OptionsUtilTest, FindOptionsMessage) {
auto path = field_path;
std::string node_extension_type = ExtensionType(std::string(split[1]));
FieldData node_options;
MP_EXPECT_OK(tool::options_field_util::GetNodeOptions(
node_data, node_extension_type, &node_options));
ASSERT_AND_ASSIGN(node_options, tool::options_field_util::GetNodeOptions(
node_data, node_extension_type));
FieldData packet_data;
MP_EXPECT_OK(tool::options_field_util::GetField(field_path, node_options,
&packet_data));
ASSERT_AND_ASSIGN(packet_data, tool::options_field_util::GetField(
node_options, field_path));
EXPECT_EQ(packet_data.value_case(), FieldData::kInt32Value);
EXPECT_EQ(packet_data.int32_value(), 33);
}
@@ -273,11 +279,11 @@ TEST_F(OptionsUtilTest, FindOptionsMessage) {
auto path = field_path;
std::string node_extension_type = ExtensionType(std::string(split[1]));
FieldData node_options;
MP_EXPECT_OK(tool::options_field_util::GetNodeOptions(
node_data, node_extension_type, &node_options));
ASSERT_AND_ASSIGN(node_options, tool::options_field_util::GetNodeOptions(
node_data, node_extension_type));
FieldData packet_data;
MP_EXPECT_OK(tool::options_field_util::GetField(field_path, node_options,
&packet_data));
ASSERT_AND_ASSIGN(packet_data, tool::options_field_util::GetField(
node_options, field_path));
EXPECT_EQ(packet_data.value_case(), FieldData::kInt32Value);
EXPECT_EQ(packet_data.int32_value(), 33);
}
@@ -285,5 +291,333 @@ TEST_F(OptionsUtilTest, FindOptionsMessage) {
// TODO: Test with specified extension_type.
}
// Constructs the field path for a string of field names.
FieldPath MakeFieldPath(std::string tag, FieldData message_data) {
tool::OptionsSyntaxUtil syntax_util;
const tool::Descriptor* descriptor =
tool::OptionsRegistry::GetProtobufDescriptor(
tool::options_field_util::ParseTypeUrl(
message_data.message_value().type_url()));
return syntax_util.OptionFieldPath(tag, descriptor);
}
// Returns the field path addressing the entire specified field.
FieldPath EntireField(FieldPath field_path) {
field_path.back().index = -1;
return field_path;
}
// Converts an int to a FieldData record.
FieldData AsFieldData(int v) {
return tool::options_field_util::AsFieldData(MakePacket<int>(v)).value();
}
// Equality comparison for field contents.
template <typename T>
absl::Status Equals(const T& v1, const T& v2) {
RET_CHECK_EQ(v1, v2);
return absl::OkStatus();
}
// Equality comparison for protobuf field contents.
// The generic Equals() fails because MessageLite lacks operator==().
// The protobuf comparison is performed using testing::EqualsProto.
using LightBundle = NightLightCalculatorOptions::LightBundle;
template <>
absl::Status Equals<LightBundle>(const LightBundle& v1, const LightBundle& v2) {
std::string s_1, s_2;
v1.SerializeToString(&s_1);
v2.SerializeToString(&s_2);
RET_CHECK(s_1 == s_2);
return absl::OkStatus();
}
// Equality comparison for FieldData vectors.
template <typename FieldType>
absl::Status Equals(std::vector<FieldData> b1, std::vector<FieldData> b2) {
using tool::options_field_util::AsPacket;
RET_CHECK_EQ(b1.size(), b2.size());
for (int i = 0; i < b1.size(); ++i) {
ASSIGN_OR_RETURN(Packet p1, AsPacket(b1.at(i)));
ASSIGN_OR_RETURN(Packet p2, AsPacket(b2.at(i)));
MP_RETURN_IF_ERROR(Equals(p1.Get<FieldType>(), p2.Get<FieldType>()));
}
return absl::OkStatus();
}
// Unit-tests for graph options feild accessors from options_field_util.
class OptionsFieldUtilTest : public ::testing::Test {
protected:
void SetUp() override {}
void TearDown() override {}
};
// Tests empty FieldPaths applied to empty options.
TEST_F(OptionsFieldUtilTest, EmptyFieldPaths) {
FieldData graph_options;
FieldData node_options;
FieldPath graph_path;
FieldPath node_path;
std::vector<FieldData> packet_data;
ASSERT_AND_ASSIGN(packet_data, GetFieldValues(graph_options, graph_path));
MP_EXPECT_OK(MergeFieldValues(node_options, node_path, packet_data));
}
// Tests GetFieldValues applied to an int field.
TEST_F(OptionsFieldUtilTest, GetFieldValuesInt) {
NightLightCalculatorOptions node_proto;
node_proto.mutable_sub_options();
node_proto.mutable_sub_options()->add_num_lights(33);
node_proto.mutable_sub_options()->add_num_lights(44);
FieldData node_data = tool::options_field_util::AsFieldData(node_proto);
// Read an entire populated repeated field.
FieldPath path = MakeFieldPath("OPTIONS/sub_options/num_lights", node_data);
MP_EXPECT_OK(Equals<int>(GetFieldValues(node_data, path).value(),
{AsFieldData(33), AsFieldData(44)}));
// Read a specific populated repeated field index.
path = MakeFieldPath("OPTIONS/sub_options/num_lights/1", node_data);
MP_EXPECT_OK(
Equals<int>(GetFieldValues(node_data, path).value(), {AsFieldData(44)}));
}
// Tests GetFieldValues applied to a protobuf field.
TEST_F(OptionsFieldUtilTest, GetFieldValuesProtobuf) {
using tool::options_field_util::AsFieldData;
using LightBundle = NightLightCalculatorOptions::LightBundle;
NightLightCalculatorOptions node_proto;
node_proto.mutable_sub_options();
node_proto.mutable_sub_options()->add_bundle();
*node_proto.mutable_sub_options()->mutable_bundle(0)->mutable_room_id() =
"111";
node_proto.mutable_sub_options()
->mutable_bundle(0)
->add_room_lights()
->set_frame_rate(11.1);
node_proto.mutable_sub_options()
->mutable_bundle(0)
->add_room_lights()
->set_frame_rate(22.1);
FieldData node_data = AsFieldData(node_proto);
// Read all values from a repeated protobuf field.
LightBundle expected_proto;
*expected_proto.mutable_room_id() = "111";
expected_proto.add_room_lights()->set_frame_rate(11.1);
expected_proto.add_room_lights()->set_frame_rate(22.1);
FieldData expected_data = AsFieldData(expected_proto);
FieldPath path = MakeFieldPath("OPTIONS/sub_options/bundle", node_data);
MP_EXPECT_OK(Equals<LightBundle>(GetFieldValues(node_data, path).value(),
{expected_data}));
// Read a specific index from a repeated protobuf field.
path = MakeFieldPath("OPTIONS/sub_options/bundle/0", node_data);
MP_EXPECT_OK(Equals<LightBundle>(GetFieldValues(node_data, path).value(),
{expected_data}));
}
// Tests SetFieldValues applied to an int field.
TEST_F(OptionsFieldUtilTest, SetFieldValuesInt) {
NightLightCalculatorOptions node_proto;
node_proto.mutable_sub_options();
FieldData node_data = tool::options_field_util::AsFieldData(node_proto);
// Replace an entire empty repeated field.
FieldPath path = MakeFieldPath("OPTIONS/sub_options/num_lights", node_data);
MP_ASSERT_OK(SetFieldValues(node_data, path, {AsFieldData(33)}));
MP_EXPECT_OK(
Equals<int>(GetFieldValues(node_data, path).value(), {AsFieldData(33)}));
// Replace an entire populated repeated field.
MP_ASSERT_OK(SetFieldValues(node_data, path, {AsFieldData(44)}));
MP_EXPECT_OK(
Equals<int>(GetFieldValues(node_data, path).value(), {AsFieldData(44)}));
// Replace an entire repeated field with a new list of values.
MP_ASSERT_OK(
SetFieldValues(node_data, path, {AsFieldData(33), AsFieldData(44)}));
MP_EXPECT_OK(Equals<int>(GetFieldValues(node_data, path).value(),
{AsFieldData(33), AsFieldData(44)}));
// Replace a single field index with a new list of values.
path = MakeFieldPath("OPTIONS/sub_options/num_lights/1", node_data);
MP_ASSERT_OK(
SetFieldValues(node_data, path, {AsFieldData(55), AsFieldData(66)}));
MP_EXPECT_OK(
Equals<int>(GetFieldValues(node_data, EntireField(path)).value(),
{AsFieldData(33), AsFieldData(55), AsFieldData(66)}));
// Replace a single field middle index with a new list of values.
path = MakeFieldPath("OPTIONS/sub_options/num_lights/1", node_data);
MP_ASSERT_OK(
SetFieldValues(node_data, path, {AsFieldData(11), AsFieldData(12)}));
MP_EXPECT_OK(Equals<int>(
GetFieldValues(node_data, EntireField(path)).value(),
{AsFieldData(33), AsFieldData(11), AsFieldData(12), AsFieldData(66)}));
// Replace field index 0 with a new value.
path = MakeFieldPath("OPTIONS/sub_options/num_lights/0", node_data);
MP_ASSERT_OK(SetFieldValues(node_data, path, {AsFieldData(77)}));
MP_EXPECT_OK(Equals<int>(
GetFieldValues(node_data, EntireField(path)).value(),
{AsFieldData(77), AsFieldData(11), AsFieldData(12), AsFieldData(66)}));
// Replace field index 0 with an empty list of values.
MP_ASSERT_OK(SetFieldValues(node_data, path, {}));
MP_EXPECT_OK(
Equals<int>(GetFieldValues(node_data, EntireField(path)).value(),
{AsFieldData(11), AsFieldData(12), AsFieldData(66)}));
// Replace an entire populated field with an empty list of values.
path = MakeFieldPath("OPTIONS/sub_options/num_lights", node_data);
MP_ASSERT_OK(SetFieldValues(node_data, path, {}));
MP_ASSERT_OK(
Equals<int>(GetFieldValues(node_data, EntireField(path)).value(), {}));
// Replace a missing field index with new values.
path = MakeFieldPath("OPTIONS/sub_options/num_lights/1", node_data);
absl::Status status =
SetFieldValues(node_data, path, {AsFieldData(55), AsFieldData(66)});
EXPECT_EQ(status.code(), absl::StatusCode::kInternal);
// TODO: status.message() appears empty on KokoroGCPDocker.
// EXPECT_THAT(status.message(),
// HasSubstr("index >= 0 && index <= v.size()"));
}
// Tests SetFieldValues applied to a protobuf field.
TEST_F(OptionsFieldUtilTest, SetFieldValuesProtobuf) {
using tool::options_field_util::AsFieldData;
using LightBundle = NightLightCalculatorOptions::LightBundle;
NightLightCalculatorOptions node_proto;
node_proto.mutable_sub_options();
FieldData node_data = AsFieldData(node_proto);
// Replace an empty repeated protobuf field.
LightBundle bundle_proto;
*bundle_proto.mutable_room_id() = "222";
bundle_proto.add_room_lights()->set_frame_rate(22.1);
FieldData bundle_data = AsFieldData(bundle_proto);
FieldData expected_data = bundle_data;
FieldPath path = MakeFieldPath("OPTIONS/sub_options/bundle", node_data);
MP_ASSERT_OK(SetFieldValues(node_data, path, {bundle_data}));
MP_EXPECT_OK(Equals<LightBundle>(
GetFieldValues(node_data, EntireField(path)).value(), {expected_data}));
// Replace a populated repeated protobuf field.
*bundle_proto.mutable_room_id() = "333";
bundle_proto.mutable_room_lights(0)->set_frame_rate(33.1);
bundle_data = AsFieldData(bundle_proto);
LightBundle expected_proto;
*expected_proto.mutable_room_id() = "333";
expected_proto.add_room_lights()->set_frame_rate(33.1);
expected_data = AsFieldData(expected_proto);
MP_ASSERT_OK(SetFieldValues(node_data, path, {bundle_data}));
MP_EXPECT_OK(Equals<LightBundle>(
GetFieldValues(node_data, EntireField(path)).value(), {expected_data}));
}
// Tests MergeFieldValues applied to an int field.
TEST_F(OptionsFieldUtilTest, MergeFieldValuesInt) {
NightLightCalculatorOptions node_proto;
node_proto.mutable_sub_options();
FieldData node_data = tool::options_field_util::AsFieldData(node_proto);
// Replace an entire empty repeated field.
FieldPath path = MakeFieldPath("OPTIONS/sub_options/num_lights", node_data);
MP_ASSERT_OK(MergeFieldValues(node_data, path, {AsFieldData(33)}));
MP_EXPECT_OK(
Equals<int>(GetFieldValues(node_data, path).value(), {AsFieldData(33)}));
// Replace an entire populated repeated field.
MP_ASSERT_OK(MergeFieldValues(node_data, path, {AsFieldData(44)}));
MP_EXPECT_OK(
Equals<int>(GetFieldValues(node_data, path).value(), {AsFieldData(44)}));
// Replace an entire repeated field with a new list of values.
MP_ASSERT_OK(
MergeFieldValues(node_data, path, {AsFieldData(33), AsFieldData(44)}));
MP_EXPECT_OK(Equals<int>(GetFieldValues(node_data, path).value(),
{AsFieldData(33), AsFieldData(44)}));
// Replace a singe field index with a new list of values.
path = MakeFieldPath("OPTIONS/sub_options/num_lights/1", node_data);
MP_ASSERT_OK(
MergeFieldValues(node_data, path, {AsFieldData(55), AsFieldData(66)}));
MP_EXPECT_OK(
Equals<int>(GetFieldValues(node_data, EntireField(path)).value(),
{AsFieldData(33), AsFieldData(55), AsFieldData(66)}));
// Replace a single field middle index with a new list of values.
path = MakeFieldPath("OPTIONS/sub_options/num_lights/1", node_data);
MP_ASSERT_OK(
MergeFieldValues(node_data, path, {AsFieldData(11), AsFieldData(12)}));
MP_EXPECT_OK(Equals<int>(
GetFieldValues(node_data, EntireField(path)).value(),
{AsFieldData(33), AsFieldData(11), AsFieldData(12), AsFieldData(66)}));
// Replace field index 0 with a new value.
path = MakeFieldPath("OPTIONS/sub_options/num_lights/0", node_data);
MP_ASSERT_OK(MergeFieldValues(node_data, path, {AsFieldData(77)}));
MP_EXPECT_OK(Equals<int>(
GetFieldValues(node_data, EntireField(path)).value(),
{AsFieldData(77), AsFieldData(11), AsFieldData(12), AsFieldData(66)}));
// Replace field index 0 with an empty list of values.
MP_ASSERT_OK(MergeFieldValues(node_data, path, {}));
MP_EXPECT_OK(
Equals<int>(GetFieldValues(node_data, EntireField(path)).value(),
{AsFieldData(11), AsFieldData(12), AsFieldData(66)}));
// Replace an entire populated field with an empty list of values.
path = MakeFieldPath("OPTIONS/sub_options/num_lights", node_data);
MP_ASSERT_OK(MergeFieldValues(node_data, path, {}));
MP_EXPECT_OK(
Equals<int>(GetFieldValues(node_data, EntireField(path)).value(), {}));
// Replace a missing field index with new values.
path = MakeFieldPath("OPTIONS/sub_options/num_lights/1", node_data);
absl::Status status =
MergeFieldValues(node_data, path, {AsFieldData(55), AsFieldData(66)});
EXPECT_EQ(status.code(), absl::StatusCode::kOutOfRange);
EXPECT_THAT(status.message(),
HasSubstr("Missing feild value: num_lights at index: 1"));
}
// Tests MergeFieldValues applied to a protobuf field.
TEST_F(OptionsFieldUtilTest, MergeFieldValuesProtobuf) {
using tool::options_field_util::AsFieldData;
using LightBundle = NightLightCalculatorOptions::LightBundle;
NightLightCalculatorOptions node_proto;
node_proto.mutable_sub_options();
FieldData node_data = AsFieldData(node_proto);
// Merge an empty repeated protobuf field.
LightBundle bundle_proto;
*bundle_proto.mutable_room_id() = "222";
bundle_proto.add_room_lights()->set_frame_rate(22.1);
FieldData bundle_data = AsFieldData(bundle_proto);
FieldData expected_data = bundle_data;
FieldPath path = MakeFieldPath("OPTIONS/sub_options/bundle", node_data);
MP_ASSERT_OK(MergeFieldValues(node_data, path, {bundle_data}));
MP_EXPECT_OK(Equals<LightBundle>(
GetFieldValues(node_data, EntireField(path)).value(), {expected_data}));
// Merge a populated repeated protobuf field.
// "LightBundle.room_id" merges to "333".
// "LightBundle.room_lights" merges to {{22.1}, {33.1}}.
*bundle_proto.mutable_room_id() = "333";
bundle_proto.mutable_room_lights(0)->set_frame_rate(33.1);
bundle_data = AsFieldData(bundle_proto);
LightBundle expected_proto;
*expected_proto.mutable_room_id() = "333";
expected_proto.add_room_lights()->set_frame_rate(22.1);
expected_proto.add_room_lights()->set_frame_rate(33.1);
expected_data = AsFieldData(expected_proto);
MP_ASSERT_OK(MergeFieldValues(node_data, path, {bundle_data}));
MP_EXPECT_OK(Equals<LightBundle>(
GetFieldValues(node_data, EntireField(path)).value(), {expected_data}));
}
} // namespace
} // namespace mediapipe
mediapipe/framework/tool/proto_util_lite.cc
+147
View File
@@ -16,11 +16,13 @@
#include <tuple>
#include "absl/strings/match.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/tool/field_data.pb.h"
#include "mediapipe/framework/type_map.h"
#define RET_CHECK_NO_LOG(cond) RET_CHECK(cond).SetNoLogging()
@@ -37,6 +39,7 @@ using FieldAccess = ProtoUtilLite::FieldAccess;
using FieldValue = ProtoUtilLite::FieldValue;
using ProtoPath = ProtoUtilLite::ProtoPath;
using FieldType = ProtoUtilLite::FieldType;
using mediapipe::FieldData;
// Returns true if a wire type includes a length indicator.
bool IsLengthDelimited(WireFormatLite::WireType wire_type) {
@@ -408,5 +411,149 @@ absl::Status ProtoUtilLite::Deserialize(
return absl::OkStatus();
}
absl::Status ProtoUtilLite::WriteValue(const FieldData& value,
FieldType field_type,
std::string* field_bytes) {
StringOutputStream sos(field_bytes);
CodedOutputStream out(&sos);
switch (field_type) {
case WireFormatLite::TYPE_INT32:
WireFormatLite::WriteInt32NoTag(value.int32_value(), &out);
break;
case WireFormatLite::TYPE_SINT32:
WireFormatLite::WriteSInt32NoTag(value.int32_value(), &out);
break;
case WireFormatLite::TYPE_INT64:
WireFormatLite::WriteInt64NoTag(value.int64_value(), &out);
break;
case WireFormatLite::TYPE_SINT64:
WireFormatLite::WriteSInt64NoTag(value.int64_value(), &out);
break;
case WireFormatLite::TYPE_UINT32:
WireFormatLite::WriteUInt32NoTag(value.uint32_value(), &out);
break;
case WireFormatLite::TYPE_UINT64:
WireFormatLite::WriteUInt64NoTag(value.uint64_value(), &out);
break;
case WireFormatLite::TYPE_DOUBLE:
WireFormatLite::WriteDoubleNoTag(value.uint64_value(), &out);
break;
case WireFormatLite::TYPE_FLOAT:
WireFormatLite::WriteFloatNoTag(value.float_value(), &out);
break;
case WireFormatLite::TYPE_BOOL:
WireFormatLite::WriteBoolNoTag(value.bool_value(), &out);
break;
case WireFormatLite::TYPE_ENUM:
WireFormatLite::WriteEnumNoTag(value.enum_value(), &out);
break;
case WireFormatLite::TYPE_STRING:
out.WriteString(value.string_value());
break;
case WireFormatLite::TYPE_MESSAGE:
out.WriteString(value.message_value().value());
break;
default:
return absl::UnimplementedError(
absl::StrCat("Cannot write type: ", field_type));
}
return absl::OkStatus();
}
template <typename ValueT, FieldType kFieldType>
static ValueT ReadValue(absl::string_view field_bytes, absl::Status* status) {
ArrayInputStream ais(field_bytes.data(), field_bytes.size());
CodedInputStream input(&ais);
ValueT result;
if (!WireFormatLite::ReadPrimitive<ValueT, kFieldType>(&input, &result)) {
status->Update(absl::InvalidArgumentError(absl::StrCat(
"Bad serialized value: ", MediaPipeTypeStringOrDemangled<ValueT>(),
".")));
}
return result;
}
absl::Status ReadValue(absl::string_view field_bytes, FieldType field_type,
absl::string_view message_type, FieldData* result) {
absl::Status status;
result->Clear();
switch (field_type) {
case WireFormatLite::TYPE_INT32:
result->set_int32_value(
ReadValue<int32, WireFormatLite::TYPE_INT32>(field_bytes, &status));
break;
case WireFormatLite::TYPE_SINT32:
result->set_int32_value(
ReadValue<int32, WireFormatLite::TYPE_SINT32>(field_bytes, &status));
break;
case WireFormatLite::TYPE_INT64:
result->set_int64_value(
ReadValue<int64, WireFormatLite::TYPE_INT64>(field_bytes, &status));
break;
case WireFormatLite::TYPE_SINT64:
result->set_int64_value(
ReadValue<int64, WireFormatLite::TYPE_SINT64>(field_bytes, &status));
break;
case WireFormatLite::TYPE_UINT32:
result->set_uint32_value(
ReadValue<uint32, WireFormatLite::TYPE_UINT32>(field_bytes, &status));
break;
case WireFormatLite::TYPE_UINT64:
result->set_uint64_value(
ReadValue<uint32, WireFormatLite::TYPE_UINT32>(field_bytes, &status));
break;
case WireFormatLite::TYPE_DOUBLE:
result->set_double_value(
ReadValue<double, WireFormatLite::TYPE_DOUBLE>(field_bytes, &status));
break;
case WireFormatLite::TYPE_FLOAT:
result->set_float_value(
ReadValue<float, WireFormatLite::TYPE_FLOAT>(field_bytes, &status));
break;
case WireFormatLite::TYPE_BOOL:
result->set_bool_value(
ReadValue<bool, WireFormatLite::TYPE_BOOL>(field_bytes, &status));
break;
case WireFormatLite::TYPE_ENUM:
result->set_enum_value(
ReadValue<int32, WireFormatLite::TYPE_ENUM>(field_bytes, &status));
break;
case WireFormatLite::TYPE_STRING:
result->set_string_value(std::string(field_bytes));
break;
case WireFormatLite::TYPE_MESSAGE:
result->mutable_message_value()->set_value(std::string(field_bytes));
result->mutable_message_value()->set_type_url(
ProtoUtilLite::TypeUrl(message_type));
break;
default:
status = absl::UnimplementedError(
absl::StrCat("Cannot read type: ", field_type));
break;
}
return status;
}
absl::Status ProtoUtilLite::ReadValue(absl::string_view field_bytes,
FieldType field_type,
absl::string_view message_type,
FieldData* result) {
return mediapipe::tool::ReadValue(field_bytes, field_type, message_type,
result);
}
std::string ProtoUtilLite::TypeUrl(absl::string_view type_name) {
constexpr std::string_view kTypeUrlPrefix = "type.googleapis.com/";
return absl::StrCat(std::string(kTypeUrlPrefix), std::string(type_name));
}
std::string ProtoUtilLite::ParseTypeUrl(absl::string_view type_url) {
constexpr std::string_view kTypeUrlPrefix = "type.googleapis.com/";
if (absl::StartsWith(std::string(type_url), std::string(kTypeUrlPrefix))) {
return std::string(type_url.substr(kTypeUrlPrefix.length()));
}
return std::string(type_url);
}
} // namespace tool
} // namespace mediapipe
mediapipe/framework/tool/proto_util_lite.h
+19
View File
@@ -23,10 +23,12 @@
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/proto_ns.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/tool/field_data.pb.h"
namespace mediapipe {
namespace tool {
// TODO: Replace this class with a namespace following Google style.
class ProtoUtilLite {
public:
// Defines field types and tag formats.
@@ -89,6 +91,23 @@ class ProtoUtilLite {
static absl::Status Deserialize(const std::vector<FieldValue>& field_values,
FieldType field_type,
std::vector<std::string>* result);
// Write a protobuf field value from a typed FieldData value.
static absl::Status WriteValue(const mediapipe::FieldData& value,
FieldType field_type,
std::string* field_bytes);
// Read a protobuf field value into a typed FieldData value.
static absl::Status ReadValue(absl::string_view field_bytes,
FieldType field_type,
absl::string_view message_type,
mediapipe::FieldData* result);
// Returns the protobuf type-url for a protobuf type-name.
static std::string TypeUrl(absl::string_view type_name);
// Returns the protobuf type-name for a protobuf type-url.
static std::string ParseTypeUrl(absl::string_view type_url);
};
} // namespace tool
mediapipe/framework/tool/status_util.cc
+2 -1
View File
@@ -59,7 +59,8 @@ absl::Status CombinedStatus(const std::string& general_comment,
}
}
if (error_code == StatusCode::kOk) return OkStatus();
Status combined = absl::Status(
Status combined;
combined = absl::Status(
error_code,
absl::StrCat(general_comment, "\n", absl::StrJoin(errors, "\n")));
return combined;
mediapipe/framework/tool/status_util_test.cc
+3
View File
@@ -28,8 +28,11 @@ namespace mediapipe {
namespace {
using testing::ContainerEq;
using testing::Eq;
using testing::HasSubstr;
using testing::IsEmpty;
using testing::Matches;
using testing::Pointwise;
TEST(StatusTest, StatusStopIsNotOk) { EXPECT_FALSE(tool::StatusStop().ok()); }
mediapipe/framework/tool/subgraph_expansion.cc
+1 -1
View File
@@ -293,7 +293,7 @@ absl::Status ExpandSubgraphs(CalculatorGraphConfig* config,
if (subgraph_nodes_start == nodes->end()) break;
std::vector<CalculatorGraphConfig> subgraphs;
for (auto it = subgraph_nodes_start; it != nodes->end(); ++it) {
const auto& node = *it;
auto& node = *it;
int node_id = it - nodes->begin();
std::string node_name = CanonicalNodeName(*config, node_id);
MP_RETURN_IF_ERROR(ValidateSubgraphFields(node));
mediapipe/framework/tool/type_util.h
+114 -64
View File
@@ -16,79 +16,129 @@
#define MEDIAPIPE_FRAMEWORK_TOOL_TYPE_UTIL_H_
#include <cstddef>
#include <string>
#include <typeinfo>
#include <utility>
#include "absl/base/attributes.h"
#include "mediapipe/framework/demangle.h"
#include "mediapipe/framework/port.h"
namespace mediapipe {
// An identifier for a type. This class is lightweight and is meant to be passed
// by value.
// To get the TypeId for SomeType, write kTypeId<SomeType>.
class TypeId {
public:
size_t hash_code() const { return impl_.hash_code(); }
std::string name() const { return impl_.name(); }
bool operator==(const TypeId& other) const { return impl_ == other.impl_; }
bool operator<(const TypeId& other) const { return impl_ < other.impl_; }
template <typename H>
friend H AbslHashValue(H h, const TypeId& r) {
return H::combine(std::move(h), r.hash_code());
}
template <class T>
static constexpr inline TypeId Of() {
return TypeId{Impl::Get<T>()};
}
private:
// This implementation uses no RTTI. It distinguishes types, but does not
// know their names.
// TODO: record compile-time type string for (some or all) types.
template <class T>
struct TypeTag {
static constexpr char dummy = 0;
};
struct NoRttiImpl {
template <class T>
static constexpr inline NoRttiImpl Get() {
return {&TypeTag<T>::dummy};
}
size_t hash_code() const { return reinterpret_cast<uintptr_t>(tag_); }
std::string name() const { return "<type name missing>"; }
bool operator==(const NoRttiImpl& other) const {
return tag_ == other.tag_;
}
bool operator<(const NoRttiImpl& other) const { return tag_ < other.tag_; }
const void* tag_;
};
#if MEDIAPIPE_HAS_RTTI
template <class T>
static const std::type_info& GetTypeInfo() {
return typeid(T);
}
// This implementation uses RTTI, and delegates all operations to
// std::type_info. In order to support constexpr construction, we don't store
// a type_info directly (which is not constexpr), but a pointer to a function
// returning it (which is). This implementation is a bit slower than the
// others. The only potential advantage would be the ability to match types
// across multiple dynamic libraries, but we don't support that setup anyway.
// This is provided for completeness.
struct FullRttiImpl {
template <class T>
static constexpr inline FullRttiImpl Get() {
return {GetTypeInfo<T>};
}
size_t hash_code() const { return get_().hash_code(); }
std::string name() const { return Demangle(get_().name()); }
bool operator==(const FullRttiImpl& other) const {
return get_ == other.get_ || get_() == other.get_();
}
bool operator<(const FullRttiImpl& other) const {
return get_().before(other.get_());
}
decltype(&GetTypeInfo<void>) get_;
};
// This implementation also stores a pointer to a std::type_info getter
// function, but it only invokes it to get the type's name. It's equivalent to
// NoRttiImpl for most operations, but it allows getting the type's name.
struct FastRttiImpl {
template <class T>
static constexpr inline FastRttiImpl Get() {
return {GetTypeInfo<T>};
}
size_t hash_code() const { return reinterpret_cast<uintptr_t>(get_); }
std::string name() const { return Demangle(get_().name()); }
bool operator==(const FastRttiImpl& other) const {
return get_ == other.get_;
}
bool operator<(const FastRttiImpl& other) const {
return reinterpret_cast<uintptr_t>(get_) <
reinterpret_cast<uintptr_t>(other.get_);
}
decltype(&GetTypeInfo<void>) get_;
};
using Impl = FastRttiImpl;
#else
using Impl = NoRttiImpl;
#endif // MEDIAPIPE_HAS_RTTI
constexpr explicit TypeId(Impl impl) : impl_(impl) {}
Impl impl_;
};
template <class T>
static constexpr TypeId kTypeId = TypeId::Of<T>();
namespace tool {
#if !MEDIAPIPE_HAS_RTTI
// A unique identifier for type T.
class TypeInfo {
public:
size_t hash_code() const { return reinterpret_cast<size_t>(this); }
bool operator==(const TypeInfo& other) const { return &other == this; }
bool operator<(const TypeInfo& other) const { return &other < this; }
const char* name() const { return "<unknown>"; }
template <typename T>
static const TypeInfo& Get() {
static TypeInfo* static_type_info = new TypeInfo;
return *static_type_info;
}
private:
TypeInfo() {}
TypeInfo(const TypeInfo&) = delete;
};
#else // MEDIAPIPE_HAS_RTTI
// The std unique identifier for type T.
class TypeInfo {
public:
size_t hash_code() const { return info_.hash_code(); }
bool operator==(const TypeInfo& o) const { return info_ == o.info_; }
bool operator<(const TypeInfo& o) const { return info_.before(o.info_); }
const char* name() const { return info_.name(); }
template <typename T>
static const TypeInfo& Get() {
static TypeInfo* static_type_info = new TypeInfo(typeid(T));
return *static_type_info;
}
private:
TypeInfo(const std::type_info& info) : info_(info) {}
TypeInfo(const TypeInfo&) = delete;
private:
const std::type_info& info_;
friend class TypeIndex;
};
#endif
// An associative key for TypeInfo.
class TypeIndex {
public:
TypeIndex(const TypeInfo& info) : info_(info) {}
size_t hash_code() const { return info_.hash_code(); }
bool operator==(const TypeIndex& other) const { return info_ == other.info_; }
bool operator<(const TypeIndex& other) const { return info_ < other.info_; }
private:
const TypeInfo& info_;
};
// Helper method that returns a hash code of the given type. This allows for
// typeid testing across multiple binaries, unlike FastTypeId which used a
// memory location that only works within the same binary. Moreover, we use this
// for supporting multiple .so binaries in a single Android app built using the
// same compiler and C++ libraries.
// Note that std::type_info may still generate the same hash code for different
// types, although the c++ standard recommends that implementations avoid this
// as much as possible.
// Helper method that returns a hash code of the given type.
// Superseded by TypeId.
template <typename T>
ABSL_DEPRECATED("Use TypeId directly instead.")
size_t GetTypeHash() {
return TypeInfo::Get<T>().hash_code();
return kTypeId<T>.hash_code();
}
} // namespace tool
mediapipe/framework/type_map.h
+7 -9
View File
@@ -361,32 +361,30 @@ DEFINE_MEDIAPIPE_TYPE_MAP(PacketTypeStringToMediaPipeTypeData, std::string);
// End define MEDIAPIPE_REGISTER_TYPE_WITH_PROXY.
// Helper functions's to retrieve registration data.
inline const std::string* MediaPipeTypeStringFromTypeId(const size_t type_id) {
inline const std::string* MediaPipeTypeStringFromTypeId(TypeId type_id) {
const MediaPipeTypeData* value =
PacketTypeIdToMediaPipeTypeData::GetValue(type_id);
PacketTypeIdToMediaPipeTypeData::GetValue(type_id.hash_code());
return (value) ? &value->type_string : nullptr;
}
// Returns string identifier of type or NULL if not registered.
template <typename T>
inline const std::string* MediaPipeTypeString() {
return MediaPipeTypeStringFromTypeId(tool::GetTypeHash<T>());
return MediaPipeTypeStringFromTypeId(kTypeId<T>);
}
inline std::string MediaPipeTypeStringOrDemangled(
const tool::TypeInfo& type_info) {
const std::string* type_string =
MediaPipeTypeStringFromTypeId(type_info.hash_code());
inline std::string MediaPipeTypeStringOrDemangled(TypeId type_id) {
const std::string* type_string = MediaPipeTypeStringFromTypeId(type_id);
if (type_string) {
return *type_string;
} else {
return mediapipe::Demangle(type_info.name());
return type_id.name();
}
}
template <typename T>
std::string MediaPipeTypeStringOrDemangled() {
return MediaPipeTypeStringOrDemangled(tool::TypeInfo::Get<T>());
return MediaPipeTypeStringOrDemangled(kTypeId<T>);
}
// Returns type hash id of type identified by type_string or NULL if not
mediapipe/framework/validated_graph_config.cc
+43 -36
View File
@@ -14,6 +14,8 @@
#include "mediapipe/framework/validated_graph_config.h"
#include <memory>
#include "absl/container/flat_hash_set.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
@@ -140,35 +142,6 @@ absl::Status AddPredefinedExecutorConfigs(CalculatorGraphConfig* graph_config) {
return absl::OkStatus();
}
absl::Status PerformBasicTransforms(
const CalculatorGraphConfig& input_graph_config,
const GraphRegistry* graph_registry,
const Subgraph::SubgraphOptions* graph_options,
const GraphServiceManager* service_manager,
CalculatorGraphConfig* output_graph_config) {
*output_graph_config = input_graph_config;
MP_RETURN_IF_ERROR(tool::ExpandSubgraphs(output_graph_config, graph_registry,
graph_options, service_manager));
MP_RETURN_IF_ERROR(AddPredefinedExecutorConfigs(output_graph_config));
// Populate each node with the graph level input stream handler if a
// stream handler wasn't explicitly provided.
// TODO Instead of pre-populating, handle the graph level
// default appropriately within CalculatorGraph.
if (output_graph_config->has_input_stream_handler()) {
const auto& graph_level_input_stream_handler =
output_graph_config->input_stream_handler();
for (auto& node : *output_graph_config->mutable_node()) {
if (!node.has_input_stream_handler()) {
*node.mutable_input_stream_handler() = graph_level_input_stream_handler;
}
}
}
return absl::OkStatus();
}
} // namespace
// static
@@ -346,8 +319,7 @@ absl::Status NodeTypeInfo::Initialize(
}
absl::Status ValidatedGraphConfig::Initialize(
const CalculatorGraphConfig& input_config,
const GraphRegistry* graph_registry,
CalculatorGraphConfig input_config, const GraphRegistry* graph_registry,
const Subgraph::SubgraphOptions* graph_options,
const GraphServiceManager* service_manager) {
RET_CHECK(!initialized_)
@@ -358,9 +330,9 @@ absl::Status ValidatedGraphConfig::Initialize(
<< input_config.DebugString();
#endif
MP_RETURN_IF_ERROR(PerformBasicTransforms(
input_config, graph_registry, graph_options, service_manager, &config_));
config_ = std::move(input_config);
MP_RETURN_IF_ERROR(
PerformBasicTransforms(graph_registry, graph_options, service_manager));
// Initialize the basic node information.
MP_RETURN_IF_ERROR(InitializeGeneratorInfo());
MP_RETURN_IF_ERROR(InitializeCalculatorInfo());
@@ -441,7 +413,12 @@ absl::Status ValidatedGraphConfig::Initialize(
const GraphServiceManager* service_manager) {
graph_registry =
graph_registry ? graph_registry : &GraphRegistry::global_graph_registry;
SubgraphContext subgraph_context(graph_options, service_manager);
Subgraph::SubgraphOptions local_graph_options;
if (graph_options) {
local_graph_options = *graph_options;
}
SubgraphContext subgraph_context =
SubgraphContext(&local_graph_options, service_manager);
auto status_or_config =
graph_registry->CreateByName("", graph_type, &subgraph_context);
MP_RETURN_IF_ERROR(status_or_config.status());
@@ -466,6 +443,32 @@ absl::Status ValidatedGraphConfig::Initialize(
service_manager);
}
absl::Status ValidatedGraphConfig::PerformBasicTransforms(
const GraphRegistry* graph_registry,
const Subgraph::SubgraphOptions* graph_options,
const GraphServiceManager* service_manager) {
MP_RETURN_IF_ERROR(tool::ExpandSubgraphs(&config_, graph_registry,
graph_options, service_manager));
MP_RETURN_IF_ERROR(AddPredefinedExecutorConfigs(&config_));
// Populate each node with the graph level input stream handler if a
// stream handler wasn't explicitly provided.
// TODO Instead of pre-populating, handle the graph level
// default appropriately within CalculatorGraph.
if (config_.has_input_stream_handler()) {
const auto& graph_level_input_stream_handler =
config_.input_stream_handler();
for (auto& node : *config_.mutable_node()) {
if (!node.has_input_stream_handler()) {
*node.mutable_input_stream_handler() = graph_level_input_stream_handler;
}
}
}
return absl::OkStatus();
}
absl::Status ValidatedGraphConfig::InitializeCalculatorInfo() {
std::vector<absl::Status> statuses;
calculators_.reserve(config_.node_size());
@@ -690,6 +693,7 @@ absl::Status ValidatedGraphConfig::AddInputStreamsForNode(
if (!need_sorting_ptr) {
LOG(WARNING) << "Input Stream \"" << name
<< "\" for node with sorted index " << node_index
<< " name " << node_type_info->Contract().GetNodeName()
<< " is marked as a back edge, but its output stream is "
"already available. This means it was not necessary "
"to mark it as a back edge.";
@@ -701,6 +705,7 @@ absl::Status ValidatedGraphConfig::AddInputStreamsForNode(
if (edge_info.back_edge) {
VLOG(1) << "Encountered expected behavior: the back edge \"" << name
<< "\" for node with (possibly sorted) index " << node_index
<< " name " << node_type_info->Contract().GetNodeName()
<< " has an output stream which we have not yet seen.";
} else if (need_sorting_ptr) {
*need_sorting_ptr = true;
@@ -709,7 +714,9 @@ absl::Status ValidatedGraphConfig::AddInputStreamsForNode(
} else {
return mediapipe::UnknownErrorBuilder(MEDIAPIPE_LOC)
<< "Input Stream \"" << name << "\" for node with sorted index "
<< node_index << " does not have a corresponding output stream.";
<< node_index << " name "
<< node_type_info->Contract().GetNodeName()
<< " does not have a corresponding output stream.";
}
}
mediapipe/framework/validated_graph_config.h
+8 -1
View File
@@ -195,7 +195,7 @@ class ValidatedGraphConfig {
// before any other functions. Subgraphs are specified through the
// global graph registry or an optional local graph registry.
absl::Status Initialize(
const CalculatorGraphConfig& input_config,
CalculatorGraphConfig input_config,
const GraphRegistry* graph_registry = nullptr,
const Subgraph::SubgraphOptions* graph_options = nullptr,
const GraphServiceManager* service_manager = nullptr);
@@ -302,6 +302,13 @@ class ValidatedGraphConfig {
}
private:
// Perform transforms such as converting legacy features, expanding
// subgraphs, and popluting input stream handler.
absl::Status PerformBasicTransforms(
const GraphRegistry* graph_registry,
const Subgraph::SubgraphOptions* graph_options,
const GraphServiceManager* service_manager);
// Initialize the PacketGenerator information.
absl::Status InitializeGeneratorInfo();
// Initialize the Calculator information.