Project import generated by Copybara.

GitOrigin-RevId: 1e13be30e2c6838d4a2ff768a39c414bc80534bb
This commit is contained in:
MediaPipe Team
2022-09-06 21:46:17 +00:00
committed by Sebastian Schmidt
parent 63e679d99c
commit 4dc4b19ddb
639 changed files with 71327 additions and 2078 deletions
+3 -2
View File
@@ -102,13 +102,13 @@ mediapipe_proto_library(
mediapipe_proto_library(
name = "packet_factory_proto",
srcs = ["packet_factory.proto"],
visibility = [":mediapipe_internal"],
visibility = ["//visibility:public"],
)
mediapipe_proto_library(
name = "packet_generator_proto",
srcs = ["packet_generator.proto"],
visibility = [":mediapipe_internal"],
visibility = ["//visibility:public"],
)
mediapipe_proto_library(
@@ -153,6 +153,7 @@ mediapipe_proto_library(
deps = ["//mediapipe/framework:mediapipe_options_proto"],
)
# It is for pure-native Android builds where the library can't have any dependency on libandroid.so
config_setting(
name = "android_no_jni",
define_values = {"MEDIAPIPE_NO_JNI": "1"},
+3
View File
@@ -460,6 +460,9 @@ class OutputShardAccessBase {
OutputShardAccessBase(const CalculatorContext& cc, OutputStreamShard* output)
: context_(cc), output_(output) {}
Timestamp NextTimestampBound() const {
return (output_) ? output_->NextTimestampBound() : Timestamp::Unset();
}
void SetNextTimestampBound(Timestamp timestamp) {
if (output_) output_->SetNextTimestampBound(timestamp);
}
@@ -20,7 +20,6 @@ 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.
//
+4 -1
View File
@@ -41,7 +41,10 @@ bzl_library(
proto_library(
name = "proto_descriptor_proto",
srcs = ["proto_descriptor.proto"],
visibility = ["//mediapipe/framework:__subpackages__"],
visibility = [
"//mediapipe/deps:__subpackages__",
"//mediapipe/framework:__subpackages__",
],
)
mediapipe_cc_proto_library(
@@ -2,6 +2,9 @@ syntax = "proto2";
package mediapipe;
option java_package = "com.google.mediapipe.proto";
option java_outer_classname = "FieldDescriptorProtoProto";
// Describes a field within a message.
message FieldDescriptorProto {
enum Type {
+6 -6
View File
@@ -129,8 +129,8 @@ namespace mediapipe {
// }));
namespace registration_internal {
constexpr char kCxxSep[] = "::";
constexpr char kNameSep[] = ".";
inline constexpr char kCxxSep[] = "::";
inline constexpr char kNameSep[] = ".";
template <typename T>
struct WrapStatusOr {
@@ -245,7 +245,7 @@ class FunctionRegistry {
// The name must be either unqualified or fully qualified with a leading "::".
// The leading "::" in a fully qualified name is stripped.
std::string GetNormalizedName(const std::string& name) {
constexpr auto kCxxSep = registration_internal::kCxxSep;
using ::mediapipe::registration_internal::kCxxSep;
std::vector<std::string> names = absl::StrSplit(name, kCxxSep);
if (names[0].empty()) {
names.erase(names.begin());
@@ -261,8 +261,8 @@ class FunctionRegistry {
// Namespaces are separated by kNameSep.
std::string GetQualifiedName(const std::string& ns,
const std::string& name) const {
constexpr auto kCxxSep = registration_internal::kCxxSep;
constexpr auto kNameSep = registration_internal::kNameSep;
using ::mediapipe::registration_internal::kCxxSep;
using ::mediapipe::registration_internal::kNameSep;
std::vector<std::string> names = absl::StrSplit(name, kNameSep);
if (names[0].empty()) {
names.erase(names.begin());
@@ -291,7 +291,7 @@ class FunctionRegistry {
// For names included in NamespaceAllowlist, strips the namespace.
std::string GetAdjustedName(const std::string& name) {
constexpr auto kCxxSep = registration_internal::kCxxSep;
using ::mediapipe::registration_internal::kCxxSep;
std::vector<std::string> names = absl::StrSplit(name, kCxxSep);
std::string base_name = names.back();
names.pop_back();
+130
View File
@@ -135,6 +135,123 @@ IsOkAndHoldsMatcher<typename std::decay<InnerMatcher>::type> IsOkAndHolds(
// Returns a gMock matcher that matches a Status or StatusOr<> which is OK.
inline IsOkMatcher IsOk() { return IsOkMatcher(); }
////////////////////////////////////////////////////////////
// Implementation of StatusIs().
//
// StatusIs() is a polymorphic matcher. This class is the common
// implementation of it shared by all types T where StatusIs() can be used as
// a Matcher<T>.
class StatusIsMatcherCommonImpl {
public:
StatusIsMatcherCommonImpl(
::testing::Matcher<const absl::StatusCode> code_matcher,
::testing::Matcher<const std::string&> message_matcher)
: code_matcher_(std::move(code_matcher)),
message_matcher_(std::move(message_matcher)) {}
void DescribeTo(std::ostream* os) const {
*os << "has a status code that ";
code_matcher_.DescribeTo(os);
*os << ", and has an error message that ";
message_matcher_.DescribeTo(os);
}
void DescribeNegationTo(std::ostream* os) const {
*os << "has a status code that ";
code_matcher_.DescribeNegationTo(os);
*os << ", or has an error message that ";
message_matcher_.DescribeNegationTo(os);
}
bool MatchAndExplain(const absl::Status& status,
::testing::MatchResultListener* result_listener) const {
::testing::StringMatchResultListener inner_listener;
inner_listener.Clear();
if (!code_matcher_.MatchAndExplain(status.code(), &inner_listener)) {
*result_listener << (inner_listener.str().empty()
? "whose status code is wrong"
: "which has a status code " +
inner_listener.str());
return false;
}
if (!message_matcher_.Matches(std::string(status.message()))) {
*result_listener << "whose error message is wrong";
return false;
}
return true;
}
private:
const ::testing::Matcher<const absl::StatusCode> code_matcher_;
const ::testing::Matcher<const std::string&> message_matcher_;
};
// Monomorphic implementation of matcher StatusIs() for a given type T. T can
// be Status, StatusOr<>, or a reference to either of them.
template <typename T>
class MonoStatusIsMatcherImpl : public ::testing::MatcherInterface<T> {
public:
explicit MonoStatusIsMatcherImpl(StatusIsMatcherCommonImpl common_impl)
: common_impl_(std::move(common_impl)) {}
void DescribeTo(std::ostream* os) const override {
common_impl_.DescribeTo(os);
}
void DescribeNegationTo(std::ostream* os) const override {
common_impl_.DescribeNegationTo(os);
}
bool MatchAndExplain(
T actual_value,
::testing::MatchResultListener* result_listener) const override {
return common_impl_.MatchAndExplain(GetStatus(actual_value),
result_listener);
}
private:
StatusIsMatcherCommonImpl common_impl_;
};
// Implements StatusIs() as a polymorphic matcher.
class StatusIsMatcher {
public:
StatusIsMatcher(::testing::Matcher<const absl::StatusCode> code_matcher,
::testing::Matcher<const std::string&> message_matcher)
: common_impl_(
::testing::MatcherCast<const absl::StatusCode>(code_matcher),
::testing::MatcherCast<const std::string&>(message_matcher)) {}
// Converts this polymorphic matcher to a monomorphic matcher of the given
// type. T can be StatusOr<>, Status, or a reference to either of them.
template <typename T>
operator ::testing::Matcher<T>() const { // NOLINT
return ::testing::MakeMatcher(new MonoStatusIsMatcherImpl<T>(common_impl_));
}
private:
const StatusIsMatcherCommonImpl common_impl_;
};
// Returns a matcher that matches a Status or StatusOr<> whose status code
// matches code_matcher, and whose error message matches message_matcher.
template <typename CodeMatcher, typename MessageMatcher>
StatusIsMatcher StatusIs(CodeMatcher code_matcher,
MessageMatcher message_matcher) {
return StatusIsMatcher(std::move(code_matcher), std::move(message_matcher));
}
// Returns a matcher that matches a Status or StatusOr<> whose status code
// matches code_matcher.
template <typename CodeMatcher>
StatusIsMatcher StatusIs(CodeMatcher code_matcher) {
return StatusIs(std::move(code_matcher), ::testing::_);
}
} // namespace mediapipe
// Macros for testing the results of functions that return absl::Status or
@@ -142,4 +259,17 @@ inline IsOkMatcher IsOk() { return IsOkMatcher(); }
#define MP_EXPECT_OK(expression) EXPECT_THAT(expression, mediapipe::IsOk())
#define MP_ASSERT_OK(expression) ASSERT_THAT(expression, mediapipe::IsOk())
#define STATUS_MACROS_IMPL_CONCAT_INNER_(x, y) x##y
#define STATUS_MACROS_IMPL_CONCAT_(x, y) STATUS_MACROS_IMPL_CONCAT_INNER_(x, y)
#undef MP_ASSERT_OK_AND_ASSIGN
#define MP_ASSERT_OK_AND_ASSIGN(lhs, rexpr) \
MP_ASSERT_OK_AND_ASSIGN_IMPL_( \
STATUS_MACROS_IMPL_CONCAT_(_status_or_value, __LINE__), lhs, rexpr)
#define MP_ASSERT_OK_AND_ASSIGN_IMPL_(statusor, lhs, rexpr) \
auto statusor = (rexpr); \
ASSERT_TRUE(statusor.ok()); \
lhs = std::move(statusor.value())
#endif // MEDIAPIPE_DEPS_STATUS_MATCHERS_H_
+23 -11
View File
@@ -206,12 +206,6 @@ cc_library(
name = "location",
srcs = ["location.cc"],
hdrs = ["location.h"],
defines = select({
"//conditions:default": [],
"//mediapipe:android": ["MEDIAPIPE_ANDROID_OPENCV"],
":portable_opencv": ["MEDIAPIPE_ANDROID_OPENCV"],
":opencv": [],
}),
visibility = ["//visibility:public"],
deps = [
"@com_google_protobuf//:protobuf",
@@ -232,11 +226,6 @@ cc_library(
"//mediapipe/framework/port:statusor",
"//mediapipe/framework/formats/annotation:rasterization_cc_proto",
] + select({
"//conditions:default": [
"//mediapipe/framework/port:opencv_imgproc",
],
"//mediapipe/framework/port:disable_opencv": [],
}) + select({
"//conditions:default": [
],
"//mediapipe:android": [],
@@ -245,6 +234,28 @@ cc_library(
alwayslink = 1,
)
cc_library(
name = "location_opencv",
srcs = ["location_opencv.cc"],
hdrs = ["location_opencv.h"],
visibility = ["//visibility:public"],
deps = [
":location",
"//mediapipe/framework/port:opencv_imgproc",
],
alwayslink = 1,
)
cc_test(
name = "location_opencv_test",
srcs = ["location_opencv_test.cc"],
deps = [
":location_opencv",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:rectangle",
],
)
cc_library(
name = "video_stream_header",
hdrs = ["video_stream_header.h"],
@@ -464,6 +475,7 @@ cc_library(
"-framework MetalKit",
],
"//conditions:default": [],
"//mediapipe/framework:android_no_jni": [],
"//mediapipe:android": [
"-landroid",
],
@@ -16,8 +16,6 @@ syntax = "proto2";
package mediapipe;
option objc_class_prefix = "MediaPipe";
// Proto for serializing Vector2 data
message Vector2Data {
optional float x = 1;
@@ -18,6 +18,8 @@ package mediapipe;
import "mediapipe/framework/formats/annotation/rasterization.proto";
option cc_enable_arenas = true;
// A way to identify a part of an image. A locus does not need to correspond to
// a subset of pixels -- e.g. for a local descriptor we might define a locus in
// terms of its location and scale, even if the support of the descriptor is the
@@ -20,7 +20,6 @@ syntax = "proto2";
package mediapipe;
option objc_class_prefix = "MediaPipe";
option java_package = "com.google.mediapipe.formats.proto";
option java_outer_classname = "ClassificationProto";
+2
View File
@@ -42,5 +42,7 @@ bool Image::ConvertToGpu() const {
MEDIAPIPE_REGISTER_TYPE(mediapipe::Image, "::mediapipe::Image", nullptr,
nullptr);
MEDIAPIPE_REGISTER_TYPE(std::vector<mediapipe::Image>,
"::std::vector<::mediapipe::Image>", nullptr, nullptr);
} // namespace mediapipe
@@ -23,6 +23,9 @@ syntax = "proto2";
package mediapipe;
option java_package = "com.google.mediapipe.formats.proto";
option java_outer_classname = "ImageFormatProto";
message ImageFormat {
enum Format {
// The format is unknown. It is not valid for an ImageFrame to be
-186
View File
@@ -32,10 +32,6 @@
#include "mediapipe/framework/tool/status_util.h"
#include "mediapipe/framework/type_map.h"
#if LOCATION_OPENCV
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
#endif
namespace mediapipe {
namespace {
@@ -61,41 +57,6 @@ Rectangle_i MaskToRectangle(const LocationData& location_data) {
return Rectangle_i(xmin, ymin, xmax - xmin + 1, ymax - ymin + 1);
}
#if LOCATION_OPENCV
std::unique_ptr<cv::Mat> MaskToMat(const LocationData::BinaryMask& mask) {
auto image = absl::make_unique<cv::Mat>();
*image = cv::Mat::zeros(cv::Size(mask.width(), mask.height()), CV_32FC1);
for (const auto& interval : mask.rasterization().interval()) {
for (int x = interval.left_x(); x <= interval.right_x(); ++x) {
image->at<float>(interval.y(), x) = 1.0f;
}
}
return image;
}
absl::StatusOr<std::unique_ptr<cv::Mat>> RectangleToMat(
int image_width, int image_height, const Rectangle_i& rect) {
// These checks prevent undefined behavior caused when setting memory for
// rectangles whose edges lie outside image edges.
if (rect.ymin() < 0 || rect.xmin() < 0 || rect.xmax() > image_width ||
rect.ymax() > image_height) {
return absl::InvalidArgumentError(absl::Substitute(
"Rectangle must be bounded by image boundaries.\nImage Width: "
"$0\nImage Height: $1\nRectangle: [($2, $3), ($4, $5)]",
image_width, image_height, rect.xmin(), rect.ymin(), rect.xmax(),
rect.ymax()));
}
// Allocate image and set pixels of foreground mask.
auto image = absl::make_unique<cv::Mat>();
*image = cv::Mat::zeros(cv::Size(image_width, image_height), CV_32FC1);
for (int y = rect.ymin(); y < rect.ymax(); ++y) {
for (int x = rect.xmin(); x < rect.xmax(); ++x) {
image->at<float>(y, x) = 1.0f;
}
}
return std::move(image);
}
#endif // OPENCV
} // namespace
Location::Location() {}
@@ -134,12 +95,6 @@ Location Location::CreateBBoxLocation(const ::mediapipe::BoundingBox& bbox) {
bbox.lower_y() - bbox.upper_y());
}
#if LOCATION_OPENCV
Location Location::CreateBBoxLocation(const cv::Rect& rect) {
return CreateBBoxLocation(rect.x, rect.y, rect.width, rect.height);
}
#endif
Location Location::CreateRelativeBBoxLocation(float relative_xmin,
float relative_ymin,
float relative_width,
@@ -159,41 +114,6 @@ Location Location::CreateRelativeBBoxLocation(const Rectangle_f& rect) {
rect.Height());
}
#if LOCATION_OPENCV
template <typename T>
Location Location::CreateCvMaskLocation(const cv::Mat_<T>& mask) {
CHECK_EQ(1, mask.channels())
<< "The specified cv::Mat mask should be single-channel.";
LocationData location_data;
location_data.set_format(LocationData::MASK);
location_data.mutable_mask()->set_width(mask.cols);
location_data.mutable_mask()->set_height(mask.rows);
auto* rasterization = location_data.mutable_mask()->mutable_rasterization();
const auto kForegroundThreshold = static_cast<T>(0);
for (int y = 0; y < mask.rows; y++) {
Rasterization::Interval* interval;
bool traversing = false;
for (int x = 0; x < mask.cols; x++) {
const bool is_foreground =
mask.template at<T>(y, x) > kForegroundThreshold;
if (is_foreground) {
if (!traversing) {
interval = rasterization->add_interval();
interval->set_y(y);
interval->set_left_x(x);
traversing = true;
}
interval->set_right_x(x);
} else {
traversing = false;
}
}
}
return Location(location_data);
}
#endif
LocationData::Format Location::GetFormat() const {
return location_data_.format();
}
@@ -274,62 +194,6 @@ Location& Location::Scale(const float scale) {
return *this;
}
#if LOCATION_OPENCV
Location& Location::Enlarge(const float factor) {
CHECK_GT(factor, 0.0f);
if (factor == 1.0f) return *this;
switch (location_data_.format()) {
case LocationData::GLOBAL: {
// Do nothing.
break;
}
case LocationData::BOUNDING_BOX: {
auto* box = location_data_.mutable_bounding_box();
const int enlarged_int_width =
static_cast<int>(std::round(factor * box->width()));
const int enlarged_int_height =
static_cast<int>(std::round(factor * box->height()));
box->set_xmin(
std::max(box->xmin() + box->width() / 2 - enlarged_int_width / 2, 0));
box->set_ymin(std::max(
box->ymin() + box->height() / 2 - enlarged_int_height / 2, 0));
box->set_width(enlarged_int_width);
box->set_height(enlarged_int_height);
break;
}
case LocationData::RELATIVE_BOUNDING_BOX: {
auto* box = location_data_.mutable_relative_bounding_box();
box->set_xmin(box->xmin() - ((factor - 1.0) * box->width()) / 2.0);
box->set_ymin(box->ymin() - ((factor - 1.0) * box->height()) / 2.0);
box->set_width(factor * box->width());
box->set_height(factor * box->height());
break;
}
case LocationData::MASK: {
auto mask_bounding_box = MaskToRectangle(location_data_);
const float scaler = std::fabs(factor - 1.0f);
const int dilation_width =
static_cast<int>(std::round(scaler * mask_bounding_box.Width()));
const int dilation_height =
static_cast<int>(std::round(scaler * mask_bounding_box.Height()));
if (dilation_width == 0 || dilation_height == 0) break;
cv::Mat morph_element(dilation_height, dilation_width, CV_8U,
cv::Scalar(1));
auto mask = GetCvMask();
if (factor > 1.0f) {
cv::dilate(*mask, *mask, morph_element);
} else {
cv::erode(*mask, *mask, morph_element);
}
Location::CreateCvMaskLocation<uint8>(*mask).ConvertToProto(
&location_data_);
break;
}
}
return *this;
}
#endif
Location& Location::Square(int image_width, int image_height) {
switch (location_data_.format()) {
case LocationData::GLOBAL: {
@@ -615,51 +479,6 @@ template <>
return bounding_box;
}
#if LOCATION_OPENCV
std::unique_ptr<cv::Mat> Location::GetCvMask() const {
CHECK_EQ(LocationData::MASK, location_data_.format());
const auto& mask = location_data_.mask();
std::unique_ptr<cv::Mat> mat(
new cv::Mat(mask.height(), mask.width(), CV_8UC1, cv::Scalar(0)));
for (const auto& interval :
location_data_.mask().rasterization().interval()) {
for (int x = interval.left_x(); x <= interval.right_x(); ++x) {
mat->at<uint8>(interval.y(), x) = 255;
}
}
return mat;
}
std::unique_ptr<cv::Mat> Location::ConvertToCvMask(int image_width,
int image_height) const {
switch (location_data_.format()) {
case LocationData::GLOBAL:
case LocationData::BOUNDING_BOX:
case LocationData::RELATIVE_BOUNDING_BOX: {
auto status_or_mat =
RectangleToMat(image_width, image_height,
ConvertToBBox<Rectangle_i>(image_width, image_height));
if (!status_or_mat.ok()) {
LOG(ERROR) << status_or_mat.status().message();
return nullptr;
}
return std::move(status_or_mat).value();
}
case LocationData::MASK: {
return MaskToMat(location_data_.mask());
}
}
// This should never happen; a new LocationData::Format enum was introduced
// without updating this function's switch(...) to support it.
#if !defined(MEDIAPIPE_MOBILE) && !defined(MEDIAPIPE_LITE)
LOG(ERROR) << "Location's LocationData has format not supported by "
"Location::ConvertToMask: "
<< location_data_.DebugString();
#endif
return nullptr;
}
#endif
std::vector<Point2_f> Location::GetRelativeKeypoints() const {
CHECK_EQ(LocationData::RELATIVE_BOUNDING_BOX, location_data_.format());
std::vector<Point2_f> keypoints;
@@ -703,9 +522,4 @@ LocationData Location::ConvertToProto() const {
return location_data;
}
#if LOCATION_OPENCV
template Location Location::CreateCvMaskLocation(const cv::Mat_<uint8>& mask);
template Location Location::CreateCvMaskLocation(const cv::Mat_<float>& mask);
#endif // LOCATION_OPENCV
} // namespace mediapipe
+1 -44
View File
@@ -30,21 +30,6 @@
#include "mediapipe/framework/port/point2.h"
#include "mediapipe/framework/port/rectangle.h"
// clang-format off
#if !defined(LOCATION_OPENCV)
# if !MEDIAPIPE_DISABLE_OPENCV && \
(!defined(MEDIAPIPE_MOBILE) || defined(MEDIAPIPE_ANDROID_OPENCV))
# define LOCATION_OPENCV 1
# else
# define LOCATION_OPENCV 0
# endif
#endif
#if LOCATION_OPENCV
#include "mediapipe/framework/port/opencv_core_inc.h"
#endif
// clang-format on
namespace mediapipe {
class BoundingBox;
} // namespace mediapipe
@@ -68,9 +53,6 @@ class Location {
// formats.
static Location CreateBBoxLocation(const Rectangle_i& rect);
static Location CreateBBoxLocation(const ::mediapipe::BoundingBox& bbox);
#if LOCATION_OPENCV
static Location CreateBBoxLocation(const cv::Rect& rect);
#endif
// Creates a location of type RELATIVE_BOUNDING_BOX, i.e. it is based on a
// bounding box defined by its upper left corner (xmin, ymin) and its width
// and height, all relative to the image dimensions.
@@ -81,14 +63,6 @@ class Location {
// Creates a location of type RELATIVE_BOUNDING_BOX from bounding boxes in
// various formats.
static Location CreateRelativeBBoxLocation(const Rectangle_f& relative_rect);
#if LOCATION_OPENCV
// Creates a location of type MASK from a single-channel uint8 or float
// cv::Mat_ (type is CV_8UC1 or CV_32FC1). Check fails if the mat is not
// single channel . All pixel with positive values are considered foreground,
// the rest background.
template <typename T>
static Location CreateCvMaskLocation(const cv::Mat_<T>& mask);
#endif
// Returns the location type describing the type of data it contains. This
// type is set at creation time based on the one of the above factory methods.
@@ -105,14 +79,6 @@ class Location {
// NOTE: it does not handle masks.
Location& Scale(float scale);
#if LOCATION_OPENCV
// Enlarges the location by the given factor. This operation keeps the center
// of the location fixed, while enlarging its dimensions by the given factor.
// Note that the location may partially lie outside the image after this
// operation. OpenCV required for mask enlargement. Returns *this.
Location& Enlarge(float factor);
#endif
// Resizes the location such that it is the tighest square location containing
// centered the original location. It supports locations of type GLOBAL,
// BOUNDING_BOX and RELATIVE_BOUNDING_BOX, otherwise it CHECK-fails. The user
@@ -154,12 +120,7 @@ class Location {
T GetBBox() const;
// Accessor for location data type RELATIVE_BOUNDING_BOX.
Rectangle_f GetRelativeBBox() const;
#if LOCATION_OPENCV
// Same as GetMask() with the difference that the return value is a cv::Mat of
// type CV_8UC1. It contains value 0 for background pixels and value 255 for
// foreground ones.
std::unique_ptr<cv::Mat> GetCvMask() const;
#endif
// Accessor for relative_keypoints in location data. Relative keypoints are
// specified with x and y coordinates, where both x and y are relative to the
// image width and height, respectively, and are in the range [0, 1]. Fails if
@@ -181,10 +142,6 @@ class Location {
template <typename T>
T ConvertToBBox(int image_width, int image_height) const;
Rectangle_f ConvertToRelativeBBox(int image_width, int image_height) const;
#if LOCATION_OPENCV
std::unique_ptr<cv::Mat> ConvertToCvMask(int image_width,
int image_height) const;
#endif
// Returns keypoints in absolute pixel coordinates.
std::vector<Point2_i> ConvertToKeypoints(int image_width,
int image_height) const;
@@ -0,0 +1,220 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "mediapipe/framework/formats/location_opencv.h"
#include "absl/memory/memory.h"
#include "absl/strings/substitute.h"
#include "mediapipe/framework/formats/annotation/rasterization.pb.h"
#include "mediapipe/framework/formats/location.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
#include "mediapipe/framework/port/statusor.h"
namespace mediapipe {
namespace {
Rectangle_i MaskToRectangle(const LocationData& location_data) {
CHECK(location_data.mask().has_rasterization());
const auto& rasterization = location_data.mask().rasterization();
if (rasterization.interval_size() == 0) {
return Rectangle_i(0, 0, 0, 0);
}
int xmin = std::numeric_limits<int>::max();
int xmax = std::numeric_limits<int>::lowest();
int ymin = std::numeric_limits<int>::max();
int ymax = std::numeric_limits<int>::lowest();
for (const auto& interval : rasterization.interval()) {
xmin = std::min(xmin, interval.left_x());
xmax = std::max(xmax, interval.right_x());
ymin = std::min(ymin, interval.y());
ymax = std::max(ymax, interval.y());
}
return Rectangle_i(xmin, ymin, xmax - xmin + 1, ymax - ymin + 1);
}
std::unique_ptr<cv::Mat> MaskToMat(const LocationData::BinaryMask& mask) {
auto image = absl::make_unique<cv::Mat>();
*image = cv::Mat::zeros(cv::Size(mask.width(), mask.height()), CV_32FC1);
for (const auto& interval : mask.rasterization().interval()) {
for (int x = interval.left_x(); x <= interval.right_x(); ++x) {
image->at<float>(interval.y(), x) = 1.0f;
}
}
return image;
}
absl::StatusOr<std::unique_ptr<cv::Mat>> RectangleToMat(
int image_width, int image_height, const Rectangle_i& rect) {
// These checks prevent undefined behavior caused when setting memory for
// rectangles whose edges lie outside image edges.
if (rect.ymin() < 0 || rect.xmin() < 0 || rect.xmax() > image_width ||
rect.ymax() > image_height) {
return absl::InvalidArgumentError(absl::Substitute(
"Rectangle must be bounded by image boundaries.\nImage Width: "
"$0\nImage Height: $1\nRectangle: [($2, $3), ($4, $5)]",
image_width, image_height, rect.xmin(), rect.ymin(), rect.xmax(),
rect.ymax()));
}
// Allocate image and set pixels of foreground mask.
auto image = absl::make_unique<cv::Mat>();
*image = cv::Mat::zeros(cv::Size(image_width, image_height), CV_32FC1);
for (int y = rect.ymin(); y < rect.ymax(); ++y) {
for (int x = rect.xmin(); x < rect.xmax(); ++x) {
image->at<float>(y, x) = 1.0f;
}
}
return std::move(image);
}
} // namespace
Location CreateBBoxLocation(const cv::Rect& rect) {
return Location::CreateBBoxLocation(rect.x, rect.y, rect.width, rect.height);
}
std::unique_ptr<cv::Mat> GetCvMask(const Location& location) {
const auto location_data = location.ConvertToProto();
CHECK_EQ(LocationData::MASK, location_data.format());
const auto& mask = location_data.mask();
std::unique_ptr<cv::Mat> mat(
new cv::Mat(mask.height(), mask.width(), CV_8UC1, cv::Scalar(0)));
for (const auto& interval : location_data.mask().rasterization().interval()) {
for (int x = interval.left_x(); x <= interval.right_x(); ++x) {
mat->at<uint8>(interval.y(), x) = 255;
}
}
return mat;
}
std::unique_ptr<cv::Mat> ConvertToCvMask(const Location& location,
int image_width, int image_height) {
const auto location_data = location.ConvertToProto();
switch (location_data.format()) {
case LocationData::GLOBAL:
case LocationData::BOUNDING_BOX:
case LocationData::RELATIVE_BOUNDING_BOX: {
auto status_or_mat = RectangleToMat(
image_width, image_height,
location.ConvertToBBox<Rectangle_i>(image_width, image_height));
if (!status_or_mat.ok()) {
LOG(ERROR) << status_or_mat.status().message();
return nullptr;
}
return std::move(status_or_mat).value();
}
case LocationData::MASK: {
return MaskToMat(location_data.mask());
}
}
// This should never happen; a new LocationData::Format enum was introduced
// without updating this function's switch(...) to support it.
#if !defined(MEDIAPIPE_MOBILE) && !defined(MEDIAPIPE_LITE)
LOG(ERROR) << "Location's LocationData has format not supported by "
"Location::ConvertToMask: "
<< location_data.DebugString();
#endif
return nullptr;
}
void EnlargeLocation(Location& location, const float factor) {
CHECK_GT(factor, 0.0f);
if (factor == 1.0f) return;
auto location_data = location.ConvertToProto();
switch (location_data.format()) {
case LocationData::GLOBAL: {
// Do nothing.
break;
}
case LocationData::BOUNDING_BOX: {
auto* box = location_data.mutable_bounding_box();
const int enlarged_int_width =
static_cast<int>(std::round(factor * box->width()));
const int enlarged_int_height =
static_cast<int>(std::round(factor * box->height()));
box->set_xmin(
std::max(box->xmin() + box->width() / 2 - enlarged_int_width / 2, 0));
box->set_ymin(std::max(
box->ymin() + box->height() / 2 - enlarged_int_height / 2, 0));
box->set_width(enlarged_int_width);
box->set_height(enlarged_int_height);
break;
}
case LocationData::RELATIVE_BOUNDING_BOX: {
auto* box = location_data.mutable_relative_bounding_box();
box->set_xmin(box->xmin() - ((factor - 1.0) * box->width()) / 2.0);
box->set_ymin(box->ymin() - ((factor - 1.0) * box->height()) / 2.0);
box->set_width(factor * box->width());
box->set_height(factor * box->height());
break;
}
case LocationData::MASK: {
auto mask_bounding_box = MaskToRectangle(location_data);
const float scaler = std::fabs(factor - 1.0f);
const int dilation_width =
static_cast<int>(std::round(scaler * mask_bounding_box.Width()));
const int dilation_height =
static_cast<int>(std::round(scaler * mask_bounding_box.Height()));
if (dilation_width == 0 || dilation_height == 0) break;
cv::Mat morph_element(dilation_height, dilation_width, CV_8U,
cv::Scalar(1));
auto mask = GetCvMask(location);
if (factor > 1.0f) {
cv::dilate(*mask, *mask, morph_element);
} else {
cv::erode(*mask, *mask, morph_element);
}
CreateCvMaskLocation<uint8>(*mask).ConvertToProto(&location_data);
break;
}
}
location.SetFromProto(location_data);
}
template <typename T>
Location CreateCvMaskLocation(const cv::Mat_<T>& mask) {
CHECK_EQ(1, mask.channels())
<< "The specified cv::Mat mask should be single-channel.";
LocationData location_data;
location_data.set_format(LocationData::MASK);
location_data.mutable_mask()->set_width(mask.cols);
location_data.mutable_mask()->set_height(mask.rows);
auto* rasterization = location_data.mutable_mask()->mutable_rasterization();
const auto kForegroundThreshold = static_cast<T>(0);
for (int y = 0; y < mask.rows; y++) {
Rasterization::Interval* interval;
bool traversing = false;
for (int x = 0; x < mask.cols; x++) {
const bool is_foreground =
mask.template at<T>(y, x) > kForegroundThreshold;
if (is_foreground) {
if (!traversing) {
interval = rasterization->add_interval();
interval->set_y(y);
interval->set_left_x(x);
traversing = true;
}
interval->set_right_x(x);
} else {
traversing = false;
}
}
}
return Location(location_data);
}
template Location CreateCvMaskLocation(const cv::Mat_<uint8>& mask);
template Location CreateCvMaskLocation(const cv::Mat_<float>& mask);
} // namespace mediapipe
@@ -0,0 +1,54 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// A collection of functions operating on MediaPipe::Location that require
// OpenCV to either convert between formats, or apply OpenCV transformations.
#ifndef MEDIAPIPE_FRAMEWORK_FORMATS_LOCATION_OPENCV_H_
#define MEDIAPIPE_FRAMEWORK_FORMATS_LOCATION_OPENCV_H_
#include "mediapipe/framework/formats/location.h"
#include "mediapipe/framework/port/opencv_core_inc.h"
namespace mediapipe {
// Creates a location of type BOUNDING_BOX from an OpenCV rectangle.
Location CreateBBoxLocation(const cv::Rect& rect);
// Creates a location of type MASK from a single-channel uint8 or float
// cv::Mat_ (type is CV_8UC1 or CV_32FC1). Check fails if the mat is not
// single channel. Pixels with positive values are treated as the foreground.
template <typename T>
Location CreateCvMaskLocation(const cv::Mat_<T>& mask);
// Enlarges the location by the given factor. This operation keeps the center
// of the location fixed, while enlarging its dimensions by the given factor.
// Note that the location may partially lie outside the image after this
// operation.
void EnlargeLocation(Location& location, float factor);
// Same as Location::GetMask() with the difference that the return value is a
// cv::Mat of type CV_8UC1. Background pixels are set to 0 and foreground pixels
// are set to 255.
std::unique_ptr<cv::Mat> GetCvMask(const Location& location);
// Returns the provided location's RELATIVE_BOUNDING_BOX or MASK location
// data as an OpenCV Mat. If the location data is in a format not directly
// convertible to the specified return type the following conversion principles
// are used:
// - Rectangle -> Mask: the rectangle is converted to a mask with all
// pixels inside the rectangle being foreground pixels.
std::unique_ptr<cv::Mat> ConvertToCvMask(const Location& location,
int image_width, int image_height);
} // namespace mediapipe
#endif // MEDIAPIPE_FRAMEWORK_FORMATS_LOCATION_OPENCV_H_
@@ -0,0 +1,167 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "mediapipe/framework/formats/location_opencv.h"
#include "mediapipe/framework/formats/annotation/rasterization.pb.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/rectangle.h"
namespace mediapipe {
// 7x3x1 test mask pattern containing the following region types: bordering left
// and right edges, multiple and single pixel lengths, multiple and single
// segments per row.
static const int kWidth = 7;
static const int kHeight = 3;
const std::vector<uint8> kTestPatternVector = {0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0,
0, 0, 0, 1, 0, 1, 0, 1, 0, 0};
// Interval {y, x_start, x_end} representation of kTestPatternVector.
const std::vector<std::vector<int>> kTestPatternIntervals = {
{0, 5, 6}, {1, 1, 2}, {2, 0, 0}, {2, 2, 2}, {2, 4, 4}};
static const float kEps = 0.0001f;
Location TestPatternIntervalsToMaskLocation() {
LocationData data;
data.set_format(LocationData::MASK);
data.mutable_mask()->set_width(kWidth);
data.mutable_mask()->set_height(kHeight);
for (const auto& test_interval : kTestPatternIntervals) {
auto interval =
data.mutable_mask()->mutable_rasterization()->add_interval();
interval->set_y(test_interval[0]);
interval->set_left_x(test_interval[1]);
interval->set_right_x(test_interval[2]);
}
return Location(data);
}
TEST(LocationOpencvTest, CreateBBoxLocation) {
const int x_start = 1;
const int y_start = 2;
const int width = 3;
const int height = 4;
const cv::Rect cv_rect(x_start, y_start, width, height);
Location location = CreateBBoxLocation(cv_rect);
auto rect = location.GetBBox<Rectangle_i>();
const std::vector<int> cv_rect_dims(
{cv_rect.x, cv_rect.y, cv_rect.width, cv_rect.height});
const std::vector<int> rect_dims(
{rect.xmin(), rect.ymin(), rect.Width(), rect.Height()});
EXPECT_EQ(cv_rect_dims, rect_dims);
}
TEST(LocationOpencvTest, CreateCvMaskLocation) {
cv::Mat_<uint8> test_mask(kHeight, kWidth,
const_cast<uint8*>(kTestPatternVector.data()));
Location location = CreateCvMaskLocation(test_mask);
auto intervals = location.ConvertToProto().mask().rasterization().interval();
EXPECT_EQ(intervals.size(), kTestPatternIntervals.size());
for (int i = 0; i < intervals.size(); ++i) {
const std::vector<int> vec = {intervals[i].y(), intervals[i].left_x(),
intervals[i].right_x()};
EXPECT_EQ(vec, kTestPatternIntervals[i]);
}
}
TEST(LocationOpenCvTest, EnlargeLocationMaskGrow) {
const float grow_factor = 1.3;
auto test_location = TestPatternIntervalsToMaskLocation();
const float sum = cv::sum(*GetCvMask(test_location))[0];
EnlargeLocation(test_location, grow_factor);
const float grown_sum = cv::sum(*GetCvMask(test_location))[0];
EXPECT_GT(grown_sum, sum);
}
TEST(LocationOpenCvTest, EnlargeMaskShrink) {
const float shrink_factor = 0.7;
auto test_location = TestPatternIntervalsToMaskLocation();
const float sum = cv::sum(*GetCvMask(test_location))[0];
EnlargeLocation(test_location, shrink_factor);
const float shrunk_sum = cv::sum(*GetCvMask(test_location))[0];
EXPECT_GT(sum, shrunk_sum);
}
TEST(LocationOpenCvTest, EnlargeBBox) {
const float test_factor = 1.2f;
auto relative_bbox =
Location::CreateRelativeBBoxLocation(0.5f, 0.3f, 0.2f, 0.6f);
EnlargeLocation(relative_bbox, test_factor);
auto enlarged_relative_bbox_rect = relative_bbox.GetRelativeBBox();
EXPECT_NEAR(enlarged_relative_bbox_rect.xmin(), 0.48f, kEps);
EXPECT_NEAR(enlarged_relative_bbox_rect.ymin(), 0.24f, kEps);
EXPECT_NEAR(enlarged_relative_bbox_rect.Width(), 0.24f, kEps);
EXPECT_NEAR(enlarged_relative_bbox_rect.Height(), 0.72f, kEps);
auto bbox = Location::CreateBBoxLocation(50, 30, 20, 60);
EnlargeLocation(bbox, test_factor);
auto enlarged_bbox_rect = bbox.GetBBox<Rectangle_i>();
EXPECT_EQ(enlarged_bbox_rect.xmin(), 48);
EXPECT_EQ(enlarged_bbox_rect.ymin(), 24);
EXPECT_EQ(enlarged_bbox_rect.Width(), 24);
EXPECT_EQ(enlarged_bbox_rect.Height(), 72);
}
TEST(LocationOpenCvTest, ConvertRelativeBBoxToCvMask) {
const float rel_x_min = 0.1;
const float rel_y_min = 0.2;
const float rel_width = 0.3;
const float rel_height = 0.6;
const int width = 10;
const int height = 20;
cv::Size expected_size(width, height);
LocationData data;
data.set_format(LocationData::RELATIVE_BOUNDING_BOX);
data.mutable_relative_bounding_box()->set_xmin(rel_x_min);
data.mutable_relative_bounding_box()->set_ymin(rel_y_min);
data.mutable_relative_bounding_box()->set_width(rel_width);
data.mutable_relative_bounding_box()->set_height(rel_height);
Location test_location(data);
const int x_start = rel_x_min * width;
const int x_end = x_start + rel_width * width;
const int y_start = rel_y_min * height;
const int y_end = y_start + rel_height * height;
const auto cv_mask = *ConvertToCvMask(test_location, width, height);
EXPECT_EQ(cv_mask.size(), expected_size);
for (int y = 0; y < cv_mask.rows; ++y) {
for (int x = 0; x < cv_mask.cols; ++x) {
bool in_mask = (x >= x_start && x < x_end && y >= y_start && y < y_end);
float expected_value = in_mask ? 1 : 0;
ASSERT_EQ(cv_mask.at<float>(y, x), expected_value);
}
}
}
TEST(LocationOpenCvTest, GetCvMask) {
auto test_location = TestPatternIntervalsToMaskLocation();
auto cv_mask = *GetCvMask(test_location);
EXPECT_EQ(cv_mask.cols * cv_mask.rows, kTestPatternVector.size());
int flat_idx = 0;
for (auto it = cv_mask.begin<uint8>(); it != cv_mask.end<uint8>(); ++it) {
const uint8 expected_value = kTestPatternVector[flat_idx] == 0 ? 0 : 255;
EXPECT_EQ(*it, expected_value);
flat_idx++;
}
}
} // namespace mediapipe
+2
View File
@@ -46,6 +46,7 @@ cc_library(
"//mediapipe/framework:type_map",
"//mediapipe/framework/deps:mathutil",
"//mediapipe/framework/formats:location",
"//mediapipe/framework/formats:location_opencv",
"//mediapipe/framework/formats/motion:optical_flow_field_data_cc_proto",
"//mediapipe/framework/port:file_helpers",
"//mediapipe/framework/port:integral_types",
@@ -67,6 +68,7 @@ cc_test(
deps = [
":optical_flow_field",
"//mediapipe/framework/deps:file_path",
"//mediapipe/framework/formats:location_opencv",
"//mediapipe/framework/port:file_helpers",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:integral_types",
@@ -22,6 +22,7 @@
#include "absl/strings/string_view.h"
#include "mediapipe/framework/deps/mathutil.h"
#include "mediapipe/framework/formats/location.h"
#include "mediapipe/framework/formats/location_opencv.h"
#include "mediapipe/framework/port/file_helpers.h"
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/logging.h"
@@ -304,6 +305,6 @@ Location OpticalFlowField::FindMotionInconsistentPixels(
}
}
}
return Location::CreateCvMaskLocation<uint8>(occluded);
return CreateCvMaskLocation<uint8>(occluded);
}
} // namespace mediapipe
@@ -20,6 +20,7 @@
#include "absl/flags/flag.h"
#include "mediapipe/framework/deps/file_path.h"
#include "mediapipe/framework/formats/location_opencv.h"
#include "mediapipe/framework/port/file_helpers.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/integral_types.h"
@@ -289,8 +290,8 @@ TEST(OpticalFlowField, Occlusions) {
OpticalFlowField::EstimateMotionConsistencyOcclusions(
OpticalFlowField(forward), OpticalFlowField(backward), 0.5,
&occlusion_mask, &disocclusion_mask);
std::unique_ptr<cv::Mat> occlusion_mat = occlusion_mask.GetCvMask();
std::unique_ptr<cv::Mat> disocclusion_mat = disocclusion_mask.GetCvMask();
std::unique_ptr<cv::Mat> occlusion_mat = GetCvMask(occlusion_mask);
std::unique_ptr<cv::Mat> disocclusion_mat = GetCvMask(disocclusion_mask);
EXPECT_EQ(3, occlusion_mat->rows);
EXPECT_EQ(3, disocclusion_mat->rows);
EXPECT_EQ(4, occlusion_mat->cols);
+39 -15
View File
@@ -338,6 +338,7 @@ Tensor::OpenGlBufferView Tensor::GetOpenGlBufferReadView() const {
void* ptr =
glMapBufferRange(GL_SHADER_STORAGE_BUFFER, 0, bytes(),
GL_MAP_INVALIDATE_BUFFER_BIT | GL_MAP_WRITE_BIT);
CHECK(ptr) << "glMapBufferRange failed: " << glGetError();
std::memcpy(ptr, cpu_buffer_, bytes());
glUnmapBuffer(GL_SHADER_STORAGE_BUFFER);
}
@@ -415,6 +416,11 @@ void Tensor::Move(Tensor* src) {
Tensor::Tensor(ElementType element_type, const Shape& shape)
: element_type_(element_type), shape_(shape) {}
Tensor::Tensor(ElementType element_type, const Shape& shape,
const QuantizationParameters& quantization_parameters)
: element_type_(element_type),
shape_(shape),
quantization_parameters_(quantization_parameters) {}
#if MEDIAPIPE_METAL_ENABLED
void Tensor::Invalidate() {
@@ -485,10 +491,15 @@ Tensor::CpuReadView Tensor::GetCpuReadView() const {
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)};
if (__builtin_available(android 26, *)) {
void* ptr = MapAhwbToCpuRead();
if (ptr) {
valid_ |= kValidCpu;
return {ptr, std::move(lock), [ahwb = ahwb_] {
auto error = AHardwareBuffer_unlock(ahwb, nullptr);
CHECK(error == 0) << "AHardwareBuffer_unlock " << error;
}};
}
}
#endif // MEDIAPIPE_TENSOR_USE_AHWB
@@ -553,11 +564,7 @@ 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 {
@@ -565,14 +572,17 @@ Tensor::CpuWriteView Tensor::GetCpuWriteView() const {
AllocateCpuBuffer();
valid_ = kValidCpu;
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
void* ptr = MapAhwbToCpuWrite();
if (ptr) {
return {ptr, ahwb_, &fence_fd_, std::move(lock)};
if (__builtin_available(android 26, *)) {
void* ptr = MapAhwbToCpuWrite();
if (ptr) {
return {ptr, std::move(lock), [ahwb = ahwb_, fence_fd = &fence_fd_] {
auto error = AHardwareBuffer_unlock(ahwb, fence_fd);
CHECK(error == 0) << "AHardwareBuffer_unlock " << error;
}};
}
}
return {cpu_buffer_, nullptr, nullptr, std::move(lock)};
#else
return {cpu_buffer_, std::move(lock)};
#endif // MEDIAPIPE_TENSOR_USE_AHWB
return {cpu_buffer_, std::move(lock)};
}
void Tensor::AllocateCpuBuffer() const {
@@ -590,7 +600,21 @@ void Tensor::AllocateCpuBuffer() const {
void Tensor::SetPreferredStorageType(StorageType type) {
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
use_ahwb_ = type == StorageType::kAhwb;
if (__builtin_available(android 26, *)) {
use_ahwb_ = type == StorageType::kAhwb;
VLOG(4) << "Tensor: use of AHardwareBuffer is "
<< (use_ahwb_ ? "allowed" : "not allowed");
}
#else
VLOG(4) << "Tensor: use of AHardwareBuffer is not allowed";
#endif // MEDIAPIPE_TENSOR_USE_AHWB
}
Tensor::StorageType Tensor::GetPreferredStorageType() {
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
return use_ahwb_ ? StorageType::kAhwb : StorageType::kDefault;
#else
return StorageType::kDefault;
#endif // MEDIAPIPE_TENSOR_USE_AHWB
}
+40 -30
View File
@@ -16,6 +16,7 @@
#define MEDIAPIPE_FRAMEWORK_FORMATS_TENSOR_H_
#include <algorithm>
#include <functional>
#include <initializer_list>
#include <tuple>
#include <type_traits>
@@ -30,10 +31,12 @@
#import <Metal/Metal.h>
#endif // MEDIAPIPE_METAL_ENABLED
#if __ANDROID_API__ >= 26 || defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__)
#define MEDIAPIPE_TENSOR_USE_AHWB 1
#endif // __ANDROID_API__ >= 26 ||
// defined(__ANDROID_UNAVAILABLE_SYMBOLS_ARE_WEAK__)
#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"
@@ -86,7 +89,7 @@ class Tensor {
public:
// No resources are allocated here.
enum class ElementType { kNone, kFloat16, kFloat32, kUInt8, kInt8 };
enum class ElementType { kNone, kFloat16, kFloat32, kUInt8, kInt8, kInt32 };
struct Shape {
Shape() = default;
Shape(std::initializer_list<int> dimensions) : dims(dimensions) {}
@@ -98,8 +101,19 @@ class Tensor {
}
std::vector<int> dims;
};
// Quantization parameters corresponding to the zero_point and scale value
// made available by TfLite quantized (uint8/int8) tensors.
struct QuantizationParameters {
QuantizationParameters() = default;
QuantizationParameters(float scale, int zero_point)
: scale(scale), zero_point(zero_point) {}
float scale = 1.0f;
int zero_point = 0;
};
Tensor(ElementType element_type, const Shape& shape);
Tensor(ElementType element_type, const Shape& shape,
const QuantizationParameters& quantization_parameters);
// Non-copyable.
Tensor(const Tensor&) = delete;
@@ -120,36 +134,21 @@ class Tensor {
}
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
release_callback_ = std::exchange(src.release_callback_, nullptr);
}
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
~CpuView() {
if (ahwb_) {
auto error = AHardwareBuffer_unlock(ahwb_, fence_fd_);
CHECK(error == 0) << "AHardwareBuffer_unlock " << error;
}
if (release_callback_) release_callback_();
}
#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)
CpuView(T* buffer, std::unique_ptr<absl::MutexLock>&& lock,
std::function<void()> release_callback = nullptr)
: 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
release_callback_(release_callback) {}
T* buffer_;
std::function<void()> release_callback_;
};
using CpuReadView = CpuView<const void>;
CpuReadView GetCpuReadView() const;
@@ -184,6 +183,7 @@ class Tensor {
#endif // MEDIAPIPE_METAL_ENABLED
#ifdef MEDIAPIPE_TENSOR_USE_AHWB
using FinishingFunc = std::function<bool(bool)>;
class AHardwareBufferView : public View {
public:
AHardwareBuffer* handle() const { return handle_; }
@@ -195,15 +195,16 @@ class Tensor {
release_callback_ = std::exchange(src.release_callback_, nullptr);
}
int file_descriptor() const { return file_descriptor_; }
void SetReadingFinishedFunc(std::function<bool()>&& func) {
void SetReadingFinishedFunc(FinishingFunc&& func) {
CHECK(ahwb_written_)
<< "AHWB write view can't accept 'reading finished callback'";
*ahwb_written_ = std::move(func);
}
void SetWritingFinishedFD(int fd) {
void SetWritingFinishedFD(int fd, FinishingFunc func = nullptr) {
CHECK(fence_fd_)
<< "AHWB read view can't accept 'writing finished file descriptor'";
*fence_fd_ = fd;
*ahwb_written_ = std::move(func);
}
// The function is called when the tensor is released.
void SetReleaseCallback(std::function<void()> callback) {
@@ -213,7 +214,7 @@ class Tensor {
protected:
friend class Tensor;
AHardwareBufferView(AHardwareBuffer* handle, int file_descriptor,
int* fence_fd, std::function<bool()>* ahwb_written,
int* fence_fd, FinishingFunc* ahwb_written,
std::function<void()>* release_callback,
std::unique_ptr<absl::MutexLock>&& lock)
: View(std::move(lock)),
@@ -226,7 +227,7 @@ class Tensor {
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_;
FinishingFunc* ahwb_written_;
std::function<void()>* release_callback_;
};
AHardwareBufferView GetAHardwareBufferReadView() const;
@@ -301,6 +302,9 @@ class Tensor {
const Shape& shape() const { return shape_; }
ElementType element_type() const { return element_type_; }
const QuantizationParameters& quantization_parameters() const {
return quantization_parameters_;
}
int element_size() const {
switch (element_type_) {
case ElementType::kNone:
@@ -313,6 +317,8 @@ class Tensor {
return 1;
case ElementType::kInt8:
return 1;
case ElementType::kInt32:
return sizeof(int32_t);
}
}
int bytes() const { return shape_.num_elements() * element_size(); }
@@ -337,6 +343,7 @@ class Tensor {
kAhwb,
};
static void SetPreferredStorageType(StorageType type);
static StorageType GetPreferredStorageType();
private:
void Move(Tensor*);
@@ -344,6 +351,7 @@ class Tensor {
ElementType element_type_;
Shape shape_;
QuantizationParameters quantization_parameters_;
// The flags describe the current source of truth resource type.
enum {
@@ -383,13 +391,15 @@ class Tensor {
// 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_;
// If the input parameter is 'true' then wait for the writing to be finished.
mutable FinishingFunc 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
// Expects the target SSBO to be already bound.
bool AllocateAhwbMapToSsbo() const;
bool InsertAhwbToSsboFence() const;
void MoveAhwbStuff(Tensor* src);
+137 -97
View File
@@ -50,8 +50,9 @@ bool IsGlSupported() {
return extensions_allowed;
}
absl::Status MapAHardwareBufferToGlBuffer(AHardwareBuffer* handle, size_t size,
GLuint name) {
// Expects the target SSBO to be already bound.
absl::Status MapAHardwareBufferToGlBuffer(AHardwareBuffer* handle,
size_t size) {
if (!IsGlSupported()) {
return absl::UnknownError(
"No GL extension functions found to bind AHardwareBuffer and "
@@ -96,33 +97,71 @@ class DelayedReleaser {
static void Add(AHardwareBuffer* ahwb, GLuint opengl_buffer,
EGLSyncKHR ssbo_sync, GLsync ssbo_read,
std::function<bool()>&& ahwb_written,
Tensor::FinishingFunc&& ahwb_written,
std::shared_ptr<mediapipe::GlContext> gl_context,
std::function<void()>&& callback) {
static absl::Mutex mutex;
absl::MutexLock lock(&mutex);
std::deque<std::unique_ptr<DelayedReleaser>> to_release_local;
using std::swap;
// IsSignaled will grab other mutexes, so we don't want to call it while
// holding the deque mutex.
{
absl::MutexLock lock(&mutex);
swap(to_release_local, to_release_);
}
// Using `new` to access a non-public constructor.
to_release_.emplace_back(absl::WrapUnique(new DelayedReleaser(
to_release_local.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();) {
for (auto it = to_release_local.begin(); it != to_release_local.end();) {
if ((*it)->IsSignaled()) {
it = to_release_.erase(it);
it = to_release_local.erase(it);
} else {
++it;
}
}
{
absl::MutexLock lock(&mutex);
to_release_.insert(to_release_.end(),
std::make_move_iterator(to_release_local.begin()),
std::make_move_iterator(to_release_local.end()));
to_release_local.clear();
}
}
~DelayedReleaser() {
AHardwareBuffer_release(ahwb_);
if (release_callback_) release_callback_();
if (__builtin_available(android 26, *)) {
AHardwareBuffer_release(ahwb_);
}
}
bool IsSignaled() {
CHECK(!(ssbo_read_ && ahwb_written_))
<< "ssbo_read_ and ahwb_written_ cannot both be set";
bool ready = true;
if (ahwb_written_) {
if (!ahwb_written_()) return false;
if (!ahwb_written_(false)) {
ready = false;
}
}
if (ssbo_read_ != 0) {
gl_context_->Run([this, &ready]() {
GLenum status = glClientWaitSync(ssbo_read_, 0,
/* timeout ns = */ 0);
if (status != GL_CONDITION_SATISFIED && status != GL_ALREADY_SIGNALED) {
ready = false;
return;
}
glDeleteSync(ssbo_read_);
ssbo_read_ = 0;
});
}
if (ready && gl_context_) {
gl_context_->Run([this]() {
if (fence_sync_ != EGL_NO_SYNC_KHR && IsGlSupported()) {
auto egl_display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
@@ -134,33 +173,9 @@ class DelayedReleaser {
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;
return ready;
}
protected:
@@ -170,14 +185,14 @@ class DelayedReleaser {
EGLSyncKHR fence_sync_;
// TODO: use wrapper instead.
GLsync ssbo_read_;
std::function<bool()> ahwb_written_;
Tensor::FinishingFunc 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,
Tensor::FinishingFunc&& ahwb_written,
std::shared_ptr<mediapipe::GlContext> gl_context,
std::function<void()>&& callback)
: ahwb_(ahwb),
@@ -240,44 +255,49 @@ Tensor::AHardwareBufferView Tensor::GetAHardwareBufferWriteView(
valid_ = kValidAHardwareBuffer;
return {ahwb_,
/*ssbo_written=*/-1,
&fence_fd_, // For SetWritingFinishedFD.
/*ahwb_written=*/nullptr, // The lifetime is managed by SSBO.
&fence_fd_, // For SetWritingFinishedFD.
&ahwb_written_,
&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);
if (__builtin_available(android 26, *)) {
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;
}
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;
}
return true;
return false;
}
bool Tensor::AllocateAhwbMapToSsbo() const {
if (AllocateAHardwareBuffer()) {
if (MapAHardwareBufferToGlBuffer(ahwb_, bytes(), opengl_buffer_).ok()) {
glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
return true;
if (__builtin_available(android 26, *)) {
if (AllocateAHardwareBuffer()) {
if (MapAHardwareBufferToGlBuffer(ahwb_, bytes()).ok()) {
glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
return true;
}
// Unable to make OpenGL <-> AHWB binding. Use regular SSBO instead.
AHardwareBuffer_release(ahwb_);
ahwb_ = nullptr;
}
// Unable to make OpenGL <-> AHWB binding. Use regular SSBO instead.
AHardwareBuffer_release(ahwb_);
ahwb_ = nullptr;
}
return false;
}
@@ -295,12 +315,19 @@ bool Tensor::InsertAhwbToSsboFence() const {
// Server-side fence.
auto egl_display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
if (egl_display == EGL_NO_DISPLAY) return true;
// EGL will take ownership of the passed fd if eglCreateSyncKHR is
// successful.
int fd_for_egl = dup(fence_fd_);
EGLint sync_attribs[] = {EGL_SYNC_NATIVE_FENCE_FD_ANDROID,
(EGLint)fence_fd_, EGL_NONE};
(EGLint)fd_for_egl, 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);
} else {
close(fd_for_egl);
}
}
return true;
@@ -321,49 +348,62 @@ void Tensor::ReleaseAhwbStuff() {
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_);
if (__builtin_available(android 26, *)) {
if (ahwb_) {
if (ssbo_read_ != 0 || fence_sync_ != EGL_NO_SYNC_KHR || ahwb_written_) {
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 {
if (release_callback_) release_callback_();
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();
if (__builtin_available(android 26, *)) {
if (ahwb_) {
if (!(valid_ & kValidCpu)) {
if ((valid_ & kValidOpenGlBuffer) && ssbo_written_ == -1) {
// EGLSync is failed. Use another synchronization method.
// TODO: Use tflite::gpu::GlBufferSync and GlActiveSync.
glFinish();
} else if (valid_ & kValidAHardwareBuffer) {
CHECK(ahwb_written_) << "Ahwb-to-Cpu synchronization requires the "
"completion function to be set";
CHECK(ahwb_written_(true))
<< "An error oqcured while waiting for the buffer to be written";
}
}
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;
}
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;
if (__builtin_available(android 26, *)) {
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;
}
@@ -1,15 +1,13 @@
#include "mediapipe/framework/formats/tensor.h"
#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"
#if !MEDIAPIPE_DISABLE_GPU
namespace mediapipe {
#if !MEDIAPIPE_DISABLE_GPU
class TensorAhwbTest : public mediapipe::GpuTestBase {
public:
};
@@ -55,5 +53,4 @@ TEST_F(TensorAhwbTest, TestCpuThenGl) {
} // namespace mediapipe
#endif // !MEDIAPIPE_DISABLE_GPU
#endif // MEDIAPIPE_TENSOR_USE_AHWB
@@ -21,8 +21,6 @@ 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.
+3
View File
@@ -22,6 +22,9 @@ syntax = "proto2";
package mediapipe;
option java_package = "com.google.mediapipe.proto";
option java_outer_classname = "PacketFactoryProto";
// Options used by a PacketFactory to create the Packet.
message PacketFactoryOptions {
extensions 20000 to max;
@@ -22,8 +22,15 @@ syntax = "proto2";
package mediapipe;
option java_package = "com.google.mediapipe.proto";
option java_outer_classname = "PacketGeneratorProto";
// Options used by a PacketGenerator.
message PacketGeneratorOptions {
// If true, this proto specifies a subset of field values,
// which should override corresponding field values.
optional bool merge_fields = 1 [default = true];
extensions 20000 to max;
}
-2
View File
@@ -18,8 +18,6 @@ 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.
+1 -2
View File
@@ -76,8 +76,7 @@
#endif
#define MEDIAPIPE_METAL_ENABLED 0
#elif defined(MEDIAPIPE_IOS)
// TODO: use MEDIAPIPE_OPENGL_ES_30 for iOS as max version.
#define MEDIAPIPE_OPENGL_ES_VERSION MEDIAPIPE_OPENGL_ES_20
#define MEDIAPIPE_OPENGL_ES_VERSION MEDIAPIPE_OPENGL_ES_30
#define MEDIAPIPE_METAL_ENABLED 1
#elif defined(MEDIAPIPE_OSX)
#define MEDIAPIPE_OPENGL_ES_VERSION 0
+7 -1
View File
@@ -32,7 +32,6 @@ typedef uint16_t uint16;
typedef uint32_t uint32;
typedef uint64_t uint64;
typedef signed int char32;
typedef unsigned long uword_t;
#define GG_LONGLONG(x) x##LL
@@ -40,6 +39,11 @@ typedef unsigned long uword_t;
#define GG_LL_FORMAT "ll" // As in "%lld". Note that "q" is poor form also.
#define GG_LL_FORMAT_W L"ll"
// Add namespace here to avoid conflict with other libraries.
namespace mediapipe {
typedef signed int char32;
const uint8 kuint8max{0xFF};
const uint16 kuint16max{0xFFFF};
const uint32 kuint32max{0xFFFFFFFF};
@@ -57,4 +61,6 @@ typedef uint64 Fprint;
static const Fprint kIllegalFprint = 0;
static const Fprint kMaxFprint = GG_ULONGLONG(0xFFFFFFFFFFFFFFFF);
} // namespace mediapipe
#endif // MEDIAPIPE_PORT_INTEGRAL_TYPES_H_
@@ -20,7 +20,9 @@
#ifdef CV_VERSION_EPOCH // for OpenCV 2.x
#include <opencv2/core/core.hpp>
#else
#if CV_VERSION_MAJOR == 3
#include <opencv2/cvconfig.h>
#endif
#include <opencv2/core.hpp>
#endif
+18
View File
@@ -195,6 +195,7 @@ cc_library(
"//mediapipe/framework:timestamp",
"//mediapipe/framework/port:integral_types",
"@com_google_absl//absl/container:node_hash_map",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
],
)
@@ -317,6 +318,23 @@ cc_library(
}),
)
# For a more maintainable build this target should not exist and the headers
# should be split into the existing cc_library targets, but this change was
# automatically done so that we can remove long standing issues and complexity
# in the build system. It's up to the OWNERS of this package to get rid of it or
# not. The use of the textual_hdrs attribute is discouraged, use hdrs instead.
# Here it is used to avoid header parsing errors in packages where the feature
# parse_headers was enabled since loose headers were not being parsed.
cc_library(
name = "loose_headers",
tags = ["avoid_dep"],
textual_hdrs = [
"graph_profiler.h",
"graph_profiler_stub.h",
],
visibility = ["//mediapipe/framework:__pkg__"],
)
cc_test(
name = "reporter_test",
srcs = ["reporter_test.cc"],
@@ -193,6 +193,7 @@ void GraphProfiler::Initialize(
"Calculator \"$0\" has already been added.", node_name);
}
profile_builder_ = std::make_unique<GraphProfileBuilder>(this);
is_initialized_ = true;
}
+9 -1
View File
@@ -14,9 +14,9 @@
#include "mediapipe/framework/profiler/graph_tracer.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "mediapipe/framework/calculator_context.h"
#include "mediapipe/framework/calculator_profile.pb.h"
#include "mediapipe/framework/input_stream_shard.h"
#include "mediapipe/framework/output_stream_shard.h"
#include "mediapipe/framework/packet.h"
@@ -117,14 +117,22 @@ Timestamp GraphTracer::TimestampAfter(absl::Time begin_time) {
return TraceBuilder::TimestampAfter(trace_buffer_, begin_time);
}
// The mutex to guard GraphTracer::trace_builder_.
absl::Mutex* trace_builder_mutex() {
static absl::Mutex trace_builder_mutex(absl::kConstInit);
return &trace_builder_mutex;
}
void GraphTracer::GetTrace(absl::Time begin_time, absl::Time end_time,
GraphTrace* result) {
absl::MutexLock lock(trace_builder_mutex());
trace_builder_.CreateTrace(trace_buffer_, begin_time, end_time, result);
trace_builder_.Clear();
}
void GraphTracer::GetLog(absl::Time begin_time, absl::Time end_time,
GraphTrace* result) {
absl::MutexLock lock(trace_builder_mutex());
trace_builder_.CreateLog(trace_buffer_, begin_time, end_time, result);
trace_builder_.Clear();
}
+3
View File
@@ -24,6 +24,9 @@ package mediapipe;
import "mediapipe/framework/mediapipe_options.proto";
option java_package = "com.google.mediapipe.proto";
option java_outer_classname = "StatusHandlerProto";
// The settings specifying a status handler and its required external inputs.
message StatusHandlerConfig {
// The name of the registered status handler class.
-1
View File
@@ -24,7 +24,6 @@ import "mediapipe/framework/mediapipe_options.proto";
option java_package = "com.google.mediapipe.proto";
option java_outer_classname = "StreamHandlerProto";
option objc_class_prefix = "MediaPipe";
// Settings specifying an input stream handler.
message InputStreamHandlerConfig {
@@ -22,8 +22,6 @@ package mediapipe;
import "mediapipe/framework/calculator.proto";
option objc_class_prefix = "MediaPipe";
message RandomMatrixCalculatorOptions {
extend CalculatorOptions {
optional RandomMatrixCalculatorOptions ext = 52056136;
+5
View File
@@ -50,3 +50,8 @@ mediapipe_proto_library(
"//mediapipe/framework:packet_generator_proto",
],
)
exports_files([
"perfetto_detailed.pbtxt",
"perfetto_minimal.pbtxt",
])
+32
View File
@@ -0,0 +1,32 @@
# Perfetto Config with minimal information needed for tracing MediaPipe functions.
# Use when benchmarking to reduce latency overhead from Perfetto trace.
# proto-file: third_party/perfetto/protos/perfetto/config/perfetto_config.proto
# proto-message: TraceConfig
buffers: {
size_kb: 150000
fill_policy: DISCARD
}
data_sources: {
config {
name: "track_event"
}
}
data_sources: {
config {
name: "linux.ftrace"
ftrace_config {
# Scheduling information & process tracking. Useful for:
# - what is happening on each CPU at each moment
ftrace_events: "power/cpu_frequency"
ftrace_events: "power/cpu_idle"
ftrace_events: "sched/sched_switch"
compact_sched {
enabled: true
}
}
}
}
write_into_file: true
file_write_period_ms: 500
+27 -6
View File
@@ -53,6 +53,7 @@ bzl_library(
"build_defs.bzl",
],
visibility = [
"//mediapipe/app/pursuit/wasm:__subpackages__",
"//mediapipe/app/xeno/catalog:__subpackages__",
"//mediapipe/framework:__subpackages__",
],
@@ -768,7 +769,6 @@ cc_library(
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/port:advanced_proto",
"//mediapipe/framework/port:file_helpers",
"//mediapipe/framework/port:gtest",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
@@ -829,24 +829,40 @@ cc_library(
],
deps = [
":container_util",
":options_util",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:collection_item_id",
"//mediapipe/framework:input_stream_shard",
"//mediapipe/framework:output_stream_shard",
"//mediapipe/framework/deps:mathutil",
"//mediapipe/framework/formats:video_stream_header",
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/stream_handler:immediate_input_stream_handler",
"//mediapipe/framework/tool:switch_container_cc_proto",
"@com_google_absl//absl/strings",
],
alwayslink = 1,
)
mediapipe_cc_test(
name = "switch_mux_calculator_test",
srcs = ["switch_mux_calculator_test.cc"],
deps = [
":container_util",
":switch_mux_calculator",
"//mediapipe/calculators/core:pass_through_calculator",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:subgraph",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/stream_handler:immediate_input_stream_handler",
"@com_google_absl//absl/strings",
],
)
mediapipe_proto_library(
name = "switch_container_proto",
srcs = ["switch_container.proto"],
@@ -867,13 +883,17 @@ cc_library(
":subgraph_expansion",
":switch_demux_calculator",
":switch_mux_calculator",
"//mediapipe/calculators/core:packet_sequencer_calculator",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:mediapipe_options_cc_proto",
"//mediapipe/framework:calculator_options_cc_proto",
"//mediapipe/framework:subgraph",
"//mediapipe/framework/deps:registration",
"//mediapipe/framework/port:core_proto",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/tool:switch_container_cc_proto",
"@com_google_absl//absl/strings",
],
alwayslink = 1,
)
@@ -890,6 +910,7 @@ cc_test(
"//mediapipe/calculators/core:pass_through_calculator",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:stream_handler_cc_proto",
"//mediapipe/framework:subgraph",
"//mediapipe/framework:test_calculators",
"//mediapipe/framework/port:gtest_main",
@@ -0,0 +1,23 @@
"""Build rule to depend on files downloaded from GCS."""
# buildifier: disable=unnamed-macro
def mediapipe_files(srcs):
"""Links file from GCS with the current directory.
Args:
srcs: the names of the mediapipe_file target, which is also the name of
the MediaPipe file in external_files.bzl. For example, if `name` is Foo,
`mediapipe_file` will create a link to the downloaded file
"@com_google_mediapipe_Foo_tfile" to the current directory as
"Foo.tflite".
"""
for src in srcs:
archive_name = "com_google_mediapipe_%s" % src.replace("/", "_").replace(".", "_")
native.genrule(
name = "%s_ln" % archive_name,
srcs = ["@%s//file" % archive_name],
outs = [src],
output_to_bindir = 1,
cmd = "ln $< $@",
)
@@ -380,7 +380,7 @@ absl::StatusOr<std::vector<FieldData>> GetFieldValues(
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() : "#",
"Missing field value: ", head.field ? head.field->name() : "#",
" at index: ", index));
}
if (!tail.empty()) {
@@ -402,7 +402,7 @@ absl::StatusOr<FieldData> GetField(const FieldData& message_data,
if (results.empty()) {
FieldPathEntry tail = field_path.back();
return absl::OutOfRangeError(absl::StrCat(
"Missing feild value: ", tail.field ? tail.field->name() : "##",
"Missing field value: ", tail.field ? tail.field->name() : "##",
" at index: ", tail.index));
}
return results[0];
@@ -414,7 +414,7 @@ absl::Status SetFieldValues(FieldData& message_data,
const std::vector<FieldData>& values) {
if (field_path.empty()) {
if (values.empty()) {
return absl::InvalidArgumentError("Missing feild value.");
return absl::InvalidArgumentError("Missing field value.");
}
message_data = values[0];
return absl::OkStatus();
@@ -345,7 +345,7 @@ absl::Status Equals(std::vector<FieldData> b1, std::vector<FieldData> b2) {
return absl::OkStatus();
}
// Unit-tests for graph options feild accessors from options_field_util.
// Unit-tests for graph options field accessors from options_field_util.
class OptionsFieldUtilTest : public ::testing::Test {
protected:
void SetUp() override {}
@@ -581,7 +581,7 @@ TEST_F(OptionsFieldUtilTest, MergeFieldValuesInt) {
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"));
HasSubstr("Missing field value: num_lights at index: 1"));
}
// Tests MergeFieldValues applied to a protobuf field.
+64 -9
View File
@@ -12,13 +12,18 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <fstream>
#include <iostream>
#include <sstream>
#include <functional>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <utility>
#include <vector>
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/mediapipe_options.pb.h"
#include "mediapipe/framework/calculator_options.pb.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
@@ -98,6 +103,18 @@ CalculatorGraphConfig::Node* BuildMuxNode(
return result;
}
// Returns a PacketSequencerCalculator node.
CalculatorGraphConfig::Node* BuildTimestampNode(CalculatorGraphConfig* config,
bool synchronize_io) {
CalculatorGraphConfig::Node* result = config->add_node();
*result->mutable_calculator() = "PacketSequencerCalculator";
if (synchronize_io) {
*result->mutable_input_stream_handler()->mutable_input_stream_handler() =
"DefaultInputStreamHandler";
}
return result;
}
// Copies options from one node to another.
void CopyOptions(const CalculatorGraphConfig::Node& source,
CalculatorGraphConfig::Node* dest) {
@@ -214,6 +231,14 @@ absl::Status ValidateContract(
return absl::OkStatus();
}
// Returns true if a set of streams references a certain tag name.
bool HasTag(const proto_ns::RepeatedPtrField<std::string>& streams,
std::string tag) {
std::map<TagIndex, std::string> tags;
ParseTags(streams, &tags);
return tags.count({tag, 0}) > 0;
}
absl::StatusOr<CalculatorGraphConfig> SwitchContainer::GetConfig(
const Subgraph::SubgraphOptions& options) {
CalculatorGraphConfig config;
@@ -232,32 +257,62 @@ absl::StatusOr<CalculatorGraphConfig> SwitchContainer::GetConfig(
ParseTags(container_streams.input_side_packet(), &side_input_tags);
ParseTags(container_streams.output_side_packet(), &side_output_tags);
CalculatorGraphConfig::Node* select_node = nullptr;
CalculatorGraphConfig::Node* enable_node = nullptr;
std::string select_stream = "SELECT:gate_select";
std::string enable_stream = "ENABLE:gate_enable";
// Add a PacketSequencerCalculator node for "SELECT" or "ENABLE" streams.
bool synchronize_io =
Subgraph::GetOptions<mediapipe::SwitchContainerOptions>(options)
.synchronize_io();
if (HasTag(container_node.input_stream(), "SELECT")) {
select_node = BuildTimestampNode(&config, synchronize_io);
select_node->add_input_stream("INPUT:gate_select");
select_node->add_output_stream("OUTPUT:gate_select_timed");
select_stream = "SELECT:gate_select_timed";
}
if (HasTag(container_node.input_stream(), "ENABLE")) {
enable_node = BuildTimestampNode(&config, synchronize_io);
enable_node->add_input_stream("INPUT:gate_enable");
enable_node->add_output_stream("OUTPUT:gate_enable_timed");
enable_stream = "ENABLE:gate_enable_timed";
}
// Add a graph node for the demux, mux.
auto demux = BuildDemuxNode(input_tags, container_node, &config);
CopyOptions(container_node, demux);
ClearContainerOptions(demux);
demux->add_input_stream("SELECT:gate_select");
demux->add_input_stream("ENABLE:gate_enable");
demux->add_input_stream(select_stream);
demux->add_input_stream(enable_stream);
demux->add_input_side_packet("SELECT:gate_select");
demux->add_input_side_packet("ENABLE:gate_enable");
auto mux = BuildMuxNode(output_tags, &config);
CopyOptions(container_node, mux);
ClearContainerOptions(mux);
mux->add_input_stream("SELECT:gate_select");
mux->add_input_stream("ENABLE:gate_enable");
mux->add_input_stream(select_stream);
mux->add_input_stream(enable_stream);
mux->add_input_side_packet("SELECT:gate_select");
mux->add_input_side_packet("ENABLE:gate_enable");
// Add input streams for graph and demux.
// Add input streams for graph and demux and the timestamper.
config.add_input_stream("SELECT:gate_select");
config.add_input_stream("ENABLE:gate_enable");
config.add_input_side_packet("SELECT:gate_select");
config.add_input_side_packet("ENABLE:gate_enable");
int tick_index = 0;
for (const auto& p : input_tags) {
std::string stream = CatStream(p.first, p.second);
config.add_input_stream(stream);
demux->add_input_stream(stream);
TagIndex tick_tag{"TICK", tick_index++};
if (select_node) {
select_node->add_input_stream(CatStream(tick_tag, p.second));
}
if (enable_node) {
enable_node->add_input_stream(CatStream(tick_tag, p.second));
}
}
// Add output streams for graph and mux.
@@ -12,7 +12,13 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/strings/str_replace.h"
#include "absl/strings/string_view.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/gmock.h"
@@ -161,7 +167,9 @@ void RunTestContainer(CalculatorGraphConfig supergraph,
// i.e. the one containing the PassThroughCalculator should output the
// input values without changing them.
EXPECT_EQ(out_bar.size(), t);
EXPECT_EQ(out_bar.back().Get<int>(), t);
if (!out_bar.empty()) {
EXPECT_EQ(out_bar.back().Get<int>(), t);
}
}
if (!send_bounds) {
@@ -187,7 +195,9 @@ void RunTestContainer(CalculatorGraphConfig supergraph,
// i.e. the one containing the TripleIntCalculator should output the values
// after tripling them.
EXPECT_EQ(out_bar.size(), t);
EXPECT_EQ(out_bar.back().Get<int>(), t * 3);
if (!out_bar.empty()) {
EXPECT_EQ(out_bar.back().Get<int>(), t * 3);
}
}
MP_ASSERT_OK(graph.CloseAllInputStreams());
@@ -236,10 +246,17 @@ TEST(SwitchContainerTest, ApplyToSubnodes) {
CalculatorGraphConfig supergraph = SubnodeContainerExample();
CalculatorGraphConfig expected_graph =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
node {
name: "switchcontainer__PacketSequencerCalculator"
calculator: "PacketSequencerCalculator"
input_stream: "INPUT:enable"
input_stream: "TICK:foo"
output_stream: "OUTPUT:switchcontainer__gate_enable_timed"
}
node {
name: "switchcontainer__SwitchDemuxCalculator"
calculator: "SwitchDemuxCalculator"
input_stream: "ENABLE:enable"
input_stream: "ENABLE:switchcontainer__gate_enable_timed"
input_stream: "foo"
output_stream: "C0__:switchcontainer__c0__foo"
output_stream: "C1__:switchcontainer__c1__foo"
@@ -262,7 +279,7 @@ TEST(SwitchContainerTest, ApplyToSubnodes) {
node {
name: "switchcontainer__SwitchMuxCalculator"
calculator: "SwitchMuxCalculator"
input_stream: "ENABLE:enable"
input_stream: "ENABLE:switchcontainer__gate_enable_timed"
input_stream: "C0__:switchcontainer__c0__bar"
input_stream: "C1__:switchcontainer__c1__bar"
output_stream: "bar"
@@ -281,7 +298,7 @@ TEST(SwitchContainerTest, ApplyToSubnodes) {
input_stream: "enable"
input_side_packet: "timezone"
)pb");
expected_graph = OrderNodes(expected_graph, {4, 0, 3, 1, 2});
expected_graph = OrderNodes(expected_graph, {5, 0, 1, 4, 2, 3});
MP_EXPECT_OK(tool::ExpandSubgraphs(&supergraph));
EXPECT_THAT(supergraph, mediapipe::EqualsProto(expected_graph));
}
+116 -82
View File
@@ -13,11 +13,15 @@
// limitations under the License.
#include <algorithm>
#include <memory>
#include <functional>
#include <iterator>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <type_traits>
#include <utility>
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/collection_item_id.h"
#include "mediapipe/framework/input_stream_shard.h"
@@ -66,20 +70,26 @@ class SwitchMuxCalculator : public CalculatorBase {
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
private:
// Stores any new input channel history.
void RecordChannel(CalculatorContext* cc);
// Temporarily enqueues every new packet or timestamp bounds.
void RecordPackets(CalculatorContext* cc);
// Immediately sends any packets or timestamp bounds for settled timestamps.
void SendActivePackets(CalculatorContext* cc);
private:
int channel_index_;
std::set<std::string> channel_tags_;
mediapipe::SwitchContainerOptions options_;
// This is used to keep around packets that we've received but not
// relayed yet (because we may not know which channel we should yet be using
// when synchronized_io flag is set).
std::map<Timestamp, std::map<CollectionItemId, Packet>> packet_history_;
// relayed yet (because we may not know which channel we should yet be using).
std::map<CollectionItemId, std::queue<Packet>> packet_queue_;
// Historical channel index values for timestamps where we don't have all
// packets available yet (when synchronized_io flag is set).
// packets available yet.
std::map<Timestamp, int> channel_history_;
// Number of output steams that we already processed for the current output
// timestamp.
int current_processed_stream_count_ = 0;
};
REGISTER_CALCULATOR(SwitchMuxCalculator);
@@ -133,10 +143,28 @@ absl::Status SwitchMuxCalculator::GetContract(CalculatorContract* cc) {
return absl::OkStatus();
}
// Returns the last delivered timestamp for an input stream.
Timestamp SettledTimestamp(const InputStreamShard& input) {
return input.Value().Timestamp();
}
// Returns the last delivered timestamp for channel selection.
Timestamp ChannelSettledTimestamp(CalculatorContext* cc) {
Timestamp result = Timestamp::Done();
if (cc->Inputs().HasTag("ENABLE")) {
result = SettledTimestamp(cc->Inputs().Tag("ENABLE"));
} else if (cc->Inputs().HasTag("SELECT")) {
result = SettledTimestamp(cc->Inputs().Tag("SELECT"));
}
return result;
}
absl::Status SwitchMuxCalculator::Open(CalculatorContext* cc) {
// Initialize channel_index_ and channel_history_.
options_ = cc->Options<mediapipe::SwitchContainerOptions>();
channel_index_ = tool::GetChannelIndex(*cc, channel_index_);
channel_tags_ = ChannelTags(cc->Inputs().TagMap());
channel_history_[Timestamp::Unset()] = channel_index_;
// Relay side packets only from channel_index_.
for (const std::string& tag : ChannelTags(cc->InputSidePackets().TagMap())) {
@@ -150,85 +178,91 @@ absl::Status SwitchMuxCalculator::Open(CalculatorContext* cc) {
return absl::OkStatus();
}
absl::Status SwitchMuxCalculator::Process(CalculatorContext* cc) {
// Update the input channel index if specified.
channel_index_ = tool::GetChannelIndex(*cc, channel_index_);
void SwitchMuxCalculator::RecordChannel(CalculatorContext* cc) {
Timestamp channel_settled = ChannelSettledTimestamp(cc);
int new_channel_index = tool::GetChannelIndex(*cc, channel_index_);
if (options_.synchronize_io()) {
// Start with adding input signals into channel_history_ and packet_history_
if (cc->Inputs().HasTag("ENABLE") &&
!cc->Inputs().Tag("ENABLE").IsEmpty()) {
channel_history_[cc->Inputs().Tag("ENABLE").Value().Timestamp()] =
channel_index_;
// Enque any new input channel and its activation timestamp.
if (channel_settled == cc->InputTimestamp() &&
new_channel_index != channel_index_) {
channel_index_ = new_channel_index;
channel_history_[channel_settled] = channel_index_;
}
}
void SwitchMuxCalculator::RecordPackets(CalculatorContext* cc) {
auto select_id = cc->Inputs().GetId("SELECT", 0);
auto enable_id = cc->Inputs().GetId("ENABLE", 0);
for (auto id = cc->Inputs().BeginId(); id < cc->Inputs().EndId(); ++id) {
if (id == select_id || id == enable_id) continue;
Packet packet = cc->Inputs().Get(id).Value();
// Enque any new packet or timestamp bound.
if (packet.Timestamp() == cc->InputTimestamp()) {
packet_queue_[id].push(packet);
}
if (cc->Inputs().HasTag("SELECT") &&
!cc->Inputs().Tag("SELECT").IsEmpty()) {
channel_history_[cc->Inputs().Tag("SELECT").Value().Timestamp()] =
channel_index_;
}
for (auto input_id = cc->Inputs().BeginId();
input_id < cc->Inputs().EndId(); ++input_id) {
auto& entry = cc->Inputs().Get(input_id);
if (entry.IsEmpty()) {
continue;
}
packet_history_[entry.Value().Timestamp()][input_id] = entry.Value();
}
// Now check if we have enough information to produce any outputs.
while (!channel_history_.empty()) {
// Look at the oldest unprocessed timestamp.
auto it = channel_history_.begin();
auto& packets = packet_history_[it->first];
int total_streams = 0;
// Loop over all outputs to see if we have anything new that we can relay.
for (const std::string& tag : channel_tags_) {
for (int index = 0; index < cc->Outputs().NumEntries(tag); ++index) {
++total_streams;
auto input_id =
cc->Inputs().GetId(tool::ChannelTag(tag, it->second), index);
auto packet_it = packets.find(input_id);
if (packet_it != packets.end()) {
cc->Outputs().Get(tag, index).AddPacket(packet_it->second);
++current_processed_stream_count_;
} else if (it->first <
cc->Inputs().Get(input_id).Value().Timestamp()) {
// Getting here means that input stream that corresponds to this
// output at the timestamp we're trying to process right now has
// already advanced beyond this timestamp. This means that we will
// shouldn't expect a packet for this timestamp anymore, and we can
// safely advance timestamp on the output.
cc->Outputs()
.Get(tag, index)
.SetNextTimestampBound(it->first.NextAllowedInStream());
++current_processed_stream_count_;
}
}
}
if (current_processed_stream_count_ == total_streams) {
// There's nothing else to wait for at the current timestamp, do the
// cleanup and move on to the next one.
packet_history_.erase(it->first);
channel_history_.erase(it);
current_processed_stream_count_ = 0;
} else {
// We're still missing some packets for the current timestamp. Clean up
// those that we just relayed and let the rest wait until the next
// Process() call.
packets.clear();
break;
}
}
} else {
// Relay packets and timestamps only from channel_index_.
}
}
void SwitchMuxCalculator::SendActivePackets(CalculatorContext* cc) {
Timestamp expired_history;
// Iterate through the recent active input channels.
for (auto it = channel_history_.begin(); it != channel_history_.end(); ++it) {
int channel = it->second;
Timestamp channel_start = it->first;
Timestamp channel_end =
(std::next(it) == channel_history_.end())
? ChannelSettledTimestamp(cc).NextAllowedInStream()
: std::next(it)->first;
Timestamp stream_settled = Timestamp::Done();
for (const std::string& tag : channel_tags_) {
for (int index = 0; index < cc->Outputs().NumEntries(tag); ++index) {
auto& output = cc->Outputs().Get(tag, index);
std::string input_tag = tool::ChannelTag(tag, channel_index_);
auto& input = cc->Inputs().Get(input_tag, index);
tool::Relay(input, &output);
std::string input_tag = tool::ChannelTag(tag, channel);
for (int index = 0; index < cc->Inputs().NumEntries(input_tag); ++index) {
CollectionItemId input_id = cc->Inputs().GetId(input_tag, index);
OutputStreamShard& output = cc->Outputs().Get(tag, index);
std::queue<Packet>& q = packet_queue_[input_id];
// Send any packets or bounds from a recent active input channel.
while (!q.empty() && q.front().Timestamp() < channel_end) {
if (q.front().Timestamp() >= channel_start) {
output.AddPacket(q.front());
}
q.pop();
}
stream_settled = std::min(stream_settled,
SettledTimestamp(cc->Inputs().Get(input_id)));
}
}
// A history entry is expired only if all streams have advanced past it.
if (stream_settled.NextAllowedInStream() < channel_end ||
std::next(it) == channel_history_.end()) {
break;
}
expired_history = channel_start;
// Discard any packets or bounds from recent inactive input channels.
for (auto id = cc->Inputs().BeginId(); id < cc->Inputs().EndId(); ++id) {
std::queue<Packet>& q = packet_queue_[id];
while (!q.empty() && q.front().Timestamp() < channel_end) {
q.pop();
}
}
}
// Discard any expired channel history entries.
if (expired_history != Timestamp::Unset()) {
channel_history_.erase(channel_history_.begin(),
std::next(channel_history_.find(expired_history)));
}
}
absl::Status SwitchMuxCalculator::Process(CalculatorContext* cc) {
// Normally packets will arrive on the active channel and will be passed
// through immediately. In the less common case in which the active input
// channel is not known for an input packet timestamp, the input packet is
// queued until the active channel becomes known.
RecordChannel(cc);
RecordPackets(cc);
SendActivePackets(cc);
return absl::OkStatus();
}
@@ -0,0 +1,329 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <string>
#include <vector>
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/framework/tool/container_util.h"
namespace mediapipe {
namespace {
// Returns a CalculatorGraph to run a single calculator.
CalculatorGraph BuildCalculatorGraph(CalculatorGraphConfig::Node node_config) {
CalculatorGraphConfig config;
*config.add_node() = node_config;
*config.mutable_input_stream() = node_config.input_stream();
*config.mutable_output_stream() = node_config.output_stream();
*config.mutable_input_side_packet() = node_config.input_side_packet();
*config.mutable_output_side_packet() = node_config.output_side_packet();
return CalculatorGraph(config);
}
// Creates a string packet.
Packet pack(std::string data, int timestamp) {
return MakePacket<std::string>(data).At(Timestamp(timestamp));
}
// Creates an int packet.
Packet pack(int data, int timestamp) {
return MakePacket<int>(data).At(Timestamp(timestamp));
}
// Tests showing packet channel synchronization through SwitchMuxCalculator.
class SwitchMuxCalculatorTest : public ::testing::Test {
protected:
SwitchMuxCalculatorTest() {}
~SwitchMuxCalculatorTest() override {}
void SetUp() override {}
void TearDown() override {}
// Defines a SwitchMuxCalculator CalculatorGraphConfig::Node.
CalculatorGraphConfig::Node BuildNodeConfig() {
CalculatorGraphConfig::Node result;
*result.mutable_calculator() = "SwitchMuxCalculator";
*result.add_input_stream() = "SELECT:select";
for (int c = 0; c < 3; ++c) {
*result.add_input_stream() =
absl::StrCat(tool::ChannelTag("FRAME", c), ":frame_", c);
*result.add_input_stream() =
absl::StrCat(tool::ChannelTag("MASK", c), ":mask_", c);
}
*result.add_output_stream() = "FRAME:frame";
*result.add_output_stream() = "MASK:mask";
return result;
}
};
// Shows the SwitchMuxCalculator is available.
TEST_F(SwitchMuxCalculatorTest, IsRegistered) {
EXPECT_TRUE(CalculatorBaseRegistry::IsRegistered("SwitchMuxCalculator"));
}
// Shows that channels are queued until packets arrive.
TEST_F(SwitchMuxCalculatorTest, ChannelEarly) {
CalculatorGraphConfig::Node node_config = BuildNodeConfig();
CalculatorGraph graph = BuildCalculatorGraph(node_config);
std::vector<Packet> output_frames;
MP_ASSERT_OK(graph.ObserveOutputStream("frame", [&](const Packet& p) {
output_frames.push_back(p);
return absl::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
// Finalize input for the "mask" input stream.
MP_ASSERT_OK(graph.CloseInputStream("mask_0"));
MP_ASSERT_OK(graph.CloseInputStream("mask_1"));
MP_ASSERT_OK(graph.CloseInputStream("mask_2"));
// All channels are specified before any frame packets arrive.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 1)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 10)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 11)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 21)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
// The packet at timestamp 10 is passed from channel 0.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t10", 10)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
EXPECT_EQ(output_frames[0].Get<std::string>(), "p0_t10");
// The packet at timestamp 20 is passed from channel 1.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t20", 20)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 2);
EXPECT_EQ(output_frames[1].Get<std::string>(), "p1_t20");
// The packet at timestamp 30 is passed from channel 0.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t30", 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 3);
EXPECT_EQ(output_frames[2].Get<std::string>(), "p0_t30");
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
// Shows that packets are queued until channels are specified.
TEST_F(SwitchMuxCalculatorTest, ChannelsLate) {
CalculatorGraphConfig::Node node_config = BuildNodeConfig();
CalculatorGraph graph = BuildCalculatorGraph(node_config);
std::vector<Packet> output_frames;
MP_ASSERT_OK(graph.ObserveOutputStream("frame", [&](const Packet& p) {
output_frames.push_back(p);
return absl::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
// Finalize input for the "mask" input stream.
MP_ASSERT_OK(graph.CloseInputStream("mask_0"));
MP_ASSERT_OK(graph.CloseInputStream("mask_1"));
MP_ASSERT_OK(graph.CloseInputStream("mask_2"));
// All frame packets arrive before any channels are specified.
// All packets are queued awaiting channel choices.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t10", 10)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t20", 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t30", 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 0);
// The packet at timestamp 10 is released from channel 0.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 1)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 10)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
EXPECT_EQ(output_frames[0].Get<std::string>(), "p0_t10");
// The packet at timestamp 20 is released from channel 1.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 11)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 20)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 2);
EXPECT_EQ(output_frames[1].Get<std::string>(), "p1_t20");
// The packet at timestamp 30 is released from channel 0.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 21)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 3);
EXPECT_EQ(output_frames[2].Get<std::string>(), "p0_t30");
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
// Shows that no packets are queued when channels are specified just in time.
TEST_F(SwitchMuxCalculatorTest, ChannelsOnTime) {
CalculatorGraphConfig::Node node_config = BuildNodeConfig();
CalculatorGraph graph = BuildCalculatorGraph(node_config);
std::vector<Packet> output_frames;
MP_ASSERT_OK(graph.ObserveOutputStream("frame", [&](const Packet& p) {
output_frames.push_back(p);
return absl::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
// Finalize input for the "mask" input stream.
MP_ASSERT_OK(graph.CloseInputStream("mask_0"));
MP_ASSERT_OK(graph.CloseInputStream("mask_1"));
MP_ASSERT_OK(graph.CloseInputStream("mask_2"));
// Channel 0 is selected just before corresponding packets arrive.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 1)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 10)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t10", 10)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
EXPECT_EQ(output_frames[0].Get<std::string>(), "p0_t10");
// Channel 1 is selected just before corresponding packets arrive.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 11)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t20", 20)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 2);
EXPECT_EQ(output_frames[1].Get<std::string>(), "p1_t20");
// Channel 0 is selected just before corresponding packets arrive.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 21)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 30)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t30", 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 3);
EXPECT_EQ(output_frames[2].Get<std::string>(), "p0_t30");
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
// Shows what happens if the last packet from a channel never arrives.
TEST_F(SwitchMuxCalculatorTest, ChannelNeverCompletes) {
CalculatorGraphConfig::Node node_config = BuildNodeConfig();
CalculatorGraph graph = BuildCalculatorGraph(node_config);
std::vector<Packet> output_frames;
MP_ASSERT_OK(graph.ObserveOutputStream("frame", [&](const Packet& p) {
output_frames.push_back(p);
return absl::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
// Finalize input for the "mask" input stream.
MP_ASSERT_OK(graph.CloseInputStream("mask_0"));
MP_ASSERT_OK(graph.CloseInputStream("mask_1"));
MP_ASSERT_OK(graph.CloseInputStream("mask_2"));
// Channel 0 is selected, but it's closing packet never arrives.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 1)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 10)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 0);
// Channel 1 is selected, but we still wait for channel 0 to finish.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 11)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t20", 20)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 0);
// Finally channel 0 advances, and channel 1 can be delivered.
// Note that "p0_t15" is discarded because its channel is deselected.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t15", 15)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
EXPECT_EQ(output_frames[0].Get<std::string>(), "p1_t20");
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
// Shows how output is sequenced when one channel is delayed substantially.
// Queues up "SELECT" packets for channel 0, 1, and 2.
// Queues up "frame" packets for channel 0 and 2.
// The output packets from channel 1, 2, and 0 wait for channel 1.
TEST_F(SwitchMuxCalculatorTest, OneChannelIsSlow) {
CalculatorGraphConfig::Node node_config = BuildNodeConfig();
CalculatorGraph graph = BuildCalculatorGraph(node_config);
std::vector<Packet> output_frames;
MP_ASSERT_OK(graph.ObserveOutputStream("frame", [&](const Packet& p) {
output_frames.push_back(p);
return absl::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
// Finalize input for the "mask" input stream.
MP_ASSERT_OK(graph.CloseInputStream("mask_0"));
MP_ASSERT_OK(graph.CloseInputStream("mask_1"));
MP_ASSERT_OK(graph.CloseInputStream("mask_2"));
// Channel 0 is selected, and some packets arrive.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 1)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 10)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t20", 10)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
// Channel 1 is selected, but its packets are delayed.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 11)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t20", 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_2", pack("p2_t20", 20)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
// Channel 2 is selected, packets arrive, but wait for channel 1.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(2, 21)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(2, 30)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_2", pack("p2_t30", 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
// Channel 0 is selected again, packets arrive, but wait for channel 1.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 31)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 40)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t40", 40)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
// Channel 1 is selected again, but its packets are still delayed.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 41)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 50)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
// Finally, the delayed packets from channel 1 arrive.
// And all packets for all five "SELECT"" inetervals are delivered.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t20", 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t50", 50)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 5);
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
} // namespace
} // namespace mediapipe
@@ -594,7 +594,7 @@ class TemplateExpanderImpl {
return AsArgument(static_cast<double>(b));
}
// Convert between a proto feild value and a template argument.
// Convert between a proto field value and a template argument.
absl::Status AsFieldValues(const std::vector<TemplateArgument>& args,
FieldType field_type,
std::vector<FieldValue>* result) {
-1
View File
@@ -34,7 +34,6 @@
#include "mediapipe/framework/formats/image_format.pb.h"
#include "mediapipe/framework/port/advanced_proto_inc.h"
#include "mediapipe/framework/port/file_helpers.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/proto_ns.h"
#include "mediapipe/framework/port/ret_check.h"