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

Project import generated by Copybara.

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/formats/BUILD
+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",
],
mediapipe/framework/formats/affine_transform_data.proto
-2
View File
@@ -16,8 +16,6 @@ syntax = "proto2";
package mediapipe;
option objc_class_prefix = "MediaPipe";
// Proto for serializing Vector2 data
message Vector2Data {
optional float x = 1;
mediapipe/framework/formats/annotation/locus.proto
+2
View File
@@ -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
mediapipe/framework/formats/classification.proto
-1
View File
@@ -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";
mediapipe/framework/formats/image.cc
+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
mediapipe/framework/formats/image_format.proto
+3
View File
@@ -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
mediapipe/framework/formats/location.cc
-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
mediapipe/framework/formats/location.h
+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;
mediapipe/framework/formats/location_opencv.cc
+220
View File
@@ -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
mediapipe/framework/formats/location_opencv.h
+54
View File
@@ -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_
mediapipe/framework/formats/location_opencv_test.cc
+167
View File
@@ -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
mediapipe/framework/formats/motion/BUILD
+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",
mediapipe/framework/formats/motion/optical_flow_field.cc
+2 -1
View File
@@ -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
mediapipe/framework/formats/motion/optical_flow_field_test.cc
+3 -2
View File
@@ -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);
mediapipe/framework/formats/tensor.cc
+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
}
mediapipe/framework/formats/tensor.h
+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);
mediapipe/framework/formats/tensor_ahwb.cc
+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;
}
mediapipe/framework/formats/tensor_ahwb_test.cc
+2 -5
View File
@@ -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
mediapipe/framework/formats/time_series_header.proto
-2
View File
@@ -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.