Project import generated by Copybara.

GitOrigin-RevId: 0517756260533d374df93679965ca662d0ec6943
This commit is contained in:
MediaPipe Team
2020-01-10 13:13:24 -08:00
committed by Hadon Nash
parent 38ee2603a7
commit ae6be10afe
192 changed files with 16410 additions and 592 deletions
+57 -5
View File
@@ -47,6 +47,13 @@ proto_library(
deps = ["//mediapipe/framework:calculator_proto"],
)
proto_library(
name = "packet_thinner_calculator_proto",
srcs = ["packet_thinner_calculator.proto"],
visibility = ["//visibility:public"],
deps = ["//mediapipe/framework:calculator_proto"],
)
proto_library(
name = "split_vector_calculator_proto",
srcs = ["split_vector_calculator.proto"],
@@ -102,6 +109,14 @@ mediapipe_cc_proto_library(
deps = [":packet_resampler_calculator_proto"],
)
mediapipe_cc_proto_library(
name = "packet_thinner_calculator_cc_proto",
srcs = ["packet_thinner_calculator.proto"],
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
visibility = ["//visibility:public"],
deps = [":packet_thinner_calculator_proto"],
)
mediapipe_cc_proto_library(
name = "split_vector_calculator_cc_proto",
srcs = ["split_vector_calculator.proto"],
@@ -284,7 +299,6 @@ cc_test(
srcs = ["concatenate_vector_calculator_test.cc"],
deps = [
":concatenate_vector_calculator",
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework:timestamp",
@@ -451,6 +465,37 @@ cc_test(
],
)
cc_library(
name = "packet_thinner_calculator",
srcs = ["packet_thinner_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/calculators/core:packet_thinner_calculator_cc_proto",
"//mediapipe/framework:calculator_context",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/formats:video_stream_header",
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:status",
],
alwayslink = 1,
)
cc_test(
name = "packet_thinner_calculator_test",
srcs = ["packet_thinner_calculator_test.cc"],
deps = [
":packet_thinner_calculator",
"//mediapipe/calculators/core:packet_thinner_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework/formats:video_stream_header",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:integral_types",
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "pass_through_calculator",
srcs = ["pass_through_calculator.cc"],
@@ -572,6 +617,7 @@ cc_test(
cc_library(
name = "packet_resampler_calculator",
srcs = ["packet_resampler_calculator.cc"],
hdrs = ["packet_resampler_calculator.h"],
visibility = [
"//visibility:public",
],
@@ -595,17 +641,17 @@ cc_library(
cc_test(
name = "packet_resampler_calculator_test",
timeout = "short",
srcs = ["packet_resampler_calculator_test.cc"],
srcs = [
"packet_resampler_calculator_test.cc",
],
deps = [
":packet_resampler_calculator",
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework:timestamp",
"//mediapipe/framework/formats:video_stream_header",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:status",
"@com_google_absl//absl/strings",
],
)
@@ -698,7 +744,13 @@ cc_library(
"//mediapipe/util:resource_util",
"@org_tensorflow//tensorflow/lite:framework",
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
],
] + select({
"//mediapipe/gpu:disable_gpu": [],
"//mediapipe:ios": [],
"//conditions:default": [
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
],
}),
alwayslink = 1,
)
@@ -12,23 +12,9 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <cstdlib>
#include <memory>
#include <string>
#include "mediapipe/calculators/core/packet_resampler_calculator.h"
#include "absl/strings/str_cat.h"
#include "mediapipe/calculators/core/packet_resampler_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/collection_item_id.h"
#include "mediapipe/framework/deps/mathutil.h"
#include "mediapipe/framework/deps/random_base.h"
#include "mediapipe/framework/formats/video_stream_header.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 "mediapipe/framework/port/status_macros.h"
#include "mediapipe/framework/tool/options_util.h"
#include <memory>
namespace {
@@ -45,120 +31,7 @@ std::unique_ptr<RandomBase> CreateSecureRandom(const std::string& seed) {
namespace mediapipe {
// This calculator is used to normalize the frequency of the packets
// out of a stream. Given a desired frame rate, packets are going to be
// removed or added to achieve it.
//
// The jitter feature is disabled by default. To enable it, you need to
// implement CreateSecureRandom(const std::string&).
//
// The data stream may be either specified as the only stream (by index)
// or as the stream with tag "DATA".
//
// The input and output streams may be accompanied by a VIDEO_HEADER
// stream. This stream includes a VideoHeader at Timestamp::PreStream().
// The input VideoHeader on the VIDEO_HEADER stream will always be updated
// with the resampler frame rate no matter what the options value for
// output_header is before being output on the output VIDEO_HEADER stream.
// If the input VideoHeader is not available, then only the frame rate
// value will be set in the output.
//
// Related:
// packet_downsampler_calculator.cc: skips packets regardless of timestamps.
class PacketResamplerCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc);
::mediapipe::Status Open(CalculatorContext* cc) override;
::mediapipe::Status Close(CalculatorContext* cc) override;
::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);
// Logic for Process() when jitter_ == 0.0.
::mediapipe::Status ProcessWithoutJitter(CalculatorContext* cc);
// Given the current count of periods that have passed, this returns
// the next valid timestamp of the middle point of the next period:
// if count is 0, it returns the first_timestamp_.
// if count is 1, it returns the first_timestamp_ + period (corresponding
// to the first tick using exact fps)
// e.g. for frame_rate=30 and first_timestamp_=0:
// 0: 0
// 1: 33333
// 2: 66667
// 3: 100000
//
// Can only be used if jitter_ equals zero.
Timestamp PeriodIndexToTimestamp(int64 index) const;
// Given a Timestamp, finds the closest sync Timestamp based on
// first_timestamp_ and the desired fps.
//
// Can only be used if jitter_ equals zero.
int64 TimestampToPeriodIndex(Timestamp timestamp) const;
// Outputs a packet if it is in range (start_time_, end_time_).
void OutputWithinLimits(CalculatorContext* cc, const Packet& packet) const;
// The timestamp of the first packet received.
Timestamp first_timestamp_;
// Number of frames per second (desired output frequency).
double frame_rate_;
// Inverse of frame_rate_.
int64 frame_time_usec_;
// Number of periods that have passed (= #packets sent to the output).
//
// Can only be used if jitter_ equals zero.
int64 period_count_;
// The last packet that was received.
Packet last_packet_;
VideoHeader video_header_;
// The "DATA" input stream.
CollectionItemId input_data_id_;
// The "DATA" output stream.
CollectionItemId output_data_id_;
// Indicator whether to flush last packet even if its timestamp is greater
// than the final stream timestamp. Set to false when jitter_ is non-zero.
bool flush_last_packet_;
// Jitter-related variables.
std::unique_ptr<RandomBase> random_;
double jitter_ = 0.0;
Timestamp next_output_timestamp_;
// If specified, output timestamps are aligned with base_timestamp.
// Otherwise, they are aligned with the first input timestamp.
Timestamp base_timestamp_;
// If specified, only outputs at/after start_time are included.
Timestamp start_time_;
// If specified, only outputs before end_time are included.
Timestamp end_time_;
// If set, the output timestamps nearest to start_time and end_time
// are included in the output, even if the nearest timestamp is not
// between start_time and end_time.
bool round_limits_;
};
REGISTER_CALCULATOR(PacketResamplerCalculator);
namespace {
// Returns a TimestampDiff (assuming microseconds) corresponding to the
// given time in seconds.
@@ -279,7 +152,10 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
"SecureRandom is not available. With \"jitter\" specified, "
"PacketResamplerCalculator processing cannot proceed.");
}
packet_reservoir_random_ = CreateSecureRandom(seed);
}
packet_reservoir_ =
std::make_unique<PacketReservoir>(packet_reservoir_random_.get());
return ::mediapipe::OkStatus();
}
@@ -294,6 +170,14 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
}
}
if (jitter_ != 0.0 && random_ != nullptr) {
// Packet reservior is used to make sure there's an output for every period,
// e.g. partial period at the end of the stream.
if (packet_reservoir_->IsEnabled() &&
(first_timestamp_ == Timestamp::Unset() ||
(cc->InputTimestamp() - next_output_timestamp_min_).Value() >= 0)) {
auto curr_packet = cc->Inputs().Get(input_data_id_).Value();
packet_reservoir_->AddSample(curr_packet);
}
MP_RETURN_IF_ERROR(ProcessWithJitter(cc));
} else {
MP_RETURN_IF_ERROR(ProcessWithoutJitter(cc));
@@ -303,11 +187,14 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
}
void PacketResamplerCalculator::InitializeNextOutputTimestampWithJitter() {
next_output_timestamp_min_ = first_timestamp_;
next_output_timestamp_ =
first_timestamp_ + frame_time_usec_ * random_->RandFloat();
}
void PacketResamplerCalculator::UpdateNextOutputTimestampWithJitter() {
packet_reservoir_->Clear();
packet_reservoir_->Disable();
next_output_timestamp_ +=
frame_time_usec_ *
((1.0 - jitter_) + 2.0 * jitter_ * random_->RandFloat());
@@ -339,10 +226,10 @@ void PacketResamplerCalculator::UpdateNextOutputTimestampWithJitter() {
while (true) {
const int64 last_diff =
(next_output_timestamp_ - last_packet_.Timestamp()).Value();
RET_CHECK_GT(last_diff, 0.0);
RET_CHECK_GT(last_diff, 0);
const int64 curr_diff =
(next_output_timestamp_ - cc->InputTimestamp()).Value();
if (curr_diff > 0.0) {
if (curr_diff > 0) {
break;
}
OutputWithinLimits(cc, (std::abs(curr_diff) > last_diff
@@ -431,6 +318,9 @@ void PacketResamplerCalculator::UpdateNextOutputTimestampWithJitter() {
OutputWithinLimits(cc,
last_packet_.At(PeriodIndexToTimestamp(period_count_)));
}
if (!packet_reservoir_->IsEmpty()) {
OutputWithinLimits(cc, packet_reservoir_->GetSample());
}
return ::mediapipe::OkStatus();
}
@@ -0,0 +1,168 @@
#ifndef MEDIAPIPE_CALCULATORS_CORE_PACKET_RESAMPLER_CALCULATOR_H_
#define MEDIAPIPE_CALCULATORS_CORE_PACKET_RESAMPLER_CALCULATOR_H_
#include <cstdlib>
#include <memory>
#include <string>
#include "absl/strings/str_cat.h"
#include "mediapipe/calculators/core/packet_resampler_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/collection_item_id.h"
#include "mediapipe/framework/deps/mathutil.h"
#include "mediapipe/framework/deps/random_base.h"
#include "mediapipe/framework/formats/video_stream_header.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 "mediapipe/framework/port/status_macros.h"
#include "mediapipe/framework/tool/options_util.h"
namespace mediapipe {
class PacketReservoir {
public:
PacketReservoir(RandomBase* rng) : rng_(rng) {}
// Replace candidate with current packet with 1/count_ probability.
void AddSample(Packet sample) {
if (rng_->UnbiasedUniform(++count_) == 0) {
reservoir_ = sample;
}
}
bool IsEnabled() { return rng_ && enabled_; }
void Disable() {
if (enabled_) enabled_ = false;
}
void Clear() { count_ = 0; }
bool IsEmpty() { return count_ == 0; }
Packet GetSample() { return reservoir_; }
private:
RandomBase* rng_;
bool enabled_ = true;
int32 count_ = 0;
Packet reservoir_;
};
// This calculator is used to normalize the frequency of the packets
// out of a stream. Given a desired frame rate, packets are going to be
// removed or added to achieve it.
//
// The jitter feature is disabled by default. To enable it, you need to
// implement CreateSecureRandom(const std::string&).
//
// The data stream may be either specified as the only stream (by index)
// or as the stream with tag "DATA".
//
// The input and output streams may be accompanied by a VIDEO_HEADER
// stream. This stream includes a VideoHeader at Timestamp::PreStream().
// The input VideoHeader on the VIDEO_HEADER stream will always be updated
// with the resampler frame rate no matter what the options value for
// output_header is before being output on the output VIDEO_HEADER stream.
// If the input VideoHeader is not available, then only the frame rate
// value will be set in the output.
//
// Related:
// packet_downsampler_calculator.cc: skips packets regardless of timestamps.
class PacketResamplerCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc);
::mediapipe::Status Open(CalculatorContext* cc) override;
::mediapipe::Status Close(CalculatorContext* cc) override;
::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);
// Logic for Process() when jitter_ == 0.0.
::mediapipe::Status ProcessWithoutJitter(CalculatorContext* cc);
// Given the current count of periods that have passed, this returns
// the next valid timestamp of the middle point of the next period:
// if count is 0, it returns the first_timestamp_.
// if count is 1, it returns the first_timestamp_ + period (corresponding
// to the first tick using exact fps)
// e.g. for frame_rate=30 and first_timestamp_=0:
// 0: 0
// 1: 33333
// 2: 66667
// 3: 100000
//
// Can only be used if jitter_ equals zero.
Timestamp PeriodIndexToTimestamp(int64 index) const;
// Given a Timestamp, finds the closest sync Timestamp based on
// first_timestamp_ and the desired fps.
//
// Can only be used if jitter_ equals zero.
int64 TimestampToPeriodIndex(Timestamp timestamp) const;
// Outputs a packet if it is in range (start_time_, end_time_).
void OutputWithinLimits(CalculatorContext* cc, const Packet& packet) const;
// The timestamp of the first packet received.
Timestamp first_timestamp_;
// Number of frames per second (desired output frequency).
double frame_rate_;
// Inverse of frame_rate_.
int64 frame_time_usec_;
// Number of periods that have passed (= #packets sent to the output).
//
// Can only be used if jitter_ equals zero.
int64 period_count_;
// The last packet that was received.
Packet last_packet_;
VideoHeader video_header_;
// The "DATA" input stream.
CollectionItemId input_data_id_;
// The "DATA" output stream.
CollectionItemId output_data_id_;
// Indicator whether to flush last packet even if its timestamp is greater
// than the final stream timestamp. Set to false when jitter_ is non-zero.
bool flush_last_packet_;
// Jitter-related variables.
std::unique_ptr<RandomBase> random_;
double jitter_ = 0.0;
Timestamp next_output_timestamp_;
Timestamp next_output_timestamp_min_;
// If specified, output timestamps are aligned with base_timestamp.
// Otherwise, they are aligned with the first input timestamp.
Timestamp base_timestamp_;
// If specified, only outputs at/after start_time are included.
Timestamp start_time_;
// If specified, only outputs before end_time are included.
Timestamp end_time_;
// If set, the output timestamps nearest to start_time and end_time
// are included in the output, even if the nearest timestamp is not
// between start_time and end_time.W
bool round_limits_;
// packet reservior used for sampling random packet out of partial
// period when jitter is enabled
std::unique_ptr<PacketReservoir> packet_reservoir_;
// random number generator used in packet_reservior_.
std::unique_ptr<RandomBase> packet_reservoir_random_;
};
} // namespace mediapipe
#endif // MEDIAPIPE_CALCULATORS_CORE_PACKET_RESAMPLER_CALCULATOR_H_
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include "mediapipe/calculators/core/packet_resampler_calculator.h"
#include <memory>
#include <string>
#include <vector>
@@ -29,7 +31,6 @@
namespace mediapipe {
namespace {
// A simple version of CalculatorRunner with built-in convenience
// methods for setting inputs from a vector and checking outputs
// against expected outputs (both timestamps and contents).
@@ -0,0 +1,304 @@
// 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.
//
// Declaration of PacketThinnerCalculator.
#include <cmath> // for ceil
#include <memory>
#include "mediapipe/calculators/core/packet_thinner_calculator.pb.h"
#include "mediapipe/framework/calculator_context.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/video_stream_header.h"
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/status.h"
namespace mediapipe {
namespace {
const double kTimebaseUs = 1000000; // Microseconds.
const char* const kPeriodTag = "PERIOD";
} // namespace
// This calculator is used to thin an input stream of Packets.
// An example application would be to sample decoded frames of video
// at a coarser temporal resolution. Unless otherwise stated, all
// timestamps are in units of microseconds.
//
// Thinning can be accomplished in one of two ways:
// 1) asynchronous thinning (known below as async):
// Algorithm does not rely on a master clock and is parameterized only
// by a single option -- the period. Once a packet is emitted, the
// thinner will discard subsequent packets for the duration of the period
// [Analogous to a refractory period during which packet emission is
// suppressed.]
// Packets arriving before start_time are discarded, as are packets
// arriving at or after end_time.
// 2) synchronous thinning (known below as sync):
// There are two variants of this algorithm, both parameterized by a
// start_time and a period. As in (1), packets arriving before start_time
// or at/after end_time are discarded. Otherwise, at most one packet is
// emitted during a period, centered at timestamps generated by the
// expression:
// start_time + i * period [where i is a non-negative integer]
// During each period, the packet closest to the generated timestamp is
// emitted (latest in the case of ties). In the first variant
// (sync_output_timestamps = true), the emitted packet is output at the
// generated timestamp. In the second variant, the packet is output at
// its original timestamp. Both variants emit exactly the same packets,
// but at different timestamps.
//
// Thinning period can be provided in the calculator options or via a
// side packet with the tag "PERIOD".
//
// Example config:
// node {
// calculator: "PacketThinnerCalculator"
// input_stream: "signal"
// output_stream: "output"
// options {
// [mediapipe.PacketThinnerCalculatorOptions.ext] {
// thinner_type: SYNC
// period: 10
// sync_output_timestamps: true
// update_frame_rate: false
// }
// }
// }
class PacketThinnerCalculator : public CalculatorBase {
public:
PacketThinnerCalculator() {}
~PacketThinnerCalculator() override {}
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).SetAny();
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
if (cc->InputSidePackets().HasTag(kPeriodTag)) {
cc->InputSidePackets().Tag(kPeriodTag).Set<int64>();
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) override;
::mediapipe::Status Close(CalculatorContext* cc) override;
::mediapipe::Status Process(CalculatorContext* cc) override {
if (cc->InputTimestamp() < start_time_) {
return ::mediapipe::OkStatus(); // Drop packets before start_time_.
} else if (cc->InputTimestamp() >= end_time_) {
if (!cc->Outputs().Index(0).IsClosed()) {
cc->Outputs()
.Index(0)
.Close(); // No more Packets will be output after end_time_.
}
return ::mediapipe::OkStatus();
} else {
return thinner_type_ == PacketThinnerCalculatorOptions::ASYNC
? AsyncThinnerProcess(cc)
: SyncThinnerProcess(cc);
}
}
private:
// Implementation of ASYNC and SYNC versions of thinner algorithm.
::mediapipe::Status AsyncThinnerProcess(CalculatorContext* cc);
::mediapipe::Status SyncThinnerProcess(CalculatorContext* cc);
// Cached option.
PacketThinnerCalculatorOptions::ThinnerType thinner_type_;
// Given a Timestamp, finds the closest sync Timestamp
// based on start_time_ and period_. This can be earlier or
// later than given Timestamp, but is guaranteed to be within
// half a period_.
Timestamp NearestSyncTimestamp(Timestamp now) const;
// Cached option used by both async and sync thinners.
TimestampDiff period_; // Interval during which only one packet is emitted.
Timestamp start_time_; // Cached option - default Timestamp::Min()
Timestamp end_time_; // Cached option - default Timestamp::Max()
// Only used by async thinner:
Timestamp next_valid_timestamp_; // Suppress packets until this timestamp.
// Only used by sync thinner:
Packet saved_packet_; // Best packet not yet emitted.
bool sync_output_timestamps_; // Cached option.
};
REGISTER_CALCULATOR(PacketThinnerCalculator);
namespace {
TimestampDiff abs(TimestampDiff t) { return t < 0 ? -t : t; }
} // namespace
::mediapipe::Status PacketThinnerCalculator::Open(CalculatorContext* cc) {
auto& options = cc->Options<PacketThinnerCalculatorOptions>();
thinner_type_ = options.thinner_type();
// This check enables us to assume only two thinner types exist in Process()
CHECK(thinner_type_ == PacketThinnerCalculatorOptions::ASYNC ||
thinner_type_ == PacketThinnerCalculatorOptions::SYNC)
<< "Unsupported thinner type.";
if (thinner_type_ == PacketThinnerCalculatorOptions::ASYNC) {
// ASYNC thinner outputs packets with the same timestamp as their input so
// its safe to SetOffset(0). SYNC thinner manipulates timestamps of its
// output so we don't do this for that case.
cc->SetOffset(0);
}
if (cc->InputSidePackets().HasTag(kPeriodTag)) {
period_ =
TimestampDiff(cc->InputSidePackets().Tag(kPeriodTag).Get<int64>());
} else {
period_ = TimestampDiff(options.period());
}
CHECK_LT(TimestampDiff(0), period_) << "Specified period must be positive.";
if (options.has_start_time()) {
start_time_ = Timestamp(options.start_time());
} else if (thinner_type_ == PacketThinnerCalculatorOptions::ASYNC) {
start_time_ = Timestamp::Min();
} else {
start_time_ = Timestamp(0);
}
end_time_ =
options.has_end_time() ? Timestamp(options.end_time()) : Timestamp::Max();
CHECK_LT(start_time_, end_time_)
<< "Invalid PacketThinner: start_time must be earlier than end_time";
sync_output_timestamps_ = options.sync_output_timestamps();
next_valid_timestamp_ = start_time_;
// Drop packets until this time.
cc->Outputs().Index(0).SetNextTimestampBound(start_time_);
if (!cc->Inputs().Index(0).Header().IsEmpty()) {
if (options.update_frame_rate()) {
const VideoHeader& video_header =
cc->Inputs().Index(0).Header().Get<VideoHeader>();
double new_frame_rate;
if (thinner_type_ == PacketThinnerCalculatorOptions::ASYNC) {
new_frame_rate =
video_header.frame_rate /
ceil(video_header.frame_rate * options.period() / kTimebaseUs);
} else {
const double sampling_rate = kTimebaseUs / options.period();
new_frame_rate = video_header.frame_rate < sampling_rate
? video_header.frame_rate
: sampling_rate;
}
std::unique_ptr<VideoHeader> header(new VideoHeader);
header->format = video_header.format;
header->width = video_header.width;
header->height = video_header.height;
header->frame_rate = new_frame_rate;
cc->Outputs().Index(0).SetHeader(Adopt(header.release()));
} else {
cc->Outputs().Index(0).SetHeader(cc->Inputs().Index(0).Header());
}
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status PacketThinnerCalculator::Close(CalculatorContext* cc) {
// Emit any saved packets before quitting.
if (!saved_packet_.IsEmpty()) {
// Only sync thinner should have saved packets.
CHECK_EQ(PacketThinnerCalculatorOptions::SYNC, thinner_type_);
if (sync_output_timestamps_) {
cc->Outputs().Index(0).AddPacket(
saved_packet_.At(NearestSyncTimestamp(saved_packet_.Timestamp())));
} else {
cc->Outputs().Index(0).AddPacket(saved_packet_);
}
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status PacketThinnerCalculator::AsyncThinnerProcess(
CalculatorContext* cc) {
if (cc->InputTimestamp() >= next_valid_timestamp_) {
cc->Outputs().Index(0).AddPacket(
cc->Inputs().Index(0).Value()); // Emit current packet.
next_valid_timestamp_ = cc->InputTimestamp() + period_;
// Guaranteed not to emit packets seen during refractory period.
cc->Outputs().Index(0).SetNextTimestampBound(next_valid_timestamp_);
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status PacketThinnerCalculator::SyncThinnerProcess(
CalculatorContext* cc) {
if (saved_packet_.IsEmpty()) {
// If no packet has been saved, store the current packet.
saved_packet_ = cc->Inputs().Index(0).Value();
cc->Outputs().Index(0).SetNextTimestampBound(
sync_output_timestamps_ ? NearestSyncTimestamp(cc->InputTimestamp())
: cc->InputTimestamp());
} else {
// Saved packet exists -- update or emit.
const Timestamp saved = saved_packet_.Timestamp();
const Timestamp saved_sync = NearestSyncTimestamp(saved);
const Timestamp now = cc->InputTimestamp();
const Timestamp now_sync = NearestSyncTimestamp(now);
CHECK_LE(saved_sync, now_sync);
if (saved_sync == now_sync) {
// Saved Packet is in same interval as current packet.
// Replace saved packet with current if it is at least as
// central as the saved packet wrt temporal interval.
// [We break ties in favor of fresher packets]
if (abs(now - now_sync) <= abs(saved - saved_sync)) {
saved_packet_ = cc->Inputs().Index(0).Value();
}
} else {
// Saved packet is the best packet from earlier interval: emit!
if (sync_output_timestamps_) {
cc->Outputs().Index(0).AddPacket(saved_packet_.At(saved_sync));
cc->Outputs().Index(0).SetNextTimestampBound(now_sync);
} else {
cc->Outputs().Index(0).AddPacket(saved_packet_);
cc->Outputs().Index(0).SetNextTimestampBound(now);
}
// Current packet is the first one we've seen from new interval -- save!
saved_packet_ = cc->Inputs().Index(0).Value();
}
}
return ::mediapipe::OkStatus();
}
Timestamp PacketThinnerCalculator::NearestSyncTimestamp(Timestamp now) const {
CHECK_NE(start_time_, Timestamp::Unset())
<< "Method only valid for sync thinner calculator.";
// Computation is done using int64 arithmetic. No easy way to avoid
// since Timestamps don't support div and multiply.
const int64 now64 = now.Value();
const int64 start64 = start_time_.Value();
const int64 period64 = period_.Value();
CHECK_LE(0, period64);
// Round now64 to its closest interval (units of period64).
int64 sync64 =
(now64 - start64 + period64 / 2) / period64 * period64 + start64;
CHECK_LE(abs(now64 - sync64), period64 / 2)
<< "start64: " << start64 << "; now64: " << now64
<< "; sync64: " << sync64;
return Timestamp(sync64);
}
} // namespace mediapipe
@@ -0,0 +1,66 @@
// 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.
syntax = "proto2";
package mediapipe;
import "mediapipe/framework/calculator.proto";
message PacketThinnerCalculatorOptions {
extend CalculatorOptions {
optional PacketThinnerCalculatorOptions ext = 288533508;
}
enum ThinnerType {
ASYNC = 1; // Asynchronous thinner, described below [default].
SYNC = 2; // Synchronous thinner, also described below.
}
optional ThinnerType thinner_type = 1 [default = ASYNC];
// The period (in microsecond) specifies the temporal interval during which
// only a single packet is emitted in the output stream. Has subtly different
// semantics depending on the thinner type, as follows.
//
// Async thinner: this option is a refractory period -- once a packet is
// emitted, we guarantee that no packets will be emitted for period ticks.
//
// Sync thinner: the period specifies a temporal interval during which
// only one packet is emitted. The emitted packet is guaranteed to be
// the one closest to the center of the temporal interval (no guarantee on
// how ties are broken). More specifically,
// intervals are centered at start_time + i * period
// (for non-negative integers i).
// Thus, each interval extends period/2 ticks before and after its center.
// Additionally, in the sync thinner any packets earlier than start_time
// are discarded and the thinner calls Close() once timestamp equals or
// exceeds end_time.
optional int64 period = 2 [default = 1];
// Packets before start_time and at/after end_time are discarded.
// Additionally, for a sync thinner, start time specifies the center of
// time invervals as described above and therefore should be set explicitly.
optional int64 start_time = 3; // If not specified, set to 0 for SYNC type,
// and set to Timestamp::Min() for ASYNC type.
optional int64 end_time = 4; // Set to Timestamp::Max() if not specified.
// Whether the timestamps of packets emitted by sync thinner should
// correspond to the center of their corresponding temporal interval.
// If false, packets emitted using original timestamp (as in async thinner).
optional bool sync_output_timestamps = 5 [default = true];
// If true, update the frame rate in the header, if it's available, to an
// estimated frame rate due to the sampling.
optional bool update_frame_rate = 6 [default = false];
}
@@ -0,0 +1,357 @@
// 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 <vector>
#include "absl/strings/str_cat.h"
#include "mediapipe/calculators/core/packet_thinner_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_runner.h"
#include "mediapipe/framework/formats/video_stream_header.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/status_matchers.h"
namespace mediapipe {
namespace {
// A simple version of CalculatorRunner with built-in convenience methods for
// setting inputs from a vector and checking outputs against a vector of
// expected outputs.
class SimpleRunner : public CalculatorRunner {
public:
explicit SimpleRunner(const CalculatorOptions& options)
: CalculatorRunner("PacketThinnerCalculator", options) {
SetNumInputs(1);
SetNumOutputs(1);
SetNumInputSidePackets(0);
}
explicit SimpleRunner(const CalculatorGraphConfig::Node& node)
: CalculatorRunner(node) {}
void SetInput(const std::vector<int>& timestamp_list) {
MutableInputs()->Index(0).packets.clear();
for (const int ts : timestamp_list) {
MutableInputs()->Index(0).packets.push_back(
MakePacket<std::string>(absl::StrCat("Frame #", ts))
.At(Timestamp(ts)));
}
}
void SetFrameRate(const double frame_rate) {
auto video_header = absl::make_unique<VideoHeader>();
video_header->frame_rate = frame_rate;
MutableInputs()->Index(0).header = Adopt(video_header.release());
}
std::vector<int64> GetOutputTimestamps() const {
std::vector<int64> timestamps;
for (const Packet& packet : Outputs().Index(0).packets) {
timestamps.emplace_back(packet.Timestamp().Value());
}
return timestamps;
}
double GetFrameRate() const {
CHECK(!Outputs().Index(0).header.IsEmpty());
return Outputs().Index(0).header.Get<VideoHeader>().frame_rate;
}
};
// Check that thinner respects start_time and end_time options.
// We only test with one thinner because the logic for start & end time
// handling is shared across both types of thinner in Process().
TEST(PacketThinnerCalculatorTest, StartAndEndTimeTest) {
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
extension->set_period(5);
extension->set_start_time(4);
extension->set_end_time(12);
SimpleRunner runner(options);
runner.SetInput({2, 3, 5, 7, 11, 13, 17, 19, 23, 29});
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {5, 11};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
}
TEST(PacketThinnerCalculatorTest, AsyncUniformStreamThinningTest) {
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
extension->set_period(5);
SimpleRunner runner(options);
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {2, 8, 14};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
}
TEST(PacketThinnerCalculatorTest, ASyncUniformStreamThinningTestBySidePacket) {
// Note: sync runner but outputting *original* timestamps.
CalculatorGraphConfig::Node node;
node.set_calculator("PacketThinnerCalculator");
node.add_input_side_packet("PERIOD:period");
node.add_input_stream("input_stream");
node.add_output_stream("output_stream");
auto* extension = node.mutable_options()->MutableExtension(
PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
extension->set_start_time(0);
extension->set_sync_output_timestamps(false);
SimpleRunner runner(node);
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
runner.MutableSidePackets()->Tag("PERIOD") = MakePacket<int64>(5);
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {2, 8, 14};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
}
TEST(PacketThinnerCalculatorTest, SyncUniformStreamThinningTest1) {
// Note: sync runner but outputting *original* timestamps.
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
extension->set_start_time(0);
extension->set_period(5);
extension->set_sync_output_timestamps(false);
SimpleRunner runner(options);
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {2, 6, 10, 14};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
}
TEST(PacketThinnerCalculatorTest, SyncUniformStreamThinningTestBySidePacket1) {
// Note: sync runner but outputting *original* timestamps.
CalculatorGraphConfig::Node node;
node.set_calculator("PacketThinnerCalculator");
node.add_input_side_packet("PERIOD:period");
node.add_input_stream("input_stream");
node.add_output_stream("output_stream");
auto* extension = node.mutable_options()->MutableExtension(
PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
extension->set_start_time(0);
extension->set_sync_output_timestamps(false);
SimpleRunner runner(node);
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
runner.MutableSidePackets()->Tag("PERIOD") = MakePacket<int64>(5);
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {2, 6, 10, 14};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
}
TEST(PacketThinnerCalculatorTest, SyncUniformStreamThinningTest2) {
// Same test but now with synced timestamps.
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
extension->set_start_time(0);
extension->set_period(5);
extension->set_sync_output_timestamps(true);
SimpleRunner runner(options);
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {0, 5, 10, 15};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
}
// Test: Given a stream with timestamps corresponding to first ten prime numbers
// and period of 5, confirm whether timestamps of thinner stream matches
// expectations.
TEST(PacketThinnerCalculatorTest, PrimeStreamThinningTest1) {
// ASYNC thinner.
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
extension->set_period(5);
SimpleRunner runner(options);
runner.SetInput({2, 3, 5, 7, 11, 13, 17, 19, 23, 29});
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {2, 7, 13, 19, 29};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
}
TEST(PacketThinnerCalculatorTest, PrimeStreamThinningTest2) {
// SYNC with original timestamps.
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
extension->set_start_time(0);
extension->set_period(5);
extension->set_sync_output_timestamps(false);
SimpleRunner runner(options);
runner.SetInput({2, 3, 5, 7, 11, 13, 17, 19, 23, 29});
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {2, 5, 11, 17, 19, 23, 29};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
}
// Confirm that Calculator correctly handles boundary cases.
TEST(PacketThinnerCalculatorTest, BoundaryTimestampTest1) {
// Odd period, negative start_time
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
extension->set_start_time(-10);
extension->set_period(5);
extension->set_sync_output_timestamps(true);
SimpleRunner runner(options);
// Two timestamps falling on either side of a period boundary.
runner.SetInput({2, 3});
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {0, 5};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
}
TEST(PacketThinnerCalculatorTest, BoundaryTimestampTest2) {
// Even period, negative start_time, negative packet timestamps.
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
extension->set_start_time(-144);
extension->set_period(6);
extension->set_sync_output_timestamps(true);
SimpleRunner runner(options);
// Two timestamps falling on either side of a period boundary.
runner.SetInput({-4, -3, 8, 9});
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {-6, 0, 6, 12};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
}
TEST(PacketThinnerCalculatorTest, FrameRateTest1) {
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
extension->set_period(5);
extension->set_update_frame_rate(true);
SimpleRunner runner(options);
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
runner.SetFrameRate(1000000.0 / 2);
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {2, 8, 14};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
// The true sampling period is 6.
EXPECT_DOUBLE_EQ(1000000.0 / 6, runner.GetFrameRate());
}
TEST(PacketThinnerCalculatorTest, FrameRateTest2) {
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
extension->set_period(5);
extension->set_update_frame_rate(true);
SimpleRunner runner(options);
runner.SetInput({8, 16, 24, 32, 40, 48, 56});
runner.SetFrameRate(1000000.0 / 8);
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {8, 16, 24, 32, 40, 48, 56};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
// The true sampling period is still 8.
EXPECT_DOUBLE_EQ(1000000.0 / 8, runner.GetFrameRate());
}
TEST(PacketThinnerCalculatorTest, FrameRateTest3) {
// Note: sync runner but outputting *original* timestamps.
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
extension->set_start_time(0);
extension->set_period(5);
extension->set_sync_output_timestamps(false);
extension->set_update_frame_rate(true);
SimpleRunner runner(options);
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
runner.SetFrameRate(1000000.0 / 2);
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {2, 6, 10, 14};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
// The true (long-run) sampling period is 5.
EXPECT_DOUBLE_EQ(1000000.0 / 5, runner.GetFrameRate());
}
TEST(PacketThinnerCalculatorTest, FrameRateTest4) {
// Same test but now with synced timestamps.
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
extension->set_start_time(0);
extension->set_period(5);
extension->set_sync_output_timestamps(true);
extension->set_update_frame_rate(true);
SimpleRunner runner(options);
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
runner.SetFrameRate(1000000.0 / 2);
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {0, 5, 10, 15};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
// The true (long-run) sampling period is 5.
EXPECT_DOUBLE_EQ(1000000.0 / 5, runner.GetFrameRate());
}
TEST(PacketThinnerCalculatorTest, FrameRateTest5) {
CalculatorOptions options;
auto* extension =
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
extension->set_start_time(0);
extension->set_period(5);
extension->set_sync_output_timestamps(true);
extension->set_update_frame_rate(true);
SimpleRunner runner(options);
runner.SetInput({8, 16, 24, 32, 40, 48, 56});
runner.SetFrameRate(1000000.0 / 8);
MP_ASSERT_OK(runner.Run());
const std::vector<int64> expected_timestamps = {10, 15, 25, 30, 40, 50, 55};
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
// The true (long-run) sampling period is 8.
EXPECT_DOUBLE_EQ(1000000.0 / 8, runner.GetFrameRate());
}
} // namespace
} // namespace mediapipe
@@ -17,6 +17,7 @@
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/timestamp.h"
namespace mediapipe {
@@ -86,6 +87,7 @@ class PreviousLoopbackCalculator : public CalculatorBase {
main_ts_.pop_front();
}
}
auto& loop_out = cc->Outputs().Get(loop_out_id_);
while (!main_ts_.empty() && !loopback_packets_.empty()) {
Timestamp main_timestamp = main_ts_.front();
@@ -95,18 +97,31 @@ class PreviousLoopbackCalculator : public CalculatorBase {
if (previous_loopback.IsEmpty()) {
// TODO: SetCompleteTimestampBound would be more useful.
cc->Outputs()
.Get(loop_out_id_)
.SetNextTimestampBound(main_timestamp + 1);
loop_out.SetNextTimestampBound(main_timestamp + 1);
} else {
cc->Outputs().Get(loop_out_id_).AddPacket(std::move(previous_loopback));
loop_out.AddPacket(std::move(previous_loopback));
}
}
// In case of an empty loopback input, the next timestamp bound for
// loopback input is the loopback timestamp + 1. The next timestamp bound
// for output is set and the main_ts_ vector is truncated accordingly.
if (loopback_packet.IsEmpty() &&
loopback_packet.Timestamp() != Timestamp::Unstarted()) {
Timestamp loopback_bound =
loopback_packet.Timestamp().NextAllowedInStream();
while (!main_ts_.empty() && main_ts_.front() <= loopback_bound) {
main_ts_.pop_front();
}
if (main_ts_.empty()) {
loop_out.SetNextTimestampBound(loopback_bound.NextAllowedInStream());
}
}
if (!main_ts_.empty()) {
cc->Outputs().Get(loop_out_id_).SetNextTimestampBound(main_ts_.front());
loop_out.SetNextTimestampBound(main_ts_.front());
}
if (cc->Inputs().Get(main_id_).IsDone() && main_ts_.empty()) {
cc->Outputs().Get(loop_out_id_).Close();
loop_out.Close();
}
return ::mediapipe::OkStatus();
}
@@ -207,5 +207,64 @@ TEST(PreviousLoopbackCalculator, ClosesCorrectly) {
MP_EXPECT_OK(graph_.WaitUntilDone());
}
// Demonstrates that downstream calculators won't be blocked by
// always-empty-LOOP-stream.
TEST(PreviousLoopbackCalculator, EmptyLoopForever) {
std::vector<Packet> outputs;
CalculatorGraphConfig graph_config_ =
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: 'in'
node {
calculator: 'PreviousLoopbackCalculator'
input_stream: 'MAIN:in'
input_stream: 'LOOP:previous'
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", 0);
MP_EXPECT_OK(graph_.WaitUntilIdle());
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{0}));
for (int main_ts = 1; main_ts < 50; ++main_ts) {
send_packet("in", main_ts);
MP_EXPECT_OK(graph_.WaitUntilIdle());
std::vector<int64> ts_values = TimestampValues(outputs);
EXPECT_EQ(ts_values.size(), main_ts + 1);
for (int j = 0; j < main_ts; ++j) {
EXPECT_EQ(ts_values[j], j);
}
}
MP_EXPECT_OK(graph_.CloseAllInputStreams());
MP_EXPECT_OK(graph_.WaitUntilIdle());
MP_EXPECT_OK(graph_.WaitUntilDone());
}
} // anonymous namespace
} // namespace mediapipe
@@ -20,6 +20,10 @@
#include "mediapipe/framework/formats/rect.pb.h"
#include "tensorflow/lite/interpreter.h"
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
#endif // !MEDIAPIPE_DISABLE_GPU
namespace mediapipe {
// Example config:
@@ -36,14 +40,21 @@ namespace mediapipe {
// }
// }
// }
typedef SplitVectorCalculator<TfLiteTensor> SplitTfLiteTensorVectorCalculator;
typedef SplitVectorCalculator<TfLiteTensor, false>
SplitTfLiteTensorVectorCalculator;
REGISTER_CALCULATOR(SplitTfLiteTensorVectorCalculator);
typedef SplitVectorCalculator<::mediapipe::NormalizedLandmark>
typedef SplitVectorCalculator<::mediapipe::NormalizedLandmark, false>
SplitLandmarkVectorCalculator;
REGISTER_CALCULATOR(SplitLandmarkVectorCalculator);
typedef SplitVectorCalculator<::mediapipe::NormalizedRect>
typedef SplitVectorCalculator<::mediapipe::NormalizedRect, false>
SplitNormalizedRectVectorCalculator;
REGISTER_CALCULATOR(SplitNormalizedRectVectorCalculator);
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
typedef SplitVectorCalculator<::tflite::gpu::gl::GlBuffer, true>
MovableSplitGlBufferVectorCalculator;
REGISTER_CALCULATOR(MovableSplitGlBufferVectorCalculator);
#endif
} // namespace mediapipe
@@ -15,12 +15,14 @@
#ifndef MEDIAPIPE_CALCULATORS_CORE_SPLIT_VECTOR_CALCULATOR_H_
#define MEDIAPIPE_CALCULATORS_CORE_SPLIT_VECTOR_CALCULATOR_H_
#include <type_traits>
#include <vector>
#include "mediapipe/calculators/core/split_vector_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"
#include "mediapipe/util/resource_util.h"
#include "tensorflow/lite/error_reporter.h"
#include "tensorflow/lite/interpreter.h"
@@ -29,6 +31,20 @@
namespace mediapipe {
template <typename T>
using IsCopyable = std::enable_if_t<std::is_copy_constructible<T>::value, bool>;
template <typename T>
using IsNotCopyable =
std::enable_if_t<!std::is_copy_constructible<T>::value, bool>;
template <typename T>
using IsMovable = std::enable_if_t<std::is_move_constructible<T>::value, bool>;
template <typename T>
using IsNotMovable =
std::enable_if_t<!std::is_move_constructible<T>::value, bool>;
// Splits an input packet with std::vector<T> into multiple std::vector<T>
// output packets using the [begin, end) ranges specified in
// SplitVectorCalculatorOptions. If the option "element_only" is set to true,
@@ -39,7 +55,7 @@ namespace mediapipe {
// combined into one vector.
// To use this class for a particular type T, register a calculator using
// SplitVectorCalculator<T>.
template <typename T>
template <typename T, bool move_elements>
class SplitVectorCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
@@ -51,23 +67,16 @@ class SplitVectorCalculator : public CalculatorBase {
const auto& options =
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
if (!std::is_copy_constructible<T>::value || move_elements) {
// Ranges of elements shouldn't overlap when the vector contains
// non-copyable elements.
RET_CHECK_OK(checkRangesDontOverlap(options));
}
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.");
}
}
}
RET_CHECK_OK(checkRangesDontOverlap(options));
} else {
if (cc->Outputs().NumEntries() != options.ranges_size()) {
return ::mediapipe::InvalidArgumentError(
@@ -117,14 +126,26 @@ class SplitVectorCalculator : public CalculatorBase {
}
::mediapipe::Status Process(CalculatorContext* cc) override {
const auto& input = cc->Inputs().Index(0).Get<std::vector<T>>();
RET_CHECK_GE(input.size(), max_range_end_);
if (cc->Inputs().Index(0).IsEmpty()) return ::mediapipe::OkStatus();
if (move_elements) {
return ProcessMovableElements<T>(cc);
} else {
return ProcessCopyableElements<T>(cc);
}
}
template <typename U, IsCopyable<U> = true>
::mediapipe::Status ProcessCopyableElements(CalculatorContext* cc) {
// static_assert(std::is_copy_constructible<U>::value,
// "Cannot copy non-copyable elements");
const auto& input = cc->Inputs().Index(0).Get<std::vector<U>>();
RET_CHECK_GE(input.size(), max_range_end_);
if (combine_outputs_) {
auto output = absl::make_unique<std::vector<T>>();
auto output = absl::make_unique<std::vector<U>>();
output->reserve(total_elements_);
for (int i = 0; i < ranges_.size(); ++i) {
auto elements = absl::make_unique<std::vector<T>>(
auto elements = absl::make_unique<std::vector<U>>(
input.begin() + ranges_[i].first,
input.begin() + ranges_[i].second);
output->insert(output->end(), elements->begin(), elements->end());
@@ -134,7 +155,7 @@ class SplitVectorCalculator : public CalculatorBase {
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()));
MakePacket<U>(input[ranges_[i].first]).At(cc->InputTimestamp()));
}
} else {
for (int i = 0; i < ranges_.size(); ++i) {
@@ -149,7 +170,78 @@ class SplitVectorCalculator : public CalculatorBase {
return ::mediapipe::OkStatus();
}
template <typename U, IsNotCopyable<U> = true>
::mediapipe::Status ProcessCopyableElements(CalculatorContext* cc) {
return ::mediapipe::InternalError("Cannot copy non-copyable elements.");
}
template <typename U, IsMovable<U> = true>
::mediapipe::Status ProcessMovableElements(CalculatorContext* cc) {
::mediapipe::StatusOr<std::unique_ptr<std::vector<U>>> input_status =
cc->Inputs().Index(0).Value().Consume<std::vector<U>>();
if (!input_status.ok()) return input_status.status();
std::unique_ptr<std::vector<U>> input_vector =
std::move(input_status).ValueOrDie();
RET_CHECK_GE(input_vector->size(), max_range_end_);
if (combine_outputs_) {
auto output = absl::make_unique<std::vector<U>>();
output->reserve(total_elements_);
for (int i = 0; i < ranges_.size(); ++i) {
output->insert(
output->end(),
std::make_move_iterator(input_vector->begin() + ranges_[i].first),
std::make_move_iterator(input_vector->begin() + ranges_[i].second));
}
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<U>(std::move(input_vector->at(ranges_[i].first)))
.At(cc->InputTimestamp()));
}
} else {
for (int i = 0; i < ranges_.size(); ++i) {
auto output = absl::make_unique<std::vector<T>>();
output->insert(
output->end(),
std::make_move_iterator(input_vector->begin() + ranges_[i].first),
std::make_move_iterator(input_vector->begin() +
ranges_[i].second));
cc->Outputs().Index(i).Add(output.release(), cc->InputTimestamp());
}
}
}
return ::mediapipe::OkStatus();
}
template <typename U, IsNotMovable<U> = true>
::mediapipe::Status ProcessMovableElements(CalculatorContext* cc) {
return ::mediapipe::InternalError("Cannot move non-movable elements.");
}
private:
static ::mediapipe::Status checkRangesDontOverlap(
const ::mediapipe::SplitVectorCalculatorOptions& options) {
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.");
}
}
}
return ::mediapipe::OkStatus();
}
std::vector<std::pair<int32, int32>> ranges_;
int32 max_range_end_ = -1;
int32 total_elements_ = 0;
@@ -452,4 +452,243 @@ TEST_F(SplitTfLiteTensorVectorCalculatorTest,
ASSERT_FALSE(graph.Initialize(graph_config).ok());
}
typedef SplitVectorCalculator<std::unique_ptr<int>, true>
MovableSplitUniqueIntPtrCalculator;
REGISTER_CALCULATOR(MovableSplitUniqueIntPtrCalculator);
class MovableSplitUniqueIntPtrCalculatorTest : public ::testing::Test {
protected:
void ValidateVectorOutput(std::vector<Packet>& output_packets,
int expected_elements, int input_begin_index) {
ASSERT_EQ(1, output_packets.size());
const std::vector<std::unique_ptr<int>>& output_vec =
output_packets[0].Get<std::vector<std::unique_ptr<int>>>();
ASSERT_EQ(expected_elements, output_vec.size());
for (int i = 0; i < expected_elements; ++i) {
const int expected_value = input_begin_index + i;
const std::unique_ptr<int>& result = output_vec[i];
ASSERT_NE(result, nullptr);
ASSERT_EQ(expected_value, *result);
}
}
void ValidateElementOutput(std::vector<Packet>& output_packets,
int expected_value) {
ASSERT_EQ(1, output_packets.size());
const std::unique_ptr<int>& result =
output_packets[0].Get<std::unique_ptr<int>>();
ASSERT_NE(result, nullptr);
ASSERT_EQ(expected_value, *result);
}
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<std::unique_ptr<int>>& output_vector =
output_packets[0].Get<std::vector<std::unique_ptr<int>>>();
ASSERT_EQ(expected_elements, output_vector.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 std::unique_ptr<int>& result = output_vector[element_id];
ASSERT_NE(result, nullptr);
ASSERT_EQ(expected_value, *result);
++element_id;
}
}
}
};
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, InvalidOverlappingRangesTest) {
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
CalculatorGraphConfig graph_config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
R"(
input_stream: "input_vector"
node {
calculator: "MovableSplitUniqueIntPtrCalculator"
input_stream: "input_vector"
output_stream: "range_0"
options {
[mediapipe.SplitVectorCalculatorOptions.ext] {
ranges: { begin: 0 end: 3 }
ranges: { begin: 1 end: 4 }
}
}
}
)");
// Run the graph.
CalculatorGraph graph;
// The graph should fail running because there are overlapping ranges.
ASSERT_FALSE(graph.Initialize(graph_config).ok());
}
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, SmokeTest) {
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
CalculatorGraphConfig graph_config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
R"(
input_stream: "input_vector"
node {
calculator: "MovableSplitUniqueIntPtrCalculator"
input_stream: "input_vector"
output_stream: "range_0"
output_stream: "range_1"
output_stream: "range_2"
options {
[mediapipe.SplitVectorCalculatorOptions.ext] {
ranges: { begin: 0 end: 1 }
ranges: { begin: 1 end: 4 }
ranges: { begin: 4 end: 5 }
}
}
}
)");
std::vector<Packet> range_0_packets;
tool::AddVectorSink("range_0", &graph_config, &range_0_packets);
std::vector<Packet> range_1_packets;
tool::AddVectorSink("range_1", &graph_config, &range_1_packets);
std::vector<Packet> range_2_packets;
tool::AddVectorSink("range_2", &graph_config, &range_2_packets);
// Run the graph.
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(graph_config));
MP_ASSERT_OK(graph.StartRun({}));
// input_vector : {0, 1, 2, 3, 4, 5}
std::unique_ptr<std::vector<std::unique_ptr<int>>> 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_ASSERT_OK(graph.AddPacketToInputStream(
"input_vector", Adopt(input_vector.release()).At(Timestamp(1))));
MP_ASSERT_OK(graph.WaitUntilIdle());
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
ValidateVectorOutput(range_0_packets, /*expected_elements=*/1,
/*input_begin_index=*/0);
ValidateVectorOutput(range_1_packets, /*expected_elements=*/3,
/*input_begin_index=*/1);
ValidateVectorOutput(range_2_packets, /*expected_elements=*/1,
/*input_begin_index=*/4);
}
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, SmokeTestElementOnly) {
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
CalculatorGraphConfig graph_config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
R"(
input_stream: "input_vector"
node {
calculator: "MovableSplitUniqueIntPtrCalculator"
input_stream: "input_vector"
output_stream: "range_0"
output_stream: "range_1"
output_stream: "range_2"
options {
[mediapipe.SplitVectorCalculatorOptions.ext] {
ranges: { begin: 0 end: 1 }
ranges: { begin: 2 end: 3 }
ranges: { begin: 4 end: 5 }
element_only: true
}
}
}
)");
std::vector<Packet> range_0_packets;
tool::AddVectorSink("range_0", &graph_config, &range_0_packets);
std::vector<Packet> range_1_packets;
tool::AddVectorSink("range_1", &graph_config, &range_1_packets);
std::vector<Packet> range_2_packets;
tool::AddVectorSink("range_2", &graph_config, &range_2_packets);
// Run the graph.
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(graph_config));
MP_ASSERT_OK(graph.StartRun({}));
// input_vector : {0, 1, 2, 3, 4, 5}
std::unique_ptr<std::vector<std::unique_ptr<int>>> 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_ASSERT_OK(graph.AddPacketToInputStream(
"input_vector", Adopt(input_vector.release()).At(Timestamp(1))));
MP_ASSERT_OK(graph.WaitUntilIdle());
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
ValidateElementOutput(range_0_packets, /*expected_value=*/0);
ValidateElementOutput(range_1_packets, /*expected_value=*/2);
ValidateElementOutput(range_2_packets, /*expected_value=*/4);
}
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, SmokeTestCombiningOutputs) {
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
CalculatorGraphConfig graph_config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
R"(
input_stream: "input_vector"
node {
calculator: "MovableSplitUniqueIntPtrCalculator"
input_stream: "input_vector"
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({}));
// input_vector : {0, 1, 2, 3, 4, 5}
std::unique_ptr<std::vector<std::unique_ptr<int>>> 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_ASSERT_OK(graph.AddPacketToInputStream(
"input_vector", Adopt(input_vector.release()).At(Timestamp(1))));
MP_ASSERT_OK(graph.WaitUntilIdle());
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
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);
}
} // namespace mediapipe
+3 -1
View File
@@ -80,7 +80,9 @@ mediapipe_cc_proto_library(
name = "opencv_image_encoder_calculator_cc_proto",
srcs = ["opencv_image_encoder_calculator.proto"],
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
visibility = ["//visibility:public"],
visibility = [
"//visibility:public",
],
deps = [":opencv_image_encoder_calculator_proto"],
)
@@ -474,13 +474,20 @@ ScaleImageCalculator::~ScaleImageCalculator() {}
input_width_, "x", input_height_));
}
if (input_format_ != image_frame.Format()) {
std::string image_frame_format_desc, input_format_desc;
#ifdef MEDIAPIPE_MOBILE
image_frame_format_desc = std::to_string(image_frame.Format());
input_format_desc = std::to_string(input_format_);
#else
const proto_ns::EnumDescriptor* desc = ImageFormat::Format_descriptor();
image_frame_format_desc =
desc->FindValueByNumber(image_frame.Format())->DebugString();
input_format_desc = desc->FindValueByNumber(input_format_)->DebugString();
#endif // MEDIAPIPE_MOBILE
return tool::StatusFail(absl::StrCat(
"If a header specifies a format, then image frames on "
"the stream must have that format. Actual format ",
desc->FindValueByNumber(image_frame.Format())->DebugString(),
" but expected ",
desc->FindValueByNumber(input_format_)->DebugString()));
image_frame_format_desc, " but expected ", input_format_desc));
}
}
return ::mediapipe::OkStatus();
@@ -264,7 +264,7 @@ class PackMediaSequenceCalculator : public CalculatorBase {
if (options.output_only_if_all_present()) {
::mediapipe::Status status = VerifySequence();
if (!status.ok()) {
cc->GetCounter(status.error_message())->Increment();
cc->GetCounter(status.ToString())->Increment();
return status;
}
}
@@ -454,7 +454,7 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
// RET_CHECK on the tf::Status object itself in order to print an
// informative error message.
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.error_message();
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.ToString();
const int64 run_end_time = absl::ToUnixMicros(clock_->TimeNow());
cc->GetCounter(kTotalSessionRunsTimeUsecsCounterSuffix)
@@ -109,7 +109,7 @@ class TensorFlowSessionFromFrozenGraphCalculator : public CalculatorBase {
RET_CHECK(graph_def.ParseFromString(graph_def_serialized));
const tf::Status tf_status = session->session->Create(graph_def);
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.error_message();
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.ToString();
for (const auto& key_value : options.tag_to_tensor_names()) {
session->tag_to_tensor_map[key_value.first] = key_value.second;
@@ -119,7 +119,7 @@ class TensorFlowSessionFromFrozenGraphCalculator : public CalculatorBase {
session->session->Run({}, {}, initialization_op_names, {});
// RET_CHECK on the tf::Status object itself in order to print an
// informative error message.
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.error_message();
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.ToString();
}
cc->OutputSidePackets().Tag("SESSION").Set(Adopt(session.release()));
@@ -109,7 +109,7 @@ class TensorFlowSessionFromFrozenGraphGenerator : public PacketGenerator {
RET_CHECK(graph_def.ParseFromString(graph_def_serialized));
const tf::Status tf_status = session->session->Create(graph_def);
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.error_message();
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.ToString();
for (const auto& key_value : options.tag_to_tensor_names()) {
session->tag_to_tensor_map[key_value.first] = key_value.second;
@@ -119,7 +119,7 @@ class TensorFlowSessionFromFrozenGraphGenerator : public PacketGenerator {
session->session->Run({}, {}, initialization_op_names, {});
// RET_CHECK on the tf::Status object itself in order to print an
// informative error message.
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.error_message();
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.ToString();
}
output_side_packets->Tag("SESSION") = Adopt(session.release());
@@ -140,7 +140,7 @@ class TensorFlowSessionFromSavedModelCalculator : public CalculatorBase {
if (!status.ok()) {
return ::mediapipe::Status(
static_cast<::mediapipe::StatusCode>(status.code()),
status.error_message());
status.ToString());
}
auto session = absl::make_unique<TensorFlowSession>();
@@ -135,7 +135,7 @@ class TensorFlowSessionFromSavedModelGenerator : public PacketGenerator {
if (!status.ok()) {
return ::mediapipe::Status(
static_cast<::mediapipe::StatusCode>(status.code()),
status.error_message());
status.ToString());
}
auto session = absl::make_unique<TensorFlowSession>();
@@ -81,7 +81,7 @@ class TFRecordReaderCalculator : public CalculatorBase {
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();
<< "Failed to open tfrecord file: " << tf_status.ToString();
tensorflow::io::RecordReader reader(file.get(),
tensorflow::io::RecordReaderOptions());
tensorflow::uint64 offset = 0;
@@ -94,7 +94,7 @@ class TFRecordReaderCalculator : public CalculatorBase {
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();
<< "Failed to read tfrecord: " << tf_status.ToString();
if (current_idx == target_idx) {
if (cc->OutputSidePackets().HasTag(kExampleTag)) {
tensorflow::Example tf_example;
@@ -294,11 +294,15 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
if (use_quantized_tensors_) {
RET_CHECK(image_frame.Format() != mediapipe::ImageFormat::VEC32F1)
<< "Only 8-bit input images are supported for quantization.";
quant.type = kTfLiteAffineQuantization;
quant.params = nullptr;
// Optional: Set 'quant' quantization params here if needed.
interpreter_->SetTensorParametersReadWrite(0, kTfLiteUInt8, "",
{channels_preserved}, quant);
} else {
// Default TfLiteQuantization used for no quantization.
// Initialize structure for no quantization.
quant.type = kTfLiteNoQuantization;
quant.params = nullptr;
interpreter_->SetTensorParametersReadWrite(0, kTfLiteFloat32, "",
{channels_preserved}, quant);
}
@@ -422,40 +426,35 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
#elif defined(MEDIAPIPE_IOS)
// GpuBuffer to id<MTLBuffer> conversion.
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
{
id<MTLTexture> src_texture = [gpu_helper_ metalTextureWithGpuBuffer:input];
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
command_buffer.label = @"TfLiteConverterCalculatorConvert";
id<MTLComputeCommandEncoder> compute_encoder =
[command_buffer computeCommandEncoder];
[compute_encoder setComputePipelineState:gpu_data_out_->pipeline_state];
[compute_encoder setTexture:src_texture atIndex:0];
[compute_encoder setBuffer:gpu_data_out_->buffer offset:0 atIndex:1];
MTLSize threads_per_group = MTLSizeMake(kWorkgroupSize, kWorkgroupSize, 1);
MTLSize threadgroups =
MTLSizeMake(NumGroups(input.width(), kWorkgroupSize),
NumGroups(input.height(), kWorkgroupSize), 1);
[compute_encoder dispatchThreadgroups:threadgroups
threadsPerThreadgroup:threads_per_group];
[compute_encoder endEncoding];
[command_buffer commit];
[command_buffer waitUntilCompleted];
}
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
id<MTLTexture> src_texture = [gpu_helper_ metalTextureWithGpuBuffer:input];
command_buffer.label = @"TfLiteConverterCalculatorConvertAndBlit";
id<MTLComputeCommandEncoder> compute_encoder =
[command_buffer computeCommandEncoder];
[compute_encoder setComputePipelineState:gpu_data_out_->pipeline_state];
[compute_encoder setTexture:src_texture atIndex:0];
[compute_encoder setBuffer:gpu_data_out_->buffer offset:0 atIndex:1];
MTLSize threads_per_group = MTLSizeMake(kWorkgroupSize, kWorkgroupSize, 1);
MTLSize threadgroups =
MTLSizeMake(NumGroups(input.width(), kWorkgroupSize),
NumGroups(input.height(), kWorkgroupSize), 1);
[compute_encoder dispatchThreadgroups:threadgroups
threadsPerThreadgroup:threads_per_group];
[compute_encoder endEncoding];
// 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;
output_tensors->at(0) =
[device newBufferWithLength:gpu_data_out_->elements * sizeof(float)
options:MTLResourceStorageModeShared];
[MPPMetalUtil blitMetalBufferTo:output_tensors->at(0)
from:gpu_data_out_->buffer
blocking:true
commandBuffer:[gpu_helper_ commandBuffer]];
}
id<MTLDevice> device = gpu_helper_.mtlDevice;
output_tensors->at(0) =
[device newBufferWithLength:gpu_data_out_->elements * sizeof(float)
options:MTLResourceStorageModeShared];
[MPPMetalUtil blitMetalBufferTo:output_tensors->at(0)
from:gpu_data_out_->buffer
blocking:false
commandBuffer:command_buffer];
cc->Outputs()
.Tag("TENSORS_GPU")
@@ -56,6 +56,10 @@
#endif // ANDROID
namespace {
// Commonly used to compute the number of blocks to launch in a kernel.
int NumGroups(const int size, const int group_size) { // NOLINT
return (size + group_size - 1) / group_size;
}
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
@@ -176,12 +180,13 @@ class TfLiteInferenceCalculator : public CalculatorBase {
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
mediapipe::GlCalculatorHelper gpu_helper_;
std::unique_ptr<GPUData> gpu_data_in_;
std::vector<std::unique_ptr<GPUData>> gpu_data_in_;
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
#elif defined(MEDIAPIPE_IOS)
MPPMetalHelper* gpu_helper_ = nullptr;
std::unique_ptr<GPUData> gpu_data_in_;
std::vector<std::unique_ptr<GPUData>> gpu_data_in_;
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
id<MTLComputePipelineState> fp32_to_fp16_program_;
TFLBufferConvert* converter_from_BPHWC4_ = nil;
#endif
@@ -308,22 +313,41 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
const auto& input_tensors =
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
RET_CHECK_EQ(input_tensors.size(), 1);
RET_CHECK_GT(input_tensors.size(), 0);
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
[this, &input_tensors]() -> ::mediapipe::Status {
// Explicit copy input.
RET_CHECK_CALL(CopyBuffer(input_tensors[0], gpu_data_in_->buffer));
gpu_data_in_.resize(input_tensors.size());
for (int i = 0; i < input_tensors.size(); ++i) {
RET_CHECK_CALL(
CopyBuffer(input_tensors[i], gpu_data_in_[i]->buffer));
}
return ::mediapipe::OkStatus();
}));
#elif defined(MEDIAPIPE_IOS)
const auto& input_tensors =
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
RET_CHECK_EQ(input_tensors.size(), 1);
// Explicit copy input.
[MPPMetalUtil blitMetalBufferTo:gpu_data_in_->buffer
from:input_tensors[0]
blocking:true
commandBuffer:[gpu_helper_ commandBuffer]];
RET_CHECK_GT(input_tensors.size(), 0);
// Explicit copy input with conversion float 32 bits to 16 bits.
gpu_data_in_.resize(input_tensors.size());
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
command_buffer.label = @"TfLiteInferenceCalculatorConvert";
id<MTLComputeCommandEncoder> compute_encoder =
[command_buffer computeCommandEncoder];
[compute_encoder setComputePipelineState:fp32_to_fp16_program_];
for (int i = 0; i < input_tensors.size(); ++i) {
[compute_encoder setBuffer:input_tensors[i] offset:0 atIndex:0];
[compute_encoder setBuffer:gpu_data_in_[i]->buffer offset:0 atIndex:1];
constexpr int kWorkgroupSize = 64; // Block size for GPU shader.
MTLSize threads_per_group = MTLSizeMake(kWorkgroupSize, 1, 1);
const int threadgroups =
NumGroups(gpu_data_in_[i]->elements, kWorkgroupSize);
[compute_encoder dispatchThreadgroups:MTLSizeMake(threadgroups, 1, 1)
threadsPerThreadgroup:threads_per_group];
}
[compute_encoder endEncoding];
[command_buffer commit];
#else
RET_CHECK_FAIL() << "GPU processing not enabled.";
#endif
@@ -404,7 +428,6 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
}
[convert_command endEncoding];
[command_buffer commit];
[command_buffer waitUntilCompleted];
cc->Outputs()
.Tag("TENSORS_GPU")
.Add(output_tensors.release(), cc->InputTimestamp());
@@ -432,7 +455,9 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
TfLiteGpuDelegateDelete(delegate_);
gpu_data_in_.reset();
for (int i = 0; i < gpu_data_in_.size(); ++i) {
gpu_data_in_[i].reset();
}
for (int i = 0; i < gpu_data_out_.size(); ++i) {
gpu_data_out_[i].reset();
}
@@ -440,7 +465,9 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
}));
#elif defined(MEDIAPIPE_IOS)
TFLGpuDelegateDelete(delegate_);
gpu_data_in_.reset();
for (int i = 0; i < gpu_data_in_.size(); ++i) {
gpu_data_in_[i].reset();
}
for (int i = 0; i < gpu_data_out_.size(); ++i) {
gpu_data_out_[i].reset();
}
@@ -545,24 +572,24 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
if (gpu_input_) {
// Get input image sizes.
gpu_data_in_ = absl::make_unique<GPUData>();
const auto& input_indices = interpreter_->inputs();
RET_CHECK_EQ(input_indices.size(), 1); // TODO accept > 1.
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[0]);
gpu_data_in_->elements = 1;
for (int d = 0; d < tensor->dims->size; ++d) {
gpu_data_in_->elements *= tensor->dims->data[d];
gpu_data_in_.resize(input_indices.size());
for (int i = 0; i < input_indices.size(); ++i) {
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[0]);
gpu_data_in_[i] = absl::make_unique<GPUData>();
gpu_data_in_[i]->elements = 1;
for (int d = 0; d < tensor->dims->size; ++d) {
gpu_data_in_[i]->elements *= tensor->dims->data[d];
}
// Create and bind input buffer.
RET_CHECK_CALL(
::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
gpu_data_in_[i]->elements, &gpu_data_in_[i]->buffer));
RET_CHECK_EQ(TfLiteGpuDelegateBindBufferToTensor(
delegate_, gpu_data_in_[i]->buffer.id(),
interpreter_->inputs()[i]),
kTfLiteOk);
}
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.
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
kTfLiteOk);
}
if (gpu_output_) {
// Get output image sizes.
@@ -594,41 +621,68 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
#endif // OpenGL
#if defined(MEDIAPIPE_IOS)
const int kHalfSize = 2; // sizeof(half)
// Configure and create the delegate.
TFLGpuDelegateOptions options;
options.allow_precision_loss = false; // Must match converter, F=float/T=half
options.allow_precision_loss = true;
options.wait_type = TFLGpuDelegateWaitType::TFLGpuDelegateWaitTypePassive;
if (!delegate_) delegate_ = TFLGpuDelegateCreate(&options);
id<MTLDevice> device = gpu_helper_.mtlDevice;
if (gpu_input_) {
// Get input image sizes.
gpu_data_in_ = absl::make_unique<GPUData>();
const auto& input_indices = interpreter_->inputs();
RET_CHECK_EQ(input_indices.size(), 1);
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[0]);
gpu_data_in_->elements = 1;
// On iOS GPU, input must be 4 channels, regardless of what model expects.
{
gpu_data_in_->elements *= tensor->dims->data[0]; // batch
gpu_data_in_->elements *= tensor->dims->data[1]; // height
gpu_data_in_->elements *= tensor->dims->data[2]; // width
gpu_data_in_->elements *= 4; // channels
gpu_data_in_.resize(input_indices.size());
for (int i = 0; i < input_indices.size(); ++i) {
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[i]);
gpu_data_in_[i] = absl::make_unique<GPUData>();
gpu_data_in_[i]->shape.b = tensor->dims->data[0];
gpu_data_in_[i]->shape.h = tensor->dims->data[1];
gpu_data_in_[i]->shape.w = tensor->dims->data[2];
// On iOS GPU, input must be 4 channels, regardless of what model expects.
gpu_data_in_[i]->shape.c = 4;
gpu_data_in_[i]->elements =
gpu_data_in_[i]->shape.b * gpu_data_in_[i]->shape.h *
gpu_data_in_[i]->shape.w * gpu_data_in_[i]->shape.c;
// Input to model can be RGBA only.
if (tensor->dims->data[3] != 4) {
LOG(WARNING) << "Please ensure input GPU tensor is 4 channels.";
}
const std::string shader_source =
absl::Substitute(R"(#include <metal_stdlib>
using namespace metal;
kernel void convertKernel(device float4* const input_buffer [[buffer(0)]],
device half4* output_buffer [[buffer(1)]],
uint gid [[thread_position_in_grid]]) {
if (gid >= $0) return;
output_buffer[gid] = half4(input_buffer[gid]);
})",
gpu_data_in_[i]->elements / 4);
NSString* library_source =
[NSString stringWithUTF8String:shader_source.c_str()];
NSError* error = nil;
id<MTLLibrary> library =
[device newLibraryWithSource:library_source options:nil error:&error];
RET_CHECK(library != nil) << "Couldn't create shader library "
<< [[error localizedDescription] UTF8String];
id<MTLFunction> kernel_func = nil;
kernel_func = [library newFunctionWithName:@"convertKernel"];
RET_CHECK(kernel_func != nil) << "Couldn't create kernel function.";
fp32_to_fp16_program_ =
[device newComputePipelineStateWithFunction:kernel_func error:&error];
RET_CHECK(fp32_to_fp16_program_ != nil)
<< "Couldn't create pipeline state "
<< [[error localizedDescription] UTF8String];
// Create and bind input buffer.
gpu_data_in_[i]->buffer =
[device newBufferWithLength:gpu_data_in_[i]->elements * kHalfSize
options:MTLResourceStorageModeShared];
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
delegate_, input_indices[i], gpu_data_in_[i]->buffer),
true);
}
// Input to model can be RGBA only.
if (tensor->dims->data[3] != 4) {
LOG(WARNING) << "Please ensure input GPU tensor is 4 channels.";
}
// Create and bind input buffer.
gpu_data_in_->buffer =
[device newBufferWithLength:gpu_data_in_->elements * sizeof(float)
options:MTLResourceStorageModeShared];
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
delegate_,
input_indices[0], // First tensor only
gpu_data_in_->buffer),
true);
}
if (gpu_output_) {
// Get output image sizes.
@@ -669,15 +723,16 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
interpreter_->SetAllowBufferHandleOutput(true);
for (int i = 0; i < gpu_data_out_.size(); ++i) {
gpu_data_out_[i]->buffer =
[device newBufferWithLength:gpu_data_out_[i]->elements * sizeof(float)
[device newBufferWithLength:gpu_data_out_[i]->elements * kHalfSize
options:MTLResourceStorageModeShared];
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
delegate_, output_indices[i], gpu_data_out_[i]->buffer),
true);
}
// Create converter for GPU output.
converter_from_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
isFloat16:false
isFloat16:true
convertToPBHWC4:false];
if (converter_from_BPHWC4_ == nil) {
return mediapipe::InternalError(
@@ -472,11 +472,11 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
// Copy inputs.
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_boxes_buffer
from:input_tensors[0]
blocking:true
blocking:false
commandBuffer:[gpu_helper_ commandBuffer]];
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_scores_buffer
from:input_tensors[1]
blocking:true
blocking:false
commandBuffer:[gpu_helper_ commandBuffer]];
if (!anchors_init_) {
if (side_packet_anchors_) {
@@ -491,48 +491,37 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
RET_CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_anchors_buffer
from:input_tensors[2]
blocking:true
blocking:false
commandBuffer:[gpu_helper_ commandBuffer]];
}
anchors_init_ = true;
}
// Run shaders.
{
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
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
command_buffer.label = @"TfLiteDecodeAndScoreBoxes";
id<MTLComputeCommandEncoder> command_encoder =
[command_buffer computeCommandEncoder];
[command_encoder setComputePipelineState:gpu_data_->decode_program];
[command_encoder setBuffer:gpu_data_->decoded_boxes_buffer
offset:0
atIndex:0];
[command_encoder setBuffer:gpu_data_->raw_boxes_buffer offset:0 atIndex:1];
[command_encoder 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);
[command_encoder dispatchThreadgroups:decode_threadgroups
threadsPerThreadgroup:decode_threads_per_group];
[command_encoder setComputePipelineState:gpu_data_->score_program];
[command_encoder setBuffer:gpu_data_->scored_boxes_buffer offset:0 atIndex:0];
[command_encoder 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);
[command_encoder dispatchThreadgroups:score_threadgroups
threadsPerThreadgroup:score_threads_per_group];
[score_command endEncoding];
[command_buffer commit];
[command_buffer waitUntilCompleted];
}
[command_encoder endEncoding];
[MPPMetalUtil commitCommandBufferAndWait:command_buffer];
// Copy decoded boxes from GPU to CPU.
std::vector<float> boxes(num_boxes_ * num_coords_);
+46
View File
@@ -65,6 +65,15 @@ proto_library(
],
)
proto_library(
name = "video_pre_stream_calculator_proto",
srcs = ["video_pre_stream_calculator.proto"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_proto",
],
)
mediapipe_cc_proto_library(
name = "motion_analysis_calculator_cc_proto",
srcs = ["motion_analysis_calculator.proto"],
@@ -98,6 +107,16 @@ mediapipe_cc_proto_library(
deps = [":box_tracker_calculator_proto"],
)
mediapipe_cc_proto_library(
name = "video_pre_stream_calculator_cc_proto",
srcs = ["video_pre_stream_calculator.proto"],
cc_deps = [
"//mediapipe/framework:calculator_cc_proto",
],
visibility = ["//visibility:public"],
deps = [":video_pre_stream_calculator_proto"],
)
mediapipe_cc_proto_library(
name = "flow_to_image_calculator_cc_proto",
srcs = ["flow_to_image_calculator.proto"],
@@ -280,6 +299,19 @@ cc_library(
alwayslink = 1,
)
cc_library(
name = "video_pre_stream_calculator",
srcs = ["video_pre_stream_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
":video_pre_stream_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:video_stream_header",
],
alwayslink = 1,
)
filegroup(
name = "test_videos",
srcs = [
@@ -411,3 +443,17 @@ cc_test(
"//mediapipe/util/tracking:tracking_cc_proto",
],
)
cc_test(
name = "video_pre_stream_calculator_test",
srcs = ["video_pre_stream_calculator_test.cc"],
deps = [
":video_pre_stream_calculator",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:video_stream_header",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:status",
],
)
@@ -72,6 +72,8 @@ ImageFormat::Format GetImageFormat(int num_channels) {
// OpenCV's VideoCapture doesn't decode audio tracks. If the audio tracks need
// to be saved, specify an output side packet with tag "SAVED_AUDIO_PATH".
// The calculator will call FFmpeg binary to save audio tracks as an aac file.
// If the audio tracks can't be extracted by FFmpeg, the output side packet
// will contain an empty std::string.
//
// Example config:
// node {
@@ -150,13 +152,23 @@ class OpenCvVideoDecoderCalculator : public CalculatorBase {
if (cc->OutputSidePackets().HasTag("SAVED_AUDIO_PATH")) {
#ifdef HAVE_FFMPEG
std::string saved_audio_path = std::tmpnam(nullptr);
system(absl::StrCat("ffmpeg -nostats -loglevel 0 -i ", input_file_path,
" -vn -f adts ", saved_audio_path)
.c_str());
cc->OutputSidePackets()
.Tag("SAVED_AUDIO_PATH")
.Set(MakePacket<std::string>(saved_audio_path));
std::string ffmpeg_command =
absl::StrCat("ffmpeg -nostats -loglevel 0 -i ", input_file_path,
" -vn -f adts ", saved_audio_path);
system(ffmpeg_command.c_str());
int status_code = system(absl::StrCat("ls ", saved_audio_path).c_str());
if (status_code == 0) {
cc->OutputSidePackets()
.Tag("SAVED_AUDIO_PATH")
.Set(MakePacket<std::string>(saved_audio_path));
} else {
LOG(WARNING) << "FFmpeg can't extract audio from " << input_file_path
<< " by executing the following command: "
<< ffmpeg_command;
cc->OutputSidePackets()
.Tag("SAVED_AUDIO_PATH")
.Set(MakePacket<std::string>(std::string()));
}
#else
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
<< "OpenCVVideoDecoderCalculator can't save the audio file "
@@ -55,8 +55,12 @@ TEST(OpenCvVideoDecoderCalculatorTest, TestMp4Avc720pVideo) {
EXPECT_EQ(640, header.height);
EXPECT_FLOAT_EQ(6.0f, header.duration);
EXPECT_FLOAT_EQ(30.0f, header.frame_rate);
EXPECT_EQ(180, runner.Outputs().Tag("VIDEO").packets.size());
for (int i = 0; i < 180; ++i) {
// The number of the output packets should be 180.
// Some OpenCV version returns the first two frames with the same timestamp on
// macos and we might miss one frame here.
int num_of_packets = runner.Outputs().Tag("VIDEO").packets.size();
EXPECT_GE(num_of_packets, 179);
for (int i = 0; i < num_of_packets; ++i) {
Packet image_frame_packet = runner.Outputs().Tag("VIDEO").packets[i];
cv::Mat output_mat =
formats::MatView(&(image_frame_packet.Get<ImageFrame>()));
@@ -141,8 +145,12 @@ TEST(OpenCvVideoDecoderCalculatorTest, TestMkvVp8Video) {
EXPECT_EQ(320, header.height);
EXPECT_FLOAT_EQ(6.0f, header.duration);
EXPECT_FLOAT_EQ(30.0f, header.frame_rate);
EXPECT_EQ(180, runner.Outputs().Tag("VIDEO").packets.size());
for (int i = 0; i < 180; ++i) {
// The number of the output packets should be 180.
// Some OpenCV version returns the first two frames with the same timestamp on
// macos and we might miss one frame here.
int num_of_packets = runner.Outputs().Tag("VIDEO").packets.size();
EXPECT_GE(num_of_packets, 179);
for (int i = 0; i < num_of_packets; ++i) {
Packet image_frame_packet = runner.Outputs().Tag("VIDEO").packets[i];
cv::Mat output_mat =
formats::MatView(&(image_frame_packet.Get<ImageFrame>()));
@@ -183,14 +183,20 @@ class OpenCvVideoEncoderCalculator : public CalculatorBase {
#ifdef HAVE_FFMPEG
const std::string& audio_file_path =
cc->InputSidePackets().Tag("AUDIO_FILE_PATH").Get<std::string>();
// A temp output file is needed because FFmpeg can't do in-place editing.
const std::string temp_file_path = std::tmpnam(nullptr);
system(absl::StrCat("mv ", output_file_path_, " ", temp_file_path,
"&& ffmpeg -nostats -loglevel 0 -i ", temp_file_path,
" -i ", audio_file_path,
" -c copy -map 0:v:0 -map 1:a:0 ", output_file_path_,
"&& rm ", temp_file_path)
.c_str());
if (audio_file_path.empty()) {
LOG(WARNING) << "OpenCvVideoEncoderCalculator isn't able to attach the "
"audio tracks to the generated video because the audio "
"file path is not specified.";
} else {
// A temp output file is needed because FFmpeg can't do in-place editing.
const std::string temp_file_path = std::tmpnam(nullptr);
system(absl::StrCat("mv ", output_file_path_, " ", temp_file_path,
"&& ffmpeg -nostats -loglevel 0 -i ", temp_file_path,
" -i ", audio_file_path,
" -c copy -map 0:v:0 -map 1:a:0 ", output_file_path_,
"&& rm ", temp_file_path)
.c_str());
}
#else
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
@@ -210,8 +210,8 @@ TEST(OpenCvVideoEncoderCalculatorTest, TestMkvVp8Video) {
EXPECT_EQ(video_header.frame_rate,
static_cast<double>(cap.get(cv::CAP_PROP_FPS)));
EXPECT_EQ(video_header.duration,
static_cast<int>(cap.get(cv::CAP_PROP_FRAME_COUNT) /
cap.get(cv::CAP_PROP_FPS)));
static_cast<int>(std::round(cap.get(cv::CAP_PROP_FRAME_COUNT) /
cap.get(cv::CAP_PROP_FPS))));
}
} // namespace
@@ -0,0 +1,142 @@
// 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/video/video_pre_stream_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/video_stream_header.h"
namespace mediapipe {
// Sets up VideoHeader based on the 1st ImageFrame and emits it with timestamp
// PreStream. Note that this calculator only fills in format, width, and height,
// i.e. frame_rate and duration will not be filled, unless:
// 1) an existing VideoHeader is provided at PreStream(). In such case, the
// frame_rate and duration, if they exist, will be copied from the existing
// VideoHeader.
// 2) you specify frame_rate and duration through the options. In this case, the
// options will overwrite the existing VideoHeader if it is available.
//
// Example config:
// node {
// calculator: "VideoPreStreamCalculator"
// input_stream: "FRAME:cropped_frames"
// input_stream: "VIDEO_PRESTREAM:original_video_header"
// output_stream: "cropped_frames_video_header"
// }
//
// or
//
// node {
// calculator: "VideoPreStreamCalculator"
// input_stream: "cropped_frames"
// output_stream: "video_header"
// }
class VideoPreStreamCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc);
::mediapipe::Status Open(CalculatorContext* cc) override;
::mediapipe::Status Process(CalculatorContext* cc) override;
private:
::mediapipe::Status ProcessWithFrameRateInPreStream(CalculatorContext* cc);
::mediapipe::Status ProcessWithFrameRateInOptions(CalculatorContext* cc);
std::unique_ptr<VideoHeader> header_;
bool frame_rate_in_prestream_ = false;
bool emitted_ = false;
};
REGISTER_CALCULATOR(VideoPreStreamCalculator);
::mediapipe::Status VideoPreStreamCalculator::GetContract(
CalculatorContract* cc) {
if (!cc->Inputs().UsesTags()) {
cc->Inputs().Index(0).Set<ImageFrame>();
} else {
cc->Inputs().Tag("FRAME").Set<ImageFrame>();
cc->Inputs().Tag("VIDEO_PRESTREAM").Set<VideoHeader>();
}
cc->Outputs().Index(0).Set<VideoHeader>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status VideoPreStreamCalculator::Open(CalculatorContext* cc) {
frame_rate_in_prestream_ = cc->Inputs().UsesTags() &&
cc->Inputs().HasTag("FRAME") &&
cc->Inputs().HasTag("VIDEO_PRESTREAM");
header_ = absl::make_unique<VideoHeader>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status VideoPreStreamCalculator::ProcessWithFrameRateInPreStream(
CalculatorContext* cc) {
cc->GetCounter("ProcessWithFrameRateInPreStream")->Increment();
if (cc->InputTimestamp() == Timestamp::PreStream()) {
RET_CHECK(cc->Inputs().Tag("FRAME").IsEmpty());
RET_CHECK(!cc->Inputs().Tag("VIDEO_PRESTREAM").IsEmpty());
*header_ = cc->Inputs().Tag("VIDEO_PRESTREAM").Get<VideoHeader>();
RET_CHECK_NE(header_->frame_rate, 0.0) << "frame rate should be non-zero";
} else {
RET_CHECK(cc->Inputs().Tag("VIDEO_PRESTREAM").IsEmpty())
<< "Packet on VIDEO_PRESTREAM must come in at Timestamp::PreStream().";
RET_CHECK(!cc->Inputs().Tag("FRAME").IsEmpty());
const auto& frame = cc->Inputs().Tag("FRAME").Get<ImageFrame>();
header_->format = frame.Format();
header_->width = frame.Width();
header_->height = frame.Height();
RET_CHECK_NE(header_->frame_rate, 0.0) << "frame rate should be non-zero";
cc->Outputs().Index(0).Add(header_.release(), Timestamp::PreStream());
emitted_ = true;
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status VideoPreStreamCalculator::Process(CalculatorContext* cc) {
cc->GetCounter("Process")->Increment();
if (emitted_) {
return ::mediapipe::OkStatus();
}
if (frame_rate_in_prestream_) {
return ProcessWithFrameRateInPreStream(cc);
} else {
return ProcessWithFrameRateInOptions(cc);
}
}
::mediapipe::Status VideoPreStreamCalculator::ProcessWithFrameRateInOptions(
CalculatorContext* cc) {
cc->GetCounter("ProcessWithFrameRateInOptions")->Increment();
RET_CHECK_NE(cc->InputTimestamp(), Timestamp::PreStream());
const auto& frame = cc->Inputs().Index(0).Get<ImageFrame>();
header_->format = frame.Format();
header_->width = frame.Width();
header_->height = frame.Height();
const auto& options = cc->Options<VideoPreStreamCalculatorOptions>();
if (options.fps().has_value()) {
header_->frame_rate = options.fps().value();
} else if (options.fps().has_ratio()) {
const VideoPreStreamCalculatorOptions::Fps::Rational32& ratio =
options.fps().ratio();
if (ratio.numerator() > 0 && ratio.denominator() > 0) {
header_->frame_rate =
static_cast<double>(ratio.numerator()) / ratio.denominator();
}
}
RET_CHECK_NE(header_->frame_rate, 0.0) << "frame rate should be non-zero";
cc->Outputs().Index(0).Add(header_.release(), Timestamp::PreStream());
emitted_ = true;
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";
message VideoPreStreamCalculatorOptions {
extend CalculatorOptions {
optional VideoPreStreamCalculatorOptions ext = 151386123;
}
// An arbitrary number of frames per second.
// Prefer the StandardFps enum to store industry-standard, safe FPS values.
message Fps {
// The possibly approximated value of the frame rate, in frames per second.
// Unsafe to use in accurate computations because prone to rounding errors.
// For example, the 23.976 FPS value has no exact representation as a
// double.
optional double value = 1;
message Rational32 {
optional int32 numerator = 1;
optional int32 denominator = 2;
}
// The exact value of the frame rate, as a rational number.
optional Rational32 ratio = 2;
}
optional Fps fps = 1;
}
@@ -0,0 +1,186 @@
// 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/formats/image_frame.h"
#include "mediapipe/framework/formats/video_stream_header.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 {
namespace {
TEST(VideoPreStreamCalculatorTest, ProcessesWithFrameRateInOptions) {
auto config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "input"
node {
calculator: "VideoPreStreamCalculator"
input_stream: "input"
output_stream: "output"
options {
[mediapipe.VideoPreStreamCalculatorOptions.ext] { fps { value: 3 } }
}
})");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
auto poller_status = graph.AddOutputStreamPoller("output");
MP_ASSERT_OK(poller_status.status());
OutputStreamPoller& poller = poller_status.ValueOrDie();
MP_ASSERT_OK(graph.StartRun({}));
MP_ASSERT_OK(graph.AddPacketToInputStream(
"input",
Adopt(new ImageFrame(ImageFormat::SRGB, 1, 2)).At(Timestamp(0))));
// It is *not* VideoPreStreamCalculator's job to detect errors in an
// ImageFrame stream. It just waits for the 1st ImageFrame, extracts info for
// VideoHeader, and emits it. Thus, the following is fine.
MP_ASSERT_OK(graph.AddPacketToInputStream(
"input",
Adopt(new ImageFrame(ImageFormat::SRGBA, 3, 4)).At(Timestamp(1))));
MP_ASSERT_OK(graph.CloseInputStream("input"));
Packet packet;
ASSERT_TRUE(poller.Next(&packet));
const auto& video_header = packet.Get<VideoHeader>();
EXPECT_EQ(video_header.format, ImageFormat::SRGB);
EXPECT_EQ(video_header.width, 1);
EXPECT_EQ(video_header.height, 2);
EXPECT_EQ(video_header.frame_rate, 3);
EXPECT_EQ(packet.Timestamp(), Timestamp::PreStream());
ASSERT_FALSE(poller.Next(&packet));
MP_EXPECT_OK(graph.WaitUntilDone());
}
TEST(VideoPreStreamCalculatorTest, ProcessesWithFrameRateInPreStream) {
auto config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "frame"
input_stream: "input_header"
node {
calculator: "VideoPreStreamCalculator"
input_stream: "FRAME:frame"
input_stream: "VIDEO_PRESTREAM:input_header"
output_stream: "output_header"
})");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
auto poller_status = graph.AddOutputStreamPoller("output_header");
MP_ASSERT_OK(poller_status.status());
OutputStreamPoller& poller = poller_status.ValueOrDie();
MP_ASSERT_OK(graph.StartRun({}));
auto input_header = absl::make_unique<VideoHeader>();
input_header->frame_rate = 3.0;
MP_ASSERT_OK(graph.AddPacketToInputStream(
"input_header",
Adopt(input_header.release()).At(Timestamp::PreStream())));
MP_ASSERT_OK(graph.CloseInputStream("input_header"));
MP_ASSERT_OK(graph.AddPacketToInputStream(
"frame",
Adopt(new ImageFrame(ImageFormat::SRGB, 1, 2)).At(Timestamp(0))));
MP_ASSERT_OK(graph.CloseInputStream("frame"));
Packet packet;
ASSERT_TRUE(poller.Next(&packet));
const auto& output_header = packet.Get<VideoHeader>();
EXPECT_EQ(output_header.format, ImageFormat::SRGB);
EXPECT_EQ(output_header.width, 1);
EXPECT_EQ(output_header.height, 2);
EXPECT_EQ(output_header.frame_rate, 3.0);
EXPECT_EQ(packet.Timestamp(), Timestamp::PreStream());
ASSERT_FALSE(poller.Next(&packet));
MP_EXPECT_OK(graph.WaitUntilDone());
}
TEST(VideoPreStreamCalculatorTest, FailsWithoutFrameRateInOptions) {
auto config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "frame"
node {
calculator: "VideoPreStreamCalculator"
input_stream: "frame"
output_stream: "output_header"
})");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.StartRun({}));
MP_ASSERT_OK(graph.AddPacketToInputStream(
"frame",
Adopt(new ImageFrame(ImageFormat::SRGB, 1, 2)).At(Timestamp(0))));
MP_ASSERT_OK(graph.CloseInputStream("frame"));
::mediapipe::Status status = graph.WaitUntilDone();
EXPECT_FALSE(status.ok());
EXPECT_THAT(status.ToString(),
testing::HasSubstr("frame rate should be non-zero"));
}
// Input header missing.
TEST(VideoPreStreamCalculatorTest, FailsWithoutFrameRateInPreStream1) {
auto config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "frame"
input_stream: "input_header"
node {
calculator: "VideoPreStreamCalculator"
input_stream: "FRAME:frame"
input_stream: "VIDEO_PRESTREAM:input_header"
output_stream: "output_header"
}
)");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.StartRun({}));
MP_ASSERT_OK(graph.AddPacketToInputStream(
"frame",
Adopt(new ImageFrame(ImageFormat::SRGB, 1, 2)).At(Timestamp(0))));
MP_ASSERT_OK(graph.CloseInputStream("frame"));
MP_ASSERT_OK(graph.CloseInputStream("input_header"));
::mediapipe::Status status = graph.WaitUntilDone();
EXPECT_FALSE(status.ok());
EXPECT_THAT(status.ToString(),
testing::HasSubstr("frame rate should be non-zero"));
}
// Input header not at prestream (before, with, and after frame data).
TEST(VideoPreStreamCalculatorTest, FailsWithoutFrameRateInPreStream2) {
auto config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "frame"
input_stream: "input_header"
node {
calculator: "VideoPreStreamCalculator"
input_stream: "FRAME:frame"
input_stream: "VIDEO_PRESTREAM:input_header"
output_stream: "output_header"
}
)");
for (int64 timestamp = -1; timestamp < 2; ++timestamp) {
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.StartRun({}));
auto input_header = absl::make_unique<VideoHeader>();
input_header->frame_rate = 3.0;
MP_ASSERT_OK(graph.AddPacketToInputStream(
"input_header",
Adopt(input_header.release()).At(Timestamp(timestamp))));
MP_ASSERT_OK(graph.CloseInputStream("input_header"));
MP_ASSERT_OK(graph.AddPacketToInputStream(
"frame",
Adopt(new ImageFrame(ImageFormat::SRGB, 1, 2)).At(Timestamp(0))));
MP_ASSERT_OK(graph.CloseInputStream("frame"));
::mediapipe::Status status = graph.WaitUntilDone();
EXPECT_FALSE(status.ok());
}
}
} // namespace
} // namespace mediapipe