Internal change
PiperOrigin-RevId: 477538515
This commit is contained in:
committed by
Sebastian Schmidt
parent
6cdc6443b6
commit
f8af41b1eb
@@ -1294,8 +1294,8 @@ cc_library(
|
||||
deps = [
|
||||
":get_vector_item_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/api2:node",
|
||||
"//mediapipe/framework/api2:packet",
|
||||
"//mediapipe/framework/api2:port",
|
||||
"//mediapipe/framework/formats:classification_cc_proto",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
@@ -1319,6 +1319,32 @@ cc_test(
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "vector_indices_calculator",
|
||||
srcs = ["vector_indices_calculator.cc"],
|
||||
hdrs = ["vector_indices_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/api2:node",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "vector_indices_calculator_test",
|
||||
srcs = ["vector_indices_calculator_test.cc"],
|
||||
deps = [
|
||||
":vector_indices_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "vector_size_calculator",
|
||||
srcs = ["vector_size_calculator.cc"],
|
||||
|
||||
@@ -40,6 +40,9 @@ REGISTER_CALCULATOR(EndLoopNormalizedLandmarkListVectorCalculator);
|
||||
typedef EndLoopCalculator<std::vector<bool>> EndLoopBooleanCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopBooleanCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<float>> EndLoopFloatCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopFloatCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<::mediapipe::RenderData>>
|
||||
EndLoopRenderDataCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopRenderDataCalculator);
|
||||
|
||||
@@ -24,6 +24,10 @@ using GetLandmarkListVectorItemCalculator =
|
||||
GetVectorItemCalculator<mediapipe::LandmarkList>;
|
||||
REGISTER_CALCULATOR(GetLandmarkListVectorItemCalculator);
|
||||
|
||||
using GetNormalizedLandmarkListVectorItemCalculator =
|
||||
GetVectorItemCalculator<mediapipe::NormalizedLandmarkList>;
|
||||
REGISTER_CALCULATOR(GetNormalizedLandmarkListVectorItemCalculator);
|
||||
|
||||
using GetClassificationListVectorItemCalculator =
|
||||
GetVectorItemCalculator<mediapipe::ClassificationList>;
|
||||
REGISTER_CALCULATOR(GetClassificationListVectorItemCalculator);
|
||||
|
||||
@@ -19,6 +19,7 @@
|
||||
|
||||
#include "mediapipe/calculators/core/get_vector_item_calculator.pb.h"
|
||||
#include "mediapipe/framework/api2/node.h"
|
||||
#include "mediapipe/framework/api2/packet.h"
|
||||
#include "mediapipe/framework/api2/port.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
@@ -58,7 +59,7 @@ template <typename T>
|
||||
class GetVectorItemCalculator : public Node {
|
||||
public:
|
||||
static constexpr Input<std::vector<T>> kIn{"VECTOR"};
|
||||
static constexpr Input<int>::Optional kIdx{"INDEX"};
|
||||
static constexpr Input<OneOf<int, uint64_t>>::Optional kIdx{"INDEX"};
|
||||
static constexpr Output<T> kOut{"ITEM"};
|
||||
|
||||
MEDIAPIPE_NODE_CONTRACT(kIn, kIdx, kOut);
|
||||
@@ -80,7 +81,9 @@ class GetVectorItemCalculator : public Node {
|
||||
|
||||
int idx = 0;
|
||||
if (kIdx(cc).IsConnected() && !kIdx(cc).IsEmpty()) {
|
||||
idx = kIdx(cc).Get();
|
||||
idx = kIdx(cc).Visit(
|
||||
[](uint64_t idx_uint64_t) { return static_cast<int>(idx_uint64_t); },
|
||||
[](int idx_int) { return idx_int; });
|
||||
} else if (options.has_item_index()) {
|
||||
idx = options.item_index();
|
||||
} else {
|
||||
|
||||
@@ -227,4 +227,15 @@ TEST(TestGetIntVectorItemCalculatorTest, IndexOptionsTwoTimestamps) {
|
||||
testing::ElementsAre(TimestampValue(1), TimestampValue(2)));
|
||||
}
|
||||
|
||||
TEST(TestGetIntVectorItemCalculatorTest, IndexUint64) {
|
||||
CalculatorRunner runner = MakeRunnerWithStream();
|
||||
const std::vector<int> inputs = {1, 2, 3};
|
||||
const uint64_t index = 1;
|
||||
AddInputVector(runner, inputs, 1);
|
||||
AddInputIndex(runner, index, 1);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Tag("ITEM").packets;
|
||||
EXPECT_THAT(outputs, testing::ElementsAre(IntPacket(inputs[index])));
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
// Copyright 2022 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "mediapipe/calculators/core/vector_indices_calculator.h"
|
||||
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace api2 {
|
||||
|
||||
using IntVectorIndicesCalculator = VectorIndicesCalculator<int>;
|
||||
REGISTER_CALCULATOR(IntVectorIndicesCalculator);
|
||||
|
||||
using Uint64tVectorIndicesCalculator = VectorIndicesCalculator<uint64_t>;
|
||||
REGISTER_CALCULATOR(Uint64tVectorIndicesCalculator);
|
||||
|
||||
using NormalizedLandmarkListVectorIndicesCalculator =
|
||||
VectorIndicesCalculator<mediapipe::NormalizedLandmarkList>;
|
||||
REGISTER_CALCULATOR(NormalizedLandmarkListVectorIndicesCalculator);
|
||||
|
||||
} // namespace api2
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,65 @@
|
||||
// Copyright 2022 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#ifndef MEDIAPIPE_CALCULATORS_CORE_VECTOR_INDICES_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_CORE_VECTOR_INDICES_CALCULATOR_H_
|
||||
|
||||
#include <optional>
|
||||
|
||||
#include "mediapipe/framework/api2/node.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace api2 {
|
||||
// Calculator that takes a vector and consturct an index range vector based on
|
||||
// the size of the input vector.
|
||||
//
|
||||
// Inputs:
|
||||
// VECTOR - std::vector<T>
|
||||
// Vector whose range of indices to return.
|
||||
//
|
||||
// Outputs:
|
||||
// INDICES - std::vector<int>
|
||||
// Indices vector of the input vector.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "{SpecificType}VectorIndicesCalculator"
|
||||
// input_stream: "VECTOR:vector"
|
||||
// output_stream: "INDICES:indices"
|
||||
// }
|
||||
//
|
||||
template <typename T>
|
||||
class VectorIndicesCalculator : public Node {
|
||||
public:
|
||||
static constexpr Input<std::vector<T>> kVector{"VECTOR"};
|
||||
static constexpr Output<std::vector<int>> kRange{"INDICES"};
|
||||
|
||||
MEDIAPIPE_NODE_CONTRACT(kVector, kRange);
|
||||
|
||||
absl::Status Process(CalculatorContext* cc) final {
|
||||
// Get the size of the input vector.
|
||||
const int vector_size = kVector(cc).Get().size();
|
||||
std::vector<int> out_idxs(vector_size);
|
||||
std::iota(out_idxs.begin(), out_idxs.end(), 0);
|
||||
kRange(cc).Send(out_idxs);
|
||||
return absl::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace api2
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_CORE_VECTOR_INDICES_CALCULATOR_H_
|
||||
@@ -0,0 +1,87 @@
|
||||
// Copyright 2022 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "mediapipe/calculators/core/vector_indices_calculator.h"
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
using ::testing::TestParamInfo;
|
||||
using ::testing::TestWithParam;
|
||||
using ::testing::Values;
|
||||
|
||||
template <typename T>
|
||||
void AddInputVector(CalculatorRunner& runner, const std::vector<T>& inputs,
|
||||
int timestamp) {
|
||||
runner.MutableInputs()->Tag("VECTOR").packets.push_back(
|
||||
MakePacket<std::vector<T>>(inputs).At(Timestamp(timestamp)));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct TestParams {
|
||||
const std::string test_name;
|
||||
const std::vector<T> inputs;
|
||||
const int timestamp;
|
||||
const std::vector<int> expected_indices;
|
||||
};
|
||||
|
||||
class IntVectorIndicesCalculatorTest
|
||||
: public testing::TestWithParam<TestParams<int>> {};
|
||||
|
||||
TEST_P(IntVectorIndicesCalculatorTest, Succeeds) {
|
||||
CalculatorRunner runner = CalculatorRunner(R"(
|
||||
calculator: "IntVectorIndicesCalculator"
|
||||
input_stream: "VECTOR:vector_stream"
|
||||
output_stream: "INDICES:indices_stream"
|
||||
)");
|
||||
const std::vector<int>& inputs = GetParam().inputs;
|
||||
std::vector<int> expected_indices(inputs.size());
|
||||
AddInputVector(runner, inputs, GetParam().timestamp);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Tag("INDICES").packets;
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
EXPECT_THAT(outputs[0].Get<std::vector<int>>(),
|
||||
testing::ElementsAreArray(GetParam().expected_indices));
|
||||
}
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(
|
||||
IntVectorIndicesCalculatorTest, IntVectorIndicesCalculatorTest,
|
||||
Values(TestParams<int>{
|
||||
/* test_name= */ "IntVectorIndices",
|
||||
/* inputs= */ {1, 2, 3},
|
||||
/* timestamp= */ 1,
|
||||
/* expected_indices= */ {0, 1, 2},
|
||||
},
|
||||
TestParams<int>{
|
||||
/* test_name= */ "EmptyVector",
|
||||
/* inputs= */ {},
|
||||
/* timestamp= */ 1,
|
||||
/* expected_indices= */ {},
|
||||
}),
|
||||
[](const TestParamInfo<IntVectorIndicesCalculatorTest::ParamType>& info) {
|
||||
return info.param.test_name;
|
||||
});
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -55,6 +55,14 @@ mediapipe_proto_library(
|
||||
cc_library(
|
||||
name = "audio_to_tensor_calculator",
|
||||
srcs = ["audio_to_tensor_calculator.cc"],
|
||||
copts = select({
|
||||
# b/215212850
|
||||
"//mediapipe:apple": [
|
||||
"-x objective-c++",
|
||||
"-fobjc-arc",
|
||||
],
|
||||
"//conditions:default": [],
|
||||
}),
|
||||
visibility = [
|
||||
"//mediapipe/framework:mediapipe_internal",
|
||||
],
|
||||
@@ -67,13 +75,16 @@ cc_library(
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
"//mediapipe/framework/formats:tensor",
|
||||
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/util:time_series_util",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/status",
|
||||
"@com_google_absl//absl/status:statusor",
|
||||
"@com_google_absl//absl/strings:str_format",
|
||||
"@com_google_audio_tools//audio/dsp:resampler_q",
|
||||
"@com_google_audio_tools//audio/dsp:window_functions",
|
||||
"@org_tensorflow//tensorflow/lite/c:common",
|
||||
"@pffft",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -83,6 +94,7 @@ cc_test(
|
||||
srcs = ["audio_to_tensor_calculator_test.cc"],
|
||||
deps = [
|
||||
":audio_to_tensor_calculator",
|
||||
":audio_to_tensor_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:timestamp",
|
||||
@@ -97,6 +109,58 @@ cc_test(
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_proto_library(
|
||||
name = "feedback_tensors_calculator_proto",
|
||||
srcs = ["feedback_tensors_calculator.proto"],
|
||||
visibility = [
|
||||
"//mediapipe/framework:mediapipe_internal",
|
||||
],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_options_proto",
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "feedback_tensors_calculator",
|
||||
srcs = ["feedback_tensors_calculator.cc"],
|
||||
copts = select({
|
||||
# b/215212850
|
||||
"//mediapipe:apple": [
|
||||
"-x objective-c++",
|
||||
"-fobjc-arc",
|
||||
],
|
||||
"//conditions:default": [],
|
||||
}),
|
||||
visibility = [
|
||||
"//mediapipe/framework:mediapipe_internal",
|
||||
],
|
||||
deps = [
|
||||
":feedback_tensors_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/api2:node",
|
||||
"//mediapipe/framework/formats:tensor",
|
||||
"@com_google_absl//absl/status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "feedback_tensors_calculator_test",
|
||||
srcs = ["feedback_tensors_calculator_test.cc"],
|
||||
deps = [
|
||||
":feedback_tensors_calculator",
|
||||
":feedback_tensors_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/formats:tensor",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"@org_tensorflow//tensorflow/lite/c:common",
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_proto_library(
|
||||
name = "inference_calculator_proto",
|
||||
srcs = ["inference_calculator.proto"],
|
||||
@@ -346,6 +410,10 @@ cc_library(
|
||||
}),
|
||||
)
|
||||
|
||||
# This target provides the InferenceCalculator and a default set of implementations tailored for the
|
||||
# current build platforms. More implementations can be added as separate dependencies to a client;
|
||||
# for clients that want a narrower set of implementations than the default should see the comment on
|
||||
# inference_calculator_interface.
|
||||
cc_library(
|
||||
name = "inference_calculator",
|
||||
visibility = ["//visibility:public"],
|
||||
|
||||
@@ -12,9 +12,8 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <math.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
@@ -26,6 +25,7 @@
|
||||
#include "absl/status/statusor.h"
|
||||
#include "absl/strings/str_format.h"
|
||||
#include "audio/dsp/resampler_q.h"
|
||||
#include "audio/dsp/window_functions.h"
|
||||
#include "mediapipe/calculators/tensor/audio_to_tensor_calculator.pb.h"
|
||||
#include "mediapipe/framework/api2/node.h"
|
||||
#include "mediapipe/framework/api2/packet.h"
|
||||
@@ -34,19 +34,60 @@
|
||||
#include "mediapipe/framework/formats/matrix.h"
|
||||
#include "mediapipe/framework/formats/tensor.h"
|
||||
#include "mediapipe/framework/formats/time_series_header.pb.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/util/time_series_util.h"
|
||||
#include "pffft.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace api2 {
|
||||
namespace {
|
||||
|
||||
using Options = ::mediapipe::AudioToTensorCalculatorOptions;
|
||||
using FlushMode = Options::FlushMode;
|
||||
|
||||
std::vector<float> HannWindow(int window_size, bool sqrt_hann) {
|
||||
std::vector<float> hann_window(window_size);
|
||||
audio_dsp::HannWindow().GetPeriodicSamples(window_size, &hann_window);
|
||||
if (sqrt_hann) {
|
||||
absl::c_transform(hann_window, hann_window.begin(),
|
||||
[](double x) { return std::sqrt(x); });
|
||||
}
|
||||
return hann_window;
|
||||
}
|
||||
|
||||
// PFFFT only supports transforms for inputs of length N of the form
|
||||
// N = (2^a)*(3^b)*(5^c) where b >=0 and c >= 0 and a >= 5 for the real FFT.
|
||||
bool IsValidFftSize(int size) {
|
||||
if (size <= 0) {
|
||||
return false;
|
||||
}
|
||||
constexpr int kFactors[] = {2, 3, 5};
|
||||
int factorization[] = {0, 0, 0};
|
||||
int n = static_cast<int>(size);
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
while (n % kFactors[i] == 0) {
|
||||
n = n / kFactors[i];
|
||||
++factorization[i];
|
||||
}
|
||||
}
|
||||
return factorization[0] >= 5 && n == 1;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// Converts audio buffers into tensors, possibly with resampling, buffering
|
||||
// and framing, according to specified inputs and options. All input audio
|
||||
// buffers will be first resampled from the input sample rate to the target
|
||||
// sample rate if they are not equal. The resampled audio data (with the
|
||||
// buffered samples from the previous runs in the streaming mode) will be broken
|
||||
// into fixed-sized, possibly overlapping frames. Finally, all frames will be
|
||||
// converted to and outputted as MediaPipe Tensors. The last output tensor will
|
||||
// be zero-padding if the remaining samples are insufficient.
|
||||
// into fixed-sized, possibly overlapping frames. If the calculator is not asked
|
||||
// to perform fft (the fft_size is not set in the calculator options), all
|
||||
// frames will be converted to and outputted as MediaPipe Tensors. The last
|
||||
// output tensor will be zero-padding if the remaining samples are insufficient.
|
||||
// Otherwise, when the fft_size is set and valid, the calculator will perform
|
||||
// fft on the fixed-sized audio frames, the complex DFT results will be
|
||||
// converted to and outputted as 2D MediaPipe float Tensors where the first
|
||||
// rows are the DFT real parts and the second rows are the DFT imagery parts.
|
||||
//
|
||||
// This calculator assumes that the input timestamps refer to the first
|
||||
// sample in each Matrix. The output timestamps follow this same convention.
|
||||
@@ -86,11 +127,15 @@ namespace api2 {
|
||||
// Outputs:
|
||||
// TENSORS - std::vector<Tensor>
|
||||
// Vector containing a single Tensor that represents a fix-sized audio
|
||||
// frame.
|
||||
// frame or the complex DFT results.
|
||||
// TIMESTAMPS - std::vector<Timestamp> @Optional
|
||||
// Vector containing the output timestamps emitted by the current Process()
|
||||
// invocation. In the non-streaming mode, the vector contains all of the
|
||||
// output timestamps for an input audio buffer.
|
||||
// DC_AND_NYQUIST - std::pair<float, float> @Optional.
|
||||
// A pair of dc component and nyquest component. Only can be connected when
|
||||
// the calculator performs fft (the fft_size is set in the calculator
|
||||
// options).
|
||||
//
|
||||
// Example:
|
||||
// node {
|
||||
@@ -116,12 +161,14 @@ class AudioToTensorCalculator : public Node {
|
||||
// such as sample rate.
|
||||
static constexpr Input<double>::Optional kAudioSampleRateIn{"SAMPLE_RATE"};
|
||||
static constexpr Output<std::vector<Tensor>> kTensorsOut{"TENSORS"};
|
||||
static constexpr Output<std::pair<float, float>>::Optional kDcAndNyquistOut{
|
||||
"DC_AND_NYQUIST"};
|
||||
// A vector of the output timestamps emitted by the current Process()
|
||||
// invocation. The packet timestamp is the last emitted timestamp.
|
||||
static constexpr Output<std::vector<Timestamp>>::Optional kTimestampsOut{
|
||||
"TIMESTAMPS"};
|
||||
MEDIAPIPE_NODE_CONTRACT(kAudioIn, kAudioSampleRateIn, kTensorsOut,
|
||||
kTimestampsOut);
|
||||
kDcAndNyquistOut, kTimestampsOut);
|
||||
|
||||
static absl::Status UpdateContract(CalculatorContract* cc);
|
||||
absl::Status Open(CalculatorContext* cc);
|
||||
@@ -138,6 +185,9 @@ class AudioToTensorCalculator : public Node {
|
||||
int frame_step_;
|
||||
bool stream_mode_;
|
||||
bool check_inconsistent_timestamps_;
|
||||
int padding_samples_before_;
|
||||
int padding_samples_after_;
|
||||
FlushMode flush_mode_;
|
||||
Timestamp initial_timestamp_ = Timestamp::Unstarted();
|
||||
int64 cumulative_input_samples_ = 0;
|
||||
Timestamp next_output_timestamp_ = Timestamp::Unstarted();
|
||||
@@ -151,22 +201,33 @@ class AudioToTensorCalculator : public Node {
|
||||
Matrix sample_buffer_;
|
||||
int processed_buffer_cols_ = 0;
|
||||
|
||||
// The internal state of the FFT library.
|
||||
PFFFT_Setup* fft_state_ = nullptr;
|
||||
int fft_size_ = 0;
|
||||
std::vector<float> fft_window_;
|
||||
std::vector<float, Eigen::aligned_allocator<float>> fft_input_buffer_;
|
||||
// pffft requires memory to work with to avoid using the stack.
|
||||
std::vector<float, Eigen::aligned_allocator<float>> fft_workplace_;
|
||||
std::vector<float, Eigen::aligned_allocator<float>> fft_output_;
|
||||
|
||||
absl::Status ProcessStreamingData(CalculatorContext* cc, const Matrix& input);
|
||||
absl::Status ProcessNonStreamingData(CalculatorContext* cc,
|
||||
const Matrix& input);
|
||||
|
||||
absl::Status SetupStreamingResampler(double input_sample_rate_);
|
||||
void AppendToSampleBuffer(Matrix buffer_to_append);
|
||||
void AppendZerosToSampleBuffer(int num_samples);
|
||||
|
||||
absl::StatusOr<std::vector<Tensor>> ConvertToTensor(
|
||||
const Matrix& frame_to_convert);
|
||||
absl::Status OutputTensors(const Matrix& buffer, bool should_flush,
|
||||
const Matrix& block, std::vector<int> tensor_dims);
|
||||
absl::Status OutputTensor(const Matrix& block, Timestamp timestamp,
|
||||
CalculatorContext* cc);
|
||||
absl::Status ProcessBuffer(const Matrix& buffer, bool should_flush,
|
||||
CalculatorContext* cc);
|
||||
};
|
||||
|
||||
absl::Status AudioToTensorCalculator::UpdateContract(CalculatorContract* cc) {
|
||||
const auto& options =
|
||||
cc->Options<mediapipe::AudioToTensorCalculatorOptions>();
|
||||
const auto& options = cc->Options<Options>();
|
||||
if (!options.has_num_channels() || !options.has_num_samples() ||
|
||||
!options.has_target_sample_rate()) {
|
||||
return absl::InvalidArgumentError(
|
||||
@@ -174,13 +235,21 @@ absl::Status AudioToTensorCalculator::UpdateContract(CalculatorContract* cc) {
|
||||
"`num_channels`, `num_samples`, and `target_sample_rate`.");
|
||||
}
|
||||
if (options.stream_mode()) {
|
||||
// Explicitly disables tiemstamp offset to disallow the timestamp bound
|
||||
// Explicitly disables timestamp offset to disallow the timestamp bound
|
||||
// from the input streams to be propagated to the output streams.
|
||||
// In the streaming mode, the output timestamp bound is based on
|
||||
// next_output_timestamp_, which can be smaller than the current input
|
||||
// timestamps.
|
||||
cc->SetTimestampOffset(TimestampDiff::Unset());
|
||||
}
|
||||
if (options.padding_samples_before() < 0 ||
|
||||
options.padding_samples_after() < 0) {
|
||||
return absl::InvalidArgumentError("Negative zero padding unsupported");
|
||||
}
|
||||
if (options.flush_mode() != Options::ENTIRE_TAIL_AT_TIMESTAMP_MAX &&
|
||||
options.flush_mode() != Options::PROCEED_AS_USUAL) {
|
||||
return absl::InvalidArgumentError("Unsupported flush mode");
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
@@ -202,6 +271,9 @@ absl::Status AudioToTensorCalculator::Open(CalculatorContext* cc) {
|
||||
check_inconsistent_timestamps_ = options.check_inconsistent_timestamps();
|
||||
sample_buffer_.resize(num_channels_, Eigen::NoChange);
|
||||
}
|
||||
padding_samples_before_ = options.padding_samples_before();
|
||||
padding_samples_after_ = options.padding_samples_after();
|
||||
flush_mode_ = options.flush_mode();
|
||||
|
||||
RET_CHECK(kAudioSampleRateIn(cc).IsConnected() ^
|
||||
!kAudioIn(cc).Header().IsEmpty())
|
||||
@@ -217,6 +289,25 @@ absl::Status AudioToTensorCalculator::Open(CalculatorContext* cc) {
|
||||
source_sample_rate_ = input_header.sample_rate();
|
||||
}
|
||||
}
|
||||
AppendZerosToSampleBuffer(padding_samples_before_);
|
||||
if (options.has_fft_size()) {
|
||||
RET_CHECK(IsValidFftSize(options.fft_size()))
|
||||
<< "FFT size must be of the form fft_size = (2^a)*(3^b)*(5^c) where b "
|
||||
">=0 and c >= 0 and a >= 5, the requested fft size is "
|
||||
<< options.fft_size();
|
||||
RET_CHECK_EQ(1, num_channels_)
|
||||
<< "Currently only support applying FFT on mono channel.";
|
||||
fft_size_ = options.fft_size();
|
||||
fft_state_ = pffft_new_setup(fft_size_, PFFFT_REAL);
|
||||
fft_window_ = HannWindow(fft_size_, /* sqrt_hann = */ false);
|
||||
fft_input_buffer_.resize(fft_size_);
|
||||
fft_workplace_.resize(fft_size_);
|
||||
fft_output_.resize(fft_size_);
|
||||
} else {
|
||||
RET_CHECK(!kDcAndNyquistOut(cc).IsConnected())
|
||||
<< "The DC_AND_NYQUIST output stream can only be connected when the "
|
||||
"calculator outputs fft tensors";
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
@@ -262,7 +353,12 @@ absl::Status AudioToTensorCalculator::Close(CalculatorContext* cc) {
|
||||
resampler_->Flush(&resampled_buffer);
|
||||
AppendToSampleBuffer(std::move(resampled_buffer));
|
||||
}
|
||||
return OutputTensors(sample_buffer_, /*should_flush=*/true, cc);
|
||||
AppendZerosToSampleBuffer(padding_samples_after_);
|
||||
MP_RETURN_IF_ERROR(ProcessBuffer(sample_buffer_, /*should_flush=*/true, cc));
|
||||
if (fft_state_) {
|
||||
pffft_destroy_setup(fft_state_);
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status AudioToTensorCalculator::ProcessStreamingData(
|
||||
@@ -303,7 +399,7 @@ absl::Status AudioToTensorCalculator::ProcessStreamingData(
|
||||
}
|
||||
}
|
||||
|
||||
MP_RETURN_IF_ERROR(OutputTensors(sample_buffer_, /*should_flush=*/false, cc));
|
||||
MP_RETURN_IF_ERROR(ProcessBuffer(sample_buffer_, /*should_flush=*/false, cc));
|
||||
// Removes the processed samples from the global sample buffer.
|
||||
sample_buffer_ = Matrix(sample_buffer_.rightCols(sample_buffer_.cols() -
|
||||
processed_buffer_cols_ - 1));
|
||||
@@ -323,9 +419,9 @@ absl::Status AudioToTensorCalculator::ProcessNonStreamingData(
|
||||
input_frame);
|
||||
Eigen::Map<const Matrix> matrix_mapping(resampled.data(), num_channels_,
|
||||
resampled.size() / num_channels_);
|
||||
return OutputTensors(matrix_mapping, /*should_flush=*/true, cc);
|
||||
return ProcessBuffer(matrix_mapping, /*should_flush=*/true, cc);
|
||||
}
|
||||
return OutputTensors(input_frame, /*should_flush=*/true, cc);
|
||||
return ProcessBuffer(input_frame, /*should_flush=*/true, cc);
|
||||
}
|
||||
|
||||
absl::Status AudioToTensorCalculator::SetupStreamingResampler(
|
||||
@@ -344,6 +440,16 @@ absl::Status AudioToTensorCalculator::SetupStreamingResampler(
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
void AudioToTensorCalculator::AppendZerosToSampleBuffer(int num_samples) {
|
||||
CHECK_GE(num_samples, 0); // Ensured by `UpdateContract`.
|
||||
if (num_samples == 0) {
|
||||
return;
|
||||
}
|
||||
sample_buffer_.conservativeResize(Eigen::NoChange,
|
||||
sample_buffer_.cols() + num_samples);
|
||||
sample_buffer_.rightCols(num_samples).setZero();
|
||||
}
|
||||
|
||||
void AudioToTensorCalculator::AppendToSampleBuffer(Matrix buffer_to_append) {
|
||||
sample_buffer_.conservativeResize(
|
||||
Eigen::NoChange, sample_buffer_.cols() + buffer_to_append.cols());
|
||||
@@ -351,49 +457,89 @@ void AudioToTensorCalculator::AppendToSampleBuffer(Matrix buffer_to_append) {
|
||||
}
|
||||
|
||||
absl::StatusOr<std::vector<Tensor>> AudioToTensorCalculator::ConvertToTensor(
|
||||
const Matrix& frame_to_convert) {
|
||||
Tensor tensor(Tensor::ElementType::kFloat32,
|
||||
Tensor::Shape({num_channels_, num_samples_}));
|
||||
const Matrix& block, std::vector<int> tensor_dims) {
|
||||
Tensor tensor(Tensor::ElementType::kFloat32, Tensor::Shape(tensor_dims));
|
||||
auto buffer_view = tensor.GetCpuWriteView();
|
||||
if (frame_to_convert.size() < num_channels_ * num_samples_) {
|
||||
int total_size = 1;
|
||||
for (int dim : tensor_dims) {
|
||||
total_size *= dim;
|
||||
}
|
||||
if (block.size() < total_size) {
|
||||
std::memset(buffer_view.buffer<float>(), 0, tensor.bytes());
|
||||
}
|
||||
std::memcpy(buffer_view.buffer<float>(), frame_to_convert.data(),
|
||||
frame_to_convert.size() * sizeof(float));
|
||||
std::memcpy(buffer_view.buffer<float>(), block.data(),
|
||||
block.size() * sizeof(float));
|
||||
std::vector<Tensor> tensor_vector;
|
||||
tensor_vector.push_back(std::move(tensor));
|
||||
return tensor_vector;
|
||||
}
|
||||
|
||||
absl::Status AudioToTensorCalculator::OutputTensors(const Matrix& buffer,
|
||||
absl::Status AudioToTensorCalculator::OutputTensor(const Matrix& block,
|
||||
Timestamp timestamp,
|
||||
CalculatorContext* cc) {
|
||||
std::vector<Tensor> output_tensor;
|
||||
if (fft_state_) {
|
||||
Eigen::VectorXf time_series_data =
|
||||
Eigen::VectorXf::Map(block.data(), block.size());
|
||||
// Window on input audio prior to FFT.
|
||||
std::transform(time_series_data.begin(), time_series_data.end(),
|
||||
fft_window_.begin(), fft_input_buffer_.begin(),
|
||||
std::multiplies<float>());
|
||||
pffft_transform_ordered(fft_state_, fft_input_buffer_.data(),
|
||||
fft_output_.data(), fft_workplace_.data(),
|
||||
PFFFT_FORWARD);
|
||||
if (kDcAndNyquistOut(cc).IsConnected()) {
|
||||
kDcAndNyquistOut(cc).Send(std::make_pair(fft_output_[0], fft_output_[1]),
|
||||
timestamp);
|
||||
}
|
||||
Matrix fft_output_matrix =
|
||||
Eigen::Map<const Matrix>(fft_output_.data() + 2, 1, fft_size_ - 2);
|
||||
fft_output_matrix.conservativeResize(Eigen::NoChange, fft_size_);
|
||||
// The last two elements are the DFT Nyquist values.
|
||||
fft_output_matrix(fft_size_ - 2) = fft_output_[1]; // Nyquist real part
|
||||
fft_output_matrix(fft_size_ - 1) = 0.0f; // Nyquist imagery part
|
||||
ASSIGN_OR_RETURN(output_tensor,
|
||||
ConvertToTensor(fft_output_matrix, {2, fft_size_ / 2}));
|
||||
} else {
|
||||
ASSIGN_OR_RETURN(output_tensor,
|
||||
ConvertToTensor(block, {num_channels_, num_samples_}));
|
||||
}
|
||||
kTensorsOut(cc).Send(std::move(output_tensor), timestamp);
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status AudioToTensorCalculator::ProcessBuffer(const Matrix& buffer,
|
||||
bool should_flush,
|
||||
CalculatorContext* cc) {
|
||||
const bool should_flush_at_timestamp_max =
|
||||
stream_mode_ && should_flush &&
|
||||
flush_mode_ == Options::ENTIRE_TAIL_AT_TIMESTAMP_MAX;
|
||||
int next_frame_first_col = 0;
|
||||
std::vector<Timestamp> timestamps;
|
||||
while ((!stream_mode_ || !should_flush) &&
|
||||
next_frame_first_col + num_samples_ <= buffer.cols()) {
|
||||
ASSIGN_OR_RETURN(auto output_tensor, ConvertToTensor(buffer.block(
|
||||
0, next_frame_first_col,
|
||||
num_channels_, num_samples_)));
|
||||
kTensorsOut(cc).Send(std::move(output_tensor), next_output_timestamp_);
|
||||
timestamps.push_back(next_output_timestamp_);
|
||||
next_output_timestamp_ += round(frame_step_ / target_sample_rate_ *
|
||||
Timestamp::kTimestampUnitsPerSecond);
|
||||
next_frame_first_col += frame_step_;
|
||||
if (!should_flush_at_timestamp_max) {
|
||||
while (next_frame_first_col + num_samples_ <= buffer.cols()) {
|
||||
MP_RETURN_IF_ERROR(OutputTensor(
|
||||
buffer.block(0, next_frame_first_col, num_channels_, num_samples_),
|
||||
next_output_timestamp_, cc));
|
||||
timestamps.push_back(next_output_timestamp_);
|
||||
next_output_timestamp_ += round(frame_step_ / target_sample_rate_ *
|
||||
Timestamp::kTimestampUnitsPerSecond);
|
||||
next_frame_first_col += frame_step_;
|
||||
}
|
||||
}
|
||||
if (should_flush && next_frame_first_col < buffer.cols()) {
|
||||
ASSIGN_OR_RETURN(auto output_tensor,
|
||||
ConvertToTensor(buffer.block(
|
||||
0, next_frame_first_col, num_channels_,
|
||||
std::min(num_samples_,
|
||||
(int)buffer.cols() - next_frame_first_col))));
|
||||
// In the streaming mode, the flush happens in Close() and a packet at
|
||||
// Timestamp::Max() will be emitted. In the non-streaming mode, each
|
||||
// Process() invocation will process the entire buffer completely.
|
||||
Timestamp timestamp =
|
||||
stream_mode_ ? Timestamp::Max() : next_output_timestamp_;
|
||||
Timestamp timestamp = should_flush_at_timestamp_max
|
||||
? Timestamp::Max()
|
||||
: next_output_timestamp_;
|
||||
MP_RETURN_IF_ERROR(OutputTensor(
|
||||
buffer.block(
|
||||
0, next_frame_first_col, num_channels_,
|
||||
std::min(num_samples_, (int)buffer.cols() - next_frame_first_col)),
|
||||
timestamp, cc));
|
||||
timestamps.push_back(timestamp);
|
||||
kTensorsOut(cc).Send(std::move(output_tensor), timestamp);
|
||||
}
|
||||
if (kTimestampsOut(cc).IsConnected()) {
|
||||
Timestamp timestamp = timestamps.back();
|
||||
|
||||
@@ -44,4 +44,28 @@ message AudioToTensorCalculatorOptions {
|
||||
// Set to false to disable checks for jitter in timestamp values. Useful with
|
||||
// live audio input.
|
||||
optional bool check_inconsistent_timestamps = 6 [default = true];
|
||||
|
||||
// Size of the fft in number of bins. If set, the calculator outputs fft
|
||||
// tensors.
|
||||
optional int64 fft_size = 7;
|
||||
|
||||
// The amount of padding samples to add before the audio after resampling.
|
||||
// Note that the timestamps shift. Currently, only zero padding is supported.
|
||||
optional int64 padding_samples_before = 8;
|
||||
|
||||
// The amount of padding samples to add after the audio after resampling.
|
||||
// Currently, only zero padding is supported.
|
||||
optional int64 padding_samples_after = 9;
|
||||
|
||||
// Determines the "flushing" behavior in stream mode.
|
||||
enum FlushMode {
|
||||
// Unspecified (causes an error). Won't be used because of the default.
|
||||
NONE = 0;
|
||||
// Emit a packet with the entire remainder at `Timestamp::Max`.
|
||||
ENTIRE_TAIL_AT_TIMESTAMP_MAX = 1;
|
||||
// Continue emitting framed packets with relevant timestamps.
|
||||
PROCEED_AS_USUAL = 2;
|
||||
}
|
||||
|
||||
optional FlushMode flush_mode = 10 [default = ENTIRE_TAIL_AT_TIMESTAMP_MAX];
|
||||
}
|
||||
|
||||
@@ -12,13 +12,13 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <cmath>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/strings/substitute.h"
|
||||
#include "audio/dsp/resampler_q.h"
|
||||
#include "mediapipe/calculators/tensor/audio_to_tensor_calculator.pb.h"
|
||||
#include "mediapipe/framework/api2/packet.h"
|
||||
#include "mediapipe/framework/calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
@@ -32,6 +32,14 @@
|
||||
namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
using ::testing::Not;
|
||||
using Options = ::mediapipe::AudioToTensorCalculatorOptions;
|
||||
using FlushMode = Options::FlushMode;
|
||||
|
||||
int DivideRoundedUp(int dividend, int divisor) {
|
||||
return (dividend + divisor - 1) / divisor;
|
||||
}
|
||||
|
||||
std::unique_ptr<Matrix> CreateTestMatrix(int num_channels, int num_samples,
|
||||
int timestamp) {
|
||||
auto matrix = std::make_unique<Matrix>(num_channels, num_samples);
|
||||
@@ -292,16 +300,17 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
|
||||
num_iterations_ = num_iterations;
|
||||
}
|
||||
|
||||
int GetExpectedNumOfSamples() {
|
||||
Matrix* expected_matrix =
|
||||
resampled_buffer_ ? resampled_buffer_.get() : sample_buffer_.get();
|
||||
return expected_matrix->cols();
|
||||
}
|
||||
int GetExpectedNumOfSamples() { return output_sample_buffer_->cols(); }
|
||||
|
||||
void Run(int num_samples, int num_overlapping_samples,
|
||||
double resampling_factor) {
|
||||
double resampling_factor, int padding_before = 0,
|
||||
int padding_after = 0, bool expect_init_error = false) {
|
||||
double input_sample_rate = 10000;
|
||||
double target_sample_rate = input_sample_rate * resampling_factor;
|
||||
FlushMode flush_mode = (padding_before != 0 || padding_after != 0)
|
||||
? Options::PROCEED_AS_USUAL
|
||||
: Options::ENTIRE_TAIL_AT_TIMESTAMP_MAX;
|
||||
|
||||
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
absl::Substitute(R"(
|
||||
input_stream: "audio"
|
||||
@@ -319,16 +328,25 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
|
||||
num_overlapping_samples: $1
|
||||
target_sample_rate: $2
|
||||
stream_mode:true
|
||||
padding_samples_before: $3
|
||||
padding_samples_after: $4
|
||||
flush_mode: $5
|
||||
}
|
||||
}
|
||||
}
|
||||
)",
|
||||
/*$0=*/num_samples, /*$1=*/num_overlapping_samples,
|
||||
/*$2=*/target_sample_rate));
|
||||
/*$2=*/target_sample_rate, /*$3=*/padding_before,
|
||||
/*$4=*/padding_after, /*$5=*/flush_mode));
|
||||
tool::AddVectorSink("tensors", &graph_config, &tensors_packets_);
|
||||
|
||||
// Run the graph.
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config));
|
||||
const absl::Status init_status = graph_.Initialize(graph_config);
|
||||
if (expect_init_error) {
|
||||
EXPECT_THAT(init_status, Not(IsOk()));
|
||||
return;
|
||||
}
|
||||
MP_ASSERT_OK(init_status);
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
for (int i = 0; i < num_iterations_; ++i) {
|
||||
Timestamp input_timestamp(Timestamp::kTimestampUnitsPerSecond * i);
|
||||
@@ -345,8 +363,18 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
|
||||
}
|
||||
MP_ASSERT_OK(graph_.CloseAllInputStreams());
|
||||
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||
if (resampling_factor != 1) {
|
||||
resampled_buffer_ = ResampleBuffer(*sample_buffer_, resampling_factor);
|
||||
if (resampling_factor == 1) {
|
||||
output_sample_buffer_ = std::make_unique<Matrix>(*sample_buffer_);
|
||||
} else {
|
||||
output_sample_buffer_ =
|
||||
ResampleBuffer(*sample_buffer_, resampling_factor);
|
||||
}
|
||||
if (padding_before != 0 || padding_after != 0) {
|
||||
Matrix padded = Matrix::Zero(
|
||||
2, padding_before + output_sample_buffer_->cols() + padding_after);
|
||||
padded.block(0, padding_before, 2, output_sample_buffer_->cols()) =
|
||||
*output_sample_buffer_;
|
||||
output_sample_buffer_->swap(padded);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -372,15 +400,13 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
|
||||
auto buffer = output_tensor.GetCpuReadView().buffer<float>();
|
||||
int num_values = output_tensor.shape().num_elements();
|
||||
std::vector<float> output_floats(buffer, buffer + num_values);
|
||||
Matrix* expected_matrix =
|
||||
resampled_buffer_ ? resampled_buffer_.get() : sample_buffer_.get();
|
||||
for (int i = 0; i < num_values; ++i) {
|
||||
if (i + sample_offset >= expected_matrix->size()) {
|
||||
if (i + sample_offset >= output_sample_buffer_->size()) {
|
||||
EXPECT_FLOAT_EQ(output_floats[i], 0);
|
||||
} else {
|
||||
EXPECT_NEAR(output_floats[i],
|
||||
expected_matrix->coeff((i + sample_offset) % 2,
|
||||
(i + sample_offset) / 2),
|
||||
output_sample_buffer_->coeff((i + sample_offset) % 2,
|
||||
(i + sample_offset) / 2),
|
||||
0.001)
|
||||
<< "i=" << i << ", sample_offset=" << sample_offset
|
||||
<< ", packet index=" << index;
|
||||
@@ -391,7 +417,8 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
|
||||
|
||||
// Fully close graph at end, otherwise calculator+tensors are destroyed
|
||||
// after calling WaitUntilDone().
|
||||
void CloseGraph() { MP_EXPECT_OK(graph_.WaitUntilDone()); }
|
||||
absl::Status TryCloseGraph() { return graph_.WaitUntilDone(); }
|
||||
void CloseGraph() { MP_EXPECT_OK(TryCloseGraph()); }
|
||||
|
||||
private:
|
||||
int input_buffer_num_samples_ = 10;
|
||||
@@ -399,7 +426,7 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
|
||||
CalculatorGraph graph_;
|
||||
std::vector<Packet> tensors_packets_;
|
||||
std::unique_ptr<Matrix> sample_buffer_;
|
||||
std::unique_ptr<Matrix> resampled_buffer_;
|
||||
std::unique_ptr<Matrix> output_sample_buffer_;
|
||||
};
|
||||
|
||||
TEST_F(AudioToTensorCalculatorStreamingModeTest,
|
||||
@@ -408,7 +435,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest,
|
||||
/*resampling_factor=*/1.0f);
|
||||
CheckTensorsOutputPackets(
|
||||
/*sample_offset=*/10,
|
||||
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 5),
|
||||
/*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 5),
|
||||
/*timestamp_interval=*/500,
|
||||
/*output_last_at_close=*/false);
|
||||
CloseGraph();
|
||||
@@ -419,7 +446,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, OutputRemainingInCloseMethod) {
|
||||
/*resampling_factor=*/1.0f);
|
||||
CheckTensorsOutputPackets(
|
||||
/*sample_offset=*/12,
|
||||
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 6),
|
||||
/*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 6),
|
||||
/*timestamp_interval=*/600,
|
||||
/*output_last_at_close=*/true);
|
||||
CloseGraph();
|
||||
@@ -431,7 +458,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, OutputOverlappingFp32Tensors) {
|
||||
/*resampling_factor=*/1.0f);
|
||||
CheckTensorsOutputPackets(
|
||||
/*sample_offset=*/16,
|
||||
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 8),
|
||||
/*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 8),
|
||||
/*timestamp_interval=*/800,
|
||||
/*output_last_at_close=*/true);
|
||||
CloseGraph();
|
||||
@@ -443,7 +470,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, Downsampling) {
|
||||
/*resampling_factor=*/0.5f);
|
||||
CheckTensorsOutputPackets(
|
||||
/*sample_offset=*/512,
|
||||
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 256),
|
||||
/*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 256),
|
||||
/*timestamp_interval=*/51200,
|
||||
/*output_last_at_close=*/true);
|
||||
CloseGraph();
|
||||
@@ -455,7 +482,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, DownsamplingWithOverlapping) {
|
||||
/*resampling_factor=*/0.5f);
|
||||
CheckTensorsOutputPackets(
|
||||
/*sample_offset=*/384,
|
||||
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 192),
|
||||
/*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 192),
|
||||
/*timestamp_interval=*/38400,
|
||||
/*output_last_at_close=*/true);
|
||||
CloseGraph();
|
||||
@@ -467,7 +494,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, Upsampling) {
|
||||
/*resampling_factor=*/2.0f);
|
||||
CheckTensorsOutputPackets(
|
||||
/*sample_offset=*/512,
|
||||
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 256),
|
||||
/*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 256),
|
||||
/*timestamp_interval=*/12800,
|
||||
/*output_last_at_close=*/true);
|
||||
CloseGraph();
|
||||
@@ -479,12 +506,33 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, UpsamplingWithOverlapping) {
|
||||
/*resampling_factor=*/2.0f);
|
||||
CheckTensorsOutputPackets(
|
||||
/*sample_offset=*/384,
|
||||
/*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 192),
|
||||
/*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 192),
|
||||
/*timestamp_interval=*/9600,
|
||||
/*output_last_at_close=*/true);
|
||||
CloseGraph();
|
||||
}
|
||||
|
||||
TEST_F(AudioToTensorCalculatorStreamingModeTest,
|
||||
UpsamplingWithOverlappingAndPadding) {
|
||||
SetInputBufferNumSamplesPerChannel(1024);
|
||||
Run(/*num_samples=*/256, /*num_overlapping_samples=*/64,
|
||||
/*resampling_factor=*/2.0f, /*padding_before=*/13, /*padding_after=*/999);
|
||||
CheckTensorsOutputPackets(
|
||||
/*sample_offset=*/384,
|
||||
/*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 192),
|
||||
/*timestamp_interval=*/9600,
|
||||
/*output_last_at_close=*/false);
|
||||
CloseGraph();
|
||||
}
|
||||
|
||||
TEST_F(AudioToTensorCalculatorStreamingModeTest, NegativePaddingUnsupported) {
|
||||
SetInputBufferNumSamplesPerChannel(1024);
|
||||
Run(/*num_samples=*/256, /*num_overlapping_samples=*/64,
|
||||
/*resampling_factor=*/2.0f, /*padding_before=*/13, /*padding_after=*/-3,
|
||||
/*expect_init_error=*/true);
|
||||
EXPECT_THAT(TryCloseGraph(), Not(IsOk()));
|
||||
}
|
||||
|
||||
TEST_F(AudioToTensorCalculatorStreamingModeTest,
|
||||
OnlyOutputInCloseIfNoSufficientSamples) {
|
||||
SetNumIterations(1);
|
||||
@@ -498,5 +546,122 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest,
|
||||
CloseGraph();
|
||||
}
|
||||
|
||||
class AudioToTensorCalculatorFftTest : public ::testing::Test {
|
||||
protected:
|
||||
// Creates an audio matrix containing a single sample of 1.0 at a specified
|
||||
// offset.
|
||||
std::unique_ptr<Matrix> CreateImpulseSignalData(int64 num_samples,
|
||||
int impulse_offset_idx) {
|
||||
Matrix impulse = Matrix::Zero(1, num_samples);
|
||||
impulse(0, impulse_offset_idx) = 1.0;
|
||||
return std::make_unique<Matrix>(std::move(impulse));
|
||||
}
|
||||
|
||||
void ConfigGraph(int num_channels, int num_samples,
|
||||
int num_overlapping_samples, double sample_rate,
|
||||
int fft_size) {
|
||||
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
absl::Substitute(R"(
|
||||
input_stream: "audio"
|
||||
input_stream: "sample_rate"
|
||||
output_stream: "tensors"
|
||||
output_stream: "dc_and_nyquist"
|
||||
node {
|
||||
calculator: "AudioToTensorCalculator"
|
||||
input_stream: "AUDIO:audio"
|
||||
input_stream: "SAMPLE_RATE:sample_rate"
|
||||
output_stream: "TENSORS:tensors"
|
||||
output_stream: "DC_AND_NYQUIST:dc_and_nyquist"
|
||||
options {
|
||||
[mediapipe.AudioToTensorCalculatorOptions.ext] {
|
||||
num_channels: $0
|
||||
num_samples: $1
|
||||
num_overlapping_samples: $2
|
||||
target_sample_rate: $3
|
||||
fft_size: $4
|
||||
}
|
||||
}
|
||||
}
|
||||
)",
|
||||
/*$0=*/num_channels,
|
||||
/*$1=*/num_samples,
|
||||
/*$2=*/num_overlapping_samples,
|
||||
/*$3=*/sample_rate, /*$4=*/fft_size));
|
||||
std::vector<Packet> tensors_packets;
|
||||
tool::AddVectorSink("tensors", &graph_config_, &tensors_packets_);
|
||||
std::vector<Packet> dc_and_nyquist_packets;
|
||||
tool::AddVectorSink("dc_and_nyquist", &graph_config_,
|
||||
&dc_and_nyquist_packets_);
|
||||
}
|
||||
|
||||
void RunGraph(std::unique_ptr<Matrix> input_data, double sample_rate) {
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config_));
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"sample_rate", MakePacket<double>(sample_rate).At(Timestamp(0))));
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"audio", MakePacket<Matrix>(*input_data).At(Timestamp(0))));
|
||||
MP_ASSERT_OK(graph_.CloseAllInputStreams());
|
||||
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||
ASSERT_EQ(tensors_packets_.size(), dc_and_nyquist_packets_.size());
|
||||
}
|
||||
|
||||
// Fully close graph at end, otherwise calculator+tensors are destroyed
|
||||
// after calling WaitUntilDone().
|
||||
void CloseGraph() { MP_EXPECT_OK(graph_.WaitUntilDone()); }
|
||||
|
||||
std::vector<Packet> tensors_packets_;
|
||||
std::vector<Packet> dc_and_nyquist_packets_;
|
||||
CalculatorGraphConfig graph_config_;
|
||||
CalculatorGraph graph_;
|
||||
};
|
||||
|
||||
TEST_F(AudioToTensorCalculatorFftTest, TestInvalidFftSize) {
|
||||
ConfigGraph(1, 320, 160, 16000, 103);
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config_));
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
auto status = graph_.WaitUntilIdle();
|
||||
EXPECT_EQ(status.code(), absl::StatusCode::kInternal);
|
||||
EXPECT_THAT(status.message(),
|
||||
::testing::HasSubstr("FFT size must be of the form"));
|
||||
}
|
||||
|
||||
TEST_F(AudioToTensorCalculatorFftTest, TestInvalidNumChannels) {
|
||||
ConfigGraph(3, 320, 160, 16000, 256);
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config_));
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
auto status = graph_.WaitUntilIdle();
|
||||
EXPECT_EQ(status.code(), absl::StatusCode::kInternal);
|
||||
EXPECT_THAT(
|
||||
status.message(),
|
||||
::testing::HasSubstr("only support applying FFT on mono channel"));
|
||||
}
|
||||
|
||||
TEST_F(AudioToTensorCalculatorFftTest, TestImpulseSignal) {
|
||||
constexpr double sample_rate = 16000;
|
||||
ConfigGraph(1, 320, 160, sample_rate, 320);
|
||||
RunGraph(CreateImpulseSignalData(320, 160), sample_rate);
|
||||
for (int i = 0; i < tensors_packets_.size(); ++i) {
|
||||
const auto& tensors = tensors_packets_[i].Get<std::vector<Tensor>>();
|
||||
ASSERT_EQ(1, tensors.size());
|
||||
const Tensor& output_tensor =
|
||||
tensors_packets_[0].Get<std::vector<Tensor>>()[0];
|
||||
auto* buffer = output_tensor.GetCpuReadView().buffer<float>();
|
||||
int num_values = output_tensor.shape().num_elements();
|
||||
const std::vector<float> output_floats(buffer, buffer + num_values);
|
||||
// Impulse signal should have (approximately) const power across all
|
||||
// frequency bins.
|
||||
const auto& pair =
|
||||
dc_and_nyquist_packets_[i].Get<std::pair<float, float>>();
|
||||
EXPECT_FLOAT_EQ(pair.first, 1.0f);
|
||||
EXPECT_FLOAT_EQ(pair.second, 1.0f);
|
||||
for (int j = 0; j < num_values / 2; ++j) {
|
||||
std::complex<float> cf(output_floats[j * 2], output_floats[j * 2 + 1]);
|
||||
EXPECT_FLOAT_EQ(std::norm(cf), 1.0f);
|
||||
}
|
||||
}
|
||||
CloseGraph();
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
// Copyright 2022 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "absl/status/status.h"
|
||||
#include "mediapipe/calculators/tensor/feedback_tensors_calculator.pb.h"
|
||||
#include "mediapipe/framework/api2/node.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/tensor.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace api2 {
|
||||
|
||||
namespace {
|
||||
constexpr char kInputTensorsTag[] = "INPUT_TENSORS";
|
||||
constexpr char kFeedbackTensorsTag[] = "FEEDBACK_TENSORS";
|
||||
constexpr char kOutputTensorsTag[] = "TENSORS";
|
||||
|
||||
using Tensors = std::vector<Tensor>;
|
||||
} // namespace
|
||||
|
||||
// FeedbackTensorsCalculator groups the input and the feedback (typically
|
||||
// recurrent neural network cell state output tensors from the previous run)
|
||||
// tensor vectors as the input tensor vector for the next recurrent model cell
|
||||
// inference. On the first step, the feedback tensor is filled with zeros to
|
||||
// jumpstart the loop.
|
||||
class FeedbackTensorsCalculator : public Node {
|
||||
public:
|
||||
static constexpr Input<Tensors> kFeedbackTensorsIn{kFeedbackTensorsTag};
|
||||
static constexpr Input<Tensors> kInputTensorsIn{kInputTensorsTag};
|
||||
static constexpr Output<Tensors> kTensorsOut{kOutputTensorsTag};
|
||||
|
||||
MEDIAPIPE_NODE_CONTRACT(kFeedbackTensorsIn, kInputTensorsIn, kTensorsOut);
|
||||
|
||||
static absl::Status GetContract(CalculatorContract* cc) {
|
||||
cc->SetProcessTimestampBounds(true);
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status Open(CalculatorContext* cc) override {
|
||||
const auto& options =
|
||||
cc->Options<mediapipe::FeedbackTensorsCalculatorOptions>();
|
||||
|
||||
const auto& shape_dims = options.feedback_tensor_shape().dims();
|
||||
feedback_tensor_shape_.dims.assign(shape_dims.begin(), shape_dims.end());
|
||||
feedback_tensor_size_ = feedback_tensor_shape_.num_elements();
|
||||
|
||||
num_feedback_tensors_ = options.num_feedback_tensors();
|
||||
|
||||
feedback_tensors_location_ = options.location();
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status Process(CalculatorContext* cc) override {
|
||||
if (feedback_tensors_location_ ==
|
||||
mediapipe::FeedbackTensorsCalculatorOptions::NONE) {
|
||||
kTensorsOut(cc).Send(kInputTensorsIn(cc).packet().As<Tensors>());
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
std::vector<Tensor> outputs;
|
||||
switch (feedback_tensors_location_) {
|
||||
case mediapipe::FeedbackTensorsCalculatorOptions::PREPENDED:
|
||||
MP_RETURN_IF_ERROR(AddFeedbackTensors(cc, outputs));
|
||||
MP_RETURN_IF_ERROR(AddInputTensors(cc, outputs));
|
||||
break;
|
||||
case mediapipe::FeedbackTensorsCalculatorOptions::APPENDED:
|
||||
MP_RETURN_IF_ERROR(AddInputTensors(cc, outputs));
|
||||
MP_RETURN_IF_ERROR(AddFeedbackTensors(cc, outputs));
|
||||
break;
|
||||
default:
|
||||
return absl::InvalidArgumentError(
|
||||
"Unsupported feedback tensors location");
|
||||
}
|
||||
kTensorsOut(cc).Send(std::move(outputs));
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
absl::Status AddInputTensors(CalculatorContext* cc,
|
||||
std::vector<Tensor>& outputs) {
|
||||
absl::StatusOr<std::unique_ptr<std::vector<Tensor>>> input_tensors =
|
||||
cc->Inputs()
|
||||
.Tag(kInputTensorsTag)
|
||||
.Value()
|
||||
.Consume<std::vector<Tensor>>();
|
||||
if (!input_tensors.ok()) {
|
||||
return absl::InternalError("The input tensors packet is not consumable");
|
||||
}
|
||||
RET_CHECK(*input_tensors);
|
||||
std::vector<Tensor>& inputs = **input_tensors;
|
||||
outputs.insert(outputs.end(), std::make_move_iterator(inputs.begin()),
|
||||
std::make_move_iterator(inputs.end()));
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status AddFeedbackTensors(CalculatorContext* cc,
|
||||
std::vector<Tensor>& outputs) {
|
||||
if (first_run_) {
|
||||
for (int index = 0; index < num_feedback_tensors_; ++index) {
|
||||
Tensor initial_feedback_tensor(Tensor::ElementType::kFloat32,
|
||||
feedback_tensor_shape_);
|
||||
float* data = initial_feedback_tensor.GetCpuWriteView().buffer<float>();
|
||||
std::fill_n(data, feedback_tensor_size_, 0.0f);
|
||||
outputs.push_back(std::move(initial_feedback_tensor));
|
||||
}
|
||||
first_run_ = false;
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
if (num_feedback_tensors_ != kFeedbackTensorsIn(cc)->size()) {
|
||||
return absl::InvalidArgumentError(
|
||||
"The number of tensors fed back differs from the configuration");
|
||||
}
|
||||
absl::StatusOr<std::unique_ptr<std::vector<Tensor>>> feedback_tensors =
|
||||
cc->Inputs()
|
||||
.Tag(kFeedbackTensorsTag)
|
||||
.Value()
|
||||
.Consume<std::vector<Tensor>>();
|
||||
if (!feedback_tensors.ok()) {
|
||||
return absl::InternalError(
|
||||
"The feedback tensors packet is not consumable");
|
||||
}
|
||||
RET_CHECK(*feedback_tensors);
|
||||
std::vector<Tensor>& feedbacks = **feedback_tensors;
|
||||
for (const auto& feedback : feedbacks) {
|
||||
if (feedback.shape().dims != feedback_tensor_shape_.dims) {
|
||||
return absl::InvalidArgumentError(
|
||||
"The shape of a tensor fed back differs from the configuration");
|
||||
}
|
||||
}
|
||||
outputs.insert(outputs.end(), std::make_move_iterator(feedbacks.begin()),
|
||||
std::make_move_iterator(feedbacks.end()));
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
Tensor::Shape feedback_tensor_shape_;
|
||||
int num_feedback_tensors_ = 0;
|
||||
mediapipe::FeedbackTensorsCalculatorOptions::FeedbackTensorsLocation
|
||||
feedback_tensors_location_;
|
||||
|
||||
int feedback_tensor_size_ = 0;
|
||||
bool first_run_ = true;
|
||||
};
|
||||
|
||||
MEDIAPIPE_REGISTER_NODE(FeedbackTensorsCalculator);
|
||||
|
||||
} // namespace api2
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,47 @@
|
||||
// Copyright 2022 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
|
||||
message FeedbackTensorsCalculatorOptions {
|
||||
extend mediapipe.CalculatorOptions {
|
||||
optional FeedbackTensorsCalculatorOptions ext = 474496252;
|
||||
}
|
||||
|
||||
// Represents the dimensions of a tensor starting from the outermost size.
|
||||
message TensorShape {
|
||||
repeated int32 dims = 1 [packed = true];
|
||||
}
|
||||
|
||||
// The shape of the feedback tensors to add.
|
||||
optional TensorShape feedback_tensor_shape = 1;
|
||||
// The number of the feedback tensors to add.
|
||||
optional int32 num_feedback_tensors = 2 [default = 1];
|
||||
|
||||
enum FeedbackTensorsLocation {
|
||||
// The feedback tensors will not be added.
|
||||
NONE = 0;
|
||||
// The feedback tensors will be added before the input tensors.
|
||||
PREPENDED = 1;
|
||||
// The feedback tensors will be added after the input tensors.
|
||||
APPENDED = 2;
|
||||
}
|
||||
|
||||
// Determines the location of the feedback tensor(s) in the output vector.
|
||||
optional FeedbackTensorsLocation location = 3 [default = APPENDED];
|
||||
}
|
||||
@@ -0,0 +1,389 @@
|
||||
// Copyright 2022 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <functional>
|
||||
#include <initializer_list>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/calculators/tensor/feedback_tensors_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/tensor.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
using ::mediapipe::CalculatorGraphConfig;
|
||||
using ::testing::ElementsAreArray;
|
||||
using ::testing::Not;
|
||||
using Tensors = std::vector<Tensor>;
|
||||
|
||||
template <typename T>
|
||||
struct TensorElementType {
|
||||
static constexpr Tensor::ElementType value = Tensor::ElementType::kNone;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct TensorElementType<float> {
|
||||
static constexpr Tensor::ElementType value = Tensor::ElementType::kFloat32;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct TensorElementType<std::int8_t> {
|
||||
static constexpr Tensor::ElementType value = Tensor::ElementType::kInt8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct TensorElementType<std::uint8_t> {
|
||||
static constexpr Tensor::ElementType value = Tensor::ElementType::kUInt8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct TensorElementType<std::int32_t> {
|
||||
static constexpr Tensor::ElementType value = Tensor::ElementType::kInt32;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
Tensor MakeTensor(std::initializer_list<int> shape,
|
||||
std::initializer_list<T> values) {
|
||||
Tensor tensor(TensorElementType<T>::value, shape);
|
||||
CHECK_EQ(values.size(), tensor.shape().num_elements())
|
||||
<< "The size of `values` is incompatible with `shape`";
|
||||
absl::c_copy(values, tensor.GetCpuWriteView().buffer<T>());
|
||||
return tensor;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void ValidateTensor(const Tensor& tensor,
|
||||
const std::vector<int>& expected_shape,
|
||||
const std::vector<T>& expected_values) {
|
||||
ASSERT_EQ(tensor.element_type(), TensorElementType<T>::value);
|
||||
EXPECT_EQ(tensor.shape().dims, expected_shape);
|
||||
EXPECT_EQ(tensor.shape().num_elements(), expected_values.size());
|
||||
|
||||
auto* tensor_buffer = tensor.GetCpuReadView().buffer<T>();
|
||||
const std::vector<T> tensor_values(
|
||||
tensor_buffer, tensor_buffer + tensor.shape().num_elements());
|
||||
EXPECT_THAT(tensor_values, ElementsAreArray(expected_values));
|
||||
}
|
||||
|
||||
TEST(FeedbackTensorsCalculatorTest, AppendsFeedback) {
|
||||
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
|
||||
input_stream: "input"
|
||||
input_stream: "feedback"
|
||||
node {
|
||||
calculator: "FeedbackTensorsCalculator"
|
||||
input_stream: "INPUT_TENSORS:input"
|
||||
input_stream: "FEEDBACK_TENSORS:feedback"
|
||||
output_stream: "TENSORS:output"
|
||||
options: {
|
||||
[mediapipe.FeedbackTensorsCalculatorOptions.ext] {
|
||||
feedback_tensor_shape: { dims: 2 dims: 3 }
|
||||
location: APPENDED
|
||||
}
|
||||
}
|
||||
}
|
||||
)pb");
|
||||
std::vector<Packet> output_packets;
|
||||
tool::AddVectorSink("output", &graph_config, &output_packets);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
auto initial_input_tensors = std::make_unique<Tensors>();
|
||||
initial_input_tensors->push_back(
|
||||
MakeTensor<std::int32_t>({2, 4}, {1, 2, 3, 4, 5, 6, 7, 8}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(initial_input_tensors.release()).At(Timestamp(1))));
|
||||
// At the beginning, the loopback packet with the model feedback is missing.
|
||||
// The calculator has to assume it's all-zero with the shape from the options.
|
||||
|
||||
auto later_input_tensors = std::make_unique<Tensors>();
|
||||
later_input_tensors->push_back(
|
||||
MakeTensor<std::int32_t>({2, 4}, {8, 7, 6, 5, 4, 3, 2, 1}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(later_input_tensors.release()).At(Timestamp(2))));
|
||||
auto later_feedback_tensors = std::make_unique<Tensors>();
|
||||
later_feedback_tensors->push_back(
|
||||
MakeTensor({2, 3}, {-1.f, -2.f, -3.f, -4.f, -5.f, -6.f}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"feedback", Adopt(later_feedback_tensors.release()).At(Timestamp(2))));
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllInputStreams())
|
||||
<< "Couldn't close the graph inputs";
|
||||
MP_ASSERT_OK(graph.WaitUntilDone()) << "Couldn't finalize the graph run";
|
||||
|
||||
ASSERT_EQ(output_packets.size(), 2);
|
||||
|
||||
const Tensors& initial_combined_tensors = output_packets[0].Get<Tensors>();
|
||||
ASSERT_EQ(initial_combined_tensors.size(), 2);
|
||||
ValidateTensor<std::int32_t>(initial_combined_tensors[0],
|
||||
/*expected_shape=*/{2, 4},
|
||||
/*expected_values=*/{1, 2, 3, 4, 5, 6, 7, 8});
|
||||
// The initial feedback is zero.
|
||||
ValidateTensor<float>(initial_combined_tensors[1], /*expected_shape=*/{2, 3},
|
||||
/*expected_values=*/{0.f, 0.f, 0.f, 0.f, 0.f, 0.f});
|
||||
|
||||
const Tensors& later_combined_tensors = output_packets[1].Get<Tensors>();
|
||||
ASSERT_EQ(later_combined_tensors.size(), 2);
|
||||
ValidateTensor<std::int32_t>(later_combined_tensors[0],
|
||||
/*expected_shape=*/{2, 4},
|
||||
/*expected_values=*/{8, 7, 6, 5, 4, 3, 2, 1});
|
||||
// Afterwards, the provided feedback is passed through.
|
||||
ValidateTensor<float>(
|
||||
later_combined_tensors[1], /*expected_shape=*/{2, 3},
|
||||
/*expected_values=*/{-1.f, -2.f, -3.f, -4.f, -5.f, -6.f});
|
||||
}
|
||||
|
||||
TEST(FeedbackTensorsCalculatorTest, PrependsFeedback) {
|
||||
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
|
||||
input_stream: "input"
|
||||
input_stream: "feedback"
|
||||
node {
|
||||
calculator: "FeedbackTensorsCalculator"
|
||||
input_stream: "INPUT_TENSORS:input"
|
||||
input_stream: "FEEDBACK_TENSORS:feedback"
|
||||
output_stream: "TENSORS:output"
|
||||
options: {
|
||||
[mediapipe.FeedbackTensorsCalculatorOptions.ext] {
|
||||
feedback_tensor_shape: { dims: 3 dims: 2 }
|
||||
location: PREPENDED
|
||||
}
|
||||
}
|
||||
}
|
||||
)pb");
|
||||
std::vector<Packet> output_packets;
|
||||
tool::AddVectorSink("output", &graph_config, &output_packets);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
auto initial_input_tensors = std::make_unique<Tensors>();
|
||||
initial_input_tensors->push_back(
|
||||
MakeTensor<std::int8_t>({2, 4}, {1, 2, 3, 4, 5, 6, 7, 8}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(initial_input_tensors.release()).At(Timestamp(1))));
|
||||
// At the beginning, the loopback packet with the model feedback is missing.
|
||||
// The calculator has to assume it's all-zero with the shape from the options.
|
||||
|
||||
auto later_input_tensors = std::make_unique<Tensors>();
|
||||
later_input_tensors->push_back(
|
||||
MakeTensor<std::int8_t>({2, 4}, {8, 7, 6, 5, 4, 3, 2, 1}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(later_input_tensors.release()).At(Timestamp(2))));
|
||||
auto later_feedback_tensors = std::make_unique<Tensors>();
|
||||
later_feedback_tensors->push_back(
|
||||
MakeTensor({3, 2}, {-1.f, -2.f, -3.f, -4.f, -5.f, -6.f}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"feedback", Adopt(later_feedback_tensors.release()).At(Timestamp(2))));
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllInputStreams())
|
||||
<< "Couldn't close the graph inputs";
|
||||
MP_ASSERT_OK(graph.WaitUntilDone()) << "Couldn't finalize the graph run";
|
||||
|
||||
ASSERT_EQ(output_packets.size(), 2);
|
||||
|
||||
const Tensors& initial_combined_tensors = output_packets[0].Get<Tensors>();
|
||||
ASSERT_EQ(initial_combined_tensors.size(), 2);
|
||||
// The initial feedback is zero.
|
||||
ValidateTensor<float>(initial_combined_tensors[0], /*expected_shape=*/{3, 2},
|
||||
/*expected_values=*/{0.f, 0.f, 0.f, 0.f, 0.f, 0.f});
|
||||
ValidateTensor<std::int8_t>(initial_combined_tensors[1],
|
||||
/*expected_shape=*/{2, 4},
|
||||
/*expected_values=*/{1, 2, 3, 4, 5, 6, 7, 8});
|
||||
|
||||
const Tensors& later_combined_tensors = output_packets[1].Get<Tensors>();
|
||||
ASSERT_EQ(later_combined_tensors.size(), 2);
|
||||
// Afterwards, the provided feedback is passed through.
|
||||
ValidateTensor<float>(
|
||||
later_combined_tensors[0], /*expected_shape=*/{3, 2},
|
||||
/*expected_values=*/{-1.f, -2.f, -3.f, -4.f, -5.f, -6.f});
|
||||
ValidateTensor<std::int8_t>(later_combined_tensors[1],
|
||||
/*expected_shape=*/{2, 4},
|
||||
/*expected_values=*/{8, 7, 6, 5, 4, 3, 2, 1});
|
||||
}
|
||||
|
||||
TEST(FeedbackTensorsCalculatorTest, NoFeedback) {
|
||||
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
|
||||
input_stream: "input"
|
||||
input_stream: "feedback"
|
||||
node {
|
||||
calculator: "FeedbackTensorsCalculator"
|
||||
input_stream: "INPUT_TENSORS:input"
|
||||
input_stream: "FEEDBACK_TENSORS:feedback"
|
||||
output_stream: "TENSORS:output"
|
||||
options: {
|
||||
[mediapipe.FeedbackTensorsCalculatorOptions.ext] {
|
||||
feedback_tensor_shape: { dims: 3 dims: 4 }
|
||||
location: NONE
|
||||
}
|
||||
}
|
||||
}
|
||||
)pb");
|
||||
std::vector<Packet> output_packets;
|
||||
tool::AddVectorSink("output", &graph_config, &output_packets);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
auto initial_input_tensors = std::make_unique<Tensors>();
|
||||
initial_input_tensors->push_back(
|
||||
MakeTensor<std::uint8_t>({2, 4}, {1, 2, 3, 4, 5, 6, 7, 8}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(initial_input_tensors.release()).At(Timestamp(1))));
|
||||
// At the beginning, the loopback packet with the model feedback is missing.
|
||||
|
||||
auto later_input_tensors = std::make_unique<Tensors>();
|
||||
later_input_tensors->push_back(
|
||||
MakeTensor<std::uint8_t>({2, 4}, {8, 7, 6, 5, 4, 3, 2, 1}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(later_input_tensors.release()).At(Timestamp(2))));
|
||||
// This feedback should be ignored due to `location: NONE`.
|
||||
auto later_feedback_tensors = std::make_unique<Tensors>();
|
||||
later_feedback_tensors->push_back(
|
||||
MakeTensor({2, 3}, {-1.f, -2.f, -3.f, -4.f, -5.f, -6.f}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"feedback", Adopt(later_feedback_tensors.release()).At(Timestamp(2))));
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllInputStreams())
|
||||
<< "Couldn't close the graph inputs";
|
||||
MP_ASSERT_OK(graph.WaitUntilDone()) << "Couldn't finalize the graph run";
|
||||
|
||||
ASSERT_EQ(output_packets.size(), 2);
|
||||
|
||||
const Tensors& initial_combined_tensors = output_packets[0].Get<Tensors>();
|
||||
ASSERT_EQ(initial_combined_tensors.size(), 1);
|
||||
ValidateTensor<std::uint8_t>(initial_combined_tensors[0],
|
||||
/*expected_shape=*/{2, 4},
|
||||
/*expected_values=*/{1, 2, 3, 4, 5, 6, 7, 8});
|
||||
// No feedback due to `location: NONE`.
|
||||
|
||||
const Tensors& later_combined_tensors = output_packets[1].Get<Tensors>();
|
||||
ASSERT_EQ(later_combined_tensors.size(), 1);
|
||||
ValidateTensor<std::uint8_t>(later_combined_tensors[0],
|
||||
/*expected_shape=*/{2, 4},
|
||||
/*expected_values=*/{8, 7, 6, 5, 4, 3, 2, 1});
|
||||
}
|
||||
|
||||
TEST(FeedbackTensorsCalculatorTest, ChecksTensorNumber) {
|
||||
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
|
||||
input_stream: "input"
|
||||
input_stream: "feedback"
|
||||
node {
|
||||
calculator: "FeedbackTensorsCalculator"
|
||||
input_stream: "INPUT_TENSORS:input"
|
||||
input_stream: "FEEDBACK_TENSORS:feedback"
|
||||
output_stream: "TENSORS:output"
|
||||
options: {
|
||||
[mediapipe.FeedbackTensorsCalculatorOptions.ext] {
|
||||
num_feedback_tensors: 2
|
||||
feedback_tensor_shape: { dims: 2 dims: 3 }
|
||||
location: PREPENDED
|
||||
}
|
||||
}
|
||||
}
|
||||
)pb");
|
||||
std::vector<Packet> output_packets;
|
||||
tool::AddVectorSink("output", &graph_config, &output_packets);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
auto initial_input_tensors = std::make_unique<Tensors>();
|
||||
initial_input_tensors->push_back(
|
||||
MakeTensor<std::uint8_t>({2, 4}, {1, 2, 3, 4, 5, 6, 7, 8}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(initial_input_tensors.release()).At(Timestamp(1))));
|
||||
// At the beginning, the loopback packet with the model feedback is missing.
|
||||
|
||||
auto later_input_tensors = std::make_unique<Tensors>();
|
||||
later_input_tensors->push_back(
|
||||
MakeTensor<std::uint8_t>({2, 4}, {8, 7, 6, 5, 4, 3, 2, 1}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(later_input_tensors.release()).At(Timestamp(2))));
|
||||
// This feedback should be ignored due to `location: NONE`.
|
||||
auto later_feedback_tensors = std::make_unique<Tensors>();
|
||||
later_feedback_tensors->push_back(
|
||||
MakeTensor({2, 3}, {-1.f, -2.f, -3.f, -4.f, -5.f, -6.f}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"feedback", Adopt(later_feedback_tensors.release()).At(Timestamp(2))));
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllInputStreams())
|
||||
<< "Couldn't close the graph inputs";
|
||||
EXPECT_THAT(graph.WaitUntilDone(), Not(IsOk()))
|
||||
<< "Tensor number mismatch missed";
|
||||
}
|
||||
|
||||
TEST(FeedbackTensorsCalculatorTest, ChecksShape) {
|
||||
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
|
||||
input_stream: "input"
|
||||
input_stream: "feedback"
|
||||
node {
|
||||
calculator: "FeedbackTensorsCalculator"
|
||||
input_stream: "INPUT_TENSORS:input"
|
||||
input_stream: "FEEDBACK_TENSORS:feedback"
|
||||
output_stream: "TENSORS:output"
|
||||
options: {
|
||||
[mediapipe.FeedbackTensorsCalculatorOptions.ext] {
|
||||
feedback_tensor_shape: { dims: 3 dims: 4 }
|
||||
location: APPENDED
|
||||
}
|
||||
}
|
||||
}
|
||||
)pb");
|
||||
std::vector<Packet> output_packets;
|
||||
tool::AddVectorSink("output", &graph_config, &output_packets);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
auto initial_input_tensors = std::make_unique<Tensors>();
|
||||
initial_input_tensors->push_back(
|
||||
MakeTensor<std::uint8_t>({2, 4}, {1, 2, 3, 4, 5, 6, 7, 8}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(initial_input_tensors.release()).At(Timestamp(1))));
|
||||
// At the beginning, the loopback packet with the model feedback is missing.
|
||||
|
||||
auto later_input_tensors = std::make_unique<Tensors>();
|
||||
later_input_tensors->push_back(
|
||||
MakeTensor<std::uint8_t>({2, 4}, {8, 7, 6, 5, 4, 3, 2, 1}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(later_input_tensors.release()).At(Timestamp(2))));
|
||||
// This feedback should be ignored due to `location: NONE`.
|
||||
auto later_feedback_tensors = std::make_unique<Tensors>();
|
||||
later_feedback_tensors->push_back(
|
||||
MakeTensor({2, 3}, {-1.f, -2.f, -3.f, -4.f, -5.f, -6.f}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"feedback", Adopt(later_feedback_tensors.release()).At(Timestamp(2))));
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllInputStreams())
|
||||
<< "Couldn't close the graph inputs";
|
||||
EXPECT_THAT(graph.WaitUntilDone(), Not(IsOk()))
|
||||
<< "Tensor shape mismatch missed";
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
+1
-1
@@ -231,7 +231,7 @@ TEST_F(TensorFlowSessionFromSavedModelCalculatorTest,
|
||||
// Session must be set.
|
||||
ASSERT_NE(session.session, nullptr);
|
||||
std::vector<tensorflow::DeviceAttributes> devices;
|
||||
ASSERT_EQ(session.session->ListDevices(&devices), tensorflow::Status::OK());
|
||||
ASSERT_EQ(session.session->ListDevices(&devices), tensorflow::OkStatus());
|
||||
EXPECT_THAT(devices.size(), 10);
|
||||
}
|
||||
|
||||
|
||||
+1
-1
@@ -220,7 +220,7 @@ TEST_F(TensorFlowSessionFromSavedModelGeneratorTest,
|
||||
// Session must be set.
|
||||
ASSERT_NE(session.session, nullptr);
|
||||
std::vector<tensorflow::DeviceAttributes> devices;
|
||||
ASSERT_EQ(session.session->ListDevices(&devices), tensorflow::Status::OK());
|
||||
ASSERT_EQ(session.session->ListDevices(&devices), tensorflow::OkStatus());
|
||||
EXPECT_THAT(devices.size(), 10);
|
||||
}
|
||||
|
||||
|
||||
@@ -135,6 +135,7 @@ filegroup(
|
||||
srcs = [
|
||||
"testdata/anchor_golden_file_0.txt",
|
||||
"testdata/anchor_golden_file_1.txt",
|
||||
"testdata/anchor_golden_file_2.txt",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
// limitations under the License.
|
||||
|
||||
#include <cmath>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/calculators/tflite/ssd_anchors_calculator.pb.h"
|
||||
@@ -24,6 +25,19 @@ namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
struct MultiScaleAnchorInfo {
|
||||
int32 level;
|
||||
std::vector<float> aspect_ratios;
|
||||
std::vector<float> scales;
|
||||
std::pair<float, float> base_anchor_size;
|
||||
std::pair<float, float> anchor_stride;
|
||||
};
|
||||
|
||||
struct FeatureMapDim {
|
||||
int height;
|
||||
int width;
|
||||
};
|
||||
|
||||
float CalculateScale(float min_scale, float max_scale, int stride_index,
|
||||
int num_strides) {
|
||||
if (num_strides == 1) {
|
||||
@@ -34,6 +48,71 @@ float CalculateScale(float min_scale, float max_scale, int stride_index,
|
||||
}
|
||||
}
|
||||
|
||||
int GetNumLayers(const SsdAnchorsCalculatorOptions& options) {
|
||||
if (options.multiscale_anchor_generation()) {
|
||||
return (options.max_level() - options.min_level() + 1);
|
||||
}
|
||||
return options.num_layers();
|
||||
}
|
||||
|
||||
FeatureMapDim GetFeatureMapDimensions(
|
||||
const SsdAnchorsCalculatorOptions& options, int index) {
|
||||
FeatureMapDim feature_map_dims;
|
||||
if (options.feature_map_height_size()) {
|
||||
feature_map_dims.height = options.feature_map_height(index);
|
||||
feature_map_dims.width = options.feature_map_width(index);
|
||||
} else {
|
||||
const int stride = options.strides(index);
|
||||
feature_map_dims.height =
|
||||
std::ceil(1.0f * options.input_size_height() / stride);
|
||||
feature_map_dims.width =
|
||||
std::ceil(1.0f * options.input_size_width() / stride);
|
||||
}
|
||||
return feature_map_dims;
|
||||
}
|
||||
|
||||
// Although we have stride for both x and y, only one value is used for offset
|
||||
// calculation. See
|
||||
// tensorflow_models/object_detection/anchor_generators/multiscale_grid_anchor_generator.py;l=121
|
||||
std::pair<float, float> GetMultiScaleAnchorOffset(
|
||||
const SsdAnchorsCalculatorOptions& options, const float stride,
|
||||
const int level) {
|
||||
std::pair<float, float> result(0., 0.);
|
||||
int denominator = std::pow(2, level);
|
||||
if (options.input_size_height() % denominator == 0 ||
|
||||
options.input_size_height() == 1) {
|
||||
result.first = stride / 2.0;
|
||||
}
|
||||
if (options.input_size_width() % denominator == 0 ||
|
||||
options.input_size_width() == 1) {
|
||||
result.second = stride / 2.0;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
void NormalizeAnchor(const int input_height, const int input_width,
|
||||
Anchor* anchor) {
|
||||
anchor->set_h(anchor->h() / (float)input_height);
|
||||
anchor->set_w(anchor->w() / (float)input_width);
|
||||
anchor->set_y_center(anchor->y_center() / (float)input_height);
|
||||
anchor->set_x_center(anchor->x_center() / (float)input_width);
|
||||
}
|
||||
|
||||
Anchor CalculateAnchorBox(const int y_center, const int x_center,
|
||||
const float scale, const float aspect_ratio,
|
||||
const std::pair<float, float> base_anchor_size,
|
||||
// y-height first
|
||||
const std::pair<float, float> anchor_stride,
|
||||
const std::pair<float, float> anchor_offset) {
|
||||
Anchor result;
|
||||
float ratio_sqrt = std::sqrt(aspect_ratio);
|
||||
result.set_h(scale * base_anchor_size.first / ratio_sqrt);
|
||||
result.set_w(scale * ratio_sqrt * base_anchor_size.second);
|
||||
result.set_y_center(y_center * anchor_stride.first + anchor_offset.first);
|
||||
result.set_x_center(x_center * anchor_stride.second + anchor_offset.second);
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// Generate anchors for SSD object detection model.
|
||||
@@ -95,9 +174,77 @@ class SsdAnchorsCalculator : public CalculatorBase {
|
||||
private:
|
||||
static absl::Status GenerateAnchors(
|
||||
std::vector<Anchor>* anchors, const SsdAnchorsCalculatorOptions& options);
|
||||
|
||||
static absl::Status GenerateMultiScaleAnchors(
|
||||
std::vector<Anchor>* anchors, const SsdAnchorsCalculatorOptions& options);
|
||||
};
|
||||
REGISTER_CALCULATOR(SsdAnchorsCalculator);
|
||||
|
||||
// Generates grid anchors on the fly corresponding to multiple CNN layers as
|
||||
// described in:
|
||||
// "Focal Loss for Dense Object Detection" (https://arxiv.org/abs/1708.02002)
|
||||
// T.-Y. Lin, P. Goyal, R. Girshick, K. He, P. Dollar
|
||||
absl::Status SsdAnchorsCalculator::GenerateMultiScaleAnchors(
|
||||
std::vector<Anchor>* anchors, const SsdAnchorsCalculatorOptions& options) {
|
||||
std::vector<MultiScaleAnchorInfo> anchor_infos;
|
||||
for (int i = options.min_level(); i <= options.max_level(); ++i) {
|
||||
MultiScaleAnchorInfo current_anchor_info;
|
||||
// level
|
||||
current_anchor_info.level = i;
|
||||
// aspect_ratios
|
||||
for (const float aspect_ratio : options.aspect_ratios()) {
|
||||
current_anchor_info.aspect_ratios.push_back(aspect_ratio);
|
||||
}
|
||||
|
||||
// scale
|
||||
for (int i = 0; i < options.scales_per_octave(); ++i) {
|
||||
current_anchor_info.scales.push_back(
|
||||
std::pow(2.0, (double)i / (double)options.scales_per_octave()));
|
||||
}
|
||||
|
||||
// anchor stride
|
||||
float anchor_stride = std::pow(2.0, i);
|
||||
current_anchor_info.anchor_stride =
|
||||
std::make_pair(anchor_stride, anchor_stride);
|
||||
|
||||
// base_anchor_size
|
||||
current_anchor_info.base_anchor_size =
|
||||
std::make_pair(anchor_stride * options.anchor_scale(),
|
||||
anchor_stride * options.anchor_scale());
|
||||
anchor_infos.push_back(current_anchor_info);
|
||||
}
|
||||
|
||||
for (unsigned int i = 0; i < anchor_infos.size(); ++i) {
|
||||
FeatureMapDim dimensions = GetFeatureMapDimensions(options, i);
|
||||
for (int y = 0; y < dimensions.height; ++y) {
|
||||
for (int x = 0; x < dimensions.width; ++x) {
|
||||
// loop over combination of scale and aspect ratio
|
||||
for (unsigned int j = 0; j < anchor_infos[i].aspect_ratios.size();
|
||||
++j) {
|
||||
for (unsigned int k = 0; k < anchor_infos[i].scales.size(); ++k) {
|
||||
Anchor anchor = CalculateAnchorBox(
|
||||
/*y_center=*/y, /*x_center=*/x, anchor_infos[i].scales[k],
|
||||
anchor_infos[i].aspect_ratios[j],
|
||||
anchor_infos[i].base_anchor_size,
|
||||
/*anchor_stride=*/anchor_infos[i].anchor_stride,
|
||||
/*anchor_offset=*/
|
||||
GetMultiScaleAnchorOffset(options,
|
||||
anchor_infos[i].anchor_stride.first,
|
||||
anchor_infos[i].level));
|
||||
if (options.normalize_coordinates()) {
|
||||
NormalizeAnchor(options.input_size_height(),
|
||||
options.input_size_width(), &anchor);
|
||||
}
|
||||
anchors->push_back(anchor);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SsdAnchorsCalculator::GenerateAnchors(
|
||||
std::vector<Anchor>* anchors, const SsdAnchorsCalculatorOptions& options) {
|
||||
// Verify the options.
|
||||
@@ -106,15 +253,21 @@ absl::Status SsdAnchorsCalculator::GenerateAnchors(
|
||||
"Both feature map shape and strides are missing. Must provide either "
|
||||
"one.");
|
||||
}
|
||||
const int kNumLayers = GetNumLayers(options);
|
||||
|
||||
if (options.feature_map_height_size()) {
|
||||
if (options.strides_size()) {
|
||||
LOG(ERROR) << "Found feature map shapes. Strides will be ignored.";
|
||||
}
|
||||
CHECK_EQ(options.feature_map_height_size(), options.num_layers());
|
||||
CHECK_EQ(options.feature_map_height_size(), kNumLayers);
|
||||
CHECK_EQ(options.feature_map_height_size(),
|
||||
options.feature_map_width_size());
|
||||
} else {
|
||||
CHECK_EQ(options.strides_size(), options.num_layers());
|
||||
CHECK_EQ(options.strides_size(), kNumLayers);
|
||||
}
|
||||
|
||||
if (options.multiscale_anchor_generation()) {
|
||||
return GenerateMultiScaleAnchors(anchors, options);
|
||||
}
|
||||
|
||||
int layer_id = 0;
|
||||
|
||||
@@ -60,4 +60,30 @@ message SsdAnchorsCalculatorOptions {
|
||||
// This option can be used when the predicted anchor width and height are in
|
||||
// pixels.
|
||||
optional bool fixed_anchor_size = 14 [default = false];
|
||||
|
||||
// Generates grid anchors on the fly corresponding to multiple CNN layers as
|
||||
// described in:
|
||||
// "Focal Loss for Dense Object Detection" (https://arxiv.org/abs/1708.02002)
|
||||
// T.-Y. Lin, P. Goyal, R. Girshick, K. He, P. Dollar
|
||||
optional bool multiscale_anchor_generation = 15 [default = false];
|
||||
|
||||
// minimum level in feature pyramid
|
||||
// for multiscale_anchor_generation only!
|
||||
optional int32 min_level = 16 [default = 3];
|
||||
|
||||
// maximum level in feature pyramid
|
||||
// for multiscale_anchor_generation only!
|
||||
optional int32 max_level = 17 [default = 7];
|
||||
|
||||
// Scale of anchor to feature stride
|
||||
// for multiscale_anchor_generation only!
|
||||
optional float anchor_scale = 18 [default = 4.0];
|
||||
|
||||
// Number of intermediate scale each scale octave
|
||||
// for multiscale_anchor_generation only!
|
||||
optional int32 scales_per_octave = 19 [default = 2];
|
||||
|
||||
// Whether to produce anchors in normalized coordinates.
|
||||
// for multiscale_anchor_generation only!
|
||||
optional bool normalize_coordinates = 20 [default = true];
|
||||
}
|
||||
|
||||
@@ -33,9 +33,6 @@ std::string GetGoldenFilePath(const std::string& filename) {
|
||||
|
||||
void ParseAnchorsFromText(const std::string& text,
|
||||
std::vector<Anchor>* anchors) {
|
||||
const std::string line_delimiter = "\n";
|
||||
const std::string number_delimiter = ",";
|
||||
|
||||
std::istringstream stream(text);
|
||||
std::string line;
|
||||
while (std::getline(stream, line)) {
|
||||
@@ -64,6 +61,8 @@ void CompareAnchors(const std::vector<Anchor>& anchors_0,
|
||||
testing::FloatNear(anchor_1.x_center(), 1e-5));
|
||||
EXPECT_THAT(anchor_0.y_center(),
|
||||
testing::FloatNear(anchor_1.y_center(), 1e-5));
|
||||
EXPECT_THAT(anchor_0.h(), testing::FloatNear(anchor_1.h(), 1e-5));
|
||||
EXPECT_THAT(anchor_0.w(), testing::FloatNear(anchor_1.w(), 1e-5));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -148,4 +147,40 @@ TEST(SsdAnchorCalculatorTest, MobileSSDConfig) {
|
||||
CompareAnchors(anchors, anchors_golden);
|
||||
}
|
||||
|
||||
TEST(SsdAnchorCalculatorTest, RetinaNetSSDConfig) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"pb(
|
||||
calculator: "SsdAnchorsCalculator"
|
||||
output_side_packet: "anchors"
|
||||
options {
|
||||
[mediapipe.SsdAnchorsCalculatorOptions.ext] {
|
||||
input_size_height: 640
|
||||
input_size_width: 640
|
||||
strides: 64
|
||||
strides: 128
|
||||
aspect_ratios: 1.0
|
||||
aspect_ratios: 2.0
|
||||
aspect_ratios: 0.5
|
||||
multiscale_anchor_generation: true
|
||||
min_level: 6
|
||||
max_level: 7
|
||||
anchor_scale: 3.0
|
||||
scales_per_octave: 3
|
||||
}
|
||||
}
|
||||
)pb"));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const auto& anchors =
|
||||
runner.OutputSidePackets().Index(0).Get<std::vector<Anchor>>();
|
||||
|
||||
std::string anchors_string;
|
||||
MP_EXPECT_OK(mediapipe::file::GetContents(
|
||||
GetGoldenFilePath("anchor_golden_file_2.txt"), &anchors_string));
|
||||
|
||||
std::vector<Anchor> anchors_golden;
|
||||
ParseAnchorsFromText(anchors_string, &anchors_golden);
|
||||
|
||||
CompareAnchors(anchors, anchors_golden);
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -19,6 +19,7 @@
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "tensorflow/lite/allocation.h"
|
||||
#include "tensorflow/lite/model.h"
|
||||
|
||||
namespace mediapipe {
|
||||
@@ -32,6 +33,8 @@ namespace mediapipe {
|
||||
// it to the graph as input side packet or you can use some of
|
||||
// calculators like LocalFileContentsCalculator to get model
|
||||
// blob and use it as input here.
|
||||
// MODEL_FD - Tflite model file descriptor std::tuple<int, size_t, size_t>
|
||||
// containing (fd, offset, size).
|
||||
//
|
||||
// Output side packets:
|
||||
// MODEL - TfLite model. (std::unique_ptr<tflite::FlatBufferModel,
|
||||
@@ -52,17 +55,42 @@ class TfLiteModelCalculator : public CalculatorBase {
|
||||
std::function<void(tflite::FlatBufferModel*)>>;
|
||||
|
||||
static absl::Status GetContract(CalculatorContract* cc) {
|
||||
cc->InputSidePackets().Tag("MODEL_BLOB").Set<std::string>();
|
||||
if (cc->InputSidePackets().HasTag("MODEL_BLOB")) {
|
||||
cc->InputSidePackets().Tag("MODEL_BLOB").Set<std::string>();
|
||||
}
|
||||
|
||||
if (cc->InputSidePackets().HasTag("MODEL_FD")) {
|
||||
cc->InputSidePackets()
|
||||
.Tag("MODEL_FD")
|
||||
.Set<std::tuple<int, size_t, size_t>>();
|
||||
}
|
||||
|
||||
cc->OutputSidePackets().Tag("MODEL").Set<TfLiteModelPtr>();
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status Open(CalculatorContext* cc) override {
|
||||
const Packet& model_packet = cc->InputSidePackets().Tag("MODEL_BLOB");
|
||||
const std::string& model_blob = model_packet.Get<std::string>();
|
||||
std::unique_ptr<tflite::FlatBufferModel> model =
|
||||
tflite::FlatBufferModel::BuildFromBuffer(model_blob.data(),
|
||||
model_blob.size());
|
||||
Packet model_packet;
|
||||
std::unique_ptr<tflite::FlatBufferModel> model;
|
||||
|
||||
if (cc->InputSidePackets().HasTag("MODEL_BLOB")) {
|
||||
model_packet = cc->InputSidePackets().Tag("MODEL_BLOB");
|
||||
const std::string& model_blob = model_packet.Get<std::string>();
|
||||
model = tflite::FlatBufferModel::BuildFromBuffer(model_blob.data(),
|
||||
model_blob.size());
|
||||
}
|
||||
|
||||
if (cc->InputSidePackets().HasTag("MODEL_FD")) {
|
||||
model_packet = cc->InputSidePackets().Tag("MODEL_FD");
|
||||
const auto& model_fd =
|
||||
model_packet.Get<std::tuple<int, size_t, size_t>>();
|
||||
auto model_allocation = std::make_unique<tflite::MMAPAllocation>(
|
||||
std::get<0>(model_fd), std::get<1>(model_fd), std::get<2>(model_fd),
|
||||
tflite::DefaultErrorReporter());
|
||||
model = tflite::FlatBufferModel::BuildFromAllocation(
|
||||
std::move(model_allocation), tflite::DefaultErrorReporter());
|
||||
}
|
||||
|
||||
RET_CHECK(model) << "Failed to load TfLite model from blob.";
|
||||
|
||||
cc->OutputSidePackets().Tag("MODEL").Set(
|
||||
|
||||
Reference in New Issue
Block a user