Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9437483827 | ||
|
|
d030c13931 | ||
|
|
fce372d153 | ||
|
|
c6fea4c9d9 | ||
|
|
259b48e082 | ||
|
|
a2a63e3876 | ||
|
|
412ab42d1f |
@@ -3,7 +3,7 @@
|
||||
# Basic build settings
|
||||
build --jobs 128
|
||||
build --define='absl=1'
|
||||
build --cxxopt='-std=c++11'
|
||||
build --cxxopt='-std=c++14'
|
||||
build --copt='-Wno-sign-compare'
|
||||
build --copt='-Wno-unused-function'
|
||||
build --copt='-Wno-uninitialized'
|
||||
|
||||
@@ -5,6 +5,8 @@
|
||||
|
||||

|
||||
|
||||
> "<em>MediaPipe has made it extremely easy to build our 3D person pose reconstruction demo app, facilitating accelerated neural network inference on device and synchronization of our result visualization with the video capture stream. Highly recommended!</em>" - George Papandreou, CTO, [Ariel AI](https://arielai.com)
|
||||
|
||||
## ML Solutions in MediaPipe
|
||||
|
||||
* [Hand Tracking](mediapipe/docs/hand_tracking_mobile_gpu.md)
|
||||
@@ -35,10 +37,18 @@ A web-based visualizer is hosted on [viz.mediapipe.dev](https://viz.mediapipe.de
|
||||
* [Discuss](https://groups.google.com/forum/#!forum/mediapipe) - General community discussion around MediaPipe
|
||||
|
||||
## Publications
|
||||
* [On-Device, Real-Time Hand Tracking with MediaPipe](https://ai.googleblog.com/2019/08/on-device-real-time-hand-tracking-with.html)
|
||||
* [MediaPipe: A Framework for Building Perception Pipelines](https://arxiv.org/abs/1906.08172)
|
||||
|
||||
## Events
|
||||
[Open sourced at CVPR 2019](https://sites.google.com/corp/view/perception-cv4arvr/mediapipe) on June 17~20 in Long Beach, CA
|
||||
* [MediaPipe Madrid Meetup, 16 Dec 2019](https://www.meetup.com/Madrid-AI-Developers-Group/events/266329088/)
|
||||
* [MediaPipe London Meetup, Google 123 Building, 12 Dec 2019](https://www.meetup.com/London-AI-Tech-Talk/events/266329038)
|
||||
* [ML Conference, Berlin, 11 Dec 2019](https://mlconference.ai/machine-learning-advanced-development/mediapipe-building-real-time-cross-platform-mobile-web-edge-desktop-video-audio-ml-pipelines/)
|
||||
* [MediaPipe Berlin Meetup, Google Berlin, 11 Dec 2019](https://www.meetup.com/Berlin-AI-Tech-Talk/events/266328794/)
|
||||
* [The 3rd Workshop on YouTube-8M Large Scale Video Understanding Workshop](https://research.google.com/youtube8m/workshop2019/index.html) Seoul, Korea ICCV 2019
|
||||
* [AI DevWorld 2019](https://aidevworld.com) on Oct 10 in San Jose, California
|
||||
* [Google Industry Workshop at ICIP 2019](http://2019.ieeeicip.org/?action=page4&id=14#Google) [Presentation](https://docs.google.com/presentation/d/e/2PACX-1vRIBBbO_LO9v2YmvbHHEt1cwyqH6EjDxiILjuT0foXy1E7g6uyh4CesB2DkkEwlRDO9_lWfuKMZx98T/pub?start=false&loop=false&delayms=3000&slide=id.g556cc1a659_0_5) on Sept 24 in Taipei, Taiwan
|
||||
* [Open sourced at CVPR 2019](https://sites.google.com/corp/view/perception-cv4arvr/mediapipe) on June 17~20 in Long Beach, CA
|
||||
|
||||
## Alpha Disclaimer
|
||||
MediaPipe is currently in alpha for v0.6. We are still making breaking API changes and expect to get to stable API by v1.0.
|
||||
|
||||
@@ -64,6 +64,12 @@ http_archive(
|
||||
sha256 = "267103f8a1e9578978aa1dc256001e6529ef593e5aea38193d31c2872ee025e8",
|
||||
strip_prefix = "glog-0.3.5",
|
||||
build_file = "@//third_party:glog.BUILD",
|
||||
patches = [
|
||||
"@//third_party:com_github_glog_glog_9779e5ea6ef59562b030248947f787d1256132ae.diff"
|
||||
],
|
||||
patch_args = [
|
||||
"-p1",
|
||||
],
|
||||
)
|
||||
|
||||
# libyuv
|
||||
@@ -97,9 +103,9 @@ http_archive(
|
||||
],
|
||||
)
|
||||
|
||||
# 2019-08-15
|
||||
_TENSORFLOW_GIT_COMMIT = "67def62936e28f97c16182dfcc467d8d1cae02b4"
|
||||
_TENSORFLOW_SHA256= "ddd4e3c056e7c0ff2ef29133b30fa62781dfbf8a903e99efb91a02d292fa9562"
|
||||
# 2019-11-12
|
||||
_TENSORFLOW_GIT_COMMIT = "a5f9bcd64453ff3d1f64cb4da4786db3d2da7f82"
|
||||
_TENSORFLOW_SHA256= "f2b6f2ab2ffe63e86eccd3ce4bea6b7197383d726638dfeeebcdc1e7de73f075"
|
||||
http_archive(
|
||||
name = "org_tensorflow",
|
||||
urls = [
|
||||
@@ -108,13 +114,6 @@ http_archive(
|
||||
],
|
||||
strip_prefix = "tensorflow-%s" % _TENSORFLOW_GIT_COMMIT,
|
||||
sha256 = _TENSORFLOW_SHA256,
|
||||
patches = [
|
||||
"@//third_party:tensorflow_065c20bf79253257c87bd4614bb9a7fdef015cbb.diff",
|
||||
"@//third_party:tensorflow_f67fcbefce906cd419e4657f0d41e21019b71abd.diff",
|
||||
],
|
||||
patch_args = [
|
||||
"-p1",
|
||||
],
|
||||
)
|
||||
|
||||
load("@org_tensorflow//tensorflow:workspace.bzl", "tf_workspace")
|
||||
@@ -248,18 +247,11 @@ android_sdk_repository(
|
||||
|
||||
# iOS basic build deps.
|
||||
|
||||
load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
|
||||
|
||||
git_repository(
|
||||
http_archive(
|
||||
name = "build_bazel_rules_apple",
|
||||
remote = "https://github.com/bazelbuild/rules_apple.git",
|
||||
tag = "0.18.0",
|
||||
patches = [
|
||||
"@//third_party:rules_apple_c0863d0596ae6b769a29fa3fb72ff036444fd249.diff",
|
||||
],
|
||||
patch_args = [
|
||||
"-p1",
|
||||
],
|
||||
sha256 = "bdc8e66e70b8a75da23b79f1f8c6207356df07d041d96d2189add7ee0780cf4e",
|
||||
strip_prefix = "rules_apple-b869b0d3868d78a1d4ffd866ccb304fb68aa12c3",
|
||||
url = "https://github.com/bazelbuild/rules_apple/archive/b869b0d3868d78a1d4ffd866ccb304fb68aa12c3.tar.gz",
|
||||
)
|
||||
|
||||
load(
|
||||
|
||||
@@ -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 =
|
||||
|
||||
@@ -113,8 +113,15 @@ class SpectrogramCalculator : public CalculatorBase {
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
Timestamp CurrentOutputTimestamp() {
|
||||
// Current output timestamp is the *center* of the next frame to be
|
||||
Timestamp CurrentOutputTimestamp(CalculatorContext* cc) {
|
||||
if (use_local_timestamp_) {
|
||||
return cc->InputTimestamp();
|
||||
}
|
||||
return CumulativeOutputTimestamp();
|
||||
}
|
||||
|
||||
Timestamp CumulativeOutputTimestamp() {
|
||||
// Cumulative output timestamp is the *center* of the next frame to be
|
||||
// emitted, hence delayed by half a window duration compared to relevant
|
||||
// input timestamp.
|
||||
return initial_input_timestamp_ +
|
||||
@@ -141,6 +148,7 @@ class SpectrogramCalculator : public CalculatorBase {
|
||||
const OutputMatrixType postprocess_output_fn(const OutputMatrixType&),
|
||||
CalculatorContext* cc);
|
||||
|
||||
bool use_local_timestamp_;
|
||||
double input_sample_rate_;
|
||||
bool pad_final_packet_;
|
||||
int frame_duration_samples_;
|
||||
@@ -173,6 +181,8 @@ const float SpectrogramCalculator::kLnPowerToDb = 4.342944819032518;
|
||||
SpectrogramCalculatorOptions spectrogram_options =
|
||||
cc->Options<SpectrogramCalculatorOptions>();
|
||||
|
||||
use_local_timestamp_ = spectrogram_options.use_local_timestamp();
|
||||
|
||||
if (spectrogram_options.frame_duration_seconds() <= 0.0) {
|
||||
::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Invalid or missing frame_duration_seconds.\n"
|
||||
@@ -351,11 +361,11 @@ template <class OutputMatrixType>
|
||||
<< "Inconsistent number of spectrogram channels.";
|
||||
if (allow_multichannel_input_) {
|
||||
cc->Outputs().Index(0).Add(spectrogram_matrices.release(),
|
||||
CurrentOutputTimestamp());
|
||||
CurrentOutputTimestamp(cc));
|
||||
} else {
|
||||
cc->Outputs().Index(0).Add(
|
||||
new OutputMatrixType(spectrogram_matrices->at(0)),
|
||||
CurrentOutputTimestamp());
|
||||
CurrentOutputTimestamp(cc));
|
||||
}
|
||||
cumulative_completed_frames_ += output_vectors.size();
|
||||
}
|
||||
|
||||
@@ -66,4 +66,11 @@ message SpectrogramCalculatorOptions {
|
||||
// uniformly regardless of output type (i.e., even dBs are multiplied, not
|
||||
// offset).
|
||||
optional double output_scale = 7 [default = 1.0];
|
||||
|
||||
// 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 = 8 [default = false];
|
||||
}
|
||||
|
||||
@@ -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()));
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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];
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -13,12 +13,12 @@
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "concatenate_vector_calculator_proto",
|
||||
srcs = ["concatenate_vector_calculator.proto"],
|
||||
@@ -26,6 +26,13 @@ proto_library(
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "dequantize_byte_array_calculator_proto",
|
||||
srcs = ["dequantize_byte_array_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "packet_cloner_calculator_proto",
|
||||
srcs = ["packet_cloner_calculator.proto"],
|
||||
@@ -72,6 +79,13 @@ proto_library(
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "clip_vector_size_calculator_proto",
|
||||
srcs = ["clip_vector_size_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "packet_cloner_calculator_cc_proto",
|
||||
srcs = ["packet_cloner_calculator.proto"],
|
||||
@@ -104,6 +118,22 @@ mediapipe_cc_proto_library(
|
||||
deps = [":concatenate_vector_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "clip_vector_size_calculator_cc_proto",
|
||||
srcs = ["clip_vector_size_calculator.proto"],
|
||||
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":clip_vector_size_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "dequantize_byte_array_calculator_cc_proto",
|
||||
srcs = ["dequantize_byte_array_calculator.proto"],
|
||||
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":dequantize_byte_array_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "quantize_float_vector_calculator_cc_proto",
|
||||
srcs = ["quantize_float_vector_calculator.proto"],
|
||||
@@ -154,6 +184,66 @@ cc_test(
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "begin_loop_calculator",
|
||||
srcs = ["begin_loop_calculator.cc"],
|
||||
hdrs = ["begin_loop_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_contract",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/memory",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "end_loop_calculator",
|
||||
srcs = ["end_loop_calculator.cc"],
|
||||
hdrs = ["end_loop_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_contract",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/util:render_data_cc_proto",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "begin_end_loop_calculator_graph_test",
|
||||
srcs = ["begin_end_loop_calculator_graph_test.cc"],
|
||||
deps = [
|
||||
":begin_loop_calculator",
|
||||
":end_loop_calculator",
|
||||
"//mediapipe/calculators/core:packet_cloner_calculator",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_contract",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/memory",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "concatenate_vector_calculator",
|
||||
srcs = ["concatenate_vector_calculator.cc"],
|
||||
@@ -166,7 +256,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,
|
||||
)
|
||||
|
||||
@@ -198,6 +294,50 @@ cc_test(
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "clip_vector_size_calculator",
|
||||
srcs = ["clip_vector_size_calculator.cc"],
|
||||
hdrs = ["clip_vector_size_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":clip_vector_size_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@org_tensorflow//tensorflow/lite:framework",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "clip_detection_vector_size_calculator",
|
||||
srcs = ["clip_detection_vector_size_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":clip_vector_size_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "clip_vector_size_calculator_test",
|
||||
srcs = ["clip_vector_size_calculator_test.cc"],
|
||||
deps = [
|
||||
":clip_vector_size_calculator",
|
||||
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "counting_source_calculator",
|
||||
srcs = ["counting_source_calculator.cc"],
|
||||
@@ -279,7 +419,7 @@ cc_library(
|
||||
"//visibility:public",
|
||||
],
|
||||
deps = [
|
||||
"//mediapipe/calculators/core:packet_cloner_calculator_cc_proto",
|
||||
":packet_cloner_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
@@ -381,6 +521,32 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "string_to_int_calculator",
|
||||
srcs = ["string_to_int_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "side_packet_to_stream_calculator",
|
||||
srcs = ["side_packet_to_stream_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "immediate_mux_calculator_test",
|
||||
srcs = ["immediate_mux_calculator_test.cc"],
|
||||
@@ -552,6 +718,32 @@ cc_test(
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "dequantize_byte_array_calculator",
|
||||
srcs = ["dequantize_byte_array_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":dequantize_byte_array_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "dequantize_byte_array_calculator_test",
|
||||
srcs = ["dequantize_byte_array_calculator_test.cc"],
|
||||
deps = [
|
||||
":dequantize_byte_array_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "quantize_float_vector_calculator",
|
||||
srcs = ["quantize_float_vector_calculator.cc"],
|
||||
@@ -688,3 +880,29 @@ cc_test(
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "stream_to_side_packet_calculator",
|
||||
srcs = ["stream_to_side_packet_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "stream_to_side_packet_calculator_test",
|
||||
srcs = ["stream_to_side_packet_calculator_test.cc"],
|
||||
deps = [
|
||||
":stream_to_side_packet_calculator",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/memory",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -0,0 +1,335 @@
|
||||
// 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 <string>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "mediapipe/calculators/core/begin_loop_calculator.h"
|
||||
#include "mediapipe/calculators/core/end_loop_calculator.h"
|
||||
#include "mediapipe/framework/calculator_contract.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h" // NOLINT
|
||||
|
||||
namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
typedef BeginLoopCalculator<std::vector<int>> BeginLoopIntegerCalculator;
|
||||
REGISTER_CALCULATOR(BeginLoopIntegerCalculator);
|
||||
|
||||
class IncrementCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Index(0).Set<int>();
|
||||
cc->Outputs().Index(0).Set<int>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
const int& input_int = cc->Inputs().Index(0).Get<int>();
|
||||
auto output_int = absl::make_unique<int>(input_int + 1);
|
||||
cc->Outputs().Index(0).Add(output_int.release(), cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(IncrementCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<int>> EndLoopIntegersCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopIntegersCalculator);
|
||||
|
||||
class BeginEndLoopCalculatorGraphTest : public ::testing::Test {
|
||||
protected:
|
||||
BeginEndLoopCalculatorGraphTest() {
|
||||
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
num_threads: 4
|
||||
input_stream: "ints"
|
||||
node {
|
||||
calculator: "BeginLoopIntegerCalculator"
|
||||
input_stream: "ITERABLE:ints"
|
||||
output_stream: "ITEM:int"
|
||||
output_stream: "BATCH_END:timestamp"
|
||||
}
|
||||
node {
|
||||
calculator: "IncrementCalculator"
|
||||
input_stream: "int"
|
||||
output_stream: "int_plus_one"
|
||||
}
|
||||
node {
|
||||
calculator: "EndLoopIntegersCalculator"
|
||||
input_stream: "ITEM:int_plus_one"
|
||||
input_stream: "BATCH_END:timestamp"
|
||||
output_stream: "ITERABLE:ints_plus_one"
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("ints_plus_one", &graph_config_, &output_packets_);
|
||||
}
|
||||
|
||||
CalculatorGraphConfig graph_config_;
|
||||
std::vector<Packet> output_packets_;
|
||||
};
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphTest, SingleEmptyVector) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||
Timestamp input_timestamp = Timestamp(0);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
// EndLoopCalc will forward the timestamp bound because there are no elements
|
||||
// in collection to output.
|
||||
ASSERT_EQ(0, output_packets_.size());
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphTest, SingleNonEmptyVector) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||
input_vector->emplace_back(0);
|
||||
input_vector->emplace_back(1);
|
||||
input_vector->emplace_back(2);
|
||||
Timestamp input_timestamp = Timestamp(0);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
ASSERT_EQ(1, output_packets_.size());
|
||||
EXPECT_EQ(input_timestamp, output_packets_[0].Timestamp());
|
||||
std::vector<int> expected_output_vector = {1, 2, 3};
|
||||
EXPECT_EQ(expected_output_vector, output_packets_[0].Get<std::vector<int>>());
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphTest, MultipleVectors) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
|
||||
auto input_vector0 = absl::make_unique<std::vector<int>>();
|
||||
input_vector0->emplace_back(0);
|
||||
input_vector0->emplace_back(1);
|
||||
Timestamp input_timestamp0 = Timestamp(0);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector0.release()).At(input_timestamp0)));
|
||||
|
||||
auto input_vector1 = absl::make_unique<std::vector<int>>();
|
||||
Timestamp input_timestamp1 = Timestamp(1);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector1.release()).At(input_timestamp1)));
|
||||
|
||||
auto input_vector2 = absl::make_unique<std::vector<int>>();
|
||||
input_vector2->emplace_back(2);
|
||||
input_vector2->emplace_back(3);
|
||||
Timestamp input_timestamp2 = Timestamp(2);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector2.release()).At(input_timestamp2)));
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
|
||||
ASSERT_EQ(2, output_packets_.size());
|
||||
|
||||
EXPECT_EQ(input_timestamp0, output_packets_[0].Timestamp());
|
||||
std::vector<int> expected_output_vector0 = {1, 2};
|
||||
EXPECT_EQ(expected_output_vector0,
|
||||
output_packets_[0].Get<std::vector<int>>());
|
||||
|
||||
// At input_timestamp1, EndLoopCalc will forward timestamp bound as there are
|
||||
// no elements in vector to process.
|
||||
|
||||
EXPECT_EQ(input_timestamp2, output_packets_[1].Timestamp());
|
||||
std::vector<int> expected_output_vector2 = {3, 4};
|
||||
EXPECT_EQ(expected_output_vector2,
|
||||
output_packets_[1].Get<std::vector<int>>());
|
||||
}
|
||||
|
||||
class MultiplierCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Index(0).Set<int>();
|
||||
cc->Inputs().Index(1).Set<int>();
|
||||
cc->Outputs().Index(0).Set<int>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
const int& input_int = cc->Inputs().Index(0).Get<int>();
|
||||
const int& multiplier_int = cc->Inputs().Index(1).Get<int>();
|
||||
auto output_int = absl::make_unique<int>(input_int * multiplier_int);
|
||||
cc->Outputs().Index(0).Add(output_int.release(), cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(MultiplierCalculator);
|
||||
|
||||
class BeginEndLoopCalculatorGraphWithClonedInputsTest : public ::testing::Test {
|
||||
protected:
|
||||
BeginEndLoopCalculatorGraphWithClonedInputsTest() {
|
||||
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
num_threads: 4
|
||||
input_stream: "ints"
|
||||
input_stream: "multiplier"
|
||||
node {
|
||||
calculator: "BeginLoopIntegerCalculator"
|
||||
input_stream: "ITERABLE:ints"
|
||||
input_stream: "CLONE:multiplier"
|
||||
output_stream: "ITEM:int_at_loop"
|
||||
output_stream: "CLONE:multiplier_cloned_at_loop"
|
||||
output_stream: "BATCH_END:timestamp"
|
||||
}
|
||||
node {
|
||||
calculator: "MultiplierCalculator"
|
||||
input_stream: "int_at_loop"
|
||||
input_stream: "multiplier_cloned_at_loop"
|
||||
output_stream: "multiplied_int_at_loop"
|
||||
}
|
||||
node {
|
||||
calculator: "EndLoopIntegersCalculator"
|
||||
input_stream: "ITEM:multiplied_int_at_loop"
|
||||
input_stream: "BATCH_END:timestamp"
|
||||
output_stream: "ITERABLE:multiplied_ints"
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("multiplied_ints", &graph_config_, &output_packets_);
|
||||
}
|
||||
|
||||
CalculatorGraphConfig graph_config_;
|
||||
std::vector<Packet> output_packets_;
|
||||
};
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, SingleEmptyVector) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||
Timestamp input_timestamp = Timestamp(42);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||
auto multiplier = absl::make_unique<int>(2);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"multiplier", Adopt(multiplier.release()).At(input_timestamp)));
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
// EndLoopCalc will forward the timestamp bound because there are no elements
|
||||
// in collection to output.
|
||||
ASSERT_EQ(0, output_packets_.size());
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, SingleNonEmptyVector) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||
input_vector->emplace_back(0);
|
||||
input_vector->emplace_back(1);
|
||||
input_vector->emplace_back(2);
|
||||
Timestamp input_timestamp = Timestamp(42);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||
auto multiplier = absl::make_unique<int>(2);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"multiplier", Adopt(multiplier.release()).At(input_timestamp)));
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
ASSERT_EQ(1, output_packets_.size());
|
||||
EXPECT_EQ(input_timestamp, output_packets_[0].Timestamp());
|
||||
std::vector<int> expected_output_vector = {0, 2, 4};
|
||||
EXPECT_EQ(expected_output_vector, output_packets_[0].Get<std::vector<int>>());
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, MultipleVectors) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
|
||||
auto input_vector0 = absl::make_unique<std::vector<int>>();
|
||||
input_vector0->emplace_back(0);
|
||||
input_vector0->emplace_back(1);
|
||||
Timestamp input_timestamp0 = Timestamp(42);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector0.release()).At(input_timestamp0)));
|
||||
auto multiplier0 = absl::make_unique<int>(2);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"multiplier", Adopt(multiplier0.release()).At(input_timestamp0)));
|
||||
|
||||
auto input_vector1 = absl::make_unique<std::vector<int>>();
|
||||
Timestamp input_timestamp1 = Timestamp(43);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector1.release()).At(input_timestamp1)));
|
||||
auto multiplier1 = absl::make_unique<int>(2);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"multiplier", Adopt(multiplier1.release()).At(input_timestamp1)));
|
||||
|
||||
auto input_vector2 = absl::make_unique<std::vector<int>>();
|
||||
input_vector2->emplace_back(2);
|
||||
input_vector2->emplace_back(3);
|
||||
Timestamp input_timestamp2 = Timestamp(44);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector2.release()).At(input_timestamp2)));
|
||||
auto multiplier2 = absl::make_unique<int>(3);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"multiplier", Adopt(multiplier2.release()).At(input_timestamp2)));
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
|
||||
ASSERT_EQ(2, output_packets_.size());
|
||||
|
||||
EXPECT_EQ(input_timestamp0, output_packets_[0].Timestamp());
|
||||
std::vector<int> expected_output_vector0 = {0, 2};
|
||||
EXPECT_EQ(expected_output_vector0,
|
||||
output_packets_[0].Get<std::vector<int>>());
|
||||
|
||||
// At input_timestamp1, EndLoopCalc will forward timestamp bound as there are
|
||||
// no elements in vector to process.
|
||||
|
||||
EXPECT_EQ(input_timestamp2, output_packets_[1].Timestamp());
|
||||
std::vector<int> expected_output_vector2 = {6, 9};
|
||||
EXPECT_EQ(expected_output_vector2,
|
||||
output_packets_[1].Get<std::vector<int>>());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,40 @@
|
||||
// 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 "mediapipe/calculators/core/begin_loop_calculator.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator to process std::vector<NormalizedLandmark>.
|
||||
typedef BeginLoopCalculator<std::vector<::mediapipe::NormalizedLandmark>>
|
||||
BeginLoopNormalizedLandmarkCalculator;
|
||||
REGISTER_CALCULATOR(BeginLoopNormalizedLandmarkCalculator);
|
||||
|
||||
// A calculator to process std::vector<std::vector<NormalizedLandmark>>.
|
||||
typedef BeginLoopCalculator<
|
||||
std::vector<std::vector<::mediapipe::NormalizedLandmark>>>
|
||||
BeginLoopNormalizedLandmarksVectorCalculator;
|
||||
REGISTER_CALCULATOR(BeginLoopNormalizedLandmarksVectorCalculator);
|
||||
|
||||
// A calculator to process std::vector<NormalizedRect>.
|
||||
typedef BeginLoopCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||
BeginLoopNormalizedRectCalculator;
|
||||
REGISTER_CALCULATOR(BeginLoopNormalizedRectCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,157 @@
|
||||
// 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_CORE_BEGIN_LOOP_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_CORE_BEGIN_LOOP_CALCULATOR_H_
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_contract.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Calculator for implementing loops on iterable collections inside a MediaPipe
|
||||
// graph.
|
||||
//
|
||||
// It is designed to be used like:
|
||||
//
|
||||
// node {
|
||||
// calculator: "BeginLoopWithIterableCalculator"
|
||||
// input_stream: "ITERABLE:input_iterable" # IterableT @ext_ts
|
||||
// output_stream: "ITEM:input_element" # ItemT @loop_internal_ts
|
||||
// output_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||
// }
|
||||
//
|
||||
// node {
|
||||
// calculator: "ElementToBlaConverterSubgraph"
|
||||
// input_stream: "ITEM:input_to_loop_body" # ItemT @loop_internal_ts
|
||||
// output_stream: "BLA:output_of_loop_body" # ItemU @loop_internal_ts
|
||||
// }
|
||||
//
|
||||
// node {
|
||||
// calculator: "EndLoopWithOutputCalculator"
|
||||
// input_stream: "ITEM:output_of_loop_body" # ItemU @loop_internal_ts
|
||||
// input_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||
// output_stream: "OUTPUT:aggregated_result" # IterableU @ext_ts
|
||||
// }
|
||||
//
|
||||
// BeginLoopCalculator accepts an optional input stream tagged with "TICK"
|
||||
// which if non-empty, wakes up the calculator and calls
|
||||
// BeginLoopCalculator::Process(). Input streams tagged with "CLONE" are cloned
|
||||
// to the corresponding output streams at loop timestamps. This ensures that a
|
||||
// MediaPipe graph or sub-graph can run multiple times, once per element in the
|
||||
// "ITERABLE" for each pakcet clone of the packets in the "CLONE" input streams.
|
||||
template <typename IterableT>
|
||||
class BeginLoopCalculator : public CalculatorBase {
|
||||
using ItemT = typename IterableT::value_type;
|
||||
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
// A non-empty packet in the optional "TICK" input stream wakes up the
|
||||
// calculator.
|
||||
if (cc->Inputs().HasTag("TICK")) {
|
||||
cc->Inputs().Tag("TICK").SetAny();
|
||||
}
|
||||
|
||||
// An iterable collection in the input stream.
|
||||
RET_CHECK(cc->Inputs().HasTag("ITERABLE"));
|
||||
cc->Inputs().Tag("ITERABLE").Set<IterableT>();
|
||||
|
||||
// An element from the collection.
|
||||
RET_CHECK(cc->Outputs().HasTag("ITEM"));
|
||||
cc->Outputs().Tag("ITEM").Set<ItemT>();
|
||||
|
||||
RET_CHECK(cc->Outputs().HasTag("BATCH_END"));
|
||||
cc->Outputs()
|
||||
.Tag("BATCH_END")
|
||||
.Set<Timestamp>(
|
||||
// A flush signal to the corresponding EndLoopCalculator for it to
|
||||
// emit the aggregated result with the timestamp contained in this
|
||||
// flush signal packet.
|
||||
);
|
||||
|
||||
// Input streams tagged with "CLONE" are cloned to the corresponding
|
||||
// "CLONE" output streams at loop timestamps.
|
||||
RET_CHECK(cc->Inputs().NumEntries("CLONE") ==
|
||||
cc->Outputs().NumEntries("CLONE"));
|
||||
if (cc->Inputs().NumEntries("CLONE") > 0) {
|
||||
for (int i = 0; i < cc->Inputs().NumEntries("CLONE"); ++i) {
|
||||
cc->Inputs().Get("CLONE", i).SetAny();
|
||||
cc->Outputs().Get("CLONE", i).SetSameAs(&cc->Inputs().Get("CLONE", i));
|
||||
}
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||
Timestamp last_timestamp = loop_internal_timestamp_;
|
||||
if (!cc->Inputs().Tag("ITERABLE").IsEmpty()) {
|
||||
const IterableT& collection =
|
||||
cc->Inputs().Tag("ITERABLE").template Get<IterableT>();
|
||||
for (const auto& item : collection) {
|
||||
cc->Outputs().Tag("ITEM").AddPacket(
|
||||
MakePacket<ItemT>(item).At(loop_internal_timestamp_));
|
||||
ForwardClonePackets(cc, loop_internal_timestamp_);
|
||||
++loop_internal_timestamp_;
|
||||
}
|
||||
}
|
||||
|
||||
// The collection was empty and nothing was processed.
|
||||
if (last_timestamp == loop_internal_timestamp_) {
|
||||
// Increment loop_internal_timestamp_ because it is used up now.
|
||||
++loop_internal_timestamp_;
|
||||
for (auto it = cc->Outputs().begin(); it < cc->Outputs().end(); ++it) {
|
||||
it->SetNextTimestampBound(loop_internal_timestamp_);
|
||||
}
|
||||
}
|
||||
|
||||
// The for loop processing the input collection already incremented
|
||||
// loop_internal_timestamp_. To emit BATCH_END packet along the last
|
||||
// non-BATCH_END packet, decrement by one.
|
||||
cc->Outputs()
|
||||
.Tag("BATCH_END")
|
||||
.AddPacket(MakePacket<Timestamp>(cc->InputTimestamp())
|
||||
.At(Timestamp(loop_internal_timestamp_ - 1)));
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
void ForwardClonePackets(CalculatorContext* cc, Timestamp output_timestamp) {
|
||||
if (cc->Inputs().NumEntries("CLONE") > 0) {
|
||||
for (int i = 0; i < cc->Inputs().NumEntries("CLONE"); ++i) {
|
||||
if (!cc->Inputs().Get("CLONE", i).IsEmpty()) {
|
||||
auto input_packet = cc->Inputs().Get("CLONE", i).Value();
|
||||
cc->Outputs()
|
||||
.Get("CLONE", i)
|
||||
.AddPacket(std::move(input_packet).At(output_timestamp));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fake timestamps generated per element in collection.
|
||||
Timestamp loop_internal_timestamp_ = Timestamp(0);
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_CORE_BEGIN_LOOP_CALCULATOR_H_
|
||||
@@ -0,0 +1,26 @@
|
||||
// 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/calculators/core/clip_vector_size_calculator.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef ClipVectorSizeCalculator<::mediapipe::Detection>
|
||||
ClipDetectionVectorSizeCalculator;
|
||||
REGISTER_CALCULATOR(ClipDetectionVectorSizeCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,28 @@
|
||||
// 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 "mediapipe/calculators/core/clip_vector_size_calculator.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef ClipVectorSizeCalculator<::mediapipe::NormalizedRect>
|
||||
ClipNormalizedRectVectorSizeCalculator;
|
||||
REGISTER_CALCULATOR(ClipNormalizedRectVectorSizeCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,137 @@
|
||||
// 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_CORE_CLIP_VECTOR_SIZE_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_CORE_CLIP_VECTOR_SIZE_CALCULATOR_H_
|
||||
|
||||
#include <type_traits>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/calculators/core/clip_vector_size_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Clips the size of the input vector of type T to a specified max_vec_size.
|
||||
// In a graph it will be used as:
|
||||
// node {
|
||||
// calculator: "ClipIntVectorSizeCalculator"
|
||||
// input_stream: "input_vector"
|
||||
// output_stream: "output_vector"
|
||||
// options {
|
||||
// [mediapipe.ClipIntVectorSizeCalculatorOptions.ext] {
|
||||
// max_vec_size: 5
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
template <typename T>
|
||||
class ClipVectorSizeCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().NumEntries() == 1);
|
||||
RET_CHECK(cc->Outputs().NumEntries() == 1);
|
||||
|
||||
if (cc->Options<::mediapipe::ClipVectorSizeCalculatorOptions>()
|
||||
.max_vec_size() < 1) {
|
||||
return ::mediapipe::InternalError(
|
||||
"max_vec_size should be greater than or equal to 1.");
|
||||
}
|
||||
|
||||
cc->Inputs().Index(0).Set<std::vector<T>>();
|
||||
cc->Outputs().Index(0).Set<std::vector<T>>();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
max_vec_size_ = cc->Options<::mediapipe::ClipVectorSizeCalculatorOptions>()
|
||||
.max_vec_size();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
if (max_vec_size_ < 1) {
|
||||
return ::mediapipe::InternalError(
|
||||
"max_vec_size should be greater than or equal to 1.");
|
||||
}
|
||||
if (cc->Inputs().Index(0).IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
return ClipVectorSize<T>(std::is_copy_constructible<T>(), cc);
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
::mediapipe::Status ClipVectorSize(std::true_type, CalculatorContext* cc) {
|
||||
auto output = absl::make_unique<std::vector<U>>();
|
||||
const std::vector<U>& input_vector =
|
||||
cc->Inputs().Index(0).Get<std::vector<U>>();
|
||||
if (max_vec_size_ >= input_vector.size()) {
|
||||
output->insert(output->end(), input_vector.begin(), input_vector.end());
|
||||
} else {
|
||||
for (int i = 0; i < max_vec_size_; ++i) {
|
||||
output->push_back(input_vector[i]);
|
||||
}
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
::mediapipe::Status ClipVectorSize(std::false_type, CalculatorContext* cc) {
|
||||
return ConsumeAndClipVectorSize<T>(std::is_move_constructible<U>(), cc);
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
::mediapipe::Status ConsumeAndClipVectorSize(std::true_type,
|
||||
CalculatorContext* cc) {
|
||||
auto output = absl::make_unique<std::vector<U>>();
|
||||
::mediapipe::StatusOr<std::unique_ptr<std::vector<U>>> input_status =
|
||||
cc->Inputs().Index(0).Value().Consume<std::vector<U>>();
|
||||
|
||||
if (input_status.ok()) {
|
||||
std::unique_ptr<std::vector<U>> input_vector =
|
||||
std::move(input_status).ValueOrDie();
|
||||
auto begin_it = input_vector->begin();
|
||||
auto end_it = input_vector->end();
|
||||
if (max_vec_size_ < input_vector->size()) {
|
||||
end_it = input_vector->begin() + max_vec_size_;
|
||||
}
|
||||
output->insert(output->end(), std::make_move_iterator(begin_it),
|
||||
std::make_move_iterator(end_it));
|
||||
} else {
|
||||
return input_status.status();
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
::mediapipe::Status ConsumeAndClipVectorSize(std::false_type,
|
||||
CalculatorContext* cc) {
|
||||
return ::mediapipe::InternalError(
|
||||
"Cannot copy or move input vectors and clip their size.");
|
||||
}
|
||||
|
||||
private:
|
||||
int max_vec_size_ = 0;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_CORE_CLIP_VECTOR_SIZE_CALCULATOR_H_
|
||||
@@ -0,0 +1,28 @@
|
||||
// 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 ClipVectorSizeCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional ClipVectorSizeCalculatorOptions ext = 274674998;
|
||||
}
|
||||
|
||||
// Maximum size of output vector.
|
||||
optional int32 max_vec_size = 1 [default = 1];
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
// 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 "mediapipe/calculators/core/clip_vector_size_calculator.h"
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#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_matchers.h" // NOLINT
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef ClipVectorSizeCalculator<int> TestClipIntVectorSizeCalculator;
|
||||
REGISTER_CALCULATOR(TestClipIntVectorSizeCalculator);
|
||||
|
||||
void AddInputVector(const std::vector<int>& input, int64 timestamp,
|
||||
CalculatorRunner* runner) {
|
||||
runner->MutableInputs()->Index(0).packets.push_back(
|
||||
MakePacket<std::vector<int>>(input).At(Timestamp(timestamp)));
|
||||
}
|
||||
|
||||
TEST(TestClipIntVectorSizeCalculatorTest, EmptyVectorInput) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "TestClipIntVectorSizeCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "output_vector"
|
||||
options {
|
||||
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 1 }
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
|
||||
std::vector<int> input = {};
|
||||
AddInputVector(input, /*timestamp=*/1, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||
EXPECT_TRUE(outputs[0].Get<std::vector<int>>().empty());
|
||||
}
|
||||
|
||||
TEST(TestClipIntVectorSizeCalculatorTest, OneTimestamp) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "TestClipIntVectorSizeCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "output_vector"
|
||||
options {
|
||||
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 2 }
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
|
||||
std::vector<int> input = {0, 1, 2, 3};
|
||||
AddInputVector(input, /*timestamp=*/1, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||
const std::vector<int>& output = outputs[0].Get<std::vector<int>>();
|
||||
EXPECT_EQ(2, output.size());
|
||||
std::vector<int> expected_vector = {0, 1};
|
||||
EXPECT_EQ(expected_vector, output);
|
||||
}
|
||||
|
||||
TEST(TestClipIntVectorSizeCalculatorTest, TwoInputsAtTwoTimestamps) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "TestClipIntVectorSizeCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "output_vector"
|
||||
options {
|
||||
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 3 }
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
|
||||
{
|
||||
std::vector<int> input = {0, 1, 2, 3};
|
||||
AddInputVector(input, /*timestamp=*/1, &runner);
|
||||
}
|
||||
{
|
||||
std::vector<int> input = {2, 3, 4, 5};
|
||||
AddInputVector(input, /*timestamp=*/2, &runner);
|
||||
}
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(2, outputs.size());
|
||||
{
|
||||
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||
const std::vector<int>& output = outputs[0].Get<std::vector<int>>();
|
||||
EXPECT_EQ(3, output.size());
|
||||
std::vector<int> expected_vector = {0, 1, 2};
|
||||
EXPECT_EQ(expected_vector, output);
|
||||
}
|
||||
{
|
||||
EXPECT_EQ(Timestamp(2), outputs[1].Timestamp());
|
||||
const std::vector<int>& output = outputs[1].Get<std::vector<int>>();
|
||||
EXPECT_EQ(3, output.size());
|
||||
std::vector<int> expected_vector = {2, 3, 4};
|
||||
EXPECT_EQ(expected_vector, output);
|
||||
}
|
||||
}
|
||||
|
||||
typedef ClipVectorSizeCalculator<std::unique_ptr<int>>
|
||||
TestClipUniqueIntPtrVectorSizeCalculator;
|
||||
REGISTER_CALCULATOR(TestClipUniqueIntPtrVectorSizeCalculator);
|
||||
|
||||
TEST(TestClipUniqueIntPtrVectorSizeCalculatorTest, 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: "input_vector"
|
||||
node {
|
||||
calculator: "TestClipUniqueIntPtrVectorSizeCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "output_vector"
|
||||
options {
|
||||
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 3 }
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
std::vector<Packet> outputs;
|
||||
tool::AddVectorSink("output_vector", &graph_config, &outputs);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
|
||||
// input1 : {0, 1, 2, 3, 4, 5}
|
||||
auto input_vector = absl::make_unique<std::vector<std::unique_ptr<int>>>(6);
|
||||
for (int i = 0; i < 6; ++i) {
|
||||
input_vector->at(i) = absl::make_unique<int>(i);
|
||||
}
|
||||
|
||||
MP_EXPECT_OK(graph.AddPacketToInputStream(
|
||||
"input_vector", Adopt(input_vector.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);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -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
|
||||
|
||||
@@ -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_;
|
||||
};
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
// 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 <cfloat>
|
||||
|
||||
#include "mediapipe/calculators/core/dequantize_byte_array_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
// Dequantizes a byte array to a vector of floats.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "DequantizeByteArrayCalculator"
|
||||
// input_stream: "ENCODED:encoded"
|
||||
// output_stream: "FLOAT_VECTOR:float_vector"
|
||||
// options {
|
||||
// [mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||
// max_quantized_value: 2
|
||||
// min_quantized_value: -2
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
namespace mediapipe {
|
||||
|
||||
class DequantizeByteArrayCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Tag("ENCODED").Set<std::string>();
|
||||
cc->Outputs().Tag("FLOAT_VECTOR").Set<std::vector<float>>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) final {
|
||||
const auto options =
|
||||
cc->Options<::mediapipe::DequantizeByteArrayCalculatorOptions>();
|
||||
if (!options.has_max_quantized_value() ||
|
||||
!options.has_min_quantized_value()) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Both max_quantized_value and min_quantized_value must be provided "
|
||||
"in DequantizeByteArrayCalculatorOptions.");
|
||||
}
|
||||
float max_quantized_value = options.max_quantized_value();
|
||||
float min_quantized_value = options.min_quantized_value();
|
||||
if (max_quantized_value < min_quantized_value + FLT_EPSILON) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"max_quantized_value must be greater than min_quantized_value.");
|
||||
}
|
||||
float range = max_quantized_value - min_quantized_value;
|
||||
scalar_ = range / 255.0;
|
||||
bias_ = (range / 512.0) + min_quantized_value;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||
const std::string& encoded =
|
||||
cc->Inputs().Tag("ENCODED").Value().Get<std::string>();
|
||||
std::vector<float> float_vector;
|
||||
float_vector.reserve(encoded.length());
|
||||
for (int i = 0; i < encoded.length(); ++i) {
|
||||
float_vector.push_back(
|
||||
static_cast<unsigned char>(encoded.at(i)) * scalar_ + bias_);
|
||||
}
|
||||
cc->Outputs()
|
||||
.Tag("FLOAT_VECTOR")
|
||||
.AddPacket(MakePacket<std::vector<float>>(float_vector)
|
||||
.At(cc->InputTimestamp()));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
float scalar_;
|
||||
float bias_;
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(DequantizeByteArrayCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,28 @@
|
||||
// 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 DequantizeByteArrayCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional DequantizeByteArrayCalculatorOptions ext = 272316343;
|
||||
}
|
||||
|
||||
optional float max_quantized_value = 1;
|
||||
optional float min_quantized_value = 2;
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
// 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 <string>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#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" // NOLINT
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
TEST(QuantizeFloatVectorCalculatorTest, WrongConfig) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "DequantizeByteArrayCalculator"
|
||||
input_stream: "ENCODED:encoded"
|
||||
output_stream: "FLOAT_VECTOR:float_vector"
|
||||
options {
|
||||
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||
max_quantized_value: 2
|
||||
}
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
std::string empty_string;
|
||||
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||
MakePacket<std::string>(empty_string).At(Timestamp(0)));
|
||||
auto status = runner.Run();
|
||||
EXPECT_FALSE(status.ok());
|
||||
EXPECT_THAT(
|
||||
status.message(),
|
||||
testing::HasSubstr(
|
||||
"Both max_quantized_value and min_quantized_value must be provided"));
|
||||
}
|
||||
|
||||
TEST(QuantizeFloatVectorCalculatorTest, WrongConfig2) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "DequantizeByteArrayCalculator"
|
||||
input_stream: "ENCODED:encoded"
|
||||
output_stream: "FLOAT_VECTOR:float_vector"
|
||||
options {
|
||||
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||
max_quantized_value: -2
|
||||
min_quantized_value: 2
|
||||
}
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
std::string empty_string;
|
||||
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||
MakePacket<std::string>(empty_string).At(Timestamp(0)));
|
||||
auto status = runner.Run();
|
||||
EXPECT_FALSE(status.ok());
|
||||
EXPECT_THAT(
|
||||
status.message(),
|
||||
testing::HasSubstr(
|
||||
"max_quantized_value must be greater than min_quantized_value"));
|
||||
}
|
||||
|
||||
TEST(QuantizeFloatVectorCalculatorTest, WrongConfig3) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "DequantizeByteArrayCalculator"
|
||||
input_stream: "ENCODED:encoded"
|
||||
output_stream: "FLOAT_VECTOR:float_vector"
|
||||
options {
|
||||
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||
max_quantized_value: 1
|
||||
min_quantized_value: 1
|
||||
}
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
std::string empty_string;
|
||||
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||
MakePacket<std::string>(empty_string).At(Timestamp(0)));
|
||||
auto status = runner.Run();
|
||||
EXPECT_FALSE(status.ok());
|
||||
EXPECT_THAT(
|
||||
status.message(),
|
||||
testing::HasSubstr(
|
||||
"max_quantized_value must be greater than min_quantized_value"));
|
||||
}
|
||||
|
||||
TEST(DequantizeByteArrayCalculatorTest, TestDequantization) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "DequantizeByteArrayCalculator"
|
||||
input_stream: "ENCODED:encoded"
|
||||
output_stream: "FLOAT_VECTOR:float_vector"
|
||||
options {
|
||||
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||
max_quantized_value: 2
|
||||
min_quantized_value: -2
|
||||
}
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
unsigned char input[4] = {0x7F, 0xFF, 0x00, 0x01};
|
||||
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||
MakePacket<std::string>(
|
||||
std::string(reinterpret_cast<char const*>(input), 4))
|
||||
.At(Timestamp(0)));
|
||||
auto status = runner.Run();
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
const std::vector<Packet>& outputs =
|
||||
runner.Outputs().Tag("FLOAT_VECTOR").packets;
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
const std::vector<float>& result = outputs[0].Get<std::vector<float>>();
|
||||
ASSERT_FALSE(result.empty());
|
||||
EXPECT_EQ(4, result.size());
|
||||
EXPECT_NEAR(0, result[0], 0.01);
|
||||
EXPECT_NEAR(2, result[1], 0.01);
|
||||
EXPECT_NEAR(-2, result[2], 0.01);
|
||||
EXPECT_NEAR(-1.976, result[3], 0.01);
|
||||
|
||||
EXPECT_EQ(Timestamp(0), outputs[0].Timestamp());
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,45 @@
|
||||
// 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 "mediapipe/calculators/core/end_loop_calculator.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/util/render_data.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef EndLoopCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||
EndLoopNormalizedRectCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopNormalizedRectCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<::mediapipe::NormalizedLandmark>>
|
||||
EndLoopNormalizedLandmarkCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopNormalizedLandmarkCalculator);
|
||||
|
||||
typedef EndLoopCalculator<
|
||||
std::vector<std::vector<::mediapipe::NormalizedLandmark>>>
|
||||
EndLoopNormalizedLandmarksVectorCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopNormalizedLandmarksVectorCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<bool>> EndLoopBooleanCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopBooleanCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<::mediapipe::RenderData>>
|
||||
EndLoopRenderDataCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopRenderDataCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,106 @@
|
||||
// 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_CORE_END_LOOP_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_CORE_END_LOOP_CALCULATOR_H_
|
||||
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_contract.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Calculator for completing the processing of loops on iterable collections
|
||||
// inside a MediaPipe graph. The EndLoopCalculator collects all input packets
|
||||
// from ITEM input_stream into a collection and upon receiving the flush signal
|
||||
// from the "BATCH_END" tagged input stream, it emits the aggregated results
|
||||
// at the original timestamp contained in the "BATCH_END" input stream.
|
||||
//
|
||||
// It is designed to be used like:
|
||||
//
|
||||
// node {
|
||||
// calculator: "BeginLoopWithIterableCalculator"
|
||||
// input_stream: "ITERABLE:input_iterable" # IterableT @ext_ts
|
||||
// output_stream: "ITEM:input_element" # ItemT @loop_internal_ts
|
||||
// output_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||
// }
|
||||
//
|
||||
// node {
|
||||
// calculator: "ElementToBlaConverterSubgraph"
|
||||
// input_stream: "ITEM:input_to_loop_body" # ItemT @loop_internal_ts
|
||||
// output_stream: "BLA:output_of_loop_body" # ItemU @loop_internal_ts
|
||||
// }
|
||||
//
|
||||
// node {
|
||||
// calculator: "EndLoopWithOutputCalculator"
|
||||
// input_stream: "ITEM:output_of_loop_body" # ItemU @loop_internal_ts
|
||||
// input_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||
// output_stream: "OUTPUT:aggregated_result" # IterableU @ext_ts
|
||||
// }
|
||||
template <typename IterableT>
|
||||
class EndLoopCalculator : public CalculatorBase {
|
||||
using ItemT = typename IterableT::value_type;
|
||||
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag("BATCH_END"))
|
||||
<< "Missing BATCH_END tagged input_stream.";
|
||||
cc->Inputs().Tag("BATCH_END").Set<Timestamp>();
|
||||
|
||||
RET_CHECK(cc->Inputs().HasTag("ITEM"));
|
||||
cc->Inputs().Tag("ITEM").Set<ItemT>();
|
||||
|
||||
RET_CHECK(cc->Outputs().HasTag("ITERABLE"));
|
||||
cc->Outputs().Tag("ITERABLE").Set<IterableT>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
if (!cc->Inputs().Tag("ITEM").IsEmpty()) {
|
||||
if (!input_stream_collection_) {
|
||||
input_stream_collection_.reset(new IterableT);
|
||||
}
|
||||
input_stream_collection_->push_back(
|
||||
cc->Inputs().Tag("ITEM").template Get<ItemT>());
|
||||
}
|
||||
|
||||
if (!cc->Inputs().Tag("BATCH_END").Value().IsEmpty()) { // flush signal
|
||||
Timestamp loop_control_ts =
|
||||
cc->Inputs().Tag("BATCH_END").template Get<Timestamp>();
|
||||
if (input_stream_collection_) {
|
||||
cc->Outputs()
|
||||
.Tag("ITERABLE")
|
||||
.Add(input_stream_collection_.release(), loop_control_ts);
|
||||
} else {
|
||||
// Since there is no collection, inform downstream calculators to not
|
||||
// expect any packet by updating the timestamp bounds.
|
||||
cc->Outputs()
|
||||
.Tag("ITERABLE")
|
||||
.SetNextTimestampBound(Timestamp(loop_control_ts.Value() + 1));
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<IterableT> input_stream_collection_;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_CORE_END_LOOP_CALCULATOR_H_
|
||||
@@ -74,6 +74,12 @@ class PacketResamplerCalculator : public CalculatorBase {
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
// Calculates the first sampled timestamp that incorporates a jittering
|
||||
// offset.
|
||||
void InitializeNextOutputTimestampWithJitter();
|
||||
// Calculates the next sampled timestamp that incorporates a jittering offset.
|
||||
void UpdateNextOutputTimestampWithJitter();
|
||||
|
||||
// Logic for Process() when jitter_ != 0.0.
|
||||
::mediapipe::Status ProcessWithJitter(CalculatorContext* cc);
|
||||
|
||||
@@ -233,6 +239,7 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
||||
<< Timestamp::kTimestampUnitsPerSecond;
|
||||
|
||||
frame_time_usec_ = static_cast<int64>(1000000.0 / frame_rate_);
|
||||
|
||||
video_header_.frame_rate = frame_rate_;
|
||||
|
||||
if (resampler_options.output_header() !=
|
||||
@@ -295,6 +302,17 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
void PacketResamplerCalculator::InitializeNextOutputTimestampWithJitter() {
|
||||
next_output_timestamp_ =
|
||||
first_timestamp_ + frame_time_usec_ * random_->RandFloat();
|
||||
}
|
||||
|
||||
void PacketResamplerCalculator::UpdateNextOutputTimestampWithJitter() {
|
||||
next_output_timestamp_ +=
|
||||
frame_time_usec_ *
|
||||
((1.0 - jitter_) + 2.0 * jitter_ * random_->RandFloat());
|
||||
}
|
||||
|
||||
::mediapipe::Status PacketResamplerCalculator::ProcessWithJitter(
|
||||
CalculatorContext* cc) {
|
||||
RET_CHECK_GT(cc->InputTimestamp(), Timestamp::PreStream());
|
||||
@@ -302,8 +320,13 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
||||
|
||||
if (first_timestamp_ == Timestamp::Unset()) {
|
||||
first_timestamp_ = cc->InputTimestamp();
|
||||
next_output_timestamp_ =
|
||||
first_timestamp_ + frame_time_usec_ * random_->RandFloat();
|
||||
InitializeNextOutputTimestampWithJitter();
|
||||
if (first_timestamp_ == next_output_timestamp_) {
|
||||
OutputWithinLimits(
|
||||
cc,
|
||||
cc->Inputs().Get(input_data_id_).Value().At(next_output_timestamp_));
|
||||
UpdateNextOutputTimestampWithJitter();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -322,9 +345,7 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
||||
? last_packet_
|
||||
: cc->Inputs().Get(input_data_id_).Value())
|
||||
.At(next_output_timestamp_));
|
||||
next_output_timestamp_ +=
|
||||
frame_time_usec_ *
|
||||
((1.0 - jitter_) + 2.0 * jitter_ * random_->RandFloat());
|
||||
UpdateNextOutputTimestampWithJitter();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
|
||||
@@ -102,6 +102,12 @@ class PreviousLoopbackCalculator : public CalculatorBase {
|
||||
cc->Outputs().Get(loop_out_id_).AddPacket(std::move(previous_loopback));
|
||||
}
|
||||
}
|
||||
if (!main_ts_.empty()) {
|
||||
cc->Outputs().Get(loop_out_id_).SetNextTimestampBound(main_ts_.front());
|
||||
}
|
||||
if (cc->Inputs().Get(main_id_).IsDone() && main_ts_.empty()) {
|
||||
cc->Outputs().Get(loop_out_id_).Close();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
|
||||
@@ -107,5 +107,96 @@ TEST(PreviousLoopbackCalculator, CorrectTimestamps) {
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
// A Calculator that outputs a summary packet in CalculatorBase::Close().
|
||||
class PacketOnCloseCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Index(0).Set<int>();
|
||||
cc->Outputs().Index(0).Set<int>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) final {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||
sum_ += cc->Inputs().Index(0).Value().Get<int>();
|
||||
cc->Outputs().Index(0).AddPacket(cc->Inputs().Index(0).Value());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Close(CalculatorContext* cc) final {
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
MakePacket<int>(sum_).At(Timestamp::Max()));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
int sum_ = 0;
|
||||
};
|
||||
REGISTER_CALCULATOR(PacketOnCloseCalculator);
|
||||
|
||||
// Demonstrates that all ouput and input streams in PreviousLoopbackCalculator
|
||||
// will close as expected when all graph input streams are closed.
|
||||
TEST(PreviousLoopbackCalculator, ClosesCorrectly) {
|
||||
std::vector<Packet> outputs;
|
||||
CalculatorGraphConfig graph_config_ =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: 'in'
|
||||
node {
|
||||
calculator: 'PreviousLoopbackCalculator'
|
||||
input_stream: 'MAIN:in'
|
||||
input_stream: 'LOOP:out'
|
||||
input_stream_info: { tag_index: 'LOOP' back_edge: true }
|
||||
output_stream: 'PREV_LOOP:previous'
|
||||
}
|
||||
# This calculator synchronizes its inputs as normal, so it is used
|
||||
# to check that both "in" and "previous" are ready.
|
||||
node {
|
||||
calculator: 'PassThroughCalculator'
|
||||
input_stream: 'in'
|
||||
input_stream: 'previous'
|
||||
output_stream: 'out'
|
||||
output_stream: 'previous2'
|
||||
}
|
||||
node {
|
||||
calculator: 'PacketOnCloseCalculator'
|
||||
input_stream: 'out'
|
||||
output_stream: 'close_out'
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("close_out", &graph_config_, &outputs);
|
||||
|
||||
CalculatorGraph graph_;
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config_, {}));
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
|
||||
auto send_packet = [&graph_](const std::string& input_name, int n) {
|
||||
MP_EXPECT_OK(graph_.AddPacketToInputStream(
|
||||
input_name, MakePacket<int>(n).At(Timestamp(n))));
|
||||
};
|
||||
|
||||
send_packet("in", 1);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1}));
|
||||
|
||||
send_packet("in", 5);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1, 5}));
|
||||
|
||||
send_packet("in", 15);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1, 5, 15}));
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs),
|
||||
(std::vector<int64>{1, 5, 15, Timestamp::Max().Value()}));
|
||||
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
// 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 <map>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <string>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
using mediapipe::PacketTypeSet;
|
||||
using mediapipe::Timestamp;
|
||||
|
||||
namespace {
|
||||
static std::map<std::string, Timestamp>* kTimestampMap = []() {
|
||||
auto* res = new std::map<std::string, Timestamp>();
|
||||
res->emplace("AT_PRESTREAM", Timestamp::PreStream());
|
||||
res->emplace("AT_POSTSTREAM", Timestamp::PostStream());
|
||||
res->emplace("AT_ZERO", Timestamp(0));
|
||||
return res;
|
||||
}();
|
||||
|
||||
} // namespace
|
||||
|
||||
// Outputs the single input_side_packet at the timestamp specified in the
|
||||
// output_stream tag. Valid tags are AT_PRESTREAM, AT_POSTSTREAM and AT_ZERO.
|
||||
class SidePacketToStreamCalculator : public CalculatorBase {
|
||||
public:
|
||||
SidePacketToStreamCalculator() = default;
|
||||
~SidePacketToStreamCalculator() override = default;
|
||||
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
};
|
||||
REGISTER_CALCULATOR(SidePacketToStreamCalculator);
|
||||
|
||||
::mediapipe::Status SidePacketToStreamCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
cc->InputSidePackets().Index(0).SetAny();
|
||||
|
||||
std::set<std::string> tags = cc->Outputs().GetTags();
|
||||
RET_CHECK_EQ(tags.size(), 1);
|
||||
|
||||
RET_CHECK_EQ(kTimestampMap->count(*tags.begin()), 1);
|
||||
cc->Outputs().Tag(*tags.begin()).SetAny();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status SidePacketToStreamCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
return mediapipe::tool::StatusStop();
|
||||
}
|
||||
|
||||
::mediapipe::Status SidePacketToStreamCalculator::Close(CalculatorContext* cc) {
|
||||
std::set<std::string> tags = cc->Outputs().GetTags();
|
||||
RET_CHECK_EQ(tags.size(), 1);
|
||||
const std::string& tag = *tags.begin();
|
||||
RET_CHECK_EQ(kTimestampMap->count(tag), 1);
|
||||
cc->Outputs().Tag(tag).AddPacket(
|
||||
cc->InputSidePackets().Index(0).At(kTimestampMap->at(tag)));
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -34,7 +34,9 @@ namespace mediapipe {
|
||||
// SplitVectorCalculatorOptions. If the option "element_only" is set to true,
|
||||
// all ranges should be of size 1 and all outputs will be elements of type T. If
|
||||
// "element_only" is false, ranges can be non-zero in size and all outputs will
|
||||
// be of type std::vector<T>.
|
||||
// be of type std::vector<T>. If the option "combine_outputs" is set to true,
|
||||
// only one output stream can be specified and all ranges of elements will be
|
||||
// combined into one vector.
|
||||
// To use this class for a particular type T, register a calculator using
|
||||
// SplitVectorCalculator<T>.
|
||||
template <typename T>
|
||||
@@ -49,28 +51,47 @@ class SplitVectorCalculator : public CalculatorBase {
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
|
||||
|
||||
if (cc->Outputs().NumEntries() != options.ranges_size()) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"The number of output streams should match the number of ranges "
|
||||
"specified in the CalculatorOptions.");
|
||||
}
|
||||
|
||||
// Set the output types for each output stream.
|
||||
for (int i = 0; i < cc->Outputs().NumEntries(); ++i) {
|
||||
if (options.ranges(i).begin() < 0 || options.ranges(i).end() < 0 ||
|
||||
options.ranges(i).begin() >= options.ranges(i).end()) {
|
||||
if (options.combine_outputs()) {
|
||||
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1);
|
||||
cc->Outputs().Index(0).Set<std::vector<T>>();
|
||||
for (int i = 0; i < options.ranges_size() - 1; ++i) {
|
||||
for (int j = i + 1; j < options.ranges_size(); ++j) {
|
||||
const auto& range_0 = options.ranges(i);
|
||||
const auto& range_1 = options.ranges(j);
|
||||
if ((range_0.begin() >= range_1.begin() &&
|
||||
range_0.begin() < range_1.end()) ||
|
||||
(range_1.begin() >= range_0.begin() &&
|
||||
range_1.begin() < range_0.end())) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Ranges must be non-overlapping when using combine_outputs "
|
||||
"option.");
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (cc->Outputs().NumEntries() != options.ranges_size()) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Indices should be non-negative and begin index should be less "
|
||||
"than the end index.");
|
||||
"The number of output streams should match the number of ranges "
|
||||
"specified in the CalculatorOptions.");
|
||||
}
|
||||
if (options.element_only()) {
|
||||
if (options.ranges(i).end() - options.ranges(i).begin() != 1) {
|
||||
|
||||
// Set the output types for each output stream.
|
||||
for (int i = 0; i < cc->Outputs().NumEntries(); ++i) {
|
||||
if (options.ranges(i).begin() < 0 || options.ranges(i).end() < 0 ||
|
||||
options.ranges(i).begin() >= options.ranges(i).end()) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Since element_only is true, all ranges should be of size 1.");
|
||||
"Indices should be non-negative and begin index should be less "
|
||||
"than the end index.");
|
||||
}
|
||||
if (options.element_only()) {
|
||||
if (options.ranges(i).end() - options.ranges(i).begin() != 1) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Since element_only is true, all ranges should be of size 1.");
|
||||
}
|
||||
cc->Outputs().Index(i).Set<T>();
|
||||
} else {
|
||||
cc->Outputs().Index(i).Set<std::vector<T>>();
|
||||
}
|
||||
cc->Outputs().Index(i).Set<T>();
|
||||
} else {
|
||||
cc->Outputs().Index(i).Set<std::vector<T>>();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -83,13 +104,15 @@ class SplitVectorCalculator : public CalculatorBase {
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
|
||||
|
||||
element_only_ = options.element_only();
|
||||
combine_outputs_ = options.combine_outputs();
|
||||
|
||||
for (const auto& range : options.ranges()) {
|
||||
ranges_.push_back({range.begin(), range.end()});
|
||||
max_range_end_ = std::max(max_range_end_, range.end());
|
||||
total_elements_ += range.end() - range.begin();
|
||||
}
|
||||
|
||||
element_only_ = options.element_only();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -97,17 +120,29 @@ class SplitVectorCalculator : public CalculatorBase {
|
||||
const auto& input = cc->Inputs().Index(0).Get<std::vector<T>>();
|
||||
RET_CHECK_GE(input.size(), max_range_end_);
|
||||
|
||||
if (element_only_) {
|
||||
if (combine_outputs_) {
|
||||
auto output = absl::make_unique<std::vector<T>>();
|
||||
output->reserve(total_elements_);
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
cc->Outputs().Index(i).AddPacket(
|
||||
MakePacket<T>(input[ranges_[i].first]).At(cc->InputTimestamp()));
|
||||
}
|
||||
} else {
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
auto output = absl::make_unique<std::vector<T>>(
|
||||
auto elements = absl::make_unique<std::vector<T>>(
|
||||
input.begin() + ranges_[i].first,
|
||||
input.begin() + ranges_[i].second);
|
||||
cc->Outputs().Index(i).Add(output.release(), cc->InputTimestamp());
|
||||
output->insert(output->end(), elements->begin(), elements->end());
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||
} else {
|
||||
if (element_only_) {
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
cc->Outputs().Index(i).AddPacket(
|
||||
MakePacket<T>(input[ranges_[i].first]).At(cc->InputTimestamp()));
|
||||
}
|
||||
} else {
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
auto output = absl::make_unique<std::vector<T>>(
|
||||
input.begin() + ranges_[i].first,
|
||||
input.begin() + ranges_[i].second);
|
||||
cc->Outputs().Index(i).Add(output.release(), cc->InputTimestamp());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -117,7 +152,9 @@ class SplitVectorCalculator : public CalculatorBase {
|
||||
private:
|
||||
std::vector<std::pair<int32, int32>> ranges_;
|
||||
int32 max_range_end_ = -1;
|
||||
int32 total_elements_ = 0;
|
||||
bool element_only_ = false;
|
||||
bool combine_outputs_ = false;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -37,4 +37,7 @@ message SplitVectorCalculatorOptions {
|
||||
// just element of type T. By default, if a range specifies only one element,
|
||||
// it is outputted as an std::vector<T>.
|
||||
optional bool element_only = 2 [default = false];
|
||||
|
||||
// Combines output elements to one vector.
|
||||
optional bool combine_outputs = 3 [default = false];
|
||||
}
|
||||
|
||||
@@ -105,6 +105,34 @@ class SplitTfLiteTensorVectorCalculatorTest : public ::testing::Test {
|
||||
}
|
||||
}
|
||||
|
||||
void ValidateCombinedVectorOutput(std::vector<Packet>& output_packets,
|
||||
int expected_elements,
|
||||
std::vector<int>& input_begin_indices,
|
||||
std::vector<int>& input_end_indices) {
|
||||
ASSERT_EQ(1, output_packets.size());
|
||||
ASSERT_EQ(input_begin_indices.size(), input_end_indices.size());
|
||||
const std::vector<TfLiteTensor>& output_vec =
|
||||
output_packets[0].Get<std::vector<TfLiteTensor>>();
|
||||
ASSERT_EQ(expected_elements, output_vec.size());
|
||||
const int num_ranges = input_begin_indices.size();
|
||||
|
||||
int element_id = 0;
|
||||
for (int range_id = 0; range_id < num_ranges; ++range_id) {
|
||||
for (int i = input_begin_indices[range_id];
|
||||
i < input_end_indices[range_id]; ++i) {
|
||||
const int expected_value = i;
|
||||
const TfLiteTensor* result = &output_vec[element_id];
|
||||
float* result_buffer = result->data.f;
|
||||
ASSERT_NE(result_buffer, nullptr);
|
||||
ASSERT_EQ(result_buffer, input_buffers_[i]);
|
||||
for (int j = 0; j < width * height * channels; ++j) {
|
||||
ASSERT_EQ(expected_value, result_buffer[j]);
|
||||
}
|
||||
element_id++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void ValidateElementOutput(std::vector<Packet>& output_packets,
|
||||
int input_begin_index) {
|
||||
ASSERT_EQ(1, output_packets.size());
|
||||
@@ -234,6 +262,65 @@ TEST_F(SplitTfLiteTensorVectorCalculatorTest, InvalidOutputStreamCountTest) {
|
||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||
}
|
||||
|
||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest,
|
||||
InvalidCombineOutputsMultipleOutputsTest) {
|
||||
ASSERT_NE(interpreter_, nullptr);
|
||||
|
||||
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "tensor_in"
|
||||
node {
|
||||
calculator: "SplitTfLiteTensorVectorCalculator"
|
||||
input_stream: "tensor_in"
|
||||
output_stream: "range_0"
|
||||
output_stream: "range_1"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 1 }
|
||||
ranges: { begin: 2 end: 3 }
|
||||
combine_outputs: true
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
// The graph should fail running because the number of output streams does not
|
||||
// match the number of range elements in the options.
|
||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||
}
|
||||
|
||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest, InvalidOverlappingRangesTest) {
|
||||
ASSERT_NE(interpreter_, nullptr);
|
||||
|
||||
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "tensor_in"
|
||||
node {
|
||||
calculator: "SplitTfLiteTensorVectorCalculator"
|
||||
input_stream: "tensor_in"
|
||||
output_stream: "range_0"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 3 }
|
||||
ranges: { begin: 1 end: 4 }
|
||||
combine_outputs: true
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
// The graph should fail running because there are overlapping ranges.
|
||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||
}
|
||||
|
||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest, SmokeTestElementOnly) {
|
||||
ASSERT_NE(interpreter_, nullptr);
|
||||
|
||||
@@ -289,6 +376,53 @@ TEST_F(SplitTfLiteTensorVectorCalculatorTest, SmokeTestElementOnly) {
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest, SmokeTestCombiningOutputs) {
|
||||
ASSERT_NE(interpreter_, nullptr);
|
||||
|
||||
PrepareTfLiteTensorVector(/*vector_size=*/5);
|
||||
ASSERT_NE(input_vec_, nullptr);
|
||||
|
||||
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "tensor_in"
|
||||
node {
|
||||
calculator: "SplitTfLiteTensorVectorCalculator"
|
||||
input_stream: "tensor_in"
|
||||
output_stream: "range_0"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 1 }
|
||||
ranges: { begin: 2 end: 3 }
|
||||
ranges: { begin: 4 end: 5 }
|
||||
combine_outputs: true
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
std::vector<Packet> range_0_packets;
|
||||
tool::AddVectorSink("range_0", &graph_config, &range_0_packets);
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"tensor_in", Adopt(input_vec_.release()).At(Timestamp(0))));
|
||||
// Wait until the calculator finishes processing.
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
std::vector<int> input_begin_indices = {0, 2, 4};
|
||||
std::vector<int> input_end_indices = {1, 3, 5};
|
||||
ValidateCombinedVectorOutput(range_0_packets, /*expected_elements=*/3,
|
||||
input_begin_indices, input_end_indices);
|
||||
|
||||
// Fully close the graph at the end.
|
||||
MP_ASSERT_OK(graph.CloseInputStream("tensor_in"));
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest,
|
||||
ElementOnlyDisablesVectorOutputs) {
|
||||
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
// 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 "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator that takes a packet of an input stream and converts it to an
|
||||
// output side packet. This calculator only works under the assumption that the
|
||||
// input stream only has a single packet passing through.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "StreamToSidePacketCalculator"
|
||||
// input_stream: "stream"
|
||||
// output_side_packet: "side_packet"
|
||||
// }
|
||||
class StreamToSidePacketCalculator : public mediapipe::CalculatorBase {
|
||||
public:
|
||||
static mediapipe::Status GetContract(mediapipe::CalculatorContract* cc) {
|
||||
cc->Inputs().Index(0).SetAny();
|
||||
cc->OutputSidePackets().Index(0).SetAny();
|
||||
return mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
mediapipe::Status Process(mediapipe::CalculatorContext* cc) override {
|
||||
mediapipe::Packet& packet = cc->Inputs().Index(0).Value();
|
||||
cc->OutputSidePackets().Index(0).Set(
|
||||
packet.At(mediapipe::Timestamp::Unset()));
|
||||
return mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
REGISTER_CALCULATOR(StreamToSidePacketCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,67 @@
|
||||
// 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 <memory>
|
||||
#include <string>
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
using ::testing::Test;
|
||||
|
||||
class StreamToSidePacketCalculatorTest : public Test {
|
||||
protected:
|
||||
StreamToSidePacketCalculatorTest() {
|
||||
const char kConfig[] = R"(
|
||||
calculator: "StreamToSidePacketCalculator"
|
||||
input_stream: "stream"
|
||||
output_side_packet: "side_packet"
|
||||
)";
|
||||
runner_ = absl::make_unique<CalculatorRunner>(kConfig);
|
||||
}
|
||||
|
||||
std::unique_ptr<CalculatorRunner> runner_;
|
||||
};
|
||||
|
||||
TEST_F(StreamToSidePacketCalculatorTest,
|
||||
StreamToSidePacketCalculatorWithEmptyStreamFails) {
|
||||
EXPECT_EQ(runner_->Run().code(), mediapipe::StatusCode::kUnavailable);
|
||||
}
|
||||
|
||||
TEST_F(StreamToSidePacketCalculatorTest,
|
||||
StreamToSidePacketCalculatorWithSinglePacketCreatesSidePacket) {
|
||||
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(new std::string("test")).At(Timestamp(1)));
|
||||
MP_ASSERT_OK(runner_->Run());
|
||||
EXPECT_EQ(runner_->OutputSidePackets().Index(0).Get<std::string>(), "test");
|
||||
}
|
||||
|
||||
TEST_F(StreamToSidePacketCalculatorTest,
|
||||
StreamToSidePacketCalculatorWithMultiplePacketsFails) {
|
||||
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(new std::string("test1")).At(Timestamp(1)));
|
||||
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(new std::string("test2")).At(Timestamp(2)));
|
||||
EXPECT_EQ(runner_->Run().code(), mediapipe::StatusCode::kAlreadyExists);
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,79 @@
|
||||
// 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 <sys/types.h>
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#include "absl/strings/numbers.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Calculator that converts a std::string into an integer type, or fails if the
|
||||
// conversion is not possible.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "StringToIntCalculator"
|
||||
// input_side_packet: "string"
|
||||
// output_side_packet: "index"
|
||||
// }
|
||||
template <typename IntType>
|
||||
class StringToIntCalculatorTemplate : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->InputSidePackets().Index(0).Set<std::string>();
|
||||
cc->OutputSidePackets().Index(0).Set<IntType>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
IntType number;
|
||||
if (!absl::SimpleAtoi(cc->InputSidePackets().Index(0).Get<std::string>(),
|
||||
&number)) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"The std::string could not be parsed as an integer.");
|
||||
}
|
||||
cc->OutputSidePackets().Index(0).Set(MakePacket<IntType>(number));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
using StringToIntCalculator = StringToIntCalculatorTemplate<int>;
|
||||
REGISTER_CALCULATOR(StringToIntCalculator);
|
||||
|
||||
using StringToUintCalculator = StringToIntCalculatorTemplate<uint>;
|
||||
REGISTER_CALCULATOR(StringToUintCalculator);
|
||||
|
||||
using StringToInt32Calculator = StringToIntCalculatorTemplate<int32>;
|
||||
REGISTER_CALCULATOR(StringToInt32Calculator);
|
||||
|
||||
using StringToUint32Calculator = StringToIntCalculatorTemplate<uint32>;
|
||||
REGISTER_CALCULATOR(StringToUint32Calculator);
|
||||
|
||||
using StringToInt64Calculator = StringToIntCalculatorTemplate<int64>;
|
||||
REGISTER_CALCULATOR(StringToInt64Calculator);
|
||||
|
||||
using StringToUint64Calculator = StringToIntCalculatorTemplate<uint64>;
|
||||
REGISTER_CALCULATOR(StringToUint64Calculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -12,15 +12,14 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
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",
|
||||
srcs = ["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,
|
||||
)
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -13,12 +13,12 @@
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "graph_tensors_packet_generator_proto",
|
||||
srcs = ["graph_tensors_packet_generator.proto"],
|
||||
@@ -104,6 +104,17 @@ proto_library(
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "unpack_media_sequence_calculator_proto",
|
||||
srcs = ["unpack_media_sequence_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/core:packet_resampler_calculator_proto",
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/util:audio_decoder_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "vector_float_to_tensor_calculator_options_proto",
|
||||
srcs = ["vector_float_to_tensor_calculator_options.proto"],
|
||||
@@ -127,7 +138,7 @@ mediapipe_cc_proto_library(
|
||||
srcs = ["image_frame_to_tensor_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":image_frame_to_tensor_calculator_proto"],
|
||||
@@ -162,7 +173,7 @@ mediapipe_cc_proto_library(
|
||||
srcs = ["pack_media_sequence_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":pack_media_sequence_calculator_proto"],
|
||||
@@ -181,7 +192,7 @@ mediapipe_cc_proto_library(
|
||||
srcs = ["tensorflow_session_from_frozen_graph_generator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:packet_generator_cc_proto",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":tensorflow_session_from_frozen_graph_generator_proto"],
|
||||
@@ -192,7 +203,7 @@ mediapipe_cc_proto_library(
|
||||
srcs = ["tensorflow_session_from_frozen_graph_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":tensorflow_session_from_frozen_graph_calculator_proto"],
|
||||
@@ -255,11 +266,23 @@ 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"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "vector_int_to_tensor_calculator_options_cc_proto",
|
||||
srcs = ["vector_int_to_tensor_calculator_options.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":vector_int_to_tensor_calculator_options_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "vector_float_to_tensor_calculator_options_cc_proto",
|
||||
srcs = ["vector_float_to_tensor_calculator_options.proto"],
|
||||
@@ -273,7 +296,7 @@ cc_library(
|
||||
srcs = ["graph_tensors_packet_generator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:graph_tensors_packet_generator_cc_proto",
|
||||
":graph_tensors_packet_generator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -288,7 +311,7 @@ cc_library(
|
||||
srcs = ["image_frame_to_tensor_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:image_frame_to_tensor_calculator_cc_proto",
|
||||
":image_frame_to_tensor_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
@@ -310,7 +333,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
||||
"//mediapipe/calculators/tensorflow:matrix_to_tensor_calculator_options_cc_proto",
|
||||
":matrix_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -331,7 +354,7 @@ cc_library(
|
||||
srcs = ["lapped_tensor_buffer_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:lapped_tensor_buffer_calculator_cc_proto",
|
||||
":lapped_tensor_buffer_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -385,7 +408,7 @@ cc_library(
|
||||
"//mediapipe/util/sequence:media_sequence",
|
||||
"//mediapipe/util/sequence:media_sequence_util",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -400,7 +423,7 @@ cc_library(
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -413,7 +436,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":tensorflow_session",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_inference_calculator_cc_proto",
|
||||
":tensorflow_inference_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/tool:status_util",
|
||||
"@com_google_absl//absl/strings",
|
||||
@@ -491,7 +514,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":tensorflow_session",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
":tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/tool:status_util",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -550,7 +573,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":tensorflow_session",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_saved_model_generator_cc_proto",
|
||||
":tensorflow_session_from_saved_model_generator_cc_proto",
|
||||
"//mediapipe/framework:packet_generator",
|
||||
"//mediapipe/framework:packet_type",
|
||||
"//mediapipe/framework/tool:status_util",
|
||||
@@ -574,7 +597,7 @@ cc_library(
|
||||
srcs = ["tensor_squeeze_dimensions_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:tensor_squeeze_dimensions_calculator_cc_proto",
|
||||
":tensor_squeeze_dimensions_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -588,7 +611,7 @@ cc_library(
|
||||
srcs = ["tensor_to_image_frame_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_image_frame_calculator_cc_proto",
|
||||
":tensor_to_image_frame_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
@@ -604,7 +627,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_matrix_calculator_cc_proto",
|
||||
":tensor_to_matrix_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -620,6 +643,22 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "tfrecord_reader_calculator",
|
||||
srcs = ["tfrecord_reader_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@org_tensorflow//tensorflow/core:lib",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "tensor_to_vector_float_calculator",
|
||||
srcs = ["tensor_to_vector_float_calculator.cc"],
|
||||
@@ -628,7 +667,7 @@ cc_library(
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_vector_float_calculator_options_cc_proto",
|
||||
":tensor_to_vector_float_calculator_options_cc_proto",
|
||||
] + select({
|
||||
"//conditions:default": [
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
@@ -653,9 +692,24 @@ 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",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "vector_int_to_tensor_calculator",
|
||||
srcs = ["vector_int_to_tensor_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":vector_int_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -665,7 +719,7 @@ cc_library(
|
||||
srcs = ["vector_float_to_tensor_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:vector_float_to_tensor_calculator_options_cc_proto",
|
||||
":vector_float_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -674,12 +728,26 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "unpack_yt8m_sequence_example_calculator",
|
||||
srcs = ["unpack_yt8m_sequence_example_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":lapped_tensor_buffer_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "graph_tensors_packet_generator_test",
|
||||
srcs = ["graph_tensors_packet_generator_test.cc"],
|
||||
deps = [
|
||||
":graph_tensors_packet_generator",
|
||||
"//mediapipe/calculators/tensorflow:graph_tensors_packet_generator_cc_proto",
|
||||
":graph_tensors_packet_generator_cc_proto",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework:packet_generator_cc_proto",
|
||||
"//mediapipe/framework:packet_set",
|
||||
@@ -711,7 +779,7 @@ cc_test(
|
||||
srcs = ["matrix_to_tensor_calculator_test.cc"],
|
||||
deps = [
|
||||
":matrix_to_tensor_calculator",
|
||||
"//mediapipe/calculators/tensorflow:matrix_to_tensor_calculator_options_cc_proto",
|
||||
":matrix_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
@@ -727,13 +795,13 @@ cc_test(
|
||||
srcs = ["lapped_tensor_buffer_calculator_test.cc"],
|
||||
deps = [
|
||||
":lapped_tensor_buffer_calculator",
|
||||
"//mediapipe/calculators/tensorflow:lapped_tensor_buffer_calculator_cc_proto",
|
||||
":lapped_tensor_buffer_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -769,11 +837,10 @@ 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",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -800,7 +867,7 @@ cc_test(
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:direct_session",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
"@org_tensorflow//tensorflow/core:testlib",
|
||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
||||
"@org_tensorflow//tensorflow/core/kernels:math",
|
||||
@@ -816,7 +883,7 @@ cc_test(
|
||||
":tensorflow_inference_calculator",
|
||||
":tensorflow_session",
|
||||
":tensorflow_session_from_frozen_graph_generator",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
":tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework:packet_generator_cc_proto",
|
||||
@@ -830,7 +897,7 @@ cc_test(
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:direct_session",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
"@org_tensorflow//tensorflow/core:testlib",
|
||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
||||
"@org_tensorflow//tensorflow/core/kernels:math",
|
||||
@@ -846,7 +913,7 @@ cc_test(
|
||||
":tensorflow_inference_calculator",
|
||||
":tensorflow_session",
|
||||
":tensorflow_session_from_saved_model_generator",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_saved_model_generator_cc_proto",
|
||||
":tensorflow_session_from_saved_model_generator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework:packet_generator_cc_proto",
|
||||
@@ -856,14 +923,8 @@ cc_test(
|
||||
"//mediapipe/framework/tool:tag_map_helper",
|
||||
"//mediapipe/framework/tool:validate_type",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:all_kernels",
|
||||
"@org_tensorflow//tensorflow/core:direct_session",
|
||||
"@org_tensorflow//tensorflow/core/kernels:array",
|
||||
"@org_tensorflow//tensorflow/core/kernels:bitcast_op",
|
||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
||||
"@org_tensorflow//tensorflow/core/kernels:io",
|
||||
"@org_tensorflow//tensorflow/core/kernels:state",
|
||||
"@org_tensorflow//tensorflow/core/kernels:string",
|
||||
"@org_tensorflow//tensorflow/core/kernels/data:tensor_dataset_op",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -887,14 +948,8 @@ cc_test(
|
||||
"//mediapipe/framework/tool:tag_map_helper",
|
||||
"//mediapipe/framework/tool:validate_type",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:all_kernels",
|
||||
"@org_tensorflow//tensorflow/core:direct_session",
|
||||
"@org_tensorflow//tensorflow/core/kernels:array",
|
||||
"@org_tensorflow//tensorflow/core/kernels:bitcast_op",
|
||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
||||
"@org_tensorflow//tensorflow/core/kernels:io",
|
||||
"@org_tensorflow//tensorflow/core/kernels:state",
|
||||
"@org_tensorflow//tensorflow/core/kernels:string",
|
||||
"@org_tensorflow//tensorflow/core/kernels/data:tensor_dataset_op",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -903,12 +958,12 @@ cc_test(
|
||||
srcs = ["tensor_squeeze_dimensions_calculator_test.cc"],
|
||||
deps = [
|
||||
":tensor_squeeze_dimensions_calculator",
|
||||
"//mediapipe/calculators/tensorflow:tensor_squeeze_dimensions_calculator_cc_proto",
|
||||
":tensor_squeeze_dimensions_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -918,13 +973,13 @@ cc_test(
|
||||
srcs = ["tensor_to_image_frame_calculator_test.cc"],
|
||||
deps = [
|
||||
":tensor_to_image_frame_calculator",
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_image_frame_calculator_cc_proto",
|
||||
":tensor_to_image_frame_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -934,14 +989,14 @@ cc_test(
|
||||
srcs = ["tensor_to_matrix_calculator_test.cc"],
|
||||
deps = [
|
||||
":tensor_to_matrix_calculator",
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_matrix_calculator_cc_proto",
|
||||
":tensor_to_matrix_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -950,12 +1005,12 @@ cc_test(
|
||||
srcs = ["tensor_to_vector_float_calculator_test.cc"],
|
||||
deps = [
|
||||
":tensor_to_vector_float_calculator",
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_vector_float_calculator_options_cc_proto",
|
||||
":tensor_to_vector_float_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -971,10 +1026,25 @@ 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",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "vector_int_to_tensor_calculator_test",
|
||||
srcs = ["vector_int_to_tensor_calculator_test.cc"],
|
||||
deps = [
|
||||
":vector_int_to_tensor_calculator",
|
||||
":vector_int_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -983,12 +1053,12 @@ cc_test(
|
||||
srcs = ["vector_float_to_tensor_calculator_test.cc"],
|
||||
deps = [
|
||||
":vector_float_to_tensor_calculator",
|
||||
"//mediapipe/calculators/tensorflow:vector_float_to_tensor_calculator_options_cc_proto",
|
||||
":vector_float_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -1012,7 +1082,7 @@ cc_test(
|
||||
":tensorflow_session",
|
||||
":tensorflow_inference_calculator",
|
||||
":tensorflow_session_from_frozen_graph_generator",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
":tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
"//mediapipe/framework/deps:file_path",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
|
||||
@@ -29,6 +29,11 @@
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
const char kBufferSize[] = "BUFFER_SIZE";
|
||||
const char kOverlap[] = "OVERLAP";
|
||||
const char kTimestampOffset[] = "TIMESTAMP_OFFSET";
|
||||
const char kCalculatorOptions[] = "CALCULATOR_OPTIONS";
|
||||
|
||||
namespace tf = tensorflow;
|
||||
|
||||
// Given an input stream of tensors, concatenates the tensors over timesteps.
|
||||
@@ -72,6 +77,9 @@ class LappedTensorBufferCalculator : public CalculatorBase {
|
||||
::mediapipe::Status AddBatchDimension(tf::Tensor* input_tensor);
|
||||
|
||||
int steps_until_output_;
|
||||
int buffer_size_;
|
||||
int overlap_;
|
||||
int timestamp_offset_;
|
||||
std::unique_ptr<CircularBuffer<Timestamp>> timestamp_buffer_;
|
||||
std::unique_ptr<CircularBuffer<tf::Tensor>> buffer_;
|
||||
LappedTensorBufferCalculatorOptions options_;
|
||||
@@ -87,6 +95,21 @@ REGISTER_CALCULATOR(LappedTensorBufferCalculator);
|
||||
);
|
||||
RET_CHECK_EQ(cc->Inputs().NumEntries(), 1)
|
||||
<< "Only one output stream is supported.";
|
||||
|
||||
if (cc->InputSidePackets().HasTag(kBufferSize)) {
|
||||
cc->InputSidePackets().Tag(kBufferSize).Set<int>();
|
||||
}
|
||||
if (cc->InputSidePackets().HasTag(kOverlap)) {
|
||||
cc->InputSidePackets().Tag(kOverlap).Set<int>();
|
||||
}
|
||||
if (cc->InputSidePackets().HasTag(kTimestampOffset)) {
|
||||
cc->InputSidePackets().Tag(kTimestampOffset).Set<int>();
|
||||
}
|
||||
if (cc->InputSidePackets().HasTag(kCalculatorOptions)) {
|
||||
cc->InputSidePackets()
|
||||
.Tag(kCalculatorOptions)
|
||||
.Set<LappedTensorBufferCalculatorOptions>();
|
||||
}
|
||||
cc->Outputs().Index(0).Set<tf::Tensor>(
|
||||
// Output tensorflow::Tensor stream with possibly overlapping steps.
|
||||
);
|
||||
@@ -95,16 +118,33 @@ REGISTER_CALCULATOR(LappedTensorBufferCalculator);
|
||||
|
||||
::mediapipe::Status LappedTensorBufferCalculator::Open(CalculatorContext* cc) {
|
||||
options_ = cc->Options<LappedTensorBufferCalculatorOptions>();
|
||||
RET_CHECK_LT(options_.overlap(), options_.buffer_size());
|
||||
RET_CHECK_GE(options_.timestamp_offset(), 0)
|
||||
if (cc->InputSidePackets().HasTag(kCalculatorOptions)) {
|
||||
options_ = cc->InputSidePackets()
|
||||
.Tag(kCalculatorOptions)
|
||||
.Get<LappedTensorBufferCalculatorOptions>();
|
||||
}
|
||||
buffer_size_ = options_.buffer_size();
|
||||
if (cc->InputSidePackets().HasTag(kBufferSize)) {
|
||||
buffer_size_ = cc->InputSidePackets().Tag(kBufferSize).Get<int>();
|
||||
}
|
||||
overlap_ = options_.overlap();
|
||||
if (cc->InputSidePackets().HasTag(kOverlap)) {
|
||||
overlap_ = cc->InputSidePackets().Tag(kOverlap).Get<int>();
|
||||
}
|
||||
timestamp_offset_ = options_.timestamp_offset();
|
||||
if (cc->InputSidePackets().HasTag(kTimestampOffset)) {
|
||||
timestamp_offset_ = cc->InputSidePackets().Tag(kTimestampOffset).Get<int>();
|
||||
}
|
||||
|
||||
RET_CHECK_LT(overlap_, buffer_size_);
|
||||
RET_CHECK_GE(timestamp_offset_, 0)
|
||||
<< "Negative timestamp_offset is not allowed.";
|
||||
RET_CHECK_LT(options_.timestamp_offset(), options_.buffer_size())
|
||||
RET_CHECK_LT(timestamp_offset_, buffer_size_)
|
||||
<< "output_frame_num_offset has to be less than buffer_size.";
|
||||
timestamp_buffer_ =
|
||||
absl::make_unique<CircularBuffer<Timestamp>>(options_.buffer_size());
|
||||
buffer_ =
|
||||
absl::make_unique<CircularBuffer<tf::Tensor>>(options_.buffer_size());
|
||||
steps_until_output_ = options_.buffer_size();
|
||||
absl::make_unique<CircularBuffer<Timestamp>>(buffer_size_);
|
||||
buffer_ = absl::make_unique<CircularBuffer<tf::Tensor>>(buffer_size_);
|
||||
steps_until_output_ = buffer_size_;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -128,11 +168,10 @@ REGISTER_CALCULATOR(LappedTensorBufferCalculator);
|
||||
concatenated.get());
|
||||
RET_CHECK(concat_status.ok()) << concat_status.ToString();
|
||||
|
||||
cc->Outputs().Index(0).Add(
|
||||
concatenated.release(),
|
||||
timestamp_buffer_->Get(options_.timestamp_offset()));
|
||||
cc->Outputs().Index(0).Add(concatenated.release(),
|
||||
timestamp_buffer_->Get(timestamp_offset_));
|
||||
|
||||
steps_until_output_ = options_.buffer_size() - options_.overlap();
|
||||
steps_until_output_ = buffer_size_ - overlap_;
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -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>();
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
// 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 <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "tensorflow/core/example/example.pb.h"
|
||||
#include "tensorflow/core/lib/core/status.h"
|
||||
#include "tensorflow/core/lib/io/record_reader.h"
|
||||
#include "tensorflow/core/platform/env.h"
|
||||
#include "tensorflow/core/platform/file_system.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
const char kTFRecordPath[] = "TFRECORD_PATH";
|
||||
const char kRecordIndex[] = "RECORD_INDEX";
|
||||
const char kExampleTag[] = "EXAMPLE";
|
||||
const char kSequenceExampleTag[] = "SEQUENCE_EXAMPLE";
|
||||
|
||||
// Reads a tensorflow example/sequence example from a tfrecord file.
|
||||
// If the "RECORD_INDEX" input side packet is provided, the calculator is going
|
||||
// to fetch the example/sequence example of the tfrecord file at the target
|
||||
// record index. Otherwise, the reader always reads the first example/sequence
|
||||
// example of the tfrecord file.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "TFRecordReaderCalculator"
|
||||
// input_side_packet: "TFRECORD_PATH:tfrecord_path"
|
||||
// input_side_packet: "RECORD_INDEX:record_index"
|
||||
// output_side_packet: "SEQUENCE_EXAMPLE:sequence_example"
|
||||
// }
|
||||
class TFRecordReaderCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
};
|
||||
|
||||
::mediapipe::Status TFRecordReaderCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
cc->InputSidePackets().Tag(kTFRecordPath).Set<std::string>();
|
||||
if (cc->InputSidePackets().HasTag(kRecordIndex)) {
|
||||
cc->InputSidePackets().Tag(kRecordIndex).Set<int>();
|
||||
}
|
||||
|
||||
RET_CHECK(cc->OutputSidePackets().HasTag(kExampleTag) ||
|
||||
cc->OutputSidePackets().HasTag(kSequenceExampleTag))
|
||||
<< "TFRecordReaderCalculator must output either Tensorflow example or "
|
||||
"sequence example.";
|
||||
if (cc->OutputSidePackets().HasTag(kExampleTag)) {
|
||||
cc->OutputSidePackets().Tag(kExampleTag).Set<tensorflow::Example>();
|
||||
} else {
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kSequenceExampleTag)
|
||||
.Set<tensorflow::SequenceExample>();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status TFRecordReaderCalculator::Open(CalculatorContext* cc) {
|
||||
std::unique_ptr<tensorflow::RandomAccessFile> file;
|
||||
auto tf_status = tensorflow::Env::Default()->NewRandomAccessFile(
|
||||
cc->InputSidePackets().Tag(kTFRecordPath).Get<std::string>(), &file);
|
||||
RET_CHECK(tf_status.ok())
|
||||
<< "Failed to open tfrecord file: " << tf_status.error_message();
|
||||
tensorflow::io::RecordReader reader(file.get(),
|
||||
tensorflow::io::RecordReaderOptions());
|
||||
tensorflow::uint64 offset = 0;
|
||||
std::string example_str;
|
||||
const int target_idx =
|
||||
cc->InputSidePackets().HasTag(kRecordIndex)
|
||||
? cc->InputSidePackets().Tag(kRecordIndex).Get<int>()
|
||||
: 0;
|
||||
int current_idx = 0;
|
||||
while (current_idx <= target_idx) {
|
||||
tf_status = reader.ReadRecord(&offset, &example_str);
|
||||
RET_CHECK(tf_status.ok())
|
||||
<< "Failed to read tfrecord: " << tf_status.error_message();
|
||||
if (current_idx == target_idx) {
|
||||
if (cc->OutputSidePackets().HasTag(kExampleTag)) {
|
||||
tensorflow::Example tf_example;
|
||||
tf_example.ParseFromString(example_str);
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kExampleTag)
|
||||
.Set(MakePacket<tensorflow::Example>(std::move(tf_example)));
|
||||
} else {
|
||||
tensorflow::SequenceExample tf_sequence_example;
|
||||
tf_sequence_example.ParseFromString(example_str);
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kSequenceExampleTag)
|
||||
.Set(MakePacket<tensorflow::SequenceExample>(
|
||||
std::move(tf_sequence_example)));
|
||||
}
|
||||
}
|
||||
++current_idx;
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status TFRecordReaderCalculator::Process(CalculatorContext* cc) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
REGISTER_CALCULATOR(TFRecordReaderCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -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(
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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".
|
||||
|
||||
@@ -0,0 +1,192 @@
|
||||
// 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 <iterator>
|
||||
|
||||
#include "mediapipe/calculators/tensorflow/lapped_tensor_buffer_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "tensorflow/core/example/example.pb.h"
|
||||
#include "tensorflow/core/example/feature.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
const char kId[] = "id";
|
||||
const char kRgb[] = "rgb";
|
||||
const char kAudio[] = "audio";
|
||||
const char kDesiredSegmentSize[] = "DESIRED_SEGMENT_SIZE";
|
||||
const char kYt8mId[] = "YT8M_ID";
|
||||
const char kYt8mSequenceExample[] = "YT8M_SEQUENCE_EXAMPLE";
|
||||
const char kQuantizedRgbFeature[] = "QUANTIZED_RGB_FEATURE";
|
||||
const char kQuantizedAudioFeature[] = "QUANTIZED_AUDIO_FEATURE";
|
||||
const char kSegmentSize[] = "SEGMENT_SIZE";
|
||||
const char kLappedTensorBufferCalculatorOptions[] =
|
||||
"LAPPED_TENSOR_BUFFER_CALCULATOR_OPTIONS";
|
||||
|
||||
std::string GetQuantizedFeature(
|
||||
const tensorflow::SequenceExample& sequence_example, const std::string& key,
|
||||
int index) {
|
||||
const auto& bytes_list = sequence_example.feature_lists()
|
||||
.feature_list()
|
||||
.at(key)
|
||||
.feature()
|
||||
.Get(index)
|
||||
.bytes_list()
|
||||
.value();
|
||||
CHECK_EQ(1, bytes_list.size());
|
||||
return bytes_list.Get(0);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
// Unpacks YT8M Sequence Example. Note that the audio feature and rgb feature
|
||||
// output are quantized. DequantizeByteArrayCalculator can do the dequantization
|
||||
// for you.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "UnpackYt8mSequenceExampleCalculator"
|
||||
// input_side_packet: "YT8M_SEQUENCE_EXAMPLE:yt8m_sequence_example"
|
||||
// output_stream: "QUANTIZED_RGB_FEATURE:quantized_rgb_feature"
|
||||
// output_stream: "QUANTIZED_AUDIO_FEATURE:quantized_audio_feature"
|
||||
// }
|
||||
class UnpackYt8mSequenceExampleCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->InputSidePackets()
|
||||
.Tag(kYt8mSequenceExample)
|
||||
.Set<tensorflow::SequenceExample>();
|
||||
if (cc->InputSidePackets().HasTag(kDesiredSegmentSize)) {
|
||||
cc->InputSidePackets().Tag(kDesiredSegmentSize).Set<int>();
|
||||
}
|
||||
cc->Outputs().Tag(kQuantizedRgbFeature).Set<std::string>();
|
||||
cc->Outputs().Tag(kQuantizedAudioFeature).Set<std::string>();
|
||||
if (cc->OutputSidePackets().HasTag(kYt8mId)) {
|
||||
cc->OutputSidePackets().Tag(kYt8mId).Set<std::string>();
|
||||
}
|
||||
if (cc->OutputSidePackets().HasTag(kLappedTensorBufferCalculatorOptions)) {
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kLappedTensorBufferCalculatorOptions)
|
||||
.Set<::mediapipe::LappedTensorBufferCalculatorOptions>();
|
||||
}
|
||||
if (cc->OutputSidePackets().HasTag(kSegmentSize)) {
|
||||
cc->OutputSidePackets().Tag(kSegmentSize).Set<int>();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
const tensorflow::SequenceExample& sequence_example =
|
||||
cc->InputSidePackets()
|
||||
.Tag(kYt8mSequenceExample)
|
||||
.Get<tensorflow::SequenceExample>();
|
||||
const std::string& yt8m_id =
|
||||
sequence_example.context().feature().at(kId).bytes_list().value().Get(
|
||||
0);
|
||||
if (cc->OutputSidePackets().HasTag(kYt8mId)) {
|
||||
cc->OutputSidePackets().Tag(kYt8mId).Set(
|
||||
MakePacket<std::string>(yt8m_id));
|
||||
}
|
||||
|
||||
int rgb_feature_list_length =
|
||||
sequence_example.feature_lists().feature_list().at(kRgb).feature_size();
|
||||
int audio_feature_list_length = sequence_example.feature_lists()
|
||||
.feature_list()
|
||||
.at(kAudio)
|
||||
.feature_size();
|
||||
|
||||
if (rgb_feature_list_length != audio_feature_list_length) {
|
||||
return ::mediapipe::FailedPreconditionError(absl::StrCat(
|
||||
"Data corruption: the length of audio features and rgb features are "
|
||||
"not equal. Please check the sequence example that contains yt8m "
|
||||
"id: ",
|
||||
yt8m_id));
|
||||
}
|
||||
feature_list_length_ = rgb_feature_list_length;
|
||||
if (cc->OutputSidePackets().HasTag(kLappedTensorBufferCalculatorOptions) ||
|
||||
cc->OutputSidePackets().HasTag(kSegmentSize)) {
|
||||
// If the desired segment size is specified, take the min of the length of
|
||||
// the feature list and the desired size to be the output segment size.
|
||||
int segment_size = feature_list_length_;
|
||||
if (cc->InputSidePackets().HasTag(kDesiredSegmentSize)) {
|
||||
int desired_segment_size =
|
||||
cc->InputSidePackets().Tag(kDesiredSegmentSize).Get<int>();
|
||||
RET_CHECK(desired_segment_size > 0)
|
||||
<< "The desired segment size must be greater than zero.";
|
||||
segment_size = std::min(
|
||||
feature_list_length_,
|
||||
cc->InputSidePackets().Tag(kDesiredSegmentSize).Get<int>());
|
||||
}
|
||||
if (cc->OutputSidePackets().HasTag(
|
||||
kLappedTensorBufferCalculatorOptions)) {
|
||||
auto lapped_tensor_buffer_calculator_options = absl::make_unique<
|
||||
::mediapipe::LappedTensorBufferCalculatorOptions>();
|
||||
lapped_tensor_buffer_calculator_options->set_add_batch_dim_to_tensors(
|
||||
true);
|
||||
lapped_tensor_buffer_calculator_options->set_buffer_size(segment_size);
|
||||
lapped_tensor_buffer_calculator_options->set_overlap(segment_size - 1);
|
||||
lapped_tensor_buffer_calculator_options->set_timestamp_offset(
|
||||
segment_size - 1);
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kLappedTensorBufferCalculatorOptions)
|
||||
.Set(Adopt(lapped_tensor_buffer_calculator_options.release()));
|
||||
}
|
||||
if (cc->OutputSidePackets().HasTag(kSegmentSize)) {
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kSegmentSize)
|
||||
.Set(MakePacket<int>(segment_size));
|
||||
}
|
||||
}
|
||||
LOG(INFO) << "Reading the sequence example that contains yt8m id: "
|
||||
<< yt8m_id << ". Feature list length: " << feature_list_length_;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
if (current_index_ >= feature_list_length_) {
|
||||
return ::mediapipe::tool::StatusStop();
|
||||
}
|
||||
const tensorflow::SequenceExample& sequence_example =
|
||||
cc->InputSidePackets()
|
||||
.Tag(kYt8mSequenceExample)
|
||||
.Get<tensorflow::SequenceExample>();
|
||||
|
||||
// Uses microsecond as the unit of time. In the YT8M dataset, each feature
|
||||
// represents a second.
|
||||
const Timestamp timestamp = Timestamp(current_index_ * 1000000);
|
||||
cc->Outputs()
|
||||
.Tag(kQuantizedRgbFeature)
|
||||
.AddPacket(
|
||||
MakePacket<std::string>(
|
||||
GetQuantizedFeature(sequence_example, kRgb, current_index_))
|
||||
.At(timestamp));
|
||||
cc->Outputs()
|
||||
.Tag(kQuantizedAudioFeature)
|
||||
.AddPacket(
|
||||
MakePacket<std::string>(
|
||||
GetQuantizedFeature(sequence_example, kAudio, current_index_))
|
||||
.At(timestamp));
|
||||
++current_index_;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
int current_index_ = 0;
|
||||
int feature_list_length_ = 0;
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(UnpackYt8mSequenceExampleCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -23,10 +23,12 @@
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace tf = ::tensorflow;
|
||||
|
||||
namespace {
|
||||
auto& INPUT_1D = VectorFloatToTensorCalculatorOptions::INPUT_1D;
|
||||
auto& INPUT_2D = VectorFloatToTensorCalculatorOptions::INPUT_2D;
|
||||
} // namespace
|
||||
|
||||
namespace tf = ::tensorflow;
|
||||
|
||||
// The calculator expects one input (a packet containing a vector<float> or
|
||||
// vector<vector<float>>) and generates one output (a packet containing a
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
// 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.
|
||||
//
|
||||
// Converts a single int or vector<int> or vector<vector<int>> to 1D (or 2D)
|
||||
// tf::Tensor.
|
||||
|
||||
#include "mediapipe/calculators/tensorflow/vector_int_to_tensor_calculator_options.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "tensorflow/core/framework/tensor.h"
|
||||
#include "tensorflow/core/framework/types.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
const char kVectorInt[] = "VECTOR_INT";
|
||||
const char kSingleInt[] = "SINGLE_INT";
|
||||
const char kTensorOut[] = "TENSOR_OUT";
|
||||
|
||||
namespace {
|
||||
auto& INPUT_1D = VectorIntToTensorCalculatorOptions::INPUT_1D;
|
||||
auto& INPUT_2D = VectorIntToTensorCalculatorOptions::INPUT_2D;
|
||||
} // namespace
|
||||
|
||||
namespace tf = ::tensorflow;
|
||||
|
||||
template <typename TensorType>
|
||||
void AssignMatrixValue(int r, int c, int value, tf::Tensor* output_tensor) {
|
||||
output_tensor->tensor<TensorType, 2>()(r, c) = value;
|
||||
}
|
||||
|
||||
// The calculator expects one input (a packet containing a single int or
|
||||
// vector<int> or vector<vector<int>>) and generates one output (a packet
|
||||
// containing a tf::Tensor containing the same data). The output tensor will be
|
||||
// either 1D or 2D with dimensions corresponding to the input vector int. It
|
||||
// will hold DT_INT32 or DT_UINT8 or DT_INT64 values.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "VectorIntToTensorCalculator"
|
||||
// input_stream: "SINGLE_INT:segment_size_int_stream"
|
||||
// output_stream: "TENSOR_OUT:segment_size_tensor"
|
||||
// }
|
||||
//
|
||||
// or
|
||||
//
|
||||
// node {
|
||||
// calculator: "VectorIntToTensorCalculator"
|
||||
// input_stream: "VECTOR_INT:vector_int_features"
|
||||
// output_stream: "TENSOR_OUT:tensor_features"
|
||||
// }
|
||||
class VectorIntToTensorCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
VectorIntToTensorCalculatorOptions options_;
|
||||
};
|
||||
REGISTER_CALCULATOR(VectorIntToTensorCalculator);
|
||||
|
||||
::mediapipe::Status VectorIntToTensorCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
const auto& options = cc->Options<VectorIntToTensorCalculatorOptions>();
|
||||
// Start with only one input packet.
|
||||
RET_CHECK_EQ(cc->Inputs().NumEntries(), 1)
|
||||
<< "Only one input stream is supported.";
|
||||
if (options.input_size() == INPUT_2D) {
|
||||
cc->Inputs().Tag(kVectorInt).Set<std::vector<std::vector<int>>>();
|
||||
} else if (options.input_size() == INPUT_1D) {
|
||||
if (cc->Inputs().HasTag(kSingleInt)) {
|
||||
cc->Inputs().Tag(kSingleInt).Set<int>();
|
||||
} else {
|
||||
cc->Inputs().Tag(kVectorInt).Set<std::vector<int>>();
|
||||
}
|
||||
} else {
|
||||
LOG(FATAL) << "input size not supported";
|
||||
}
|
||||
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1)
|
||||
<< "Only one output stream is supported.";
|
||||
cc->Outputs().Tag(kTensorOut).Set<tf::Tensor>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status VectorIntToTensorCalculator::Open(CalculatorContext* cc) {
|
||||
options_ = cc->Options<VectorIntToTensorCalculatorOptions>();
|
||||
RET_CHECK(options_.tensor_data_type() == tf::DT_UINT8 ||
|
||||
options_.tensor_data_type() == tf::DT_INT32 ||
|
||||
options_.tensor_data_type() == tf::DT_INT64)
|
||||
<< "Output tensor data type is not supported.";
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status VectorIntToTensorCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
tf::TensorShape tensor_shape;
|
||||
if (options_.input_size() == INPUT_2D) {
|
||||
const std::vector<std::vector<int>>& input =
|
||||
cc->Inputs()
|
||||
.Tag(kVectorInt)
|
||||
.Value()
|
||||
.Get<std::vector<std::vector<int>>>();
|
||||
|
||||
const int32 rows = input.size();
|
||||
CHECK_GE(rows, 1);
|
||||
const int32 cols = input[0].size();
|
||||
CHECK_GE(cols, 1);
|
||||
for (int i = 1; i < rows; ++i) {
|
||||
CHECK_EQ(input[i].size(), cols);
|
||||
}
|
||||
if (options_.transpose()) {
|
||||
tensor_shape = tf::TensorShape({cols, rows});
|
||||
} else {
|
||||
tensor_shape = tf::TensorShape({rows, cols});
|
||||
}
|
||||
auto output = ::absl::make_unique<tf::Tensor>(options_.tensor_data_type(),
|
||||
tensor_shape);
|
||||
if (options_.transpose()) {
|
||||
for (int r = 0; r < rows; ++r) {
|
||||
for (int c = 0; c < cols; ++c) {
|
||||
switch (options_.tensor_data_type()) {
|
||||
case tf::DT_INT64:
|
||||
AssignMatrixValue<tf::int64>(c, r, input[r][c], output.get());
|
||||
break;
|
||||
case tf::DT_UINT8:
|
||||
AssignMatrixValue<uint8>(c, r, input[r][c], output.get());
|
||||
break;
|
||||
case tf::DT_INT32:
|
||||
AssignMatrixValue<int>(c, r, input[r][c], output.get());
|
||||
break;
|
||||
default:
|
||||
LOG(FATAL) << "tensor data type is not supported.";
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for (int r = 0; r < rows; ++r) {
|
||||
for (int c = 0; c < cols; ++c) {
|
||||
switch (options_.tensor_data_type()) {
|
||||
case tf::DT_INT64:
|
||||
AssignMatrixValue<tf::int64>(r, c, input[r][c], output.get());
|
||||
break;
|
||||
case tf::DT_UINT8:
|
||||
AssignMatrixValue<uint8>(r, c, input[r][c], output.get());
|
||||
break;
|
||||
case tf::DT_INT32:
|
||||
AssignMatrixValue<int>(r, c, input[r][c], output.get());
|
||||
break;
|
||||
default:
|
||||
LOG(FATAL) << "tensor data type is not supported.";
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
cc->Outputs().Tag(kTensorOut).Add(output.release(), cc->InputTimestamp());
|
||||
} else if (options_.input_size() == INPUT_1D) {
|
||||
std::vector<int> input;
|
||||
if (cc->Inputs().HasTag(kSingleInt)) {
|
||||
input.push_back(cc->Inputs().Tag(kSingleInt).Get<int>());
|
||||
} else {
|
||||
input = cc->Inputs().Tag(kVectorInt).Value().Get<std::vector<int>>();
|
||||
}
|
||||
CHECK_GE(input.size(), 1);
|
||||
const int32 length = input.size();
|
||||
tensor_shape = tf::TensorShape({length});
|
||||
auto output = ::absl::make_unique<tf::Tensor>(options_.tensor_data_type(),
|
||||
tensor_shape);
|
||||
for (int i = 0; i < length; ++i) {
|
||||
switch (options_.tensor_data_type()) {
|
||||
case tf::DT_INT64:
|
||||
output->tensor<tf::int64, 1>()(i) = input.at(i);
|
||||
break;
|
||||
case tf::DT_UINT8:
|
||||
output->tensor<uint8, 1>()(i) = input.at(i);
|
||||
break;
|
||||
case tf::DT_INT32:
|
||||
output->tensor<int, 1>()(i) = input.at(i);
|
||||
break;
|
||||
default:
|
||||
LOG(FATAL) << "tensor data type is not supported.";
|
||||
}
|
||||
}
|
||||
cc->Outputs().Tag(kTensorOut).Add(output.release(), cc->InputTimestamp());
|
||||
} else {
|
||||
LOG(FATAL) << "input size not supported";
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,43 @@
|
||||
// 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";
|
||||
import "tensorflow/core/framework/types.proto";
|
||||
|
||||
message VectorIntToTensorCalculatorOptions {
|
||||
extend mediapipe.CalculatorOptions {
|
||||
optional VectorIntToTensorCalculatorOptions ext = 275364184;
|
||||
}
|
||||
enum InputSize {
|
||||
UNKNOWN = 0;
|
||||
INPUT_1D = 1;
|
||||
INPUT_2D = 2;
|
||||
}
|
||||
|
||||
// If input_size is INPUT_2D, unpack a vector<vector<int>> to a
|
||||
// 2d tensor (matrix). If INPUT_1D, convert a single int or vector<int>
|
||||
// into a 1d tensor (vector).
|
||||
optional InputSize input_size = 1 [default = INPUT_1D];
|
||||
|
||||
// If true, the output tensor is transposed.
|
||||
// Otherwise, the output tensor is not transposed.
|
||||
// It will be ignored if tensor_is_2d is INPUT_1D.
|
||||
optional bool transpose = 2 [default = false];
|
||||
|
||||
optional tensorflow.DataType tensor_data_type = 3 [default = DT_INT32];
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
// Copyright 2018 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "mediapipe/calculators/tensorflow/vector_int_to_tensor_calculator_options.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "tensorflow/core/framework/tensor.h"
|
||||
#include "tensorflow/core/framework/types.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
namespace tf = ::tensorflow;
|
||||
|
||||
class VectorIntToTensorCalculatorTest : public ::testing::Test {
|
||||
protected:
|
||||
void SetUpRunner(
|
||||
const VectorIntToTensorCalculatorOptions::InputSize input_size,
|
||||
const tensorflow::DataType tensor_data_type, const bool transpose,
|
||||
const bool single_value) {
|
||||
CalculatorGraphConfig::Node config;
|
||||
config.set_calculator("VectorIntToTensorCalculator");
|
||||
if (single_value) {
|
||||
config.add_input_stream("SINGLE_INT:input_int");
|
||||
} else {
|
||||
config.add_input_stream("VECTOR_INT:input_int");
|
||||
}
|
||||
config.add_output_stream("TENSOR_OUT:output_tensor");
|
||||
auto options = config.mutable_options()->MutableExtension(
|
||||
VectorIntToTensorCalculatorOptions::ext);
|
||||
options->set_input_size(input_size);
|
||||
options->set_transpose(transpose);
|
||||
options->set_tensor_data_type(tensor_data_type);
|
||||
runner_ = ::absl::make_unique<CalculatorRunner>(config);
|
||||
}
|
||||
|
||||
void TestConvertFromVectoVectorInt(const bool transpose) {
|
||||
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_2D,
|
||||
tensorflow::DT_INT32, transpose, false);
|
||||
auto input = ::absl::make_unique<std::vector<std::vector<int>>>(
|
||||
2, std::vector<int>(2));
|
||||
for (int i = 0; i < 2; ++i) {
|
||||
for (int j = 0; j < 2; ++j) {
|
||||
input->at(i).at(j) = i * 2 + j;
|
||||
}
|
||||
}
|
||||
|
||||
const int64 time = 1234;
|
||||
runner_->MutableInputs()
|
||||
->Tag("VECTOR_INT")
|
||||
.packets.push_back(Adopt(input.release()).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||
|
||||
EXPECT_EQ(2, output_tensor.dims());
|
||||
EXPECT_EQ(tf::DT_INT32, output_tensor.dtype());
|
||||
const auto matrix = output_tensor.matrix<int>();
|
||||
|
||||
for (int i = 0; i < 2; ++i) {
|
||||
for (int j = 0; j < 2; ++j) {
|
||||
if (!transpose) {
|
||||
EXPECT_EQ(i * 2 + j, matrix(i, j));
|
||||
} else {
|
||||
EXPECT_EQ(j * 2 + i, matrix(i, j));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::unique_ptr<CalculatorRunner> runner_;
|
||||
};
|
||||
|
||||
TEST_F(VectorIntToTensorCalculatorTest, TestSingleValue) {
|
||||
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||
tensorflow::DT_INT32, false, true);
|
||||
const int64 time = 1234;
|
||||
runner_->MutableInputs()
|
||||
->Tag("SINGLE_INT")
|
||||
.packets.push_back(MakePacket<int>(1).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||
|
||||
EXPECT_EQ(1, output_tensor.dims());
|
||||
EXPECT_EQ(tf::DT_INT32, output_tensor.dtype());
|
||||
const auto vec = output_tensor.vec<int32>();
|
||||
EXPECT_EQ(1, vec(0));
|
||||
}
|
||||
|
||||
TEST_F(VectorIntToTensorCalculatorTest, TesOneDim) {
|
||||
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||
tensorflow::DT_INT32, false, false);
|
||||
auto input = ::absl::make_unique<std::vector<int>>(5);
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
input->at(i) = i;
|
||||
}
|
||||
const int64 time = 1234;
|
||||
runner_->MutableInputs()
|
||||
->Tag("VECTOR_INT")
|
||||
.packets.push_back(Adopt(input.release()).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||
|
||||
EXPECT_EQ(1, output_tensor.dims());
|
||||
EXPECT_EQ(tf::DT_INT32, output_tensor.dtype());
|
||||
const auto vec = output_tensor.vec<int32>();
|
||||
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
EXPECT_EQ(i, vec(i));
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(VectorIntToTensorCalculatorTest, TestTwoDims) {
|
||||
for (bool transpose : {false, true}) {
|
||||
TestConvertFromVectoVectorInt(transpose);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(VectorIntToTensorCalculatorTest, TestInt64) {
|
||||
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||
tensorflow::DT_INT64, false, true);
|
||||
const int64 time = 1234;
|
||||
runner_->MutableInputs()
|
||||
->Tag("SINGLE_INT")
|
||||
.packets.push_back(MakePacket<int>(2 ^ 31).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||
|
||||
EXPECT_EQ(1, output_tensor.dims());
|
||||
EXPECT_EQ(tf::DT_INT64, output_tensor.dtype());
|
||||
const auto vec = output_tensor.vec<tf::int64>();
|
||||
EXPECT_EQ(2 ^ 31, vec(0));
|
||||
}
|
||||
|
||||
TEST_F(VectorIntToTensorCalculatorTest, TestUint8) {
|
||||
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||
tensorflow::DT_UINT8, false, false);
|
||||
auto input = ::absl::make_unique<std::vector<int>>(5);
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
input->at(i) = i;
|
||||
}
|
||||
const int64 time = 1234;
|
||||
runner_->MutableInputs()
|
||||
->Tag("VECTOR_INT")
|
||||
.packets.push_back(Adopt(input.release()).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||
|
||||
EXPECT_EQ(1, output_tensor.dims());
|
||||
EXPECT_EQ(tf::DT_UINT8, output_tensor.dtype());
|
||||
const auto vec = output_tensor.vec<uint8>();
|
||||
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
EXPECT_EQ(i, vec(i));
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -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,26 @@ 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",
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate_internal",
|
||||
],
|
||||
"//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 +272,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 +308,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 +322,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 +335,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 +361,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 +389,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,
|
||||
)
|
||||
|
||||
@@ -16,23 +16,24 @@
|
||||
#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(__EMSCRIPTEN__) && \
|
||||
!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 +41,13 @@
|
||||
#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(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
typedef id<MTLBuffer> GpuTensor;
|
||||
@@ -66,26 +69,28 @@ 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(__EMSCRIPTEN__) && \
|
||||
!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 +107,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 +149,8 @@ class TfLiteConverterCalculator : public CalculatorBase {
|
||||
|
||||
std::unique_ptr<tflite::Interpreter> interpreter_ = nullptr;
|
||||
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||
std::unique_ptr<GPUData> gpu_data_out_;
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
@@ -175,25 +181,34 @@ 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) && !defined(__EMSCRIPTEN__)
|
||||
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) && !defined(__EMSCRIPTEN__)
|
||||
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(__EMSCRIPTEN__) && \
|
||||
!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 +223,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) && !defined(__EMSCRIPTEN__)
|
||||
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 +236,8 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
cc->Outputs().HasTag("TENSORS_GPU"));
|
||||
// Cannot use quantization.
|
||||
use_quantized_tensors_ = false;
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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 +254,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 +270,8 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
}
|
||||
|
||||
::mediapipe::Status TfLiteConverterCalculator::Close(CalculatorContext* cc) {
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
gpu_helper_.RunInGlContext([this] { gpu_data_out_.reset(); });
|
||||
#endif
|
||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
@@ -372,7 +390,8 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteConverterCalculator::ProcessGPU(
|
||||
CalculatorContext* cc) {
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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 +400,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 +413,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 +451,61 @@ 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) && !defined(__EMSCRIPTEN__)
|
||||
// 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(__EMSCRIPTEN__) && \
|
||||
!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 +513,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 +552,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 +616,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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
#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,45 @@
|
||||
#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"
|
||||
#include "tensorflow/lite/delegates/gpu/metal_delegate_internal.h"
|
||||
#endif // iOS
|
||||
|
||||
#if defined(__ANDROID__)
|
||||
namespace {
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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(__EMSCRIPTEN__) && \
|
||||
!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) && !defined(__EMSCRIPTEN__)
|
||||
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 +150,8 @@ class TfLiteInferenceCalculator : public CalculatorBase {
|
||||
std::unique_ptr<tflite::FlatBufferModel> model_;
|
||||
TfLiteDelegate* delegate_ = nullptr;
|
||||
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||
std::unique_ptr<GPUData> gpu_data_in_;
|
||||
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
|
||||
@@ -142,6 +159,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 +179,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) && !defined(__EMSCRIPTEN__)
|
||||
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) && !defined(__EMSCRIPTEN__)
|
||||
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 +205,18 @@ 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(__EMSCRIPTEN__) && \
|
||||
!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,37 +230,49 @@ 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) && !defined(__EMSCRIPTEN__)
|
||||
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) && !defined(__EMSCRIPTEN__)
|
||||
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(__EMSCRIPTEN__) && \
|
||||
!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(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||
[this, &cc]() -> ::mediapipe::Status { return LoadDelegate(cc); }));
|
||||
#else
|
||||
MP_RETURN_IF_ERROR(LoadDelegate(cc));
|
||||
#endif
|
||||
}
|
||||
|
||||
#if defined(__EMSCRIPTEN__)
|
||||
MP_RETURN_IF_ERROR(LoadDelegate(cc));
|
||||
#endif // __EMSCRIPTEN__
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -237,35 +280,28 @@ 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(__EMSCRIPTEN__) && \
|
||||
!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 +314,21 @@ 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(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
MP_RETURN_IF_ERROR(
|
||||
gpu_helper_.RunInGlContext([this]() -> ::mediapipe::Status {
|
||||
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
|
||||
@@ -304,52 +343,52 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
// 3. Output processed tensors.
|
||||
if (gpu_output_) {
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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 +406,8 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteInferenceCalculator::Close(CalculatorContext* cc) {
|
||||
if (delegate_) {
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
|
||||
TfLiteGpuDelegateDelete(delegate_);
|
||||
gpu_data_in_.reset();
|
||||
@@ -431,6 +471,10 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
RET_CHECK(interpreter_);
|
||||
|
||||
#if defined(__EMSCRIPTEN__)
|
||||
interpreter_->SetNumThreads(1);
|
||||
#endif // __EMSCRIPTEN__
|
||||
|
||||
if (gpu_output_) {
|
||||
use_quantized_tensors_ = false;
|
||||
} else {
|
||||
@@ -446,7 +490,8 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteInferenceCalculator::LoadDelegate(
|
||||
CalculatorContext* cc) {
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
// Configure and create the delegate.
|
||||
TfLiteGpuDelegateOptions options = TfLiteGpuDelegateOptionsDefault();
|
||||
options.compile_options.precision_loss_allowed = 1;
|
||||
@@ -466,15 +511,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 +538,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 +549,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;
|
||||
TFLGpuDelegateOptions options;
|
||||
options.allow_precision_loss = false; // Must match converter, F=float/T=half
|
||||
options.wait_type = GpuDelegateOptions::WaitType::kActive;
|
||||
options.wait_type = TFLGpuDelegateWaitType::TFLGpuDelegateWaitTypePassive;
|
||||
if (!delegate_) delegate_ = TFLGpuDelegateCreate(&options);
|
||||
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
||||
|
||||
if (gpu_input_) {
|
||||
// Get input image sizes.
|
||||
@@ -539,11 +578,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 +598,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 +633,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
|
||||
|
||||
|
||||
@@ -24,7 +24,8 @@
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
||||
#if defined(__EMSCRIPTEN__) || defined(__ANDROID__) || \
|
||||
(defined(__APPLE__) && !TARGET_OS_OSX)
|
||||
#include "mediapipe/util/android/file/base/file.h"
|
||||
#include "mediapipe/util/android/file/base/helpers.h"
|
||||
#else
|
||||
@@ -66,8 +67,8 @@ class TfLiteTensorsToClassificationCalculator : public CalculatorBase {
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
::mediapipe::TfLiteTensorsToClassificationCalculatorOptions options_;
|
||||
int top_k_ = 0;
|
||||
double min_score_threshold_ = 0;
|
||||
std::unordered_map<int, std::string> label_map_;
|
||||
bool label_map_loaded_ = false;
|
||||
};
|
||||
@@ -93,15 +94,14 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
||||
CalculatorContext* cc) {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
auto options = cc->Options<
|
||||
options_ = cc->Options<
|
||||
::mediapipe::TfLiteTensorsToClassificationCalculatorOptions>();
|
||||
|
||||
top_k_ = options.top_k();
|
||||
min_score_threshold_ = options.min_score_threshold();
|
||||
if (options.has_label_map_path()) {
|
||||
top_k_ = options_.top_k();
|
||||
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(file::GetContents(string_path, &label_map_string));
|
||||
|
||||
@@ -125,9 +125,11 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
||||
RET_CHECK_EQ(input_tensors.size(), 1);
|
||||
|
||||
const TfLiteTensor* raw_score_tensor = &input_tensors[0];
|
||||
RET_CHECK_EQ(raw_score_tensor->dims->size, 2);
|
||||
RET_CHECK_EQ(raw_score_tensor->dims->data[0], 1);
|
||||
int num_classes = raw_score_tensor->dims->data[1];
|
||||
int num_classes = 1;
|
||||
for (int i = 0; i < raw_score_tensor->dims->size; ++i) {
|
||||
num_classes *= raw_score_tensor->dims->data[i];
|
||||
}
|
||||
|
||||
if (label_map_loaded_) {
|
||||
RET_CHECK_EQ(num_classes, label_map_.size());
|
||||
}
|
||||
@@ -135,7 +137,8 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
||||
|
||||
auto classification_list = absl::make_unique<ClassificationList>();
|
||||
for (int i = 0; i < num_classes; ++i) {
|
||||
if (raw_scores[i] < min_score_threshold_) {
|
||||
if (options_.has_min_score_threshold() &&
|
||||
raw_scores[i] < options_.min_score_threshold()) {
|
||||
continue;
|
||||
}
|
||||
Classification* classification = classification_list->add_classification();
|
||||
@@ -148,6 +151,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
||||
|
||||
// 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_) {
|
||||
std::partial_sort(raw_classification_list->begin(),
|
||||
|
||||
@@ -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,64 @@
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
||||
using ::tflite::gpu::gl::GlShader;
|
||||
#endif
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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) && !defined(__EMSCRIPTEN__)
|
||||
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 +125,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 +163,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 +183,13 @@ class TfLiteTensorsToDetectionsCalculator : public CalculatorBase {
|
||||
std::vector<Anchor> anchors_;
|
||||
bool side_packet_anchors_{};
|
||||
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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 +202,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) && !defined(__EMSCRIPTEN__)
|
||||
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 +225,14 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(__ANDROID__)
|
||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||
if (use_gpu) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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 +243,12 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
gpu_input_ = true;
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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 +256,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 +275,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 +292,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 +406,85 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
}
|
||||
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::ProcessGPU(
|
||||
CalculatorContext* cc, std::vector<Detection>* output_detections) {
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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 +492,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 +560,21 @@ 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(__EMSCRIPTEN__) && \
|
||||
!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 +664,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -586,12 +721,17 @@ 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(__EMSCRIPTEN__) && \
|
||||
!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 +805,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 +819,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 +903,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 +1132,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();
|
||||
}
|
||||
|
||||
|
||||
@@ -21,7 +21,8 @@
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator for converting TFLite tensors from regression models into
|
||||
// landmarks.
|
||||
// landmarks. Note that if the landmarks in the tensor has more than 3
|
||||
// dimensions, only the first 3 dimensions will be converted to x,y,z.
|
||||
//
|
||||
// Input:
|
||||
// TENSORS - Vector of TfLiteTensor of type kTfLiteFloat32. Only the first
|
||||
@@ -96,7 +97,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 "
|
||||
@@ -121,9 +123,6 @@ REGISTER_CALCULATOR(TfLiteTensorsToLandmarksCalculator);
|
||||
num_values *= raw_tensor->dims->data[i];
|
||||
}
|
||||
const int num_dimensions = num_values / num_landmarks_;
|
||||
// Landmarks must have less than 3 dimensions. Otherwise please consider
|
||||
// using matrix.
|
||||
CHECK_LE(num_dimensions, 3);
|
||||
CHECK_GT(num_dimensions, 0);
|
||||
|
||||
const float* raw_landmarks = raw_tensor->data.f;
|
||||
@@ -133,7 +132,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() -
|
||||
|
||||
@@ -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];
|
||||
}
|
||||
|
||||
@@ -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,8 @@
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||
#include "mediapipe/gpu/gl_simple_shaders.h"
|
||||
#include "mediapipe/gpu/shader_util.h"
|
||||
@@ -36,7 +38,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 +54,15 @@ float Clamp(float val, float min, float max) {
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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 +131,14 @@ class TfLiteTensorsToSegmentationCalculator : public CalculatorBase {
|
||||
int tensor_channels_ = 0;
|
||||
|
||||
bool use_gpu_ = false;
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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 +148,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 +162,40 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
}
|
||||
|
||||
// Inputs GPU.
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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(__EMSCRIPTEN__) && \
|
||||
!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(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -189,24 +205,25 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
use_gpu_ = true;
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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(__EMSCRIPTEN__) && \
|
||||
!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 +232,14 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
if (use_gpu_) {
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!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 +249,8 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::Close(
|
||||
CalculatorContext* cc) {
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
gpu_helper_.RunInGlContext([this] {
|
||||
if (upsample_program_) glDeleteProgram(upsample_program_);
|
||||
upsample_program_ = 0;
|
||||
@@ -239,7 +258,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
mask_program_no_prev_.reset();
|
||||
tensor_buffer_.reset();
|
||||
});
|
||||
#endif // __ANDROID__
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -358,7 +377,8 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
if (cc->Inputs().Tag("TENSORS_GPU").IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
#if defined(__ANDROID__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
// Get input streams.
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GlBuffer>>();
|
||||
@@ -379,9 +399,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 +409,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 +417,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 +458,14 @@ 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(__EMSCRIPTEN__) && \
|
||||
!defined(__APPLE__)
|
||||
static const GLfloat square_vertices[] = {
|
||||
-1.0f, -1.0f, // bottom left
|
||||
1.0f, -1.0f, // bottom right
|
||||
@@ -492,7 +513,7 @@ void TfLiteTensorsToSegmentationCalculator::GlRender() {
|
||||
glBindVertexArray(0);
|
||||
glDeleteVertexArrays(1, &vao);
|
||||
glDeleteBuffers(2, vbo);
|
||||
#endif // __ANDROID__
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
}
|
||||
|
||||
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::LoadOptions(
|
||||
@@ -516,14 +537,16 @@ 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(__EMSCRIPTEN__) && \
|
||||
!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 +612,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 +690,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();
|
||||
}
|
||||
|
||||
@@ -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_
|
||||
@@ -12,14 +12,14 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
package(default_visibility = ["//visibility:public"])
|
||||
|
||||
exports_files(["LICENSE"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "annotation_overlay_calculator_proto",
|
||||
srcs = ["annotation_overlay_calculator.proto"],
|
||||
@@ -72,6 +72,24 @@ proto_library(
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "collection_has_min_size_calculator_proto",
|
||||
srcs = ["collection_has_min_size_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "association_calculator_proto",
|
||||
srcs = ["association_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "annotation_overlay_calculator_cc_proto",
|
||||
srcs = ["annotation_overlay_calculator.proto"],
|
||||
@@ -141,6 +159,26 @@ mediapipe_cc_proto_library(
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "collection_has_min_size_calculator_cc_proto",
|
||||
srcs = ["collection_has_min_size_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
],
|
||||
visibility = ["//mediapipe:__subpackages__"],
|
||||
deps = [":collection_has_min_size_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "association_calculator_cc_proto",
|
||||
srcs = ["association_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
],
|
||||
visibility = ["//mediapipe:__subpackages__"],
|
||||
deps = [":association_calculator_proto"],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "packet_frequency_calculator",
|
||||
srcs = ["packet_frequency_calculator.cc"],
|
||||
@@ -234,20 +272,15 @@ cc_library(
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:vector",
|
||||
"//mediapipe/util:annotation_renderer",
|
||||
"//mediapipe/util:render_data_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: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,
|
||||
)
|
||||
@@ -366,6 +399,16 @@ mediapipe_cc_proto_library(
|
||||
deps = [":landmark_projection_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "landmarks_to_floats_calculator_cc_proto",
|
||||
srcs = ["landmarks_to_floats_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":landmarks_to_floats_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "rect_transformation_calculator_cc_proto",
|
||||
srcs = ["rect_transformation_calculator.proto"],
|
||||
@@ -378,7 +421,12 @@ mediapipe_cc_proto_library(
|
||||
|
||||
cc_library(
|
||||
name = "detections_to_rects_calculator",
|
||||
srcs = ["detections_to_rects_calculator.cc"],
|
||||
srcs = [
|
||||
"detections_to_rects_calculator.cc",
|
||||
],
|
||||
hdrs = [
|
||||
"detections_to_rects_calculator.h",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":detections_to_rects_calculator_cc_proto",
|
||||
@@ -460,6 +508,17 @@ proto_library(
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "labels_to_render_data_calculator_proto",
|
||||
srcs = ["labels_to_render_data_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/util:color_proto",
|
||||
"//mediapipe/util:render_data_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "thresholding_calculator_proto",
|
||||
srcs = ["thresholding_calculator.proto"],
|
||||
@@ -489,6 +548,15 @@ proto_library(
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "landmarks_to_floats_calculator_proto",
|
||||
srcs = ["landmarks_to_floats_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "rect_transformation_calculator_proto",
|
||||
srcs = ["rect_transformation_calculator.proto"],
|
||||
@@ -583,6 +651,26 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "labels_to_render_data_calculator",
|
||||
srcs = ["labels_to_render_data_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":labels_to_render_data_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_options_cc_proto",
|
||||
"//mediapipe/framework/formats:classification_cc_proto",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:statusor",
|
||||
"//mediapipe/util:color_cc_proto",
|
||||
"//mediapipe/util:render_data_cc_proto",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "rect_to_render_data_calculator",
|
||||
srcs = ["rect_to_render_data_calculator.cc"],
|
||||
@@ -664,6 +752,22 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "landmarks_to_floats_calculator",
|
||||
srcs = ["landmarks_to_floats_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":landmarks_to_floats_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@eigen_archive//:eigen",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "detection_letterbox_removal_calculator_test",
|
||||
srcs = ["detection_letterbox_removal_calculator_test.cc"],
|
||||
@@ -694,3 +798,188 @@ 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/formats:classification_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_cc_proto_library(
|
||||
name = "labels_to_render_data_calculator_cc_proto",
|
||||
srcs = ["labels_to_render_data_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/util:color_cc_proto",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":labels_to_render_data_calculator_proto"],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "local_file_contents_calculator",
|
||||
srcs = ["local_file_contents_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:file_helpers",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "filter_collection_calculator",
|
||||
srcs = ["filter_collection_calculator.cc"],
|
||||
hdrs = ["filter_collection_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "collection_has_min_size_calculator",
|
||||
srcs = ["collection_has_min_size_calculator.cc"],
|
||||
hdrs = ["collection_has_min_size_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":collection_has_min_size_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "association_calculator",
|
||||
hdrs = ["association_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":association_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework/port:rectangle",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/memory",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "association_norm_rect_calculator",
|
||||
srcs = ["association_norm_rect_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":association_calculator",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:rectangle",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "association_detection_calculator",
|
||||
srcs = ["association_detection_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":association_calculator",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/formats:location",
|
||||
"//mediapipe/framework/port:rectangle",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "association_calculator_test",
|
||||
srcs = ["association_calculator_test.cc"],
|
||||
deps = [
|
||||
":association_detection_calculator",
|
||||
":association_norm_rect_calculator",
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/deps:message_matchers",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/formats:location_data_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -26,13 +26,14 @@
|
||||
#include "mediapipe/framework/port/vector.h"
|
||||
#include "mediapipe/util/annotation_renderer.h"
|
||||
#include "mediapipe/util/color.pb.h"
|
||||
#include "mediapipe/util/render_data.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 {
|
||||
|
||||
@@ -41,6 +42,8 @@ namespace {
|
||||
constexpr char kInputFrameTag[] = "INPUT_FRAME";
|
||||
constexpr char kOutputFrameTag[] = "OUTPUT_FRAME";
|
||||
|
||||
constexpr char kInputVectorTag[] = "VECTOR";
|
||||
|
||||
constexpr char kInputFrameTagGpu[] = "INPUT_FRAME_GPU";
|
||||
constexpr char kOutputFrameTagGpu[] = "OUTPUT_FRAME_GPU";
|
||||
|
||||
@@ -65,6 +68,9 @@ constexpr int kAnnotationBackgroundColor[] = {100, 101, 102};
|
||||
// 2. RenderData proto on variable number of input streams. All the RenderData
|
||||
// at a particular timestamp is drawn on the image in the order of their
|
||||
// input streams. No tags required.
|
||||
// 3. std::vector<RenderData> on variable number of input streams. RenderData
|
||||
// objects at a particular timestamp are drawn on the image in order of the
|
||||
// input vector items. These input streams are tagged with "VECTOR".
|
||||
//
|
||||
// Output:
|
||||
// 1. OUTPUT_FRAME or OUTPUT_FRAME_GPU: A rendered ImageFrame (or GpuBuffer).
|
||||
@@ -85,6 +91,8 @@ constexpr int kAnnotationBackgroundColor[] = {100, 101, 102};
|
||||
// input_stream: "render_data_1"
|
||||
// input_stream: "render_data_2"
|
||||
// input_stream: "render_data_3"
|
||||
// input_stream: "VECTOR:0:render_data_vec_0"
|
||||
// input_stream: "VECTOR:1:render_data_vec_1"
|
||||
// output_stream: "OUTPUT_FRAME:decorated_frames"
|
||||
// options {
|
||||
// [mediapipe.AnnotationOverlayCalculatorOptions.ext] {
|
||||
@@ -99,6 +107,8 @@ constexpr int kAnnotationBackgroundColor[] = {100, 101, 102};
|
||||
// input_stream: "render_data_1"
|
||||
// input_stream: "render_data_2"
|
||||
// input_stream: "render_data_3"
|
||||
// input_stream: "VECTOR:0:render_data_vec_0"
|
||||
// input_stream: "VECTOR:1:render_data_vec_1"
|
||||
// output_stream: "OUTPUT_FRAME_GPU:decorated_frames"
|
||||
// options {
|
||||
// [mediapipe.AnnotationOverlayCalculatorOptions.ext] {
|
||||
@@ -138,21 +148,18 @@ class AnnotationOverlayCalculator : public CalculatorBase {
|
||||
// Underlying helper renderer library.
|
||||
std::unique_ptr<AnnotationRenderer> renderer_;
|
||||
|
||||
// Number of input streams with render data.
|
||||
int num_render_streams_;
|
||||
|
||||
// Indicates if image frame is available as input.
|
||||
bool image_frame_available_ = false;
|
||||
|
||||
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 +167,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.");
|
||||
@@ -169,39 +178,46 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
|
||||
return ::mediapipe::InternalError("GPU output must have GPU input.");
|
||||
}
|
||||
|
||||
// Assume all inputs are render streams; adjust below.
|
||||
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;
|
||||
}
|
||||
|
||||
// Data streams to render.
|
||||
for (int i = 0; i < num_render_streams; ++i) {
|
||||
cc->Inputs().Index(i).Set<RenderData>();
|
||||
for (CollectionItemId id = cc->Inputs().BeginId(); id < cc->Inputs().EndId();
|
||||
++id) {
|
||||
auto tag_and_index = cc->Inputs().TagAndIndexFromId(id);
|
||||
std::string tag = tag_and_index.first;
|
||||
if (tag == kInputVectorTag) {
|
||||
cc->Inputs().Get(id).Set<std::vector<RenderData>>();
|
||||
} else if (tag.empty()) {
|
||||
// Empty tag defaults to accepting a single object of RenderData type.
|
||||
cc->Inputs().Get(id).Set<RenderData>();
|
||||
}
|
||||
}
|
||||
|
||||
// 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,22 +228,20 @@ 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) ||
|
||||
cc->Inputs().HasTag(kInputFrameTag)) {
|
||||
image_frame_available_ = true;
|
||||
num_render_streams_ = cc->Inputs().NumEntries() - 1;
|
||||
} else {
|
||||
image_frame_available_ = false;
|
||||
RET_CHECK(options_.has_canvas_width_px());
|
||||
RET_CHECK(options_.has_canvas_height_px());
|
||||
num_render_streams_ = cc->Inputs().NumEntries();
|
||||
}
|
||||
|
||||
// Initialize the helper renderer library.
|
||||
@@ -246,9 +260,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 +274,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 +283,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));
|
||||
@@ -279,16 +293,32 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
|
||||
renderer_->AdoptImage(image_mat.get());
|
||||
|
||||
// Render streams onto render target.
|
||||
for (int i = 0; i < num_render_streams_; ++i) {
|
||||
if (cc->Inputs().Index(i).IsEmpty()) {
|
||||
for (CollectionItemId id = cc->Inputs().BeginId(); id < cc->Inputs().EndId();
|
||||
++id) {
|
||||
auto tag_and_index = cc->Inputs().TagAndIndexFromId(id);
|
||||
std::string tag = tag_and_index.first;
|
||||
if (!tag.empty() && tag != kInputVectorTag) {
|
||||
continue;
|
||||
}
|
||||
const RenderData& render_data = cc->Inputs().Index(i).Get<RenderData>();
|
||||
renderer_->RenderDataOnImage(render_data);
|
||||
if (cc->Inputs().Get(id).IsEmpty()) {
|
||||
continue;
|
||||
}
|
||||
if (tag.empty()) {
|
||||
// Empty tag defaults to accepting a single object of RenderData type.
|
||||
const RenderData& render_data = cc->Inputs().Get(id).Get<RenderData>();
|
||||
renderer_->RenderDataOnImage(render_data);
|
||||
} else {
|
||||
RET_CHECK_EQ(kInputVectorTag, tag);
|
||||
const std::vector<RenderData>& render_data_vec =
|
||||
cc->Inputs().Get(id).Get<std::vector<RenderData>>();
|
||||
for (const RenderData& render_data : render_data_vec) {
|
||||
renderer_->RenderDataOnImage(render_data);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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 +326,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 +337,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 +355,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 +363,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 +374,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 +420,7 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
|
||||
// Cleanup
|
||||
input_texture.Release();
|
||||
output_texture.Release();
|
||||
#endif // __ANDROID__ or iOS
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -451,15 +481,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 +502,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 +557,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 +640,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();
|
||||
}
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
// 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_UTIL_ASSOCIATION_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_UTIL_ASSOCIATION_CALCULATOR_H_
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "mediapipe/calculators/util/association_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/rectangle.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Computes the overlap similarity based on Intersection over Union (IoU) of
|
||||
// two rectangles.
|
||||
inline float OverlapSimilarity(const Rectangle_f& rect1,
|
||||
const Rectangle_f& rect2) {
|
||||
if (!rect1.Intersects(rect2)) return 0.0f;
|
||||
// Compute IoU similarity score.
|
||||
const float intersection_area = Rectangle_f(rect1).Intersect(rect2).Area();
|
||||
const float normalization = rect1.Area() + rect2.Area() - intersection_area;
|
||||
return normalization > 0.0f ? intersection_area / normalization : 0.0f;
|
||||
}
|
||||
|
||||
// AssocationCalculator<T> accepts multiple inputs of vectors of type T that can
|
||||
// be converted to Rectangle_f. The output is a vector of type T that contains
|
||||
// elements from the input vectors that don't overlap with each other. When
|
||||
// two elements overlap, the element that comes in from a later input stream
|
||||
// is kept in the output. This association operation is useful for multiple
|
||||
// instance inference pipelines in MediaPipe.
|
||||
// If an input stream is tagged with "PREV" tag, IDs of overlapping elements
|
||||
// from "PREV" input stream are propagated to the output. Elements in the "PREV"
|
||||
// input stream that don't overlap with other elements are not added to the
|
||||
// output. This stream is designed to take detections from previous timestamp,
|
||||
// e.g. output of PreviousLoopbackCalculator to provide temporal association.
|
||||
// See AssociationDetectionCalculator and AssociationNormRectCalculator for
|
||||
// example uses.
|
||||
template <typename T>
|
||||
class AssociationCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
// Atmost one input stream can be tagged with "PREV".
|
||||
RET_CHECK_LE(cc->Inputs().NumEntries("PREV"), 1);
|
||||
|
||||
if (cc->Inputs().HasTag("PREV")) {
|
||||
RET_CHECK_GE(cc->Inputs().NumEntries(), 2);
|
||||
}
|
||||
|
||||
for (CollectionItemId id = cc->Inputs().BeginId();
|
||||
id < cc->Inputs().EndId(); ++id) {
|
||||
cc->Inputs().Get(id).Set<std::vector<T>>();
|
||||
}
|
||||
|
||||
cc->Outputs().Index(0).Set<std::vector<T>>();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
has_prev_input_stream_ = cc->Inputs().HasTag("PREV");
|
||||
if (has_prev_input_stream_) {
|
||||
prev_input_stream_id_ = cc->Inputs().GetId("PREV", 0);
|
||||
}
|
||||
options_ = cc->Options<::mediapipe::AssociationCalculatorOptions>();
|
||||
CHECK_GE(options_.min_similarity_threshold(), 0);
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
auto get_non_overlapping_elements = GetNonOverlappingElements(cc);
|
||||
if (!get_non_overlapping_elements.ok()) {
|
||||
return get_non_overlapping_elements.status();
|
||||
}
|
||||
std::list<T> result = get_non_overlapping_elements.ValueOrDie();
|
||||
|
||||
if (has_prev_input_stream_ &&
|
||||
!cc->Inputs().Get(prev_input_stream_id_).IsEmpty()) {
|
||||
// Processed all regular input streams. Now compare the result list
|
||||
// elements with those in the PREV input stream, and propagate IDs from
|
||||
// PREV input stream as appropriate.
|
||||
const std::vector<T>& prev_input_vec =
|
||||
cc->Inputs()
|
||||
.Get(prev_input_stream_id_)
|
||||
.template Get<std::vector<T>>();
|
||||
|
||||
MP_RETURN_IF_ERROR(
|
||||
PropagateIdsFromPreviousToCurrent(prev_input_vec, &result));
|
||||
}
|
||||
|
||||
auto output = absl::make_unique<std::vector<T>>();
|
||||
for (auto it = result.begin(); it != result.end(); ++it) {
|
||||
output->push_back(*it);
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
protected:
|
||||
::mediapipe::AssociationCalculatorOptions options_;
|
||||
|
||||
bool has_prev_input_stream_;
|
||||
CollectionItemId prev_input_stream_id_;
|
||||
|
||||
virtual ::mediapipe::StatusOr<Rectangle_f> GetRectangle(const T& input) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
virtual std::pair<bool, int> GetId(const T& input) { return {false, -1}; }
|
||||
|
||||
virtual void SetId(T* input, int id) {}
|
||||
|
||||
private:
|
||||
// Get a list of non-overlapping elements from all input streams, with
|
||||
// increasing order of priority based on input stream index.
|
||||
mediapipe::StatusOr<std::list<T>> GetNonOverlappingElements(
|
||||
CalculatorContext* cc) {
|
||||
std::list<T> result;
|
||||
|
||||
// Initialize result with the first non-empty input vector.
|
||||
CollectionItemId non_empty_id = cc->Inputs().BeginId();
|
||||
for (CollectionItemId id = cc->Inputs().BeginId();
|
||||
id < cc->Inputs().EndId(); ++id) {
|
||||
if (id == prev_input_stream_id_ || cc->Inputs().Get(id).IsEmpty()) {
|
||||
continue;
|
||||
}
|
||||
const std::vector<T>& input_vec =
|
||||
cc->Inputs().Get(id).Get<std::vector<T>>();
|
||||
if (!input_vec.empty()) {
|
||||
non_empty_id = id;
|
||||
result.push_back(input_vec[0]);
|
||||
for (int j = 1; j < input_vec.size(); ++j) {
|
||||
MP_RETURN_IF_ERROR(AddElementToList(input_vec[j], &result));
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Compare remaining input vectors with the non-empty result vector,
|
||||
// remove lower-priority overlapping elements from the result vector and
|
||||
// had corresponding higher-priority elements as necessary.
|
||||
for (CollectionItemId id = non_empty_id + 1; id < cc->Inputs().EndId();
|
||||
++id) {
|
||||
if (id == prev_input_stream_id_ || cc->Inputs().Get(id).IsEmpty()) {
|
||||
continue;
|
||||
}
|
||||
const std::vector<T>& input_vec =
|
||||
cc->Inputs().Get(id).Get<std::vector<T>>();
|
||||
|
||||
for (int vi = 0; vi < input_vec.size(); ++vi) {
|
||||
MP_RETURN_IF_ERROR(AddElementToList(input_vec[vi], &result));
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
::mediapipe::Status AddElementToList(T element, std::list<T>* current) {
|
||||
// Compare this element with elements of the input collection. If this
|
||||
// element has high overlap with elements of the collection, remove
|
||||
// those elements from the collection and add this element.
|
||||
ASSIGN_OR_RETURN(auto cur_rect, GetRectangle(element));
|
||||
|
||||
bool change_id = false;
|
||||
int new_elem_id = -1;
|
||||
|
||||
for (auto uit = current->begin(); uit != current->end();) {
|
||||
ASSIGN_OR_RETURN(auto prev_rect, GetRectangle(*uit));
|
||||
if (OverlapSimilarity(cur_rect, prev_rect) >
|
||||
options_.min_similarity_threshold()) {
|
||||
std::pair<bool, int> prev_id = GetId(*uit);
|
||||
// If prev_id.first is false when some element doesn't have an ID,
|
||||
// change_id and new_elem_id will not be updated.
|
||||
if (prev_id.first) {
|
||||
change_id = prev_id.first;
|
||||
new_elem_id = prev_id.second;
|
||||
}
|
||||
uit = current->erase(uit);
|
||||
} else {
|
||||
++uit;
|
||||
}
|
||||
}
|
||||
|
||||
if (change_id) {
|
||||
SetId(&element, new_elem_id);
|
||||
}
|
||||
current->push_back(element);
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
// Compare elements of the current list with elements in from the collection
|
||||
// of elements from the previous input stream, and propagate IDs from the
|
||||
// previous input stream as appropriate.
|
||||
::mediapipe::Status PropagateIdsFromPreviousToCurrent(
|
||||
const std::vector<T>& prev_input_vec, std::list<T>* current) {
|
||||
for (auto vit = current->begin(); vit != current->end(); ++vit) {
|
||||
auto get_cur_rectangle = GetRectangle(*vit);
|
||||
if (!get_cur_rectangle.ok()) {
|
||||
return get_cur_rectangle.status();
|
||||
}
|
||||
const Rectangle_f& cur_rect = get_cur_rectangle.ValueOrDie();
|
||||
|
||||
bool change_id = false;
|
||||
int id_for_vi = -1;
|
||||
|
||||
for (int ui = 0; ui < prev_input_vec.size(); ++ui) {
|
||||
auto get_prev_rectangle = GetRectangle(prev_input_vec[ui]);
|
||||
if (!get_prev_rectangle.ok()) {
|
||||
return get_prev_rectangle.status();
|
||||
}
|
||||
const Rectangle_f& prev_rect = get_prev_rectangle.ValueOrDie();
|
||||
|
||||
if (OverlapSimilarity(cur_rect, prev_rect) >
|
||||
options_.min_similarity_threshold()) {
|
||||
std::pair<bool, int> prev_id = GetId(prev_input_vec[ui]);
|
||||
// If prev_id.first is false when some element doesn't have an ID,
|
||||
// change_id and id_for_vi will not be updated.
|
||||
if (prev_id.first) {
|
||||
change_id = prev_id.first;
|
||||
id_for_vi = prev_id.second;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (change_id) {
|
||||
T element = *vit;
|
||||
SetId(&element, id_for_vi);
|
||||
*vit = element;
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_UTIL_ASSOCIATION_CALCULATOR_H_
|
||||
@@ -0,0 +1,27 @@
|
||||
// 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 AssociationCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional AssociationCalculatorOptions ext = 275124847;
|
||||
}
|
||||
|
||||
optional float min_similarity_threshold = 1 [default = 1.0];
|
||||
}
|
||||
@@ -0,0 +1,476 @@
|
||||
// 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 "mediapipe/framework/calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/deps/message_matchers.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/location_data.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/packet.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_matchers.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
::mediapipe::Detection DetectionWithRelativeLocationData(double xmin,
|
||||
double ymin,
|
||||
double width,
|
||||
double height) {
|
||||
::mediapipe::Detection detection;
|
||||
::mediapipe::LocationData* location_data = detection.mutable_location_data();
|
||||
location_data->set_format(::mediapipe::LocationData::RELATIVE_BOUNDING_BOX);
|
||||
location_data->mutable_relative_bounding_box()->set_xmin(xmin);
|
||||
location_data->mutable_relative_bounding_box()->set_ymin(ymin);
|
||||
location_data->mutable_relative_bounding_box()->set_width(width);
|
||||
location_data->mutable_relative_bounding_box()->set_height(height);
|
||||
return detection;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
class AssociationDetectionCalculatorTest : public ::testing::Test {
|
||||
protected:
|
||||
AssociationDetectionCalculatorTest() {
|
||||
// 0.4 ================
|
||||
// | | | |
|
||||
// 0.3 ===================== | DET2 | |
|
||||
// | | | DET1 | | | DET4 |
|
||||
// 0.2 | DET0 | =========== ================
|
||||
// | | | | | |
|
||||
// 0.1 =====|=============== |
|
||||
// | DET3 | | |
|
||||
// 0.0 ================ |
|
||||
// | DET5 |
|
||||
// -0.1 ===========
|
||||
// 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 1.1 1.2
|
||||
|
||||
// Detection det_0.
|
||||
det_0 = DetectionWithRelativeLocationData(/*xmin=*/0.1, /*ymin=*/0.1,
|
||||
/*width=*/0.2, /*height=*/0.2);
|
||||
det_0.set_detection_id(0);
|
||||
|
||||
// Detection det_1.
|
||||
det_1 = DetectionWithRelativeLocationData(/*xmin=*/0.3, /*ymin=*/0.1,
|
||||
/*width=*/0.2, /*height=*/0.2);
|
||||
det_1.set_detection_id(1);
|
||||
|
||||
// Detection det_2.
|
||||
det_2 = DetectionWithRelativeLocationData(/*xmin=*/0.9, /*ymin=*/0.2,
|
||||
/*width=*/0.2, /*height=*/0.2);
|
||||
det_2.set_detection_id(2);
|
||||
|
||||
// Detection det_3.
|
||||
det_3 = DetectionWithRelativeLocationData(/*xmin=*/0.2, /*ymin=*/0.0,
|
||||
/*width=*/0.3, /*height=*/0.3);
|
||||
det_3.set_detection_id(3);
|
||||
|
||||
// Detection det_4.
|
||||
det_4 = DetectionWithRelativeLocationData(/*xmin=*/1.0, /*ymin=*/0.2,
|
||||
/*width=*/0.2, /*height=*/0.2);
|
||||
det_4.set_detection_id(4);
|
||||
|
||||
// Detection det_5.
|
||||
det_5 = DetectionWithRelativeLocationData(/*xmin=*/0.3, /*ymin=*/-0.1,
|
||||
/*width=*/0.3, /*height=*/0.3);
|
||||
det_5.set_detection_id(5);
|
||||
}
|
||||
|
||||
::mediapipe::Detection det_0, det_1, det_2, det_3, det_4, det_5;
|
||||
};
|
||||
|
||||
TEST_F(AssociationDetectionCalculatorTest, DetectionAssocTest) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationDetectionCalculator"
|
||||
input_stream: "input_vec_0"
|
||||
input_stream: "input_vec_1"
|
||||
input_stream: "input_vec_2"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream 0: det_0, det_1, det_2.
|
||||
auto input_vec_0 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_0->push_back(det_0);
|
||||
input_vec_0->push_back(det_1);
|
||||
input_vec_0->push_back(det_2);
|
||||
runner.MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(input_vec_0.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 1: det_3, det_4.
|
||||
auto input_vec_1 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_1->push_back(det_3);
|
||||
input_vec_1->push_back(det_4);
|
||||
runner.MutableInputs()->Index(1).packets.push_back(
|
||||
Adopt(input_vec_1.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 2: det_5.
|
||||
auto input_vec_2 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_2->push_back(det_5);
|
||||
runner.MutableInputs()->Index(2).packets.push_back(
|
||||
Adopt(input_vec_2.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::Detection>>();
|
||||
|
||||
// det_3 overlaps with det_0, det_1 and det_5 overlaps with det_3. Since det_5
|
||||
// is in the highest priority, we remove other rects. det_4 overlaps with
|
||||
// det_2, and det_4 is higher priority, so we keep it. The final output
|
||||
// therefore contains 2 elements.
|
||||
EXPECT_EQ(2, assoc_rects.size());
|
||||
// Outputs are in order of inputs, so det_4 is before det_5 in output vector.
|
||||
|
||||
// det_4 overlaps with det_2, so new id for det_4 is 2.
|
||||
EXPECT_TRUE(assoc_rects[0].has_detection_id());
|
||||
EXPECT_EQ(2, assoc_rects[0].detection_id());
|
||||
det_4.set_detection_id(2);
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(det_4));
|
||||
|
||||
// det_3 overlaps with det_0, so new id for det_3 is 0.
|
||||
// det_3 overlaps with det_1, so new id for det_3 is 1.
|
||||
// det_5 overlaps with det_3, so new id for det_5 is 1.
|
||||
EXPECT_TRUE(assoc_rects[1].has_detection_id());
|
||||
EXPECT_EQ(1, assoc_rects[1].detection_id());
|
||||
det_5.set_detection_id(1);
|
||||
EXPECT_THAT(assoc_rects[1], EqualsProto(det_5));
|
||||
}
|
||||
|
||||
TEST_F(AssociationDetectionCalculatorTest, DetectionAssocTestWithPrev) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationDetectionCalculator"
|
||||
input_stream: "PREV:input_vec_0"
|
||||
input_stream: "input_vec_1"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream 0: det_3, det_4.
|
||||
auto input_vec_0 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_0->push_back(det_3);
|
||||
input_vec_0->push_back(det_4);
|
||||
CollectionItemId prev_input_stream_id =
|
||||
runner.MutableInputs()->GetId("PREV", 0);
|
||||
runner.MutableInputs()
|
||||
->Get(prev_input_stream_id)
|
||||
.packets.push_back(Adopt(input_vec_0.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 1: det_5.
|
||||
auto input_vec_1 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_1->push_back(det_5);
|
||||
CollectionItemId input_stream_id = runner.MutableInputs()->GetId("", 0);
|
||||
runner.MutableInputs()
|
||||
->Get(input_stream_id)
|
||||
.packets.push_back(Adopt(input_vec_1.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::Detection>>();
|
||||
|
||||
// det_5 overlaps with det_3 and doesn't overlap with det_4. Since det_4 is
|
||||
// in the PREV input stream, it doesn't get copied to the output, so the final
|
||||
// output contains 1 element.
|
||||
EXPECT_EQ(1, assoc_rects.size());
|
||||
|
||||
// det_5 overlaps with det_3, det_3 is in PREV, so new id for det_5 is 3.
|
||||
EXPECT_TRUE(assoc_rects[0].has_detection_id());
|
||||
EXPECT_EQ(3, assoc_rects[0].detection_id());
|
||||
det_5.set_detection_id(3);
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(det_5));
|
||||
}
|
||||
|
||||
TEST_F(AssociationDetectionCalculatorTest, DetectionAssocTestReverse) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationDetectionCalculator"
|
||||
input_stream: "input_vec_0"
|
||||
input_stream: "input_vec_1"
|
||||
input_stream: "input_vec_2"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream 0: det_5.
|
||||
auto input_vec_0 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_0->push_back(det_5);
|
||||
runner.MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(input_vec_0.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 1: det_3, det_4.
|
||||
auto input_vec_1 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_1->push_back(det_3);
|
||||
input_vec_1->push_back(det_4);
|
||||
runner.MutableInputs()->Index(1).packets.push_back(
|
||||
Adopt(input_vec_1.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 2: det_0, det_1, det_2.
|
||||
auto input_vec_2 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_2->push_back(det_0);
|
||||
input_vec_2->push_back(det_1);
|
||||
input_vec_2->push_back(det_2);
|
||||
runner.MutableInputs()->Index(2).packets.push_back(
|
||||
Adopt(input_vec_2.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::Detection>>();
|
||||
|
||||
// det_3 overlaps with det_5, so det_5 is removed. det_0 overlaps with det_3,
|
||||
// so det_3 is removed as det_0 is in higher priority for keeping. det_2
|
||||
// overlaps with det_4 so det_4 is removed as det_2 is higher priority for
|
||||
// keeping. The final output therefore contains 3 elements.
|
||||
EXPECT_EQ(3, assoc_rects.size());
|
||||
// Outputs are in same order as inputs.
|
||||
|
||||
// det_3 overlaps with det_5, so new id for det_3 is 5.
|
||||
// det_0 overlaps with det_3, so new id for det_0 is 5.
|
||||
EXPECT_TRUE(assoc_rects[0].has_detection_id());
|
||||
EXPECT_EQ(5, assoc_rects[0].detection_id());
|
||||
det_0.set_detection_id(5);
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(det_0));
|
||||
|
||||
// det_1 stays with id 1.
|
||||
EXPECT_TRUE(assoc_rects[1].has_detection_id());
|
||||
EXPECT_EQ(1, assoc_rects[1].detection_id());
|
||||
EXPECT_THAT(assoc_rects[1], EqualsProto(det_1));
|
||||
|
||||
// det_2 overlaps with det_4, so new id for det_2 is 4.
|
||||
EXPECT_TRUE(assoc_rects[2].has_detection_id());
|
||||
EXPECT_EQ(4, assoc_rects[2].detection_id());
|
||||
det_2.set_detection_id(4);
|
||||
EXPECT_THAT(assoc_rects[2], EqualsProto(det_2));
|
||||
}
|
||||
|
||||
class AssociationNormRectCalculatorTest : public ::testing::Test {
|
||||
protected:
|
||||
AssociationNormRectCalculatorTest() {
|
||||
// 0.4 ================
|
||||
// | | | |
|
||||
// 0.3 ===================== | NR2 | |
|
||||
// | | | NR1 | | | NR4 |
|
||||
// 0.2 | NR0 | =========== ================
|
||||
// | | | | | |
|
||||
// 0.1 =====|=============== |
|
||||
// | NR3 | | |
|
||||
// 0.0 ================ |
|
||||
// | NR5 |
|
||||
// -0.1 ===========
|
||||
// 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 1.1 1.2
|
||||
|
||||
// NormalizedRect nr_0.
|
||||
nr_0.set_x_center(0.2);
|
||||
nr_0.set_y_center(0.2);
|
||||
nr_0.set_width(0.2);
|
||||
nr_0.set_height(0.2);
|
||||
|
||||
// NormalizedRect nr_1.
|
||||
nr_1.set_x_center(0.4);
|
||||
nr_1.set_y_center(0.2);
|
||||
nr_1.set_width(0.2);
|
||||
nr_1.set_height(0.2);
|
||||
|
||||
// NormalizedRect nr_2.
|
||||
nr_2.set_x_center(1.0);
|
||||
nr_2.set_y_center(0.3);
|
||||
nr_2.set_width(0.2);
|
||||
nr_2.set_height(0.2);
|
||||
|
||||
// NormalizedRect nr_3.
|
||||
nr_3.set_x_center(0.35);
|
||||
nr_3.set_y_center(0.15);
|
||||
nr_3.set_width(0.3);
|
||||
nr_3.set_height(0.3);
|
||||
|
||||
// NormalizedRect nr_4.
|
||||
nr_4.set_x_center(1.1);
|
||||
nr_4.set_y_center(0.3);
|
||||
nr_4.set_width(0.2);
|
||||
nr_4.set_height(0.2);
|
||||
|
||||
// NormalizedRect nr_5.
|
||||
nr_5.set_x_center(0.45);
|
||||
nr_5.set_y_center(0.05);
|
||||
nr_5.set_width(0.3);
|
||||
nr_5.set_height(0.3);
|
||||
}
|
||||
|
||||
::mediapipe::NormalizedRect nr_0, nr_1, nr_2, nr_3, nr_4, nr_5;
|
||||
};
|
||||
|
||||
TEST_F(AssociationNormRectCalculatorTest, NormRectAssocTest) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationNormRectCalculator"
|
||||
input_stream: "input_vec_0"
|
||||
input_stream: "input_vec_1"
|
||||
input_stream: "input_vec_2"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream 0: nr_0, nr_1, nr_2.
|
||||
auto input_vec_0 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_0->push_back(nr_0);
|
||||
input_vec_0->push_back(nr_1);
|
||||
input_vec_0->push_back(nr_2);
|
||||
runner.MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(input_vec_0.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 1: nr_3, nr_4.
|
||||
auto input_vec_1 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_1->push_back(nr_3);
|
||||
input_vec_1->push_back(nr_4);
|
||||
runner.MutableInputs()->Index(1).packets.push_back(
|
||||
Adopt(input_vec_1.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 2: nr_5.
|
||||
auto input_vec_2 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_2->push_back(nr_5);
|
||||
runner.MutableInputs()->Index(2).packets.push_back(
|
||||
Adopt(input_vec_2.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::NormalizedRect>>();
|
||||
|
||||
// nr_3 overlaps with nr_0, nr_1 and nr_5 overlaps with nr_3. Since nr_5 is
|
||||
// in the highest priority, we remove other rects.
|
||||
// nr_4 overlaps with nr_2, and nr_4 is higher priority, so we keep it.
|
||||
// The final output therefore contains 2 elements.
|
||||
EXPECT_EQ(2, assoc_rects.size());
|
||||
// Outputs are in order of inputs, so nr_4 is before nr_5 in output vector.
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(nr_4));
|
||||
EXPECT_THAT(assoc_rects[1], EqualsProto(nr_5));
|
||||
}
|
||||
|
||||
TEST_F(AssociationNormRectCalculatorTest, NormRectAssocTestReverse) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationNormRectCalculator"
|
||||
input_stream: "input_vec_0"
|
||||
input_stream: "input_vec_1"
|
||||
input_stream: "input_vec_2"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream 0: nr_5.
|
||||
auto input_vec_0 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_0->push_back(nr_5);
|
||||
runner.MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(input_vec_0.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 1: nr_3, nr_4.
|
||||
auto input_vec_1 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_1->push_back(nr_3);
|
||||
input_vec_1->push_back(nr_4);
|
||||
runner.MutableInputs()->Index(1).packets.push_back(
|
||||
Adopt(input_vec_1.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 2: nr_0, nr_1, nr_2.
|
||||
auto input_vec_2 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_2->push_back(nr_0);
|
||||
input_vec_2->push_back(nr_1);
|
||||
input_vec_2->push_back(nr_2);
|
||||
runner.MutableInputs()->Index(2).packets.push_back(
|
||||
Adopt(input_vec_2.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::NormalizedRect>>();
|
||||
|
||||
// nr_3 overlaps with nr_5, so nr_5 is removed. nr_0 overlaps with nr_3, so
|
||||
// nr_3 is removed as nr_0 is in higher priority for keeping. nr_2 overlaps
|
||||
// with nr_4 so nr_4 is removed as nr_2 is higher priority for keeping.
|
||||
// The final output therefore contains 3 elements.
|
||||
EXPECT_EQ(3, assoc_rects.size());
|
||||
// Outputs are in same order as inputs.
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(nr_0));
|
||||
EXPECT_THAT(assoc_rects[1], EqualsProto(nr_1));
|
||||
EXPECT_THAT(assoc_rects[2], EqualsProto(nr_2));
|
||||
}
|
||||
|
||||
TEST_F(AssociationNormRectCalculatorTest, NormRectAssocSingleInputStream) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationNormRectCalculator"
|
||||
input_stream: "input_vec"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream : nr_3, nr_5.
|
||||
auto input_vec =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec->push_back(nr_3);
|
||||
input_vec->push_back(nr_5);
|
||||
runner.MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(input_vec.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::NormalizedRect>>();
|
||||
|
||||
// nr_5 overlaps with nr_3. Since nr_5 is after nr_3 in the same input stream
|
||||
// we remove nr_3 and keep nr_5.
|
||||
// The final output therefore contains 1 elements.
|
||||
EXPECT_EQ(1, assoc_rects.size());
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(nr_5));
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,77 @@
|
||||
// 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 "mediapipe/calculators/util/association_calculator.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/location.h"
|
||||
#include "mediapipe/framework/port/rectangle.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A subclass of AssociationCalculator<T> for Detection. Example:
|
||||
// node {
|
||||
// calculator: "AssociationDetectionCalculator"
|
||||
// input_stream: "PREV:input_vec_0"
|
||||
// input_stream: "input_vec_1"
|
||||
// input_stream: "input_vec_2"
|
||||
// output_stream: "output_vec"
|
||||
// options {
|
||||
// [mediapipe.AssociationCalculatorOptions.ext] {
|
||||
// min_similarity_threshold: 0.1
|
||||
// }
|
||||
// }
|
||||
class AssociationDetectionCalculator
|
||||
: public AssociationCalculator<::mediapipe::Detection> {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
return AssociationCalculator<::mediapipe::Detection>::GetContract(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::Detection>::Open(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::Detection>::Process(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::Detection>::Close(cc);
|
||||
}
|
||||
|
||||
protected:
|
||||
::mediapipe::StatusOr<Rectangle_f> GetRectangle(
|
||||
const ::mediapipe::Detection& input) override {
|
||||
if (!input.has_location_data()) {
|
||||
return ::mediapipe::InternalError("Missing location_data in Detection");
|
||||
}
|
||||
const Location location(input.location_data());
|
||||
return location.GetRelativeBBox();
|
||||
}
|
||||
|
||||
std::pair<bool, int> GetId(const ::mediapipe::Detection& input) override {
|
||||
return {input.has_detection_id(), input.detection_id()};
|
||||
}
|
||||
|
||||
void SetId(::mediapipe::Detection* input, int id) override {
|
||||
input->set_detection_id(id);
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(AssociationDetectionCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,72 @@
|
||||
// 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 "mediapipe/calculators/util/association_calculator.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/port/rectangle.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A subclass of AssociationCalculator<T> for NormalizedRect. Example use case:
|
||||
// node {
|
||||
// calculator: "AssociationNormRectCalculator"
|
||||
// input_stream: "input_vec_0"
|
||||
// input_stream: "input_vec_1"
|
||||
// input_stream: "input_vec_2"
|
||||
// output_stream: "output_vec"
|
||||
// options {
|
||||
// [mediapipe.AssociationCalculatorOptions.ext] {
|
||||
// min_similarity_threshold: 0.1
|
||||
// }
|
||||
// }
|
||||
class AssociationNormRectCalculator
|
||||
: public AssociationCalculator<::mediapipe::NormalizedRect> {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
return AssociationCalculator<::mediapipe::NormalizedRect>::GetContract(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::NormalizedRect>::Open(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::NormalizedRect>::Process(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::NormalizedRect>::Close(cc);
|
||||
}
|
||||
|
||||
protected:
|
||||
::mediapipe::StatusOr<Rectangle_f> GetRectangle(
|
||||
const ::mediapipe::NormalizedRect& input) override {
|
||||
if (!input.has_x_center() || !input.has_y_center() || !input.has_width() ||
|
||||
!input.has_height()) {
|
||||
return ::mediapipe::InternalError(
|
||||
"Missing dimensions in NormalizedRect.");
|
||||
}
|
||||
const float xmin = input.x_center() - input.width() / 2.0;
|
||||
const float ymin = input.y_center() - input.height() / 2.0;
|
||||
// TODO: Support rotation for rectangle.
|
||||
return Rectangle_f(xmin, ymin, input.width(), input.height());
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(AssociationNormRectCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,26 @@
|
||||
|
||||
// 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 "mediapipe/calculators/util/collection_has_min_size_calculator.h"
|
||||
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef CollectionHasMinSizeCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||
NormalizedRectVectorHasMinSizeCalculator;
|
||||
REGISTER_CALCULATOR(NormalizedRectVectorHasMinSizeCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,84 @@
|
||||
// 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_UTIL_COLLECTION_HAS_MIN_SIZE_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_UTIL_COLLECTION_HAS_MIN_SIZE_CALCULATOR_H_
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/calculators/util/collection_has_min_size_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Deterimines if an input iterable collection has a minimum size, specified
|
||||
// in CollectionHasMinSizeCalculatorOptions. Example usage:
|
||||
// node {
|
||||
// calculator: "IntVectorHasMinSizeCalculator"
|
||||
// input_stream: "ITERABLE:input_int_vector"
|
||||
// output_stream: "has_min_ints"
|
||||
// options {
|
||||
// [mediapipe.CollectionHasMinSizeCalculatorOptions.ext] {
|
||||
// min_size: 2
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
template <typename IterableT>
|
||||
class CollectionHasMinSizeCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag("ITERABLE"));
|
||||
RET_CHECK_EQ(1, cc->Inputs().NumEntries());
|
||||
|
||||
RET_CHECK_EQ(1, cc->Outputs().NumEntries());
|
||||
|
||||
RET_CHECK_GE(
|
||||
cc->Options<::mediapipe::CollectionHasMinSizeCalculatorOptions>()
|
||||
.min_size(),
|
||||
0);
|
||||
|
||||
cc->Inputs().Tag("ITERABLE").Set<IterableT>();
|
||||
cc->Outputs().Index(0).Set<bool>();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
min_size_ =
|
||||
cc->Options<::mediapipe::CollectionHasMinSizeCalculatorOptions>()
|
||||
.min_size();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
const IterableT& input = cc->Inputs().Tag("ITERABLE").Get<IterableT>();
|
||||
bool has_min_size = input.size() >= min_size_;
|
||||
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
MakePacket<bool>(has_min_size).At(cc->InputTimestamp()));
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
int min_size_ = 0;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_UTIL_COLLECTION_HAS_MIN_SIZE_CALCULATOR_H_
|
||||
@@ -0,0 +1,29 @@
|
||||
// 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 CollectionHasMinSizeCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional CollectionHasMinSizeCalculatorOptions ext = 259397840;
|
||||
}
|
||||
|
||||
// The minimum size an input iterable collection should have for the
|
||||
// calculator to output true.
|
||||
optional int32 min_size = 1 [default = 0];
|
||||
}
|
||||
@@ -19,8 +19,8 @@
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
|
||||
#if defined(MEDIAPIPE_LITE) || defined(__ANDROID__) || \
|
||||
(defined(__APPLE__) && !TARGET_OS_OSX)
|
||||
#if defined(MEDIAPIPE_LITE) || defined(__EMSCRIPTEN__) || \
|
||||
defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
||||
#include "mediapipe/util/android/file/base/file.h"
|
||||
#include "mediapipe/util/android/file/base/helpers.h"
|
||||
#else
|
||||
|
||||
@@ -11,6 +11,8 @@
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
#include "mediapipe/calculators/util/detections_to_rects_calculator.h"
|
||||
|
||||
#include <cmath>
|
||||
|
||||
#include "mediapipe/calculators/util/detections_to_rects_calculator.pb.h"
|
||||
@@ -24,8 +26,6 @@
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
using mediapipe::DetectionsToRectsCalculatorOptions;
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr char kDetectionTag[] = "DETECTION";
|
||||
@@ -36,7 +36,10 @@ constexpr char kNormRectTag[] = "NORM_RECT";
|
||||
constexpr char kRectsTag[] = "RECTS";
|
||||
constexpr char kNormRectsTag[] = "NORM_RECTS";
|
||||
|
||||
::mediapipe::Status DetectionToRect(const Detection& detection, Rect* rect) {
|
||||
} // namespace
|
||||
|
||||
::mediapipe::Status DetectionsToRectsCalculator::DetectionToRect(
|
||||
const Detection& detection, Rect* rect) {
|
||||
const LocationData location_data = detection.location_data();
|
||||
RET_CHECK(location_data.format() == LocationData::BOUNDING_BOX)
|
||||
<< "Only Detection with formats of BOUNDING_BOX can be converted to Rect";
|
||||
@@ -48,8 +51,8 @@ constexpr char kNormRectsTag[] = "NORM_RECTS";
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status DetectionToNormalizedRect(const Detection& detection,
|
||||
NormalizedRect* rect) {
|
||||
::mediapipe::Status DetectionsToRectsCalculator::DetectionToNormalizedRect(
|
||||
const Detection& detection, NormalizedRect* rect) {
|
||||
const LocationData location_data = detection.location_data();
|
||||
RET_CHECK(location_data.format() == LocationData::RELATIVE_BOUNDING_BOX)
|
||||
<< "Only Detection with formats of RELATIVE_BOUNDING_BOX can be "
|
||||
@@ -63,79 +66,6 @@ constexpr char kNormRectsTag[] = "NORM_RECTS";
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
// Wraps around an angle in radians to within -M_PI and M_PI.
|
||||
inline float NormalizeRadians(float angle) {
|
||||
return angle - 2 * M_PI * std::floor((angle - (-M_PI)) / (2 * M_PI));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// A calculator that converts Detection proto to Rect proto.
|
||||
//
|
||||
// Detection is the format for encoding one or more detections in an image.
|
||||
// The input can be a single Detection or std::vector<Detection>. The output can
|
||||
// be either a single Rect or NormalizedRect, or std::vector<Rect> or
|
||||
// std::vector<NormalizedRect>. If Rect is used, the LocationData format is
|
||||
// expected to be BOUNDING_BOX, and if NormalizedRect is used it is expected to
|
||||
// be RELATIVE_BOUNDING_BOX.
|
||||
//
|
||||
// When the input is std::vector<Detection> and the output is a Rect or
|
||||
// NormalizedRect, only the first detection is converted. When the input is a
|
||||
// single Detection and the output is a std::vector<Rect> or
|
||||
// std::vector<NormalizedRect>, the output is a vector of size 1.
|
||||
//
|
||||
// Inputs:
|
||||
//
|
||||
// One of the following:
|
||||
// DETECTION: A Detection proto.
|
||||
// DETECTIONS: An std::vector<Detection>.
|
||||
//
|
||||
// IMAGE_SIZE (optional): A std::pair<int, int> represention image width and
|
||||
// height. This is required only when rotation needs to be computed (see
|
||||
// calculator options).
|
||||
//
|
||||
// Output:
|
||||
// One of the following:
|
||||
// RECT: A Rect proto.
|
||||
// NORM_RECT: A NormalizedRect proto.
|
||||
// RECTS: An std::vector<Rect>.
|
||||
// NORM_RECTS: An std::vector<NormalizedRect>.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "DetectionsToRectsCalculator"
|
||||
// input_stream: "DETECTIONS:detections"
|
||||
// input_stream: "IMAGE_SIZE:image_size"
|
||||
// output_stream: "NORM_RECT:rect"
|
||||
// options: {
|
||||
// [mediapipe.DetectionsToRectCalculatorOptions.ext] {
|
||||
// rotation_vector_start_keypoint_index: 0
|
||||
// rotation_vector_end_keypoint_index: 2
|
||||
// rotation_vector_target_angle_degrees: 90
|
||||
// output_zero_rect_for_empty_detections: true
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
class DetectionsToRectsCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
float ComputeRotation(const Detection& detection,
|
||||
const std::pair<int, int> image_size);
|
||||
|
||||
DetectionsToRectsCalculatorOptions options_;
|
||||
int start_keypoint_index_;
|
||||
int end_keypoint_index_;
|
||||
float target_angle_; // In radians.
|
||||
bool rotate_;
|
||||
bool output_zero_rect_for_empty_detections_;
|
||||
};
|
||||
REGISTER_CALCULATOR(DetectionsToRectsCalculator);
|
||||
|
||||
::mediapipe::Status DetectionsToRectsCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag(kDetectionTag) ^
|
||||
@@ -232,6 +162,13 @@ REGISTER_CALCULATOR(DetectionsToRectsCalculator);
|
||||
.Tag(kNormRectTag)
|
||||
.AddPacket(MakePacket<NormalizedRect>().At(cc->InputTimestamp()));
|
||||
}
|
||||
if (cc->Outputs().HasTag(kNormRectsTag)) {
|
||||
auto rect_vector = absl::make_unique<std::vector<NormalizedRect>>();
|
||||
rect_vector->emplace_back(NormalizedRect());
|
||||
cc->Outputs()
|
||||
.Tag(kNormRectsTag)
|
||||
.Add(rect_vector.release(), cc->InputTimestamp());
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -312,4 +249,6 @@ float DetectionsToRectsCalculator::ComputeRotation(
|
||||
return NormalizeRadians(rotation);
|
||||
}
|
||||
|
||||
REGISTER_CALCULATOR(DetectionsToRectsCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
// 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_UTIL_DETECTIONS_TO_RECTS_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_UTIL_DETECTIONS_TO_RECTS_CALCULATOR_H_
|
||||
|
||||
#include <cmath>
|
||||
|
||||
#include "mediapipe/calculators/util/detections_to_rects_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_options.pb.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/location_data.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator that converts Detection proto to Rect proto.
|
||||
//
|
||||
// Detection is the format for encoding one or more detections in an image.
|
||||
// The input can be a single Detection or std::vector<Detection>. The output can
|
||||
// be either a single Rect or NormalizedRect, or std::vector<Rect> or
|
||||
// std::vector<NormalizedRect>. If Rect is used, the LocationData format is
|
||||
// expected to be BOUNDING_BOX, and if NormalizedRect is used it is expected to
|
||||
// be RELATIVE_BOUNDING_BOX.
|
||||
//
|
||||
// When the input is std::vector<Detection> and the output is a Rect or
|
||||
// NormalizedRect, only the first detection is converted. When the input is a
|
||||
// single Detection and the output is a std::vector<Rect> or
|
||||
// std::vector<NormalizedRect>, the output is a vector of size 1.
|
||||
//
|
||||
// Inputs:
|
||||
//
|
||||
// One of the following:
|
||||
// DETECTION: A Detection proto.
|
||||
// DETECTIONS: An std::vector<Detection>.
|
||||
//
|
||||
// IMAGE_SIZE (optional): A std::pair<int, int> represention image width and
|
||||
// height. This is required only when rotation needs to be computed (see
|
||||
// calculator options).
|
||||
//
|
||||
// Output:
|
||||
// One of the following:
|
||||
// RECT: A Rect proto.
|
||||
// NORM_RECT: A NormalizedRect proto.
|
||||
// RECTS: An std::vector<Rect>.
|
||||
// NORM_RECTS: An std::vector<NormalizedRect>.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "DetectionsToRectsCalculator"
|
||||
// input_stream: "DETECTIONS:detections"
|
||||
// input_stream: "IMAGE_SIZE:image_size"
|
||||
// output_stream: "NORM_RECT:rect"
|
||||
// options: {
|
||||
// [mediapipe.DetectionsToRectCalculatorOptions.ext] {
|
||||
// rotation_vector_start_keypoint_index: 0
|
||||
// rotation_vector_end_keypoint_index: 2
|
||||
// rotation_vector_target_angle_degrees: 90
|
||||
// output_zero_rect_for_empty_detections: true
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
class DetectionsToRectsCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
protected:
|
||||
virtual float ComputeRotation(const ::mediapipe::Detection& detection,
|
||||
const std::pair<int, int> image_size);
|
||||
virtual ::mediapipe::Status DetectionToRect(
|
||||
const ::mediapipe::Detection& detection, ::mediapipe::Rect* rect);
|
||||
virtual ::mediapipe::Status DetectionToNormalizedRect(
|
||||
const ::mediapipe::Detection& detection,
|
||||
::mediapipe::NormalizedRect* rect);
|
||||
|
||||
static inline float NormalizeRadians(float angle) {
|
||||
return angle - 2 * M_PI * std::floor((angle - (-M_PI)) / (2 * M_PI));
|
||||
}
|
||||
|
||||
::mediapipe::DetectionsToRectsCalculatorOptions options_;
|
||||
int start_keypoint_index_;
|
||||
int end_keypoint_index_;
|
||||
float target_angle_ = 0.0f; // In radians.
|
||||
bool rotate_;
|
||||
bool output_zero_rect_for_empty_detections_;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
#endif // MEDIAPIPE_CALCULATORS_UTIL_DETECTIONS_TO_RECTS_CALCULATOR_H_
|
||||
@@ -0,0 +1,34 @@
|
||||
|
||||
// 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 "mediapipe/calculators/util/filter_collection_calculator.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef FilterCollectionCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||
FilterNormalizedRectCollectionCalculator;
|
||||
REGISTER_CALCULATOR(FilterNormalizedRectCollectionCalculator);
|
||||
|
||||
typedef FilterCollectionCalculator<
|
||||
std::vector<std::vector<::mediapipe::NormalizedLandmark>>>
|
||||
FilterLandmarksCollectionCalculator;
|
||||
REGISTER_CALCULATOR(FilterLandmarksCollectionCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,109 @@
|
||||
// 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_UTIL_FILTER_VECTOR_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_UTIL_FILTER_VECTOR_CALCULATOR_H_
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator that gates elements of an input collection based on
|
||||
// corresponding boolean values of the "CONDITION" vector. If there is no input
|
||||
// collection or "CONDITION" vector, the calculator forwards timestamp bounds
|
||||
// for downstream calculators. If the "CONDITION" vector has false values for
|
||||
// all elements of the input collection, the calculator outputs a packet
|
||||
// containing an empty collection.
|
||||
// Example usage:
|
||||
// node {
|
||||
// calculator: "FilterCollectionCalculator"
|
||||
// input_stream: "ITERABLE:input_collection"
|
||||
// input_stream: "CONDITION:condition_vector"
|
||||
// output_stream: "ITERABLE:output_collection"
|
||||
// }
|
||||
// This calculator is able to handle collections of copyable types T.
|
||||
template <typename IterableT>
|
||||
class FilterCollectionCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag("ITERABLE"));
|
||||
RET_CHECK(cc->Inputs().HasTag("CONDITION"));
|
||||
RET_CHECK(cc->Outputs().HasTag("ITERABLE"));
|
||||
|
||||
cc->Inputs().Tag("ITERABLE").Set<IterableT>();
|
||||
cc->Inputs().Tag("CONDITION").Set<std::vector<bool>>();
|
||||
|
||||
cc->Outputs().Tag("ITERABLE").Set<IterableT>();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
if (cc->Inputs().Tag("ITERABLE").IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
if (cc->Inputs().Tag("CONDITION").IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
const std::vector<bool>& filter_by =
|
||||
cc->Inputs().Tag("CONDITION").Get<std::vector<bool>>();
|
||||
|
||||
return FilterCollection<IterableT>(
|
||||
std::is_copy_constructible<typename IterableT::value_type>(), cc,
|
||||
filter_by);
|
||||
}
|
||||
|
||||
template <typename IterableU>
|
||||
::mediapipe::Status FilterCollection(std::true_type, CalculatorContext* cc,
|
||||
const std::vector<bool>& filter_by) {
|
||||
const IterableU& input = cc->Inputs().Tag("ITERABLE").Get<IterableU>();
|
||||
if (input.size() != filter_by.size()) {
|
||||
return ::mediapipe::InternalError(absl::StrCat(
|
||||
"Input vector size: ", input.size(),
|
||||
" doesn't mach condition vector size: ", filter_by.size()));
|
||||
}
|
||||
|
||||
auto output = absl::make_unique<IterableU>();
|
||||
for (int i = 0; i < input.size(); ++i) {
|
||||
if (filter_by[i]) {
|
||||
output->push_back(input[i]);
|
||||
}
|
||||
}
|
||||
cc->Outputs().Tag("ITERABLE").Add(output.release(), cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
template <typename IterableU>
|
||||
::mediapipe::Status FilterCollection(std::false_type, CalculatorContext* cc,
|
||||
const std::vector<bool>& filter_by) {
|
||||
return ::mediapipe::InternalError(
|
||||
"Cannot copy input collection to filter it.");
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_UTIL_FILTER_VECTOR_CALCULATOR_H_
|
||||
@@ -0,0 +1,181 @@
|
||||
// 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 <math.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "mediapipe/calculators/util/labels_to_render_data_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/classification.pb.h"
|
||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/statusor.h"
|
||||
#include "mediapipe/util/color.pb.h"
|
||||
#include "mediapipe/util/render_data.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
constexpr float kFontHeightScale = 1.25f;
|
||||
|
||||
// A calculator takes in pairs of labels and scores or classifications, outputs
|
||||
// generates render data. Either both "LABELS" and "SCORES" or "CLASSIFICATIONS"
|
||||
// must be present.
|
||||
//
|
||||
// Usage example:
|
||||
// node {
|
||||
// calculator: "LabelsToRenderDataCalculator"
|
||||
// input_stream: "LABELS:labels"
|
||||
// input_stream: "SCORES:scores"
|
||||
// output_stream: "VIDEO_PRESTREAM:video_header"
|
||||
// options {
|
||||
// [LabelsToRenderDataCalculatorOptions.ext] {
|
||||
// color { r: 255 g: 0 b: 0 }
|
||||
// color { r: 0 g: 255 b: 0 }
|
||||
// color { r: 0 g: 0 b: 255 }
|
||||
// thickness: 2.0
|
||||
// font_height_px: 20
|
||||
// max_num_labels: 3
|
||||
// font_face: 1
|
||||
// location: TOP_LEFT
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
class LabelsToRenderDataCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
LabelsToRenderDataCalculatorOptions options_;
|
||||
int num_colors_ = 0;
|
||||
int video_width_ = 0;
|
||||
int video_height_ = 0;
|
||||
int label_height_px_ = 0;
|
||||
int label_left_px_ = 0;
|
||||
};
|
||||
REGISTER_CALCULATOR(LabelsToRenderDataCalculator);
|
||||
|
||||
::mediapipe::Status LabelsToRenderDataCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
if (cc->Inputs().HasTag("CLASSIFICATIONS")) {
|
||||
cc->Inputs().Tag("CLASSIFICATIONS").Set<ClassificationList>();
|
||||
} else {
|
||||
RET_CHECK(cc->Inputs().HasTag("LABELS"))
|
||||
<< "Must provide input stream \"LABELS\"";
|
||||
cc->Inputs().Tag("LABELS").Set<std::vector<std::string>>();
|
||||
if (cc->Inputs().HasTag("SCORES")) {
|
||||
cc->Inputs().Tag("SCORES").Set<std::vector<float>>();
|
||||
}
|
||||
}
|
||||
if (cc->Inputs().HasTag("VIDEO_PRESTREAM")) {
|
||||
cc->Inputs().Tag("VIDEO_PRESTREAM").Set<VideoHeader>();
|
||||
}
|
||||
cc->Outputs().Tag("RENDER_DATA").Set<RenderData>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status LabelsToRenderDataCalculator::Open(CalculatorContext* cc) {
|
||||
options_ = cc->Options<LabelsToRenderDataCalculatorOptions>();
|
||||
num_colors_ = options_.color_size();
|
||||
label_height_px_ = std::ceil(options_.font_height_px() * kFontHeightScale);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status LabelsToRenderDataCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
if (cc->Inputs().HasTag("VIDEO_PRESTREAM") &&
|
||||
cc->InputTimestamp() == Timestamp::PreStream()) {
|
||||
const VideoHeader& video_header =
|
||||
cc->Inputs().Tag("VIDEO_PRESTREAM").Get<VideoHeader>();
|
||||
video_width_ = video_header.width;
|
||||
video_height_ = video_header.height;
|
||||
return ::mediapipe::OkStatus();
|
||||
} else {
|
||||
CHECK_EQ(options_.location(), LabelsToRenderDataCalculatorOptions::TOP_LEFT)
|
||||
<< "Only TOP_LEFT is supported without VIDEO_PRESTREAM.";
|
||||
}
|
||||
|
||||
std::vector<std::string> labels;
|
||||
std::vector<float> scores;
|
||||
if (cc->Inputs().HasTag("CLASSIFICATIONS")) {
|
||||
const ClassificationList& classifications =
|
||||
cc->Inputs().Tag("CLASSIFICATIONS").Get<ClassificationList>();
|
||||
labels.resize(classifications.classification_size());
|
||||
scores.resize(classifications.classification_size());
|
||||
for (int i = 0; i < classifications.classification_size(); ++i) {
|
||||
labels[i] = classifications.classification(i).label();
|
||||
scores[i] = classifications.classification(i).score();
|
||||
}
|
||||
} else {
|
||||
const std::vector<std::string>& label_vector =
|
||||
cc->Inputs().Tag("LABELS").Get<std::vector<std::string>>();
|
||||
std::vector<float> score_vector;
|
||||
if (cc->Inputs().HasTag("SCORES")) {
|
||||
score_vector = cc->Inputs().Tag("SCORES").Get<std::vector<float>>();
|
||||
}
|
||||
CHECK_EQ(label_vector.size(), score_vector.size());
|
||||
labels.resize(label_vector.size());
|
||||
scores.resize(label_vector.size());
|
||||
for (int i = 0; i < label_vector.size(); ++i) {
|
||||
labels[i] = label_vector[i];
|
||||
scores[i] = score_vector[i];
|
||||
}
|
||||
}
|
||||
|
||||
RenderData render_data;
|
||||
int num_label = std::min((int)labels.size(), options_.max_num_labels());
|
||||
int label_baseline_px = options_.vertical_offset_px();
|
||||
if (options_.location() == LabelsToRenderDataCalculatorOptions::TOP_LEFT) {
|
||||
label_baseline_px += label_height_px_;
|
||||
} else if (options_.location() ==
|
||||
LabelsToRenderDataCalculatorOptions::BOTTOM_LEFT) {
|
||||
label_baseline_px += video_height_ - label_height_px_ * (num_label - 1);
|
||||
}
|
||||
label_left_px_ = options_.horizontal_offset_px();
|
||||
for (int i = 0; i < num_label; ++i) {
|
||||
auto* label_annotation = render_data.add_render_annotations();
|
||||
label_annotation->set_thickness(options_.thickness());
|
||||
if (num_colors_ > 0) {
|
||||
*(label_annotation->mutable_color()) = options_.color(i % num_colors_);
|
||||
} else {
|
||||
label_annotation->mutable_color()->set_r(255);
|
||||
label_annotation->mutable_color()->set_g(0);
|
||||
label_annotation->mutable_color()->set_b(0);
|
||||
}
|
||||
|
||||
auto* text = label_annotation->mutable_text();
|
||||
std::string display_text = labels[i];
|
||||
if (cc->Inputs().HasTag("SCORES")) {
|
||||
absl::StrAppend(&display_text, ":", scores[i]);
|
||||
}
|
||||
text->set_display_text(display_text);
|
||||
text->set_font_height(options_.font_height_px());
|
||||
text->set_left(label_left_px_);
|
||||
text->set_baseline(label_baseline_px + i * label_height_px_);
|
||||
text->set_font_face(options_.font_face());
|
||||
}
|
||||
cc->Outputs()
|
||||
.Tag("RENDER_DATA")
|
||||
.AddPacket(MakePacket<RenderData>(render_data).At(cc->InputTimestamp()));
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,62 @@
|
||||
// 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";
|
||||
import "mediapipe/util/color.proto";
|
||||
|
||||
message LabelsToRenderDataCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional LabelsToRenderDataCalculatorOptions ext = 271660364;
|
||||
}
|
||||
|
||||
// Colors for drawing the label(s).
|
||||
repeated Color color = 1;
|
||||
|
||||
// Thickness for drawing the label(s).
|
||||
optional double thickness = 2 [default = 2];
|
||||
|
||||
// The font height in absolute pixels.
|
||||
optional int32 font_height_px = 3 [default = 50];
|
||||
|
||||
// The offset of the starting text in horizontal direction in absolute pixels.
|
||||
optional int32 horizontal_offset_px = 7 [default = 0];
|
||||
// The offset of the starting text in vertical direction in absolute pixels.
|
||||
optional int32 vertical_offset_px = 8 [default = 0];
|
||||
|
||||
// The maximum number of labels to display.
|
||||
optional int32 max_num_labels = 4 [default = 1];
|
||||
|
||||
// Specifies the font for the text. Font must be one of the following from
|
||||
// OpenCV:
|
||||
// cv::FONT_HERSHEY_SIMPLEX (0)
|
||||
// cv::FONT_HERSHEY_PLAIN (1)
|
||||
// cv::FONT_HERSHEY_DUPLEX (2)
|
||||
// cv::FONT_HERSHEY_COMPLEX (3)
|
||||
// cv::FONT_HERSHEY_TRIPLEX (4)
|
||||
// cv::FONT_HERSHEY_COMPLEX_SMALL (5)
|
||||
// cv::FONT_HERSHEY_SCRIPT_SIMPLEX (6)
|
||||
// cv::FONT_HERSHEY_SCRIPT_COMPLEX (7)
|
||||
optional int32 font_face = 5 [default = 0];
|
||||
|
||||
// Label location.
|
||||
enum Location {
|
||||
TOP_LEFT = 0;
|
||||
BOTTOM_LEFT = 1;
|
||||
}
|
||||
optional Location location = 6 [default = TOP_LEFT];
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
// 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.
|
||||
|
||||
// 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 <cmath>
|
||||
#include <vector>
|
||||
|
||||
#include "Eigen/Core"
|
||||
#include "mediapipe/calculators/util/landmarks_to_floats_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/matrix.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr char kLandmarksTag[] = "NORM_LANDMARKS";
|
||||
constexpr char kFloatsTag[] = "FLOATS";
|
||||
constexpr char kMatrixTag[] = "MATRIX";
|
||||
|
||||
} // namespace
|
||||
|
||||
// Converts a vector of landmarks to a vector of floats or a matrix.
|
||||
// Input:
|
||||
// NORM_LANDMARKS: An std::vector<NormalizedLandmark>.
|
||||
//
|
||||
// Output:
|
||||
// FLOATS(optional): A vector of floats from flattened landmarks.
|
||||
// MATRIX(optional): A matrix of floats of the landmarks.
|
||||
//
|
||||
// Usage example:
|
||||
// node {
|
||||
// calculator: "LandmarksToFloatsCalculator"
|
||||
// input_stream: "NORM_LANDMARKS:landmarks"
|
||||
// output_stream: "MATRIX:landmark_matrix"
|
||||
// }
|
||||
class LandmarksToFloatsCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Tag(kLandmarksTag).Set<std::vector<NormalizedLandmark>>();
|
||||
RET_CHECK(cc->Outputs().HasTag(kFloatsTag) ||
|
||||
cc->Outputs().HasTag(kMatrixTag));
|
||||
if (cc->Outputs().HasTag(kFloatsTag)) {
|
||||
cc->Outputs().Tag(kFloatsTag).Set<std::vector<float>>();
|
||||
}
|
||||
if (cc->Outputs().HasTag(kMatrixTag)) {
|
||||
cc->Outputs().Tag(kMatrixTag).Set<Matrix>();
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::LandmarksToFloatsCalculatorOptions>();
|
||||
num_dimensions_ = options.num_dimensions();
|
||||
// Currently number of dimensions must be within [1, 3].
|
||||
RET_CHECK_GE(num_dimensions_, 1);
|
||||
RET_CHECK_LE(num_dimensions_, 3);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
// Only process if there's input landmarks.
|
||||
if (cc->Inputs().Tag(kLandmarksTag).IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
const auto& input_landmarks =
|
||||
cc->Inputs().Tag(kLandmarksTag).Get<std::vector<NormalizedLandmark>>();
|
||||
|
||||
if (cc->Outputs().HasTag(kFloatsTag)) {
|
||||
auto output_floats = absl::make_unique<std::vector<float>>();
|
||||
for (const auto& landmark : input_landmarks) {
|
||||
output_floats->emplace_back(landmark.x());
|
||||
if (num_dimensions_ > 1) {
|
||||
output_floats->emplace_back(landmark.y());
|
||||
}
|
||||
if (num_dimensions_ > 2) {
|
||||
output_floats->emplace_back(landmark.z());
|
||||
}
|
||||
}
|
||||
|
||||
cc->Outputs()
|
||||
.Tag(kFloatsTag)
|
||||
.Add(output_floats.release(), cc->InputTimestamp());
|
||||
} else {
|
||||
auto output_matrix = absl::make_unique<Matrix>();
|
||||
output_matrix->setZero(num_dimensions_, input_landmarks.size());
|
||||
for (int i = 0; i < input_landmarks.size(); ++i) {
|
||||
(*output_matrix)(0, i) = input_landmarks[i].x();
|
||||
if (num_dimensions_ > 1) {
|
||||
(*output_matrix)(1, i) = input_landmarks[i].y();
|
||||
}
|
||||
if (num_dimensions_ > 2) {
|
||||
(*output_matrix)(2, i) = input_landmarks[i].z();
|
||||
}
|
||||
}
|
||||
cc->Outputs()
|
||||
.Tag(kMatrixTag)
|
||||
.Add(output_matrix.release(), cc->InputTimestamp());
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
int num_dimensions_ = 0;
|
||||
};
|
||||
REGISTER_CALCULATOR(LandmarksToFloatsCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,28 @@
|
||||
// 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 LandmarksToFloatsCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional LandmarksToFloatsCalculatorOptions ext = 274035660;
|
||||
}
|
||||
|
||||
// Number of dimensions to convert. Must within [1, 3].
|
||||
optional int32 num_dimensions = 1 [default = 2];
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
// 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 <memory>
|
||||
#include <string>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/file_helpers.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
// The calculator takes the path to the local file as an input side packet and
|
||||
// outputs the contents of that file.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "LocalFileContentsCalculator"
|
||||
// input_side_packet: "FILE_PATH:file_path"
|
||||
// output_side_packet: "CONTENTS:contents"
|
||||
// }
|
||||
class LocalFileContentsCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->InputSidePackets().Tag("FILE_PATH").Set<std::string>();
|
||||
cc->OutputSidePackets().Tag("CONTENTS").Set<std::string>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
std::string contents;
|
||||
MP_RETURN_IF_ERROR(mediapipe::file::GetContents(
|
||||
cc->InputSidePackets().Tag("FILE_PATH").Get<std::string>(), &contents));
|
||||
cc->OutputSidePackets()
|
||||
.Tag("CONTENTS")
|
||||
.Set(MakePacket<std::string>(std::move(contents)));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(LocalFileContentsCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -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)
|
||||
|
||||
@@ -23,7 +23,9 @@ namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
constexpr char kNormRectTag[] = "NORM_RECT";
|
||||
constexpr char kNormRectsTag[] = "NORM_RECTS";
|
||||
constexpr char kRectTag[] = "RECT";
|
||||
constexpr char kRectsTag[] = "RECTS";
|
||||
constexpr char kImageSizeTag[] = "IMAGE_SIZE";
|
||||
|
||||
// Wraps around an angle in radians to within -M_PI and M_PI.
|
||||
@@ -72,17 +74,31 @@ REGISTER_CALCULATOR(RectTransformationCalculator);
|
||||
|
||||
::mediapipe::Status RectTransformationCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag(kNormRectTag) ^ cc->Inputs().HasTag(kRectTag));
|
||||
RET_CHECK_EQ((cc->Inputs().HasTag(kNormRectTag) ? 1 : 0) +
|
||||
(cc->Inputs().HasTag(kNormRectsTag) ? 1 : 0) +
|
||||
(cc->Inputs().HasTag(kRectTag) ? 1 : 0) +
|
||||
(cc->Inputs().HasTag(kRectsTag) ? 1 : 0),
|
||||
1);
|
||||
if (cc->Inputs().HasTag(kRectTag)) {
|
||||
cc->Inputs().Tag(kRectTag).Set<Rect>();
|
||||
cc->Outputs().Index(0).Set<Rect>();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kRectsTag)) {
|
||||
cc->Inputs().Tag(kRectsTag).Set<std::vector<Rect>>();
|
||||
cc->Outputs().Index(0).Set<std::vector<Rect>>();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kNormRectTag)) {
|
||||
RET_CHECK(cc->Inputs().HasTag(kImageSizeTag));
|
||||
cc->Inputs().Tag(kNormRectTag).Set<NormalizedRect>();
|
||||
cc->Inputs().Tag(kImageSizeTag).Set<std::pair<int, int>>();
|
||||
cc->Outputs().Index(0).Set<NormalizedRect>();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kNormRectsTag)) {
|
||||
RET_CHECK(cc->Inputs().HasTag(kImageSizeTag));
|
||||
cc->Inputs().Tag(kNormRectsTag).Set<std::vector<NormalizedRect>>();
|
||||
cc->Inputs().Tag(kImageSizeTag).Set<std::pair<int, int>>();
|
||||
cc->Outputs().Index(0).Set<std::vector<NormalizedRect>>();
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -105,7 +121,17 @@ REGISTER_CALCULATOR(RectTransformationCalculator);
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
MakePacket<Rect>(rect).At(cc->InputTimestamp()));
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag(kRectsTag) &&
|
||||
!cc->Inputs().Tag(kRectsTag).IsEmpty()) {
|
||||
auto rects = cc->Inputs().Tag(kRectsTag).Get<std::vector<Rect>>();
|
||||
auto output_rects = absl::make_unique<std::vector<Rect>>(rects.size());
|
||||
for (int i = 0; i < rects.size(); ++i) {
|
||||
output_rects->at(i) = rects[i];
|
||||
auto it = output_rects->begin() + i;
|
||||
TransformRect(&(*it));
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output_rects.release(), cc->InputTimestamp());
|
||||
}
|
||||
if (cc->Inputs().HasTag(kNormRectTag) &&
|
||||
!cc->Inputs().Tag(kNormRectTag).IsEmpty()) {
|
||||
auto rect = cc->Inputs().Tag(kNormRectTag).Get<NormalizedRect>();
|
||||
@@ -115,6 +141,21 @@ REGISTER_CALCULATOR(RectTransformationCalculator);
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
MakePacket<NormalizedRect>(rect).At(cc->InputTimestamp()));
|
||||
}
|
||||
if (cc->Inputs().HasTag(kNormRectsTag) &&
|
||||
!cc->Inputs().Tag(kNormRectsTag).IsEmpty()) {
|
||||
auto rects =
|
||||
cc->Inputs().Tag(kNormRectsTag).Get<std::vector<NormalizedRect>>();
|
||||
const auto& image_size =
|
||||
cc->Inputs().Tag(kImageSizeTag).Get<std::pair<int, int>>();
|
||||
auto output_rects =
|
||||
absl::make_unique<std::vector<NormalizedRect>>(rects.size());
|
||||
for (int i = 0; i < rects.size(); ++i) {
|
||||
output_rects->at(i) = rects[i];
|
||||
auto it = output_rects->begin() + i;
|
||||
TransformNormalizedRect(&(*it), image_size.first, image_size.second);
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output_rects.release(), cc->InputTimestamp());
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -0,0 +1,236 @@
|
||||
// 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/formats/classification.pb.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(__EMSCRIPTEN__) || \
|
||||
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, classification protos, and summary std::string
|
||||
// (in csv format).
|
||||
//
|
||||
// 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"
|
||||
// output_stream: "TOP_K_CLASSIFICATIONS:top_k_classes"
|
||||
// output_stream: "SUMMARY:summary"
|
||||
// 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_;
|
||||
bool label_map_loaded_ = false;
|
||||
};
|
||||
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>>();
|
||||
}
|
||||
if (cc->Outputs().HasTag("CLASSIFICATIONS")) {
|
||||
cc->Outputs().Tag("CLASSIFICATIONS").Set<ClassificationList>();
|
||||
}
|
||||
if (cc->Outputs().HasTag("SUMMARY")) {
|
||||
cc->Outputs().Tag("SUMMARY").Set<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 (label_map_loaded_) {
|
||||
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()));
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("SUMMARY")) {
|
||||
std::vector<std::string> results;
|
||||
for (int index = 0; index < top_k_indexes.size(); ++index) {
|
||||
if (label_map_loaded_) {
|
||||
results.push_back(
|
||||
absl::StrCat(top_k_labels[index], ":", top_k_scores[index]));
|
||||
} else {
|
||||
results.push_back(
|
||||
absl::StrCat(top_k_indexes[index], ":", top_k_scores[index]));
|
||||
}
|
||||
}
|
||||
cc->Outputs().Tag("SUMMARY").AddPacket(
|
||||
MakePacket<std::string>(absl::StrJoin(results, ","))
|
||||
.At(cc->InputTimestamp()));
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("TOP_K_CLASSIFICATION")) {
|
||||
auto classification_list = absl::make_unique<ClassificationList>();
|
||||
for (int index = 0; index < top_k_indexes.size(); ++index) {
|
||||
Classification* classification =
|
||||
classification_list->add_classification();
|
||||
classification->set_index(top_k_indexes[index]);
|
||||
classification->set_score(top_k_scores[index]);
|
||||
if (label_map_loaded_) {
|
||||
classification->set_label(top_k_labels[index]);
|
||||
}
|
||||
}
|
||||
}
|
||||
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;
|
||||
}
|
||||
label_map_loaded_ = true;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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
|
||||
@@ -13,12 +13,16 @@
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
load(
|
||||
"//mediapipe/framework/tool:mediapipe_graph.bzl",
|
||||
"mediapipe_binary_graph",
|
||||
)
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "flow_to_image_calculator_proto",
|
||||
srcs = ["flow_to_image_calculator.proto"],
|
||||
@@ -52,9 +56,7 @@ mediapipe_cc_proto_library(
|
||||
cc_library(
|
||||
name = "flow_to_image_calculator",
|
||||
srcs = ["flow_to_image_calculator.cc"],
|
||||
visibility = [
|
||||
"//visibility:public",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/video:flow_to_image_calculator_cc_proto",
|
||||
"//mediapipe/calculators/video/tool:flow_quantizer_model",
|
||||
@@ -129,10 +131,20 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
filegroup(
|
||||
name = "test_videos",
|
||||
srcs = [
|
||||
"testdata/format_FLV_H264_AAC.video",
|
||||
"testdata/format_MKV_VP8_VORBIS.video",
|
||||
"testdata/format_MP4_AVC720P_AAC.video",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "opencv_video_decoder_calculator_test",
|
||||
srcs = ["opencv_video_decoder_calculator_test.cc"],
|
||||
data = ["//mediapipe/calculators/video/testdata:test_videos"],
|
||||
data = [":test_videos"],
|
||||
deps = [
|
||||
":opencv_video_decoder_calculator",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
@@ -151,7 +163,7 @@ cc_test(
|
||||
cc_test(
|
||||
name = "opencv_video_encoder_calculator_test",
|
||||
srcs = ["opencv_video_encoder_calculator_test.cc"],
|
||||
data = ["//mediapipe/calculators/video/testdata:test_videos"],
|
||||
data = [":test_videos"],
|
||||
deps = [
|
||||
":opencv_video_decoder_calculator",
|
||||
":opencv_video_encoder_calculator",
|
||||
@@ -175,7 +187,6 @@ cc_test(
|
||||
cc_test(
|
||||
name = "tvl1_optical_flow_calculator_test",
|
||||
srcs = ["tvl1_optical_flow_calculator_test.cc"],
|
||||
data = ["//mediapipe/calculators/image/testdata:test_images"],
|
||||
deps = [
|
||||
":tvl1_optical_flow_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
|
||||
@@ -123,6 +123,7 @@ class OpenCvVideoDecoderCalculator : public CalculatorBase {
|
||||
cc->Outputs()
|
||||
.Tag("VIDEO_PRESTREAM")
|
||||
.Add(header.release(), Timestamp::PreStream());
|
||||
cc->Outputs().Tag("VIDEO_PRESTREAM").Close();
|
||||
}
|
||||
// Rewind to the very first frame.
|
||||
cap_->set(cv::CAP_PROP_POS_AVI_RATIO, 0);
|
||||
|
||||
@@ -13,24 +13,24 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//mediapipe/calculators/video:__subpackages__"])
|
||||
|
||||
exports_files(["LICENSE"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "flow_quantizer_model_proto",
|
||||
srcs = ["flow_quantizer_model.proto"],
|
||||
visibility = ["//mediapipe:__subpackages__"],
|
||||
visibility = ["//visibility:public"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "flow_quantizer_model_cc_proto",
|
||||
srcs = ["flow_quantizer_model.proto"],
|
||||
visibility = ["//mediapipe:__subpackages__"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":flow_quantizer_model_proto"],
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
## MediaPipe Android Archive Library
|
||||
|
||||
***Experimental Only***
|
||||
|
||||
The MediaPipe Android archive library is a convenient way to use MediaPipe with
|
||||
Android Studio and Gradle. MediaPipe doesn't publish a general AAR that can be
|
||||
used by all projects. Instead, developers need to add a mediapipe_aar() target
|
||||
to generate a custom AAR file for their own projects. This is necessary in order
|
||||
to include specific resources such as MediaPipe calculators needed for each
|
||||
project.
|
||||
|
||||
### Steps to build a MediaPipe AAR
|
||||
|
||||
1. Create a mediapipe_aar() target.
|
||||
|
||||
In the MediaPipe directory, create a new mediapipe_aar() target in a BUILD
|
||||
file. You need to figure out what calculators are used in the graph and
|
||||
provide the calculator dependencies to the mediapipe_aar(). For example, to
|
||||
build an AAR for [face detection gpu](./face_detection_mobile_gpu.md), you
|
||||
can put the following code into
|
||||
mediapipe/examples/android/src/java/com/google/mediapipe/apps/aar_example/BUILD.
|
||||
|
||||
```
|
||||
load("//mediapipe/java/com/google/mediapipe:mediapipe_aar.bzl", "mediapipe_aar")
|
||||
|
||||
mediapipe_aar(
|
||||
name = "mp_face_detection_aar",
|
||||
calculators = ["//mediapipe/graphs/face_detection:mobile_calculators"],
|
||||
)
|
||||
```
|
||||
|
||||
2. Run the Bazel build command to generate the AAR.
|
||||
|
||||
```bash
|
||||
bazel build -c opt --fat_apk_cpu=arm64-v8a,armeabi-v7a //path/to/the/aar/build/file:aar_name
|
||||
```
|
||||
|
||||
For the face detection AAR target we made in the step 1, run:
|
||||
|
||||
```bash
|
||||
bazel build -c opt --fat_apk_cpu=arm64-v8a,armeabi-v7a \
|
||||
//mediapipe/examples/android/src/java/com/google/mediapipe/apps/aar_example:mp_face_detection_aar
|
||||
|
||||
# It should print:
|
||||
# Target //mediapipe/examples/android/src/java/com/google/mediapipe/apps/aar_example:mp_face_detection_aar up-to-date:
|
||||
# bazel-bin/mediapipe/examples/android/src/java/com/google/mediapipe/apps/aar_example/mp_face_detection_aar.aar
|
||||
```
|
||||
|
||||
3. (Optional) Save the AAR to your preferred location.
|
||||
|
||||
```bash
|
||||
cp bazel-bin/mediapipe/examples/android/src/java/com/google/mediapipe/apps/aar_example/mp_face_detection_aar.aar
|
||||
/absolute/path/to/your/preferred/location
|
||||
```
|
||||
|
||||
### Steps to use a MediaPipe AAR in Android Studio with Gradle
|
||||
|
||||
1. Start Android Studio and go to your project.
|
||||
|
||||
2. Copy the AAR into app/libs.
|
||||
|
||||
```bash
|
||||
cp bazel-bin/mediapipe/examples/android/src/java/com/google/mediapipe/apps/aar_example/mp_face_detection_aar.aar
|
||||
/path/to/your/app/libs/
|
||||
```
|
||||
|
||||

|
||||
|
||||
3. Make app/src/main/assets and copy assets (graph, model, and etc) into
|
||||
app/src/main/assets.
|
||||
|
||||
Build the MediaPipe binary graph and copy the assets into
|
||||
app/src/main/assets, e.g., for the face detection graph, you need to build
|
||||
and copy
|
||||
[the binary graph](https://github.com/google/mediapipe/blob/master/mediapipe/examples/android/src/java/com/google/mediapipe/apps/facedetectiongpu/BUILD#L41),
|
||||
[the tflite model](https://github.com/google/mediapipe/tree/master/mediapipe/models/face_detection_front.tflite),
|
||||
and
|
||||
[the label map](https://github.com/google/mediapipe/blob/master/mediapipe/models/face_detection_front_labelmap.txt).
|
||||
|
||||
```bash
|
||||
bazel build -c opt mediapipe/examples/android/src/java/com/google/mediapipe/apps/facedetectiongpu:binary_graph
|
||||
cp bazel-bin/mediapipe/examples/android/src/java/com/google/mediapipe/apps/facedetectiongpu/facedetectiongpu.binarypb /path/to/your/app/src/main/assets/
|
||||
cp mediapipe/models/face_detection_front.tflite /path/to/your/app/src/main/assets/
|
||||
cp mediapipe/models/face_detection_front_labelmap.txt /path/to/your/app/src/main/assets/
|
||||
```
|
||||
|
||||

|
||||
|
||||
4. Make app/src/main/jniLibs and copy OpenCV JNI libraries into
|
||||
app/src/main/jniLibs.
|
||||
|
||||
MediaPipe depends on OpenCV, you will need to copy the precompiled OpenCV so
|
||||
files into app/src/main/jniLibs. You can download the official OpenCV
|
||||
Android SDK from
|
||||
[here](https://github.com/opencv/opencv/releases/download/4.1.0/opencv-4.1.0-android-sdk.zip)
|
||||
and run:
|
||||
|
||||
```bash
|
||||
cp -R ~/Downloads/OpenCV-android-sdk/sdk/native/libs/arm* /path/to/your/app/src/main/jniLibs/
|
||||
```
|
||||
|
||||

|
||||
|
||||
5. Modify app/build.gradle to add MediaPipe dependencies and MediaPipe AAR.
|
||||
|
||||
```
|
||||
dependencies {
|
||||
implementation fileTree(dir: 'libs', include: ['*.jar', '*.aar'])
|
||||
implementation 'androidx.appcompat:appcompat:1.0.2'
|
||||
implementation 'androidx.constraintlayout:constraintlayout:1.1.3'
|
||||
testImplementation 'junit:junit:4.12'
|
||||
androidTestImplementation 'androidx.test.ext:junit:1.1.0'
|
||||
androidTestImplementation 'androidx.test.espresso:espresso-core:3.1.1'
|
||||
// MediaPipe deps
|
||||
implementation 'com.google.flogger:flogger:0.3.1'
|
||||
implementation 'com.google.flogger:flogger-system-backend:0.3.1'
|
||||
implementation 'com.google.code.findbugs:jsr305:3.0.2'
|
||||
implementation 'com.google.guava:guava:27.0.1-android'
|
||||
implementation 'com.google.guava:guava:27.0.1-android'
|
||||
implementation 'com.google.protobuf:protobuf-lite:3.0.0'
|
||||
// CameraX core library
|
||||
def camerax_version = "1.0.0-alpha06"
|
||||
implementation "androidx.camera:camera-core:$camerax_version"
|
||||
implementation "androidx.camera:camera-camera2:$camerax_version"
|
||||
}
|
||||
```
|
||||
|
||||
6. Follow our Android app examples to use MediaPipe in Android Studio for your
|
||||
use case. If you are looking for an example, a working face detection
|
||||
example can be found
|
||||
[here](https://github.com/jiuqiant/mediapipe_aar_example).
|
||||
@@ -51,6 +51,15 @@ and model details are described in the
|
||||
* [Android](./face_detection_mobile_gpu.md)
|
||||
* [iOS](./face_detection_mobile_gpu.md)
|
||||
|
||||
### Face Detection with CPU
|
||||
|
||||
[Face Detection with CPU](./face_detection_mobile_cpu.md) illustrates using the
|
||||
same TFLite model in a CPU-based pipeline. This example highlights how graphs
|
||||
can be easily adapted to run on CPU v.s. GPU.
|
||||
|
||||
* [Android](./face_detection_mobile_cpu.md)
|
||||
* [iOS](./face_detection_mobile_cpu.md)
|
||||
|
||||
### Hand Detection with GPU
|
||||
|
||||
[Hand Detection with GPU](./hand_detection_mobile_gpu.md) illustrates how to use
|
||||
@@ -67,6 +76,15 @@ MediaPipe with a TFLite model for hand tracking in a GPU-accelerated pipeline.
|
||||
* [Android](./hand_tracking_mobile_gpu.md)
|
||||
* [iOS](./hand_tracking_mobile_gpu.md)
|
||||
|
||||
### Multi-Hand Tracking with GPU
|
||||
|
||||
[Multi-Hand Tracking with GPU](./multi_hand_tracking_mobile_gpu.md) illustrates
|
||||
how to use MediaPipe with a TFLite model for multi-hand tracking in a
|
||||
GPU-accelerated pipeline.
|
||||
|
||||
* [Android](./multi_hand_tracking_mobile_gpu.md)
|
||||
* [iOS](./multi_hand_tracking_mobile_gpu.md)
|
||||
|
||||
### Hair Segmentation with GPU
|
||||
|
||||
[Hair Segmentation on GPU](./hair_segmentation_mobile_gpu.md) illustrates how to
|
||||
@@ -87,8 +105,9 @@ using the MediaPipe C++ APIs.
|
||||
|
||||
### Feature Extration for YouTube-8M Challenge
|
||||
|
||||
[Feature Extration for YouTube-8M Challenge](./youtube_8m.md) shows how to use
|
||||
MediaPipe to prepare training data for the YouTube-8M Challenge.
|
||||
[Feature Extration and Model Inference for YouTube-8M Challenge](./youtube_8m.md)
|
||||
shows how to use MediaPipe to prepare training data for the YouTube-8M Challenge
|
||||
and do the model inference with the baseline model.
|
||||
|
||||
### Preparing Data Sets with MediaSequence
|
||||
|
||||
@@ -103,3 +122,38 @@ object detection models (TensorFlow and TFLite) using the MediaPipe C++ APIs.
|
||||
|
||||
[Sobel edge detection]:https://en.wikipedia.org/wiki/Sobel_operator
|
||||
[CameraX]:https://developer.android.com/training/camerax
|
||||
|
||||
### Face Detection on Desktop with Webcam
|
||||
|
||||
[Face Detection on Desktop with Webcam](./face_detection_desktop.md) shows how
|
||||
to use MediaPipe with a TFLite model for face detection on desktop using CPU or
|
||||
GPU with live video from a webcam.
|
||||
|
||||
* [Desktop GPU](./face_detection_desktop.md)
|
||||
* [Desktop CPU](./face_detection_desktop.md)
|
||||
|
||||
### Hand Tracking on Desktop with Webcam
|
||||
|
||||
[Hand Tracking on Desktop with Webcam](./hand_tracking_desktop.md) shows how to
|
||||
use MediaPipe with a TFLite model for hand tracking on desktop using CPU or GPU
|
||||
with live video from a webcam.
|
||||
|
||||
* [Desktop GPU](./hand_tracking_desktop.md)
|
||||
* [Desktop CPU](./hand_tracking_desktop.md)
|
||||
|
||||
### Multi-Hand Tracking on Desktop with Webcam
|
||||
|
||||
[Multi-Hand Tracking on Desktop with Webcam](./multi_hand_tracking_desktop.md)
|
||||
shows how to use MediaPipe with a TFLite model for multi-hand tracking on
|
||||
desktop using CPU or GPU with live video from a webcam.
|
||||
|
||||
* [Desktop GPU](./multi_hand_tracking_desktop.md)
|
||||
* [Desktop CPU](./multi_hand_tracking_desktop.md)
|
||||
|
||||
### Hair Segmentation on Desktop with Webcam
|
||||
|
||||
[Hair Segmentation on Desktop with Webcam](./hair_segmentation_desktop.md) shows
|
||||
how to use MediaPipe with a TFLite model for hair segmentation on desktop using
|
||||
GPU with live video from a webcam.
|
||||
|
||||
* [Desktop GPU](./hair_segmentation_desktop.md)
|
||||
|
||||
@@ -0,0 +1,263 @@
|
||||
## Face Detection on Desktop
|
||||
|
||||
This is an example of using MediaPipe to run face detection models (TensorFlow
|
||||
Lite) and render bounding boxes on the detected faces. To know more about the
|
||||
face detection models, please refer to the model [`README file`]. Moreover, if
|
||||
you are interested in running the same TensorfFlow Lite model on Android/iOS,
|
||||
please see the
|
||||
[Face Detection on GPU on Android/iOS](face_detection_mobile_gpu.md) and
|
||||
[Face Detection on CPU on Android/iOS](face_detection_mobile_cpu.md) examples.
|
||||
|
||||
We show the face detection demos with TensorFlow Lite model using the Webcam:
|
||||
|
||||
- [TensorFlow Lite Face Detection Demo with Webcam (CPU)](#tensorflow-lite-face-detection-demo-with-webcam-cpu)
|
||||
|
||||
- [TensorFlow Lite Face Detection Demo with Webcam (GPU)](#tensorflow-lite-face-detection-demo-with-webcam-gpu)
|
||||
|
||||
Note: Desktop GPU works only on Linux. Mesa drivers need to be installed. Please
|
||||
see
|
||||
[step 4 of "Installing on Debian and Ubuntu" in the installation guide](./install.md).
|
||||
|
||||
Note: If MediaPipe depends on OpenCV 2, please see the [known issues with OpenCV 2](#known-issues-with-opencv-2) section.
|
||||
|
||||
### TensorFlow Lite Face Detection Demo with Webcam (CPU)
|
||||
|
||||
To build and run the TensorFlow Lite example on desktop (CPU) with Webcam, run:
|
||||
|
||||
```bash
|
||||
# Video from webcam running on desktop CPU
|
||||
$ bazel build -c opt --define MEDIAPIPE_DISABLE_GPU=1 \
|
||||
mediapipe/examples/desktop/face_detection:face_detection_cpu
|
||||
|
||||
# It should print:
|
||||
# Target //mediapipe/examples/desktop/face_detection:face_detection_cpu up-to-date:
|
||||
# bazel-bin/mediapipe/examples/desktop/face_detection/face_detection_cpu
|
||||
# INFO: Elapsed time: 36.417s, Critical Path: 23.22s
|
||||
# INFO: 711 processes: 710 linux-sandbox, 1 local.
|
||||
# INFO: Build completed successfully, 734 total actions
|
||||
|
||||
# This will open up your webcam as long as it is connected and on
|
||||
# Any errors is likely due to your webcam being not accessible
|
||||
$ GLOG_logtostderr=1 bazel-bin/mediapipe/examples/desktop/face_detection/face_detection_cpu \
|
||||
--calculator_graph_config_file=mediapipe/graphs/face_detection/face_detection_desktop_live.pbtxt
|
||||
```
|
||||
|
||||
### TensorFlow Lite Face Detection Demo with Webcam (GPU)
|
||||
|
||||
To build and run the TensorFlow Lite example on desktop (GPU) with Webcam, run:
|
||||
|
||||
```bash
|
||||
# Video from webcam running on desktop GPU
|
||||
# This works only for linux currently
|
||||
$ bazel build -c opt --copt -DMESA_EGL_NO_X11_HEADERS \
|
||||
mediapipe/examples/desktop/face_detection:face_detection_gpu
|
||||
|
||||
# It should print:
|
||||
# Target //mediapipe/examples/desktop/face_detection:face_detection_gpu up-to-date:
|
||||
# bazel-bin/mediapipe/examples/desktop/face_detection/face_detection_gpu
|
||||
# INFO: Elapsed time: 36.417s, Critical Path: 23.22s
|
||||
# INFO: 711 processes: 710 linux-sandbox, 1 local.
|
||||
# INFO: Build completed successfully, 734 total actions
|
||||
|
||||
# This will open up your webcam as long as it is connected and on
|
||||
# Any errors is likely due to your webcam being not accessible,
|
||||
# or GPU drivers not setup properly.
|
||||
$ GLOG_logtostderr=1 bazel-bin/mediapipe/examples/desktop/face_detection/face_detection_gpu \
|
||||
--calculator_graph_config_file=mediapipe/graphs/face_detection/face_detection_mobile_gpu.pbtxt
|
||||
```
|
||||
|
||||
#### Graph
|
||||
|
||||

|
||||
|
||||
To visualize the graph as shown above, copy the text specification of the graph
|
||||
below and paste it into
|
||||
[MediaPipe Visualizer](https://viz.mediapipe.dev).
|
||||
|
||||
```bash
|
||||
# MediaPipe graph that performs face detection with TensorFlow Lite on CPU & GPU.
|
||||
# Used in the examples in
|
||||
# mediapipe/examples/desktop/face_detection:face_detection_cpu.
|
||||
|
||||
# Images on CPU coming into and out of the graph.
|
||||
input_stream: "input_video"
|
||||
output_stream: "output_video"
|
||||
|
||||
# Throttles the images flowing downstream for flow control. It passes through
|
||||
# the very first incoming image unaltered, and waits for
|
||||
# TfLiteTensorsToDetectionsCalculator downstream in the graph to finish
|
||||
# generating the corresponding detections before it passes through another
|
||||
# image. All images that come in while waiting are dropped, limiting the number
|
||||
# of in-flight images between this calculator and
|
||||
# TfLiteTensorsToDetectionsCalculator to 1. This prevents the nodes in between
|
||||
# from queuing up incoming images and data excessively, which leads to increased
|
||||
# latency and memory usage, unwanted in real-time mobile applications. It also
|
||||
# eliminates unnecessarily computation, e.g., a transformed image produced by
|
||||
# ImageTransformationCalculator may get dropped downstream if the subsequent
|
||||
# TfLiteConverterCalculator or TfLiteInferenceCalculator is still busy
|
||||
# processing previous inputs.
|
||||
node {
|
||||
calculator: "FlowLimiterCalculator"
|
||||
input_stream: "input_video"
|
||||
input_stream: "FINISHED:detections"
|
||||
input_stream_info: {
|
||||
tag_index: "FINISHED"
|
||||
back_edge: true
|
||||
}
|
||||
output_stream: "throttled_input_video"
|
||||
}
|
||||
|
||||
# Transforms the input image on CPU to a 128x128 image. To scale the input
|
||||
# image, the scale_mode option is set to FIT to preserve the aspect ratio,
|
||||
# resulting in potential letterboxing in the transformed image.
|
||||
node: {
|
||||
calculator: "ImageTransformationCalculator"
|
||||
input_stream: "IMAGE:throttled_input_video"
|
||||
output_stream: "IMAGE:transformed_input_video_cpu"
|
||||
output_stream: "LETTERBOX_PADDING:letterbox_padding"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.ImageTransformationCalculatorOptions] {
|
||||
output_width: 128
|
||||
output_height: 128
|
||||
scale_mode: FIT
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Converts the transformed input image on CPU into an image tensor stored as a
|
||||
# TfLiteTensor.
|
||||
node {
|
||||
calculator: "TfLiteConverterCalculator"
|
||||
input_stream: "IMAGE:transformed_input_video_cpu"
|
||||
output_stream: "TENSORS:image_tensor"
|
||||
}
|
||||
|
||||
# Runs a TensorFlow Lite model on CPU that takes an image tensor and outputs a
|
||||
# vector of tensors representing, for instance, detection boxes/keypoints and
|
||||
# scores.
|
||||
node {
|
||||
calculator: "TfLiteInferenceCalculator"
|
||||
input_stream: "TENSORS:image_tensor"
|
||||
output_stream: "TENSORS:detection_tensors"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.TfLiteInferenceCalculatorOptions] {
|
||||
model_path: "mediapipe/models/face_detection_front.tflite"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Generates a single side packet containing a vector of SSD anchors based on
|
||||
# the specification in the options.
|
||||
node {
|
||||
calculator: "SsdAnchorsCalculator"
|
||||
output_side_packet: "anchors"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.SsdAnchorsCalculatorOptions] {
|
||||
num_layers: 4
|
||||
min_scale: 0.1484375
|
||||
max_scale: 0.75
|
||||
input_size_height: 128
|
||||
input_size_width: 128
|
||||
anchor_offset_x: 0.5
|
||||
anchor_offset_y: 0.5
|
||||
strides: 8
|
||||
strides: 16
|
||||
strides: 16
|
||||
strides: 16
|
||||
aspect_ratios: 1.0
|
||||
fixed_anchor_size: true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Decodes the detection tensors generated by the TensorFlow Lite model, based on
|
||||
# the SSD anchors and the specification in the options, into a vector of
|
||||
# detections. Each detection describes a detected object.
|
||||
node {
|
||||
calculator: "TfLiteTensorsToDetectionsCalculator"
|
||||
input_stream: "TENSORS:detection_tensors"
|
||||
input_side_packet: "ANCHORS:anchors"
|
||||
output_stream: "DETECTIONS:detections"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.TfLiteTensorsToDetectionsCalculatorOptions] {
|
||||
num_classes: 1
|
||||
num_boxes: 896
|
||||
num_coords: 16
|
||||
box_coord_offset: 0
|
||||
keypoint_coord_offset: 4
|
||||
num_keypoints: 6
|
||||
num_values_per_keypoint: 2
|
||||
sigmoid_score: true
|
||||
score_clipping_thresh: 100.0
|
||||
reverse_output_order: true
|
||||
x_scale: 128.0
|
||||
y_scale: 128.0
|
||||
h_scale: 128.0
|
||||
w_scale: 128.0
|
||||
min_score_thresh: 0.75
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Performs non-max suppression to remove excessive detections.
|
||||
node {
|
||||
calculator: "NonMaxSuppressionCalculator"
|
||||
input_stream: "detections"
|
||||
output_stream: "filtered_detections"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.NonMaxSuppressionCalculatorOptions] {
|
||||
min_suppression_threshold: 0.3
|
||||
overlap_type: INTERSECTION_OVER_UNION
|
||||
algorithm: WEIGHTED
|
||||
return_empty_detections: true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Maps detection label IDs to the corresponding label text ("Face"). The label
|
||||
# map is provided in the label_map_path option.
|
||||
node {
|
||||
calculator: "DetectionLabelIdToTextCalculator"
|
||||
input_stream: "filtered_detections"
|
||||
output_stream: "labeled_detections"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.DetectionLabelIdToTextCalculatorOptions] {
|
||||
label_map_path: "mediapipe/models/face_detection_front_labelmap.txt"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Adjusts detection locations (already normalized to [0.f, 1.f]) on the
|
||||
# letterboxed image (after image transformation with the FIT scale mode) to the
|
||||
# corresponding locations on the same image with the letterbox removed (the
|
||||
# input image to the graph before image transformation).
|
||||
node {
|
||||
calculator: "DetectionLetterboxRemovalCalculator"
|
||||
input_stream: "DETECTIONS:labeled_detections"
|
||||
input_stream: "LETTERBOX_PADDING:letterbox_padding"
|
||||
output_stream: "DETECTIONS:output_detections"
|
||||
}
|
||||
|
||||
# Converts the detections to drawing primitives for annotation overlay.
|
||||
node {
|
||||
calculator: "DetectionsToRenderDataCalculator"
|
||||
input_stream: "DETECTIONS:output_detections"
|
||||
output_stream: "RENDER_DATA:render_data"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.DetectionsToRenderDataCalculatorOptions] {
|
||||
thickness: 4.0
|
||||
color { r: 255 g: 0 b: 0 }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Draws annotations and overlays them on top of the input images.
|
||||
node {
|
||||
calculator: "AnnotationOverlayCalculator"
|
||||
input_stream: "INPUT_FRAME:throttled_input_video"
|
||||
input_stream: "render_data"
|
||||
output_stream: "OUTPUT_FRAME:output_video"
|
||||
}
|
||||
```
|
||||
|
||||
[`README file`]:https://github.com/google/mediapipe/tree/master/mediapipe/models/object_detection_saved_model/README.md
|
||||
@@ -0,0 +1,244 @@
|
||||
# Face Detection (CPU)
|
||||
|
||||
This doc focuses on the
|
||||
[example graph](https://github.com/google/mediapipe/tree/master/mediapipe/graphs/face_detection/face_detection_mobile_cpu.pbtxt)
|
||||
that performs face detection with TensorFlow Lite on CPU.
|
||||
|
||||
## Android
|
||||
|
||||
[Source](https://github.com/google/mediapipe/tree/master/mediapipe/examples/android/src/java/com/google/mediapipe/apps/facedetectioncpu)
|
||||
|
||||
To build and install the app:
|
||||
|
||||
```bash
|
||||
bazel build -c opt --config=android_arm64 mediapipe/examples/android/src/java/com/google/mediapipe/apps/facedetectioncpu
|
||||
adb install bazel-bin/mediapipe/examples/android/src/java/com/google/mediapipe/apps/facedetectioncpu/facedetectioncpu.apk
|
||||
```
|
||||
|
||||
## iOS
|
||||
|
||||
[Source](https://github.com/google/mediapipe/tree/master/mediapipe/examples/ios/facedetectioncpu).
|
||||
|
||||
See the general [instructions](./mediapipe_ios_setup.md) for building iOS
|
||||
examples and generating an Xcode project. This will be the FaceDetectionCpuApp
|
||||
target.
|
||||
|
||||
To build on the command line:
|
||||
|
||||
```bash
|
||||
bazel build -c opt --config=ios_arm64 mediapipe/examples/ios/facedetectioncpu:FaceDetectionCpuApp
|
||||
```
|
||||
|
||||
## Graph
|
||||
|
||||

|
||||
|
||||
To visualize the graph as shown above, copy the text specification of the graph
|
||||
below and paste it into [MediaPipe Visualizer](https://viz.mediapipe.dev/).
|
||||
|
||||
[Source pbtxt file](https://github.com/google/mediapipe/tree/master/mediapipe/graphs/face_detection/face_detection_mobile_cpu.pbtxt)
|
||||
|
||||
```bash
|
||||
# MediaPipe graph that performs face detection with TensorFlow Lite on CPU.
|
||||
# Used in the examples in
|
||||
# mediapipie/examples/android/src/java/com/mediapipe/apps/facedetectioncpu and
|
||||
# mediapipie/examples/ios/facedetectioncpu.
|
||||
|
||||
# Images on GPU coming into and out of the graph.
|
||||
input_stream: "input_video"
|
||||
output_stream: "output_video"
|
||||
|
||||
# Throttles the images flowing downstream for flow control. It passes through
|
||||
# the very first incoming image unaltered, and waits for
|
||||
# TfLiteTensorsToDetectionsCalculator downstream in the graph to finish
|
||||
# generating the corresponding detections before it passes through another
|
||||
# image. All images that come in while waiting are dropped, limiting the number
|
||||
# of in-flight images between this calculator and
|
||||
# TfLiteTensorsToDetectionsCalculator to 1. This prevents the nodes in between
|
||||
# from queuing up incoming images and data excessively, which leads to increased
|
||||
# latency and memory usage, unwanted in real-time mobile applications. It also
|
||||
# eliminates unnecessarily computation, e.g., a transformed image produced by
|
||||
# ImageTransformationCalculator may get dropped downstream if the subsequent
|
||||
# TfLiteConverterCalculator or TfLiteInferenceCalculator is still busy
|
||||
# processing previous inputs.
|
||||
node {
|
||||
calculator: "FlowLimiterCalculator"
|
||||
input_stream: "input_video"
|
||||
input_stream: "FINISHED:detections"
|
||||
input_stream_info: {
|
||||
tag_index: "FINISHED"
|
||||
back_edge: true
|
||||
}
|
||||
output_stream: "throttled_input_video"
|
||||
}
|
||||
|
||||
# Transfers the input image from GPU to CPU memory for the purpose of
|
||||
# demonstrating a CPU-based pipeline. Note that the input image on GPU has the
|
||||
# origin defined at the bottom-left corner (OpenGL convention). As a result,
|
||||
# the transferred image on CPU also shares the same representation.
|
||||
node: {
|
||||
calculator: "GpuBufferToImageFrameCalculator"
|
||||
input_stream: "throttled_input_video"
|
||||
output_stream: "input_video_cpu"
|
||||
}
|
||||
|
||||
# Transforms the input image on CPU to a 128x128 image. To scale the input
|
||||
# image, the scale_mode option is set to FIT to preserve the aspect ratio,
|
||||
# resulting in potential letterboxing in the transformed image.
|
||||
node: {
|
||||
calculator: "ImageTransformationCalculator"
|
||||
input_stream: "IMAGE:input_video_cpu"
|
||||
output_stream: "IMAGE:transformed_input_video_cpu"
|
||||
output_stream: "LETTERBOX_PADDING:letterbox_padding"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.ImageTransformationCalculatorOptions] {
|
||||
output_width: 128
|
||||
output_height: 128
|
||||
scale_mode: FIT
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Converts the transformed input image on CPU into an image tensor stored as a
|
||||
# TfLiteTensor.
|
||||
node {
|
||||
calculator: "TfLiteConverterCalculator"
|
||||
input_stream: "IMAGE:transformed_input_video_cpu"
|
||||
output_stream: "TENSORS:image_tensor"
|
||||
}
|
||||
|
||||
# Runs a TensorFlow Lite model on CPU that takes an image tensor and outputs a
|
||||
# vector of tensors representing, for instance, detection boxes/keypoints and
|
||||
# scores.
|
||||
node {
|
||||
calculator: "TfLiteInferenceCalculator"
|
||||
input_stream: "TENSORS:image_tensor"
|
||||
output_stream: "TENSORS:detection_tensors"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.TfLiteInferenceCalculatorOptions] {
|
||||
model_path: "mediapipe/models/face_detection_front.tflite"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Generates a single side packet containing a vector of SSD anchors based on
|
||||
# the specification in the options.
|
||||
node {
|
||||
calculator: "SsdAnchorsCalculator"
|
||||
output_side_packet: "anchors"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.SsdAnchorsCalculatorOptions] {
|
||||
num_layers: 4
|
||||
min_scale: 0.1484375
|
||||
max_scale: 0.75
|
||||
input_size_height: 128
|
||||
input_size_width: 128
|
||||
anchor_offset_x: 0.5
|
||||
anchor_offset_y: 0.5
|
||||
strides: 8
|
||||
strides: 16
|
||||
strides: 16
|
||||
strides: 16
|
||||
aspect_ratios: 1.0
|
||||
fixed_anchor_size: true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Decodes the detection tensors generated by the TensorFlow Lite model, based on
|
||||
# the SSD anchors and the specification in the options, into a vector of
|
||||
# detections. Each detection describes a detected object.
|
||||
node {
|
||||
calculator: "TfLiteTensorsToDetectionsCalculator"
|
||||
input_stream: "TENSORS:detection_tensors"
|
||||
input_side_packet: "ANCHORS:anchors"
|
||||
output_stream: "DETECTIONS:detections"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.TfLiteTensorsToDetectionsCalculatorOptions] {
|
||||
num_classes: 1
|
||||
num_boxes: 896
|
||||
num_coords: 16
|
||||
box_coord_offset: 0
|
||||
keypoint_coord_offset: 4
|
||||
num_keypoints: 6
|
||||
num_values_per_keypoint: 2
|
||||
sigmoid_score: true
|
||||
score_clipping_thresh: 100.0
|
||||
reverse_output_order: true
|
||||
x_scale: 128.0
|
||||
y_scale: 128.0
|
||||
h_scale: 128.0
|
||||
w_scale: 128.0
|
||||
min_score_thresh: 0.75
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Performs non-max suppression to remove excessive detections.
|
||||
node {
|
||||
calculator: "NonMaxSuppressionCalculator"
|
||||
input_stream: "detections"
|
||||
output_stream: "filtered_detections"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.NonMaxSuppressionCalculatorOptions] {
|
||||
min_suppression_threshold: 0.3
|
||||
overlap_type: INTERSECTION_OVER_UNION
|
||||
algorithm: WEIGHTED
|
||||
return_empty_detections: true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Maps detection label IDs to the corresponding label text ("Face"). The label
|
||||
# map is provided in the label_map_path option.
|
||||
node {
|
||||
calculator: "DetectionLabelIdToTextCalculator"
|
||||
input_stream: "filtered_detections"
|
||||
output_stream: "labeled_detections"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.DetectionLabelIdToTextCalculatorOptions] {
|
||||
label_map_path: "mediapipe/models/face_detection_front_labelmap.txt"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Adjusts detection locations (already normalized to [0.f, 1.f]) on the
|
||||
# letterboxed image (after image transformation with the FIT scale mode) to the
|
||||
# corresponding locations on the same image with the letterbox removed (the
|
||||
# input image to the graph before image transformation).
|
||||
node {
|
||||
calculator: "DetectionLetterboxRemovalCalculator"
|
||||
input_stream: "DETECTIONS:labeled_detections"
|
||||
input_stream: "LETTERBOX_PADDING:letterbox_padding"
|
||||
output_stream: "DETECTIONS:output_detections"
|
||||
}
|
||||
|
||||
# Converts the detections to drawing primitives for annotation overlay.
|
||||
node {
|
||||
calculator: "DetectionsToRenderDataCalculator"
|
||||
input_stream: "DETECTIONS:output_detections"
|
||||
output_stream: "RENDER_DATA:render_data"
|
||||
node_options: {
|
||||
[type.googleapis.com/mediapipe.DetectionsToRenderDataCalculatorOptions] {
|
||||
thickness: 4.0
|
||||
color { r: 255 g: 0 b: 0 }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Draws annotations and overlays them on top of the input images.
|
||||
node {
|
||||
calculator: "AnnotationOverlayCalculator"
|
||||
input_stream: "INPUT_FRAME:input_video_cpu"
|
||||
input_stream: "render_data"
|
||||
output_stream: "OUTPUT_FRAME:output_video_cpu"
|
||||
}
|
||||
|
||||
# Transfers the annotated image from CPU back to GPU memory, to be sent out of
|
||||
# the graph.
|
||||
node: {
|
||||
calculator: "ImageFrameToGpuBufferCalculator"
|
||||
input_stream: "output_video_cpu"
|
||||
output_stream: "output_video"
|
||||
}
|
||||
```
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user