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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 796203faee20d7aae2876aac8ca5a1827dee4fe3
This commit is contained in:
MediaPipe Team
2019-09-30 11:26:36 -07:00
committed by jqtang
parent 412ab42d1f
commit a2a63e3876
122 changed files with 7330 additions and 2016 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/calculators/audio/audio_decoder_calculator.cc
+6 -2
View File
@@ -61,7 +61,9 @@ class AudioDecoderCalculator : public CalculatorBase {
::mediapipe::Status AudioDecoderCalculator::GetContract(
CalculatorContract* cc) {
cc->InputSidePackets().Tag("INPUT_FILE_PATH").Set<std::string>();
if (cc->InputSidePackets().HasTag("OPTIONS")) {
cc->InputSidePackets().Tag("OPTIONS").Set<mediapipe::AudioDecoderOptions>();
}
cc->Outputs().Tag("AUDIO").Set<Matrix>();
if (cc->Outputs().HasTag("AUDIO_HEADER")) {
cc->Outputs().Tag("AUDIO_HEADER").SetNone();
@@ -72,7 +74,9 @@ class AudioDecoderCalculator : public CalculatorBase {
::mediapipe::Status AudioDecoderCalculator::Open(CalculatorContext* cc) {
const std::string& input_file_path =
cc->InputSidePackets().Tag("INPUT_FILE_PATH").Get<std::string>();
const auto& decoder_options = cc->Options<mediapipe::AudioDecoderOptions>();
const auto& decoder_options =
tool::RetrieveOptions(cc->Options<mediapipe::AudioDecoderOptions>(),
cc->InputSidePackets(), "OPTIONS");
decoder_ = absl::make_unique<AudioDecoder>();
MP_RETURN_IF_ERROR(decoder_->Initialize(input_file_path, decoder_options));
std::unique_ptr<mediapipe::TimeSeriesHeader> header =
mediapipe/calculators/audio/stabilized_log_calculator.cc
+6 -1
View File
@@ -75,8 +75,13 @@ class StabilizedLogCalculator : public CalculatorBase {
::mediapipe::Status Process(CalculatorContext* cc) override {
auto input_matrix = cc->Inputs().Index(0).Get<Matrix>();
if (input_matrix.array().isNaN().any()) {
return ::mediapipe::InvalidArgumentError("NaN input to log operation.");
}
if (check_nonnegativity_) {
CHECK_GE(input_matrix.minCoeff(), 0);
if (input_matrix.minCoeff() < 0.0) {
return ::mediapipe::OutOfRangeError("Negative input to log operation.");
}
}
std::unique_ptr<Matrix> output_frame(new Matrix(
output_scale_ * (input_matrix.array() + stabilizer_).log().matrix()));
mediapipe/calculators/audio/stabilized_log_calculator_test.cc
+12 -2
View File
@@ -11,6 +11,7 @@
// 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 <cmath>
#include "Eigen/Core"
#include "mediapipe/calculators/audio/stabilized_log_calculator.pb.h"
@@ -108,13 +109,22 @@ TEST_F(StabilizedLogCalculatorTest, ZerosAreStabilized) {
runner_->Outputs().Index(0).packets[0].Get<Matrix>());
}
TEST_F(StabilizedLogCalculatorTest, NegativeValuesCheckFail) {
TEST_F(StabilizedLogCalculatorTest, NanValuesReturnError) {
InitializeGraph();
FillInputHeader();
AppendInputPacket(
new Matrix(Matrix::Constant(kNumChannels, kNumSamples, std::nanf(""))),
0 /* timestamp */);
ASSERT_FALSE(RunGraph().ok());
}
TEST_F(StabilizedLogCalculatorTest, NegativeValuesReturnError) {
InitializeGraph();
FillInputHeader();
AppendInputPacket(
new Matrix(Matrix::Constant(kNumChannels, kNumSamples, -1.0)),
0 /* timestamp */);
ASSERT_DEATH(RunGraphNoReturn(), "");
ASSERT_FALSE(RunGraph().ok());
}
TEST_F(StabilizedLogCalculatorTest, NegativeValuesDoNotCheckFailIfCheckIsOff) {
mediapipe/calculators/audio/time_series_framer_calculator.cc
+40 -6
View File
@@ -56,6 +56,14 @@ namespace mediapipe {
// If pad_final_packet is true, all input samples will be emitted and the final
// packet will be zero padded as necessary. If pad_final_packet is false, some
// samples may be dropped at the end of the stream.
//
// If use_local_timestamp is true, the output packet's timestamp is based on the
// last sample of the packet. The timestamp of this sample is inferred by
// input_packet_timesamp + local_sample_index / sampling_rate_. If false, the
// output packet's timestamp is based on the cumulative timestamping, which is
// done by adopting the timestamp of the first sample of the packet and this
// sample's timestamp is inferred by initial_input_timestamp_ +
// cumulative_completed_samples / sample_rate_.
class TimeSeriesFramerCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
@@ -86,11 +94,26 @@ class TimeSeriesFramerCalculator : public CalculatorBase {
void FrameOutput(CalculatorContext* cc);
Timestamp CurrentOutputTimestamp() {
if (use_local_timestamp_) {
return current_timestamp_;
}
return CumulativeOutputTimestamp();
}
Timestamp CumulativeOutputTimestamp() {
return initial_input_timestamp_ +
round(cumulative_completed_samples_ / sample_rate_ *
Timestamp::kTimestampUnitsPerSecond);
}
// Returns the timestamp of a sample on a base, which is usually the time
// stamp of a packet.
Timestamp CurrentSampleTimestamp(const Timestamp& timestamp_base,
int64 number_of_samples) {
return timestamp_base + round(number_of_samples / sample_rate_ *
Timestamp::kTimestampUnitsPerSecond);
}
// The number of input samples to advance after the current output frame is
// emitted.
int next_frame_step_samples() const {
@@ -118,14 +141,18 @@ class TimeSeriesFramerCalculator : public CalculatorBase {
// any overlap).
int64 cumulative_completed_samples_;
Timestamp initial_input_timestamp_;
// The current timestamp is updated along with the incoming packets.
Timestamp current_timestamp_;
int num_channels_;
// Each entry in this deque consists of a single sample, i.e. a
// single column vector.
std::deque<Matrix> sample_buffer_;
// single column vector, and its timestamp.
std::deque<std::pair<Matrix, Timestamp>> sample_buffer_;
bool use_window_;
Matrix window_;
bool use_local_timestamp_;
};
REGISTER_CALCULATOR(TimeSeriesFramerCalculator);
@@ -133,7 +160,8 @@ void TimeSeriesFramerCalculator::EnqueueInput(CalculatorContext* cc) {
const Matrix& input_frame = cc->Inputs().Index(0).Get<Matrix>();
for (int i = 0; i < input_frame.cols(); ++i) {
sample_buffer_.emplace_back(input_frame.col(i));
sample_buffer_.emplace_back(std::make_pair(
input_frame.col(i), CurrentSampleTimestamp(cc->InputTimestamp(), i)));
}
cumulative_input_samples_ += input_frame.cols();
@@ -151,14 +179,16 @@ void TimeSeriesFramerCalculator::FrameOutput(CalculatorContext* cc) {
new Matrix(num_channels_, frame_duration_samples_));
for (int i = 0; i < std::min(frame_step_samples, frame_duration_samples_);
++i) {
output_frame->col(i) = sample_buffer_.front();
output_frame->col(i) = sample_buffer_.front().first;
current_timestamp_ = sample_buffer_.front().second;
sample_buffer_.pop_front();
}
const int frame_overlap_samples =
frame_duration_samples_ - frame_step_samples;
if (frame_overlap_samples > 0) {
for (int i = 0; i < frame_overlap_samples; ++i) {
output_frame->col(i + frame_step_samples) = sample_buffer_[i];
output_frame->col(i + frame_step_samples) = sample_buffer_[i].first;
current_timestamp_ = sample_buffer_[i].second;
}
} else {
samples_still_to_drop_ = -frame_overlap_samples;
@@ -178,6 +208,7 @@ void TimeSeriesFramerCalculator::FrameOutput(CalculatorContext* cc) {
::mediapipe::Status TimeSeriesFramerCalculator::Process(CalculatorContext* cc) {
if (initial_input_timestamp_ == Timestamp::Unstarted()) {
initial_input_timestamp_ = cc->InputTimestamp();
current_timestamp_ = initial_input_timestamp_;
}
EnqueueInput(cc);
@@ -195,7 +226,8 @@ void TimeSeriesFramerCalculator::FrameOutput(CalculatorContext* cc) {
std::unique_ptr<Matrix> output_frame(new Matrix);
output_frame->setZero(num_channels_, frame_duration_samples_);
for (int i = 0; i < sample_buffer_.size(); ++i) {
output_frame->col(i) = sample_buffer_[i];
output_frame->col(i) = sample_buffer_[i].first;
current_timestamp_ = sample_buffer_[i].second;
}
cc->Outputs().Index(0).Add(output_frame.release(),
@@ -258,6 +290,7 @@ void TimeSeriesFramerCalculator::FrameOutput(CalculatorContext* cc) {
cumulative_output_frames_ = 0;
samples_still_to_drop_ = 0;
initial_input_timestamp_ = Timestamp::Unstarted();
current_timestamp_ = Timestamp::Unstarted();
std::vector<double> window_vector;
use_window_ = false;
@@ -282,6 +315,7 @@ void TimeSeriesFramerCalculator::FrameOutput(CalculatorContext* cc) {
frame_duration_samples_)
.cast<float>();
}
use_local_timestamp_ = framer_options.use_local_timestamp();
return ::mediapipe::OkStatus();
}
mediapipe/calculators/audio/time_series_framer_calculator.proto
+7
View File
@@ -62,4 +62,11 @@ message TimeSeriesFramerCalculatorOptions {
HANN = 2;
}
optional WindowFunction window_function = 4 [default = NONE];
// If use_local_timestamp is true, the output packet's timestamp is based on
// the last sample of the packet and it's inferred from the latest input
// packet's timestamp. If false, the output packet's timestamp is based on
// the cumulative timestamping, which is inferred from the intial input
// timestamp and the cumulative number of samples.
optional bool use_local_timestamp = 6 [default = false];
}
mediapipe/calculators/audio/time_series_framer_calculator_test.cc
+91 -1
View File
@@ -35,6 +35,8 @@ namespace mediapipe {
namespace {
const int kInitialTimestampOffsetMicroseconds = 4;
const int kGapBetweenPacketsInSeconds = 1;
const int kUniversalInputPacketSize = 50;
class TimeSeriesFramerCalculatorTest
: public TimeSeriesCalculatorTest<TimeSeriesFramerCalculatorOptions> {
@@ -391,5 +393,93 @@ TEST_F(TimeSeriesFramerCalculatorWindowingSanityTest, HannWindowSanityCheck) {
RunAndTestSinglePacketAverage(0.5f);
}
} // anonymous namespace
// A simple test class that checks the local packet time stamp. This class
// generate a series of packets with and without gaps between packets and tests
// the behavior with cumulative timestamping and local packet timestamping.
class TimeSeriesFramerCalculatorTimestampingTest
: public TimeSeriesFramerCalculatorTest {
protected:
// Creates test input and saves a reference copy.
void InitializeInputForTimeStampingTest() {
concatenated_input_samples_.resize(0, num_input_channels_);
num_input_samples_ = 0;
for (int i = 0; i < 10; ++i) {
// This range of packet sizes was chosen such that some input
// packets will be smaller than the output packet size and other
// input packets will be larger.
int packet_size = kUniversalInputPacketSize;
double timestamp_seconds = kInitialTimestampOffsetMicroseconds * 1.0e-6 +
num_input_samples_ / input_sample_rate_;
if (options_.use_local_timestamp()) {
timestamp_seconds += kGapBetweenPacketsInSeconds * i;
}
Matrix* data_frame =
NewTestFrame(num_input_channels_, packet_size, timestamp_seconds);
AppendInputPacket(data_frame, round(timestamp_seconds *
Timestamp::kTimestampUnitsPerSecond));
num_input_samples_ += packet_size;
}
}
void CheckOutputTimestamps() {
int num_full_packets = output().packets.size();
if (options_.pad_final_packet()) {
num_full_packets -= 1;
}
int64 num_samples = 0;
for (int packet_num = 0; packet_num < num_full_packets; ++packet_num) {
const Packet& packet = output().packets[packet_num];
num_samples += FrameDurationSamples();
double expected_timestamp =
options_.use_local_timestamp()
? GetExpectedLocalTimestampForSample(num_samples - 1)
: GetExpectedCumulativeTimestamp(num_samples - 1);
ASSERT_NEAR(packet.Timestamp().Seconds(), expected_timestamp, 1e-10);
}
}
::mediapipe::Status RunTimestampTest() {
InitializeGraph();
InitializeInputForTimeStampingTest();
FillInputHeader();
return RunGraph();
}
private:
// Returns the timestamp in seconds based on local timestamping.
double GetExpectedLocalTimestampForSample(int sample_index) {
return kInitialTimestampOffsetMicroseconds * 1.0e-6 +
sample_index / input_sample_rate_ +
(sample_index / kUniversalInputPacketSize) *
kGapBetweenPacketsInSeconds;
}
// Returns the timestamp inseconds based on cumulative timestamping.
double GetExpectedCumulativeTimestamp(int sample_index) {
return kInitialTimestampOffsetMicroseconds * 1.0e-6 +
sample_index / FrameDurationSamples() * FrameDurationSamples() /
input_sample_rate_;
}
};
TEST_F(TimeSeriesFramerCalculatorTimestampingTest, UseLocalTimeStamp) {
options_.set_frame_duration_seconds(100.0 / input_sample_rate_);
options_.set_use_local_timestamp(true);
MP_ASSERT_OK(RunTimestampTest());
CheckOutputTimestamps();
}
TEST_F(TimeSeriesFramerCalculatorTimestampingTest, UseCumulativeTimeStamp) {
options_.set_frame_duration_seconds(100.0 / input_sample_rate_);
options_.set_use_local_timestamp(false);
MP_ASSERT_OK(RunTimestampTest());
CheckOutputTimestamps();
}
} // namespace
} // namespace mediapipe
mediapipe/calculators/core/BUILD
+7 -1
View File
@@ -166,7 +166,13 @@ cc_library(
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"@org_tensorflow//tensorflow/lite:framework",
],
] + select({
"//mediapipe/gpu:disable_gpu": [],
"//mediapipe:ios": [],
"//conditions:default": [
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
],
}),
alwayslink = 1,
)
mediapipe/calculators/core/concatenate_vector_calculator.cc
+11
View File
@@ -19,6 +19,10 @@
#include "mediapipe/framework/formats/landmark.pb.h"
#include "tensorflow/lite/interpreter.h"
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
#endif // !MEDIAPIPE_DISABLE_GPU
namespace mediapipe {
// Example config:
@@ -45,4 +49,11 @@ REGISTER_CALCULATOR(ConcatenateTfLiteTensorVectorCalculator);
typedef ConcatenateVectorCalculator<::mediapipe::NormalizedLandmark>
ConcatenateLandmarkVectorCalculator;
REGISTER_CALCULATOR(ConcatenateLandmarkVectorCalculator);
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
typedef ConcatenateVectorCalculator<::tflite::gpu::gl::GlBuffer>
ConcatenateGlBufferVectorCalculator;
REGISTER_CALCULATOR(ConcatenateGlBufferVectorCalculator);
#endif
} // namespace mediapipe
mediapipe/calculators/core/concatenate_vector_calculator.h
+45 -2
View File
@@ -15,6 +15,7 @@
#ifndef MEDIAPIPE_CALCULATORS_CORE_CONCATENATE_VECTOR_CALCULATOR_H_
#define MEDIAPIPE_CALCULATORS_CORE_CONCATENATE_VECTOR_CALCULATOR_H_
#include <type_traits>
#include <vector>
#include "mediapipe/calculators/core/concatenate_vector_calculator.pb.h"
@@ -59,16 +60,58 @@ class ConcatenateVectorCalculator : public CalculatorBase {
if (cc->Inputs().Index(i).IsEmpty()) return ::mediapipe::OkStatus();
}
}
auto output = absl::make_unique<std::vector<T>>();
return ConcatenateVectors<T>(std::is_copy_constructible<T>(), cc);
}
template <typename U>
::mediapipe::Status ConcatenateVectors(std::true_type,
CalculatorContext* cc) {
auto output = absl::make_unique<std::vector<U>>();
for (int i = 0; i < cc->Inputs().NumEntries(); ++i) {
if (cc->Inputs().Index(i).IsEmpty()) continue;
const std::vector<T>& input = cc->Inputs().Index(i).Get<std::vector<T>>();
const std::vector<U>& input = cc->Inputs().Index(i).Get<std::vector<U>>();
output->insert(output->end(), input.begin(), input.end());
}
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
return ::mediapipe::OkStatus();
}
template <typename U>
::mediapipe::Status ConcatenateVectors(std::false_type,
CalculatorContext* cc) {
return ConsumeAndConcatenateVectors<T>(std::is_move_constructible<U>(), cc);
}
template <typename U>
::mediapipe::Status ConsumeAndConcatenateVectors(std::true_type,
CalculatorContext* cc) {
auto output = absl::make_unique<std::vector<U>>();
for (int i = 0; i < cc->Inputs().NumEntries(); ++i) {
if (cc->Inputs().Index(i).IsEmpty()) continue;
::mediapipe::StatusOr<std::unique_ptr<std::vector<U>>> input_status =
cc->Inputs().Index(i).Value().Consume<std::vector<U>>();
if (input_status.ok()) {
std::unique_ptr<std::vector<U>> input_vector =
std::move(input_status).ValueOrDie();
output->insert(output->end(),
std::make_move_iterator(input_vector->begin()),
std::make_move_iterator(input_vector->end()));
} else {
return input_status.status();
}
}
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
return ::mediapipe::OkStatus();
}
template <typename U>
::mediapipe::Status ConsumeAndConcatenateVectors(std::false_type,
CalculatorContext* cc) {
return ::mediapipe::InternalError(
"Cannot copy or move input vectors to concatenate them");
}
private:
bool only_emit_if_all_present_;
};
mediapipe/calculators/core/concatenate_vector_calculator_test.cc
+163
View File
@@ -235,4 +235,167 @@ TEST(ConcatenateFloatVectorCalculatorTest, OneEmptyStreamNoOutput) {
EXPECT_EQ(0, outputs.size());
}
typedef ConcatenateVectorCalculator<std::unique_ptr<int>>
TestConcatenateUniqueIntPtrCalculator;
REGISTER_CALCULATOR(TestConcatenateUniqueIntPtrCalculator);
TEST(TestConcatenateUniqueIntVectorCalculatorTest, ConsumeOneTimestamp) {
/* Note: We don't use CalculatorRunner for this test because it keeps copies
* of input packets, so packets sent to the graph don't have sole ownership.
* The test needs to send packets that own the data.
*/
CalculatorGraphConfig graph_config =
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "in_1"
input_stream: "in_2"
input_stream: "in_3"
node {
calculator: "TestConcatenateUniqueIntPtrCalculator"
input_stream: "in_1"
input_stream: "in_2"
input_stream: "in_3"
output_stream: "out"
}
)");
std::vector<Packet> outputs;
tool::AddVectorSink("out", &graph_config, &outputs);
CalculatorGraph graph;
MP_EXPECT_OK(graph.Initialize(graph_config));
MP_EXPECT_OK(graph.StartRun({}));
// input1 : {0, 1, 2}
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_1 =
absl::make_unique<std::vector<std::unique_ptr<int>>>(3);
for (int i = 0; i < 3; ++i) {
input_1->at(i) = absl::make_unique<int>(i);
}
// input2: {3}
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_2 =
absl::make_unique<std::vector<std::unique_ptr<int>>>(1);
input_2->at(0) = absl::make_unique<int>(3);
// input3: {4, 5}
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_3 =
absl::make_unique<std::vector<std::unique_ptr<int>>>(2);
input_3->at(0) = absl::make_unique<int>(4);
input_3->at(1) = absl::make_unique<int>(5);
MP_EXPECT_OK(graph.AddPacketToInputStream(
"in_1", Adopt(input_1.release()).At(Timestamp(1))));
MP_EXPECT_OK(graph.AddPacketToInputStream(
"in_2", Adopt(input_2.release()).At(Timestamp(1))));
MP_EXPECT_OK(graph.AddPacketToInputStream(
"in_3", Adopt(input_3.release()).At(Timestamp(1))));
MP_EXPECT_OK(graph.WaitUntilIdle());
MP_EXPECT_OK(graph.CloseAllPacketSources());
MP_EXPECT_OK(graph.WaitUntilDone());
EXPECT_EQ(1, outputs.size());
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
const std::vector<std::unique_ptr<int>>& result =
outputs[0].Get<std::vector<std::unique_ptr<int>>>();
EXPECT_EQ(6, result.size());
for (int i = 0; i < 6; ++i) {
const std::unique_ptr<int>& v = result[i];
EXPECT_EQ(i, *v);
}
}
TEST(TestConcatenateUniqueIntVectorCalculatorTest, OneEmptyStreamStillOutput) {
/* Note: We don't use CalculatorRunner for this test because it keeps copies
* of input packets, so packets sent to the graph don't have sole ownership.
* The test needs to send packets that own the data.
*/
CalculatorGraphConfig graph_config =
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "in_1"
input_stream: "in_2"
node {
calculator: "TestConcatenateUniqueIntPtrCalculator"
input_stream: "in_1"
input_stream: "in_2"
output_stream: "out"
}
)");
std::vector<Packet> outputs;
tool::AddVectorSink("out", &graph_config, &outputs);
CalculatorGraph graph;
MP_EXPECT_OK(graph.Initialize(graph_config));
MP_EXPECT_OK(graph.StartRun({}));
// input1 : {0, 1, 2}
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_1 =
absl::make_unique<std::vector<std::unique_ptr<int>>>(3);
for (int i = 0; i < 3; ++i) {
input_1->at(i) = absl::make_unique<int>(i);
}
MP_EXPECT_OK(graph.AddPacketToInputStream(
"in_1", Adopt(input_1.release()).At(Timestamp(1))));
MP_EXPECT_OK(graph.WaitUntilIdle());
MP_EXPECT_OK(graph.CloseAllPacketSources());
MP_EXPECT_OK(graph.WaitUntilDone());
EXPECT_EQ(1, outputs.size());
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
const std::vector<std::unique_ptr<int>>& result =
outputs[0].Get<std::vector<std::unique_ptr<int>>>();
EXPECT_EQ(3, result.size());
for (int i = 0; i < 3; ++i) {
const std::unique_ptr<int>& v = result[i];
EXPECT_EQ(i, *v);
}
}
TEST(TestConcatenateUniqueIntVectorCalculatorTest, OneEmptyStreamNoOutput) {
/* Note: We don't use CalculatorRunner for this test because it keeps copies
* of input packets, so packets sent to the graph don't have sole ownership.
* The test needs to send packets that own the data.
*/
CalculatorGraphConfig graph_config =
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "in_1"
input_stream: "in_2"
node {
calculator: "TestConcatenateUniqueIntPtrCalculator"
input_stream: "in_1"
input_stream: "in_2"
output_stream: "out"
options {
[mediapipe.ConcatenateVectorCalculatorOptions.ext] {
only_emit_if_all_present: true
}
}
}
)");
std::vector<Packet> outputs;
tool::AddVectorSink("out", &graph_config, &outputs);
CalculatorGraph graph;
MP_EXPECT_OK(graph.Initialize(graph_config));
MP_EXPECT_OK(graph.StartRun({}));
// input1 : {0, 1, 2}
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_1 =
absl::make_unique<std::vector<std::unique_ptr<int>>>(3);
for (int i = 0; i < 3; ++i) {
input_1->at(i) = absl::make_unique<int>(i);
}
MP_EXPECT_OK(graph.AddPacketToInputStream(
"in_1", Adopt(input_1.release()).At(Timestamp(1))));
MP_EXPECT_OK(graph.WaitUntilIdle());
MP_EXPECT_OK(graph.CloseAllPacketSources());
MP_EXPECT_OK(graph.WaitUntilDone());
EXPECT_EQ(0, outputs.size());
}
} // namespace mediapipe
mediapipe/calculators/image/BUILD
+19 -31
View File
@@ -19,7 +19,6 @@ package(default_visibility = ["//visibility:private"])
exports_files(["LICENSE"])
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
load("@bazel_skylib//lib:selects.bzl", "selects")
proto_library(
name = "opencv_image_encoder_calculator_proto",
@@ -227,19 +226,13 @@ cc_library(
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:vector",
] + select({
"//mediapipe:android": [
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:gl_quad_renderer",
"//mediapipe/gpu:shader_util",
],
"//mediapipe:ios": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:shader_util",
],
"//conditions:default": [],
}),
alwayslink = 1,
)
@@ -263,13 +256,13 @@ cc_library(
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:vector",
] + select({
"//mediapipe:android": [
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:gl_quad_renderer",
"//mediapipe/gpu:shader_util",
],
"//conditions:default": [],
}),
alwayslink = 1,
)
@@ -322,14 +315,14 @@ cc_library(
"//mediapipe/framework/port:opencv_imgproc",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
] + selects.with_or({
("//mediapipe:android", "//mediapipe:ios"): [
] + select({
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gl_quad_renderer",
"//mediapipe/gpu:shader_util",
],
"//conditions:default": [],
}),
alwayslink = 1,
)
@@ -363,14 +356,15 @@ cc_library(
"//mediapipe/framework/port:opencv_imgproc",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
] + selects.with_or({
("//mediapipe:android", "//mediapipe:ios"): [
"//mediapipe/gpu:gpu_buffer",
] + select({
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gl_quad_renderer",
"//mediapipe/gpu:shader_util",
],
"//conditions:default": [],
}),
alwayslink = 1,
)
@@ -415,19 +409,13 @@ cc_library(
"//mediapipe/framework/port:ret_check",
"//mediapipe/util:color_cc_proto",
] + select({
"//mediapipe:android": [
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:gl_quad_renderer",
"//mediapipe/gpu:shader_util",
],
"//mediapipe:ios": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:shader_util",
],
"//conditions:default": [],
}),
alwayslink = 1,
)
@@ -486,11 +474,11 @@ cc_library(
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
] + selects.with_or({
("//mediapipe:android", "//mediapipe:ios"): [
] + select({
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": [
"//mediapipe/gpu:gpu_buffer",
],
"//conditions:default": [],
}),
alwayslink = 1,
)
mediapipe/calculators/image/bilateral_filter_calculator.cc
+35 -28
View File
@@ -27,11 +27,11 @@
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/vector.h"
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gl_simple_shaders.h"
#include "mediapipe/gpu/shader_util.h"
#endif // __ANDROID__ || __EMSCRIPTEN__
#endif // !MEDIAPIPE_DISABLE_GPU
namespace mediapipe {
@@ -101,11 +101,11 @@ class BilateralFilterCalculator : public CalculatorBase {
bool use_gpu_ = false;
bool gpu_initialized_ = false;
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
mediapipe::GlCalculatorHelper gpu_helper_;
GLuint program_ = 0;
GLuint program_joint_ = 0;
#endif // __ANDROID__ || __EMSCRIPTEN__
#endif // !MEDIAPIPE_DISABLE_GPU
};
REGISTER_CALCULATOR(BilateralFilterCalculator);
@@ -122,39 +122,46 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
return ::mediapipe::InternalError("GPU output must have GPU input.");
}
bool use_gpu = false;
// Input image to filter.
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag(kInputFrameTagGpu)) {
cc->Inputs().Tag(kInputFrameTagGpu).Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__ || __EMSCRIPTEN__
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Inputs().HasTag(kInputFrameTag)) {
cc->Inputs().Tag(kInputFrameTag).Set<ImageFrame>();
}
// Input guide image mask (optional)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag(kInputGuideTagGpu)) {
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
cc->Inputs().Tag(kInputGuideTagGpu).Set<mediapipe::GpuBuffer>();
#endif // __ANDROID__ || __EMSCRIPTEN__
use_gpu |= true;
}
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Inputs().HasTag(kInputGuideTag)) {
cc->Inputs().Tag(kInputGuideTag).Set<ImageFrame>();
}
// Output image.
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Outputs().HasTag(kOutputFrameTagGpu)) {
cc->Outputs().Tag(kOutputFrameTagGpu).Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__ || __EMSCRIPTEN__
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Outputs().HasTag(kOutputFrameTag)) {
cc->Outputs().Tag(kOutputFrameTag).Set<ImageFrame>();
}
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // __ANDROID__ || __EMSCRIPTEN__
if (use_gpu) {
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
}
@@ -166,11 +173,11 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
if (cc->Inputs().HasTag(kInputFrameTagGpu) &&
cc->Outputs().HasTag(kOutputFrameTagGpu)) {
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
use_gpu_ = true;
#else
RET_CHECK_FAIL() << "GPU processing on non-Android not supported yet.";
#endif // __ANDROID__ || __EMSCRIPTEN__
RET_CHECK_FAIL() << "GPU processing not enabled.";
#endif
}
sigma_color_ = options_.sigma_color();
@@ -180,9 +187,9 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
if (!use_gpu_) sigma_color_ *= 255.0;
if (use_gpu_) {
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#endif
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
@@ -190,7 +197,7 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
::mediapipe::Status BilateralFilterCalculator::Process(CalculatorContext* cc) {
if (use_gpu_) {
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
if (!gpu_initialized_) {
@@ -200,7 +207,7 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
MP_RETURN_IF_ERROR(RenderGpu(cc));
return ::mediapipe::OkStatus();
}));
#endif // __ANDROID__ || __EMSCRIPTEN__
#endif // !MEDIAPIPE_DISABLE_GPU
} else {
MP_RETURN_IF_ERROR(RenderCpu(cc));
}
@@ -209,14 +216,14 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
}
::mediapipe::Status BilateralFilterCalculator::Close(CalculatorContext* cc) {
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
gpu_helper_.RunInGlContext([this] {
if (program_) glDeleteProgram(program_);
program_ = 0;
if (program_joint_) glDeleteProgram(program_joint_);
program_joint_ = 0;
});
#endif // __ANDROID__ || __EMSCRIPTEN__
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
@@ -263,7 +270,7 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
if (cc->Inputs().Tag(kInputFrameTagGpu).IsEmpty()) {
return ::mediapipe::OkStatus();
}
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
const auto& input_frame =
cc->Inputs().Tag(kInputFrameTagGpu).Get<mediapipe::GpuBuffer>();
auto input_texture = gpu_helper_.CreateSourceTexture(input_frame);
@@ -321,13 +328,13 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
// Cleanup
input_texture.Release();
output_texture.Release();
#endif // __ANDROID__ || __EMSCRIPTEN__
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
void BilateralFilterCalculator::GlRender(CalculatorContext* cc) {
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
static const GLfloat square_vertices[] = {
-1.0f, -1.0f, // bottom left
1.0f, -1.0f, // bottom right
@@ -373,11 +380,11 @@ void BilateralFilterCalculator::GlRender(CalculatorContext* cc) {
glDeleteVertexArrays(1, &vao);
glDeleteBuffers(2, vbo);
#endif // __ANDROID__ || __EMSCRIPTEN__
#endif // !MEDIAPIPE_DISABLE_GPU
}
::mediapipe::Status BilateralFilterCalculator::GlSetup(CalculatorContext* cc) {
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
const GLint attr_location[NUM_ATTRIBUTES] = {
ATTRIB_VERTEX,
ATTRIB_TEXTURE_POSITION,
@@ -545,7 +552,7 @@ void BilateralFilterCalculator::GlRender(CalculatorContext* cc) {
glUniform1i(glGetUniformLocation(program_joint_, "input_frame"), 1);
glUniform1i(glGetUniformLocation(program_joint_, "guide_frame"), 2);
#endif // __ANDROID__ || __EMSCRIPTEN__
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
mediapipe/calculators/image/image_cropping_calculator.cc
+94 -63
View File
@@ -24,12 +24,12 @@
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gl_simple_shaders.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/gpu/shader_util.h"
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
namespace {
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
@@ -37,9 +37,20 @@ enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
namespace mediapipe {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
namespace {
#endif // __ANDROID__ or iOS
#if !defined(MEDIAPIPE_DISABLE_GPU)
#endif // !MEDIAPIPE_DISABLE_GPU
constexpr char kRectTag[] = "RECT";
constexpr char kNormRectTag[] = "NORM_RECT";
constexpr char kHeightTag[] = "HEIGHT";
constexpr char kImageTag[] = "IMAGE";
constexpr char kImageGpuTag[] = "IMAGE_GPU";
constexpr char kWidthTag[] = "WIDTH";
} // namespace
// Crops the input texture to the given rectangle region. The rectangle can
// be at arbitrary location on the image with rotation. If there's rotation, the
@@ -91,48 +102,55 @@ class ImageCroppingCalculator : public CalculatorBase {
bool use_gpu_ = false;
// Output texture corners (4) after transoformation in normalized coordinates.
float transformed_points_[8];
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
bool gpu_initialized_ = false;
mediapipe::GlCalculatorHelper gpu_helper_;
GLuint program_ = 0;
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
};
REGISTER_CALCULATOR(ImageCroppingCalculator);
::mediapipe::Status ImageCroppingCalculator::GetContract(
CalculatorContract* cc) {
RET_CHECK(cc->Inputs().HasTag("IMAGE") ^ cc->Inputs().HasTag("IMAGE_GPU"));
RET_CHECK(cc->Outputs().HasTag("IMAGE") ^ cc->Outputs().HasTag("IMAGE_GPU"));
RET_CHECK(cc->Inputs().HasTag(kImageTag) ^ cc->Inputs().HasTag(kImageGpuTag));
RET_CHECK(cc->Outputs().HasTag(kImageTag) ^
cc->Outputs().HasTag(kImageGpuTag));
if (cc->Inputs().HasTag("IMAGE")) {
RET_CHECK(cc->Outputs().HasTag("IMAGE"));
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
if (cc->Inputs().HasTag("IMAGE_GPU")) {
RET_CHECK(cc->Outputs().HasTag("IMAGE_GPU"));
cc->Inputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
cc->Outputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
}
#endif // __ANDROID__ or iOS
bool use_gpu = false;
if (cc->Inputs().HasTag("RECT")) {
cc->Inputs().Tag("RECT").Set<Rect>();
if (cc->Inputs().HasTag(kImageTag)) {
RET_CHECK(cc->Outputs().HasTag(kImageTag));
cc->Inputs().Tag(kImageTag).Set<ImageFrame>();
cc->Outputs().Tag(kImageTag).Set<ImageFrame>();
}
if (cc->Inputs().HasTag("NORM_RECT")) {
cc->Inputs().Tag("NORM_RECT").Set<NormalizedRect>();
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag(kImageGpuTag)) {
RET_CHECK(cc->Outputs().HasTag(kImageGpuTag));
cc->Inputs().Tag(kImageGpuTag).Set<GpuBuffer>();
cc->Outputs().Tag(kImageGpuTag).Set<GpuBuffer>();
use_gpu |= true;
}
if (cc->Inputs().HasTag("WIDTH")) {
cc->Inputs().Tag("WIDTH").Set<int>();
#endif // !MEDIAPIPE_DISABLE_GPU
RET_CHECK(cc->Inputs().HasTag(kRectTag) ^ cc->Inputs().HasTag(kNormRectTag));
if (cc->Inputs().HasTag(kRectTag)) {
cc->Inputs().Tag(kRectTag).Set<Rect>();
}
if (cc->Inputs().HasTag("HEIGHT")) {
cc->Inputs().Tag("HEIGHT").Set<int>();
if (cc->Inputs().HasTag(kNormRectTag)) {
cc->Inputs().Tag(kNormRectTag).Set<NormalizedRect>();
}
if (cc->Inputs().HasTag(kWidthTag)) {
cc->Inputs().Tag(kWidthTag).Set<int>();
}
if (cc->Inputs().HasTag(kHeightTag)) {
cc->Inputs().Tag(kHeightTag).Set<int>();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // __ANDROID__ or iOS
if (use_gpu) {
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
}
@@ -140,26 +158,35 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
::mediapipe::Status ImageCroppingCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
if (cc->Inputs().HasTag("IMAGE_GPU")) {
if (cc->Inputs().HasTag(kImageGpuTag)) {
use_gpu_ = true;
}
options_ = cc->Options<mediapipe::ImageCroppingCalculatorOptions>();
if (use_gpu_) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#else
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status ImageCroppingCalculator::Process(CalculatorContext* cc) {
if (cc->Inputs().HasTag(kRectTag) && cc->Inputs().Tag(kRectTag).IsEmpty()) {
VLOG(1) << "RECT is empty for timestamp: " << cc->InputTimestamp();
return ::mediapipe::OkStatus();
}
if (cc->Inputs().HasTag(kNormRectTag) &&
cc->Inputs().Tag(kNormRectTag).IsEmpty()) {
VLOG(1) << "NORM_RECT is empty for timestamp: " << cc->InputTimestamp();
return ::mediapipe::OkStatus();
}
if (use_gpu_) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
if (!gpu_initialized_) {
@@ -169,7 +196,7 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
MP_RETURN_IF_ERROR(RenderGpu(cc));
return ::mediapipe::OkStatus();
}));
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
} else {
MP_RETURN_IF_ERROR(RenderCpu(cc));
}
@@ -177,19 +204,22 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
}
::mediapipe::Status ImageCroppingCalculator::Close(CalculatorContext* cc) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
gpu_helper_.RunInGlContext([this] {
if (program_) glDeleteProgram(program_);
program_ = 0;
});
gpu_initialized_ = false;
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
::mediapipe::Status ImageCroppingCalculator::RenderCpu(CalculatorContext* cc) {
const auto& input_img = cc->Inputs().Tag("IMAGE").Get<ImageFrame>();
if (cc->Inputs().Tag(kImageTag).IsEmpty()) {
return ::mediapipe::OkStatus();
}
const auto& input_img = cc->Inputs().Tag(kImageTag).Get<ImageFrame>();
cv::Mat input_mat = formats::MatView(&input_img);
float rect_center_x = input_img.Width() / 2.0f;
@@ -197,8 +227,8 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
float rotation = 0.0f;
int target_width = input_img.Width();
int target_height = input_img.Height();
if (cc->Inputs().HasTag("RECT")) {
const auto& rect = cc->Inputs().Tag("RECT").Get<Rect>();
if (cc->Inputs().HasTag(kRectTag)) {
const auto& rect = cc->Inputs().Tag(kRectTag).Get<Rect>();
if (rect.width() > 0 && rect.height() > 0 && rect.x_center() >= 0 &&
rect.y_center() >= 0) {
rect_center_x = rect.x_center();
@@ -207,8 +237,8 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
target_height = rect.height();
rotation = rect.rotation();
}
} else if (cc->Inputs().HasTag("NORM_RECT")) {
const auto& rect = cc->Inputs().Tag("NORM_RECT").Get<NormalizedRect>();
} else if (cc->Inputs().HasTag(kNormRectTag)) {
const auto& rect = cc->Inputs().Tag(kNormRectTag).Get<NormalizedRect>();
if (rect.width() > 0.0 && rect.height() > 0.0 && rect.x_center() >= 0.0 &&
rect.y_center() >= 0.0) {
rect_center_x = std::round(rect.x_center() * input_img.Width());
@@ -218,9 +248,9 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
rotation = rect.rotation();
}
} else {
if (cc->Inputs().HasTag("WIDTH") && cc->Inputs().HasTag("HEIGHT")) {
target_width = cc->Inputs().Tag("WIDTH").Get<int>();
target_height = cc->Inputs().Tag("HEIGHT").Get<int>();
if (cc->Inputs().HasTag(kWidthTag) && cc->Inputs().HasTag(kHeightTag)) {
target_width = cc->Inputs().Tag(kWidthTag).Get<int>();
target_height = cc->Inputs().Tag(kHeightTag).Get<int>();
} else if (options_.has_width() && options_.has_height()) {
target_width = options_.width();
target_height = options_.height();
@@ -253,16 +283,17 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
input_img.Format(), cropped_image.cols, cropped_image.rows));
cv::Mat output_mat = formats::MatView(output_frame.get());
cropped_image.copyTo(output_mat);
cc->Outputs().Tag("IMAGE").Add(output_frame.release(), cc->InputTimestamp());
cc->Outputs().Tag(kImageTag).Add(output_frame.release(),
cc->InputTimestamp());
return ::mediapipe::OkStatus();
}
::mediapipe::Status ImageCroppingCalculator::RenderGpu(CalculatorContext* cc) {
if (cc->Inputs().Tag("IMAGE_GPU").IsEmpty()) {
if (cc->Inputs().Tag(kImageGpuTag).IsEmpty()) {
return ::mediapipe::OkStatus();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
const Packet& input_packet = cc->Inputs().Tag("IMAGE_GPU").Value();
#if !defined(MEDIAPIPE_DISABLE_GPU)
const Packet& input_packet = cc->Inputs().Tag(kImageGpuTag).Value();
const auto& input_buffer = input_packet.Get<mediapipe::GpuBuffer>();
auto src_tex = gpu_helper_.CreateSourceTexture(input_buffer);
@@ -287,18 +318,18 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
// Send result image in GPU packet.
auto output = dst_tex.GetFrame<mediapipe::GpuBuffer>();
cc->Outputs().Tag("IMAGE_GPU").Add(output.release(), cc->InputTimestamp());
cc->Outputs().Tag(kImageGpuTag).Add(output.release(), cc->InputTimestamp());
// Cleanup
src_tex.Release();
dst_tex.Release();
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
void ImageCroppingCalculator::GlRender() {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
static const GLfloat square_vertices[] = {
-1.0f, -1.0f, // bottom left
1.0f, -1.0f, // bottom right
@@ -342,11 +373,11 @@ void ImageCroppingCalculator::GlRender() {
glDeleteVertexArrays(1, &vao);
glDeleteBuffers(2, vbo);
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
}
::mediapipe::Status ImageCroppingCalculator::InitGpu(CalculatorContext* cc) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
const GLint attr_location[NUM_ATTRIBUTES] = {
ATTRIB_VERTEX,
ATTRIB_TEXTURE_POSITION,
@@ -392,7 +423,7 @@ void ImageCroppingCalculator::GlRender() {
// Parameters
glUseProgram(program_);
glUniform1i(glGetUniformLocation(program_, "input_frame"), 1);
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
@@ -410,8 +441,8 @@ void ImageCroppingCalculator::GetOutputDimensions(CalculatorContext* cc,
int y_center = src_height / 2;
// Get the rotation of the cropping box.
float rotation = 0.0f;
if (cc->Inputs().HasTag("RECT")) {
const auto& rect = cc->Inputs().Tag("RECT").Get<Rect>();
if (cc->Inputs().HasTag(kRectTag)) {
const auto& rect = cc->Inputs().Tag(kRectTag).Get<Rect>();
// Only use the rect if it is valid.
if (rect.width() > 0 && rect.height() > 0 && rect.x_center() >= 0 &&
rect.y_center() >= 0) {
@@ -421,8 +452,8 @@ void ImageCroppingCalculator::GetOutputDimensions(CalculatorContext* cc,
crop_height = rect.height();
rotation = rect.rotation();
}
} else if (cc->Inputs().HasTag("NORM_RECT")) {
const auto& rect = cc->Inputs().Tag("NORM_RECT").Get<NormalizedRect>();
} else if (cc->Inputs().HasTag(kNormRectTag)) {
const auto& rect = cc->Inputs().Tag(kNormRectTag).Get<NormalizedRect>();
// Only use the rect if it is valid.
if (rect.width() > 0.0 && rect.height() > 0.0 && rect.x_center() >= 0.0 &&
rect.y_center() >= 0.0) {
@@ -433,9 +464,9 @@ void ImageCroppingCalculator::GetOutputDimensions(CalculatorContext* cc,
rotation = rect.rotation();
}
} else {
if (cc->Inputs().HasTag("WIDTH") && cc->Inputs().HasTag("HEIGHT")) {
crop_width = cc->Inputs().Tag("WIDTH").Get<int>();
crop_height = cc->Inputs().Tag("HEIGHT").Get<int>();
if (cc->Inputs().HasTag(kWidthTag) && cc->Inputs().HasTag(kHeightTag)) {
crop_width = cc->Inputs().Tag(kWidthTag).Get<int>();
crop_height = cc->Inputs().Tag(kHeightTag).Get<int>();
} else if (options_.has_width() && options_.has_height()) {
crop_width = options_.width();
crop_height = options_.height();
mediapipe/calculators/image/image_properties_calculator.cc
+6 -6
View File
@@ -15,9 +15,9 @@
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_frame.h"
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
#include "mediapipe/gpu/gpu_buffer.h"
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
namespace mediapipe {
@@ -44,11 +44,11 @@ class ImagePropertiesCalculator : public CalculatorBase {
if (cc->Inputs().HasTag("IMAGE")) {
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag("IMAGE_GPU")) {
cc->Inputs().Tag("IMAGE_GPU").Set<::mediapipe::GpuBuffer>();
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Outputs().HasTag("SIZE")) {
cc->Outputs().Tag("SIZE").Set<std::pair<int, int>>();
@@ -71,7 +71,7 @@ class ImagePropertiesCalculator : public CalculatorBase {
width = image.Width();
height = image.Height();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag("IMAGE_GPU") &&
!cc->Inputs().Tag("IMAGE_GPU").IsEmpty()) {
const auto& image =
@@ -79,7 +79,7 @@ class ImagePropertiesCalculator : public CalculatorBase {
width = image.width();
height = image.height();
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
cc->Outputs().Tag("SIZE").AddPacket(
MakePacket<std::pair<int, int>>(width, height)
mediapipe/calculators/image/image_transformation_calculator.cc
+29 -22
View File
@@ -22,12 +22,12 @@
#include "mediapipe/framework/port/status.h"
#include "mediapipe/gpu/scale_mode.pb.h"
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
#if !defined(MEDIAPIPE_DISABLE_GPU)
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gl_quad_renderer.h"
#include "mediapipe/gpu/gl_simple_shaders.h"
#include "mediapipe/gpu/shader_util.h"
#endif // __ANDROID__ || iOS
#endif // !MEDIAPIPE_DISABLE_GPU
#if defined(__ANDROID__)
// The size of Java arrays is dynamic, which makes it difficult to
@@ -42,9 +42,9 @@ typedef int DimensionsPacketType[2];
namespace mediapipe {
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
#if !defined(MEDIAPIPE_DISABLE_GPU)
#endif // __ANDROID__ || iOS
#endif // !MEDIAPIPE_DISABLE_GPU
namespace {
int RotationModeToDegrees(mediapipe::RotationMode_Mode rotation) {
@@ -170,12 +170,12 @@ class ImageTransformationCalculator : public CalculatorBase {
mediapipe::ScaleMode_Mode scale_mode_;
bool use_gpu_ = false;
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
#if !defined(MEDIAPIPE_DISABLE_GPU)
GlCalculatorHelper helper_;
std::unique_ptr<QuadRenderer> rgb_renderer_;
std::unique_ptr<QuadRenderer> yuv_renderer_;
std::unique_ptr<QuadRenderer> ext_rgb_renderer_;
#endif // __ANDROID__ || iOS
#endif // !MEDIAPIPE_DISABLE_GPU
};
REGISTER_CALCULATOR(ImageTransformationCalculator);
@@ -185,18 +185,22 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
RET_CHECK(cc->Inputs().HasTag("IMAGE") ^ cc->Inputs().HasTag("IMAGE_GPU"));
RET_CHECK(cc->Outputs().HasTag("IMAGE") ^ cc->Outputs().HasTag("IMAGE_GPU"));
bool use_gpu = false;
if (cc->Inputs().HasTag("IMAGE")) {
RET_CHECK(cc->Outputs().HasTag("IMAGE"));
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
}
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag("IMAGE_GPU")) {
RET_CHECK(cc->Outputs().HasTag("IMAGE_GPU"));
cc->Inputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
cc->Outputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__ || iOS
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Inputs().HasTag("ROTATION_DEGREES")) {
cc->Inputs().Tag("ROTATION_DEGREES").Set<int>();
}
@@ -212,9 +216,11 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
cc->Outputs().Tag("LETTERBOX_PADDING").Set<std::array<float, 4>>();
}
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
MP_RETURN_IF_ERROR(GlCalculatorHelper::UpdateContract(cc));
#endif // __ANDROID__ || iOS
if (use_gpu) {
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(GlCalculatorHelper::UpdateContract(cc));
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
}
@@ -250,12 +256,12 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
scale_mode_ = ParseScaleMode(options_.scale_mode(), DEFAULT_SCALE_MODE);
if (use_gpu_) {
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
#if !defined(MEDIAPIPE_DISABLE_GPU)
// Let the helper access the GL context information.
MP_RETURN_IF_ERROR(helper_.Open(cc));
#else
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
#endif // __ANDROID__ || iOS
RET_CHECK_FAIL() << "GPU processing not enabled.";
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
@@ -264,10 +270,10 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
::mediapipe::Status ImageTransformationCalculator::Process(
CalculatorContext* cc) {
if (use_gpu_) {
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
#if !defined(MEDIAPIPE_DISABLE_GPU)
return helper_.RunInGlContext(
[this, cc]() -> ::mediapipe::Status { return RenderGpu(cc); });
#endif // __ANDROID__ || iOS
#endif // !MEDIAPIPE_DISABLE_GPU
} else {
return RenderCpu(cc);
}
@@ -277,7 +283,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
::mediapipe::Status ImageTransformationCalculator::Close(
CalculatorContext* cc) {
if (use_gpu_) {
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
#if !defined(MEDIAPIPE_DISABLE_GPU)
QuadRenderer* rgb_renderer = rgb_renderer_.release();
QuadRenderer* yuv_renderer = yuv_renderer_.release();
QuadRenderer* ext_rgb_renderer = ext_rgb_renderer_.release();
@@ -295,8 +301,9 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
delete yuv_renderer;
}
});
#endif // __ANDROID__ || iOS
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
}
@@ -371,7 +378,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
::mediapipe::Status ImageTransformationCalculator::RenderGpu(
CalculatorContext* cc) {
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
#if !defined(MEDIAPIPE_DISABLE_GPU)
int input_width = cc->Inputs().Tag("IMAGE_GPU").Get<GpuBuffer>().width();
int input_height = cc->Inputs().Tag("IMAGE_GPU").Get<GpuBuffer>().height();
@@ -408,7 +415,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
#endif // iOS
{
src1 = helper_.CreateSourceTexture(input);
#if defined(__ANDROID__)
#if defined(TEXTURE_EXTERNAL_OES)
if (src1.target() == GL_TEXTURE_EXTERNAL_OES) {
if (!ext_rgb_renderer_) {
ext_rgb_renderer_ = absl::make_unique<QuadRenderer>();
@@ -417,7 +424,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
}
renderer = ext_rgb_renderer_.get();
} else // NOLINT(readability/braces)
#endif // __ANDROID__
#endif // TEXTURE_EXTERNAL_OES
{
if (!rgb_renderer_) {
rgb_renderer_ = absl::make_unique<QuadRenderer>();
@@ -460,7 +467,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
auto output = dst.GetFrame<GpuBuffer>();
cc->Outputs().Tag("IMAGE_GPU").Add(output.release(), cc->InputTimestamp());
#endif // __ANDROID__ || iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
mediapipe/calculators/image/recolor_calculator.cc
+32 -26
View File
@@ -21,12 +21,11 @@
#include "mediapipe/framework/port/status.h"
#include "mediapipe/util/color.pb.h"
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gl_simple_shaders.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/gpu/shader_util.h"
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
namespace {
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
@@ -95,10 +94,10 @@ class RecolorCalculator : public CalculatorBase {
mediapipe::RecolorCalculatorOptions::MaskChannel mask_channel_;
bool use_gpu_ = false;
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
mediapipe::GlCalculatorHelper gpu_helper_;
GLuint program_ = 0;
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
};
REGISTER_CALCULATOR(RecolorCalculator);
@@ -107,36 +106,43 @@ REGISTER_CALCULATOR(RecolorCalculator);
RET_CHECK(!cc->Inputs().GetTags().empty());
RET_CHECK(!cc->Outputs().GetTags().empty());
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
bool use_gpu = false;
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag("IMAGE_GPU")) {
cc->Inputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Inputs().HasTag("IMAGE")) {
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag("MASK_GPU")) {
cc->Inputs().Tag("MASK_GPU").Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Inputs().HasTag("MASK")) {
cc->Inputs().Tag("MASK").Set<ImageFrame>();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Outputs().HasTag("IMAGE_GPU")) {
cc->Outputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Outputs().HasTag("IMAGE")) {
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // __ANDROID__ or iOS
if (use_gpu) {
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
}
@@ -146,9 +152,9 @@ REGISTER_CALCULATOR(RecolorCalculator);
if (cc->Inputs().HasTag("IMAGE_GPU")) {
use_gpu_ = true;
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
}
MP_RETURN_IF_ERROR(LoadOptions(cc));
@@ -158,7 +164,7 @@ REGISTER_CALCULATOR(RecolorCalculator);
::mediapipe::Status RecolorCalculator::Process(CalculatorContext* cc) {
if (use_gpu_) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, &cc]() -> ::mediapipe::Status {
if (!initialized_) {
@@ -168,7 +174,7 @@ REGISTER_CALCULATOR(RecolorCalculator);
MP_RETURN_IF_ERROR(RenderGpu(cc));
return ::mediapipe::OkStatus();
}));
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
} else {
MP_RETURN_IF_ERROR(RenderCpu(cc));
}
@@ -176,12 +182,12 @@ REGISTER_CALCULATOR(RecolorCalculator);
}
::mediapipe::Status RecolorCalculator::Close(CalculatorContext* cc) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
gpu_helper_.RunInGlContext([this] {
if (program_) glDeleteProgram(program_);
program_ = 0;
});
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
@@ -194,7 +200,7 @@ REGISTER_CALCULATOR(RecolorCalculator);
if (cc->Inputs().Tag("MASK_GPU").IsEmpty()) {
return ::mediapipe::OkStatus();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
// Get inputs and setup output.
const Packet& input_packet = cc->Inputs().Tag("IMAGE_GPU").Value();
const Packet& mask_packet = cc->Inputs().Tag("MASK_GPU").Value();
@@ -233,13 +239,13 @@ REGISTER_CALCULATOR(RecolorCalculator);
img_tex.Release();
mask_tex.Release();
dst_tex.Release();
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
void RecolorCalculator::GlRender() {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
static const GLfloat square_vertices[] = {
-1.0f, -1.0f, // bottom left
1.0f, -1.0f, // bottom right
@@ -287,7 +293,7 @@ void RecolorCalculator::GlRender() {
glBindVertexArray(0);
glDeleteVertexArrays(1, &vao);
glDeleteBuffers(2, vbo);
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
}
::mediapipe::Status RecolorCalculator::LoadOptions(CalculatorContext* cc) {
@@ -305,7 +311,7 @@ void RecolorCalculator::GlRender() {
}
::mediapipe::Status RecolorCalculator::InitGpu(CalculatorContext* cc) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
const GLint attr_location[NUM_ATTRIBUTES] = {
ATTRIB_VERTEX,
ATTRIB_TEXTURE_POSITION,
@@ -374,7 +380,7 @@ void RecolorCalculator::GlRender() {
glUniform1i(glGetUniformLocation(program_, "mask"), 2);
glUniform3f(glGetUniformLocation(program_, "recolor"), color_[0], color_[1],
color_[2]);
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
mediapipe/calculators/image/set_alpha_calculator.cc
+35 -29
View File
@@ -25,12 +25,11 @@
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/vector.h"
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gl_simple_shaders.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/gpu/shader_util.h"
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
namespace mediapipe {
@@ -107,16 +106,18 @@ class SetAlphaCalculator : public CalculatorBase {
bool use_gpu_ = false;
bool gpu_initialized_ = false;
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
mediapipe::GlCalculatorHelper gpu_helper_;
GLuint program_ = 0;
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
};
REGISTER_CALCULATOR(SetAlphaCalculator);
::mediapipe::Status SetAlphaCalculator::GetContract(CalculatorContract* cc) {
CHECK_GE(cc->Inputs().NumEntries(), 1);
bool use_gpu = false;
if (cc->Inputs().HasTag(kInputFrameTag) &&
cc->Inputs().HasTag(kInputFrameTagGpu)) {
return ::mediapipe::InternalError("Cannot have multiple input images.");
@@ -127,38 +128,43 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
}
// Input image to add/edit alpha channel.
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag(kInputFrameTagGpu)) {
cc->Inputs().Tag(kInputFrameTagGpu).Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Inputs().HasTag(kInputFrameTag)) {
cc->Inputs().Tag(kInputFrameTag).Set<ImageFrame>();
}
// Input alpha image mask (optional)
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag(kInputAlphaTagGpu)) {
cc->Inputs().Tag(kInputAlphaTagGpu).Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Inputs().HasTag(kInputAlphaTag)) {
cc->Inputs().Tag(kInputAlphaTag).Set<ImageFrame>();
}
// RGBA output image.
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Outputs().HasTag(kOutputFrameTagGpu)) {
cc->Outputs().Tag(kOutputFrameTagGpu).Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Outputs().HasTag(kOutputFrameTag)) {
cc->Outputs().Tag(kOutputFrameTag).Set<ImageFrame>();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // __ANDROID__ or iOS
if (use_gpu) {
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
}
@@ -170,11 +176,11 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
if (cc->Inputs().HasTag(kInputFrameTagGpu) &&
cc->Outputs().HasTag(kOutputFrameTagGpu)) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
use_gpu_ = true;
#else
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
#endif // __ANDROID__ or iOS
RET_CHECK_FAIL() << "GPU processing not enabled.";
#endif // !MEDIAPIPE_DISABLE_GPU
}
// Get global value from options (-1 if not set).
@@ -187,17 +193,17 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
RET_CHECK_FAIL() << "Must use either image mask or options alpha value.";
if (use_gpu_) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#endif
}
} // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
::mediapipe::Status SetAlphaCalculator::Process(CalculatorContext* cc) {
if (use_gpu_) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
if (!gpu_initialized_) {
@@ -207,7 +213,7 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
MP_RETURN_IF_ERROR(RenderGpu(cc));
return ::mediapipe::OkStatus();
}));
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
} else {
MP_RETURN_IF_ERROR(RenderCpu(cc));
}
@@ -216,12 +222,12 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
}
::mediapipe::Status SetAlphaCalculator::Close(CalculatorContext* cc) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
gpu_helper_.RunInGlContext([this] {
if (program_) glDeleteProgram(program_);
program_ = 0;
});
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
@@ -295,7 +301,7 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
if (cc->Inputs().Tag(kInputFrameTagGpu).IsEmpty()) {
return ::mediapipe::OkStatus();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
// Setup source texture.
const auto& input_frame =
cc->Inputs().Tag(kInputFrameTagGpu).Get<mediapipe::GpuBuffer>();
@@ -348,13 +354,13 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
// Cleanup
input_texture.Release();
output_texture.Release();
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
void SetAlphaCalculator::GlRender(CalculatorContext* cc) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
static const GLfloat square_vertices[] = {
-1.0f, -1.0f, // bottom left
1.0f, -1.0f, // bottom right
@@ -403,11 +409,11 @@ void SetAlphaCalculator::GlRender(CalculatorContext* cc) {
glDeleteVertexArrays(1, &vao);
glDeleteBuffers(2, vbo);
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
}
::mediapipe::Status SetAlphaCalculator::GlSetup(CalculatorContext* cc) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
const GLint attr_location[NUM_ATTRIBUTES] = {
ATTRIB_VERTEX,
ATTRIB_TEXTURE_POSITION,
@@ -460,7 +466,7 @@ void SetAlphaCalculator::GlRender(CalculatorContext* cc) {
glUniform1i(glGetUniformLocation(program_, "alpha_mask"), 2);
glUniform1f(glGetUniformLocation(program_, "alpha_value"), alpha_value_);
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
mediapipe/calculators/tensorflow/BUILD
+3 -1
View File
@@ -255,6 +255,7 @@ mediapipe_cc_proto_library(
cc_deps = [
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/util:audio_decoder_cc_proto",
],
visibility = ["//visibility:public"],
deps = [":unpack_media_sequence_calculator_proto"],
@@ -653,6 +654,7 @@ cc_library(
"//mediapipe/framework/formats:location",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/util:audio_decoder_cc_proto",
"//mediapipe/util/sequence:media_sequence",
"@com_google_absl//absl/strings",
"@org_tensorflow//tensorflow/core:protos_all_cc",
@@ -769,7 +771,6 @@ cc_test(
"//mediapipe/framework/formats:location",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:opencv_imgcodecs",
"//mediapipe/framework/port:status",
"//mediapipe/util/sequence:media_sequence",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
@@ -971,6 +972,7 @@ cc_test(
"//mediapipe/framework/formats:location",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:rectangle",
"//mediapipe/util:audio_decoder_cc_proto",
"//mediapipe/util/sequence:media_sequence",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
mediapipe/calculators/tensorflow/pack_media_sequence_calculator.cc
+24 -4
View File
@@ -285,6 +285,10 @@ class PackMediaSequenceCalculator : public CalculatorBase {
}
::mediapipe::Status Process(CalculatorContext* cc) override {
int image_height = -1;
int image_width = -1;
// Because the tag order may vary, we need to loop through tags to get
// image information before processing other tag types.
for (const auto& tag : cc->Inputs().GetTags()) {
if (!cc->Inputs().Tag(tag).IsEmpty()) {
features_present_[tag] = true;
@@ -306,14 +310,21 @@ class PackMediaSequenceCalculator : public CalculatorBase {
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
<< "No encoded image";
}
image_height = image.height();
image_width = image.width();
mpms::AddImageTimestamp(key, cc->InputTimestamp().Value(),
sequence_.get());
mpms::AddImageEncoded(key, image.encoded_image(), sequence_.get());
}
}
for (const auto& tag : cc->Inputs().GetTags()) {
if (!cc->Inputs().Tag(tag).IsEmpty()) {
features_present_[tag] = true;
}
if (absl::StartsWith(tag, kKeypointsTag) &&
!cc->Inputs().Tag(tag).IsEmpty()) {
std::string key = "";
if (tag != kImageTag) {
if (tag != kKeypointsTag) {
int tag_length = sizeof(kKeypointsTag) / sizeof(*kKeypointsTag) - 1;
if (tag[tag_length] == '_') {
key = tag.substr(tag_length + 1);
@@ -363,11 +374,20 @@ class PackMediaSequenceCalculator : public CalculatorBase {
LocationData::BOUNDING_BOX ||
detection.location_data().format() ==
LocationData::RELATIVE_BOUNDING_BOX) {
int height = mpms::GetImageHeight(*sequence_);
int width = mpms::GetImageWidth(*sequence_);
if (mpms::HasImageHeight(*sequence_) &&
mpms::HasImageWidth(*sequence_)) {
image_height = mpms::GetImageHeight(*sequence_);
image_width = mpms::GetImageWidth(*sequence_);
}
if (image_height == -1 || image_width == -1) {
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
<< "Images must be provided with bounding boxes or the "
"image "
<< "height and width must already be in the example.";
}
Location relative_bbox = Location::CreateRelativeBBoxLocation(
Location(detection.location_data())
.ConvertToRelativeBBox(width, height));
.ConvertToRelativeBBox(image_width, image_height));
predicted_locations.push_back(relative_bbox);
if (detection.label_size() > 0) {
predicted_class_strings.push_back(detection.label(0));
mediapipe/calculators/tensorflow/pack_media_sequence_calculator_test.cc
+142
View File
@@ -357,6 +357,148 @@ TEST_F(PackMediaSequenceCalculatorTest, PacksTwoBBoxDetections) {
}
}
TEST_F(PackMediaSequenceCalculatorTest, PacksBBoxWithoutImageDims) {
SetUpCalculator({"BBOX_PREDICTED:detections"}, {}, false, true);
auto input_sequence = ::absl::make_unique<tf::SequenceExample>();
std::string test_video_id = "test_video_id";
mpms::SetClipMediaId(test_video_id, input_sequence.get());
int height = 480;
int width = 640;
int num_vectors = 2;
for (int i = 0; i < num_vectors; ++i) {
auto detections = ::absl::make_unique<::std::vector<Detection>>();
Detection detection;
detection.add_label("absolute bbox");
detection.add_label_id(0);
detection.add_score(0.5);
Location::CreateBBoxLocation(0, height / 2, width / 2, height / 2)
.ConvertToProto(detection.mutable_location_data());
detections->push_back(detection);
detection = Detection();
detection.add_label("relative bbox");
detection.add_label_id(1);
detection.add_score(0.75);
Location::CreateRelativeBBoxLocation(0, 0.5, 0.5, 0.5)
.ConvertToProto(detection.mutable_location_data());
detections->push_back(detection);
// The mask detection should be ignored in the output.
detection = Detection();
detection.add_label("mask");
detection.add_score(1.0);
cv::Mat image(2, 3, CV_8UC1, cv::Scalar(0));
Location::CreateCvMaskLocation<uint8>(image).ConvertToProto(
detection.mutable_location_data());
detections->push_back(detection);
runner_->MutableInputs()
->Tag("BBOX_PREDICTED")
.packets.push_back(Adopt(detections.release()).At(Timestamp(i)));
}
runner_->MutableSidePackets()->Tag("SEQUENCE_EXAMPLE") =
Adopt(input_sequence.release());
auto status = runner_->Run();
EXPECT_EQ(::mediapipe::StatusCode::kInvalidArgument, status.code());
}
TEST_F(PackMediaSequenceCalculatorTest, PacksBBoxWithImages) {
SetUpCalculator({"BBOX_PREDICTED:detections", "IMAGE:images"}, {}, false,
true);
auto input_sequence = ::absl::make_unique<tf::SequenceExample>();
std::string test_video_id = "test_video_id";
mpms::SetClipMediaId(test_video_id, input_sequence.get());
int height = 480;
int width = 640;
int num_vectors = 2;
for (int i = 0; i < num_vectors; ++i) {
auto detections = ::absl::make_unique<::std::vector<Detection>>();
Detection detection;
detection.add_label("absolute bbox");
detection.add_label_id(0);
detection.add_score(0.5);
Location::CreateBBoxLocation(0, height / 2, width / 2, height / 2)
.ConvertToProto(detection.mutable_location_data());
detections->push_back(detection);
detection = Detection();
detection.add_label("relative bbox");
detection.add_label_id(1);
detection.add_score(0.75);
Location::CreateRelativeBBoxLocation(0, 0.5, 0.5, 0.5)
.ConvertToProto(detection.mutable_location_data());
detections->push_back(detection);
// The mask detection should be ignored in the output.
detection = Detection();
detection.add_label("mask");
detection.add_score(1.0);
cv::Mat image(2, 3, CV_8UC1, cv::Scalar(0));
Location::CreateCvMaskLocation<uint8>(image).ConvertToProto(
detection.mutable_location_data());
detections->push_back(detection);
runner_->MutableInputs()
->Tag("BBOX_PREDICTED")
.packets.push_back(Adopt(detections.release()).At(Timestamp(i)));
}
cv::Mat image(height, width, CV_8UC3, cv::Scalar(0, 0, 255));
std::vector<uchar> bytes;
ASSERT_TRUE(cv::imencode(".jpg", image, bytes, {80}));
std::string test_image_string(bytes.begin(), bytes.end());
OpenCvImageEncoderCalculatorResults encoded_image;
encoded_image.set_encoded_image(test_image_string);
encoded_image.set_width(width);
encoded_image.set_height(height);
int num_images = 2;
for (int i = 0; i < num_images; ++i) {
auto image_ptr =
::absl::make_unique<OpenCvImageEncoderCalculatorResults>(encoded_image);
runner_->MutableInputs()->Tag("IMAGE").packets.push_back(
Adopt(image_ptr.release()).At(Timestamp(i)));
}
runner_->MutableSidePackets()->Tag("SEQUENCE_EXAMPLE") =
Adopt(input_sequence.release());
MP_ASSERT_OK(runner_->Run());
const std::vector<Packet>& output_packets =
runner_->Outputs().Tag("SEQUENCE_EXAMPLE").packets;
ASSERT_EQ(1, output_packets.size());
const tf::SequenceExample& output_sequence =
output_packets[0].Get<tf::SequenceExample>();
ASSERT_EQ(test_video_id, mpms::GetClipMediaId(output_sequence));
ASSERT_EQ(height, mpms::GetImageHeight(output_sequence));
ASSERT_EQ(width, mpms::GetImageWidth(output_sequence));
ASSERT_EQ(num_vectors, mpms::GetPredictedBBoxSize(output_sequence));
ASSERT_EQ(num_vectors, mpms::GetPredictedBBoxTimestampSize(output_sequence));
ASSERT_EQ(0, mpms::GetClassSegmentationEncodedSize(output_sequence));
ASSERT_EQ(0, mpms::GetClassSegmentationTimestampSize(output_sequence));
for (int i = 0; i < num_vectors; ++i) {
ASSERT_EQ(i, mpms::GetPredictedBBoxTimestampAt(output_sequence, i));
auto bboxes = mpms::GetPredictedBBoxAt(output_sequence, i);
ASSERT_EQ(2, bboxes.size());
for (int j = 0; j < bboxes.size(); ++j) {
auto rect = bboxes[j].GetRelativeBBox();
ASSERT_NEAR(0, rect.xmin(), 0.001);
ASSERT_NEAR(0.5, rect.ymin(), 0.001);
ASSERT_NEAR(0.5, rect.xmax(), 0.001);
ASSERT_NEAR(1.0, rect.ymax(), 0.001);
}
auto class_strings =
mpms::GetPredictedBBoxLabelStringAt(output_sequence, i);
ASSERT_EQ("absolute bbox", class_strings[0]);
ASSERT_EQ("relative bbox", class_strings[1]);
auto class_indices = mpms::GetPredictedBBoxLabelIndexAt(output_sequence, i);
ASSERT_EQ(0, class_indices[0]);
ASSERT_EQ(1, class_indices[1]);
}
}
TEST_F(PackMediaSequenceCalculatorTest, PacksTwoKeypoints) {
SetUpCalculator({"KEYPOINTS_TEST:keypoints"}, {}, false, true);
auto input_sequence = ::absl::make_unique<tf::SequenceExample>();
mediapipe/calculators/tensorflow/unpack_media_sequence_calculator.cc
+34 -4
View File
@@ -19,6 +19,7 @@
#include "mediapipe/framework/formats/location.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/util/audio_decoder.pb.h"
#include "mediapipe/util/sequence/media_sequence.h"
#include "tensorflow/core/example/example.pb.h"
#include "tensorflow/core/example/feature.pb.h"
@@ -37,6 +38,7 @@ const char kDatasetRootDirTag[] = "DATASET_ROOT";
const char kDataPath[] = "DATA_PATH";
const char kPacketResamplerOptions[] = "RESAMPLER_OPTIONS";
const char kImagesFrameRateTag[] = "IMAGE_FRAME_RATE";
const char kAudioDecoderOptions[] = "AUDIO_DECODER_OPTIONS";
namespace tf = ::tensorflow;
namespace mpms = ::mediapipe::mediasequence;
@@ -126,6 +128,11 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
if (cc->OutputSidePackets().HasTag(kDataPath)) {
cc->OutputSidePackets().Tag(kDataPath).Set<std::string>();
}
if (cc->OutputSidePackets().HasTag(kAudioDecoderOptions)) {
cc->OutputSidePackets()
.Tag(kAudioDecoderOptions)
.Set<AudioDecoderOptions>();
}
if (cc->OutputSidePackets().HasTag(kImagesFrameRateTag)) {
cc->OutputSidePackets().Tag(kImagesFrameRateTag).Set<double>();
}
@@ -136,10 +143,11 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
}
if ((options.has_padding_before_label() ||
options.has_padding_after_label()) &&
!(cc->OutputSidePackets().HasTag(kPacketResamplerOptions))) {
!(cc->OutputSidePackets().HasTag(kAudioDecoderOptions) ||
cc->OutputSidePackets().HasTag(kPacketResamplerOptions))) {
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
<< "If specifying padding, must output "
<< kPacketResamplerOptions;
<< "If specifying padding, must output " << kPacketResamplerOptions
<< "or" << kAudioDecoderOptions;
}
// Optional streams.
@@ -260,7 +268,8 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
// Set the start and end of the clip in the appropriate options protos.
double start_time = 0;
double end_time = 0;
if (cc->OutputSidePackets().HasTag(kPacketResamplerOptions)) {
if (cc->OutputSidePackets().HasTag(kAudioDecoderOptions) ||
cc->OutputSidePackets().HasTag(kPacketResamplerOptions)) {
if (mpms::HasClipStartTimestamp(sequence)) {
start_time =
Timestamp(mpms::GetClipStartTimestamp(sequence)).Seconds() -
@@ -271,6 +280,27 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
options.padding_after_label();
}
}
if (cc->OutputSidePackets().HasTag(kAudioDecoderOptions)) {
auto audio_decoder_options = absl::make_unique<AudioDecoderOptions>(
options.base_audio_decoder_options());
if (mpms::HasClipStartTimestamp(sequence)) {
if (options.force_decoding_from_start_of_media()) {
audio_decoder_options->set_start_time(0);
} else {
audio_decoder_options->set_start_time(
start_time - options.extra_padding_from_media_decoder());
}
}
if (mpms::HasClipEndTimestamp(sequence)) {
audio_decoder_options->set_end_time(
end_time + options.extra_padding_from_media_decoder());
}
LOG(INFO) << "Created AudioDecoderOptions:\n"
<< audio_decoder_options->DebugString();
cc->OutputSidePackets()
.Tag(kAudioDecoderOptions)
.Set(Adopt(audio_decoder_options.release()));
}
if (cc->OutputSidePackets().HasTag(kPacketResamplerOptions)) {
auto resampler_options = absl::make_unique<CalculatorOptions>();
*(resampler_options->MutableExtension(
mediapipe/calculators/tensorflow/unpack_media_sequence_calculator.proto
+7
View File
@@ -18,6 +18,7 @@ package mediapipe;
import "mediapipe/calculators/core/packet_resampler_calculator.proto";
import "mediapipe/framework/calculator.proto";
import "mediapipe/util/audio_decoder.proto";
message UnpackMediaSequenceCalculatorOptions {
extend mediapipe.CalculatorOptions {
@@ -49,4 +50,10 @@ message UnpackMediaSequenceCalculatorOptions {
// parameters for the MediaDecoderCalculator. End time parameters are still
// respected.
optional bool force_decoding_from_start_of_media = 7;
// Stores the audio decoder settings for the graph. (e.g. which audio
// stream to pull from the video.) The sequence's metadata overrides
// the clip start and end times and outputs these for the
// AudioDecoderCalculator to consume.
optional AudioDecoderOptions base_audio_decoder_options = 9;
}
mediapipe/calculators/tensorflow/unpack_media_sequence_calculator_test.cc
+57
View File
@@ -23,6 +23,7 @@
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/rectangle.h"
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/util/audio_decoder.pb.h"
#include "mediapipe/util/sequence/media_sequence.h"
#include "tensorflow/core/example/example.pb.h"
@@ -459,6 +460,62 @@ TEST_F(UnpackMediaSequenceCalculatorTest, GetDatasetFromExample) {
data_path_);
}
TEST_F(UnpackMediaSequenceCalculatorTest, GetAudioDecoderOptions) {
CalculatorOptions options;
options.MutableExtension(UnpackMediaSequenceCalculatorOptions::ext)
->set_padding_before_label(1);
options.MutableExtension(UnpackMediaSequenceCalculatorOptions::ext)
->set_padding_after_label(2);
SetUpCalculator({}, {"AUDIO_DECODER_OPTIONS:audio_decoder_options"}, {},
&options);
runner_->MutableSidePackets()->Tag("SEQUENCE_EXAMPLE") =
Adopt(sequence_.release());
MP_ASSERT_OK(runner_->Run());
MP_EXPECT_OK(runner_->OutputSidePackets()
.Tag("AUDIO_DECODER_OPTIONS")
.ValidateAsType<AudioDecoderOptions>());
EXPECT_NEAR(runner_->OutputSidePackets()
.Tag("AUDIO_DECODER_OPTIONS")
.Get<AudioDecoderOptions>()
.start_time(),
2.0, 1e-5);
EXPECT_NEAR(runner_->OutputSidePackets()
.Tag("AUDIO_DECODER_OPTIONS")
.Get<AudioDecoderOptions>()
.end_time(),
7.0, 1e-5);
}
TEST_F(UnpackMediaSequenceCalculatorTest, GetAudioDecoderOptionsOverride) {
CalculatorOptions options;
options.MutableExtension(UnpackMediaSequenceCalculatorOptions::ext)
->set_padding_before_label(1);
options.MutableExtension(UnpackMediaSequenceCalculatorOptions::ext)
->set_padding_after_label(2);
options.MutableExtension(UnpackMediaSequenceCalculatorOptions::ext)
->set_force_decoding_from_start_of_media(true);
SetUpCalculator({}, {"AUDIO_DECODER_OPTIONS:audio_decoder_options"}, {},
&options);
runner_->MutableSidePackets()->Tag("SEQUENCE_EXAMPLE") =
Adopt(sequence_.release());
MP_ASSERT_OK(runner_->Run());
MP_EXPECT_OK(runner_->OutputSidePackets()
.Tag("AUDIO_DECODER_OPTIONS")
.ValidateAsType<AudioDecoderOptions>());
EXPECT_NEAR(runner_->OutputSidePackets()
.Tag("AUDIO_DECODER_OPTIONS")
.Get<AudioDecoderOptions>()
.start_time(),
0.0, 1e-5);
EXPECT_NEAR(runner_->OutputSidePackets()
.Tag("AUDIO_DECODER_OPTIONS")
.Get<AudioDecoderOptions>()
.end_time(),
7.0, 1e-5);
}
TEST_F(UnpackMediaSequenceCalculatorTest, GetPacketResamplingOptions) {
// TODO: Suport proto3 proto.Any in CalculatorOptions.
// TODO: Avoid proto2 extensions in "RESAMPLER_OPTIONS".
mediapipe/calculators/tflite/BUILD
+58 -22
View File
@@ -195,6 +195,12 @@ cc_test(
],
)
cc_library(
name = "util",
hdrs = ["util.h"],
alwayslink = 1,
)
cc_library(
name = "tflite_inference_calculator",
srcs = ["tflite_inference_calculator.cc"],
@@ -214,6 +220,7 @@ cc_library(
}),
visibility = ["//visibility:public"],
deps = [
":util",
":tflite_inference_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/util:resource_util",
@@ -222,20 +229,25 @@ cc_library(
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
"//mediapipe/framework/port:ret_check",
] + select({
"//mediapipe:android": [
"//mediapipe/gpu:disable_gpu": [],
"//mediapipe:ios": [
"//mediapipe/gpu:MPPMetalHelper",
"//mediapipe/gpu:MPPMetalUtil",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/objc:mediapipe_framework_ios",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
"@org_tensorflow//tensorflow/lite/delegates/gpu/metal:buffer_convert",
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
],
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gpu_buffer",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
],
"//mediapipe:ios": [
"//mediapipe/gpu:MPPMetalHelper",
"//mediapipe/objc:mediapipe_framework_ios",
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
],
"//conditions:default": [],
}),
alwayslink = 1,
)
@@ -259,33 +271,33 @@ cc_library(
}),
visibility = ["//visibility:public"],
deps = [
":util",
":tflite_converter_calculator_cc_proto",
"//mediapipe/util:resource_util",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:matrix",
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
"//mediapipe/framework/tool:status_util",
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:ret_check",
"@org_tensorflow//tensorflow/lite:framework",
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
] + select({
"//mediapipe:android": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gpu_buffer",
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
],
"//mediapipe/gpu:disable_gpu": [],
"//mediapipe:ios": [
"//mediapipe/gpu:MPPMetalUtil",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:MPPMetalHelper",
"//mediapipe/objc:mediapipe_framework_ios",
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
],
"//conditions:default": [],
"//conditions:default": [
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:gl_calculator_helper",
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
],
}),
alwayslink = 1,
)
@@ -295,6 +307,7 @@ cc_library(
srcs = ["tflite_tensors_to_segmentation_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
":util",
":tflite_tensors_to_segmentation_calculator_cc_proto",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
@@ -308,7 +321,9 @@ cc_library(
"//mediapipe/util:resource_util",
"@org_tensorflow//tensorflow/lite:framework",
] + select({
"//mediapipe:android": [
"//mediapipe/gpu:disable_gpu": [],
"//mediapipe:ios": [],
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gpu_buffer",
@@ -319,7 +334,6 @@ cc_library(
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_texture",
],
"//conditions:default": [],
}),
alwayslink = 1,
)
@@ -346,8 +360,23 @@ cc_test(
cc_library(
name = "tflite_tensors_to_detections_calculator",
srcs = ["tflite_tensors_to_detections_calculator.cc"],
copts = select({
"//mediapipe:ios": [
"-x objective-c++",
"-fobjc-arc", # enable reference-counting
],
"//conditions:default": [],
}),
linkopts = select({
"//mediapipe:ios": [
"-framework CoreVideo",
"-framework MetalKit",
],
"//conditions:default": [],
}),
visibility = ["//visibility:public"],
deps = [
":util",
":tflite_tensors_to_detections_calculator_cc_proto",
"//mediapipe/framework/formats:detection_cc_proto",
"@com_google_absl//absl/strings:str_format",
@@ -359,14 +388,21 @@ cc_library(
"//mediapipe/framework/port:ret_check",
"@org_tensorflow//tensorflow/lite:framework",
] + select({
"//mediapipe:android": [
"//mediapipe/gpu:disable_gpu": [],
"//mediapipe:ios": [
"//mediapipe/gpu:MPPMetalUtil",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:MPPMetalHelper",
"//mediapipe/objc:mediapipe_framework_ios",
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
],
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
],
"//conditions:default": [],
}),
alwayslink = 1,
)
mediapipe/calculators/tflite/tflite_converter_calculator.cc
+97 -102
View File
@@ -16,23 +16,23 @@
#include <vector>
#include "mediapipe/calculators/tflite/tflite_converter_calculator.pb.h"
#include "mediapipe/calculators/tflite/util.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/matrix.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/util/resource_util.h"
#include "tensorflow/lite/error_reporter.h"
#include "tensorflow/lite/interpreter.h"
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
#import <CoreVideo/CoreVideo.h>
@@ -40,11 +40,12 @@
#import <MetalKit/MetalKit.h>
#import "mediapipe/gpu/MPPMetalHelper.h"
#include "mediapipe/gpu/MPPMetalUtil.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
#endif // iOS
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
typedef id<MTLBuffer> GpuTensor;
@@ -66,26 +67,27 @@ typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>
namespace mediapipe {
#if defined(__ANDROID__)
using ::tflite::gpu::gl::GlBuffer;
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
using ::tflite::gpu::gl::GlProgram;
using ::tflite::gpu::gl::GlShader;
struct GPUData {
int elements = 1;
GlBuffer buffer;
GpuTensor buffer;
GlShader shader;
GlProgram program;
};
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
struct GPUData {
int elements = 1;
id<MTLBuffer> buffer;
GpuTensor buffer;
id<MTLComputePipelineState> pipeline_state;
};
#endif
// Calculator for normalizing and converting an ImageFrame or Matrix
// into a TfLiteTensor (float 32) or a GpuBuffer to a tflite::gpu::GlBuffer.
// into a TfLiteTensor (float 32) or a GpuBuffer to a tflite::gpu::GlBuffer
// or MTLBuffer.
//
// This calculator is designed to be used with the TfLiteInferenceCalcualtor,
// as a pre-processing step for calculator inputs.
@@ -102,7 +104,7 @@ struct GPUData {
// Output:
// One of the following tags:
// TENSORS - Vector of TfLiteTensor of type kTfLiteFloat32, or kTfLiteUint8.
// TENSORS_GPU - vector of GlBuffer.
// TENSORS_GPU - vector of GlBuffer or MTLBuffer.
//
// Example use:
// node {
@@ -144,7 +146,7 @@ class TfLiteConverterCalculator : public CalculatorBase {
std::unique_ptr<tflite::Interpreter> interpreter_ = nullptr;
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
mediapipe::GlCalculatorHelper gpu_helper_;
std::unique_ptr<GPUData> gpu_data_out_;
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
@@ -175,25 +177,33 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
RET_CHECK(cc->Outputs().HasTag("TENSORS") ^
cc->Outputs().HasTag("TENSORS_GPU"));
bool use_gpu = false;
if (cc->Inputs().HasTag("IMAGE")) cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
if (cc->Inputs().HasTag("MATRIX")) cc->Inputs().Tag("MATRIX").Set<Matrix>();
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
if (cc->Inputs().HasTag("IMAGE_GPU"))
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag("IMAGE_GPU")) {
cc->Inputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
#endif
use_gpu |= true;
}
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Outputs().HasTag("TENSORS"))
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
if (cc->Outputs().HasTag("TENSORS_GPU"))
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Outputs().HasTag("TENSORS_GPU")) {
cc->Outputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
#endif
use_gpu |= true;
}
#endif // !MEDIAPIPE_DISABLE_GPU
#if defined(__ANDROID__)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
if (use_gpu) {
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
#endif
}
// Assign this calculator's default InputStreamHandler.
cc->SetInputStreamHandler("FixedSizeInputStreamHandler");
@@ -208,10 +218,10 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
if (cc->Inputs().HasTag("IMAGE_GPU") ||
cc->Outputs().HasTag("IMAGE_OUT_GPU")) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
use_gpu_ = true;
#else
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
RET_CHECK_FAIL() << "GPU processing not enabled.";
#endif
}
@@ -221,7 +231,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
cc->Outputs().HasTag("TENSORS_GPU"));
// Cannot use quantization.
use_quantized_tensors_ = false;
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
@@ -238,6 +248,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
::mediapipe::Status TfLiteConverterCalculator::Process(CalculatorContext* cc) {
if (use_gpu_) {
// GpuBuffer to tflite::gpu::GlBuffer conversion.
if (!initialized_) {
MP_RETURN_IF_ERROR(InitGpu(cc));
initialized_ = true;
@@ -253,7 +264,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
}
::mediapipe::Status TfLiteConverterCalculator::Close(CalculatorContext* cc) {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
gpu_helper_.RunInGlContext([this] { gpu_data_out_.reset(); });
#endif
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
@@ -372,7 +383,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
::mediapipe::Status TfLiteConverterCalculator::ProcessGPU(
CalculatorContext* cc) {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
// GpuBuffer to tflite::gpu::GlBuffer conversion.
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
MP_RETURN_IF_ERROR(
@@ -381,17 +392,11 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
auto src = gpu_helper_.CreateSourceTexture(input);
glActiveTexture(GL_TEXTURE0 + 0);
glBindTexture(GL_TEXTURE_2D, src.name());
auto status = gpu_data_out_->buffer.BindToIndex(1);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
RET_CHECK_CALL(gpu_data_out_->buffer.BindToIndex(1));
const tflite::gpu::uint3 workgroups = {
NumGroups(input.width(), kWorkgroupSize),
NumGroups(input.height(), kWorkgroupSize), 1};
status = gpu_data_out_->program.Dispatch(workgroups);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
RET_CHECK_CALL(gpu_data_out_->program.Dispatch(workgroups));
glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
glBindTexture(GL_TEXTURE_2D, 0);
src.Release();
@@ -400,17 +405,17 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
// Copy into outputs.
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
output_tensors->resize(1);
{
GlBuffer& tensor = output_tensors->at(0);
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
auto status = CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_->elements, &tensor);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
tflite::gpu::gl::CopyBuffer(gpu_data_out_->buffer, tensor);
}
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &output_tensors]() -> ::mediapipe::Status {
output_tensors->resize(1);
{
GpuTensor& tensor = output_tensors->at(0);
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_->elements, &tensor));
RET_CHECK_CALL(CopyBuffer(gpu_data_out_->buffer, tensor));
}
return ::mediapipe::OkStatus();
}));
cc->Outputs()
.Tag("TENSORS_GPU")
.Add(output_tensors.release(), cc->InputTimestamp());
@@ -438,66 +443,60 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
}
// Copy into outputs.
// TODO Avoid this copy.
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
output_tensors->resize(1);
{
id<MTLDevice> device = gpu_helper_.mtlDevice;
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
command_buffer.label = @"TfLiteConverterCalculatorCopy";
id<MTLBuffer> tensor =
output_tensors->at(0) =
[device newBufferWithLength:gpu_data_out_->elements * sizeof(float)
options:MTLResourceStorageModeShared];
id<MTLBlitCommandEncoder> blit_command =
[command_buffer blitCommandEncoder];
[blit_command copyFromBuffer:gpu_data_out_->buffer
sourceOffset:0
toBuffer:tensor
destinationOffset:0
size:gpu_data_out_->elements * sizeof(float)];
[blit_command endEncoding];
[command_buffer commit];
[command_buffer waitUntilCompleted];
output_tensors->push_back(tensor);
[MPPMetalUtil blitMetalBufferTo:output_tensors->at(0)
from:gpu_data_out_->buffer
blocking:true
commandBuffer:[gpu_helper_ commandBuffer]];
}
cc->Outputs()
.Tag("TENSORS_GPU")
.Add(output_tensors.release(), cc->InputTimestamp());
#else
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
RET_CHECK_FAIL() << "GPU processing is not enabled.";
#endif
return ::mediapipe::OkStatus();
}
::mediapipe::Status TfLiteConverterCalculator::InitGpu(CalculatorContext* cc) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
// Configure inputs.
#if !defined(MEDIAPIPE_DISABLE_GPU)
// Get input image sizes.
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
mediapipe::ImageFormat::Format format =
mediapipe::ImageFormatForGpuBufferFormat(input.format());
gpu_data_out_ = absl::make_unique<GPUData>();
gpu_data_out_->elements = input.height() * input.width() * max_num_channels_;
const bool include_alpha = (max_num_channels_ == 4);
if (!(format == mediapipe::ImageFormat::SRGB ||
const bool single_channel = (max_num_channels_ == 1);
if (!(format == mediapipe::ImageFormat::GRAY8 ||
format == mediapipe::ImageFormat::SRGB ||
format == mediapipe::ImageFormat::SRGBA))
RET_CHECK_FAIL() << "Unsupported GPU input format.";
if (include_alpha && (format != mediapipe::ImageFormat::SRGBA))
RET_CHECK_FAIL() << "Num input channels is less than desired output.";
#endif
#endif // !MEDIAPIPE_DISABLE_GPU
#if defined(__ANDROID__)
// Device memory.
auto status = ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_->elements, &gpu_data_out_->buffer);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &include_alpha, &input, &single_channel]() -> ::mediapipe::Status {
// Device memory.
RET_CHECK_CALL(
::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_->elements, &gpu_data_out_->buffer));
// Shader to convert GL Texture to Shader Storage Buffer Object (SSBO),
// with normalization to either: [0,1] or [-1,1].
const std::string shader_source = absl::Substitute(
R"( #version 310 es
// Shader to convert GL Texture to Shader Storage Buffer Object (SSBO),
// with normalization to either: [0,1] or [-1,1].
const std::string shader_source = absl::Substitute(
R"( #version 310 es
layout(local_size_x = $0, local_size_y = $0) in;
layout(binding = 0) uniform sampler2D input_texture;
layout(std430, binding = 1) buffer Output {float elements[];} output_data;
@@ -505,33 +504,31 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
void main() {
ivec2 gid = ivec2(gl_GlobalInvocationID.xy);
if (gid.x >= width_height.x || gid.y >= width_height.y) return;
$5 // pixel fetch
vec4 pixel = texelFetch(input_texture, gid, 0);
$3 // normalize [-1,1]
int linear_index = $7 * ($4 * width_height.x + gid.x);
output_data.elements[linear_index + 0] = pixel.x;
output_data.elements[linear_index + 1] = pixel.y;
output_data.elements[linear_index + 2] = pixel.z;
output_data.elements[linear_index + 0] = pixel.x; // r channel
$5 // g & b channels
$6 // alpha channel
})",
/*$0=*/kWorkgroupSize, /*$1=*/input.width(), /*$2=*/input.height(),
/*$3=*/zero_center_ ? "pixel = (pixel - 0.5) * 2.0;" : "",
/*$4=*/flip_vertically_ ? "(width_height.y - 1 - gid.y)" : "gid.y",
/*$5=*/
include_alpha ? "vec4 pixel = texelFetch(input_texture, gid, 0);"
: "vec3 pixel = texelFetch(input_texture, gid, 0).xyz;",
/*$6=*/
include_alpha ? "output_data.elements[linear_index + 3] = pixel.w;" : "",
/*$7=*/include_alpha ? 4 : 3);
status = GlShader::CompileShader(GL_COMPUTE_SHADER, shader_source,
&gpu_data_out_->shader);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
status = GlProgram::CreateWithShader(gpu_data_out_->shader,
&gpu_data_out_->program);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
/*$0=*/kWorkgroupSize, /*$1=*/input.width(), /*$2=*/input.height(),
/*$3=*/zero_center_ ? "pixel = (pixel - 0.5) * 2.0;" : "",
/*$4=*/flip_vertically_ ? "(width_height.y - 1 - gid.y)" : "gid.y",
/*$5=*/
single_channel
? ""
: R"(output_data.elements[linear_index + 1] = pixel.y;
output_data.elements[linear_index + 2] = pixel.z;)",
/*$6=*/
include_alpha ? "output_data.elements[linear_index + 3] = pixel.w;"
: "",
/*$7=*/max_num_channels_);
RET_CHECK_CALL(GlShader::CompileShader(GL_COMPUTE_SHADER, shader_source,
&gpu_data_out_->shader));
RET_CHECK_CALL(GlProgram::CreateWithShader(gpu_data_out_->shader,
&gpu_data_out_->program));
return ::mediapipe::OkStatus();
}));
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
RET_CHECK(include_alpha)
<< "iOS GPU inference currently accepts only RGBA input.";
@@ -546,8 +543,6 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
// with normalization to either: [0,1] or [-1,1].
const std::string shader_source = absl::Substitute(
R"(
#include <simd/simd.h>
#include <metal_stdlib>
using namespace metal;
@@ -612,9 +607,9 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
// Get desired way to handle input channels.
max_num_channels_ = options.max_num_channels();
// Currently only alpha channel toggling is suppored.
CHECK_GE(max_num_channels_, 3);
CHECK_GE(max_num_channels_, 1);
CHECK_LE(max_num_channels_, 4);
CHECK_NE(max_num_channels_, 2);
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
if (cc->Inputs().HasTag("IMAGE_GPU"))
// Currently on iOS, tflite gpu input tensor must be 4 channels,
mediapipe/calculators/tflite/tflite_converter_calculator.proto
+1 -2
View File
@@ -36,8 +36,7 @@ message TfLiteConverterCalculatorOptions {
optional bool flip_vertically = 2 [default = false];
// Controls how many channels of the input image get passed through to the
// tensor. Currently this only controls whether or not to ignore alpha
// channel, so it must be 3 or 4.
// tensor. Valid values are 1,3,4 only. Ignored for iOS GPU.
optional int32 max_num_channels = 3 [default = 3];
// The calculator expects Matrix inputs to be in column-major order. Set
mediapipe/calculators/tflite/tflite_inference_calculator.cc
+144 -98
View File
@@ -12,10 +12,13 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <cstring>
#include <memory>
#include <string>
#include <vector>
#include "mediapipe/calculators/tflite/tflite_inference_calculator.pb.h"
#include "mediapipe/calculators/tflite/util.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/util/resource_util.h"
@@ -24,14 +27,15 @@
#include "tensorflow/lite/kernels/register.h"
#include "tensorflow/lite/model.h"
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
#import <CoreVideo/CoreVideo.h>
@@ -39,33 +43,42 @@
#import <MetalKit/MetalKit.h>
#import "mediapipe/gpu/MPPMetalHelper.h"
#include "mediapipe/gpu/MPPMetalUtil.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/metal/buffer_convert.h"
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
#endif // iOS
#if defined(__ANDROID__)
namespace {
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
typedef id<MTLBuffer> GpuTensor;
#endif
// Round up n to next multiple of m.
size_t RoundUp(size_t n, size_t m) { return ((n + m - 1) / m) * m; } // NOLINT
} // namespace
// TfLiteInferenceCalculator File Layout:
// * Header
// * Core
// * Aux
namespace mediapipe {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
using ::tflite::gpu::gl::CopyBuffer;
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
using ::tflite::gpu::gl::GlBuffer;
using ::tflite::gpu::gl::GlProgram;
using ::tflite::gpu::gl::GlShader;
#endif
#if !defined(MEDIAPIPE_DISABLE_GPU)
struct GPUData {
int elements = 1;
GlBuffer buffer;
};
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
struct GPUData {
int elements = 1;
id<MTLBuffer> buffer;
GpuTensor buffer;
::tflite::gpu::BHWC shape;
};
#endif
@@ -134,7 +147,7 @@ class TfLiteInferenceCalculator : public CalculatorBase {
std::unique_ptr<tflite::FlatBufferModel> model_;
TfLiteDelegate* delegate_ = nullptr;
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
mediapipe::GlCalculatorHelper gpu_helper_;
std::unique_ptr<GPUData> gpu_data_in_;
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
@@ -142,6 +155,7 @@ class TfLiteInferenceCalculator : public CalculatorBase {
MPPMetalHelper* gpu_helper_ = nullptr;
std::unique_ptr<GPUData> gpu_data_in_;
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
TFLBufferConvert* converter_from_BPHWC4_ = nil;
#endif
std::string model_path_ = "";
@@ -161,19 +175,25 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
RET_CHECK(cc->Outputs().HasTag("TENSORS") ^
cc->Outputs().HasTag("TENSORS_GPU"));
bool use_gpu = false;
if (cc->Inputs().HasTag("TENSORS"))
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
if (cc->Inputs().HasTag("TENSORS_GPU"))
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag("TENSORS_GPU")) {
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
#endif
use_gpu |= true;
}
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Outputs().HasTag("TENSORS"))
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
if (cc->Outputs().HasTag("TENSORS_GPU"))
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Outputs().HasTag("TENSORS_GPU")) {
cc->Outputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
#endif
use_gpu |= true;
}
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
cc->InputSidePackets()
@@ -181,11 +201,17 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
.Set<tflite::ops::builtin::BuiltinOpResolver>();
}
#if defined(__ANDROID__)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
const auto& options =
cc->Options<::mediapipe::TfLiteInferenceCalculatorOptions>();
use_gpu |= options.use_gpu();
if (use_gpu) {
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
#endif
}
// Assign this calculator's default InputStreamHandler.
cc->SetInputStreamHandler("FixedSizeInputStreamHandler");
@@ -199,35 +225,41 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
MP_RETURN_IF_ERROR(LoadOptions(cc));
if (cc->Inputs().HasTag("TENSORS_GPU")) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
gpu_input_ = true;
gpu_inference_ = true; // Inference must be on GPU also.
#else
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
#endif
RET_CHECK(!cc->Inputs().HasTag("TENSORS_GPU"))
<< "GPU processing not enabled.";
#endif // !MEDIAPIPE_DISABLE_GPU
}
if (cc->Outputs().HasTag("TENSORS_GPU")) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
gpu_output_ = true;
RET_CHECK(cc->Inputs().HasTag("TENSORS_GPU"))
<< "GPU output must also have GPU Input.";
#else
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
#endif
RET_CHECK(!cc->Inputs().HasTag("TENSORS_GPU"))
<< "GPU processing not enabled.";
#endif // !MEDIAPIPE_DISABLE_GPU
}
MP_RETURN_IF_ERROR(LoadModel(cc));
if (gpu_inference_) {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
RET_CHECK(gpu_helper_);
#endif
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &cc]() -> ::mediapipe::Status { return LoadDelegate(cc); }));
#else
MP_RETURN_IF_ERROR(LoadDelegate(cc));
#endif
}
return ::mediapipe::OkStatus();
@@ -237,35 +269,27 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
// 1. Receive pre-processed tensor inputs.
if (gpu_input_) {
// Read GPU input into SSBO.
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
const auto& input_tensors =
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
RET_CHECK_EQ(input_tensors.size(), 1);
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &input_tensors]() -> ::mediapipe::Status {
// Explicit copy input.
tflite::gpu::gl::CopyBuffer(input_tensors[0], gpu_data_in_->buffer);
RET_CHECK_CALL(CopyBuffer(input_tensors[0], gpu_data_in_->buffer));
return ::mediapipe::OkStatus();
}));
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
const auto& input_tensors =
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
RET_CHECK_EQ(input_tensors.size(), 1);
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
command_buffer.label = @"TfLiteInferenceCalculatorInput";
id<MTLBlitCommandEncoder> blit_command =
[command_buffer blitCommandEncoder];
// Explicit copy input.
[blit_command copyFromBuffer:input_tensors[0]
sourceOffset:0
toBuffer:gpu_data_in_->buffer
destinationOffset:0
size:gpu_data_in_->elements * sizeof(float)];
[blit_command endEncoding];
[command_buffer commit];
[command_buffer waitUntilCompleted];
[MPPMetalUtil blitMetalBufferTo:gpu_data_in_->buffer
from:input_tensors[0]
blocking:true
commandBuffer:[gpu_helper_ commandBuffer]];
#else
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
RET_CHECK_FAIL() << "GPU processing not enabled.";
#endif
} else {
// Read CPU input into tensors.
@@ -278,18 +302,20 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
if (use_quantized_tensors_) {
const uint8* input_tensor_buffer = input_tensor->data.uint8;
uint8* local_tensor_buffer = interpreter_->typed_input_tensor<uint8>(i);
memcpy(local_tensor_buffer, input_tensor_buffer, input_tensor->bytes);
std::memcpy(local_tensor_buffer, input_tensor_buffer,
input_tensor->bytes);
} else {
const float* input_tensor_buffer = input_tensor->data.f;
float* local_tensor_buffer = interpreter_->typed_input_tensor<float>(i);
memcpy(local_tensor_buffer, input_tensor_buffer, input_tensor->bytes);
std::memcpy(local_tensor_buffer, input_tensor_buffer,
input_tensor->bytes);
}
}
}
// 2. Run inference.
if (gpu_inference_) {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this]() -> ::mediapipe::Status {
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
@@ -304,52 +330,51 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
// 3. Output processed tensors.
if (gpu_output_) {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
// Output result tensors (GPU).
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
output_tensors->resize(gpu_data_out_.size());
for (int i = 0; i < gpu_data_out_.size(); ++i) {
GlBuffer& tensor = output_tensors->at(i);
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
auto status = CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_[i]->elements, &tensor);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
tflite::gpu::gl::CopyBuffer(gpu_data_out_[i]->buffer, tensor);
}
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &output_tensors]() -> ::mediapipe::Status {
output_tensors->resize(gpu_data_out_.size());
for (int i = 0; i < gpu_data_out_.size(); ++i) {
GpuTensor& tensor = output_tensors->at(i);
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_[i]->elements, &tensor));
RET_CHECK_CALL(CopyBuffer(gpu_data_out_[i]->buffer, tensor));
}
return ::mediapipe::OkStatus();
}));
cc->Outputs()
.Tag("TENSORS_GPU")
.Add(output_tensors.release(), cc->InputTimestamp());
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
// Output result tensors (GPU).
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
output_tensors->resize(gpu_data_out_.size());
id<MTLDevice> device = gpu_helper_.mtlDevice;
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
command_buffer.label = @"TfLiteInferenceCalculatorOutput";
command_buffer.label = @"TfLiteInferenceBPHWC4Convert";
id<MTLComputeCommandEncoder> convert_command =
[command_buffer computeCommandEncoder];
for (int i = 0; i < gpu_data_out_.size(); ++i) {
id<MTLBuffer> tensor =
output_tensors->at(i) =
[device newBufferWithLength:gpu_data_out_[i]->elements * sizeof(float)
options:MTLResourceStorageModeShared];
id<MTLBlitCommandEncoder> blit_command =
[command_buffer blitCommandEncoder];
// Explicit copy input.
[blit_command copyFromBuffer:gpu_data_out_[i]->buffer
sourceOffset:0
toBuffer:tensor
destinationOffset:0
size:gpu_data_out_[i]->elements * sizeof(float)];
[blit_command endEncoding];
[command_buffer commit];
[command_buffer waitUntilCompleted];
output_tensors->push_back(tensor);
// Reshape tensor.
[converter_from_BPHWC4_ convertWithEncoder:convert_command
shape:gpu_data_out_[i]->shape
sourceBuffer:gpu_data_out_[i]->buffer
convertedBuffer:output_tensors->at(i)];
}
[convert_command endEncoding];
[command_buffer commit];
[command_buffer waitUntilCompleted];
cc->Outputs()
.Tag("TENSORS_GPU")
.Add(output_tensors.release(), cc->InputTimestamp());
#else
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
#endif
RET_CHECK_FAIL() << "GPU processing not enabled.";
#endif // !MEDIAPIPE_DISABLE_GPU
} else {
// Output result tensors (CPU).
const auto& tensor_indexes = interpreter_->outputs();
@@ -367,7 +392,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
::mediapipe::Status TfLiteInferenceCalculator::Close(CalculatorContext* cc) {
if (delegate_) {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
TfLiteGpuDelegateDelete(delegate_);
gpu_data_in_.reset();
@@ -446,7 +471,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
::mediapipe::Status TfLiteInferenceCalculator::LoadDelegate(
CalculatorContext* cc) {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
// Configure and create the delegate.
TfLiteGpuDelegateOptions options = TfLiteGpuDelegateOptionsDefault();
options.compile_options.precision_loss_allowed = 1;
@@ -466,15 +491,12 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
for (int d = 0; d < tensor->dims->size; ++d) {
gpu_data_in_->elements *= tensor->dims->data[d];
}
// Input to model can be either RGB/RGBA only.
RET_CHECK_GE(tensor->dims->data[3], 3);
RET_CHECK_LE(tensor->dims->data[3], 4);
CHECK_GE(tensor->dims->data[3], 1);
CHECK_LE(tensor->dims->data[3], 4);
CHECK_NE(tensor->dims->data[3], 2);
// Create and bind input buffer.
auto status = ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
gpu_data_in_->elements, &gpu_data_in_->buffer);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
RET_CHECK_CALL(::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
gpu_data_in_->elements, &gpu_data_in_->buffer));
RET_CHECK_EQ(TfLiteGpuDelegateBindBufferToTensor(
delegate_, gpu_data_in_->buffer.id(),
interpreter_->inputs()[0]), // First tensor only
@@ -496,12 +518,8 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
// Create and bind output buffers.
interpreter_->SetAllowBufferHandleOutput(true);
for (int i = 0; i < gpu_data_out_.size(); ++i) {
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
auto status = CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_[i]->elements, &gpu_data_out_[i]->buffer);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
gpu_data_out_[i]->elements, &gpu_data_out_[i]->buffer));
RET_CHECK_EQ(
TfLiteGpuDelegateBindBufferToTensor(
delegate_, gpu_data_out_[i]->buffer.id(), output_indices[i]),
@@ -511,14 +529,15 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
// Must call this last.
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
#endif // __ANDROID__
#endif // OpenGL
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
// Configure and create the delegate.
GpuDelegateOptions options;
options.allow_precision_loss = false; // Must match converter, F=float/T=half
options.wait_type = GpuDelegateOptions::WaitType::kActive;
options.wait_type = GpuDelegateOptions::WaitType::kPassive;
if (!delegate_) delegate_ = TFLGpuDelegateCreate(&options);
id<MTLDevice> device = gpu_helper_.mtlDevice;
if (gpu_input_) {
// Get input image sizes.
@@ -539,11 +558,9 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
LOG(WARNING) << "Please ensure input GPU tensor is 4 channels.";
}
// Create and bind input buffer.
id<MTLDevice> device = gpu_helper_.mtlDevice;
gpu_data_in_->buffer =
[device newBufferWithLength:gpu_data_in_->elements * sizeof(float)
options:MTLResourceStorageModeShared];
// Must call this before TFLGpuDelegateBindMetalBufferToTensor.
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
delegate_,
@@ -561,12 +578,33 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
gpu_data_out_[i]->elements = 1;
// TODO handle *2 properly on some dialated models
for (int d = 0; d < tensor->dims->size; ++d) {
gpu_data_out_[i]->elements *= tensor->dims->data[d];
// Pad each dim for BHWC4 conversion inside delegate.
gpu_data_out_[i]->elements *= RoundUp(tensor->dims->data[d], 4);
}
// Save dimensions for reshaping back later.
gpu_data_out_[i]->shape.b = tensor->dims->data[0];
switch (tensor->dims->size) {
case 2:
gpu_data_out_[i]->shape.h = 1;
gpu_data_out_[i]->shape.w = 1;
gpu_data_out_[i]->shape.c = tensor->dims->data[1];
break;
case 3:
gpu_data_out_[i]->shape.h = 1;
gpu_data_out_[i]->shape.w = tensor->dims->data[1];
gpu_data_out_[i]->shape.c = tensor->dims->data[2];
break;
case 4:
gpu_data_out_[i]->shape.h = tensor->dims->data[1];
gpu_data_out_[i]->shape.w = tensor->dims->data[2];
gpu_data_out_[i]->shape.c = tensor->dims->data[3];
break;
default:
return mediapipe::InternalError("Unsupported tensor shape.");
}
}
// Create and bind output buffers.
interpreter_->SetAllowBufferHandleOutput(true);
id<MTLDevice> device = gpu_helper_.mtlDevice;
for (int i = 0; i < gpu_data_out_.size(); ++i) {
gpu_data_out_[i]->buffer =
[device newBufferWithLength:gpu_data_out_[i]->elements * sizeof(float)
@@ -575,6 +613,14 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
delegate_, output_indices[i], gpu_data_out_[i]->buffer),
true);
}
// Create converter for GPU output.
converter_from_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
isFloat16:false
convertToPBHWC4:false];
if (converter_from_BPHWC4_ == nil) {
return mediapipe::InternalError(
"Error initializating output buffer converter");
}
}
#endif // iOS
mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.cc
+495 -158
View File
@@ -18,6 +18,7 @@
#include "absl/strings/str_format.h"
#include "absl/types/span.h"
#include "mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.pb.h"
#include "mediapipe/calculators/tflite/util.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/deps/file_path.h"
#include "mediapipe/framework/formats/detection.pb.h"
@@ -26,28 +27,61 @@
#include "mediapipe/framework/port/ret_check.h"
#include "tensorflow/lite/interpreter.h"
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
#endif // ANDROID
#endif // !MEDIAPIPE_DISABLE_GPU
#if defined(__ANDROID__)
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
using ::tflite::gpu::gl::GlBuffer;
using ::tflite::gpu::gl::GlProgram;
using ::tflite::gpu::gl::GlShader;
#endif // ANDROID
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
#import <CoreVideo/CoreVideo.h>
#import <Metal/Metal.h>
#import <MetalKit/MetalKit.h>
namespace mediapipe {
#import "mediapipe/gpu/MPPMetalHelper.h"
#include "mediapipe/gpu/MPPMetalUtil.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
#endif // iOS
namespace {
constexpr int kNumInputTensorsWithAnchors = 3;
constexpr int kNumCoordsPerBox = 4;
} // namespace
namespace mediapipe {
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
using ::tflite::gpu::gl::GlShader;
#endif
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
typedef ::tflite::gpu::gl::GlProgram GpuProgram;
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
typedef id<MTLBuffer> GpuTensor;
typedef id<MTLComputePipelineState> GpuProgram;
#endif
namespace {
#if !defined(MEDIAPIPE_DISABLE_GPU)
struct GPUData {
GpuProgram decode_program;
GpuProgram score_program;
GpuTensor decoded_boxes_buffer;
GpuTensor raw_boxes_buffer;
GpuTensor raw_anchors_buffer;
GpuTensor scored_boxes_buffer;
GpuTensor raw_scores_buffer;
};
#endif
void ConvertRawValuesToAnchors(const float* raw_anchors, int num_boxes,
std::vector<Anchor>* anchors) {
anchors->clear();
@@ -88,7 +122,7 @@ void ConvertAnchorsToRawValues(const std::vector<Anchor>& anchors,
// optional to pass in a third tensor for anchors (e.g. for SSD
// models) depend on the outputs of the detection model. The size
// of anchor tensor must be (num_boxes * 4).
// TENSORS_GPU - vector of GlBuffer.
// TENSORS_GPU - vector of GlBuffer of MTLBuffer.
// Output:
// DETECTIONS - Result MediaPipe detections.
//
@@ -126,7 +160,7 @@ class TfLiteTensorsToDetectionsCalculator : public CalculatorBase {
std::vector<Detection>* output_detections);
::mediapipe::Status LoadOptions(CalculatorContext* cc);
::mediapipe::Status GlSetup(CalculatorContext* cc);
::mediapipe::Status GpuInit(CalculatorContext* cc);
::mediapipe::Status DecodeBoxes(const float* raw_boxes,
const std::vector<Anchor>& anchors,
std::vector<float>* boxes);
@@ -146,15 +180,12 @@ class TfLiteTensorsToDetectionsCalculator : public CalculatorBase {
std::vector<Anchor> anchors_;
bool side_packet_anchors_{};
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
mediapipe::GlCalculatorHelper gpu_helper_;
std::unique_ptr<GlProgram> decode_program_;
std::unique_ptr<GlProgram> score_program_;
std::unique_ptr<GlBuffer> decoded_boxes_buffer_;
std::unique_ptr<GlBuffer> raw_boxes_buffer_;
std::unique_ptr<GlBuffer> raw_anchors_buffer_;
std::unique_ptr<GlBuffer> scored_boxes_buffer_;
std::unique_ptr<GlBuffer> raw_scores_buffer_;
std::unique_ptr<GPUData> gpu_data_;
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
MPPMetalHelper* gpu_helper_ = nullptr;
std::unique_ptr<GPUData> gpu_data_;
#endif
bool gpu_input_ = false;
@@ -167,15 +198,18 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
RET_CHECK(!cc->Inputs().GetTags().empty());
RET_CHECK(!cc->Outputs().GetTags().empty());
bool use_gpu = false;
if (cc->Inputs().HasTag("TENSORS")) {
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
}
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag("TENSORS_GPU")) {
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GlBuffer>>();
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
use_gpu |= true;
}
#endif
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Outputs().HasTag("DETECTIONS")) {
cc->Outputs().Tag("DETECTIONS").Set<std::vector<Detection>>();
@@ -187,9 +221,13 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
}
}
#if defined(__ANDROID__)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
if (use_gpu) {
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
#endif
}
return ::mediapipe::OkStatus();
}
@@ -200,8 +238,11 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
if (cc->Inputs().HasTag("TENSORS_GPU")) {
gpu_input_ = true;
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
RET_CHECK(gpu_helper_);
#endif
}
@@ -209,7 +250,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
side_packet_anchors_ = cc->InputSidePackets().HasTag("ANCHORS");
if (gpu_input_) {
MP_RETURN_IF_ERROR(GlSetup(cc));
MP_RETURN_IF_ERROR(GpuInit(cc));
}
return ::mediapipe::OkStatus();
@@ -228,7 +269,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
MP_RETURN_IF_ERROR(ProcessGPU(cc, output_detections.get()));
} else {
MP_RETURN_IF_ERROR(ProcessCPU(cc, output_detections.get()));
} // if gpu_input_
}
// Output
if (cc->Outputs().HasTag("DETECTIONS")) {
@@ -245,7 +286,8 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
const auto& input_tensors =
cc->Inputs().Tag("TENSORS").Get<std::vector<TfLiteTensor>>();
if (input_tensors.size() == 2) {
if (input_tensors.size() == 2 ||
input_tensors.size() == kNumInputTensorsWithAnchors) {
// Postprocessing on CPU for model without postprocessing op. E.g. output
// raw score tensor and box tensor. Anchor decoding will be handled below.
const TfLiteTensor* raw_box_tensor = &input_tensors[0];
@@ -358,13 +400,84 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
}
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::ProcessGPU(
CalculatorContext* cc, std::vector<Detection>* output_detections) {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
const auto& input_tensors =
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GlBuffer>>();
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
RET_CHECK_GE(input_tensors.size(), 2);
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, &input_tensors, &cc,
&output_detections]()
-> ::mediapipe::Status {
// Copy inputs.
RET_CHECK_CALL(CopyBuffer(input_tensors[0], gpu_data_->raw_boxes_buffer));
RET_CHECK_CALL(CopyBuffer(input_tensors[1], gpu_data_->raw_scores_buffer));
if (!anchors_init_) {
if (side_packet_anchors_) {
CHECK(!cc->InputSidePackets().Tag("ANCHORS").IsEmpty());
const auto& anchors =
cc->InputSidePackets().Tag("ANCHORS").Get<std::vector<Anchor>>();
std::vector<float> raw_anchors(num_boxes_ * kNumCoordsPerBox);
ConvertAnchorsToRawValues(anchors, num_boxes_, raw_anchors.data());
RET_CHECK_CALL(gpu_data_->raw_anchors_buffer.Write<float>(
absl::MakeSpan(raw_anchors)));
} else {
CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
RET_CHECK_CALL(
CopyBuffer(input_tensors[2], gpu_data_->raw_anchors_buffer));
}
anchors_init_ = true;
}
// Run shaders.
// Decode boxes.
RET_CHECK_CALL(gpu_data_->decoded_boxes_buffer.BindToIndex(0));
RET_CHECK_CALL(gpu_data_->raw_boxes_buffer.BindToIndex(1));
RET_CHECK_CALL(gpu_data_->raw_anchors_buffer.BindToIndex(2));
const tflite::gpu::uint3 decode_workgroups = {num_boxes_, 1, 1};
RET_CHECK_CALL(gpu_data_->decode_program.Dispatch(decode_workgroups));
// Score boxes.
RET_CHECK_CALL(gpu_data_->scored_boxes_buffer.BindToIndex(0));
RET_CHECK_CALL(gpu_data_->raw_scores_buffer.BindToIndex(1));
const tflite::gpu::uint3 score_workgroups = {num_boxes_, 1, 1};
RET_CHECK_CALL(gpu_data_->score_program.Dispatch(score_workgroups));
// Copy decoded boxes from GPU to CPU.
std::vector<float> boxes(num_boxes_ * num_coords_);
RET_CHECK_CALL(gpu_data_->decoded_boxes_buffer.Read(absl::MakeSpan(boxes)));
std::vector<float> score_class_id_pairs(num_boxes_ * 2);
RET_CHECK_CALL(gpu_data_->scored_boxes_buffer.Read(
absl::MakeSpan(score_class_id_pairs)));
// TODO: b/138851969. Is it possible to output a float vector
// for score and an int vector for class so that we can avoid copying twice?
std::vector<float> detection_scores(num_boxes_);
std::vector<int> detection_classes(num_boxes_);
for (int i = 0; i < num_boxes_; ++i) {
detection_scores[i] = score_class_id_pairs[i * 2];
detection_classes[i] = static_cast<int>(score_class_id_pairs[i * 2 + 1]);
}
MP_RETURN_IF_ERROR(
ConvertToDetections(boxes.data(), detection_scores.data(),
detection_classes.data(), output_detections));
return ::mediapipe::OkStatus();
}));
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
const auto& input_tensors =
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
RET_CHECK_GE(input_tensors.size(), 2);
// Copy inputs.
tflite::gpu::gl::CopyBuffer(input_tensors[0], *raw_boxes_buffer_.get());
tflite::gpu::gl::CopyBuffer(input_tensors[1], *raw_scores_buffer_.get());
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_boxes_buffer
from:input_tensors[0]
blocking:true
commandBuffer:[gpu_helper_ commandBuffer]];
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_scores_buffer
from:input_tensors[1]
blocking:true
commandBuffer:[gpu_helper_ commandBuffer]];
if (!anchors_init_) {
if (side_packet_anchors_) {
CHECK(!cc->InputSidePackets().Tag("ANCHORS").IsEmpty());
@@ -372,47 +485,65 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
cc->InputSidePackets().Tag("ANCHORS").Get<std::vector<Anchor>>();
std::vector<float> raw_anchors(num_boxes_ * kNumCoordsPerBox);
ConvertAnchorsToRawValues(anchors, num_boxes_, raw_anchors.data());
raw_anchors_buffer_->Write<float>(absl::MakeSpan(raw_anchors));
memcpy([gpu_data_->raw_anchors_buffer contents], raw_anchors.data(),
raw_anchors.size() * sizeof(float));
} else {
CHECK_EQ(input_tensors.size(), 3);
tflite::gpu::gl::CopyBuffer(input_tensors[2], *raw_anchors_buffer_.get());
RET_CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_anchors_buffer
from:input_tensors[2]
blocking:true
commandBuffer:[gpu_helper_ commandBuffer]];
}
anchors_init_ = true;
}
// Run shaders.
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &input_tensors]() -> ::mediapipe::Status {
// Decode boxes.
decoded_boxes_buffer_->BindToIndex(0);
raw_boxes_buffer_->BindToIndex(1);
raw_anchors_buffer_->BindToIndex(2);
const tflite::gpu::uint3 decode_workgroups = {num_boxes_, 1, 1};
decode_program_->Dispatch(decode_workgroups);
// Score boxes.
scored_boxes_buffer_->BindToIndex(0);
raw_scores_buffer_->BindToIndex(1);
const tflite::gpu::uint3 score_workgroups = {num_boxes_, 1, 1};
score_program_->Dispatch(score_workgroups);
return ::mediapipe::OkStatus();
}));
{
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
command_buffer.label = @"TfLiteDecodeBoxes";
id<MTLComputeCommandEncoder> decode_command =
[command_buffer computeCommandEncoder];
[decode_command setComputePipelineState:gpu_data_->decode_program];
[decode_command setBuffer:gpu_data_->decoded_boxes_buffer
offset:0
atIndex:0];
[decode_command setBuffer:gpu_data_->raw_boxes_buffer offset:0 atIndex:1];
[decode_command setBuffer:gpu_data_->raw_anchors_buffer offset:0 atIndex:2];
MTLSize decode_threads_per_group = MTLSizeMake(1, 1, 1);
MTLSize decode_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
[decode_command dispatchThreadgroups:decode_threadgroups
threadsPerThreadgroup:decode_threads_per_group];
[decode_command endEncoding];
[command_buffer commit];
[command_buffer waitUntilCompleted];
}
{
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
command_buffer.label = @"TfLiteScoreBoxes";
id<MTLComputeCommandEncoder> score_command =
[command_buffer computeCommandEncoder];
[score_command setComputePipelineState:gpu_data_->score_program];
[score_command setBuffer:gpu_data_->scored_boxes_buffer offset:0 atIndex:0];
[score_command setBuffer:gpu_data_->raw_scores_buffer offset:0 atIndex:1];
MTLSize score_threads_per_group = MTLSizeMake(1, num_classes_, 1);
MTLSize score_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
[score_command dispatchThreadgroups:score_threadgroups
threadsPerThreadgroup:score_threads_per_group];
[score_command endEncoding];
[command_buffer commit];
[command_buffer waitUntilCompleted];
}
// Copy decoded boxes from GPU to CPU.
std::vector<float> boxes(num_boxes_ * num_coords_);
auto status = decoded_boxes_buffer_->Read(absl::MakeSpan(boxes));
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
memcpy(boxes.data(), [gpu_data_->decoded_boxes_buffer contents],
num_boxes_ * num_coords_ * sizeof(float));
std::vector<float> score_class_id_pairs(num_boxes_ * 2);
status = scored_boxes_buffer_->Read(absl::MakeSpan(score_class_id_pairs));
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
memcpy(score_class_id_pairs.data(), [gpu_data_->scored_boxes_buffer contents],
num_boxes_ * 2 * sizeof(float));
// TODO: b/138851969. Is it possible to output a float vector
// for score and an int vector for class so that we can avoid copying twice?
// Output detections.
// TODO Adjust shader to avoid copying shader output twice.
std::vector<float> detection_scores(num_boxes_);
std::vector<int> detection_classes(num_boxes_);
for (int i = 0; i < num_boxes_; ++i) {
@@ -422,25 +553,20 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
MP_RETURN_IF_ERROR(ConvertToDetections(boxes.data(), detection_scores.data(),
detection_classes.data(),
output_detections));
#else
LOG(ERROR) << "GPU input on non-Android not supported yet.";
#endif // defined(__ANDROID__)
#endif
return ::mediapipe::OkStatus();
}
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::Close(
CalculatorContext* cc) {
#if defined(__ANDROID__)
gpu_helper_.RunInGlContext([this] {
decode_program_.reset();
score_program_.reset();
decoded_boxes_buffer_.reset();
raw_boxes_buffer_.reset();
raw_anchors_buffer_.reset();
scored_boxes_buffer_.reset();
raw_scores_buffer_.reset();
});
#endif // __ANDROID__
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
gpu_helper_.RunInGlContext([this] { gpu_data_.reset(); });
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
gpu_data_.reset();
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
@@ -530,6 +656,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
}
}
}
return ::mediapipe::OkStatus();
}
@@ -586,12 +713,16 @@ Detection TfLiteTensorsToDetectionsCalculator::ConvertToDetection(
return detection;
}
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::GlSetup(
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::GpuInit(
CalculatorContext* cc) {
#if defined(__ANDROID__)
// A shader to decode detection boxes.
const std::string decode_src = absl::Substitute(
R"( #version 310 es
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]()
-> ::mediapipe::Status {
gpu_data_ = absl::make_unique<GPUData>();
// A shader to decode detection boxes.
const std::string decode_src = absl::Substitute(
R"( #version 310 es
layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
@@ -665,7 +796,7 @@ void main() {
if (num_keypoints > int(0)){
for (int k = 0; k < num_keypoints; ++k) {
int kp_offset =
int(g_idx * num_coords) + keypt_coord_offset + k * num_values_per_keypt;
int(g_idx * num_coords) + keypt_coord_offset + k * num_values_per_keypt;
float kp_y, kp_x;
if (reverse_output_order == int(0)) {
kp_y = raw_boxes.data[kp_offset + int(0)];
@@ -679,55 +810,37 @@ void main() {
}
}
})",
options_.num_coords(), // box xywh
options_.reverse_output_order() ? 1 : 0,
options_.apply_exponential_on_box_size() ? 1 : 0,
options_.box_coord_offset(), options_.num_keypoints(),
options_.keypoint_coord_offset(), options_.num_values_per_keypoint());
options_.num_coords(), // box xywh
options_.reverse_output_order() ? 1 : 0,
options_.apply_exponential_on_box_size() ? 1 : 0,
options_.box_coord_offset(), options_.num_keypoints(),
options_.keypoint_coord_offset(), options_.num_values_per_keypoint());
// Shader program
GlShader decode_shader;
auto status =
GlShader::CompileShader(GL_COMPUTE_SHADER, decode_src, &decode_shader);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
decode_program_ = absl::make_unique<GlProgram>();
status = GlProgram::CreateWithShader(decode_shader, decode_program_.get());
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
// Outputs
size_t decoded_boxes_length = num_boxes_ * num_coords_;
decoded_boxes_buffer_ = absl::make_unique<GlBuffer>();
status = CreateReadWriteShaderStorageBuffer<float>(
decoded_boxes_length, decoded_boxes_buffer_.get());
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
// Inputs
size_t raw_boxes_length = num_boxes_ * num_coords_;
raw_boxes_buffer_ = absl::make_unique<GlBuffer>();
status = CreateReadWriteShaderStorageBuffer<float>(raw_boxes_length,
raw_boxes_buffer_.get());
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
size_t raw_anchors_length = num_boxes_ * kNumCoordsPerBox;
raw_anchors_buffer_ = absl::make_unique<GlBuffer>();
status = CreateReadWriteShaderStorageBuffer<float>(raw_anchors_length,
raw_anchors_buffer_.get());
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
// Parameters
glUseProgram(decode_program_->id());
glUniform4f(0, options_.x_scale(), options_.y_scale(), options_.w_scale(),
options_.h_scale());
// Shader program
GlShader decode_shader;
RET_CHECK_CALL(
GlShader::CompileShader(GL_COMPUTE_SHADER, decode_src, &decode_shader));
RET_CHECK_CALL(GpuProgram::CreateWithShader(decode_shader,
&gpu_data_->decode_program));
// Outputs
size_t decoded_boxes_length = num_boxes_ * num_coords_;
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
decoded_boxes_length, &gpu_data_->decoded_boxes_buffer));
// Inputs
size_t raw_boxes_length = num_boxes_ * num_coords_;
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
raw_boxes_length, &gpu_data_->raw_boxes_buffer));
size_t raw_anchors_length = num_boxes_ * kNumCoordsPerBox;
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
raw_anchors_length, &gpu_data_->raw_anchors_buffer));
// Parameters
glUseProgram(gpu_data_->decode_program.id());
glUniform4f(0, options_.x_scale(), options_.y_scale(), options_.w_scale(),
options_.h_scale());
// A shader to score detection boxes.
const std::string score_src = absl::Substitute(
R"( #version 310 es
// A shader to score detection boxes.
const std::string score_src = absl::Substitute(
R"( #version 310 es
layout(local_size_x = 1, local_size_y = $0, local_size_z = 1) in;
@@ -781,6 +894,228 @@ void main() {
scored_boxes.data[g_idx * uint(2) + uint(0)] = max_score;
scored_boxes.data[g_idx * uint(2) + uint(1)] = max_class;
}
})",
num_classes_, options_.sigmoid_score() ? 1 : 0,
options_.has_score_clipping_thresh() ? 1 : 0,
options_.has_score_clipping_thresh() ? options_.score_clipping_thresh()
: 0,
!ignore_classes_.empty() ? 1 : 0);
// # filter classes supported is hardware dependent.
int max_wg_size; // typically <= 1024
glGetIntegeri_v(GL_MAX_COMPUTE_WORK_GROUP_SIZE, 1,
&max_wg_size); // y-dim
CHECK_LT(num_classes_, max_wg_size)
<< "# classes must be < " << max_wg_size;
// TODO support better filtering.
CHECK_LE(ignore_classes_.size(), 1) << "Only ignore class 0 is allowed";
// Shader program
GlShader score_shader;
RET_CHECK_CALL(
GlShader::CompileShader(GL_COMPUTE_SHADER, score_src, &score_shader));
RET_CHECK_CALL(
GpuProgram::CreateWithShader(score_shader, &gpu_data_->score_program));
// Outputs
size_t scored_boxes_length = num_boxes_ * 2; // score, class
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
scored_boxes_length, &gpu_data_->scored_boxes_buffer));
// Inputs
size_t raw_scores_length = num_boxes_ * num_classes_;
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
raw_scores_length, &gpu_data_->raw_scores_buffer));
return ::mediapipe::OkStatus();
}));
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
// TODO consolidate Metal and OpenGL shaders via vulkan.
gpu_data_ = absl::make_unique<GPUData>();
id<MTLDevice> device = gpu_helper_.mtlDevice;
// A shader to decode detection boxes.
std::string decode_src = absl::Substitute(
R"(
#include <metal_stdlib>
using namespace metal;
kernel void decodeKernel(
device float* boxes [[ buffer(0) ]],
device float* raw_boxes [[ buffer(1) ]],
device float* raw_anchors [[ buffer(2) ]],
uint2 gid [[ thread_position_in_grid ]]) {
uint num_coords = uint($0);
int reverse_output_order = int($1);
int apply_exponential = int($2);
int box_coord_offset = int($3);
int num_keypoints = int($4);
int keypt_coord_offset = int($5);
int num_values_per_keypt = int($6);
)",
options_.num_coords(), // box xywh
options_.reverse_output_order() ? 1 : 0,
options_.apply_exponential_on_box_size() ? 1 : 0,
options_.box_coord_offset(), options_.num_keypoints(),
options_.keypoint_coord_offset(), options_.num_values_per_keypoint());
decode_src += absl::Substitute(
R"(
float4 scale = float4(($0),($1),($2),($3));
)",
options_.x_scale(), options_.y_scale(), options_.w_scale(),
options_.h_scale());
decode_src += R"(
uint g_idx = gid.x;
uint box_offset = g_idx * num_coords + uint(box_coord_offset);
uint anchor_offset = g_idx * uint(4); // check kNumCoordsPerBox
float y_center, x_center, h, w;
if (reverse_output_order == int(0)) {
y_center = raw_boxes[box_offset + uint(0)];
x_center = raw_boxes[box_offset + uint(1)];
h = raw_boxes[box_offset + uint(2)];
w = raw_boxes[box_offset + uint(3)];
} else {
x_center = raw_boxes[box_offset + uint(0)];
y_center = raw_boxes[box_offset + uint(1)];
w = raw_boxes[box_offset + uint(2)];
h = raw_boxes[box_offset + uint(3)];
}
float anchor_yc = raw_anchors[anchor_offset + uint(0)];
float anchor_xc = raw_anchors[anchor_offset + uint(1)];
float anchor_h = raw_anchors[anchor_offset + uint(2)];
float anchor_w = raw_anchors[anchor_offset + uint(3)];
x_center = x_center / scale.x * anchor_w + anchor_xc;
y_center = y_center / scale.y * anchor_h + anchor_yc;
if (apply_exponential == int(1)) {
h = exp(h / scale.w) * anchor_h;
w = exp(w / scale.z) * anchor_w;
} else {
h = (h / scale.w) * anchor_h;
w = (w / scale.z) * anchor_w;
}
float ymin = y_center - h / 2.0;
float xmin = x_center - w / 2.0;
float ymax = y_center + h / 2.0;
float xmax = x_center + w / 2.0;
boxes[box_offset + uint(0)] = ymin;
boxes[box_offset + uint(1)] = xmin;
boxes[box_offset + uint(2)] = ymax;
boxes[box_offset + uint(3)] = xmax;
if (num_keypoints > int(0)){
for (int k = 0; k < num_keypoints; ++k) {
int kp_offset =
int(g_idx * num_coords) + keypt_coord_offset + k * num_values_per_keypt;
float kp_y, kp_x;
if (reverse_output_order == int(0)) {
kp_y = raw_boxes[kp_offset + int(0)];
kp_x = raw_boxes[kp_offset + int(1)];
} else {
kp_x = raw_boxes[kp_offset + int(0)];
kp_y = raw_boxes[kp_offset + int(1)];
}
boxes[kp_offset + int(0)] = kp_x / scale.x * anchor_w + anchor_xc;
boxes[kp_offset + int(1)] = kp_y / scale.y * anchor_h + anchor_yc;
}
}
})";
{
// Shader program
NSString* library_source =
[NSString stringWithUTF8String:decode_src.c_str()];
NSError* error = nil;
id<MTLLibrary> library = [device newLibraryWithSource:library_source
options:nullptr
error:&error];
RET_CHECK(library != nil) << "Couldn't create shader library "
<< [[error localizedDescription] UTF8String];
id<MTLFunction> kernel_func = nil;
kernel_func = [library newFunctionWithName:@"decodeKernel"];
RET_CHECK(kernel_func != nil) << "Couldn't create kernel function.";
gpu_data_->decode_program =
[device newComputePipelineStateWithFunction:kernel_func error:&error];
RET_CHECK(gpu_data_->decode_program != nil)
<< "Couldn't create pipeline state "
<< [[error localizedDescription] UTF8String];
// Outputs
size_t decoded_boxes_length = num_boxes_ * num_coords_ * sizeof(float);
gpu_data_->decoded_boxes_buffer =
[device newBufferWithLength:decoded_boxes_length
options:MTLResourceStorageModeShared];
// Inputs
size_t raw_boxes_length = num_boxes_ * num_coords_ * sizeof(float);
gpu_data_->raw_boxes_buffer =
[device newBufferWithLength:raw_boxes_length
options:MTLResourceStorageModeShared];
size_t raw_anchors_length = num_boxes_ * kNumCoordsPerBox * sizeof(float);
gpu_data_->raw_anchors_buffer =
[device newBufferWithLength:raw_anchors_length
options:MTLResourceStorageModeShared];
}
// A shader to score detection boxes.
const std::string score_src = absl::Substitute(
R"(
#include <metal_stdlib>
using namespace metal;
float optional_sigmoid(float x) {
int apply_sigmoid = int($1);
int apply_clipping_thresh = int($2);
float clipping_thresh = float($3);
if (apply_sigmoid == int(0)) return x;
if (apply_clipping_thresh == int(1)) {
x = clamp(x, -clipping_thresh, clipping_thresh);
}
x = 1.0 / (1.0 + exp(-x));
return x;
}
kernel void scoreKernel(
device float* scored_boxes [[ buffer(0) ]],
device float* raw_scores [[ buffer(1) ]],
uint2 tid [[ thread_position_in_threadgroup ]],
uint2 gid [[ thread_position_in_grid ]]) {
uint num_classes = uint($0);
int apply_sigmoid = int($1);
int apply_clipping_thresh = int($2);
float clipping_thresh = float($3);
int ignore_class_0 = int($4);
uint g_idx = gid.x; // box idx
uint s_idx = tid.y; // score/class idx
// load all scores into shared memory
threadgroup float local_scores[$0];
float score = raw_scores[g_idx * num_classes + s_idx];
local_scores[s_idx] = optional_sigmoid(score);
threadgroup_barrier(mem_flags::mem_threadgroup);
// find max score in shared memory
if (s_idx == uint(0)) {
float max_score = -FLT_MAX;
float max_class = -1.0;
for (int i=ignore_class_0; i<int(num_classes); ++i) {
if (local_scores[i] > max_score) {
max_score = local_scores[i];
max_class = float(i);
}
}
scored_boxes[g_idx * uint(2) + uint(0)] = max_score;
scored_boxes[g_idx * uint(2) + uint(1)] = max_class;
}
})",
num_classes_, options_.sigmoid_score() ? 1 : 0,
options_.has_score_clipping_thresh() ? 1 : 0,
@@ -788,42 +1123,44 @@ void main() {
: 0,
ignore_classes_.size() ? 1 : 0);
// # filter classes supported is hardware dependent.
int max_wg_size; // typically <= 1024
glGetIntegeri_v(GL_MAX_COMPUTE_WORK_GROUP_SIZE, 1, &max_wg_size); // y-dim
CHECK_LT(num_classes_, max_wg_size) << "# classes must be < " << max_wg_size;
// TODO support better filtering.
CHECK_LE(ignore_classes_.size(), 1) << "Only ignore class 0 is allowed";
// Shader program
GlShader score_shader;
status = GlShader::CompileShader(GL_COMPUTE_SHADER, score_src, &score_shader);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
score_program_ = absl::make_unique<GlProgram>();
status = GlProgram::CreateWithShader(score_shader, score_program_.get());
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
// Outputs
size_t scored_boxes_length = num_boxes_ * 2; // score, class
scored_boxes_buffer_ = absl::make_unique<GlBuffer>();
status = CreateReadWriteShaderStorageBuffer<float>(
scored_boxes_length, scored_boxes_buffer_.get());
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
// Inputs
size_t raw_scores_length = num_boxes_ * num_classes_;
raw_scores_buffer_ = absl::make_unique<GlBuffer>();
status = CreateReadWriteShaderStorageBuffer<float>(raw_scores_length,
raw_scores_buffer_.get());
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
{
// Shader program
NSString* library_source =
[NSString stringWithUTF8String:score_src.c_str()];
NSError* error = nil;
id<MTLLibrary> library = [device newLibraryWithSource:library_source
options:nullptr
error:&error];
RET_CHECK(library != nil) << "Couldn't create shader library "
<< [[error localizedDescription] UTF8String];
id<MTLFunction> kernel_func = nil;
kernel_func = [library newFunctionWithName:@"scoreKernel"];
RET_CHECK(kernel_func != nil) << "Couldn't create kernel function.";
gpu_data_->score_program =
[device newComputePipelineStateWithFunction:kernel_func error:&error];
RET_CHECK(gpu_data_->score_program != nil)
<< "Couldn't create pipeline state "
<< [[error localizedDescription] UTF8String];
// Outputs
size_t scored_boxes_length = num_boxes_ * 2 * sizeof(float); // score,class
gpu_data_->scored_boxes_buffer =
[device newBufferWithLength:scored_boxes_length
options:MTLResourceStorageModeShared];
// Inputs
size_t raw_scores_length = num_boxes_ * num_classes_ * sizeof(float);
gpu_data_->raw_scores_buffer =
[device newBufferWithLength:raw_scores_length
options:MTLResourceStorageModeShared];
// # filter classes supported is hardware dependent.
int max_wg_size = gpu_data_->score_program.maxTotalThreadsPerThreadgroup;
CHECK_LT(num_classes_, max_wg_size) << "# classes must be <" << max_wg_size;
}
#endif // defined(__ANDROID__)
#endif // __ANDROID__ or iOS
return ::mediapipe::OkStatus();
}
mediapipe/calculators/tflite/tflite_tensors_to_landmarks_calculator.cc
+8 -2
View File
@@ -96,7 +96,8 @@ REGISTER_CALCULATOR(TfLiteTensorsToLandmarksCalculator);
options_.has_input_image_width())
<< "Must provide input with/height for getting normalized landmarks.";
}
if (cc->Outputs().HasTag("LANDMARKS") && options_.flip_vertically()) {
if (cc->Outputs().HasTag("LANDMARKS") &&
(options_.flip_vertically() || options_.flip_horizontally())) {
RET_CHECK(options_.has_input_image_height() &&
options_.has_input_image_width())
<< "Must provide input with/height for using flip_vertically option "
@@ -133,7 +134,12 @@ REGISTER_CALCULATOR(TfLiteTensorsToLandmarksCalculator);
for (int ld = 0; ld < num_landmarks_; ++ld) {
const int offset = ld * num_dimensions;
Landmark landmark;
landmark.set_x(raw_landmarks[offset]);
if (options_.flip_horizontally()) {
landmark.set_x(options_.input_image_width() - raw_landmarks[offset]);
} else {
landmark.set_x(raw_landmarks[offset]);
}
if (num_dimensions > 1) {
if (options_.flip_vertically()) {
landmark.set_y(options_.input_image_height() -
mediapipe/calculators/tflite/tflite_tensors_to_landmarks_calculator.proto
+6
View File
@@ -40,6 +40,12 @@ message TfLiteTensorsToLandmarksCalculatorOptions {
// representation has a bottom-left origin (e.g., in OpenGL).
optional bool flip_vertically = 4 [default = false];
// Whether the detection coordinates from the input tensors should be flipped
// horizontally (along the x-direction). This is useful, for example, when the
// input image is horizontally flipped in ImageTransformationCalculator
// beforehand.
optional bool flip_horizontally = 6 [default = false];
// A value that z values should be divided by.
optional float normalize_z = 5 [default = 1.0];
}
mediapipe/calculators/tflite/tflite_tensors_to_segmentation_calculator.cc
+113 -116
View File
@@ -17,6 +17,7 @@
#include "absl/strings/str_format.h"
#include "absl/types/span.h"
#include "mediapipe/calculators/tflite/tflite_tensors_to_segmentation_calculator.pb.h"
#include "mediapipe/calculators/tflite/util.h"
#include "mediapipe/framework/calculator_context.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_frame.h"
@@ -27,7 +28,7 @@
#include "mediapipe/util/resource_util.h"
#include "tensorflow/lite/interpreter.h"
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gl_simple_shaders.h"
#include "mediapipe/gpu/shader_util.h"
@@ -36,7 +37,7 @@
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_texture.h"
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
namespace {
constexpr int kWorkgroupSize = 8; // Block size for GPU shader.
@@ -52,12 +53,14 @@ float Clamp(float val, float min, float max) {
namespace mediapipe {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
using ::tflite::gpu::gl::CopyBuffer;
using ::tflite::gpu::gl::CreateReadWriteRgbaImageTexture;
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
using ::tflite::gpu::gl::GlBuffer;
using ::tflite::gpu::gl::GlProgram;
using ::tflite::gpu::gl::GlShader;
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
// Converts TFLite tensors from a tflite segmentation model to an image mask.
//
@@ -126,13 +129,13 @@ class TfLiteTensorsToSegmentationCalculator : public CalculatorBase {
int tensor_channels_ = 0;
bool use_gpu_ = false;
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
mediapipe::GlCalculatorHelper gpu_helper_;
std::unique_ptr<GlProgram> mask_program_with_prev_;
std::unique_ptr<GlProgram> mask_program_no_prev_;
std::unique_ptr<GlBuffer> tensor_buffer_;
GLuint upsample_program_;
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
};
REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
@@ -142,6 +145,8 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
RET_CHECK(!cc->Inputs().GetTags().empty());
RET_CHECK(!cc->Outputs().GetTags().empty());
bool use_gpu = false;
// Inputs CPU.
if (cc->Inputs().HasTag("TENSORS")) {
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
@@ -154,32 +159,37 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
}
// Inputs GPU.
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
if (cc->Inputs().HasTag("TENSORS_GPU")) {
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GlBuffer>>();
use_gpu |= true;
}
if (cc->Inputs().HasTag("PREV_MASK_GPU")) {
cc->Inputs().Tag("PREV_MASK_GPU").Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
if (cc->Inputs().HasTag("REFERENCE_IMAGE_GPU")) {
cc->Inputs().Tag("REFERENCE_IMAGE_GPU").Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
// Outputs.
if (cc->Outputs().HasTag("MASK")) {
cc->Outputs().Tag("MASK").Set<ImageFrame>();
}
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
if (cc->Outputs().HasTag("MASK_GPU")) {
cc->Outputs().Tag("MASK_GPU").Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__
#if defined(__ANDROID__)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
if (use_gpu) {
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
}
@@ -189,24 +199,23 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
if (cc->Inputs().HasTag("TENSORS_GPU")) {
use_gpu_ = true;
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
}
MP_RETURN_IF_ERROR(LoadOptions(cc));
if (use_gpu_) {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
MP_RETURN_IF_ERROR(InitGpu(cc));
return ::mediapipe::OkStatus();
}));
#else
RET_CHECK_FAIL()
<< "GPU processing on non-Android devices is not supported yet.";
#endif // __ANDROID__
RET_CHECK_FAIL() << "GPU processing not enabled.";
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
@@ -215,13 +224,13 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::Process(
CalculatorContext* cc) {
if (use_gpu_) {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
MP_RETURN_IF_ERROR(ProcessGpu(cc));
return ::mediapipe::OkStatus();
}));
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
} else {
MP_RETURN_IF_ERROR(ProcessCpu(cc));
}
@@ -231,7 +240,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::Close(
CalculatorContext* cc) {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
gpu_helper_.RunInGlContext([this] {
if (upsample_program_) glDeleteProgram(upsample_program_);
upsample_program_ = 0;
@@ -239,7 +248,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
mask_program_no_prev_.reset();
tensor_buffer_.reset();
});
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
@@ -358,7 +367,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
if (cc->Inputs().Tag("TENSORS_GPU").IsEmpty()) {
return ::mediapipe::OkStatus();
}
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
// Get input streams.
const auto& input_tensors =
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GlBuffer>>();
@@ -379,9 +388,9 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
// Create initial working mask texture.
::tflite::gpu::gl::GlTexture small_mask_texture;
::tflite::gpu::gl::CreateReadWriteRgbaImageTexture(
RET_CHECK_CALL(CreateReadWriteRgbaImageTexture(
tflite::gpu::DataType::UINT8, // GL_RGBA8
{tensor_width_, tensor_height_}, &small_mask_texture);
{tensor_width_, tensor_height_}, &small_mask_texture));
// Get input previous mask.
auto input_mask_texture = has_prev_mask
@@ -389,7 +398,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
: mediapipe::GlTexture();
// Copy input tensor.
tflite::gpu::gl::CopyBuffer(input_tensors[0], *tensor_buffer_);
RET_CHECK_CALL(CopyBuffer(input_tensors[0], *tensor_buffer_));
// Run shader, process mask tensor.
// Run softmax over tensor output and blend with previous mask.
@@ -397,18 +406,18 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
const int output_index = 0;
glBindImageTexture(output_index, small_mask_texture.id(), 0, GL_FALSE, 0,
GL_WRITE_ONLY, GL_RGBA8);
tensor_buffer_->BindToIndex(2);
RET_CHECK_CALL(tensor_buffer_->BindToIndex(2));
const tflite::gpu::uint3 workgroups = {
NumGroups(tensor_width_, kWorkgroupSize),
NumGroups(tensor_height_, kWorkgroupSize), 1};
if (!has_prev_mask) {
mask_program_no_prev_->Dispatch(workgroups);
RET_CHECK_CALL(mask_program_no_prev_->Dispatch(workgroups));
} else {
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D, input_mask_texture.name());
mask_program_with_prev_->Dispatch(workgroups);
RET_CHECK_CALL(mask_program_with_prev_->Dispatch(workgroups));
glActiveTexture(GL_TEXTURE1);
glBindTexture(GL_TEXTURE_2D, 0);
}
@@ -438,13 +447,13 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
// Cleanup
input_mask_texture.Release();
output_texture.Release();
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
void TfLiteTensorsToSegmentationCalculator::GlRender() {
#if defined(__ANDROID__)
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
static const GLfloat square_vertices[] = {
-1.0f, -1.0f, // bottom left
1.0f, -1.0f, // bottom right
@@ -492,7 +501,7 @@ void TfLiteTensorsToSegmentationCalculator::GlRender() {
glBindVertexArray(0);
glDeleteVertexArrays(1, &vao);
glDeleteBuffers(2, vbo);
#endif // __ANDROID__
#endif // !MEDIAPIPE_DISABLE_GPU
}
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::LoadOptions(
@@ -516,14 +525,15 @@ void TfLiteTensorsToSegmentationCalculator::GlRender() {
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::InitGpu(
CalculatorContext* cc) {
#if defined(__ANDROID__)
// A shader to process a segmentation tensor into an output mask,
// and use an optional previous mask as input.
// Currently uses 4 channels for output,
// and sets both R and A channels as mask value.
const std::string shader_src_template =
R"( #version 310 es
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]()
-> ::mediapipe::Status {
// A shader to process a segmentation tensor into an output mask,
// and use an optional previous mask as input.
// Currently uses 4 channels for output,
// and sets both R and A channels as mask value.
const std::string shader_src_template =
R"( #version 310 es
layout(local_size_x = $0, local_size_y = $0, local_size_z = 1) in;
@@ -589,76 +599,60 @@ void main() {
imageStore(output_texture, output_coordinate, out_value);
})";
const std::string shader_src_no_previous = absl::Substitute(
shader_src_template, kWorkgroupSize, options_.output_layer_index(),
options_.combine_with_previous_ratio(), "",
options_.flip_vertically() ? "out_height - gid.y - 1" : "gid.y");
const std::string shader_src_with_previous = absl::Substitute(
shader_src_template, kWorkgroupSize, options_.output_layer_index(),
options_.combine_with_previous_ratio(), "#define READ_PREVIOUS",
options_.flip_vertically() ? "out_height - gid.y - 1" : "gid.y");
const std::string shader_src_no_previous = absl::Substitute(
shader_src_template, kWorkgroupSize, options_.output_layer_index(),
options_.combine_with_previous_ratio(), "",
options_.flip_vertically() ? "out_height - gid.y - 1" : "gid.y");
const std::string shader_src_with_previous = absl::Substitute(
shader_src_template, kWorkgroupSize, options_.output_layer_index(),
options_.combine_with_previous_ratio(), "#define READ_PREVIOUS",
options_.flip_vertically() ? "out_height - gid.y - 1" : "gid.y");
auto status = ::tflite::gpu::OkStatus();
// Shader programs.
GlShader shader_without_previous;
RET_CHECK_CALL(GlShader::CompileShader(
GL_COMPUTE_SHADER, shader_src_no_previous, &shader_without_previous));
mask_program_no_prev_ = absl::make_unique<GlProgram>();
RET_CHECK_CALL(GlProgram::CreateWithShader(shader_without_previous,
mask_program_no_prev_.get()));
GlShader shader_with_previous;
RET_CHECK_CALL(GlShader::CompileShader(
GL_COMPUTE_SHADER, shader_src_with_previous, &shader_with_previous));
mask_program_with_prev_ = absl::make_unique<GlProgram>();
RET_CHECK_CALL(GlProgram::CreateWithShader(shader_with_previous,
mask_program_with_prev_.get()));
// Shader programs.
GlShader shader_without_previous;
status = GlShader::CompileShader(GL_COMPUTE_SHADER, shader_src_no_previous,
&shader_without_previous);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
mask_program_no_prev_ = absl::make_unique<GlProgram>();
status = GlProgram::CreateWithShader(shader_without_previous,
mask_program_no_prev_.get());
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
GlShader shader_with_previous;
status = GlShader::CompileShader(GL_COMPUTE_SHADER, shader_src_with_previous,
&shader_with_previous);
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
mask_program_with_prev_ = absl::make_unique<GlProgram>();
status = GlProgram::CreateWithShader(shader_with_previous,
mask_program_with_prev_.get());
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
// Buffer storage for input tensor.
size_t tensor_length = tensor_width_ * tensor_height_ * tensor_channels_;
tensor_buffer_ = absl::make_unique<GlBuffer>();
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
tensor_length, tensor_buffer_.get()));
// Buffer storage for input tensor.
size_t tensor_length = tensor_width_ * tensor_height_ * tensor_channels_;
tensor_buffer_ = absl::make_unique<GlBuffer>();
status = CreateReadWriteShaderStorageBuffer<float>(tensor_length,
tensor_buffer_.get());
if (!status.ok()) {
return ::mediapipe::InternalError(status.error_message());
}
// Parameters.
glUseProgram(mask_program_with_prev_->id());
glUniform2i(glGetUniformLocation(mask_program_with_prev_->id(), "out_size"),
tensor_width_, tensor_height_);
glUniform1i(
glGetUniformLocation(mask_program_with_prev_->id(), "input_texture"),
1);
glUseProgram(mask_program_no_prev_->id());
glUniform2i(glGetUniformLocation(mask_program_no_prev_->id(), "out_size"),
tensor_width_, tensor_height_);
glUniform1i(
glGetUniformLocation(mask_program_no_prev_->id(), "input_texture"), 1);
// Parameters.
glUseProgram(mask_program_with_prev_->id());
glUniform2i(glGetUniformLocation(mask_program_with_prev_->id(), "out_size"),
tensor_width_, tensor_height_);
glUniform1i(
glGetUniformLocation(mask_program_with_prev_->id(), "input_texture"), 1);
glUseProgram(mask_program_no_prev_->id());
glUniform2i(glGetUniformLocation(mask_program_no_prev_->id(), "out_size"),
tensor_width_, tensor_height_);
glUniform1i(
glGetUniformLocation(mask_program_no_prev_->id(), "input_texture"), 1);
// Vertex shader attributes.
const GLint attr_location[NUM_ATTRIBUTES] = {
ATTRIB_VERTEX,
ATTRIB_TEXTURE_POSITION,
};
const GLchar* attr_name[NUM_ATTRIBUTES] = {
"position",
"texture_coordinate",
};
// Vertex shader attributes.
const GLint attr_location[NUM_ATTRIBUTES] = {
ATTRIB_VERTEX,
ATTRIB_TEXTURE_POSITION,
};
const GLchar* attr_name[NUM_ATTRIBUTES] = {
"position",
"texture_coordinate",
};
// Simple pass-through shader, used for hardware upsampling.
std::string upsample_shader_base = R"(
// Simple pass-through shader, used for hardware upsampling.
std::string upsample_shader_base = R"(
#if __VERSION__ < 130
#define in varying
#endif // __VERSION__ < 130
@@ -683,16 +677,19 @@ void main() {
}
)";
// Program
mediapipe::GlhCreateProgram(mediapipe::kBasicVertexShader,
upsample_shader_base.c_str(), NUM_ATTRIBUTES,
&attr_name[0], attr_location, &upsample_program_);
RET_CHECK(upsample_program_) << "Problem initializing the program.";
// Program
mediapipe::GlhCreateProgram(
mediapipe::kBasicVertexShader, upsample_shader_base.c_str(),
NUM_ATTRIBUTES, &attr_name[0], attr_location, &upsample_program_);
RET_CHECK(upsample_program_) << "Problem initializing the program.";
// Parameters
glUseProgram(upsample_program_);
glUniform1i(glGetUniformLocation(upsample_program_, "input_data"), 1);
#endif // __ANDROID__
// Parameters
glUseProgram(upsample_program_);
glUniform1i(glGetUniformLocation(upsample_program_, "input_data"), 1);
return ::mediapipe::OkStatus();
}));
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
mediapipe/calculators/tflite/util.h
+25
View File
@@ -0,0 +1,25 @@
// Copyright 2019 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.
#ifndef MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
#define MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
#define RET_CHECK_CALL(call) \
do { \
const auto status = (call); \
if (ABSL_PREDICT_FALSE(!status.ok())) \
return ::mediapipe::InternalError(status.error_message()); \
} while (0);
#endif // MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
mediapipe/calculators/util/BUILD
+64 -8
View File
@@ -235,19 +235,13 @@ cc_library(
"//mediapipe/framework/port:vector",
"//mediapipe/util:annotation_renderer",
] + select({
"//mediapipe:android": [
"//mediapipe/gpu:disable_gpu": [],
"//conditions:default": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:shader_util",
],
"//mediapipe:ios": [
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gl_simple_shaders",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/gpu:shader_util",
],
"//conditions:default": [],
}),
alwayslink = 1,
)
@@ -694,3 +688,65 @@ cc_test(
"//mediapipe/framework/tool:validate_type",
],
)
proto_library(
name = "top_k_scores_calculator_proto",
srcs = ["top_k_scores_calculator.proto"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_proto",
],
)
mediapipe_cc_proto_library(
name = "top_k_scores_calculator_cc_proto",
srcs = ["top_k_scores_calculator.proto"],
cc_deps = [
"//mediapipe/framework:calculator_cc_proto",
],
visibility = ["//visibility:public"],
deps = [":top_k_scores_calculator_proto"],
)
cc_library(
name = "top_k_scores_calculator",
srcs = ["top_k_scores_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
":top_k_scores_calculator_cc_proto",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:statusor",
"//mediapipe/framework:calculator_framework",
"//mediapipe/util:resource_util",
] + select({
"//mediapipe:android": [
"//mediapipe/util/android/file/base",
],
"//mediapipe:apple": [
"//mediapipe/util/android/file/base",
],
"//mediapipe:macos": [
"//mediapipe/framework/port:file_helpers",
],
"//conditions:default": [
"//mediapipe/framework/port:file_helpers",
],
}),
alwayslink = 1,
)
cc_test(
name = "top_k_scores_calculator_test",
srcs = ["top_k_scores_calculator_test.cc"],
deps = [
":top_k_scores_calculator",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework:packet",
"//mediapipe/framework/deps:message_matchers",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:status",
],
)
mediapipe/calculators/util/annotation_overlay_calculator.cc
+41 -34
View File
@@ -27,12 +27,12 @@
#include "mediapipe/util/annotation_renderer.h"
#include "mediapipe/util/color.pb.h"
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gl_simple_shaders.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/gpu/shader_util.h"
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
namespace mediapipe {
@@ -146,13 +146,13 @@ class AnnotationOverlayCalculator : public CalculatorBase {
bool use_gpu_ = false;
bool gpu_initialized_ = false;
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
mediapipe::GlCalculatorHelper gpu_helper_;
GLuint program_ = 0;
GLuint image_mat_tex_ = 0; // Overlay drawing image for GPU.
int width_ = 0;
int height_ = 0;
#endif // __ANDROID__ or iOS
#endif // MEDIAPIPE_DISABLE_GPU
};
REGISTER_CALCULATOR(AnnotationOverlayCalculator);
@@ -160,6 +160,8 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
CalculatorContract* cc) {
CHECK_GE(cc->Inputs().NumEntries(), 1);
bool use_gpu = false;
if (cc->Inputs().HasTag(kInputFrameTag) &&
cc->Inputs().HasTag(kInputFrameTagGpu)) {
return ::mediapipe::InternalError("Cannot have multiple input images.");
@@ -173,12 +175,13 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
int num_render_streams = cc->Inputs().NumEntries();
// Input image to render onto copy of.
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Inputs().HasTag(kInputFrameTagGpu)) {
cc->Inputs().Tag(kInputFrameTagGpu).Set<mediapipe::GpuBuffer>();
num_render_streams = cc->Inputs().NumEntries() - 1;
use_gpu |= true;
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Inputs().HasTag(kInputFrameTag)) {
cc->Inputs().Tag(kInputFrameTag).Set<ImageFrame>();
num_render_streams = cc->Inputs().NumEntries() - 1;
@@ -190,18 +193,21 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
}
// Rendered image.
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (cc->Outputs().HasTag(kOutputFrameTagGpu)) {
cc->Outputs().Tag(kOutputFrameTagGpu).Set<mediapipe::GpuBuffer>();
use_gpu |= true;
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
if (cc->Outputs().HasTag(kOutputFrameTag)) {
cc->Outputs().Tag(kOutputFrameTag).Set<ImageFrame>();
}
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // __ANDROID__ or iOS
if (use_gpu) {
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
}
@@ -212,11 +218,11 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
options_ = cc->Options<AnnotationOverlayCalculatorOptions>();
if (cc->Inputs().HasTag(kInputFrameTagGpu) &&
cc->Outputs().HasTag(kOutputFrameTagGpu)) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
use_gpu_ = true;
#else
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
#endif // __ANDROID__ or iOS
RET_CHECK_FAIL() << "GPU processing not enabled.";
#endif // !MEDIAPIPE_DISABLE_GPU
}
if (cc->Inputs().HasTag(kInputFrameTagGpu) ||
@@ -246,9 +252,9 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
}
if (use_gpu_) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
}
return ::mediapipe::OkStatus();
@@ -260,7 +266,7 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
std::unique_ptr<cv::Mat> image_mat;
ImageFormat::Format target_format;
if (use_gpu_) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (!gpu_initialized_) {
MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
@@ -269,7 +275,7 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
}));
gpu_initialized_ = true;
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
MP_RETURN_IF_ERROR(CreateRenderTargetGpu(cc, image_mat));
} else {
MP_RETURN_IF_ERROR(CreateRenderTargetCpu(cc, image_mat, &target_format));
@@ -288,7 +294,7 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
}
if (use_gpu_) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
// Overlay rendered image in OpenGL, onto a copy of input.
uchar* image_mat_ptr = image_mat->data;
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
@@ -296,7 +302,7 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
MP_RETURN_IF_ERROR(RenderToGpu(cc, image_mat_ptr));
return ::mediapipe::OkStatus();
}));
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
} else {
// Copy the rendered image to output.
uchar* image_mat_ptr = image_mat->data;
@@ -307,14 +313,14 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
}
::mediapipe::Status AnnotationOverlayCalculator::Close(CalculatorContext* cc) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
gpu_helper_.RunInGlContext([this] {
if (program_) glDeleteProgram(program_);
program_ = 0;
if (image_mat_tex_) glDeleteTextures(1, &image_mat_tex_);
image_mat_tex_ = 0;
});
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
@@ -325,7 +331,7 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
auto output_frame = absl::make_unique<ImageFrame>(
target_format, renderer_->GetImageWidth(), renderer_->GetImageHeight());
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
output_frame->CopyPixelData(target_format, renderer_->GetImageWidth(),
renderer_->GetImageHeight(), data_image,
ImageFrame::kGlDefaultAlignmentBoundary);
@@ -333,7 +339,7 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
output_frame->CopyPixelData(target_format, renderer_->GetImageWidth(),
renderer_->GetImageHeight(), data_image,
ImageFrame::kDefaultAlignmentBoundary);
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
cc->Outputs()
.Tag(kOutputFrameTag)
@@ -344,7 +350,7 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
::mediapipe::Status AnnotationOverlayCalculator::RenderToGpu(
CalculatorContext* cc, uchar* overlay_image) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
// Source and destination textures.
const auto& input_frame =
cc->Inputs().Tag(kInputFrameTagGpu).Get<mediapipe::GpuBuffer>();
@@ -390,7 +396,7 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
// Cleanup
input_texture.Release();
output_texture.Release();
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
@@ -451,15 +457,16 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
::mediapipe::Status AnnotationOverlayCalculator::CreateRenderTargetGpu(
CalculatorContext* cc, std::unique_ptr<cv::Mat>& image_mat) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
if (image_frame_available_) {
const auto& input_frame =
cc->Inputs().Tag(kInputFrameTagGpu).Get<mediapipe::GpuBuffer>();
const mediapipe::ImageFormat::Format format =
mediapipe::ImageFormatForGpuBufferFormat(input_frame.format());
if (format != mediapipe::ImageFormat::SRGBA)
RET_CHECK_FAIL() << "Unsupported GPU input format.";
if (format != mediapipe::ImageFormat::SRGBA &&
format != mediapipe::ImageFormat::SRGB)
RET_CHECK_FAIL() << "Unsupported GPU input format: " << format;
image_mat = absl::make_unique<cv::Mat>(
height_, width_, CV_8UC3,
@@ -471,14 +478,14 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
cv::Scalar(options_.canvas_color().r(), options_.canvas_color().g(),
options_.canvas_color().b()));
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
::mediapipe::Status AnnotationOverlayCalculator::GlRender(
CalculatorContext* cc) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
static const GLfloat square_vertices[] = {
-1.0f, -1.0f, // bottom left
1.0f, -1.0f, // bottom right
@@ -526,14 +533,14 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
glBindVertexArray(0);
glDeleteVertexArrays(1, &vao);
glDeleteBuffers(2, vbo);
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
::mediapipe::Status AnnotationOverlayCalculator::GlSetup(
CalculatorContext* cc) {
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
#if !defined(MEDIAPIPE_DISABLE_GPU)
const GLint attr_location[NUM_ATTRIBUTES] = {
ATTRIB_VERTEX,
ATTRIB_TEXTURE_POSITION,
@@ -609,7 +616,7 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
glBindTexture(GL_TEXTURE_2D, 0);
}
#endif // __ANDROID__ or iOS
#endif // !MEDIAPIPE_DISABLE_GPU
return ::mediapipe::OkStatus();
}
mediapipe/calculators/util/rect_to_render_data_calculator.cc
+59 -19
View File
@@ -23,10 +23,25 @@ namespace mediapipe {
namespace {
constexpr char kNormalizedRectTag[] = "NORM_RECT";
constexpr char kNormRectTag[] = "NORM_RECT";
constexpr char kRectTag[] = "RECT";
constexpr char kNormRectsTag[] = "NORM_RECTS";
constexpr char kRectsTag[] = "RECTS";
constexpr char kRenderDataTag[] = "RENDER_DATA";
RenderAnnotation::Rectangle* NewRect(
const RectToRenderDataCalculatorOptions& options, RenderData* render_data) {
auto* annotation = render_data->add_render_annotations();
annotation->mutable_color()->set_r(options.color().r());
annotation->mutable_color()->set_g(options.color().g());
annotation->mutable_color()->set_b(options.color().b());
annotation->set_thickness(options.thickness());
return options.filled()
? annotation->mutable_filled_rectangle()->mutable_rectangle()
: annotation->mutable_rectangle();
}
void SetRect(bool normalized, double xmin, double ymin, double width,
double height, double rotation,
RenderAnnotation::Rectangle* rect) {
@@ -51,6 +66,8 @@ void SetRect(bool normalized, double xmin, double ymin, double width,
// One of the following:
// NORM_RECT: A NormalizedRect
// RECT: A Rect
// NORM_RECTS: An std::vector<NormalizedRect>
// RECTS: An std::vector<Rect>
//
// Output:
// RENDER_DATA: A RenderData
@@ -83,16 +100,27 @@ REGISTER_CALCULATOR(RectToRenderDataCalculator);
::mediapipe::Status RectToRenderDataCalculator::GetContract(
CalculatorContract* cc) {
RET_CHECK(cc->Inputs().HasTag(kNormalizedRectTag) ^
cc->Inputs().HasTag(kRectTag));
RET_CHECK_EQ((cc->Inputs().HasTag(kNormRectTag) ? 1 : 0) +
(cc->Inputs().HasTag(kRectTag) ? 1 : 0) +
(cc->Inputs().HasTag(kNormRectsTag) ? 1 : 0) +
(cc->Inputs().HasTag(kRectsTag) ? 1 : 0),
1)
<< "Exactly one of NORM_RECT, RECT, NORM_RECTS or RECTS input stream "
"should be provided.";
RET_CHECK(cc->Outputs().HasTag(kRenderDataTag));
if (cc->Inputs().HasTag(kNormalizedRectTag)) {
cc->Inputs().Tag(kNormalizedRectTag).Set<NormalizedRect>();
if (cc->Inputs().HasTag(kNormRectTag)) {
cc->Inputs().Tag(kNormRectTag).Set<NormalizedRect>();
}
if (cc->Inputs().HasTag(kRectTag)) {
cc->Inputs().Tag(kRectTag).Set<Rect>();
}
if (cc->Inputs().HasTag(kNormRectsTag)) {
cc->Inputs().Tag(kNormRectsTag).Set<std::vector<NormalizedRect>>();
}
if (cc->Inputs().HasTag(kRectsTag)) {
cc->Inputs().Tag(kRectsTag).Set<std::vector<Rect>>();
}
cc->Outputs().Tag(kRenderDataTag).Set<RenderData>();
return ::mediapipe::OkStatus();
@@ -108,31 +136,43 @@ REGISTER_CALCULATOR(RectToRenderDataCalculator);
::mediapipe::Status RectToRenderDataCalculator::Process(CalculatorContext* cc) {
auto render_data = absl::make_unique<RenderData>();
auto* annotation = render_data->add_render_annotations();
annotation->mutable_color()->set_r(options_.color().r());
annotation->mutable_color()->set_g(options_.color().g());
annotation->mutable_color()->set_b(options_.color().b());
annotation->set_thickness(options_.thickness());
auto* rectangle =
options_.filled()
? annotation->mutable_filled_rectangle()->mutable_rectangle()
: annotation->mutable_rectangle();
if (cc->Inputs().HasTag(kNormalizedRectTag) &&
!cc->Inputs().Tag(kNormalizedRectTag).IsEmpty()) {
const auto& rect =
cc->Inputs().Tag(kNormalizedRectTag).Get<NormalizedRect>();
if (cc->Inputs().HasTag(kNormRectTag) &&
!cc->Inputs().Tag(kNormRectTag).IsEmpty()) {
const auto& rect = cc->Inputs().Tag(kNormRectTag).Get<NormalizedRect>();
auto* rectangle = NewRect(options_, render_data.get());
SetRect(/*normalized=*/true, rect.x_center() - rect.width() / 2.f,
rect.y_center() - rect.height() / 2.f, rect.width(), rect.height(),
rect.rotation(), rectangle);
}
if (cc->Inputs().HasTag(kRectTag) && !cc->Inputs().Tag(kRectTag).IsEmpty()) {
const auto& rect = cc->Inputs().Tag(kRectTag).Get<Rect>();
auto* rectangle = NewRect(options_, render_data.get());
SetRect(/*normalized=*/false, rect.x_center() - rect.width() / 2.f,
rect.y_center() - rect.height() / 2.f, rect.width(), rect.height(),
rect.rotation(), rectangle);
}
if (cc->Inputs().HasTag(kNormRectsTag) &&
!cc->Inputs().Tag(kNormRectsTag).IsEmpty()) {
const auto& rects =
cc->Inputs().Tag(kNormRectsTag).Get<std::vector<NormalizedRect>>();
for (auto& rect : rects) {
auto* rectangle = NewRect(options_, render_data.get());
SetRect(/*normalized=*/true, rect.x_center() - rect.width() / 2.f,
rect.y_center() - rect.height() / 2.f, rect.width(),
rect.height(), rect.rotation(), rectangle);
}
}
if (cc->Inputs().HasTag(kRectsTag) &&
!cc->Inputs().Tag(kRectsTag).IsEmpty()) {
const auto& rects = cc->Inputs().Tag(kRectsTag).Get<std::vector<Rect>>();
for (auto& rect : rects) {
auto* rectangle = NewRect(options_, render_data.get());
SetRect(/*normalized=*/false, rect.x_center() - rect.width() / 2.f,
rect.y_center() - rect.height() / 2.f, rect.width(),
rect.height(), rect.rotation(), rectangle);
}
}
cc->Outputs()
.Tag(kRenderDataTag)
mediapipe/calculators/util/top_k_scores_calculator.cc
+194
View File
@@ -0,0 +1,194 @@
// Copyright 2019 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 <algorithm>
#include <istream>
#include <iterator>
#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include "mediapipe/calculators/util/top_k_scores_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/statusor.h"
#include "mediapipe/util/resource_util.h"
#if defined(MEDIAPIPE_LITE) || defined(__ANDROID__) || \
(defined(__APPLE__) && !TARGET_OS_OSX)
#include "mediapipe/util/android/file/base/file.h"
#include "mediapipe/util/android/file/base/helpers.h"
#else
#include "mediapipe/framework/port/file_helpers.h"
#endif
namespace mediapipe {
// A calculator that takes a vector of scores and returns the indexes, scores,
// labels of the top k elements.
//
// Usage example:
// node {
// calculator: "TopKScoresCalculator"
// input_stream: "SCORES:score_vector"
// output_stream: "TOP_K_INDEXES:top_k_indexes"
// output_stream: "TOP_K_SCORES:top_k_scores"
// output_stream: "TOP_K_LABELS:top_k_labels"
// options: {
// [mediapipe.TopKScoresCalculatorOptions.ext] {
// top_k: 5
// threshold: 0.1
// label_map_path: "/path/to/label/map"
// }
// }
// }
class TopKScoresCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc);
::mediapipe::Status Open(CalculatorContext* cc) override;
::mediapipe::Status Process(CalculatorContext* cc) override;
private:
::mediapipe::Status LoadLabelmap(std::string label_map_path);
int top_k_ = -1;
float threshold_ = 0.0;
std::unordered_map<int, std::string> label_map_;
};
REGISTER_CALCULATOR(TopKScoresCalculator);
::mediapipe::Status TopKScoresCalculator::GetContract(CalculatorContract* cc) {
RET_CHECK(cc->Inputs().HasTag("SCORES"));
cc->Inputs().Tag("SCORES").Set<std::vector<float>>();
if (cc->Outputs().HasTag("TOP_K_INDEXES")) {
cc->Outputs().Tag("TOP_K_INDEXES").Set<std::vector<int>>();
}
if (cc->Outputs().HasTag("TOP_K_SCORES")) {
cc->Outputs().Tag("TOP_K_SCORES").Set<std::vector<float>>();
}
if (cc->Outputs().HasTag("TOP_K_LABELS")) {
cc->Outputs().Tag("TOP_K_LABELS").Set<std::vector<std::string>>();
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status TopKScoresCalculator::Open(CalculatorContext* cc) {
const auto& options = cc->Options<::mediapipe::TopKScoresCalculatorOptions>();
RET_CHECK(options.has_top_k() || options.has_threshold())
<< "Must specify at least one of the top_k and threshold fields in "
"TopKScoresCalculatorOptions.";
if (options.has_top_k()) {
RET_CHECK(options.top_k() > 0) << "top_k must be greater than zero.";
top_k_ = options.top_k();
}
if (options.has_threshold()) {
threshold_ = options.threshold();
}
if (options.has_label_map_path()) {
MP_RETURN_IF_ERROR(LoadLabelmap(options.label_map_path()));
}
if (cc->Outputs().HasTag("TOP_K_LABELS")) {
RET_CHECK(!label_map_.empty());
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status TopKScoresCalculator::Process(CalculatorContext* cc) {
const std::vector<float>& input_vector =
cc->Inputs().Tag("SCORES").Get<std::vector<float>>();
std::vector<int> top_k_indexes;
std::vector<float> top_k_scores;
std::vector<std::string> top_k_labels;
if (top_k_ > 0) {
top_k_indexes.reserve(top_k_);
top_k_scores.reserve(top_k_);
top_k_labels.reserve(top_k_);
}
std::priority_queue<std::pair<float, int>, std::vector<std::pair<float, int>>,
std::greater<std::pair<float, int>>>
pq;
for (int i = 0; i < input_vector.size(); ++i) {
if (input_vector[i] < threshold_) {
continue;
}
if (top_k_ > 0) {
if (pq.size() < top_k_) {
pq.push(std::pair<float, int>(input_vector[i], i));
} else if (pq.top().first < input_vector[i]) {
pq.pop();
pq.push(std::pair<float, int>(input_vector[i], i));
}
} else {
pq.push(std::pair<float, int>(input_vector[i], i));
}
}
while (!pq.empty()) {
top_k_indexes.push_back(pq.top().second);
top_k_scores.push_back(pq.top().first);
pq.pop();
}
reverse(top_k_indexes.begin(), top_k_indexes.end());
reverse(top_k_scores.begin(), top_k_scores.end());
if (cc->Outputs().HasTag("TOP_K_LABELS")) {
for (int index : top_k_indexes) {
top_k_labels.push_back(label_map_[index]);
}
}
if (cc->Outputs().HasTag("TOP_K_INDEXES")) {
cc->Outputs()
.Tag("TOP_K_INDEXES")
.AddPacket(MakePacket<std::vector<int>>(top_k_indexes)
.At(cc->InputTimestamp()));
}
if (cc->Outputs().HasTag("TOP_K_SCORES")) {
cc->Outputs()
.Tag("TOP_K_SCORES")
.AddPacket(MakePacket<std::vector<float>>(top_k_scores)
.At(cc->InputTimestamp()));
}
if (cc->Outputs().HasTag("TOP_K_LABELS")) {
cc->Outputs()
.Tag("TOP_K_LABELS")
.AddPacket(MakePacket<std::vector<std::string>>(top_k_labels)
.At(cc->InputTimestamp()));
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status TopKScoresCalculator::LoadLabelmap(
std::string label_map_path) {
std::string string_path;
ASSIGN_OR_RETURN(string_path, PathToResourceAsFile(label_map_path));
std::string label_map_string;
MP_RETURN_IF_ERROR(file::GetContents(string_path, &label_map_string));
std::istringstream stream(label_map_string);
std::string line;
int i = 0;
while (std::getline(stream, line)) {
label_map_[i++] = line;
}
return ::mediapipe::OkStatus();
}
} // namespace mediapipe
mediapipe/calculators/util/top_k_scores_calculator.proto
+33
View File
@@ -0,0 +1,33 @@
// Copyright 2019 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.
syntax = "proto2";
package mediapipe;
import "mediapipe/framework/calculator.proto";
message TopKScoresCalculatorOptions {
extend CalculatorOptions {
optional TopKScoresCalculatorOptions ext = 271211788;
}
// How many highest scoring packets to output.
optional int32 top_k = 1;
// If set, only keep the scores that are greater than the threshold.
optional float threshold = 2;
// Path to a label map file for getting the actual name of classes.
optional string label_map_path = 3;
}
mediapipe/calculators/util/top_k_scores_calculator_test.cc
+150
View File
@@ -0,0 +1,150 @@
// Copyright 2019 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 <vector>
#include "mediapipe/framework/calculator_runner.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/status_matchers.h"
namespace mediapipe {
TEST(TopKScoresCalculatorTest, TestNodeConfig) {
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
calculator: "TopKScoresCalculator"
input_stream: "SCORES:score_vector"
output_stream: "TOP_K_INDEXES:top_k_indexes"
output_stream: "TOP_K_SCORES:top_k_scores"
options: {
[mediapipe.TopKScoresCalculatorOptions.ext] {}
}
)"));
auto status = runner.Run();
ASSERT_TRUE(!status.ok());
EXPECT_THAT(
status.ToString(),
testing::HasSubstr(
"Must specify at least one of the top_k and threshold fields"));
}
TEST(TopKScoresCalculatorTest, TestTopKOnly) {
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
calculator: "TopKScoresCalculator"
input_stream: "SCORES:score_vector"
output_stream: "TOP_K_INDEXES:top_k_indexes"
output_stream: "TOP_K_SCORES:top_k_scores"
options: {
[mediapipe.TopKScoresCalculatorOptions.ext] { top_k: 2 }
}
)"));
std::vector<float> score_vector{0.9, 0.2, 0.3, 1.0, 0.1};
runner.MutableInputs()->Tag("SCORES").packets.push_back(
MakePacket<std::vector<float>>(score_vector).At(Timestamp(0)));
MP_ASSERT_OK(runner.Run());
const std::vector<Packet>& indexes_outputs =
runner.Outputs().Tag("TOP_K_INDEXES").packets;
ASSERT_EQ(1, indexes_outputs.size());
const auto& indexes = indexes_outputs[0].Get<std::vector<int>>();
EXPECT_EQ(2, indexes.size());
EXPECT_EQ(3, indexes[0]);
EXPECT_EQ(0, indexes[1]);
const std::vector<Packet>& scores_outputs =
runner.Outputs().Tag("TOP_K_SCORES").packets;
ASSERT_EQ(1, scores_outputs.size());
const auto& scores = scores_outputs[0].Get<std::vector<float>>();
EXPECT_EQ(2, scores.size());
EXPECT_NEAR(1, scores[0], 1e-5);
EXPECT_NEAR(0.9, scores[1], 1e-5);
}
TEST(TopKScoresCalculatorTest, TestThresholdOnly) {
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
calculator: "TopKScoresCalculator"
input_stream: "SCORES:score_vector"
output_stream: "TOP_K_INDEXES:top_k_indexes"
output_stream: "TOP_K_SCORES:top_k_scores"
options: {
[mediapipe.TopKScoresCalculatorOptions.ext] { threshold: 0.2 }
}
)"));
std::vector<float> score_vector{0.9, 0.2, 0.3, 1.0, 0.1};
runner.MutableInputs()->Tag("SCORES").packets.push_back(
MakePacket<std::vector<float>>(score_vector).At(Timestamp(0)));
MP_ASSERT_OK(runner.Run());
const std::vector<Packet>& indexes_outputs =
runner.Outputs().Tag("TOP_K_INDEXES").packets;
ASSERT_EQ(1, indexes_outputs.size());
const auto& indexes = indexes_outputs[0].Get<std::vector<int>>();
EXPECT_EQ(4, indexes.size());
EXPECT_EQ(3, indexes[0]);
EXPECT_EQ(0, indexes[1]);
EXPECT_EQ(2, indexes[2]);
EXPECT_EQ(1, indexes[3]);
const std::vector<Packet>& scores_outputs =
runner.Outputs().Tag("TOP_K_SCORES").packets;
ASSERT_EQ(1, scores_outputs.size());
const auto& scores = scores_outputs[0].Get<std::vector<float>>();
EXPECT_EQ(4, scores.size());
EXPECT_NEAR(1.0, scores[0], 1e-5);
EXPECT_NEAR(0.9, scores[1], 1e-5);
EXPECT_NEAR(0.3, scores[2], 1e-5);
EXPECT_NEAR(0.2, scores[3], 1e-5);
}
TEST(TopKScoresCalculatorTest, TestBothTopKAndThreshold) {
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
calculator: "TopKScoresCalculator"
input_stream: "SCORES:score_vector"
output_stream: "TOP_K_INDEXES:top_k_indexes"
output_stream: "TOP_K_SCORES:top_k_scores"
options: {
[mediapipe.TopKScoresCalculatorOptions.ext] { top_k: 4 threshold: 0.3 }
}
)"));
std::vector<float> score_vector{0.9, 0.2, 0.3, 1.0, 0.1};
runner.MutableInputs()->Tag("SCORES").packets.push_back(
MakePacket<std::vector<float>>(score_vector).At(Timestamp(0)));
MP_ASSERT_OK(runner.Run());
const std::vector<Packet>& indexes_outputs =
runner.Outputs().Tag("TOP_K_INDEXES").packets;
ASSERT_EQ(1, indexes_outputs.size());
const auto& indexes = indexes_outputs[0].Get<std::vector<int>>();
EXPECT_EQ(3, indexes.size());
EXPECT_EQ(3, indexes[0]);
EXPECT_EQ(0, indexes[1]);
EXPECT_EQ(2, indexes[2]);
const std::vector<Packet>& scores_outputs =
runner.Outputs().Tag("TOP_K_SCORES").packets;
ASSERT_EQ(1, scores_outputs.size());
const auto& scores = scores_outputs[0].Get<std::vector<float>>();
EXPECT_EQ(3, scores.size());
EXPECT_NEAR(1.0, scores[0], 1e-5);
EXPECT_NEAR(0.9, scores[1], 1e-5);
EXPECT_NEAR(0.3, scores[2], 1e-5);
}
} // namespace mediapipe