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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 6e5aa035cd1f6a9333962df5d3ab97a05bd5744e
This commit is contained in:
MediaPipe Team
2022-06-28 12:11:05 +00:00
committed by Sebastian Schmidt
parent 4a20e9909d
commit c688862570
144 changed files with 5772 additions and 2118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/calculators/tensor/BUILD
+102 -10
View File
@@ -40,6 +40,63 @@ selects.config_setting_group(
],
)
mediapipe_proto_library(
name = "audio_to_tensor_calculator_proto",
srcs = ["audio_to_tensor_calculator.proto"],
visibility = [
"//mediapipe/framework:mediapipe_internal",
],
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
],
)
cc_library(
name = "audio_to_tensor_calculator",
srcs = ["audio_to_tensor_calculator.cc"],
visibility = [
"//mediapipe/framework:mediapipe_internal",
],
deps = [
":audio_to_tensor_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/api2:packet",
"//mediapipe/framework/api2:port",
"//mediapipe/framework/formats:matrix",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/formats:time_series_header_cc_proto",
"//mediapipe/util:time_series_util",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
"@com_google_audio_tools//audio/dsp:resampler_q",
"@org_tensorflow//tensorflow/lite/c:common",
],
alwayslink = 1,
)
cc_test(
name = "audio_to_tensor_calculator_test",
srcs = ["audio_to_tensor_calculator_test.cc"],
deps = [
":audio_to_tensor_calculator",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:timestamp",
"//mediapipe/framework/api2:packet",
"//mediapipe/framework/formats:matrix",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:parse_text_proto",
"@com_google_absl//absl/strings",
"@com_google_audio_tools//audio/dsp:resampler_q",
"@org_tensorflow//tensorflow/lite/c:common",
],
)
mediapipe_proto_library(
name = "inference_calculator_proto",
srcs = ["inference_calculator.proto"],
@@ -50,6 +107,14 @@ mediapipe_proto_library(
],
)
# This target defines the "InferenceCalculator" component, which looks for the available concrete
# implementations linked into the current binary and picks the one to use.
# You can depend on :inference_calculator instead if you want to automatically include a default
# set of implementations tailored for the current build configuration.
# If you want to have precise control of which implementations to include (e.g. for strict binary
# size concerns), depend on those implementations directly, and do not depend on
# :inference_calculator.
# In all cases, use "InferenceCalulator" in your graphs.
cc_library(
name = "inference_calculator_interface",
srcs = ["inference_calculator.cc"],
@@ -62,8 +127,9 @@ cc_library(
],
"//conditions:default": [],
}),
visibility = ["//visibility:public"],
deps = [
":inference_calculator_cc_proto",
":inference_calculator_options_lib",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/api2:packet",
@@ -85,18 +151,31 @@ cc_library(
name = "inference_calculator_gl",
srcs = ["inference_calculator_gl.cc"],
tags = ["nomac"], # config problem with cpuinfo via TF
visibility = ["//visibility:public"],
deps = [
"inference_calculator_interface",
"//mediapipe/framework/deps:file_path",
":inference_calculator_interface",
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/util/tflite:config",
"//mediapipe/util/tflite:tflite_gpu_runner",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@org_tensorflow//tensorflow/lite:framework_stable",
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
],
alwayslink = 1,
)
cc_library(
name = "inference_calculator_gl_advanced",
srcs = ["inference_calculator_gl_advanced.cc"],
tags = ["nomac"],
visibility = ["//visibility:public"],
deps = [
":inference_calculator_interface",
"//mediapipe/framework/deps:file_path",
"//mediapipe/gpu:gl_calculator_helper",
"//mediapipe/util/tflite:tflite_gpu_runner",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@org_tensorflow//tensorflow/lite:framework_stable",
],
alwayslink = 1,
)
@@ -113,6 +192,7 @@ cc_library(
"-framework MetalKit",
],
tags = ["ios"],
visibility = ["//visibility:public"],
deps = [
"inference_calculator_interface",
"//mediapipe/gpu:MPPMetalHelper",
@@ -142,6 +222,7 @@ cc_library(
],
"//conditions:default": [],
}),
visibility = ["//visibility:public"],
deps = [
":inference_calculator_interface",
"@com_google_absl//absl/memory",
@@ -161,9 +242,13 @@ cc_library(
cc_library(
name = "inference_calculator_gl_if_compute_shader_available",
visibility = ["//visibility:public"],
deps = selects.with_or({
":compute_shader_unavailable": [],
"//conditions:default": [":inference_calculator_gl"],
"//conditions:default": [
":inference_calculator_gl",
":inference_calculator_gl_advanced",
],
}),
)
@@ -484,6 +569,7 @@ cc_library(
"//mediapipe/framework/formats:location",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/formats:tensor",
"//mediapipe/util:label_map_cc_proto",
"//mediapipe/util:resource_util",
] + select({
"//mediapipe:android": [
@@ -506,6 +592,7 @@ mediapipe_proto_library(
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
"//mediapipe/util:label_map_proto",
],
)
@@ -672,6 +759,7 @@ cc_library(
],
"//conditions:default": [],
}),
visibility = ["//visibility:public"],
deps = [
":image_to_tensor_converter",
":image_to_tensor_utils",
@@ -858,9 +946,7 @@ cc_library(
"@com_google_absl//absl/types:span",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:image_opencv",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:opencv_imgproc",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework:calculator_context",
"//mediapipe/framework:calculator_framework",
@@ -890,6 +976,12 @@ cc_library(
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_texture",
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl/converters:util",
],
}) + select({
"//mediapipe/framework/port:disable_opencv": [],
"//conditions:default": [
"//mediapipe/framework/formats:image_opencv",
"//mediapipe/framework/port:opencv_imgproc",
],
}),
alwayslink = 1,
)
mediapipe/calculators/tensor/audio_to_tensor_calculator.cc
+401
View File
@@ -0,0 +1,401 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <math.h>
#include <algorithm>
#include <cstring>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "audio/dsp/resampler_q.h"
#include "mediapipe/calculators/tensor/audio_to_tensor_calculator.pb.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/api2/packet.h"
#include "mediapipe/framework/api2/port.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/matrix.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/formats/time_series_header.pb.h"
#include "mediapipe/util/time_series_util.h"
namespace mediapipe {
namespace api2 {
// Converts audio buffers into tensors, possibly with resampling, buffering
// and framing, according to specified inputs and options. All input audio
// buffers will be first resampled from the input sample rate to the target
// sample rate if they are not equal. The resampled audio data (with the
// buffered samples from the previous runs in the streaming mode) will be broken
// into fixed-sized, possibly overlapping frames. Finally, all frames will be
// converted to and outputted as MediaPipe Tensors. The last output tensor will
// be zero-padding if the remaining samples are insufficient.
//
// This calculator assumes that the input timestamps refer to the first
// sample in each Matrix. The output timestamps follow this same convention.
// One Process() call may output multiple tensors packets. The timestamps of
// the output packets are determined by the timestamp of the previous output
// packet, the target sample rate, and the number of samples advanced after the
// previous output.
//
// The calculator has two running modes:
// Streaming mode: when "streaming_mode" is set to true in the calculator
// options, the calculator treats the input audio stream as a continuous
// stream. Thus, any samples that are not consumed in the previous runs will
// be cached in a global sample buffer. The audio data resampled from the
// current raw audio input will be appended to the global sample buffer.
// The calculator will process the global sample buffer and output as many
// tensors as possible.
// Non-streaming mode: when "streaming_mode" is set to false in the calculator
// options, the calculators treats the packets in the input audio stream as
// a batch of unrelated audio buffers. In each Process() call, the input
// buffer will be frist resampled, and framed as fixed-sized, possibly
// overlapping tensors. The last tensor produced by a Process() invocation
// will be zero-padding if the remaining samples are insufficient. As the
// calculator treats the input packets as unrelated, all samples will be
// processed immediately and no samples will be cached in the global sample
// buffer.
//
// Inputs:
// AUDIO - mediapipe::Matrix
// The audio data represented as mediapipe::Matrix.
// SAMPLE_RATE - double @Optional
// The sample rate of the corresponding audio data in the "AUDIO" stream.
// If a sample rate packet is provided at Timestamp::PreStream(), the sample
// rate will be used as the sample rate of every audio packets in the
// "AUDIO" stream. Note that one and only one of the "AUDIO" stream's time
// series header or the "SAMPLE_RATE" stream can exist.
//
// Outputs:
// TENSORS - std::vector<Tensor>
// Vector containing a single Tensor that represents a fix-sized audio
// frame.
// TIMESTAMPS - std::vector<Timestamp> @Optional
// Vector containing the output timestamps emitted by the current Process()
// invocation. In the non-streaming mode, the vector contains all of the
// output timestamps for an input audio buffer.
//
// Example:
// node {
// calculator: "AudioToTensorCalculator"
// input_stream: "AUDIO:audio"
// output_stream: "TENSORS:tensors"
// output_stream: "TIMESTAMPS:timestamps"
// options {
// [mediapipe.AudioToTensorCalculatorOptions.ext] {
// num_channels: 2
// num_samples: 512
// num_overlapping_samples: 64
// target_sample_rate: 16000
// streaming_mode: true # or false
// }
// }
// }
class AudioToTensorCalculator : public Node {
public:
static constexpr Input<Matrix> kAudioIn{"AUDIO"};
// TODO: Removes this optional input stream when the "AUDIO" stream
// uses the new mediapipe audio data containers that carry audio metatdata,
// such as sample rate.
static constexpr Input<double>::Optional kAudioSampleRateIn{"SAMPLE_RATE"};
static constexpr Output<std::vector<Tensor>> kTensorsOut{"TENSORS"};
// A vector of the output timestamps emitted by the current Process()
// invocation. The packet timestamp is the last emitted timestamp.
static constexpr Output<std::vector<Timestamp>>::Optional kTimestampsOut{
"TIMESTAMPS"};
MEDIAPIPE_NODE_CONTRACT(kAudioIn, kAudioSampleRateIn, kTensorsOut,
kTimestampsOut);
static absl::Status UpdateContract(CalculatorContract* cc);
absl::Status Open(CalculatorContext* cc);
absl::Status Process(CalculatorContext* cc);
absl::Status Close(CalculatorContext* cc);
private:
// The target number of channels.
int num_channels_;
// The target number of samples per channel.
int num_samples_;
// The number of samples per channel to advance after the current frame is
// processed.
int frame_step_;
bool streaming_mode_;
bool check_inconsistent_timestamps_;
Timestamp initial_timestamp_ = Timestamp::Unstarted();
int64 cumulative_input_samples_ = 0;
Timestamp next_output_timestamp_ = Timestamp::Unstarted();
double source_sample_rate_ = -1;
double target_sample_rate_ = -1;
// TODO: Configures QResamplerParams through calculator options.
audio_dsp::QResamplerParams params_;
// A QResampler instance to resample an audio stream.
std::unique_ptr<audio_dsp::QResampler<float>> resampler_;
Matrix sample_buffer_;
int processed_buffer_cols_ = 0;
absl::Status ProcessStreamingData(CalculatorContext* cc);
absl::Status ProcessNonStreamingData(CalculatorContext* cc);
absl::Status SetupStreamingResampler(double input_sample_rate_);
void AppendToSampleBuffer(Matrix buffer_to_append);
absl::StatusOr<std::vector<Tensor>> ConvertToTensor(
const Matrix& frame_to_convert);
absl::Status OutputTensors(const Matrix& buffer, bool should_flush,
CalculatorContext* cc);
};
absl::Status AudioToTensorCalculator::UpdateContract(CalculatorContract* cc) {
const auto& options =
cc->Options<mediapipe::AudioToTensorCalculatorOptions>();
if (!options.has_num_channels() || !options.has_num_samples() ||
!options.has_target_sample_rate()) {
return absl::InvalidArgumentError(
"AudioToTensorCalculatorOptions must specifiy "
"`num_channels`, `num_samples`, and `target_sample_rate`.");
}
if (options.streaming_mode()) {
// Explicitly disables tiemstamp offset to disallow the timestamp bound
// from the input streams to be propagated to the output streams.
// In the streaming mode, the output timestamp bound is based on
// next_output_timestamp_, which can be smaller than the current input
// timestamps.
cc->SetTimestampOffset(TimestampDiff::Unset());
}
return absl::OkStatus();
}
absl::Status AudioToTensorCalculator::Open(CalculatorContext* cc) {
const auto& options =
cc->Options<mediapipe::AudioToTensorCalculatorOptions>();
num_channels_ = options.num_channels();
num_samples_ = options.num_samples();
if (options.has_num_overlapping_samples()) {
RET_CHECK_GE(options.num_overlapping_samples(), 0);
RET_CHECK_LT(options.num_overlapping_samples(), num_samples_);
frame_step_ = num_samples_ - options.num_overlapping_samples();
} else {
frame_step_ = num_samples_;
}
target_sample_rate_ = options.target_sample_rate();
streaming_mode_ = options.streaming_mode();
if (streaming_mode_) {
check_inconsistent_timestamps_ = options.check_inconsistent_timestamps();
sample_buffer_.resize(num_channels_, Eigen::NoChange);
}
RET_CHECK(kAudioSampleRateIn(cc).IsConnected() ^
!kAudioIn(cc).Header().IsEmpty())
<< "Must either specify the time series header of the \"AUDIO\" stream "
"or have the \"SAMPLE_RATE\" stream connected.";
if (!kAudioIn(cc).Header().IsEmpty()) {
mediapipe::TimeSeriesHeader input_header;
MP_RETURN_IF_ERROR(mediapipe::time_series_util::FillTimeSeriesHeaderIfValid(
kAudioIn(cc).Header(), &input_header));
if (streaming_mode_) {
MP_RETURN_IF_ERROR(SetupStreamingResampler(input_header.sample_rate()));
} else {
source_sample_rate_ = input_header.sample_rate();
}
}
return absl::OkStatus();
}
absl::Status AudioToTensorCalculator::Process(CalculatorContext* cc) {
if (cc->InputTimestamp() == Timestamp::PreStream()) {
double current_source_sample_rate = kAudioSampleRateIn(cc).Get();
if (cc->Options<mediapipe::AudioToTensorCalculatorOptions>()
.streaming_mode()) {
return SetupStreamingResampler(current_source_sample_rate);
} else {
source_sample_rate_ = current_source_sample_rate;
return absl::OkStatus();
}
}
// Sanity checks.
const auto& input_frame = kAudioIn(cc).Get();
if (input_frame.rows() != num_channels_) {
return absl::InvalidArgumentError(absl::StrFormat(
"Audio input has %d channel(s) but the model requires %d channel(s).",
input_frame.rows(), num_channels_));
}
if (num_channels_ > 1 && input_frame.IsRowMajor) {
return absl::InvalidArgumentError(
"The audio data should be stored in column-major.");
}
return streaming_mode_ ? ProcessStreamingData(cc)
: ProcessNonStreamingData(cc);
}
absl::Status AudioToTensorCalculator::Close(CalculatorContext* cc) {
if (!streaming_mode_) {
return absl::OkStatus();
}
if (resampler_) {
Matrix resampled_buffer(num_channels_, 0);
resampler_->Flush(&resampled_buffer);
AppendToSampleBuffer(std::move(resampled_buffer));
}
return OutputTensors(sample_buffer_, /*should_flush=*/true, cc);
}
absl::Status AudioToTensorCalculator::ProcessStreamingData(
CalculatorContext* cc) {
const auto& input_buffer = kAudioIn(cc).Get();
if (initial_timestamp_ == Timestamp::Unstarted()) {
initial_timestamp_ = cc->InputTimestamp();
next_output_timestamp_ = initial_timestamp_;
}
if (source_sample_rate_ != -1 && check_inconsistent_timestamps_) {
mediapipe::time_series_util::LogWarningIfTimestampIsInconsistent(
cc->InputTimestamp(), initial_timestamp_, cumulative_input_samples_,
source_sample_rate_);
cumulative_input_samples_ += input_buffer.cols();
}
if (!kAudioSampleRateIn(cc).IsEmpty()) {
double current_source_sample_rate = kAudioSampleRateIn(cc).Get();
if (resampler_) {
RET_CHECK_EQ(current_source_sample_rate, source_sample_rate_);
} else {
MP_RETURN_IF_ERROR(SetupStreamingResampler(current_source_sample_rate));
}
}
if (resampler_) {
Matrix resampled_buffer(num_channels_, 0);
resampler_->ProcessSamples(input_buffer, &resampled_buffer);
AppendToSampleBuffer(std::move(resampled_buffer));
} else {
// Tries to consume the input matrix first to avoid extra data copy.
auto status_or_matrix = kAudioIn(cc).packet().Consume<Matrix>();
if (status_or_matrix.ok()) {
Matrix local_matrix(num_channels_, 0);
local_matrix.swap(*status_or_matrix.value());
AppendToSampleBuffer(std::move(local_matrix));
} else {
AppendToSampleBuffer(input_buffer);
}
}
MP_RETURN_IF_ERROR(OutputTensors(sample_buffer_, /*should_flush=*/false, cc));
// Removes the processed samples from the global sample buffer.
sample_buffer_ = Matrix(sample_buffer_.rightCols(sample_buffer_.cols() -
processed_buffer_cols_ - 1));
return absl::OkStatus();
}
absl::Status AudioToTensorCalculator::ProcessNonStreamingData(
CalculatorContext* cc) {
initial_timestamp_ = cc->InputTimestamp();
next_output_timestamp_ = initial_timestamp_;
const auto& input_frame = kAudioIn(cc).Get();
double source_sample_rate = kAudioSampleRateIn(cc).GetOr(source_sample_rate_);
if (source_sample_rate != -1 && source_sample_rate != target_sample_rate_) {
std::vector<float> resampled = audio_dsp::QResampleSignal<float>(
source_sample_rate, target_sample_rate_, num_channels_, params_,
input_frame);
Eigen::Map<const Matrix> matrix_mapping(resampled.data(), num_channels_,
resampled.size() / num_channels_);
return OutputTensors(matrix_mapping, /*should_flush=*/true, cc);
}
return OutputTensors(input_frame, /*should_flush=*/true, cc);
}
absl::Status AudioToTensorCalculator::SetupStreamingResampler(
double input_sample_rate) {
if (input_sample_rate == source_sample_rate_) {
return absl::OkStatus();
}
source_sample_rate_ = input_sample_rate;
if (source_sample_rate_ != target_sample_rate_) {
resampler_ = absl::make_unique<audio_dsp::QResampler<float>>(
source_sample_rate_, target_sample_rate_, num_channels_, params_);
if (!resampler_) {
return absl::InternalError("Failed to initialize resampler.");
}
}
return absl::OkStatus();
}
void AudioToTensorCalculator::AppendToSampleBuffer(Matrix buffer_to_append) {
sample_buffer_.conservativeResize(
Eigen::NoChange, sample_buffer_.cols() + buffer_to_append.cols());
sample_buffer_.rightCols(buffer_to_append.cols()).swap(buffer_to_append);
}
absl::StatusOr<std::vector<Tensor>> AudioToTensorCalculator::ConvertToTensor(
const Matrix& frame_to_convert) {
Tensor tensor(Tensor::ElementType::kFloat32,
Tensor::Shape({num_channels_, num_samples_}));
auto buffer_view = tensor.GetCpuWriteView();
if (frame_to_convert.size() < num_channels_ * num_samples_) {
std::memset(buffer_view.buffer<float>(), 0, tensor.bytes());
}
std::memcpy(buffer_view.buffer<float>(), frame_to_convert.data(),
frame_to_convert.size() * sizeof(float));
std::vector<Tensor> tensor_vector;
tensor_vector.push_back(std::move(tensor));
return tensor_vector;
}
absl::Status AudioToTensorCalculator::OutputTensors(const Matrix& buffer,
bool should_flush,
CalculatorContext* cc) {
int next_frame_first_col = 0;
std::vector<Timestamp> timestamps;
while ((!streaming_mode_ || !should_flush) &&
next_frame_first_col + num_samples_ <= buffer.cols()) {
ASSIGN_OR_RETURN(auto output_tensor, ConvertToTensor(buffer.block(
0, next_frame_first_col,
num_channels_, num_samples_)));
kTensorsOut(cc).Send(std::move(output_tensor), next_output_timestamp_);
timestamps.push_back(next_output_timestamp_);
next_output_timestamp_ += round(frame_step_ / target_sample_rate_ *
Timestamp::kTimestampUnitsPerSecond);
next_frame_first_col += frame_step_;
}
if (should_flush && next_frame_first_col < buffer.cols()) {
ASSIGN_OR_RETURN(auto output_tensor,
ConvertToTensor(buffer.block(
0, next_frame_first_col, num_channels_,
std::min(num_samples_,
(int)buffer.cols() - next_frame_first_col))));
// In the streaming mode, the flush happens in Close() and a packet at
// Timestamp::Max() will be emitted. In the non-streaming mode, each
// Process() invocation will process the entire buffer completely.
Timestamp timestamp =
streaming_mode_ ? Timestamp::Max() : next_output_timestamp_;
timestamps.push_back(timestamp);
kTensorsOut(cc).Send(std::move(output_tensor), timestamp);
}
if (kTimestampsOut(cc).IsConnected()) {
Timestamp timestamp = timestamps.back();
kTimestampsOut(cc).Send(std::move(timestamps), timestamp);
}
processed_buffer_cols_ = next_frame_first_col - 1;
return absl::OkStatus();
}
MEDIAPIPE_REGISTER_NODE(AudioToTensorCalculator);
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/tensor/audio_to_tensor_calculator.proto
+46
View File
@@ -0,0 +1,46 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
syntax = "proto2";
package mediapipe;
import "mediapipe/framework/calculator.proto";
message AudioToTensorCalculatorOptions {
extend mediapipe.CalculatorOptions {
optional AudioToTensorCalculatorOptions ext = 448635064;
}
// The required number of channels the output audio tensor has.
optional int64 num_channels = 1;
// The required number of samples per channel the output audio tensor has.
optional int64 num_samples = 2;
// The number of overlapping samples per channel the output audio tensor has.
optional int64 num_overlapping_samples = 3 [default = 0];
// The target number of samples per second (hertz) of the audio buffers that
// will be converted into tensors.
optional double target_sample_rate = 4;
// Whether to treat the input audio stream as a continous stream or a batch
// of unrelated audio buffers.
optional bool streaming_mode = 5 [default = true];
// Set to false to disable checks for jitter in timestamp values. Useful with
// live audio input.
optional bool check_inconsistent_timestamps = 6 [default = true];
}
mediapipe/calculators/tensor/audio_to_tensor_calculator_test.cc
+483
View File
@@ -0,0 +1,483 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <cmath>
#include <memory>
#include <string>
#include <vector>
#include "absl/strings/substitute.h"
#include "audio/dsp/resampler_q.h"
#include "mediapipe/framework/api2/packet.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/matrix.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/framework/timestamp.h"
namespace mediapipe {
namespace {
std::unique_ptr<Matrix> CreateTestMatrix(int num_channels, int num_samples,
int timestamp) {
auto matrix = std::make_unique<Matrix>(num_channels, num_samples);
for (int c = 0; c < num_channels; ++c) {
for (int i = 0; i < num_samples; ++i) {
// A float value with the sample, channel, and timestamp separated by a
// few orders of magnitude, for easy parsing by humans.
(*matrix)(c, i) = timestamp / 10000 + i + c / 100.0;
}
}
return matrix;
}
std::unique_ptr<Matrix> ResampleBuffer(const Matrix& input_matrix,
double resampling_factor) {
audio_dsp::QResamplerParams params;
std::vector<float> resampled;
int num_channels = input_matrix.rows();
std::vector<float> input_data(input_matrix.data(),
input_matrix.data() + input_matrix.size());
resampled = audio_dsp::QResampleSignal<float>(
1, resampling_factor, num_channels, params, input_data);
Matrix res = Eigen::Map<Matrix>(resampled.data(), num_channels,
resampled.size() / num_channels);
return std::make_unique<Matrix>(std::move(res));
}
class AudioToTensorCalculatorNonStreamingModeTest : public ::testing::Test {
protected:
void SetUp() override {}
void Run(int num_samples, int num_overlapping_samples,
double resampling_factor, const Matrix& input_matrix) {
double input_sample_rate = 10000;
double target_sample_rate = input_sample_rate * resampling_factor;
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(
absl::Substitute(R"(
input_stream: "audio"
input_stream: "sample_rate"
output_stream: "tensors"
output_stream: "timestamps"
node {
calculator: "AudioToTensorCalculator"
input_stream: "AUDIO:audio"
input_stream: "SAMPLE_RATE:sample_rate"
output_stream: "TENSORS:tensors"
output_stream: "TIMESTAMPS:timestamps"
options {
[mediapipe.AudioToTensorCalculatorOptions.ext] {
num_channels: $0
num_samples: $1
num_overlapping_samples: $2
target_sample_rate: $3
streaming_mode: false
}
}
}
)",
/*$0=*/input_matrix.rows(),
/*$1=*/num_samples, /*$2=*/num_overlapping_samples,
/*$3=*/target_sample_rate));
tool::AddVectorSink("tensors", &graph_config, &tensors_packets_);
tool::AddVectorSink("timestamps", &graph_config, &timestamps_packets_);
// Run the graph.
MP_ASSERT_OK(graph_.Initialize(graph_config));
MP_ASSERT_OK(graph_.StartRun({}));
// Run with the input matrix multiple times.
for (int i = 0; i < num_iterations_; ++i) {
MP_ASSERT_OK(graph_.AddPacketToInputStream(
"audio",
MakePacket<Matrix>(input_matrix)
.At(Timestamp(i * Timestamp::kTimestampUnitsPerSecond))));
MP_ASSERT_OK(graph_.AddPacketToInputStream(
"sample_rate",
MakePacket<double>(input_sample_rate)
.At(Timestamp(i * Timestamp::kTimestampUnitsPerSecond))));
}
MP_ASSERT_OK(graph_.CloseAllInputStreams());
MP_ASSERT_OK(graph_.WaitUntilIdle());
}
void CheckTensorsOutputPackets(const Matrix& expected_matrix,
int sample_offset, int num_tensors_per_input) {
ASSERT_EQ(num_iterations_ * num_tensors_per_input, tensors_packets_.size());
for (int i = 0; i < num_iterations_; ++i) {
for (int j = 0; j < num_tensors_per_input; ++j) {
CheckTensorsOutputPacket(
expected_matrix, tensors_packets_[i * num_tensors_per_input + j],
/*sample_offset*/ sample_offset * j, /*index=*/j);
}
}
}
void CheckTensorsOutputPacket(const Matrix& expected_matrix,
const Packet& packet, int sample_offset,
int index) {
MP_ASSERT_OK(packet.ValidateAsType<std::vector<Tensor>>());
ASSERT_EQ(1, packet.Get<std::vector<Tensor>>().size());
const Tensor& output_tensor = packet.Get<std::vector<Tensor>>()[0];
auto* buffer = output_tensor.GetCpuReadView().buffer<float>();
int num_values = output_tensor.shape().num_elements();
const std::vector<float> output_floats(buffer, buffer + num_values);
for (int i = 0; i < num_values; ++i) {
if (i + sample_offset >= expected_matrix.size()) {
EXPECT_FLOAT_EQ(output_floats[i], 0);
} else {
EXPECT_FLOAT_EQ(output_floats[i],
expected_matrix.coeff((i + sample_offset) % 2,
(i + sample_offset) / 2))
<< "i=" << i << ", sample_offset=" << sample_offset;
}
}
}
void CheckTimestampsOutputPackets(
std::vector<int64> expected_timestamp_values) {
ASSERT_EQ(num_iterations_, timestamps_packets_.size());
for (int i = 0; i < timestamps_packets_.size(); ++i) {
const auto& p = timestamps_packets_[i];
MP_ASSERT_OK(p.ValidateAsType<std::vector<Timestamp>>());
auto output_timestamps = p.Get<std::vector<Timestamp>>();
int64 base_timestamp = i * Timestamp::kTimestampUnitsPerSecond;
std::vector<Timestamp> expected_timestamps;
expected_timestamps.resize(expected_timestamp_values.size());
std::transform(
expected_timestamp_values.begin(), expected_timestamp_values.end(),
expected_timestamps.begin(), [base_timestamp](int64 v) -> Timestamp {
return Timestamp(v + base_timestamp);
});
EXPECT_EQ(expected_timestamps, output_timestamps);
EXPECT_EQ(p.Timestamp(), expected_timestamps.back());
}
}
// Fully close graph at end, otherwise calculator+tensors are destroyed
// after calling WaitUntilDone().
void CloseGraph() { MP_EXPECT_OK(graph_.WaitUntilDone()); }
private:
CalculatorGraph graph_;
int num_iterations_ = 10;
std::vector<Packet> tensors_packets_;
std::vector<Packet> timestamps_packets_;
};
TEST_F(AudioToTensorCalculatorNonStreamingModeTest,
ConvertToNoOverlappingFp32Tensors) {
auto input_matrix = CreateTestMatrix(2, 8, 0);
Run(/*num_samples=*/4, /*num_overlapping_samples=*/0,
/*resampling_factor=*/1.0f, *input_matrix);
CheckTensorsOutputPackets(*input_matrix, /*sample_offset=*/8,
/*num_tensors_per_input=*/2);
CheckTimestampsOutputPackets({0, 400});
CloseGraph();
}
TEST_F(AudioToTensorCalculatorNonStreamingModeTest,
ConvertToOverlappingFp32Tensors) {
auto input_matrix = CreateTestMatrix(2, 8, 0);
Run(/*num_samples=*/4, /*num_overlapping_samples=*/2,
/*resampling_factor=*/1.0f, *input_matrix);
CheckTensorsOutputPackets(*input_matrix, /*sample_offset=*/4,
/*num_tensors_per_input=*/4);
CheckTimestampsOutputPackets({0, 200, 400, 600});
CloseGraph();
}
TEST_F(AudioToTensorCalculatorNonStreamingModeTest, TensorsWithZeroPadding) {
auto input_matrix = CreateTestMatrix(2, 7, 0);
Run(/*num_samples=*/4, /*num_overlapping_samples=*/2,
/*resampling_factor=*/1.0f, *input_matrix);
CheckTensorsOutputPackets(*input_matrix, /*sample_offset=*/4,
/*num_tensors_per_input=*/3);
CheckTimestampsOutputPackets({0, 200, 400});
CloseGraph();
}
TEST_F(AudioToTensorCalculatorNonStreamingModeTest, Downsampling) {
auto input_matrix = CreateTestMatrix(2, 1024, 0);
Run(/*num_samples=*/256, /*num_overlapping_samples=*/0,
/*resampling_factor=*/0.5f, *input_matrix);
auto expected_matrix =
ResampleBuffer(*input_matrix, /*resampling_factor=*/0.5f);
CheckTensorsOutputPackets(*expected_matrix, /*sample_offset=*/512,
/*num_tensors_per_input=*/3);
CheckTimestampsOutputPackets({0, 51200, 102400});
CloseGraph();
}
TEST_F(AudioToTensorCalculatorNonStreamingModeTest,
DownsamplingWithOverlapping) {
auto input_matrix = CreateTestMatrix(2, 1024, 0);
Run(/*num_samples=*/256, /*num_overlapping_samples=*/64,
/*resampling_factor=*/0.5f, *input_matrix);
auto expected_matrix =
ResampleBuffer(*input_matrix, /*resampling_factor=*/0.5f);
CheckTensorsOutputPackets(*expected_matrix, /*sample_offset=*/384,
/*num_tensors_per_input=*/3);
CheckTimestampsOutputPackets({0, 38400, 76800});
CloseGraph();
}
TEST_F(AudioToTensorCalculatorNonStreamingModeTest, Upsampling) {
auto input_matrix = CreateTestMatrix(2, 1024, 0);
Run(/*num_samples=*/256, /*num_overlapping_samples=*/0,
/*resampling_factor=*/2.0f, *input_matrix);
auto expected_matrix =
ResampleBuffer(*input_matrix, /*resampling_factor=*/2.0f);
CheckTensorsOutputPackets(*expected_matrix,
/*sample_offset=*/512,
/*num_tensors_per_input=*/9);
CheckTimestampsOutputPackets(
{0, 12800, 25600, 38400, 51200, 64000, 76800, 89600, 102400});
CloseGraph();
}
TEST_F(AudioToTensorCalculatorNonStreamingModeTest, UpsamplingWithOverlapping) {
auto input_matrix = CreateTestMatrix(2, 256, 0);
Run(/*num_samples=*/256, /*num_overlapping_samples=*/64,
/*resampling_factor=*/2.0f, *input_matrix);
auto expected_matrix =
ResampleBuffer(*input_matrix, /*resampling_factor=*/2.0f);
CheckTensorsOutputPackets(*expected_matrix,
/*sample_offset=*/384,
/*num_tensors_per_input=*/3);
CheckTimestampsOutputPackets({0, 9600, 19200});
CloseGraph();
}
class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
protected:
void SetUp() override { sample_buffer_ = std::make_unique<Matrix>(2, 0); }
void SetInputBufferNumSamplesPerChannel(int num_samples) {
input_buffer_num_samples_ = num_samples;
}
void SetNumIterations(int num_iterations) {
num_iterations_ = num_iterations;
}
int GetExpectedNumOfSamples() {
Matrix* expected_matrix =
resampled_buffer_ ? resampled_buffer_.get() : sample_buffer_.get();
return expected_matrix->cols();
}
void Run(int num_samples, int num_overlapping_samples,
double resampling_factor) {
double input_sample_rate = 10000;
double target_sample_rate = input_sample_rate * resampling_factor;
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(
absl::Substitute(R"(
input_stream: "audio"
input_stream: "sample_rate"
output_stream: "tensors"
node {
calculator: "AudioToTensorCalculator"
input_stream: "AUDIO:audio"
input_stream: "SAMPLE_RATE:sample_rate"
output_stream: "TENSORS:tensors"
options {
[mediapipe.AudioToTensorCalculatorOptions.ext] {
num_channels: 2
num_samples: $0
num_overlapping_samples: $1
target_sample_rate: $2
streaming_mode:true
}
}
}
)",
/*$0=*/num_samples, /*$1=*/num_overlapping_samples,
/*$2=*/target_sample_rate));
tool::AddVectorSink("tensors", &graph_config, &tensors_packets_);
// Run the graph.
MP_ASSERT_OK(graph_.Initialize(graph_config));
MP_ASSERT_OK(graph_.StartRun({}));
for (int i = 0; i < num_iterations_; ++i) {
Timestamp input_timestamp(Timestamp::kTimestampUnitsPerSecond * i);
auto new_data = CreateTestMatrix(2, input_buffer_num_samples_,
input_timestamp.Value());
MP_ASSERT_OK(graph_.AddPacketToInputStream(
"audio", MakePacket<Matrix>(*new_data).At(input_timestamp)));
MP_ASSERT_OK(graph_.AddPacketToInputStream(
"sample_rate",
MakePacket<double>(input_sample_rate).At(input_timestamp)));
sample_buffer_->conservativeResize(
Eigen::NoChange, sample_buffer_->cols() + new_data->cols());
sample_buffer_->rightCols(new_data->cols()).swap(*new_data);
}
MP_ASSERT_OK(graph_.CloseAllInputStreams());
MP_ASSERT_OK(graph_.WaitUntilIdle());
if (resampling_factor != 1) {
resampled_buffer_ = ResampleBuffer(*sample_buffer_, resampling_factor);
}
}
void CheckTensorsOutputPackets(int sample_offset, int num_packets,
int64 timestamp_interval,
bool output_last_at_close) {
ASSERT_EQ(num_packets, tensors_packets_.size());
for (int i = 0; i < num_packets; ++i) {
if (i == num_packets - 1 && output_last_at_close) {
CheckTensorsOutputPacket(sample_offset * i, i, Timestamp::Max());
} else {
CheckTensorsOutputPacket(sample_offset * i, i,
Timestamp(timestamp_interval * i));
}
}
}
void CheckTensorsOutputPacket(int sample_offset, int index,
Timestamp expected_timestamp) {
const Packet& p = tensors_packets_[index];
MP_ASSERT_OK(p.ValidateAsType<std::vector<Tensor>>());
const Tensor& output_tensor = p.Get<std::vector<Tensor>>()[0];
auto buffer = output_tensor.GetCpuReadView().buffer<float>();
int num_values = output_tensor.shape().num_elements();
std::vector<float> output_floats(buffer, buffer + num_values);
Matrix* expected_matrix =
resampled_buffer_ ? resampled_buffer_.get() : sample_buffer_.get();
for (int i = 0; i < num_values; ++i) {
if (i + sample_offset >= expected_matrix->size()) {
EXPECT_FLOAT_EQ(output_floats[i], 0);
} else {
EXPECT_NEAR(output_floats[i],
expected_matrix->coeff((i + sample_offset) % 2,
(i + sample_offset) / 2),
0.001)
<< "i=" << i << ", sample_offset=" << sample_offset
<< ", packet index=" << index;
}
}
EXPECT_EQ(p.Timestamp(), expected_timestamp);
}
// Fully close graph at end, otherwise calculator+tensors are destroyed
// after calling WaitUntilDone().
void CloseGraph() { MP_EXPECT_OK(graph_.WaitUntilDone()); }
private:
int input_buffer_num_samples_ = 10;
int num_iterations_ = 10;
CalculatorGraph graph_;
std::vector<Packet> tensors_packets_;
std::unique_ptr<Matrix> sample_buffer_;
std::unique_ptr<Matrix> resampled_buffer_;
};
TEST_F(AudioToTensorCalculatorStreamingModeTest,
OutputNoOverlappingFp32Tensors) {
Run(/*num_samples=*/5, /*num_overlapping_samples=*/0,
/*resampling_factor=*/1.0f);
CheckTensorsOutputPackets(
/*sample_offset=*/10,
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 5),
/*timestamp_interval=*/500,
/*output_last_at_close=*/false);
CloseGraph();
}
TEST_F(AudioToTensorCalculatorStreamingModeTest, OutputRemainingInCloseMethod) {
Run(/*num_samples=*/6, /*num_overlapping_samples=*/0,
/*resampling_factor=*/1.0f);
CheckTensorsOutputPackets(
/*sample_offset=*/12,
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 6),
/*timestamp_interval=*/600,
/*output_last_at_close=*/true);
CloseGraph();
}
TEST_F(AudioToTensorCalculatorStreamingModeTest, OutputOverlappingFp32Tensors) {
SetInputBufferNumSamplesPerChannel(12);
Run(/*num_samples=*/10, /*num_overlapping_samples=*/2,
/*resampling_factor=*/1.0f);
CheckTensorsOutputPackets(
/*sample_offset=*/16,
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 8),
/*timestamp_interval=*/800,
/*output_last_at_close=*/true);
CloseGraph();
}
TEST_F(AudioToTensorCalculatorStreamingModeTest, Downsampling) {
SetInputBufferNumSamplesPerChannel(1000);
Run(/*num_samples=*/256, /*num_overlapping_samples=*/0,
/*resampling_factor=*/0.5f);
CheckTensorsOutputPackets(
/*sample_offset=*/512,
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 256),
/*timestamp_interval=*/51200,
/*output_last_at_close=*/true);
CloseGraph();
}
TEST_F(AudioToTensorCalculatorStreamingModeTest, DownsamplingWithOverlapping) {
SetInputBufferNumSamplesPerChannel(1024);
Run(/*num_samples=*/256, /*num_overlapping_samples=*/64,
/*resampling_factor=*/0.5f);
CheckTensorsOutputPackets(
/*sample_offset=*/384,
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 192),
/*timestamp_interval=*/38400,
/*output_last_at_close=*/true);
CloseGraph();
}
TEST_F(AudioToTensorCalculatorStreamingModeTest, Upsampling) {
SetInputBufferNumSamplesPerChannel(1000);
Run(/*num_samples=*/256, /*num_overlapping_samples=*/0,
/*resampling_factor=*/2.0f);
CheckTensorsOutputPackets(
/*sample_offset=*/512,
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 256),
/*timestamp_interval=*/12800,
/*output_last_at_close=*/true);
CloseGraph();
}
TEST_F(AudioToTensorCalculatorStreamingModeTest, UpsamplingWithOverlapping) {
SetInputBufferNumSamplesPerChannel(1024);
Run(/*num_samples=*/256, /*num_overlapping_samples=*/64,
/*resampling_factor=*/2.0f);
CheckTensorsOutputPackets(
/*sample_offset=*/384,
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 192),
/*timestamp_interval=*/9600,
/*output_last_at_close=*/true);
CloseGraph();
}
TEST_F(AudioToTensorCalculatorStreamingModeTest,
OnlyOutputInCloseIfNoSufficientSamples) {
SetNumIterations(1);
Run(/*num_samples=*/8, /*num_overlapping_samples=*/0,
/*resampling_factor=*/0.5f);
CheckTensorsOutputPackets(
/*sample_offset=*/0,
/*num_packets=*/1,
/*timestamp_interval=*/0,
/*output_last_at_close=*/true);
CloseGraph();
}
} // namespace
} // namespace mediapipe
mediapipe/calculators/tensor/inference_calculator.cc
+13 -2
View File
@@ -19,8 +19,8 @@
#include <string>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/strings/string_view.h"
#include "mediapipe/calculators/tensor/inference_calculator.pb.h"
#include "mediapipe/framework/api2/packet.h"
#include "mediapipe/framework/tool/subgraph_expansion.h"
#include "tensorflow/lite/core/api/op_resolver.h"
@@ -43,8 +43,19 @@ class InferenceCalculatorSelectorImpl
!options.has_delegate() || // Use GPU delegate if not specified
(options.has_delegate() && options.delegate().has_gpu());
if (should_use_gpu) {
const auto& api = options.delegate().gpu().api();
using Gpu = ::mediapipe::InferenceCalculatorOptions::Delegate::Gpu;
impls.emplace_back("Metal");
impls.emplace_back("Gl");
const bool prefer_gl_advanced =
options.delegate().gpu().use_advanced_gpu_api() &&
(api == Gpu::ANY || api == Gpu::OPENGL || api == Gpu::OPENCL);
if (prefer_gl_advanced) {
impls.emplace_back("GlAdvanced");
impls.emplace_back("Gl");
} else {
impls.emplace_back("Gl");
impls.emplace_back("GlAdvanced");
}
}
impls.emplace_back("Cpu");
for (const auto& suffix : impls) {
mediapipe/calculators/tensor/inference_calculator.h
+4
View File
@@ -134,6 +134,10 @@ struct InferenceCalculatorGl : public InferenceCalculator {
static constexpr char kCalculatorName[] = "InferenceCalculatorGl";
};
struct InferenceCalculatorGlAdvanced : public InferenceCalculator {
static constexpr char kCalculatorName[] = "InferenceCalculatorGlAdvanced";
};
struct InferenceCalculatorMetal : public InferenceCalculator {
static constexpr char kCalculatorName[] = "InferenceCalculatorMetal";
};
mediapipe/calculators/tensor/inference_calculator_face_detection_test.cc
+4 -3
View File
@@ -75,9 +75,10 @@ const std::vector<Param>& GetParams() {
class InferenceCalculatorTest : public testing::TestWithParam<Param> {
protected:
void SetDelegateForParam(mediapipe::CalculatorGraphConfig_Node* node) {
*node->mutable_options()
->MutableExtension(mediapipe::InferenceCalculatorOptions::ext)
->mutable_delegate() = GetParam().delegate;
auto options_map = tool::MutableOptionsMap().Initialize(*node);
auto options = options_map.Get<mediapipe::InferenceCalculatorOptions>();
*options.mutable_delegate() = GetParam().delegate;
options_map.Set(options);
}
};
mediapipe/calculators/tensor/inference_calculator_gl.cc
+40 -265
View File
@@ -20,22 +20,8 @@
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "mediapipe/calculators/tensor/inference_calculator.h"
#include "mediapipe/framework/deps/file_path.h"
#include "mediapipe/util/tflite/config.h"
#if MEDIAPIPE_TFLITE_GL_INFERENCE
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/util/tflite/tflite_gpu_runner.h"
#include "tensorflow/lite/delegates/gpu/common/shape.h"
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
#if defined(MEDIAPIPE_ANDROID)
#include "mediapipe/util/android/file/base/file.h"
#include "mediapipe/util/android/file/base/filesystem.h"
#include "mediapipe/util/android/file/base/helpers.h"
#endif // ANDROID
namespace mediapipe {
namespace api2 {
@@ -50,42 +36,22 @@ class InferenceCalculatorGlImpl
absl::Status Close(CalculatorContext* cc) override;
private:
absl::Status ReadGpuCaches();
absl::Status SaveGpuCaches();
absl::Status LoadModel(CalculatorContext* cc);
absl::Status LoadDelegate(CalculatorContext* cc);
absl::Status LoadDelegateAndAllocateTensors(CalculatorContext* cc);
absl::Status InitTFLiteGPURunner(CalculatorContext* cc);
// TfLite requires us to keep the model alive as long as the interpreter is.
Packet<TfLiteModelPtr> model_packet_;
#if MEDIAPIPE_TFLITE_GL_INFERENCE
mediapipe::GlCalculatorHelper gpu_helper_;
std::unique_ptr<tflite::gpu::TFLiteGPURunner> tflite_gpu_runner_;
bool allow_precision_loss_ = false;
mediapipe::InferenceCalculatorOptions::Delegate::Gpu::Api
tflite_gpu_runner_api_;
mediapipe::InferenceCalculatorOptions::Delegate::Gpu::InferenceUsage
tflite_gpu_runner_usage_;
#endif // MEDIAPIPE_TFLITE_GL_INFERENCE
TfLiteDelegatePtr delegate_;
std::unique_ptr<tflite::Interpreter> interpreter_;
#if MEDIAPIPE_TFLITE_GPU_SUPPORTED
std::vector<Tensor::Shape> output_shapes_;
std::vector<std::unique_ptr<Tensor>> gpu_buffers_in_;
std::vector<std::unique_ptr<Tensor>> gpu_buffers_out_;
#endif // MEDIAPIPE_TFLITE_GPU_SUPPORTED
bool use_advanced_gpu_api_ = false;
bool use_gpu_delegate_ = false;
bool use_kernel_caching_ = false;
std::string cached_kernel_filename_;
bool use_serialized_model_ = false;
std::string serialized_model_path_;
};
absl::Status InferenceCalculatorGlImpl::UpdateContract(CalculatorContract* cc) {
@@ -93,8 +59,7 @@ absl::Status InferenceCalculatorGlImpl::UpdateContract(CalculatorContract* cc) {
RET_CHECK(!options.model_path().empty() ^ kSideInModel(cc).IsConnected())
<< "Either model as side packet or model path in options is required.";
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
return absl::OkStatus();
return mediapipe::GlCalculatorHelper::UpdateContract(cc);
}
absl::Status InferenceCalculatorGlImpl::Open(CalculatorContext* cc) {
@@ -110,46 +75,12 @@ absl::Status InferenceCalculatorGlImpl::Open(CalculatorContext* cc) {
<< "for Gpu";
delegate.MergeFrom(input_side_packet_delegate);
}
const bool has_delegate = options.has_delegate() || !kDelegate(cc).IsEmpty();
use_advanced_gpu_api_ = has_delegate && delegate.has_gpu() &&
delegate.gpu().use_advanced_gpu_api();
allow_precision_loss_ = delegate.gpu().allow_precision_loss();
tflite_gpu_runner_api_ = delegate.gpu().api();
tflite_gpu_runner_usage_ = delegate.gpu().usage();
use_kernel_caching_ =
use_advanced_gpu_api_ && delegate.gpu().has_cached_kernel_path();
use_serialized_model_ = use_advanced_gpu_api_ &&
delegate.gpu().has_serialized_model_dir() &&
delegate.gpu().has_model_token();
use_gpu_delegate_ = !use_advanced_gpu_api_;
if (use_kernel_caching_) {
#ifdef MEDIAPIPE_ANDROID
cached_kernel_filename_ = delegate.gpu().cached_kernel_path() +
mediapipe::File::Basename(options.model_path()) +
".ker";
#endif // MEDIAPIPE_ANDROID
}
if (use_serialized_model_) {
#ifdef MEDIAPIPE_ANDROID
serialized_model_path_ = mediapipe::file::JoinPath(
delegate.gpu().serialized_model_dir(), delegate.gpu().model_token());
#endif // MEDIAPIPE_ANDROID
}
// When use_advanced_gpu_api_, model loading is handled in InitTFLiteGPURunner
// for everything.
if (!use_advanced_gpu_api_) {
MP_RETURN_IF_ERROR(LoadModel(cc));
}
MP_RETURN_IF_ERROR(LoadModel(cc));
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
MP_RETURN_IF_ERROR(
gpu_helper_.RunInGlContext([this, &cc]() -> ::mediapipe::Status {
return use_advanced_gpu_api_ ? InitTFLiteGPURunner(cc)
: LoadDelegateAndAllocateTensors(cc);
}));
return absl::OkStatus();
return gpu_helper_.RunInGlContext([this, &cc]() -> ::mediapipe::Status {
return LoadDelegateAndAllocateTensors(cc);
});
}
absl::Status InferenceCalculatorGlImpl::Process(CalculatorContext* cc) {
@@ -160,205 +91,53 @@ absl::Status InferenceCalculatorGlImpl::Process(CalculatorContext* cc) {
RET_CHECK(!input_tensors.empty());
auto output_tensors = absl::make_unique<std::vector<Tensor>>();
if (use_advanced_gpu_api_) {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &input_tensors, &output_tensors]() -> ::mediapipe::Status {
for (int i = 0; i < input_tensors.size(); ++i) {
MP_RETURN_IF_ERROR(tflite_gpu_runner_->BindSSBOToInputTensor(
input_tensors[i].GetOpenGlBufferReadView().name(), i));
}
output_tensors->reserve(output_shapes_.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
output_tensors->emplace_back(Tensor::ElementType::kFloat32,
output_shapes_[i]);
MP_RETURN_IF_ERROR(tflite_gpu_runner_->BindSSBOToOutputTensor(
output_tensors->back().GetOpenGlBufferWriteView().name(), i));
}
return absl::OkStatus();
}));
} else {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &input_tensors]() -> ::mediapipe::Status {
// Explicitly copy input.
for (int i = 0; i < input_tensors.size(); ++i) {
glBindBuffer(GL_COPY_READ_BUFFER,
input_tensors[i].GetOpenGlBufferReadView().name());
glBindBuffer(GL_COPY_WRITE_BUFFER,
gpu_buffers_in_[i]->GetOpenGlBufferWriteView().name());
glCopyBufferSubData(GL_COPY_READ_BUFFER, GL_COPY_WRITE_BUFFER, 0, 0,
input_tensors[i].bytes());
}
return absl::OkStatus();
}));
}
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &input_tensors]() -> ::mediapipe::Status {
// Explicitly copy input.
for (int i = 0; i < input_tensors.size(); ++i) {
glBindBuffer(GL_COPY_READ_BUFFER,
input_tensors[i].GetOpenGlBufferReadView().name());
glBindBuffer(GL_COPY_WRITE_BUFFER,
gpu_buffers_in_[i]->GetOpenGlBufferWriteView().name());
glCopyBufferSubData(GL_COPY_READ_BUFFER, GL_COPY_WRITE_BUFFER, 0, 0,
input_tensors[i].bytes());
}
return absl::OkStatus();
}));
// Run inference.
if (use_advanced_gpu_api_) {
RET_CHECK(tflite_gpu_runner_->Invoke().ok());
} else {
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
}
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
if (use_gpu_delegate_) {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &output_tensors]() -> ::mediapipe::Status {
output_tensors->reserve(output_shapes_.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
const auto& t = gpu_buffers_out_[i];
output_tensors->emplace_back(Tensor::ElementType::kFloat32,
gpu_buffers_out_[i]->shape());
auto read_view = t->GetOpenGlBufferReadView();
glBindBuffer(GL_COPY_READ_BUFFER, read_view.name());
auto write_view = output_tensors->back().GetOpenGlBufferWriteView();
glBindBuffer(GL_COPY_WRITE_BUFFER, write_view.name());
glCopyBufferSubData(GL_COPY_READ_BUFFER, GL_COPY_WRITE_BUFFER, 0, 0,
t->bytes());
}
return absl::OkStatus();
}));
}
// Output tensors are already bound if use_advanced_gpu_api_ is true.
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &output_tensors]() -> ::mediapipe::Status {
output_tensors->reserve(output_shapes_.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
const auto& t = gpu_buffers_out_[i];
output_tensors->emplace_back(Tensor::ElementType::kFloat32,
gpu_buffers_out_[i]->shape());
auto read_view = t->GetOpenGlBufferReadView();
glBindBuffer(GL_COPY_READ_BUFFER, read_view.name());
auto write_view = output_tensors->back().GetOpenGlBufferWriteView();
glBindBuffer(GL_COPY_WRITE_BUFFER, write_view.name());
glCopyBufferSubData(GL_COPY_READ_BUFFER, GL_COPY_WRITE_BUFFER, 0, 0,
t->bytes());
}
return absl::OkStatus();
}));
kOutTensors(cc).Send(std::move(output_tensors));
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::SaveGpuCaches() {
#ifdef MEDIAPIPE_ANDROID
if (use_kernel_caching_) {
// Save kernel file.
auto kernel_cache = absl::make_unique<std::vector<uint8_t>>(
tflite_gpu_runner_->GetSerializedBinaryCache());
std::string cache_str(kernel_cache->begin(), kernel_cache->end());
MP_RETURN_IF_ERROR(
mediapipe::file::SetContents(cached_kernel_filename_, cache_str));
}
if (use_serialized_model_) {
// Save serialized model file.
ASSIGN_OR_RETURN(std::vector<uint8_t> serialized_model_vec,
tflite_gpu_runner_->GetSerializedModel());
absl::string_view serialized_model(
reinterpret_cast<char*>(serialized_model_vec.data()),
serialized_model_vec.size());
MP_RETURN_IF_ERROR(
mediapipe::file::SetContents(serialized_model_path_, serialized_model));
}
#endif // MEDIAPIPE_ANDROID
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::Close(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(SaveGpuCaches());
if (use_gpu_delegate_) {
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
gpu_buffers_in_.clear();
gpu_buffers_out_.clear();
// Delegate must outlive the interpreter, hence the order is important.
interpreter_ = nullptr;
delegate_ = nullptr;
return absl::OkStatus();
}));
} else {
return gpu_helper_.RunInGlContext([this]() -> absl::Status {
gpu_buffers_in_.clear();
gpu_buffers_out_.clear();
// Delegate must outlive the interpreter, hence the order is important.
interpreter_ = nullptr;
delegate_ = nullptr;
}
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::ReadGpuCaches() {
#ifdef MEDIAPIPE_ANDROID
if (use_kernel_caching_ && File::Exists(cached_kernel_filename_)) {
// Load pre-compiled kernel file.
std::string cache_str;
MP_RETURN_IF_ERROR(
mediapipe::file::GetContents(cached_kernel_filename_, &cache_str));
std::vector<uint8_t> cache_vec(cache_str.begin(), cache_str.end());
tflite_gpu_runner_->SetSerializedBinaryCache(std::move(cache_vec));
}
if (use_serialized_model_ && File::Exists(serialized_model_path_)) {
// Load serialized model file.
std::string serialized_model_str;
MP_RETURN_IF_ERROR(
file::GetContents(serialized_model_path_, &serialized_model_str));
std::vector<uint8_t> serialized_model_vec(serialized_model_str.begin(),
serialized_model_str.end());
tflite_gpu_runner_->SetSerializedModel(std::move(serialized_model_vec));
}
#endif // MEDIAPIPE_ANDROID
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlImpl::InitTFLiteGPURunner(
CalculatorContext* cc) {
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(cc));
const auto& model = *model_packet_.Get();
// Create runner
tflite::gpu::InferenceOptions options;
options.priority1 = allow_precision_loss_
? tflite::gpu::InferencePriority::MIN_LATENCY
: tflite::gpu::InferencePriority::MAX_PRECISION;
options.priority2 = tflite::gpu::InferencePriority::AUTO;
options.priority3 = tflite::gpu::InferencePriority::AUTO;
switch (tflite_gpu_runner_usage_) {
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::
FAST_SINGLE_ANSWER: {
options.usage = tflite::gpu::InferenceUsage::FAST_SINGLE_ANSWER;
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::
SUSTAINED_SPEED: {
options.usage = tflite::gpu::InferenceUsage::SUSTAINED_SPEED;
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::UNSPECIFIED: {
return absl::InternalError("inference usage need to be specified.");
}
}
tflite_gpu_runner_ = std::make_unique<tflite::gpu::TFLiteGPURunner>(options);
switch (tflite_gpu_runner_api_) {
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::ANY: {
// Do not need to force any specific API.
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::OPENGL: {
tflite_gpu_runner_->ForceOpenGL();
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::OPENCL: {
tflite_gpu_runner_->ForceOpenCL();
break;
}
}
if (kSideInOpResolver(cc).IsConnected()) {
const tflite::OpResolver& op_resolver = kSideInOpResolver(cc).Get();
MP_RETURN_IF_ERROR(tflite_gpu_runner_->InitializeWithModel(
model, op_resolver, /*allow_quant_ops=*/true));
} else {
tflite::ops::builtin::BuiltinOpResolver op_resolver =
kSideInCustomOpResolver(cc).GetOr(
tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates());
MP_RETURN_IF_ERROR(tflite_gpu_runner_->InitializeWithModel(
model, op_resolver, /*allow_quant_ops=*/true));
}
// Create and bind OpenGL buffers for outputs.
// The buffers are created once and their ids are passed to calculator outputs
output_shapes_.resize(tflite_gpu_runner_->outputs_size());
for (int i = 0; i < tflite_gpu_runner_->outputs_size(); ++i) {
output_shapes_[i] = {tflite_gpu_runner_->GetOutputShapes()[i].b,
tflite_gpu_runner_->GetOutputShapes()[i].h,
tflite_gpu_runner_->GetOutputShapes()[i].w,
tflite_gpu_runner_->GetOutputShapes()[i].c};
}
MP_RETURN_IF_ERROR(ReadGpuCaches());
MP_RETURN_IF_ERROR(tflite_gpu_runner_->Build());
return absl::OkStatus();
return absl::OkStatus();
});
}
absl::Status InferenceCalculatorGlImpl::LoadModel(CalculatorContext* cc) {
@@ -375,12 +154,8 @@ absl::Status InferenceCalculatorGlImpl::LoadModel(CalculatorContext* cc) {
}
RET_CHECK(interpreter_);
#if defined(__EMSCRIPTEN__)
interpreter_->SetNumThreads(1);
#else
interpreter_->SetNumThreads(
cc->Options<mediapipe::InferenceCalculatorOptions>().cpu_num_thread());
#endif // __EMSCRIPTEN__
return absl::OkStatus();
}
mediapipe/calculators/tensor/inference_calculator_gl_advanced.cc
+285
View File
@@ -0,0 +1,285 @@
// Copyright 2022 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <cstring>
#include <memory>
#include <string>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "mediapipe/calculators/tensor/inference_calculator.h"
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/util/tflite/tflite_gpu_runner.h"
#if defined(MEDIAPIPE_ANDROID)
#include "mediapipe/framework/deps/file_path.h"
#include "mediapipe/util/android/file/base/file.h"
#include "mediapipe/util/android/file/base/filesystem.h"
#include "mediapipe/util/android/file/base/helpers.h"
#endif // ANDROID
namespace mediapipe {
namespace api2 {
// Runs TFLite GPU delegate API2 directly, bypassing interpreter usage, and
// allows choosing specific API.
//
// To trigger this code path:
// [mediapipe.InferenceCalculatorOptions.ext] {
// delegate {
// gpu {
// use_advanced_gpu_api: true
// api: OPENCL # or OPENGL or ANY
// }
// }
// }
class InferenceCalculatorGlAdvancedImpl
: public NodeImpl<InferenceCalculatorGlAdvanced,
InferenceCalculatorGlAdvancedImpl> {
public:
static absl::Status UpdateContract(CalculatorContract* cc);
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
absl::Status Close(CalculatorContext* cc) override;
private:
absl::Status ReadGpuCaches();
absl::Status SaveGpuCaches();
absl::Status InitTFLiteGPURunner(CalculatorContext* cc);
// TfLite requires us to keep the model alive as long as the interpreter is.
Packet<TfLiteModelPtr> model_packet_;
mediapipe::GlCalculatorHelper gpu_helper_;
std::unique_ptr<tflite::gpu::TFLiteGPURunner> tflite_gpu_runner_;
bool allow_precision_loss_ = false;
mediapipe::InferenceCalculatorOptions::Delegate::Gpu::Api
tflite_gpu_runner_api_;
mediapipe::InferenceCalculatorOptions::Delegate::Gpu::InferenceUsage
tflite_gpu_runner_usage_;
std::vector<Tensor::Shape> output_shapes_;
bool use_kernel_caching_ = false;
std::string cached_kernel_filename_;
bool use_serialized_model_ = false;
std::string serialized_model_path_;
};
absl::Status InferenceCalculatorGlAdvancedImpl::UpdateContract(
CalculatorContract* cc) {
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
RET_CHECK(!options.model_path().empty() ^ kSideInModel(cc).IsConnected())
<< "Either model as side packet or model path in options is required.";
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlAdvancedImpl::Open(CalculatorContext* cc) {
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
mediapipe::InferenceCalculatorOptions::Delegate delegate = options.delegate();
if (!kDelegate(cc).IsEmpty()) {
mediapipe::InferenceCalculatorOptions::Delegate input_side_packet_delegate =
kDelegate(cc).Get();
CHECK(input_side_packet_delegate.has_gpu() ||
input_side_packet_delegate.delegate_case() ==
mediapipe::InferenceCalculatorOptions::Delegate::DELEGATE_NOT_SET)
<< "inference_calculator_gl_advanced only supports delegate input side "
"packet for Gpu";
delegate.MergeFrom(input_side_packet_delegate);
}
allow_precision_loss_ = delegate.gpu().allow_precision_loss();
tflite_gpu_runner_api_ = delegate.gpu().api();
tflite_gpu_runner_usage_ = delegate.gpu().usage();
use_kernel_caching_ = delegate.gpu().has_cached_kernel_path();
use_serialized_model_ = delegate.gpu().has_serialized_model_dir() &&
delegate.gpu().has_model_token();
if (use_kernel_caching_) {
#ifdef MEDIAPIPE_ANDROID
cached_kernel_filename_ = delegate.gpu().cached_kernel_path() +
mediapipe::File::Basename(options.model_path()) +
".ker";
#endif // MEDIAPIPE_ANDROID
}
if (use_serialized_model_) {
#ifdef MEDIAPIPE_ANDROID
serialized_model_path_ = mediapipe::file::JoinPath(
delegate.gpu().serialized_model_dir(), delegate.gpu().model_token());
#endif // MEDIAPIPE_ANDROID
}
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
return gpu_helper_.RunInGlContext(
[this, &cc]() -> absl::Status { return InitTFLiteGPURunner(cc); });
}
absl::Status InferenceCalculatorGlAdvancedImpl::Process(CalculatorContext* cc) {
if (kInTensors(cc).IsEmpty()) {
return absl::OkStatus();
}
const auto& input_tensors = *kInTensors(cc);
RET_CHECK(!input_tensors.empty());
auto output_tensors = absl::make_unique<std::vector<Tensor>>();
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &input_tensors, &output_tensors]() -> absl::Status {
for (int i = 0; i < input_tensors.size(); ++i) {
MP_RETURN_IF_ERROR(tflite_gpu_runner_->BindSSBOToInputTensor(
input_tensors[i].GetOpenGlBufferReadView().name(), i));
}
output_tensors->reserve(output_shapes_.size());
for (int i = 0; i < output_shapes_.size(); ++i) {
output_tensors->emplace_back(Tensor::ElementType::kFloat32,
output_shapes_[i]);
MP_RETURN_IF_ERROR(tflite_gpu_runner_->BindSSBOToOutputTensor(
output_tensors->back().GetOpenGlBufferWriteView().name(), i));
}
return absl::OkStatus();
}));
// Run inference.
MP_RETURN_IF_ERROR(tflite_gpu_runner_->Invoke());
kOutTensors(cc).Send(std::move(output_tensors));
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlAdvancedImpl::SaveGpuCaches() {
#ifdef MEDIAPIPE_ANDROID
if (use_kernel_caching_) {
// Save kernel file.
auto kernel_cache = absl::make_unique<std::vector<uint8_t>>(
tflite_gpu_runner_->GetSerializedBinaryCache());
std::string cache_str(kernel_cache->begin(), kernel_cache->end());
MP_RETURN_IF_ERROR(
mediapipe::file::SetContents(cached_kernel_filename_, cache_str));
}
if (use_serialized_model_) {
// Save serialized model file.
ASSIGN_OR_RETURN(std::vector<uint8_t> serialized_model_vec,
tflite_gpu_runner_->GetSerializedModel());
absl::string_view serialized_model(
reinterpret_cast<char*>(serialized_model_vec.data()),
serialized_model_vec.size());
MP_RETURN_IF_ERROR(
mediapipe::file::SetContents(serialized_model_path_, serialized_model));
}
#endif // MEDIAPIPE_ANDROID
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlAdvancedImpl::Close(CalculatorContext* cc) {
MP_RETURN_IF_ERROR(SaveGpuCaches());
return gpu_helper_.RunInGlContext([this]() -> absl::Status {
tflite_gpu_runner_.reset();
return absl::OkStatus();
});
}
absl::Status InferenceCalculatorGlAdvancedImpl::ReadGpuCaches() {
#ifdef MEDIAPIPE_ANDROID
if (use_kernel_caching_ && File::Exists(cached_kernel_filename_)) {
// Load pre-compiled kernel file.
std::string cache_str;
MP_RETURN_IF_ERROR(
mediapipe::file::GetContents(cached_kernel_filename_, &cache_str));
std::vector<uint8_t> cache_vec(cache_str.begin(), cache_str.end());
tflite_gpu_runner_->SetSerializedBinaryCache(std::move(cache_vec));
}
if (use_serialized_model_ && File::Exists(serialized_model_path_)) {
// Load serialized model file.
std::string serialized_model_str;
MP_RETURN_IF_ERROR(
file::GetContents(serialized_model_path_, &serialized_model_str));
std::vector<uint8_t> serialized_model_vec(serialized_model_str.begin(),
serialized_model_str.end());
tflite_gpu_runner_->SetSerializedModel(std::move(serialized_model_vec));
}
#endif // MEDIAPIPE_ANDROID
return absl::OkStatus();
}
absl::Status InferenceCalculatorGlAdvancedImpl::InitTFLiteGPURunner(
CalculatorContext* cc) {
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(cc));
const auto& model = *model_packet_.Get();
// Create runner
tflite::gpu::InferenceOptions options;
options.priority1 = allow_precision_loss_
? tflite::gpu::InferencePriority::MIN_LATENCY
: tflite::gpu::InferencePriority::MAX_PRECISION;
options.priority2 = tflite::gpu::InferencePriority::AUTO;
options.priority3 = tflite::gpu::InferencePriority::AUTO;
switch (tflite_gpu_runner_usage_) {
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::
FAST_SINGLE_ANSWER: {
options.usage = tflite::gpu::InferenceUsage::FAST_SINGLE_ANSWER;
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::
SUSTAINED_SPEED: {
options.usage = tflite::gpu::InferenceUsage::SUSTAINED_SPEED;
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::UNSPECIFIED: {
return absl::InternalError("inference usage need to be specified.");
}
}
tflite_gpu_runner_ = std::make_unique<tflite::gpu::TFLiteGPURunner>(options);
switch (tflite_gpu_runner_api_) {
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::ANY: {
// Do not need to force any specific API.
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::OPENGL: {
tflite_gpu_runner_->ForceOpenGL();
break;
}
case mediapipe::InferenceCalculatorOptions::Delegate::Gpu::OPENCL: {
tflite_gpu_runner_->ForceOpenCL();
break;
}
}
if (kSideInOpResolver(cc).IsConnected()) {
const tflite::OpResolver& op_resolver = kSideInOpResolver(cc).Get();
MP_RETURN_IF_ERROR(tflite_gpu_runner_->InitializeWithModel(
model, op_resolver, /*allow_quant_ops=*/true));
} else {
tflite::ops::builtin::BuiltinOpResolver op_resolver =
kSideInCustomOpResolver(cc).GetOr(
tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates());
MP_RETURN_IF_ERROR(tflite_gpu_runner_->InitializeWithModel(
model, op_resolver, /*allow_quant_ops=*/true));
}
// Create and bind OpenGL buffers for outputs.
// The buffers are created once and their ids are passed to calculator outputs
output_shapes_.resize(tflite_gpu_runner_->outputs_size());
for (int i = 0; i < tflite_gpu_runner_->outputs_size(); ++i) {
output_shapes_[i] = {tflite_gpu_runner_->GetOutputShapes()[i].b,
tflite_gpu_runner_->GetOutputShapes()[i].h,
tflite_gpu_runner_->GetOutputShapes()[i].w,
tflite_gpu_runner_->GetOutputShapes()[i].c};
}
MP_RETURN_IF_ERROR(ReadGpuCaches());
return tflite_gpu_runner_->Build();
}
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/tensor/inference_calculator_test.cc
+84 -66
View File
@@ -38,61 +38,13 @@
#endif // defined(__APPLE__)
namespace mediapipe {
namespace {
void DoSmokeTest(const std::string& graph_proto) {
const int width = 8;
const int height = 8;
const int channels = 3;
// Prepare input tensor.
auto input_vec = absl::make_unique<std::vector<Tensor>>();
input_vec->emplace_back(Tensor::ElementType::kFloat32,
Tensor::Shape{1, height, width, channels});
{
auto view1 = input_vec->back().GetCpuWriteView();
auto tensor_buffer = view1.buffer<float>();
ASSERT_NE(tensor_buffer, nullptr);
for (int i = 0; i < width * height * channels - 1; i++) {
tensor_buffer[i] = 1;
}
}
constexpr int kTensorWidth = 8;
constexpr int kTensorHeight = 8;
constexpr int kTensorChannels = 3;
// Prepare single calculator graph to and wait for packets.
CalculatorGraphConfig graph_config =
ParseTextProtoOrDie<CalculatorGraphConfig>(graph_proto);
std::vector<Packet> output_packets;
tool::AddVectorSink("tensor_out", &graph_config, &output_packets);
CalculatorGraph graph(graph_config);
MP_ASSERT_OK(graph.StartRun({}));
// Push the tensor into the graph.
MP_ASSERT_OK(graph.AddPacketToInputStream(
"tensor_in", Adopt(input_vec.release()).At(Timestamp(0))));
// Wait until the calculator done processing.
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(1, output_packets.size());
// Get and process results.
const std::vector<Tensor>& result_vec =
output_packets[0].Get<std::vector<Tensor>>();
ASSERT_EQ(1, result_vec.size());
const Tensor& result = result_vec[0];
auto view = result.GetCpuReadView();
auto result_buffer = view.buffer<float>();
ASSERT_NE(result_buffer, nullptr);
for (int i = 0; i < width * height * channels - 1; i++) {
ASSERT_EQ(3, result_buffer[i]);
}
// Fully close graph at end, otherwise calculator+tensors are destroyed
// after calling WaitUntilDone().
MP_ASSERT_OK(graph.CloseInputStream("tensor_in"));
MP_ASSERT_OK(graph.WaitUntilDone());
}
// Tests a simple add model that adds an input tensor to itself.
TEST(InferenceCalculatorTest, SmokeTest) {
std::string graph_proto = R"(
constexpr char kGraphWithModelPathInOption[] = R"(
input_stream: "tensor_in"
node {
calculator: "InferenceCalculator"
@@ -106,18 +58,7 @@ TEST(InferenceCalculatorTest, SmokeTest) {
}
}
)";
// Test CPU inference only.
DoSmokeTest(/*graph_proto=*/absl::StrReplaceAll(
graph_proto, {{"$delegate", "delegate { tflite {} }"}}));
DoSmokeTest(absl::StrReplaceAll(graph_proto,
{{"$delegate", "delegate { xnnpack {} }"}}));
DoSmokeTest(absl::StrReplaceAll(
graph_proto,
{{"$delegate", "delegate { xnnpack { num_threads: 10 } }"}}));
}
TEST(InferenceCalculatorTest, SmokeTest_ModelAsInputSidePacket) {
std::string graph_proto = R"(
constexpr char kGraphWithModelAsInputSidePacket[] = R"(
input_stream: "tensor_in"
node {
@@ -154,7 +95,84 @@ TEST(InferenceCalculatorTest, SmokeTest_ModelAsInputSidePacket) {
}
}
)";
DoSmokeTest(graph_proto);
std::vector<Tensor> CreateInputs() {
std::vector<Tensor> input_vec;
// Prepare input tensor.
input_vec.emplace_back(
Tensor::ElementType::kFloat32,
Tensor::Shape{1, kTensorHeight, kTensorWidth, kTensorChannels});
{
auto view = input_vec.back().GetCpuWriteView();
auto num_elements = input_vec.back().shape().num_elements();
auto tensor_buffer = view.buffer<float>();
for (int i = 0; i < num_elements; i++) {
tensor_buffer[i] = 1;
}
}
return input_vec;
}
void RunGraphThenClose(CalculatorGraph& graph, std::vector<Tensor> input_vec) {
MP_ASSERT_OK(graph.StartRun({}));
// Push the tensor into the graph.
MP_ASSERT_OK(graph.AddPacketToInputStream(
"tensor_in",
MakePacket<std::vector<Tensor>>(std::move(input_vec)).At(Timestamp(0))));
// Wait until the calculator done processing.
MP_ASSERT_OK(graph.WaitUntilIdle());
// Fully close graph at end, otherwise calculator+tensors are destroyed
// after calling WaitUntilDone().
MP_ASSERT_OK(graph.CloseInputStream("tensor_in"));
MP_ASSERT_OK(graph.WaitUntilDone());
}
void DoSmokeTest(const std::string& graph_proto) {
auto input_vec = CreateInputs();
// Prepare single calculator graph to and wait for packets.
CalculatorGraphConfig graph_config =
ParseTextProtoOrDie<CalculatorGraphConfig>(graph_proto);
std::vector<Packet> output_packets;
tool::AddVectorSink("tensor_out", &graph_config, &output_packets);
CalculatorGraph graph(graph_config);
RunGraphThenClose(graph, std::move(input_vec));
ASSERT_EQ(1, output_packets.size());
// Get and process results.
const std::vector<Tensor>& result_vec =
output_packets[0].Get<std::vector<Tensor>>();
ASSERT_EQ(1, result_vec.size());
const Tensor& result = result_vec[0];
auto view = result.GetCpuReadView();
auto result_buffer = view.buffer<float>();
ASSERT_NE(result_buffer, nullptr);
for (int i = 0; i < result.shape().num_elements(); i++) {
ASSERT_EQ(3, result_buffer[i]);
}
}
// Tests a simple add model that adds an input tensor to itself.
TEST(InferenceCalculatorTest, SmokeTest) {
// Test CPU inference only.
DoSmokeTest(/*graph_proto=*/absl::StrReplaceAll(
kGraphWithModelPathInOption, {{"$delegate", "delegate { tflite {} }"}}));
DoSmokeTest(absl::StrReplaceAll(kGraphWithModelPathInOption,
{{"$delegate", "delegate { xnnpack {} }"}}));
DoSmokeTest(absl::StrReplaceAll(
kGraphWithModelPathInOption,
{{"$delegate", "delegate { xnnpack { num_threads: 10 } }"}}));
}
TEST(InferenceCalculatorTest, ModelAsInputSidePacketSmokeTest) {
DoSmokeTest(kGraphWithModelAsInputSidePacket);
}
} // namespace
} // namespace mediapipe
mediapipe/calculators/tensor/tensors_to_classification_calculator.cc
+108 -32
View File
@@ -16,7 +16,6 @@
#include <unordered_map>
#include <vector>
#include "absl/container/node_hash_map.h"
#include "absl/strings/str_format.h"
#include "absl/types/span.h"
#include "mediapipe/calculators/tensor/tensors_to_classification_calculator.pb.h"
@@ -25,6 +24,7 @@
#include "mediapipe/framework/formats/classification.pb.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/util/label_map.pb.h"
#include "mediapipe/util/resource_util.h"
#if defined(MEDIAPIPE_MOBILE)
#include "mediapipe/util/android/file/base/file.h"
@@ -35,6 +35,17 @@
namespace mediapipe {
namespace api2 {
namespace {
void SetClassificationLabel(const LabelMapItem label_map_item,
Classification* classification) {
classification->set_label(label_map_item.name());
if (label_map_item.has_display_name()) {
classification->set_display_name(label_map_item.display_name());
}
}
} // namespace
// Convert result tensors from classification models into MediaPipe
// classifications.
@@ -54,7 +65,6 @@ namespace api2 {
// output_stream: "CLASSIFICATIONS:classifications"
// options: {
// [mediapipe.TensorsToClassificationCalculatorOptions.ext] {
// num_classes: 1024
// min_score_threshold: 0.1
// label_map_path: "labelmap.txt"
// }
@@ -72,22 +82,35 @@ class TensorsToClassificationCalculator : public Node {
absl::Status Close(CalculatorContext* cc) override;
private:
::mediapipe::TensorsToClassificationCalculatorOptions options_;
int top_k_ = 0;
absl::node_hash_map<int, std::string> label_map_;
bool sort_by_descending_score_ = false;
proto_ns::Map<int64, LabelMapItem> local_label_map_;
bool label_map_loaded_ = false;
bool is_binary_classification_ = false;
float min_score_threshold_ = std::numeric_limits<float>::lowest();
// Set of allowed or ignored class indices.
struct ClassIndexSet {
absl::flat_hash_set<int> values;
bool is_allowlist;
};
// Allowed or ignored class indices based on provided options.
// These are used to filter out the output classification results.
ClassIndexSet class_index_set_;
bool IsClassIndexAllowed(int class_index);
const proto_ns::Map<int64, LabelMapItem>& GetLabelMap(CalculatorContext* cc);
};
MEDIAPIPE_REGISTER_NODE(TensorsToClassificationCalculator);
absl::Status TensorsToClassificationCalculator::Open(CalculatorContext* cc) {
options_ =
cc->Options<::mediapipe::TensorsToClassificationCalculatorOptions>();
const auto& options = cc->Options<TensorsToClassificationCalculatorOptions>();
top_k_ = options_.top_k();
if (options_.has_label_map_path()) {
top_k_ = options.top_k();
sort_by_descending_score_ = options.sort_by_descending_score();
if (options.has_label_map_path()) {
std::string string_path;
ASSIGN_OR_RETURN(string_path,
PathToResourceAsFile(options_.label_map_path()));
PathToResourceAsFile(options.label_map_path()));
std::string label_map_string;
MP_RETURN_IF_ERROR(
mediapipe::GetResourceContents(string_path, &label_map_string));
@@ -96,18 +119,45 @@ absl::Status TensorsToClassificationCalculator::Open(CalculatorContext* cc) {
std::string line;
int i = 0;
while (std::getline(stream, line)) {
label_map_[i++] = line;
LabelMapItem item;
item.set_name(line);
local_label_map_[i++] = item;
}
label_map_loaded_ = true;
} else if (options_.has_label_map()) {
for (int i = 0; i < options_.label_map().entries_size(); ++i) {
const auto& entry = options_.label_map().entries(i);
RET_CHECK(!label_map_.contains(entry.id()))
} else if (!options.label_items().empty()) {
label_map_loaded_ = true;
} else if (options.has_label_map()) {
for (int i = 0; i < options.label_map().entries_size(); ++i) {
const auto& entry = options.label_map().entries(i);
RET_CHECK(!local_label_map_.contains(entry.id()))
<< "Duplicate id found: " << entry.id();
label_map_[entry.id()] = entry.label();
LabelMapItem item;
item.set_name(entry.label());
local_label_map_[entry.id()] = item;
}
label_map_loaded_ = true;
}
if (options.has_min_score_threshold()) {
min_score_threshold_ = options.min_score_threshold();
}
is_binary_classification_ = options.binary_classification();
if (is_binary_classification_) {
RET_CHECK(options.allow_classes().empty() &&
options.ignore_classes().empty());
}
if (!options.allow_classes().empty()) {
RET_CHECK(options.ignore_classes().empty());
class_index_set_.is_allowlist = true;
for (int i = 0; i < options.allow_classes_size(); ++i) {
class_index_set_.values.insert(options.allow_classes(i));
}
} else {
class_index_set_.is_allowlist = false;
for (int i = 0; i < options.ignore_classes_size(); ++i) {
class_index_set_.values.insert(options.ignore_classes(i));
}
}
return absl::OkStatus();
}
@@ -118,19 +168,19 @@ absl::Status TensorsToClassificationCalculator::Process(CalculatorContext* cc) {
int num_classes = input_tensors[0].shape().num_elements();
if (options_.binary_classification()) {
if (is_binary_classification_) {
RET_CHECK_EQ(num_classes, 1);
// Number of classes for binary classification.
num_classes = 2;
}
if (label_map_loaded_) {
RET_CHECK_EQ(num_classes, label_map_.size());
RET_CHECK_EQ(num_classes, GetLabelMap(cc).size());
}
auto view = input_tensors[0].GetCpuReadView();
auto raw_scores = view.buffer<float>();
auto classification_list = absl::make_unique<ClassificationList>();
if (options_.binary_classification()) {
if (is_binary_classification_) {
Classification* class_first = classification_list->add_classification();
Classification* class_second = classification_list->add_classification();
class_first->set_index(0);
@@ -139,41 +189,48 @@ absl::Status TensorsToClassificationCalculator::Process(CalculatorContext* cc) {
class_second->set_score(1. - raw_scores[0]);
if (label_map_loaded_) {
class_first->set_label(label_map_[0]);
class_second->set_label(label_map_[1]);
SetClassificationLabel(GetLabelMap(cc).at(0), class_first);
SetClassificationLabel(GetLabelMap(cc).at(1), class_second);
}
} else {
for (int i = 0; i < num_classes; ++i) {
if (options_.has_min_score_threshold() &&
raw_scores[i] < options_.min_score_threshold()) {
if (!IsClassIndexAllowed(i)) {
continue;
}
if (raw_scores[i] < min_score_threshold_) {
continue;
}
Classification* classification =
classification_list->add_classification();
classification->set_index(i);
classification->set_score(raw_scores[i]);
if (label_map_loaded_) {
classification->set_label(label_map_[i]);
SetClassificationLabel(GetLabelMap(cc).at(i), classification);
}
}
}
// Note that partial_sort will raise error when top_k_ >
// classification_list->classification_size().
CHECK_GE(classification_list->classification_size(), top_k_);
auto raw_classification_list = classification_list->mutable_classification();
if (top_k_ > 0 && classification_list->classification_size() >= top_k_) {
if (top_k_ > 0) {
int desired_size =
std::min(classification_list->classification_size(), top_k_);
std::partial_sort(raw_classification_list->begin(),
raw_classification_list->begin() + top_k_,
raw_classification_list->begin() + desired_size,
raw_classification_list->end(),
[](const Classification a, const Classification b) {
return a.score() > b.score();
});
// Resizes the underlying list to have only top_k_ classifications.
raw_classification_list->DeleteSubrange(
top_k_, raw_classification_list->size() - top_k_);
if (desired_size >= top_k_) {
// Resizes the underlying list to have only top_k_ classifications.
raw_classification_list->DeleteSubrange(
top_k_, raw_classification_list->size() - top_k_);
}
} else if (sort_by_descending_score_) {
std::sort(raw_classification_list->begin(), raw_classification_list->end(),
[](const Classification a, const Classification b) {
return a.score() > b.score();
});
}
kOutClassificationList(cc).Send(std::move(classification_list));
return absl::OkStatus();
@@ -183,5 +240,24 @@ absl::Status TensorsToClassificationCalculator::Close(CalculatorContext* cc) {
return absl::OkStatus();
}
bool TensorsToClassificationCalculator::IsClassIndexAllowed(int class_index) {
if (class_index_set_.values.empty()) {
return true;
}
if (class_index_set_.is_allowlist) {
return class_index_set_.values.contains(class_index);
} else {
return !class_index_set_.values.contains(class_index);
}
}
const proto_ns::Map<int64, LabelMapItem>&
TensorsToClassificationCalculator::GetLabelMap(CalculatorContext* cc) {
return !local_label_map_.empty()
? local_label_map_
: cc->Options<TensorsToClassificationCalculatorOptions>()
.label_items();
}
} // namespace api2
} // namespace mediapipe
mediapipe/calculators/tensor/tensors_to_classification_calculator.proto
+23 -1
View File
@@ -19,6 +19,7 @@ syntax = "proto2";
package mediapipe;
import "mediapipe/framework/calculator.proto";
import "mediapipe/util/label_map.proto";
message TensorsToClassificationCalculatorOptions {
extend .mediapipe.CalculatorOptions {
@@ -38,16 +39,37 @@ message TensorsToClassificationCalculatorOptions {
// Number of highest scoring labels to output. If top_k is not positive then
// all labels are used.
optional int32 top_k = 2;
// Whether results should be sorted by descending score. By default, results
// may or may not be sorted: setting this to true guarantees that the returned
// results will be sorted by descending score.
optional bool sort_by_descending_score = 9;
// Path to a label map file for getting the actual name of class ids.
optional string label_map_path = 3;
// Label map. (Can be used instead of label_map_path.)
// NOTE: "label_map_path", if specified, takes precedence over "label_map".
// NOTE: either "label_map_path" or "label_items", if specified, takes
// precedence over "label_map".
// Deprecated: please use `label_items` instead.
optional LabelMap label_map = 5;
// Label items. (Can be used instead of label_map_path.)
// NOTE: "label_map_path", if specified, takes precedence over "label_items".
map<int64, LabelMapItem> label_items = 6;
// Whether the input is a single float for binary classification.
// When true, only a single float is expected in the input tensor and the
// label map, if provided, is expected to have exactly two labels.
// The single score(float) represent the probability of first label, and
// 1 - score is the probabilility of the second label.
optional bool binary_classification = 4;
// The ids of classes that should be ignored during decoding the score for
// each classification. If `ignore_classes` is specified, all the other
// classes that are not in the `ignore_class` field will be considered during
// decoding. `ignore_classes` and `allow_classes` are mutually exclusive.
repeated int32 ignore_classes = 7 [packed = true];
// The ids of classes that will be allowed during decoding the score for
// each classification. If `allow_classes` is specified, all the other classes
// that are not in the `allow_classes` field will be completely ignored.
// `ignore_classes` and `allow_classes` are mutually exclusive.
repeated int32 allow_classes = 8 [packed = true];
}
mediapipe/calculators/tensor/tensors_to_classification_calculator_test.cc
+117
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <algorithm>
#include <limits>
#include <vector>
#include "absl/memory/memory.h"
@@ -206,4 +208,119 @@ TEST_F(TensorsToClassificationCalculatorTest, CorrectOutputWithTopK) {
}
}
TEST_F(TensorsToClassificationCalculatorTest,
CorrectOutputWithSortByDescendingScore) {
mediapipe::CalculatorRunner runner(ParseTextProtoOrDie<Node>(R"pb(
calculator: "TensorsToClassificationCalculator"
input_stream: "TENSORS:tensors"
output_stream: "CLASSIFICATIONS:classifications"
options {
[mediapipe.TensorsToClassificationCalculatorOptions.ext] {
sort_by_descending_score: true
}
}
)pb"));
BuildGraph(&runner, {0, 0.5, 1});
MP_ASSERT_OK(runner.Run());
const auto& output_packets_ = runner.Outputs().Tag("CLASSIFICATIONS").packets;
EXPECT_EQ(1, output_packets_.size());
const auto& classification_list =
output_packets_[0].Get<ClassificationList>();
// Verify results are sorted by descending score.
EXPECT_EQ(3, classification_list.classification_size());
float score = std::numeric_limits<float>::max();
for (int i = 0; i < classification_list.classification_size(); ++i) {
EXPECT_LE(classification_list.classification(i).score(), score);
score = classification_list.classification(i).score();
}
}
TEST_F(TensorsToClassificationCalculatorTest,
ClassNameAllowlistWithLabelItems) {
mediapipe::CalculatorRunner runner(ParseTextProtoOrDie<Node>(R"pb(
calculator: "TensorsToClassificationCalculator"
input_stream: "TENSORS:tensors"
output_stream: "CLASSIFICATIONS:classifications"
options {
[mediapipe.TensorsToClassificationCalculatorOptions.ext] {
label_items {
key: 0
value { name: "ClassA" }
}
label_items {
key: 1
value { name: "ClassB" }
}
label_items {
key: 2
value { name: "ClassC" }
}
allow_classes: 1
}
}
)pb"));
BuildGraph(&runner, {0, 0.5, 1});
MP_ASSERT_OK(runner.Run());
const auto& output_packets_ = runner.Outputs().Tag("CLASSIFICATIONS").packets;
EXPECT_EQ(1, output_packets_.size());
const auto& classification_list =
output_packets_[0].Get<ClassificationList>();
EXPECT_EQ(1, classification_list.classification_size());
EXPECT_EQ(1, classification_list.classification(0).index());
EXPECT_EQ(0.5, classification_list.classification(0).score());
ASSERT_TRUE(classification_list.classification(0).has_label());
}
TEST_F(TensorsToClassificationCalculatorTest,
ClassNameIgnorelistWithLabelItems) {
mediapipe::CalculatorRunner runner(ParseTextProtoOrDie<Node>(R"pb(
calculator: "TensorsToClassificationCalculator"
input_stream: "TENSORS:tensors"
output_stream: "CLASSIFICATIONS:classifications"
options {
[mediapipe.TensorsToClassificationCalculatorOptions.ext] {
label_items {
key: 0
value { name: "ClassA" }
}
label_items {
key: 1
value { name: "ClassB" }
}
label_items {
key: 2
value { name: "ClassC" }
}
ignore_classes: 1
}
}
)pb"));
BuildGraph(&runner, {0, 0.5, 1});
MP_ASSERT_OK(runner.Run());
const auto& output_packets_ = runner.Outputs().Tag("CLASSIFICATIONS").packets;
EXPECT_EQ(1, output_packets_.size());
const auto& classification_list =
output_packets_[0].Get<ClassificationList>();
EXPECT_EQ(2, classification_list.classification_size());
EXPECT_EQ(0, classification_list.classification(0).index());
EXPECT_EQ(0, classification_list.classification(0).score());
ASSERT_TRUE(classification_list.classification(0).has_label());
EXPECT_EQ(2, classification_list.classification(1).index());
EXPECT_EQ(1, classification_list.classification(1).score());
ASSERT_TRUE(classification_list.classification(1).has_label());
}
} // namespace mediapipe
mediapipe/calculators/tensor/tensors_to_segmentation_calculator.cc
+15 -3
View File
@@ -20,10 +20,8 @@
#include "mediapipe/framework/calculator_context.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/image_opencv.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port.h"
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/statusor.h"
#include "mediapipe/gpu/gpu_origin.pb.h"
@@ -37,6 +35,11 @@
#include "mediapipe/gpu/shader_util.h"
#endif // !MEDIAPIPE_DISABLE_GPU
#if !MEDIAPIPE_DISABLE_OPENCV
#include "mediapipe/framework/formats/image_opencv.h"
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
#endif // !MEDIAPIPE_DISABLE_OPENCV
#if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_31
#include "tensorflow/lite/delegates/gpu/gl/converters/util.h"
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
@@ -159,9 +162,10 @@ class TensorsToSegmentationCalculator : public CalculatorBase {
return options_.gpu_origin() != mediapipe::GpuOrigin_Mode_TOP_LEFT;
}
#if !MEDIAPIPE_DISABLE_OPENCV
template <class T>
absl::Status ApplyActivation(cv::Mat& tensor_mat, cv::Mat* small_mask_mat);
#endif // !MEDIAPIPE_DISABLE_OPENCV
::mediapipe::TensorsToSegmentationCalculatorOptions options_;
#if !MEDIAPIPE_DISABLE_GPU
@@ -283,7 +287,11 @@ absl::Status TensorsToSegmentationCalculator::Process(CalculatorContext* cc) {
RET_CHECK_FAIL() << "GPU processing disabled.";
#endif // !MEDIAPIPE_DISABLE_GPU
} else {
#if !MEDIAPIPE_DISABLE_OPENCV
MP_RETURN_IF_ERROR(ProcessCpu(cc));
#else
RET_CHECK_FAIL() << "OpenCV processing disabled.";
#endif // !MEDIAPIPE_DISABLE_OPENCV
}
return absl::OkStatus();
@@ -311,6 +319,7 @@ absl::Status TensorsToSegmentationCalculator::Close(CalculatorContext* cc) {
absl::Status TensorsToSegmentationCalculator::ProcessCpu(
CalculatorContext* cc) {
#if !MEDIAPIPE_DISABLE_OPENCV
// Get input streams, and dimensions.
const auto& input_tensors =
cc->Inputs().Tag(kTensorsTag).Get<std::vector<Tensor>>();
@@ -360,10 +369,12 @@ absl::Status TensorsToSegmentationCalculator::ProcessCpu(
cv::resize(small_mask_mat, *output_mat,
cv::Size(output_width, output_height));
cc->Outputs().Tag(kMaskTag).Add(output_mask.release(), cc->InputTimestamp());
#endif // !MEDIAPIPE_DISABLE_OPENCV
return absl::OkStatus();
}
#if !MEDIAPIPE_DISABLE_OPENCV
template <class T>
absl::Status TensorsToSegmentationCalculator::ApplyActivation(
cv::Mat& tensor_mat, cv::Mat* small_mask_mat) {
@@ -411,6 +422,7 @@ absl::Status TensorsToSegmentationCalculator::ApplyActivation(
return absl::OkStatus();
}
#endif // !MEDIAPIPE_DISABLE_OPENCV
// Steps:
// 1. receive tensor