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

Merge branch 'google:master' into image-embedder-python

Browse Source
This commit is contained in:
Kinar R
2022-11-05 00:17:26 +05:30
committed by GitHub
parent e2d50745ac 5e1a2fcdbb
commit b1d1ef0124
79 changed files with 2407 additions and 669 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/tasks/cc/audio/audio_classifier/BUILD
+1
View File
@@ -65,6 +65,7 @@ cc_library(
"//mediapipe/tasks/cc/audio/core:audio_task_api_factory",
"//mediapipe/tasks/cc/audio/core:base_audio_task_api",
"//mediapipe/tasks/cc/audio/core:running_mode",
"//mediapipe/tasks/cc/components/containers:classification_result",
"//mediapipe/tasks/cc/components/containers/proto:classifications_cc_proto",
"//mediapipe/tasks/cc/components/processors:classifier_options",
"//mediapipe/tasks/cc/components/processors/proto:classifier_options_cc_proto",
mediapipe/tasks/cc/audio/audio_classifier/audio_classifier.cc
+35 -10
View File
@@ -18,12 +18,14 @@ limitations under the License.
#include <map>
#include <memory>
#include <utility>
#include <vector>
#include "absl/status/statusor.h"
#include "mediapipe/framework/api2/builder.h"
#include "mediapipe/framework/formats/matrix.h"
#include "mediapipe/tasks/cc/audio/audio_classifier/proto/audio_classifier_graph_options.pb.h"
#include "mediapipe/tasks/cc/audio/core/audio_task_api_factory.h"
#include "mediapipe/tasks/cc/components/containers/classification_result.h"
#include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h"
#include "mediapipe/tasks/cc/components/processors/classifier_options.h"
#include "mediapipe/tasks/cc/components/processors/proto/classifier_options.pb.h"
@@ -38,12 +40,16 @@ namespace audio_classifier {
namespace {
using ::mediapipe::tasks::components::containers::ConvertToClassificationResult;
using ::mediapipe::tasks::components::containers::proto::ClassificationResult;
constexpr char kAudioStreamName[] = "audio_in";
constexpr char kAudioTag[] = "AUDIO";
constexpr char kClassificationResultStreamName[] = "classification_result_out";
constexpr char kClassificationResultTag[] = "CLASSIFICATION_RESULT";
constexpr char kClassificationsTag[] = "CLASSIFICATIONS";
constexpr char kClassificationsName[] = "classifications_out";
constexpr char kTimestampedClassificationsTag[] = "TIMESTAMPED_CLASSIFICATIONS";
constexpr char kTimestampedClassificationsName[] =
"timestamped_classifications_out";
constexpr char kSampleRateName[] = "sample_rate_in";
constexpr char kSampleRateTag[] = "SAMPLE_RATE";
constexpr char kSubgraphTypeName[] =
@@ -63,9 +69,11 @@ CalculatorGraphConfig CreateGraphConfig(
}
subgraph.GetOptions<proto::AudioClassifierGraphOptions>().Swap(
options_proto.get());
subgraph.Out(kClassificationResultTag)
.SetName(kClassificationResultStreamName) >>
graph.Out(kClassificationResultTag);
subgraph.Out(kClassificationsTag).SetName(kClassificationsName) >>
graph.Out(kClassificationsTag);
subgraph.Out(kTimestampedClassificationsTag)
.SetName(kTimestampedClassificationsName) >>
graph.Out(kTimestampedClassificationsTag);
return graph.GetConfig();
}
@@ -91,13 +99,30 @@ ConvertAudioClassifierOptionsToProto(AudioClassifierOptions* options) {
return options_proto;
}
absl::StatusOr<ClassificationResult> ConvertOutputPackets(
absl::StatusOr<std::vector<AudioClassifierResult>> ConvertOutputPackets(
absl::StatusOr<tasks::core::PacketMap> status_or_packets) {
if (!status_or_packets.ok()) {
return status_or_packets.status();
}
return status_or_packets.value()[kClassificationResultStreamName]
.Get<ClassificationResult>();
auto classification_results =
status_or_packets.value()[kTimestampedClassificationsName]
.Get<std::vector<ClassificationResult>>();
std::vector<AudioClassifierResult> results;
results.reserve(classification_results.size());
for (const auto& classification_result : classification_results) {
results.emplace_back(ConvertToClassificationResult(classification_result));
}
return results;
}
absl::StatusOr<AudioClassifierResult> ConvertAsyncOutputPackets(
absl::StatusOr<tasks::core::PacketMap> status_or_packets) {
if (!status_or_packets.ok()) {
return status_or_packets.status();
}
return ConvertToClassificationResult(
status_or_packets.value()[kClassificationsName]
.Get<ClassificationResult>());
}
} // namespace
@@ -118,7 +143,7 @@ absl::StatusOr<std::unique_ptr<AudioClassifier>> AudioClassifier::Create(
auto result_callback = options->result_callback;
packets_callback =
[=](absl::StatusOr<tasks::core::PacketMap> status_or_packets) {
result_callback(ConvertOutputPackets(status_or_packets));
result_callback(ConvertAsyncOutputPackets(status_or_packets));
};
}
return core::AudioTaskApiFactory::Create<AudioClassifier,
@@ -128,7 +153,7 @@ absl::StatusOr<std::unique_ptr<AudioClassifier>> AudioClassifier::Create(
std::move(packets_callback));
}
absl::StatusOr<ClassificationResult> AudioClassifier::Classify(
absl::StatusOr<std::vector<AudioClassifierResult>> AudioClassifier::Classify(
Matrix audio_clip, double audio_sample_rate) {
return ConvertOutputPackets(ProcessAudioClip(
{{kAudioStreamName, MakePacket<Matrix>(std::move(audio_clip))},
mediapipe/tasks/cc/audio/audio_classifier/audio_classifier.h
+40 -30
View File
@@ -18,12 +18,13 @@ limitations under the License.
#include <memory>
#include <utility>
#include <vector>
#include "absl/status/statusor.h"
#include "mediapipe/framework/formats/matrix.h"
#include "mediapipe/tasks/cc/audio/core/base_audio_task_api.h"
#include "mediapipe/tasks/cc/audio/core/running_mode.h"
#include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h"
#include "mediapipe/tasks/cc/components/containers/classification_result.h"
#include "mediapipe/tasks/cc/components/processors/classifier_options.h"
#include "mediapipe/tasks/cc/core/base_options.h"
@@ -32,6 +33,10 @@ namespace tasks {
namespace audio {
namespace audio_classifier {
// Alias the shared ClassificationResult struct as result type.
using AudioClassifierResult =
::mediapipe::tasks::components::containers::ClassificationResult;
// The options for configuring a mediapipe audio classifier task.
struct AudioClassifierOptions {
// Base options for configuring Task library, such as specifying the TfLite
@@ -59,9 +64,8 @@ struct AudioClassifierOptions {
// The user-defined result callback for processing audio stream data.
// The result callback should only be specified when the running mode is set
// to RunningMode::AUDIO_STREAM.
std::function<void(
absl::StatusOr<components::containers::proto::ClassificationResult>)>
result_callback = nullptr;
std::function<void(absl::StatusOr<AudioClassifierResult>)> result_callback =
nullptr;
};
// Performs audio classification on audio clips or audio stream.
@@ -117,23 +121,36 @@ class AudioClassifier : tasks::audio::core::BaseAudioTaskApi {
// required to provide the corresponding audio sample rate along with the
// input audio clips.
//
// For each audio clip, the output classifications are grouped in a
// ClassificationResult object that has three dimensions:
// Classification head:
// The prediction heads targeting different audio classification tasks
// such as audio event classification and bird sound classification.
// Classification timestamp:
// The start time (in milliseconds) of each audio clip that is sent to the
// model for audio classification. As the audio classification models take
// a fixed number of audio samples, long audio clips will be framed to
// multiple buffers (with the desired number of audio samples) during
// preprocessing.
// Classification category:
// The list of the classification categories that model predicts per
// framed audio clip.
// The input audio clip may be longer than what the model is able to process
// in a single inference. When this occurs, the input audio clip is split into
// multiple chunks starting at different timestamps. For this reason, this
// function returns a vector of ClassificationResult objects, each associated
// with a timestamp corresponding to the start (in milliseconds) of the chunk
// data that was classified, e.g:
//
// ClassificationResult #0 (first chunk of data):
// timestamp_ms: 0 (starts at 0ms)
// classifications #0 (single head model):
// category #0:
// category_name: "Speech"
// score: 0.6
// category #1:
// category_name: "Music"
// score: 0.2
// ClassificationResult #1 (second chunk of data):
// timestamp_ms: 800 (starts at 800ms)
// classifications #0 (single head model):
// category #0:
// category_name: "Speech"
// score: 0.5
// category #1:
// category_name: "Silence"
// score: 0.1
// ...
//
// TODO: Use `sample_rate` in AudioClassifierOptions by default
// and makes `audio_sample_rate` optional.
absl::StatusOr<components::containers::proto::ClassificationResult> Classify(
absl::StatusOr<std::vector<AudioClassifierResult>> Classify(
mediapipe::Matrix audio_clip, double audio_sample_rate);
// Sends audio data (a block in a continuous audio stream) to perform audio
@@ -147,17 +164,10 @@ class AudioClassifier : tasks::audio::core::BaseAudioTaskApi {
// milliseconds) to indicate the start time of the input audio block. The
// timestamps must be monotonically increasing.
//
// The output classifications are grouped in a ClassificationResult object
// that has three dimensions:
// Classification head:
// The prediction heads targeting different audio classification tasks
// such as audio event classification and bird sound classification.
// Classification timestamp :
// The start time (in milliseconds) of the framed audio block that is sent
// to the model for audio classification.
// Classification category:
// The list of the classification categories that model predicts per
// framed audio clip.
// The input audio block may be longer than what the model is able to process
// in a single inference. When this occurs, the input audio block is split
// into multiple chunks. For this reason, the callback may be called multiple
// times (once per chunk) for each call to this function.
absl::Status ClassifyAsync(mediapipe::Matrix audio_block, int64 timestamp_ms);
// Shuts down the AudioClassifier when all works are done.
mediapipe/tasks/cc/audio/audio_classifier/audio_classifier_graph.cc
+34 -14
View File
@@ -16,6 +16,7 @@ limitations under the License.
#include <stdint.h>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
@@ -57,12 +58,20 @@ using ::mediapipe::tasks::components::containers::proto::ClassificationResult;
constexpr char kAtPrestreamTag[] = "AT_PRESTREAM";
constexpr char kAudioTag[] = "AUDIO";
constexpr char kClassificationResultTag[] = "CLASSIFICATION_RESULT";
constexpr char kClassificationsTag[] = "CLASSIFICATIONS";
constexpr char kTimestampedClassificationsTag[] = "TIMESTAMPED_CLASSIFICATIONS";
constexpr char kPacketTag[] = "PACKET";
constexpr char kSampleRateTag[] = "SAMPLE_RATE";
constexpr char kTensorsTag[] = "TENSORS";
constexpr char kTimestampsTag[] = "TIMESTAMPS";
// Struct holding the different output streams produced by the audio classifier
// graph.
struct AudioClassifierOutputStreams {
Source<ClassificationResult> classifications;
Source<std::vector<ClassificationResult>> timestamped_classifications;
};
absl::Status SanityCheckOptions(
const proto::AudioClassifierGraphOptions& options) {
if (options.base_options().use_stream_mode() &&
@@ -124,16 +133,20 @@ void ConfigureAudioToTensorCalculator(
// series stream header with sample rate info.
//
// Outputs:
// CLASSIFICATION_RESULT - ClassificationResult
// The aggregated classification result object that has 3 dimensions:
// (classification head, classification timestamp, classification category).
// CLASSIFICATIONS - ClassificationResult @Optional
// The classification results aggregated by head. Only produces results if
// the graph if the 'use_stream_mode' option is true.
// TIMESTAMPED_CLASSIFICATIONS - std::vector<ClassificationResult> @Optional
// The classification result aggregated by timestamp, then by head. Only
// produces results if the graph if the 'use_stream_mode' option is false.
//
// Example:
// node {
// calculator: "mediapipe.tasks.audio.audio_classifier.AudioClassifierGraph"
// input_stream: "AUDIO:audio_in"
// input_stream: "SAMPLE_RATE:sample_rate_in"
// output_stream: "CLASSIFICATION_RESULT:classification_result_out"
// output_stream: "CLASSIFICATIONS:classifications"
// output_stream: "TIMESTAMPED_CLASSIFICATIONS:timestamped_classifications"
// options {
// [mediapipe.tasks.audio.audio_classifier.proto.AudioClassifierGraphOptions.ext]
// {
@@ -162,7 +175,7 @@ class AudioClassifierGraph : public core::ModelTaskGraph {
.base_options()
.use_stream_mode();
ASSIGN_OR_RETURN(
auto classification_result_out,
auto output_streams,
BuildAudioClassificationTask(
sc->Options<proto::AudioClassifierGraphOptions>(), *model_resources,
graph[Input<Matrix>(kAudioTag)],
@@ -170,8 +183,11 @@ class AudioClassifierGraph : public core::ModelTaskGraph {
? absl::nullopt
: absl::make_optional(graph[Input<double>(kSampleRateTag)]),
graph));
classification_result_out >>
graph[Output<ClassificationResult>(kClassificationResultTag)];
output_streams.classifications >>
graph[Output<ClassificationResult>(kClassificationsTag)];
output_streams.timestamped_classifications >>
graph[Output<std::vector<ClassificationResult>>(
kTimestampedClassificationsTag)];
return graph.GetConfig();
}
@@ -187,7 +203,7 @@ class AudioClassifierGraph : public core::ModelTaskGraph {
// audio_in: (mediapipe::Matrix) stream to run audio classification on.
// sample_rate_in: (double) optional stream of the input audio sample rate.
// graph: the mediapipe builder::Graph instance to be updated.
absl::StatusOr<Source<ClassificationResult>> BuildAudioClassificationTask(
absl::StatusOr<AudioClassifierOutputStreams> BuildAudioClassificationTask(
const proto::AudioClassifierGraphOptions& task_options,
const core::ModelResources& model_resources, Source<Matrix> audio_in,
absl::optional<Source<double>> sample_rate_in, Graph& graph) {
@@ -250,16 +266,20 @@ class AudioClassifierGraph : public core::ModelTaskGraph {
inference.Out(kTensorsTag) >> postprocessing.In(kTensorsTag);
// Time aggregation is only needed for performing audio classification on
// audio files. Disables time aggregration by not connecting the
// audio files. Disables timestamp aggregation by not connecting the
// "TIMESTAMPS" streams.
if (!use_stream_mode) {
audio_to_tensor.Out(kTimestampsTag) >> postprocessing.In(kTimestampsTag);
}
// Outputs the aggregated classification result as the subgraph output
// stream.
return postprocessing[Output<ClassificationResult>(
kClassificationResultTag)];
// Output both streams as graph output streams/
return AudioClassifierOutputStreams{
/*classifications=*/postprocessing[Output<ClassificationResult>(
kClassificationsTag)],
/*timestamped_classifications=*/
postprocessing[Output<std::vector<ClassificationResult>>(
kTimestampedClassificationsTag)],
};
}
};
mediapipe/tasks/cc/audio/audio_classifier/audio_classifier_test.cc
+112 -124
View File
@@ -32,13 +32,11 @@ limitations under the License.
#include "mediapipe/framework/formats/matrix.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/tasks/cc/audio/core/running_mode.h"
#include "mediapipe/tasks/cc/audio/utils/test_utils.h"
#include "mediapipe/tasks/cc/common.h"
#include "mediapipe/tasks/cc/components/containers/proto/category.pb.h"
#include "mediapipe/tasks/cc/components/containers/proto/classifications.pb.h"
#include "mediapipe/tasks/cc/components/containers/category.h"
#include "mediapipe/tasks/cc/components/containers/classification_result.h"
#include "tensorflow/lite/core/shims/cc/shims_test_util.h"
namespace mediapipe {
@@ -49,7 +47,6 @@ namespace {
using ::absl::StatusOr;
using ::mediapipe::file::JoinPath;
using ::mediapipe::tasks::components::containers::proto::ClassificationResult;
using ::testing::HasSubstr;
using ::testing::Optional;
@@ -73,95 +70,86 @@ Matrix GetAudioData(absl::string_view filename) {
return matrix_mapping.matrix();
}
void CheckSpeechClassificationResult(const ClassificationResult& result) {
EXPECT_THAT(result.classifications_size(), testing::Eq(1));
EXPECT_EQ(result.classifications(0).head_name(), "scores");
EXPECT_EQ(result.classifications(0).head_index(), 0);
EXPECT_THAT(result.classifications(0).entries_size(), testing::Eq(5));
void CheckSpeechResult(const std::vector<AudioClassifierResult>& result,
int expected_num_categories = 521) {
EXPECT_EQ(result.size(), 5);
// Ignore last result, which operates on a too small chunk to return relevant
// results.
std::vector<int64> timestamps_ms = {0, 975, 1950, 2925};
for (int i = 0; i < timestamps_ms.size(); i++) {
EXPECT_THAT(result.classifications(0).entries(0).categories_size(),
testing::Eq(521));
const auto* top_category =
&result.classifications(0).entries(0).categories(0);
EXPECT_THAT(top_category->category_name(), testing::Eq("Speech"));
EXPECT_GT(top_category->score(), 0.9f);
EXPECT_EQ(result.classifications(0).entries(i).timestamp_ms(),
timestamps_ms[i]);
EXPECT_EQ(result[i].timestamp_ms, timestamps_ms[i]);
EXPECT_EQ(result[i].classifications.size(), 1);
auto classifications = result[i].classifications[0];
EXPECT_EQ(classifications.head_index, 0);
EXPECT_EQ(classifications.head_name, "scores");
EXPECT_EQ(classifications.categories.size(), expected_num_categories);
auto category = classifications.categories[0];
EXPECT_EQ(category.index, 0);
EXPECT_EQ(category.category_name, "Speech");
EXPECT_GT(category.score, 0.9f);
}
}
void CheckTwoHeadsClassificationResult(const ClassificationResult& result) {
EXPECT_THAT(result.classifications_size(), testing::Eq(2));
// Checks classification head #1.
EXPECT_EQ(result.classifications(0).head_name(), "yamnet_classification");
EXPECT_EQ(result.classifications(0).head_index(), 0);
EXPECT_THAT(result.classifications(0).entries(0).categories_size(),
testing::Eq(521));
const auto* top_category =
&result.classifications(0).entries(0).categories(0);
EXPECT_THAT(top_category->category_name(),
testing::Eq("Environmental noise"));
EXPECT_GT(top_category->score(), 0.5f);
EXPECT_EQ(result.classifications(0).entries(0).timestamp_ms(), 0);
if (result.classifications(0).entries_size() == 2) {
top_category = &result.classifications(0).entries(1).categories(0);
EXPECT_THAT(top_category->category_name(), testing::Eq("Silence"));
EXPECT_GT(top_category->score(), 0.99f);
EXPECT_EQ(result.classifications(0).entries(1).timestamp_ms(), 975);
void CheckTwoHeadsResult(const std::vector<AudioClassifierResult>& result) {
EXPECT_GE(result.size(), 1);
EXPECT_LE(result.size(), 2);
// Check first result.
EXPECT_EQ(result[0].timestamp_ms, 0);
EXPECT_EQ(result[0].classifications.size(), 2);
// Check first head.
EXPECT_EQ(result[0].classifications[0].head_index, 0);
EXPECT_EQ(result[0].classifications[0].head_name, "yamnet_classification");
EXPECT_EQ(result[0].classifications[0].categories.size(), 521);
EXPECT_EQ(result[0].classifications[0].categories[0].index, 508);
EXPECT_EQ(result[0].classifications[0].categories[0].category_name,
"Environmental noise");
EXPECT_GT(result[0].classifications[0].categories[0].score, 0.5f);
// Check second head.
EXPECT_EQ(result[0].classifications[1].head_index, 1);
EXPECT_EQ(result[0].classifications[1].head_name, "bird_classification");
EXPECT_EQ(result[0].classifications[1].categories.size(), 5);
EXPECT_EQ(result[0].classifications[1].categories[0].index, 4);
EXPECT_EQ(result[0].classifications[1].categories[0].category_name,
"Chestnut-crowned Antpitta");
EXPECT_GT(result[0].classifications[1].categories[0].score, 0.9f);
// Check second result, if present.
if (result.size() == 2) {
EXPECT_EQ(result[1].timestamp_ms, 975);
EXPECT_EQ(result[1].classifications.size(), 2);
// Check first head.
EXPECT_EQ(result[1].classifications[0].head_index, 0);
EXPECT_EQ(result[1].classifications[0].head_name, "yamnet_classification");
EXPECT_EQ(result[1].classifications[0].categories.size(), 521);
EXPECT_EQ(result[1].classifications[0].categories[0].index, 494);
EXPECT_EQ(result[1].classifications[0].categories[0].category_name,
"Silence");
EXPECT_GT(result[1].classifications[0].categories[0].score, 0.99f);
// Check second head.
EXPECT_EQ(result[1].classifications[1].head_index, 1);
EXPECT_EQ(result[1].classifications[1].head_name, "bird_classification");
EXPECT_EQ(result[1].classifications[1].categories.size(), 5);
EXPECT_EQ(result[1].classifications[1].categories[0].index, 1);
EXPECT_EQ(result[1].classifications[1].categories[0].category_name,
"White-breasted Wood-Wren");
EXPECT_GT(result[1].classifications[1].categories[0].score, 0.99f);
}
// Checks classification head #2.
EXPECT_EQ(result.classifications(1).head_name(), "bird_classification");
EXPECT_EQ(result.classifications(1).head_index(), 1);
EXPECT_THAT(result.classifications(1).entries(0).categories_size(),
testing::Eq(5));
top_category = &result.classifications(1).entries(0).categories(0);
EXPECT_THAT(top_category->category_name(),
testing::Eq("Chestnut-crowned Antpitta"));
EXPECT_GT(top_category->score(), 0.9f);
EXPECT_EQ(result.classifications(1).entries(0).timestamp_ms(), 0);
}
ClassificationResult GenerateSpeechClassificationResult() {
return ParseTextProtoOrDie<ClassificationResult>(
R"pb(classifications {
head_index: 0
head_name: "scores"
entries {
categories { index: 0 score: 0.94140625 category_name: "Speech" }
timestamp_ms: 0
}
entries {
categories { index: 0 score: 0.9921875 category_name: "Speech" }
timestamp_ms: 975
}
entries {
categories { index: 0 score: 0.98828125 category_name: "Speech" }
timestamp_ms: 1950
}
entries {
categories { index: 0 score: 0.99609375 category_name: "Speech" }
timestamp_ms: 2925
}
entries {
# categories are filtered out due to the low scores.
timestamp_ms: 3900
}
})pb");
}
void CheckStreamingModeClassificationResult(
std::vector<ClassificationResult> outputs) {
ASSERT_TRUE(outputs.size() == 5 || outputs.size() == 6);
auto expected_results = GenerateSpeechClassificationResult();
for (int i = 0; i < outputs.size() - 1; ++i) {
EXPECT_THAT(outputs[i].classifications(0).entries(0),
EqualsProto(expected_results.classifications(0).entries(i)));
void CheckStreamingModeResults(std::vector<AudioClassifierResult> outputs) {
EXPECT_EQ(outputs.size(), 5);
// Ignore last result, which operates on a too small chunk to return relevant
// results.
for (int i = 0; i < outputs.size() - 1; i++) {
EXPECT_FALSE(outputs[i].timestamp_ms.has_value());
EXPECT_EQ(outputs[i].classifications.size(), 1);
EXPECT_EQ(outputs[i].classifications[0].head_index, 0);
EXPECT_EQ(outputs[i].classifications[0].head_name, "scores");
EXPECT_EQ(outputs[i].classifications[0].categories.size(), 1);
EXPECT_EQ(outputs[i].classifications[0].categories[0].index, 0);
EXPECT_EQ(outputs[i].classifications[0].categories[0].category_name,
"Speech");
EXPECT_GT(outputs[i].classifications[0].categories[0].score, 0.9f);
}
int last_elem_index = outputs.size() - 1;
EXPECT_EQ(
mediapipe::Timestamp::Done().Value() / 1000,
outputs[last_elem_index].classifications(0).entries(0).timestamp_ms());
}
class CreateFromOptionsTest : public tflite_shims::testing::Test {};
@@ -264,7 +252,7 @@ TEST_F(CreateFromOptionsTest, FailsWithUnnecessaryCallback) {
options->base_options.model_asset_path =
JoinPath("./", kTestDataDirectory, kModelWithoutMetadata);
options->result_callback =
[](absl::StatusOr<ClassificationResult> status_or_result) {};
[](absl::StatusOr<AudioClassifierResult> status_or_result) {};
StatusOr<std::unique_ptr<AudioClassifier>> audio_classifier_or =
AudioClassifier::Create(std::move(options));
@@ -284,7 +272,7 @@ TEST_F(CreateFromOptionsTest, FailsWithMissingDefaultInputAudioSampleRate) {
JoinPath("./", kTestDataDirectory, kModelWithoutMetadata);
options->running_mode = core::RunningMode::AUDIO_STREAM;
options->result_callback =
[](absl::StatusOr<ClassificationResult> status_or_result) {};
[](absl::StatusOr<AudioClassifierResult> status_or_result) {};
StatusOr<std::unique_ptr<AudioClassifier>> audio_classifier_or =
AudioClassifier::Create(std::move(options));
@@ -310,7 +298,7 @@ TEST_F(ClassifyTest, Succeeds) {
auto result, audio_classifier->Classify(std::move(audio_buffer),
/*audio_sample_rate=*/16000));
MP_ASSERT_OK(audio_classifier->Close());
CheckSpeechClassificationResult(result);
CheckSpeechResult(result);
}
TEST_F(ClassifyTest, SucceedsWithResampling) {
@@ -324,7 +312,7 @@ TEST_F(ClassifyTest, SucceedsWithResampling) {
auto result, audio_classifier->Classify(std::move(audio_buffer),
/*audio_sample_rate=*/48000));
MP_ASSERT_OK(audio_classifier->Close());
CheckSpeechClassificationResult(result);
CheckSpeechResult(result);
}
TEST_F(ClassifyTest, SucceedsWithInputsAtDifferentSampleRates) {
@@ -339,13 +327,13 @@ TEST_F(ClassifyTest, SucceedsWithInputsAtDifferentSampleRates) {
auto result_16k_hz,
audio_classifier->Classify(std::move(audio_buffer_16k_hz),
/*audio_sample_rate=*/16000));
CheckSpeechClassificationResult(result_16k_hz);
CheckSpeechResult(result_16k_hz);
MP_ASSERT_OK_AND_ASSIGN(
auto result_48k_hz,
audio_classifier->Classify(std::move(audio_buffer_48k_hz),
/*audio_sample_rate=*/48000));
MP_ASSERT_OK(audio_classifier->Close());
CheckSpeechClassificationResult(result_48k_hz);
CheckSpeechResult(result_48k_hz);
}
TEST_F(ClassifyTest, SucceedsWithInsufficientData) {
@@ -361,15 +349,16 @@ TEST_F(ClassifyTest, SucceedsWithInsufficientData) {
MP_ASSERT_OK_AND_ASSIGN(
auto result, audio_classifier->Classify(std::move(zero_matrix), 16000));
MP_ASSERT_OK(audio_classifier->Close());
EXPECT_THAT(result.classifications_size(), testing::Eq(1));
EXPECT_THAT(result.classifications(0).entries_size(), testing::Eq(1));
EXPECT_THAT(result.classifications(0).entries(0).categories_size(),
testing::Eq(521));
EXPECT_THAT(
result.classifications(0).entries(0).categories(0).category_name(),
testing::Eq("Silence"));
EXPECT_THAT(result.classifications(0).entries(0).categories(0).score(),
testing::FloatEq(.800781f));
EXPECT_EQ(result.size(), 1);
EXPECT_EQ(result[0].timestamp_ms, 0);
EXPECT_EQ(result[0].classifications.size(), 1);
EXPECT_EQ(result[0].classifications[0].head_index, 0);
EXPECT_EQ(result[0].classifications[0].head_name, "scores");
EXPECT_EQ(result[0].classifications[0].categories.size(), 521);
EXPECT_EQ(result[0].classifications[0].categories[0].index, 494);
EXPECT_EQ(result[0].classifications[0].categories[0].category_name,
"Silence");
EXPECT_FLOAT_EQ(result[0].classifications[0].categories[0].score, 0.800781f);
}
TEST_F(ClassifyTest, SucceedsWithMultiheadsModel) {
@@ -383,7 +372,7 @@ TEST_F(ClassifyTest, SucceedsWithMultiheadsModel) {
auto result, audio_classifier->Classify(std::move(audio_buffer),
/*audio_sample_rate=*/16000));
MP_ASSERT_OK(audio_classifier->Close());
CheckTwoHeadsClassificationResult(result);
CheckTwoHeadsResult(result);
}
TEST_F(ClassifyTest, SucceedsWithMultiheadsModelAndResampling) {
@@ -397,7 +386,7 @@ TEST_F(ClassifyTest, SucceedsWithMultiheadsModelAndResampling) {
auto result, audio_classifier->Classify(std::move(audio_buffer),
/*audio_sample_rate=*/44100));
MP_ASSERT_OK(audio_classifier->Close());
CheckTwoHeadsClassificationResult(result);
CheckTwoHeadsResult(result);
}
TEST_F(ClassifyTest,
@@ -413,13 +402,13 @@ TEST_F(ClassifyTest,
auto result_44k_hz,
audio_classifier->Classify(std::move(audio_buffer_44k_hz),
/*audio_sample_rate=*/44100));
CheckTwoHeadsClassificationResult(result_44k_hz);
CheckTwoHeadsResult(result_44k_hz);
MP_ASSERT_OK_AND_ASSIGN(
auto result_16k_hz,
audio_classifier->Classify(std::move(audio_buffer_16k_hz),
/*audio_sample_rate=*/16000));
MP_ASSERT_OK(audio_classifier->Close());
CheckTwoHeadsClassificationResult(result_16k_hz);
CheckTwoHeadsResult(result_16k_hz);
}
TEST_F(ClassifyTest, SucceedsWithMaxResultOption) {
@@ -428,14 +417,13 @@ TEST_F(ClassifyTest, SucceedsWithMaxResultOption) {
options->base_options.model_asset_path =
JoinPath("./", kTestDataDirectory, kModelWithMetadata);
options->classifier_options.max_results = 1;
options->classifier_options.score_threshold = 0.35f;
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<AudioClassifier> audio_classifier,
AudioClassifier::Create(std::move(options)));
MP_ASSERT_OK_AND_ASSIGN(
auto result, audio_classifier->Classify(std::move(audio_buffer),
/*audio_sample_rate=*/48000));
MP_ASSERT_OK(audio_classifier->Close());
EXPECT_THAT(result, EqualsProto(GenerateSpeechClassificationResult()));
CheckSpeechResult(result, /*expected_num_categories=*/1);
}
TEST_F(ClassifyTest, SucceedsWithScoreThresholdOption) {
@@ -450,7 +438,7 @@ TEST_F(ClassifyTest, SucceedsWithScoreThresholdOption) {
auto result, audio_classifier->Classify(std::move(audio_buffer),
/*audio_sample_rate=*/48000));
MP_ASSERT_OK(audio_classifier->Close());
EXPECT_THAT(result, EqualsProto(GenerateSpeechClassificationResult()));
CheckSpeechResult(result, /*expected_num_categories=*/1);
}
TEST_F(ClassifyTest, SucceedsWithCategoryAllowlist) {
@@ -466,7 +454,7 @@ TEST_F(ClassifyTest, SucceedsWithCategoryAllowlist) {
auto result, audio_classifier->Classify(std::move(audio_buffer),
/*audio_sample_rate=*/48000));
MP_ASSERT_OK(audio_classifier->Close());
EXPECT_THAT(result, EqualsProto(GenerateSpeechClassificationResult()));
CheckSpeechResult(result, /*expected_num_categories=*/1);
}
TEST_F(ClassifyTest, SucceedsWithCategoryDenylist) {
@@ -482,16 +470,16 @@ TEST_F(ClassifyTest, SucceedsWithCategoryDenylist) {
auto result, audio_classifier->Classify(std::move(audio_buffer),
/*audio_sample_rate=*/48000));
MP_ASSERT_OK(audio_classifier->Close());
// All categroies with the "Speech" label are filtered out.
EXPECT_THAT(result, EqualsProto(R"pb(classifications {
head_index: 0
head_name: "scores"
entries { timestamp_ms: 0 }
entries { timestamp_ms: 975 }
entries { timestamp_ms: 1950 }
entries { timestamp_ms: 2925 }
entries { timestamp_ms: 3900 }
})pb"));
// All categories with the "Speech" label are filtered out.
std::vector<int64> timestamps_ms = {0, 975, 1950, 2925};
for (int i = 0; i < timestamps_ms.size(); i++) {
EXPECT_EQ(result[i].timestamp_ms, timestamps_ms[i]);
EXPECT_EQ(result[i].classifications.size(), 1);
auto classifications = result[i].classifications[0];
EXPECT_EQ(classifications.head_index, 0);
EXPECT_EQ(classifications.head_name, "scores");
EXPECT_TRUE(classifications.categories.empty());
}
}
class ClassifyAsyncTest : public tflite_shims::testing::Test {};
@@ -506,9 +494,9 @@ TEST_F(ClassifyAsyncTest, Succeeds) {
options->classifier_options.score_threshold = 0.3f;
options->running_mode = core::RunningMode::AUDIO_STREAM;
options->sample_rate = kSampleRateHz;
std::vector<ClassificationResult> outputs;
std::vector<AudioClassifierResult> outputs;
options->result_callback =
[&outputs](absl::StatusOr<ClassificationResult> status_or_result) {
[&outputs](absl::StatusOr<AudioClassifierResult> status_or_result) {
MP_ASSERT_OK_AND_ASSIGN(outputs.emplace_back(), status_or_result);
};
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<AudioClassifier> audio_classifier,
@@ -523,7 +511,7 @@ TEST_F(ClassifyAsyncTest, Succeeds) {
start_col += kYamnetNumOfAudioSamples * 3;
}
MP_ASSERT_OK(audio_classifier->Close());
CheckStreamingModeClassificationResult(outputs);
CheckStreamingModeResults(outputs);
}
TEST_F(ClassifyAsyncTest, SucceedsWithNonDeterministicNumAudioSamples) {
@@ -536,9 +524,9 @@ TEST_F(ClassifyAsyncTest, SucceedsWithNonDeterministicNumAudioSamples) {
options->classifier_options.score_threshold = 0.3f;
options->running_mode = core::RunningMode::AUDIO_STREAM;
options->sample_rate = kSampleRateHz;
std::vector<ClassificationResult> outputs;
std::vector<AudioClassifierResult> outputs;
options->result_callback =
[&outputs](absl::StatusOr<ClassificationResult> status_or_result) {
[&outputs](absl::StatusOr<AudioClassifierResult> status_or_result) {
MP_ASSERT_OK_AND_ASSIGN(outputs.emplace_back(), status_or_result);
};
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<AudioClassifier> audio_classifier,
@@ -555,7 +543,7 @@ TEST_F(ClassifyAsyncTest, SucceedsWithNonDeterministicNumAudioSamples) {
start_col += num_samples;
}
MP_ASSERT_OK(audio_classifier->Close());
CheckStreamingModeClassificationResult(outputs);
CheckStreamingModeResults(outputs);
}
} // namespace
mediapipe/tasks/cc/components/BUILD
+1 -35
View File
@@ -40,6 +40,7 @@ cc_library(
"//mediapipe/calculators/image:image_properties_calculator",
"//mediapipe/calculators/tensor:image_to_tensor_calculator",
"//mediapipe/calculators/tensor:image_to_tensor_calculator_cc_proto",
"//mediapipe/calculators/tensor:inference_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:builder",
"//mediapipe/framework/api2:port",
@@ -60,41 +61,6 @@ cc_library(
# TODO: Enable this test
cc_library(
name = "embedder_options",
srcs = ["embedder_options.cc"],
hdrs = ["embedder_options.h"],
deps = ["//mediapipe/tasks/cc/components/proto:embedder_options_cc_proto"],
)
cc_library(
name = "embedding_postprocessing_graph",
srcs = ["embedding_postprocessing_graph.cc"],
hdrs = ["embedding_postprocessing_graph.h"],
deps = [
"//mediapipe/calculators/tensor:tensors_dequantization_calculator",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:builder",
"//mediapipe/framework/api2:port",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/tool:options_map",
"//mediapipe/tasks/cc:common",
"//mediapipe/tasks/cc/components/calculators:tensors_to_embeddings_calculator",
"//mediapipe/tasks/cc/components/calculators:tensors_to_embeddings_calculator_cc_proto",
"//mediapipe/tasks/cc/components/containers/proto:embeddings_cc_proto",
"//mediapipe/tasks/cc/components/proto:embedder_options_cc_proto",
"//mediapipe/tasks/cc/components/proto:embedding_postprocessing_graph_options_cc_proto",
"//mediapipe/tasks/cc/components/utils:source_or_node_output",
"//mediapipe/tasks/cc/core:model_resources",
"//mediapipe/tasks/cc/metadata:metadata_extractor",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
"@org_tensorflow//tensorflow/lite/schema:schema_fbs",
],
alwayslink = 1,
)
# TODO: Investigate rewriting the build rule to only link
# the Bert Preprocessor if it's needed.
cc_library(
mediapipe/tasks/cc/components/calculators/BUILD
+2 -2
View File
@@ -163,7 +163,7 @@ mediapipe_proto_library(
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
"//mediapipe/tasks/cc/components/proto:embedder_options_proto",
"//mediapipe/tasks/cc/components/processors/proto:embedder_options_proto",
],
)
@@ -178,7 +178,7 @@ cc_library(
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:ret_check",
"//mediapipe/tasks/cc/components/containers/proto:embeddings_cc_proto",
"//mediapipe/tasks/cc/components/proto:embedder_options_cc_proto",
"//mediapipe/tasks/cc/components/processors/proto:embedder_options_cc_proto",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:str_format",
],
mediapipe/tasks/cc/components/calculators/tensors_to_embeddings_calculator.cc
+16 -16
View File
@@ -26,14 +26,14 @@
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/tasks/cc/components/calculators/tensors_to_embeddings_calculator.pb.h"
#include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h"
#include "mediapipe/tasks/cc/components/proto/embedder_options.pb.h"
#include "mediapipe/tasks/cc/components/processors/proto/embedder_options.pb.h"
namespace mediapipe {
namespace api2 {
namespace {
using ::mediapipe::tasks::components::containers::proto::EmbeddingEntry;
using ::mediapipe::tasks::components::containers::proto::Embedding;
using ::mediapipe::tasks::components::containers::proto::EmbeddingResult;
// Computes the inverse L2 norm of the provided array of values. Returns 1.0 in
@@ -66,7 +66,7 @@ float GetInverseL2Norm(const float* values, int size) {
class TensorsToEmbeddingsCalculator : public Node {
public:
static constexpr Input<std::vector<Tensor>> kTensorsIn{"TENSORS"};
static constexpr Output<EmbeddingResult> kEmbeddingsOut{"EMBEDDING_RESULT"};
static constexpr Output<EmbeddingResult> kEmbeddingsOut{"EMBEDDINGS"};
MEDIAPIPE_NODE_CONTRACT(kTensorsIn, kEmbeddingsOut);
absl::Status Open(CalculatorContext* cc) override;
@@ -77,8 +77,8 @@ class TensorsToEmbeddingsCalculator : public Node {
bool quantize_;
std::vector<std::string> head_names_;
void FillFloatEmbeddingEntry(const Tensor& tensor, EmbeddingEntry* entry);
void FillQuantizedEmbeddingEntry(const Tensor& tensor, EmbeddingEntry* entry);
void FillFloatEmbedding(const Tensor& tensor, Embedding* embedding);
void FillQuantizedEmbedding(const Tensor& tensor, Embedding* embedding);
};
absl::Status TensorsToEmbeddingsCalculator::Open(CalculatorContext* cc) {
@@ -104,42 +104,42 @@ absl::Status TensorsToEmbeddingsCalculator::Process(CalculatorContext* cc) {
for (int i = 0; i < tensors.size(); ++i) {
const auto& tensor = tensors[i];
RET_CHECK(tensor.element_type() == Tensor::ElementType::kFloat32);
auto* embeddings = result.add_embeddings();
embeddings->set_head_index(i);
auto* embedding = result.add_embeddings();
embedding->set_head_index(i);
if (!head_names_.empty()) {
embeddings->set_head_name(head_names_[i]);
embedding->set_head_name(head_names_[i]);
}
if (quantize_) {
FillQuantizedEmbeddingEntry(tensor, embeddings->add_entries());
FillQuantizedEmbedding(tensor, embedding);
} else {
FillFloatEmbeddingEntry(tensor, embeddings->add_entries());
FillFloatEmbedding(tensor, embedding);
}
}
kEmbeddingsOut(cc).Send(result);
return absl::OkStatus();
}
void TensorsToEmbeddingsCalculator::FillFloatEmbeddingEntry(
const Tensor& tensor, EmbeddingEntry* entry) {
void TensorsToEmbeddingsCalculator::FillFloatEmbedding(const Tensor& tensor,
Embedding* embedding) {
int size = tensor.shape().num_elements();
auto tensor_view = tensor.GetCpuReadView();
const float* tensor_buffer = tensor_view.buffer<float>();
float inv_l2_norm =
l2_normalize_ ? GetInverseL2Norm(tensor_buffer, size) : 1.0f;
auto* float_embedding = entry->mutable_float_embedding();
auto* float_embedding = embedding->mutable_float_embedding();
for (int i = 0; i < size; ++i) {
float_embedding->add_values(tensor_buffer[i] * inv_l2_norm);
}
}
void TensorsToEmbeddingsCalculator::FillQuantizedEmbeddingEntry(
const Tensor& tensor, EmbeddingEntry* entry) {
void TensorsToEmbeddingsCalculator::FillQuantizedEmbedding(
const Tensor& tensor, Embedding* embedding) {
int size = tensor.shape().num_elements();
auto tensor_view = tensor.GetCpuReadView();
const float* tensor_buffer = tensor_view.buffer<float>();
float inv_l2_norm =
l2_normalize_ ? GetInverseL2Norm(tensor_buffer, size) : 1.0f;
auto* values = entry->mutable_quantized_embedding()->mutable_values();
auto* values = embedding->mutable_quantized_embedding()->mutable_values();
values->resize(size);
for (int i = 0; i < size; ++i) {
// Normalize.
mediapipe/tasks/cc/components/calculators/tensors_to_embeddings_calculator.proto
+3 -3
View File
@@ -18,7 +18,7 @@ syntax = "proto2";
package mediapipe;
import "mediapipe/framework/calculator.proto";
import "mediapipe/tasks/cc/components/proto/embedder_options.proto";
import "mediapipe/tasks/cc/components/processors/proto/embedder_options.proto";
message TensorsToEmbeddingsCalculatorOptions {
extend mediapipe.CalculatorOptions {
@@ -27,8 +27,8 @@ message TensorsToEmbeddingsCalculatorOptions {
// The embedder options defining whether to L2-normalize or scalar-quantize
// the outputs.
optional mediapipe.tasks.components.proto.EmbedderOptions embedder_options =
1;
optional mediapipe.tasks.components.processors.proto.EmbedderOptions
embedder_options = 1;
// The embedder head names.
repeated string head_names = 2;
mediapipe/tasks/cc/components/calculators/tensors_to_embeddings_calculator_test.cc
+50 -77
View File
@@ -55,7 +55,7 @@ TEST(TensorsToEmbeddingsCalculatorTest, FailsWithInvalidHeadNamesNumber) {
CalculatorRunner runner(ParseTextProtoOrDie<Node>(R"pb(
calculator: "TensorsToEmbeddingsCalculator"
input_stream: "TENSORS:tensors"
output_stream: "EMBEDDING_RESULT:embeddings"
output_stream: "EMBEDDINGS:embeddings"
options {
[mediapipe.TensorsToEmbeddingsCalculatorOptions.ext] { head_names: "foo" }
}
@@ -73,7 +73,7 @@ TEST(TensorsToEmbeddingsCalculatorTest, SucceedsWithoutHeadNames) {
CalculatorRunner runner(ParseTextProtoOrDie<Node>(R"pb(
calculator: "TensorsToEmbeddingsCalculator"
input_stream: "TENSORS:tensors"
output_stream: "EMBEDDING_RESULT:embeddings"
output_stream: "EMBEDDINGS:embeddings"
options {
[mediapipe.TensorsToEmbeddingsCalculatorOptions.ext] {
embedder_options { l2_normalize: false quantize: false }
@@ -84,28 +84,24 @@ TEST(TensorsToEmbeddingsCalculatorTest, SucceedsWithoutHeadNames) {
BuildGraph(&runner, {{0.1, 0.2}, {-0.2, -0.3}});
MP_ASSERT_OK(runner.Run());
const EmbeddingResult& result = runner.Outputs()
.Get("EMBEDDING_RESULT", 0)
.packets[0]
.Get<EmbeddingResult>();
EXPECT_THAT(
result,
EqualsProto(ParseTextProtoOrDie<EmbeddingResult>(
R"pb(embeddings {
entries { float_embedding { values: 0.1 values: 0.2 } }
head_index: 0
}
embeddings {
entries { float_embedding { values: -0.2 values: -0.3 } }
head_index: 1
})pb")));
const EmbeddingResult& result =
runner.Outputs().Get("EMBEDDINGS", 0).packets[0].Get<EmbeddingResult>();
EXPECT_THAT(result, EqualsProto(ParseTextProtoOrDie<EmbeddingResult>(
R"pb(embeddings {
float_embedding { values: 0.1 values: 0.2 }
head_index: 0
}
embeddings {
float_embedding { values: -0.2 values: -0.3 }
head_index: 1
})pb")));
}
TEST(TensorsToEmbeddingsCalculatorTest, SucceedsWithHeadNames) {
CalculatorRunner runner(ParseTextProtoOrDie<Node>(R"pb(
calculator: "TensorsToEmbeddingsCalculator"
input_stream: "TENSORS:tensors"
output_stream: "EMBEDDING_RESULT:embeddings"
output_stream: "EMBEDDINGS:embeddings"
options {
[mediapipe.TensorsToEmbeddingsCalculatorOptions.ext] {
embedder_options { l2_normalize: false quantize: false }
@@ -118,30 +114,26 @@ TEST(TensorsToEmbeddingsCalculatorTest, SucceedsWithHeadNames) {
BuildGraph(&runner, {{0.1, 0.2}, {-0.2, -0.3}});
MP_ASSERT_OK(runner.Run());
const EmbeddingResult& result = runner.Outputs()
.Get("EMBEDDING_RESULT", 0)
.packets[0]
.Get<EmbeddingResult>();
EXPECT_THAT(
result,
EqualsProto(ParseTextProtoOrDie<EmbeddingResult>(
R"pb(embeddings {
entries { float_embedding { values: 0.1 values: 0.2 } }
head_index: 0
head_name: "foo"
}
embeddings {
entries { float_embedding { values: -0.2 values: -0.3 } }
head_index: 1
head_name: "bar"
})pb")));
const EmbeddingResult& result =
runner.Outputs().Get("EMBEDDINGS", 0).packets[0].Get<EmbeddingResult>();
EXPECT_THAT(result, EqualsProto(ParseTextProtoOrDie<EmbeddingResult>(
R"pb(embeddings {
float_embedding { values: 0.1 values: 0.2 }
head_index: 0
head_name: "foo"
}
embeddings {
float_embedding { values: -0.2 values: -0.3 }
head_index: 1
head_name: "bar"
})pb")));
}
TEST(TensorsToEmbeddingsCalculatorTest, SucceedsWithNormalization) {
CalculatorRunner runner(ParseTextProtoOrDie<Node>(R"pb(
calculator: "TensorsToEmbeddingsCalculator"
input_stream: "TENSORS:tensors"
output_stream: "EMBEDDING_RESULT:embeddings"
output_stream: "EMBEDDINGS:embeddings"
options {
[mediapipe.TensorsToEmbeddingsCalculatorOptions.ext] {
embedder_options { l2_normalize: true quantize: false }
@@ -152,23 +144,17 @@ TEST(TensorsToEmbeddingsCalculatorTest, SucceedsWithNormalization) {
BuildGraph(&runner, {{0.1, 0.2}, {-0.2, -0.3}});
MP_ASSERT_OK(runner.Run());
const EmbeddingResult& result = runner.Outputs()
.Get("EMBEDDING_RESULT", 0)
.packets[0]
.Get<EmbeddingResult>();
const EmbeddingResult& result =
runner.Outputs().Get("EMBEDDINGS", 0).packets[0].Get<EmbeddingResult>();
EXPECT_THAT(
result,
EqualsProto(ParseTextProtoOrDie<EmbeddingResult>(
R"pb(embeddings {
entries {
float_embedding { values: 0.44721356 values: 0.8944271 }
}
float_embedding { values: 0.44721356 values: 0.8944271 }
head_index: 0
}
embeddings {
entries {
float_embedding { values: -0.5547002 values: -0.8320503 }
}
float_embedding { values: -0.5547002 values: -0.8320503 }
head_index: 1
})pb")));
}
@@ -177,7 +163,7 @@ TEST(TensorsToEmbeddingsCalculatorTest, SucceedsWithQuantization) {
CalculatorRunner runner(ParseTextProtoOrDie<Node>(R"pb(
calculator: "TensorsToEmbeddingsCalculator"
input_stream: "TENSORS:tensors"
output_stream: "EMBEDDING_RESULT:embeddings"
output_stream: "EMBEDDINGS:embeddings"
options {
[mediapipe.TensorsToEmbeddingsCalculatorOptions.ext] {
embedder_options { l2_normalize: false quantize: true }
@@ -188,22 +174,16 @@ TEST(TensorsToEmbeddingsCalculatorTest, SucceedsWithQuantization) {
BuildGraph(&runner, {{0.1, 0.2}, {-0.2, -0.3}});
MP_ASSERT_OK(runner.Run());
const EmbeddingResult& result = runner.Outputs()
.Get("EMBEDDING_RESULT", 0)
.packets[0]
.Get<EmbeddingResult>();
const EmbeddingResult& result =
runner.Outputs().Get("EMBEDDINGS", 0).packets[0].Get<EmbeddingResult>();
EXPECT_THAT(result,
EqualsProto(ParseTextProtoOrDie<EmbeddingResult>(
R"pb(embeddings {
entries {
quantized_embedding { values: "\x0d\x1a" } # 13,26
}
quantized_embedding { values: "\x0d\x1a" } # 13,26
head_index: 0
}
embeddings {
entries {
quantized_embedding { values: "\xe6\xda" } # -26,-38
}
quantized_embedding { values: "\xe6\xda" } # -26,-38
head_index: 1
})pb")));
}
@@ -213,7 +193,7 @@ TEST(TensorsToEmbeddingsCalculatorTest,
CalculatorRunner runner(ParseTextProtoOrDie<Node>(R"pb(
calculator: "TensorsToEmbeddingsCalculator"
input_stream: "TENSORS:tensors"
output_stream: "EMBEDDING_RESULT:embeddings"
output_stream: "EMBEDDINGS:embeddings"
options {
[mediapipe.TensorsToEmbeddingsCalculatorOptions.ext] {
embedder_options { l2_normalize: true quantize: true }
@@ -224,25 +204,18 @@ TEST(TensorsToEmbeddingsCalculatorTest,
BuildGraph(&runner, {{0.1, 0.2}, {-0.2, -0.3}});
MP_ASSERT_OK(runner.Run());
const EmbeddingResult& result = runner.Outputs()
.Get("EMBEDDING_RESULT", 0)
.packets[0]
.Get<EmbeddingResult>();
EXPECT_THAT(
result,
EqualsProto(ParseTextProtoOrDie<EmbeddingResult>(
R"pb(embeddings {
entries {
quantized_embedding { values: "\x39\x72" } # 57,114
}
head_index: 0
}
embeddings {
entries {
quantized_embedding { values: "\xb9\x95" } # -71,-107
}
head_index: 1
})pb")));
const EmbeddingResult& result =
runner.Outputs().Get("EMBEDDINGS", 0).packets[0].Get<EmbeddingResult>();
EXPECT_THAT(result,
EqualsProto(ParseTextProtoOrDie<EmbeddingResult>(
R"pb(embeddings {
quantized_embedding { values: "\x39\x72" } # 57,114
head_index: 0
}
embeddings {
quantized_embedding { values: "\xb9\x95" } # -71,-107
head_index: 1
})pb")));
}
} // namespace
mediapipe/tasks/cc/components/containers/BUILD
+9
View File
@@ -49,3 +49,12 @@ cc_library(
"//mediapipe/tasks/cc/components/containers/proto:classifications_cc_proto",
],
)
cc_library(
name = "embedding_result",
srcs = ["embedding_result.cc"],
hdrs = ["embedding_result.h"],
deps = [
"//mediapipe/tasks/cc/components/containers/proto:embeddings_cc_proto",
],
)
mediapipe/tasks/cc/components/containers/embedding_result.cc
+57
View File
@@ -0,0 +1,57 @@
/* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
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/tasks/cc/components/containers/embedding_result.h"
#include <iterator>
#include <optional>
#include <string>
#include <vector>
#include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h"
namespace mediapipe::tasks::components::containers {
Embedding ConvertToEmbedding(const proto::Embedding& proto) {
Embedding embedding;
if (proto.has_float_embedding()) {
embedding.float_embedding = {
std::make_move_iterator(proto.float_embedding().values().begin()),
std::make_move_iterator(proto.float_embedding().values().end())};
} else {
embedding.quantized_embedding = {
std::make_move_iterator(proto.quantized_embedding().values().begin()),
std::make_move_iterator(proto.quantized_embedding().values().end())};
}
embedding.head_index = proto.head_index();
if (proto.has_head_name()) {
embedding.head_name = proto.head_name();
}
return embedding;
}
EmbeddingResult ConvertToEmbeddingResult(const proto::EmbeddingResult& proto) {
EmbeddingResult embedding_result;
embedding_result.embeddings.reserve(proto.embeddings_size());
for (const auto& embedding : proto.embeddings()) {
embedding_result.embeddings.push_back(ConvertToEmbedding(embedding));
}
if (proto.has_timestamp_ms()) {
embedding_result.timestamp_ms = proto.timestamp_ms();
}
return embedding_result;
}
} // namespace mediapipe::tasks::components::containers
mediapipe/tasks/cc/components/containers/embedding_result.h
+72
View File
@@ -0,0 +1,72 @@
/* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
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_TASKS_CC_COMPONENTS_CONTAINERS_EMBEDDING_RESULT_H_
#define MEDIAPIPE_TASKS_CC_COMPONENTS_CONTAINERS_EMBEDDING_RESULT_H_
#include <optional>
#include <string>
#include <vector>
#include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h"
namespace mediapipe::tasks::components::containers {
// Embedding result for a given embedder head.
//
// One and only one of the two 'float_embedding' and 'quantized_embedding' will
// contain data, based on whether or not the embedder was configured to perform
// scalar quantization.
struct Embedding {
// Floating-point embedding. Empty if the embedder was configured to perform
// scalar-quantization.
std::vector<float> float_embedding;
// Scalar-quantized embedding. Empty if the embedder was not configured to
// perform scalar quantization.
std::string quantized_embedding;
// The index of the embedder head (i.e. output tensor) this embedding comes
// from. This is useful for multi-head models.
int head_index;
// The optional name of the embedder head, as provided in the TFLite Model
// Metadata [1] if present. This is useful for multi-head models.
//
// [1]: https://www.tensorflow.org/lite/convert/metadata
std::optional<std::string> head_name = std::nullopt;
};
// Defines embedding results of a model.
struct EmbeddingResult {
// The embedding results for each head of the model.
std::vector<Embedding> embeddings;
// The optional timestamp (in milliseconds) of the start of the chunk of data
// corresponding to these results.
//
// This is only used for embedding extraction on time series (e.g. audio
// embedding). In these use cases, the amount of data to process might
// exceed the maximum size that the model can process: to solve this, the
// input data is split into multiple chunks starting at different timestamps.
std::optional<int64_t> timestamp_ms = std::nullopt;
};
// Utility function to convert from Embedding proto to Embedding struct.
Embedding ConvertToEmbedding(const proto::Embedding& proto);
// Utility function to convert from EmbeddingResult proto to EmbeddingResult
// struct.
EmbeddingResult ConvertToEmbeddingResult(const proto::EmbeddingResult& proto);
} // namespace mediapipe::tasks::components::containers
#endif // MEDIAPIPE_TASKS_CC_COMPONENTS_CONTAINERS_EMBEDDING_RESULT_H_
mediapipe/tasks/cc/components/containers/proto/embeddings.proto
+16 -15
View File
@@ -30,30 +30,31 @@ message QuantizedEmbedding {
optional bytes values = 1;
}
// Floating-point or scalar-quantized embedding with an optional timestamp.
message EmbeddingEntry {
// The actual embedding, either floating-point or scalar-quantized.
// Embedding result for a given embedder head.
message Embedding {
// The actual embedding, either floating-point or quantized.
oneof embedding {
FloatEmbedding float_embedding = 1;
QuantizedEmbedding quantized_embedding = 2;
}
// The optional timestamp (in milliseconds) associated to the embedding entry.
// This is useful for time series use cases, e.g. audio embedding.
optional int64 timestamp_ms = 3;
}
// Embeddings for a given embedder head.
message Embeddings {
repeated EmbeddingEntry entries = 1;
// The index of the embedder head that produced this embedding. This is useful
// for multi-head models.
optional int32 head_index = 2;
optional int32 head_index = 3;
// The name of the embedder head, which is the corresponding tensor metadata
// name (if any). This is useful for multi-head models.
optional string head_name = 3;
optional string head_name = 4;
}
// Contains one set of results per embedder head.
// Embedding results for a given embedder model.
message EmbeddingResult {
repeated Embeddings embeddings = 1;
// The embedding results for each model head, i.e. one for each output tensor.
repeated Embedding embeddings = 1;
// The optional timestamp (in milliseconds) of the start of the chunk of data
// corresponding to these results.
//
// This is only used for embedding extraction on time series (e.g. audio
// embedding). In these use cases, the amount of data to process might
// exceed the maximum size that the model can process: to solve this, the
// input data is split into multiple chunks starting at different timestamps.
optional int64 timestamp_ms = 2;
}
mediapipe/tasks/cc/components/image_preprocessing.cc
+1
View File
@@ -24,6 +24,7 @@ limitations under the License.
#include "absl/status/statusor.h"
#include "mediapipe/calculators/image/image_clone_calculator.pb.h"
#include "mediapipe/calculators/tensor/image_to_tensor_calculator.pb.h"
#include "mediapipe/calculators/tensor/inference_calculator.pb.h"
#include "mediapipe/framework/api2/builder.h"
#include "mediapipe/framework/api2/port.h"
#include "mediapipe/framework/calculator_framework.h"
mediapipe/tasks/cc/components/processors/BUILD
+35
View File
@@ -62,3 +62,38 @@ cc_library(
],
alwayslink = 1,
)
cc_library(
name = "embedder_options",
srcs = ["embedder_options.cc"],
hdrs = ["embedder_options.h"],
deps = ["//mediapipe/tasks/cc/components/processors/proto:embedder_options_cc_proto"],
)
cc_library(
name = "embedding_postprocessing_graph",
srcs = ["embedding_postprocessing_graph.cc"],
hdrs = ["embedding_postprocessing_graph.h"],
deps = [
"//mediapipe/calculators/tensor:tensors_dequantization_calculator",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:builder",
"//mediapipe/framework/api2:port",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/tool:options_map",
"//mediapipe/tasks/cc:common",
"//mediapipe/tasks/cc/components/calculators:tensors_to_embeddings_calculator",
"//mediapipe/tasks/cc/components/calculators:tensors_to_embeddings_calculator_cc_proto",
"//mediapipe/tasks/cc/components/containers/proto:embeddings_cc_proto",
"//mediapipe/tasks/cc/components/processors/proto:embedder_options_cc_proto",
"//mediapipe/tasks/cc/components/processors/proto:embedding_postprocessing_graph_options_cc_proto",
"//mediapipe/tasks/cc/components/utils:source_or_node_output",
"//mediapipe/tasks/cc/core:model_resources",
"//mediapipe/tasks/cc/metadata:metadata_extractor",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
"@org_tensorflow//tensorflow/lite/schema:schema_fbs",
],
alwayslink = 1,
)
mediapipe/tasks/cc/components/embedder_options.cc → mediapipe/tasks/cc/components/processors/embedder_options.cc
+6 -4
View File
@@ -13,22 +13,24 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "mediapipe/tasks/cc/components/embedder_options.h"
#include "mediapipe/tasks/cc/components/processors/embedder_options.h"
#include "mediapipe/tasks/cc/components/proto/embedder_options.pb.h"
#include "mediapipe/tasks/cc/components/processors/proto/embedder_options.pb.h"
namespace mediapipe {
namespace tasks {
namespace components {
namespace processors {
tasks::components::proto::EmbedderOptions ConvertEmbedderOptionsToProto(
proto::EmbedderOptions ConvertEmbedderOptionsToProto(
EmbedderOptions* embedder_options) {
tasks::components::proto::EmbedderOptions options_proto;
proto::EmbedderOptions options_proto;
options_proto.set_l2_normalize(embedder_options->l2_normalize);
options_proto.set_quantize(embedder_options->quantize);
return options_proto;
}
} // namespace processors
} // namespace components
} // namespace tasks
} // namespace mediapipe
mediapipe/tasks/cc/components/embedder_options.h → mediapipe/tasks/cc/components/processors/embedder_options.h
+7 -5
View File
@@ -13,14 +13,15 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef MEDIAPIPE_TASKS_CC_COMPONENTS_EMBEDDER_OPTIONS_H_
#define MEDIAPIPE_TASKS_CC_COMPONENTS_EMBEDDER_OPTIONS_H_
#ifndef MEDIAPIPE_TASKS_CC_COMPONENTS_PROCESSORS_EMBEDDER_OPTIONS_H_
#define MEDIAPIPE_TASKS_CC_COMPONENTS_PROCESSORS_EMBEDDER_OPTIONS_H_
#include "mediapipe/tasks/cc/components/proto/embedder_options.pb.h"
#include "mediapipe/tasks/cc/components/processors/proto/embedder_options.pb.h"
namespace mediapipe {
namespace tasks {
namespace components {
namespace processors {
// Embedder options for MediaPipe C++ embedding extraction tasks.
struct EmbedderOptions {
@@ -37,11 +38,12 @@ struct EmbedderOptions {
bool quantize;
};
tasks::components::proto::EmbedderOptions ConvertEmbedderOptionsToProto(
proto::EmbedderOptions ConvertEmbedderOptionsToProto(
EmbedderOptions* embedder_options);
} // namespace processors
} // namespace components
} // namespace tasks
} // namespace mediapipe
#endif // MEDIAPIPE_TASKS_CC_COMPONENTS_EMBEDDER_OPTIONS_H_
#endif // MEDIAPIPE_TASKS_CC_COMPONENTS_PROCESSORS_EMBEDDER_OPTIONS_H_
mediapipe/tasks/cc/components/embedding_postprocessing_graph.cc → mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph.cc
+10 -10
View File
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "mediapipe/tasks/cc/components/embedding_postprocessing_graph.h"
#include "mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph.h"
#include <string>
#include <vector>
@@ -29,8 +29,8 @@ limitations under the License.
#include "mediapipe/tasks/cc/common.h"
#include "mediapipe/tasks/cc/components/calculators/tensors_to_embeddings_calculator.pb.h"
#include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h"
#include "mediapipe/tasks/cc/components/proto/embedder_options.pb.h"
#include "mediapipe/tasks/cc/components/proto/embedding_postprocessing_graph_options.pb.h"
#include "mediapipe/tasks/cc/components/processors/proto/embedder_options.pb.h"
#include "mediapipe/tasks/cc/components/processors/proto/embedding_postprocessing_graph_options.pb.h"
#include "mediapipe/tasks/cc/components/utils/source_or_node_output.h"
#include "mediapipe/tasks/cc/core/model_resources.h"
#include "mediapipe/tasks/cc/metadata/metadata_extractor.h"
@@ -39,6 +39,7 @@ limitations under the License.
namespace mediapipe {
namespace tasks {
namespace components {
namespace processors {
namespace {
@@ -49,13 +50,12 @@ using ::mediapipe::api2::builder::GenericNode;
using ::mediapipe::api2::builder::Graph;
using ::mediapipe::api2::builder::Source;
using ::mediapipe::tasks::components::containers::proto::EmbeddingResult;
using ::mediapipe::tasks::components::proto::EmbedderOptions;
using ::mediapipe::tasks::core::ModelResources;
using TensorsSource =
::mediapipe::tasks::SourceOrNodeOutput<std::vector<Tensor>>;
constexpr char kTensorsTag[] = "TENSORS";
constexpr char kEmbeddingResultTag[] = "EMBEDDING_RESULT";
constexpr char kEmbeddingsTag[] = "EMBEDDINGS";
// Identifies whether or not the model has quantized outputs, and performs
// sanity checks.
@@ -144,7 +144,7 @@ absl::StatusOr<std::vector<std::string>> GetHeadNames(
absl::Status ConfigureEmbeddingPostprocessing(
const ModelResources& model_resources,
const EmbedderOptions& embedder_options,
const proto::EmbedderOptions& embedder_options,
proto::EmbeddingPostprocessingGraphOptions* options) {
ASSIGN_OR_RETURN(bool has_quantized_outputs,
HasQuantizedOutputs(model_resources));
@@ -188,7 +188,7 @@ class EmbeddingPostprocessingGraph : public mediapipe::Subgraph {
BuildEmbeddingPostprocessing(
sc->Options<proto::EmbeddingPostprocessingGraphOptions>(),
graph[Input<std::vector<Tensor>>(kTensorsTag)], graph));
embedding_result_out >> graph[Output<EmbeddingResult>(kEmbeddingResultTag)];
embedding_result_out >> graph[Output<EmbeddingResult>(kEmbeddingsTag)];
return graph.GetConfig();
}
@@ -220,13 +220,13 @@ class EmbeddingPostprocessingGraph : public mediapipe::Subgraph {
.GetOptions<mediapipe::TensorsToEmbeddingsCalculatorOptions>()
.CopyFrom(options.tensors_to_embeddings_options());
dequantized_tensors >> tensors_to_embeddings_node.In(kTensorsTag);
return tensors_to_embeddings_node[Output<EmbeddingResult>(
kEmbeddingResultTag)];
return tensors_to_embeddings_node[Output<EmbeddingResult>(kEmbeddingsTag)];
}
};
REGISTER_MEDIAPIPE_GRAPH(
::mediapipe::tasks::components::EmbeddingPostprocessingGraph);
::mediapipe::tasks::components::processors::EmbeddingPostprocessingGraph);
} // namespace processors
} // namespace components
} // namespace tasks
} // namespace mediapipe
mediapipe/tasks/cc/components/embedding_postprocessing_graph.h → mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph.h
+9 -7
View File
@@ -13,17 +13,18 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef MEDIAPIPE_TASKS_CC_COMPONENTS_EMBEDDING_POSTPROCESSING_GRAPH_H_
#define MEDIAPIPE_TASKS_CC_COMPONENTS_EMBEDDING_POSTPROCESSING_GRAPH_H_
#ifndef MEDIAPIPE_TASKS_CC_COMPONENTS_PROCESSORS_EMBEDDING_POSTPROCESSING_GRAPH_H_
#define MEDIAPIPE_TASKS_CC_COMPONENTS_PROCESSORS_EMBEDDING_POSTPROCESSING_GRAPH_H_
#include "absl/status/status.h"
#include "mediapipe/tasks/cc/components/proto/embedder_options.pb.h"
#include "mediapipe/tasks/cc/components/proto/embedding_postprocessing_graph_options.pb.h"
#include "mediapipe/tasks/cc/components/processors/proto/embedder_options.pb.h"
#include "mediapipe/tasks/cc/components/processors/proto/embedding_postprocessing_graph_options.pb.h"
#include "mediapipe/tasks/cc/core/model_resources.h"
namespace mediapipe {
namespace tasks {
namespace components {
namespace processors {
// Configures an EmbeddingPostprocessingGraph using the provided model resources
// and EmbedderOptions.
@@ -44,18 +45,19 @@ namespace components {
// The output tensors of an InferenceCalculator, to convert into
// EmbeddingResult objects. Expected to be of type kFloat32 or kUInt8.
// Outputs:
// EMBEDDING_RESULT - EmbeddingResult
// EMBEDDINGS - EmbeddingResult
// The output EmbeddingResult.
//
// TODO: add support for additional optional "TIMESTAMPS" input for
// embeddings aggregation.
absl::Status ConfigureEmbeddingPostprocessing(
const tasks::core::ModelResources& model_resources,
const tasks::components::proto::EmbedderOptions& embedder_options,
const proto::EmbedderOptions& embedder_options,
proto::EmbeddingPostprocessingGraphOptions* options);
} // namespace processors
} // namespace components
} // namespace tasks
} // namespace mediapipe
#endif // MEDIAPIPE_TASKS_CC_COMPONENTS_EMBEDDING_POSTPROCESSING_GRAPH_H_
#endif // MEDIAPIPE_TASKS_CC_COMPONENTS_PROCESSORS_EMBEDDING_POSTPROCESSING_GRAPH_H_
mediapipe/tasks/cc/components/embedding_postprocessing_graph_test.cc → mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph_test.cc
+31 -29
View File
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "mediapipe/tasks/cc/components/embedding_postprocessing_graph.h"
#include "mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph.h"
#include <memory>
@@ -25,8 +25,8 @@ limitations under the License.
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/tasks/cc/components/proto/embedder_options.pb.h"
#include "mediapipe/tasks/cc/components/proto/embedding_postprocessing_graph_options.pb.h"
#include "mediapipe/tasks/cc/components/processors/proto/embedder_options.pb.h"
#include "mediapipe/tasks/cc/components/processors/proto/embedding_postprocessing_graph_options.pb.h"
#include "mediapipe/tasks/cc/core/model_resources.h"
#include "mediapipe/tasks/cc/core/proto/external_file.pb.h"
#include "tensorflow/lite/core/shims/cc/shims_test_util.h"
@@ -34,12 +34,10 @@ limitations under the License.
namespace mediapipe {
namespace tasks {
namespace components {
namespace processors {
namespace {
using ::mediapipe::file::JoinPath;
using ::mediapipe::tasks::components::proto::EmbedderOptions;
using ::mediapipe::tasks::components::proto::
EmbeddingPostprocessingGraphOptions;
using ::mediapipe::tasks::core::ModelResources;
constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/";
@@ -69,68 +67,72 @@ TEST_F(ConfigureTest, SucceedsWithQuantizedModelWithMetadata) {
MP_ASSERT_OK_AND_ASSIGN(
auto model_resources,
CreateModelResourcesForModel(kQuantizedImageClassifierWithMetadata));
EmbedderOptions options_in;
proto::EmbedderOptions options_in;
options_in.set_l2_normalize(true);
EmbeddingPostprocessingGraphOptions options_out;
proto::EmbeddingPostprocessingGraphOptions options_out;
MP_ASSERT_OK(ConfigureEmbeddingPostprocessing(*model_resources, options_in,
&options_out));
EXPECT_THAT(
options_out,
EqualsProto(ParseTextProtoOrDie<EmbeddingPostprocessingGraphOptions>(
R"pb(tensors_to_embeddings_options {
embedder_options { l2_normalize: true }
head_names: "probability"
}
has_quantized_outputs: true)pb")));
EqualsProto(
ParseTextProtoOrDie<proto::EmbeddingPostprocessingGraphOptions>(
R"pb(tensors_to_embeddings_options {
embedder_options { l2_normalize: true }
head_names: "probability"
}
has_quantized_outputs: true)pb")));
}
TEST_F(ConfigureTest, SucceedsWithQuantizedModelWithoutMetadata) {
MP_ASSERT_OK_AND_ASSIGN(
auto model_resources,
CreateModelResourcesForModel(kQuantizedImageClassifierWithoutMetadata));
EmbedderOptions options_in;
proto::EmbedderOptions options_in;
options_in.set_quantize(true);
EmbeddingPostprocessingGraphOptions options_out;
proto::EmbeddingPostprocessingGraphOptions options_out;
MP_ASSERT_OK(ConfigureEmbeddingPostprocessing(*model_resources, options_in,
&options_out));
EXPECT_THAT(
options_out,
EqualsProto(ParseTextProtoOrDie<EmbeddingPostprocessingGraphOptions>(
R"pb(tensors_to_embeddings_options {
embedder_options { quantize: true }
}
has_quantized_outputs: true)pb")));
EqualsProto(
ParseTextProtoOrDie<proto::EmbeddingPostprocessingGraphOptions>(
R"pb(tensors_to_embeddings_options {
embedder_options { quantize: true }
}
has_quantized_outputs: true)pb")));
}
TEST_F(ConfigureTest, SucceedsWithFloatModelWithMetadata) {
MP_ASSERT_OK_AND_ASSIGN(auto model_resources,
CreateModelResourcesForModel(kMobileNetV3Embedder));
EmbedderOptions options_in;
proto::EmbedderOptions options_in;
options_in.set_quantize(true);
options_in.set_l2_normalize(true);
EmbeddingPostprocessingGraphOptions options_out;
proto::EmbeddingPostprocessingGraphOptions options_out;
MP_ASSERT_OK(ConfigureEmbeddingPostprocessing(*model_resources, options_in,
&options_out));
EXPECT_THAT(
options_out,
EqualsProto(ParseTextProtoOrDie<EmbeddingPostprocessingGraphOptions>(
R"pb(tensors_to_embeddings_options {
embedder_options { quantize: true l2_normalize: true }
head_names: "feature"
}
has_quantized_outputs: false)pb")));
EqualsProto(
ParseTextProtoOrDie<proto::EmbeddingPostprocessingGraphOptions>(
R"pb(tensors_to_embeddings_options {
embedder_options { quantize: true l2_normalize: true }
head_names: "feature"
}
has_quantized_outputs: false)pb")));
}
// TODO: add E2E Postprocessing tests once timestamp aggregation is
// supported.
} // namespace
} // namespace processors
} // namespace components
} // namespace tasks
} // namespace mediapipe
mediapipe/tasks/cc/components/processors/proto/BUILD
+15
View File
@@ -34,3 +34,18 @@ mediapipe_proto_library(
"//mediapipe/tasks/cc/components/calculators:score_calibration_calculator_proto",
],
)
mediapipe_proto_library(
name = "embedder_options_proto",
srcs = ["embedder_options.proto"],
)
mediapipe_proto_library(
name = "embedding_postprocessing_graph_options_proto",
srcs = ["embedding_postprocessing_graph_options.proto"],
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
"//mediapipe/tasks/cc/components/calculators:tensors_to_embeddings_calculator_proto",
],
)
mediapipe/tasks/cc/components/proto/embedder_options.proto → mediapipe/tasks/cc/components/processors/proto/embedder_options.proto
+4 -1
View File
@@ -15,7 +15,10 @@ limitations under the License.
syntax = "proto2";
package mediapipe.tasks.components.proto;
package mediapipe.tasks.components.processors.proto;
option java_package = "com.google.mediapipe.tasks.components.processors.proto";
option java_outer_classname = "EmbedderOptionsProto";
// Shared options used by all embedding extraction tasks.
message EmbedderOptions {
mediapipe/tasks/cc/components/proto/embedding_postprocessing_graph_options.proto → mediapipe/tasks/cc/components/processors/proto/embedding_postprocessing_graph_options.proto
+1 -1
View File
@@ -15,7 +15,7 @@ limitations under the License.
syntax = "proto2";
package mediapipe.tasks.components.proto;
package mediapipe.tasks.components.processors.proto;
import "mediapipe/framework/calculator.proto";
import "mediapipe/tasks/cc/components/calculators/tensors_to_embeddings_calculator.proto";
mediapipe/tasks/cc/components/proto/BUILD
-15
View File
@@ -23,21 +23,6 @@ mediapipe_proto_library(
srcs = ["segmenter_options.proto"],
)
mediapipe_proto_library(
name = "embedder_options_proto",
srcs = ["embedder_options.proto"],
)
mediapipe_proto_library(
name = "embedding_postprocessing_graph_options_proto",
srcs = ["embedding_postprocessing_graph_options.proto"],
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
"//mediapipe/tasks/cc/components/calculators:tensors_to_embeddings_calculator_proto",
],
)
mediapipe_proto_library(
name = "text_preprocessing_graph_options_proto",
srcs = ["text_preprocessing_graph_options.proto"],
mediapipe/tasks/cc/components/utils/BUILD
+2 -2
View File
@@ -26,7 +26,7 @@ cc_library(
hdrs = ["cosine_similarity.h"],
deps = [
"//mediapipe/tasks/cc:common",
"//mediapipe/tasks/cc/components/containers/proto:embeddings_cc_proto",
"//mediapipe/tasks/cc/components/containers:embedding_result",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
@@ -39,7 +39,7 @@ cc_test(
deps = [
":cosine_similarity",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/tasks/cc/components/containers/proto:embeddings_cc_proto",
"//mediapipe/tasks/cc/components/containers:embedding_result",
],
)
mediapipe/tasks/cc/components/utils/cosine_similarity.cc
+18 -22
View File
@@ -21,7 +21,7 @@ limitations under the License.
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "mediapipe/tasks/cc/common.h"
#include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h"
#include "mediapipe/tasks/cc/components/containers/embedding_result.h"
namespace mediapipe {
namespace tasks {
@@ -30,7 +30,7 @@ namespace utils {
namespace {
using ::mediapipe::tasks::components::containers::proto::EmbeddingEntry;
using ::mediapipe::tasks::components::containers::Embedding;
template <typename T>
absl::StatusOr<double> ComputeCosineSimilarity(const T& u, const T& v,
@@ -66,39 +66,35 @@ absl::StatusOr<double> ComputeCosineSimilarity(const T& u, const T& v,
// an L2-norm of 0.
//
// [1]: https://en.wikipedia.org/wiki/Cosine_similarity
absl::StatusOr<double> CosineSimilarity(const EmbeddingEntry& u,
const EmbeddingEntry& v) {
if (u.has_float_embedding() && v.has_float_embedding()) {
if (u.float_embedding().values().size() !=
v.float_embedding().values().size()) {
absl::StatusOr<double> CosineSimilarity(const Embedding& u,
const Embedding& v) {
if (!u.float_embedding.empty() && !v.float_embedding.empty()) {
if (u.float_embedding.size() != v.float_embedding.size()) {
return CreateStatusWithPayload(
absl::StatusCode::kInvalidArgument,
absl::StrFormat("Cannot compute cosine similarity between embeddings "
"of different sizes (%d vs. %d)",
u.float_embedding().values().size(),
v.float_embedding().values().size()),
u.float_embedding.size(), v.float_embedding.size()),
MediaPipeTasksStatus::kInvalidArgumentError);
}
return ComputeCosineSimilarity(u.float_embedding().values().data(),
v.float_embedding().values().data(),
u.float_embedding().values().size());
return ComputeCosineSimilarity(u.float_embedding.data(),
v.float_embedding.data(),
u.float_embedding.size());
}
if (u.has_quantized_embedding() && v.has_quantized_embedding()) {
if (u.quantized_embedding().values().size() !=
v.quantized_embedding().values().size()) {
if (!u.quantized_embedding.empty() && !v.quantized_embedding.empty()) {
if (u.quantized_embedding.size() != v.quantized_embedding.size()) {
return CreateStatusWithPayload(
absl::StatusCode::kInvalidArgument,
absl::StrFormat("Cannot compute cosine similarity between embeddings "
"of different sizes (%d vs. %d)",
u.quantized_embedding().values().size(),
v.quantized_embedding().values().size()),
u.quantized_embedding.size(),
v.quantized_embedding.size()),
MediaPipeTasksStatus::kInvalidArgumentError);
}
return ComputeCosineSimilarity(reinterpret_cast<const int8_t*>(
u.quantized_embedding().values().data()),
reinterpret_cast<const int8_t*>(
v.quantized_embedding().values().data()),
u.quantized_embedding().values().size());
return ComputeCosineSimilarity(
reinterpret_cast<const int8_t*>(u.quantized_embedding.data()),
reinterpret_cast<const int8_t*>(v.quantized_embedding.data()),
u.quantized_embedding.size());
}
return CreateStatusWithPayload(
absl::StatusCode::kInvalidArgument,
mediapipe/tasks/cc/components/utils/cosine_similarity.h
+6 -8
View File
@@ -17,22 +17,20 @@ limitations under the License.
#define MEDIAPIPE_TASKS_CC_COMPONENTS_UTILS_COSINE_SIMILARITY_H_
#include "absl/status/statusor.h"
#include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h"
#include "mediapipe/tasks/cc/components/containers/embedding_result.h"
namespace mediapipe {
namespace tasks {
namespace components {
namespace utils {
// Utility function to compute cosine similarity [1] between two embedding
// entries. May return an InvalidArgumentError if e.g. the feature vectors are
// of different types (quantized vs. float), have different sizes, or have a
// an L2-norm of 0.
// Utility function to compute cosine similarity [1] between two embeddings. May
// return an InvalidArgumentError if e.g. the embeddings are of different types
// (quantized vs. float), have different sizes, or have a an L2-norm of 0.
//
// [1]: https://en.wikipedia.org/wiki/Cosine_similarity
absl::StatusOr<double> CosineSimilarity(
const containers::proto::EmbeddingEntry& u,
const containers::proto::EmbeddingEntry& v);
absl::StatusOr<double> CosineSimilarity(const containers::Embedding& u,
const containers::Embedding& v);
} // namespace utils
} // namespace components
mediapipe/tasks/cc/components/utils/cosine_similarity_test.cc
+23 -25
View File
@@ -22,7 +22,7 @@ limitations under the License.
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h"
#include "mediapipe/tasks/cc/components/containers/embedding_result.h"
namespace mediapipe {
namespace tasks {
@@ -30,29 +30,27 @@ namespace components {
namespace utils {
namespace {
using ::mediapipe::tasks::components::containers::proto::EmbeddingEntry;
using ::mediapipe::tasks::components::containers::Embedding;
using ::testing::HasSubstr;
// Helper function to generate float EmbeddingEntry.
EmbeddingEntry BuildFloatEntry(std::vector<float> values) {
EmbeddingEntry entry;
for (const float value : values) {
entry.mutable_float_embedding()->add_values(value);
}
return entry;
// Helper function to generate float Embedding.
Embedding BuildFloatEmbedding(std::vector<float> values) {
Embedding embedding;
embedding.float_embedding = values;
return embedding;
}
// Helper function to generate quantized EmbeddingEntry.
EmbeddingEntry BuildQuantizedEntry(std::vector<int8_t> values) {
EmbeddingEntry entry;
entry.mutable_quantized_embedding()->set_values(
reinterpret_cast<uint8_t*>(values.data()), values.size());
return entry;
// Helper function to generate quantized Embedding.
Embedding BuildQuantizedEmbedding(std::vector<int8_t> values) {
Embedding embedding;
uint8_t* data = reinterpret_cast<uint8_t*>(values.data());
embedding.quantized_embedding = {data, data + values.size()};
return embedding;
}
TEST(CosineSimilarity, FailsWithQuantizedAndFloatEmbeddings) {
auto u = BuildFloatEntry({0.1, 0.2});
auto v = BuildQuantizedEntry({0, 1});
auto u = BuildFloatEmbedding({0.1, 0.2});
auto v = BuildQuantizedEmbedding({0, 1});
auto status = CosineSimilarity(u, v);
@@ -63,8 +61,8 @@ TEST(CosineSimilarity, FailsWithQuantizedAndFloatEmbeddings) {
}
TEST(CosineSimilarity, FailsWithZeroNorm) {
auto u = BuildFloatEntry({0.1, 0.2});
auto v = BuildFloatEntry({0.0, 0.0});
auto u = BuildFloatEmbedding({0.1, 0.2});
auto v = BuildFloatEmbedding({0.0, 0.0});
auto status = CosineSimilarity(u, v);
@@ -75,8 +73,8 @@ TEST(CosineSimilarity, FailsWithZeroNorm) {
}
TEST(CosineSimilarity, FailsWithDifferentSizes) {
auto u = BuildFloatEntry({0.1, 0.2});
auto v = BuildFloatEntry({0.1, 0.2, 0.3});
auto u = BuildFloatEmbedding({0.1, 0.2});
auto v = BuildFloatEmbedding({0.1, 0.2, 0.3});
auto status = CosineSimilarity(u, v);
@@ -87,8 +85,8 @@ TEST(CosineSimilarity, FailsWithDifferentSizes) {
}
TEST(CosineSimilarity, SucceedsWithFloatEntries) {
auto u = BuildFloatEntry({1.0, 0.0, 0.0, 0.0});
auto v = BuildFloatEntry({0.5, 0.5, 0.5, 0.5});
auto u = BuildFloatEmbedding({1.0, 0.0, 0.0, 0.0});
auto v = BuildFloatEmbedding({0.5, 0.5, 0.5, 0.5});
MP_ASSERT_OK_AND_ASSIGN(auto result, CosineSimilarity(u, v));
@@ -96,8 +94,8 @@ TEST(CosineSimilarity, SucceedsWithFloatEntries) {
}
TEST(CosineSimilarity, SucceedsWithQuantizedEntries) {
auto u = BuildQuantizedEntry({127, 0, 0, 0});
auto v = BuildQuantizedEntry({-128, 0, 0, 0});
auto u = BuildQuantizedEmbedding({127, 0, 0, 0});
auto v = BuildQuantizedEmbedding({-128, 0, 0, 0});
MP_ASSERT_OK_AND_ASSIGN(auto result, CosineSimilarity(u, v));
mediapipe/tasks/cc/vision/image_embedder/BUILD
+5 -4
View File
@@ -26,12 +26,12 @@ cc_library(
"//mediapipe/framework/api2:port",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:rect_cc_proto",
"//mediapipe/tasks/cc/components:embedding_postprocessing_graph",
"//mediapipe/tasks/cc/components:image_preprocessing",
"//mediapipe/tasks/cc/components:image_preprocessing_options_cc_proto",
"//mediapipe/tasks/cc/components/calculators:tensors_to_embeddings_calculator",
"//mediapipe/tasks/cc/components/containers/proto:embeddings_cc_proto",
"//mediapipe/tasks/cc/components/proto:embedding_postprocessing_graph_options_cc_proto",
"//mediapipe/tasks/cc/components/processors:embedding_postprocessing_graph",
"//mediapipe/tasks/cc/components/processors/proto:embedding_postprocessing_graph_options_cc_proto",
"//mediapipe/tasks/cc/core:model_task_graph",
"//mediapipe/tasks/cc/vision/image_embedder/proto:image_embedder_graph_options_cc_proto",
"@com_google_absl//absl/status:statusor",
@@ -49,9 +49,10 @@ cc_library(
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:rect_cc_proto",
"//mediapipe/framework/tool:options_map",
"//mediapipe/tasks/cc/components:embedder_options",
"//mediapipe/tasks/cc/components/containers:embedding_result",
"//mediapipe/tasks/cc/components/containers/proto:embeddings_cc_proto",
"//mediapipe/tasks/cc/components/proto:embedder_options_cc_proto",
"//mediapipe/tasks/cc/components/processors:embedder_options",
"//mediapipe/tasks/cc/components/processors/proto:embedder_options_cc_proto",
"//mediapipe/tasks/cc/components/utils:cosine_similarity",
"//mediapipe/tasks/cc/core:base_options",
"//mediapipe/tasks/cc/core:task_runner",
mediapipe/tasks/cc/vision/image_embedder/image_embedder.cc
+22 -17
View File
@@ -21,9 +21,10 @@ limitations under the License.
#include "mediapipe/framework/api2/builder.h"
#include "mediapipe/framework/formats/rect.pb.h"
#include "mediapipe/framework/tool/options_map.h"
#include "mediapipe/tasks/cc/components/containers/embedding_result.h"
#include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h"
#include "mediapipe/tasks/cc/components/embedder_options.h"
#include "mediapipe/tasks/cc/components/proto/embedder_options.pb.h"
#include "mediapipe/tasks/cc/components/processors/embedder_options.h"
#include "mediapipe/tasks/cc/components/processors/proto/embedder_options.pb.h"
#include "mediapipe/tasks/cc/components/utils/cosine_similarity.h"
#include "mediapipe/tasks/cc/core/base_options.h"
#include "mediapipe/tasks/cc/core/proto/base_options.pb.h"
@@ -41,8 +42,8 @@ namespace image_embedder {
namespace {
constexpr char kEmbeddingResultStreamName[] = "embedding_result_out";
constexpr char kEmbeddingResultTag[] = "EMBEDDING_RESULT";
constexpr char kEmbeddingsStreamName[] = "embeddings_out";
constexpr char kEmbeddingsTag[] = "EMBEDDINGS";
constexpr char kImageInStreamName[] = "image_in";
constexpr char kImageOutStreamName[] = "image_out";
constexpr char kImageTag[] = "IMAGE";
@@ -53,7 +54,7 @@ constexpr char kGraphTypeName[] =
"mediapipe.tasks.vision.image_embedder.ImageEmbedderGraph";
constexpr int kMicroSecondsPerMilliSecond = 1000;
using ::mediapipe::tasks::components::containers::proto::EmbeddingEntry;
using ::mediapipe::tasks::components::containers::ConvertToEmbeddingResult;
using ::mediapipe::tasks::components::containers::proto::EmbeddingResult;
using ::mediapipe::tasks::core::PacketMap;
using ::mediapipe::tasks::vision::image_embedder::proto::
@@ -71,13 +72,13 @@ CalculatorGraphConfig CreateGraphConfig(
graph.In(kNormRectTag).SetName(kNormRectStreamName);
auto& task_graph = graph.AddNode(kGraphTypeName);
task_graph.GetOptions<ImageEmbedderGraphOptions>().Swap(options_proto.get());
task_graph.Out(kEmbeddingResultTag).SetName(kEmbeddingResultStreamName) >>
graph.Out(kEmbeddingResultTag);
task_graph.Out(kEmbeddingsTag).SetName(kEmbeddingsStreamName) >>
graph.Out(kEmbeddingsTag);
task_graph.Out(kImageTag).SetName(kImageOutStreamName) >>
graph.Out(kImageTag);
if (enable_flow_limiting) {
return tasks::core::AddFlowLimiterCalculator(
graph, task_graph, {kImageTag, kNormRectTag}, kEmbeddingResultTag);
graph, task_graph, {kImageTag, kNormRectTag}, kEmbeddingsTag);
}
graph.In(kImageTag) >> task_graph.In(kImageTag);
graph.In(kNormRectTag) >> task_graph.In(kNormRectTag);
@@ -95,8 +96,8 @@ std::unique_ptr<ImageEmbedderGraphOptions> ConvertImageEmbedderOptionsToProto(
options_proto->mutable_base_options()->set_use_stream_mode(
options->running_mode != core::RunningMode::IMAGE);
auto embedder_options_proto =
std::make_unique<tasks::components::proto::EmbedderOptions>(
components::ConvertEmbedderOptionsToProto(
std::make_unique<components::processors::proto::EmbedderOptions>(
components::processors::ConvertEmbedderOptionsToProto(
&(options->embedder_options)));
options_proto->mutable_embedder_options()->Swap(embedder_options_proto.get());
return options_proto;
@@ -121,9 +122,10 @@ absl::StatusOr<std::unique_ptr<ImageEmbedder>> ImageEmbedder::Create(
return;
}
Packet embedding_result_packet =
status_or_packets.value()[kEmbeddingResultStreamName];
status_or_packets.value()[kEmbeddingsStreamName];
Packet image_packet = status_or_packets.value()[kImageOutStreamName];
result_callback(embedding_result_packet.Get<EmbeddingResult>(),
result_callback(ConvertToEmbeddingResult(
embedding_result_packet.Get<EmbeddingResult>()),
image_packet.Get<Image>(),
embedding_result_packet.Timestamp().Value() /
kMicroSecondsPerMilliSecond);
@@ -138,7 +140,7 @@ absl::StatusOr<std::unique_ptr<ImageEmbedder>> ImageEmbedder::Create(
std::move(packets_callback));
}
absl::StatusOr<EmbeddingResult> ImageEmbedder::Embed(
absl::StatusOr<ImageEmbedderResult> ImageEmbedder::Embed(
Image image,
std::optional<core::ImageProcessingOptions> image_processing_options) {
if (image.UsesGpu()) {
@@ -155,10 +157,11 @@ absl::StatusOr<EmbeddingResult> ImageEmbedder::Embed(
{{kImageInStreamName, MakePacket<Image>(std::move(image))},
{kNormRectStreamName,
MakePacket<NormalizedRect>(std::move(norm_rect))}}));
return output_packets[kEmbeddingResultStreamName].Get<EmbeddingResult>();
return ConvertToEmbeddingResult(
output_packets[kEmbeddingsStreamName].Get<EmbeddingResult>());
}
absl::StatusOr<EmbeddingResult> ImageEmbedder::EmbedForVideo(
absl::StatusOr<ImageEmbedderResult> ImageEmbedder::EmbedForVideo(
Image image, int64 timestamp_ms,
std::optional<core::ImageProcessingOptions> image_processing_options) {
if (image.UsesGpu()) {
@@ -178,7 +181,8 @@ absl::StatusOr<EmbeddingResult> ImageEmbedder::EmbedForVideo(
{kNormRectStreamName,
MakePacket<NormalizedRect>(std::move(norm_rect))
.At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}}));
return output_packets[kEmbeddingResultStreamName].Get<EmbeddingResult>();
return ConvertToEmbeddingResult(
output_packets[kEmbeddingsStreamName].Get<EmbeddingResult>());
}
absl::Status ImageEmbedder::EmbedAsync(
@@ -202,7 +206,8 @@ absl::Status ImageEmbedder::EmbedAsync(
}
absl::StatusOr<double> ImageEmbedder::CosineSimilarity(
const EmbeddingEntry& u, const EmbeddingEntry& v) {
const components::containers::Embedding& u,
const components::containers::Embedding& v) {
return components::utils::CosineSimilarity(u, v);
}
mediapipe/tasks/cc/vision/image_embedder/image_embedder.h
+16 -14
View File
@@ -21,8 +21,8 @@ limitations under the License.
#include "absl/status/statusor.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h"
#include "mediapipe/tasks/cc/components/embedder_options.h"
#include "mediapipe/tasks/cc/components/containers/embedding_result.h"
#include "mediapipe/tasks/cc/components/processors/embedder_options.h"
#include "mediapipe/tasks/cc/core/base_options.h"
#include "mediapipe/tasks/cc/vision/core/base_vision_task_api.h"
#include "mediapipe/tasks/cc/vision/core/image_processing_options.h"
@@ -33,6 +33,10 @@ namespace tasks {
namespace vision {
namespace image_embedder {
// Alias the shared EmbeddingResult struct as result typo.
using ImageEmbedderResult =
::mediapipe::tasks::components::containers::EmbeddingResult;
// The options for configuring a MediaPipe image embedder task.
struct ImageEmbedderOptions {
// Base options for configuring MediaPipe Tasks, such as specifying the model
@@ -50,14 +54,12 @@ struct ImageEmbedderOptions {
// Options for configuring the embedder behavior, such as L2-normalization or
// scalar-quantization.
components::EmbedderOptions embedder_options;
components::processors::EmbedderOptions embedder_options;
// The user-defined result callback for processing live stream data.
// The result callback should only be specified when the running mode is set
// to RunningMode::LIVE_STREAM.
std::function<void(
absl::StatusOr<components::containers::proto::EmbeddingResult>,
const Image&, int64)>
std::function<void(absl::StatusOr<ImageEmbedderResult>, const Image&, int64)>
result_callback = nullptr;
};
@@ -104,7 +106,7 @@ class ImageEmbedder : core::BaseVisionTaskApi {
// running mode.
//
// The image can be of any size with format RGB or RGBA.
absl::StatusOr<components::containers::proto::EmbeddingResult> Embed(
absl::StatusOr<ImageEmbedderResult> Embed(
mediapipe::Image image,
std::optional<core::ImageProcessingOptions> image_processing_options =
std::nullopt);
@@ -127,7 +129,7 @@ class ImageEmbedder : core::BaseVisionTaskApi {
// The image can be of any size with format RGB or RGBA. It's required to
// provide the video frame's timestamp (in milliseconds). The input timestamps
// must be monotonically increasing.
absl::StatusOr<components::containers::proto::EmbeddingResult> EmbedForVideo(
absl::StatusOr<ImageEmbedderResult> EmbedForVideo(
mediapipe::Image image, int64 timestamp_ms,
std::optional<core::ImageProcessingOptions> image_processing_options =
std::nullopt);
@@ -168,15 +170,15 @@ class ImageEmbedder : core::BaseVisionTaskApi {
// Shuts down the ImageEmbedder when all works are done.
absl::Status Close() { return runner_->Close(); }
// Utility function to compute cosine similarity [1] between two embedding
// entries. May return an InvalidArgumentError if e.g. the feature vectors are
// of different types (quantized vs. float), have different sizes, or have a
// an L2-norm of 0.
// Utility function to compute cosine similarity [1] between two embeddings.
// May return an InvalidArgumentError if e.g. the embeddings are of different
// types (quantized vs. float), have different sizes, or have a an L2-norm of
// 0.
//
// [1]: https://en.wikipedia.org/wiki/Cosine_similarity
static absl::StatusOr<double> CosineSimilarity(
const components::containers::proto::EmbeddingEntry& u,
const components::containers::proto::EmbeddingEntry& v);
const components::containers::Embedding& u,
const components::containers::Embedding& v);
};
} // namespace image_embedder
mediapipe/tasks/cc/vision/image_embedder/image_embedder_graph.cc
+11 -12
View File
@@ -20,10 +20,10 @@ limitations under the License.
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/rect.pb.h"
#include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h"
#include "mediapipe/tasks/cc/components/embedding_postprocessing_graph.h"
#include "mediapipe/tasks/cc/components/image_preprocessing.h"
#include "mediapipe/tasks/cc/components/image_preprocessing_options.pb.h"
#include "mediapipe/tasks/cc/components/proto/embedding_postprocessing_graph_options.pb.h"
#include "mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph.h"
#include "mediapipe/tasks/cc/components/processors/proto/embedding_postprocessing_graph_options.pb.h"
#include "mediapipe/tasks/cc/core/model_task_graph.h"
#include "mediapipe/tasks/cc/vision/image_embedder/proto/image_embedder_graph_options.pb.h"
@@ -40,10 +40,8 @@ using ::mediapipe::api2::builder::GenericNode;
using ::mediapipe::api2::builder::Graph;
using ::mediapipe::api2::builder::Source;
using ::mediapipe::tasks::components::containers::proto::EmbeddingResult;
using ::mediapipe::tasks::components::proto::
EmbeddingPostprocessingGraphOptions;
constexpr char kEmbeddingResultTag[] = "EMBEDDING_RESULT";
constexpr char kEmbeddingsTag[] = "EMBEDDINGS";
constexpr char kImageTag[] = "IMAGE";
constexpr char kNormRectTag[] = "NORM_RECT";
constexpr char kTensorsTag[] = "TENSORS";
@@ -67,7 +65,7 @@ struct ImageEmbedderOutputStreams {
// Describes region of image to perform embedding extraction on.
// @Optional: rect covering the whole image is used if not specified.
// Outputs:
// EMBEDDING_RESULT - EmbeddingResult
// EMBEDDINGS - EmbeddingResult
// The embedding result.
// IMAGE - Image
// The image that embedding extraction runs on.
@@ -76,7 +74,7 @@ struct ImageEmbedderOutputStreams {
// node {
// calculator: "mediapipe.tasks.vision.image_embedder.ImageEmbedderGraph"
// input_stream: "IMAGE:image_in"
// output_stream: "EMBEDDING_RESULT:embedding_result_out"
// output_stream: "EMBEDDINGS:embedding_result_out"
// output_stream: "IMAGE:image_out"
// options {
// [mediapipe.tasks.vision.image_embedder.proto.ImageEmbedderOptions.ext]
@@ -107,7 +105,7 @@ class ImageEmbedderGraph : public core::ModelTaskGraph {
graph[Input<Image>(kImageTag)],
graph[Input<NormalizedRect>::Optional(kNormRectTag)], graph));
output_streams.embedding_result >>
graph[Output<EmbeddingResult>(kEmbeddingResultTag)];
graph[Output<EmbeddingResult>(kEmbeddingsTag)];
output_streams.image >> graph[Output<Image>(kImageTag)];
return graph.GetConfig();
}
@@ -152,16 +150,17 @@ class ImageEmbedderGraph : public core::ModelTaskGraph {
// Adds postprocessing calculators and connects its input stream to the
// inference results.
auto& postprocessing = graph.AddNode(
"mediapipe.tasks.components.EmbeddingPostprocessingGraph");
MP_RETURN_IF_ERROR(components::ConfigureEmbeddingPostprocessing(
"mediapipe.tasks.components.processors.EmbeddingPostprocessingGraph");
MP_RETURN_IF_ERROR(components::processors::ConfigureEmbeddingPostprocessing(
model_resources, task_options.embedder_options(),
&postprocessing.GetOptions<EmbeddingPostprocessingGraphOptions>()));
&postprocessing.GetOptions<components::processors::proto::
EmbeddingPostprocessingGraphOptions>()));
inference.Out(kTensorsTag) >> postprocessing.In(kTensorsTag);
// Outputs the embedding results.
return ImageEmbedderOutputStreams{
/*embedding_result=*/postprocessing[Output<EmbeddingResult>(
kEmbeddingResultTag)],
kEmbeddingsTag)],
/*image=*/preprocessing[Output<Image>(kImageTag)]};
}
};
mediapipe/tasks/cc/vision/image_embedder/image_embedder_test.cc
+49 -60
View File
@@ -26,7 +26,7 @@ limitations under the License.
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/tasks/cc/components/containers/proto/embeddings.pb.h"
#include "mediapipe/tasks/cc/components/containers/embedding_result.h"
#include "mediapipe/tasks/cc/vision/core/running_mode.h"
#include "mediapipe/tasks/cc/vision/utils/image_utils.h"
#include "tensorflow/lite/core/api/op_resolver.h"
@@ -42,7 +42,6 @@ namespace {
using ::mediapipe::file::JoinPath;
using ::mediapipe::tasks::components::containers::Rect;
using ::mediapipe::tasks::components::containers::proto::EmbeddingResult;
using ::mediapipe::tasks::vision::core::ImageProcessingOptions;
using ::testing::HasSubstr;
using ::testing::Optional;
@@ -54,18 +53,14 @@ constexpr double kSimilarityTolerancy = 1e-6;
// Utility function to check the sizes, head_index and head_names of a result
// procuded by kMobileNetV3Embedder.
void CheckMobileNetV3Result(const EmbeddingResult& result, bool quantized) {
EXPECT_EQ(result.embeddings().size(), 1);
EXPECT_EQ(result.embeddings(0).head_index(), 0);
EXPECT_EQ(result.embeddings(0).head_name(), "feature");
EXPECT_EQ(result.embeddings(0).entries().size(), 1);
void CheckMobileNetV3Result(const ImageEmbedderResult& result, bool quantized) {
EXPECT_EQ(result.embeddings.size(), 1);
EXPECT_EQ(result.embeddings[0].head_index, 0);
EXPECT_EQ(result.embeddings[0].head_name, "feature");
if (quantized) {
EXPECT_EQ(
result.embeddings(0).entries(0).quantized_embedding().values().size(),
1024);
EXPECT_EQ(result.embeddings[0].quantized_embedding.size(), 1024);
} else {
EXPECT_EQ(result.embeddings(0).entries(0).float_embedding().values().size(),
1024);
EXPECT_EQ(result.embeddings[0].float_embedding.size(), 1024);
}
}
@@ -154,7 +149,7 @@ TEST_F(CreateTest, FailsWithIllegalCallbackInImageOrVideoMode) {
options->base_options.model_asset_path =
JoinPath("./", kTestDataDirectory, kMobileNetV3Embedder);
options->running_mode = running_mode;
options->result_callback = [](absl::StatusOr<EmbeddingResult>,
options->result_callback = [](absl::StatusOr<ImageEmbedderResult>,
const Image& image, int64 timestamp_ms) {};
auto image_embedder = ImageEmbedder::Create(std::move(options));
@@ -231,19 +226,18 @@ TEST_F(ImageModeTest, SucceedsWithoutL2Normalization) {
JoinPath("./", kTestDataDirectory, "burger_crop.jpg")));
// Extract both embeddings.
MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& image_result,
MP_ASSERT_OK_AND_ASSIGN(const ImageEmbedderResult& image_result,
image_embedder->Embed(image));
MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& crop_result,
MP_ASSERT_OK_AND_ASSIGN(const ImageEmbedderResult& crop_result,
image_embedder->Embed(crop));
// Check results.
CheckMobileNetV3Result(image_result, false);
CheckMobileNetV3Result(crop_result, false);
// CheckCosineSimilarity.
MP_ASSERT_OK_AND_ASSIGN(
double similarity,
ImageEmbedder::CosineSimilarity(image_result.embeddings(0).entries(0),
crop_result.embeddings(0).entries(0)));
MP_ASSERT_OK_AND_ASSIGN(double similarity, ImageEmbedder::CosineSimilarity(
image_result.embeddings[0],
crop_result.embeddings[0]));
double expected_similarity = 0.925519;
EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy);
}
@@ -264,19 +258,18 @@ TEST_F(ImageModeTest, SucceedsWithL2Normalization) {
JoinPath("./", kTestDataDirectory, "burger_crop.jpg")));
// Extract both embeddings.
MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& image_result,
MP_ASSERT_OK_AND_ASSIGN(const ImageEmbedderResult& image_result,
image_embedder->Embed(image));
MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& crop_result,
MP_ASSERT_OK_AND_ASSIGN(const ImageEmbedderResult& crop_result,
image_embedder->Embed(crop));
// Check results.
CheckMobileNetV3Result(image_result, false);
CheckMobileNetV3Result(crop_result, false);
// CheckCosineSimilarity.
MP_ASSERT_OK_AND_ASSIGN(
double similarity,
ImageEmbedder::CosineSimilarity(image_result.embeddings(0).entries(0),
crop_result.embeddings(0).entries(0)));
MP_ASSERT_OK_AND_ASSIGN(double similarity, ImageEmbedder::CosineSimilarity(
image_result.embeddings[0],
crop_result.embeddings[0]));
double expected_similarity = 0.925519;
EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy);
}
@@ -297,19 +290,18 @@ TEST_F(ImageModeTest, SucceedsWithQuantization) {
JoinPath("./", kTestDataDirectory, "burger_crop.jpg")));
// Extract both embeddings.
MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& image_result,
MP_ASSERT_OK_AND_ASSIGN(const ImageEmbedderResult& image_result,
image_embedder->Embed(image));
MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& crop_result,
MP_ASSERT_OK_AND_ASSIGN(const ImageEmbedderResult& crop_result,
image_embedder->Embed(crop));
// Check results.
CheckMobileNetV3Result(image_result, true);
CheckMobileNetV3Result(crop_result, true);
// CheckCosineSimilarity.
MP_ASSERT_OK_AND_ASSIGN(
double similarity,
ImageEmbedder::CosineSimilarity(image_result.embeddings(0).entries(0),
crop_result.embeddings(0).entries(0)));
MP_ASSERT_OK_AND_ASSIGN(double similarity, ImageEmbedder::CosineSimilarity(
image_result.embeddings[0],
crop_result.embeddings[0]));
double expected_similarity = 0.926791;
EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy);
}
@@ -333,19 +325,18 @@ TEST_F(ImageModeTest, SucceedsWithRegionOfInterest) {
// Extract both embeddings.
MP_ASSERT_OK_AND_ASSIGN(
const EmbeddingResult& image_result,
const ImageEmbedderResult& image_result,
image_embedder->Embed(image, image_processing_options));
MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& crop_result,
MP_ASSERT_OK_AND_ASSIGN(const ImageEmbedderResult& crop_result,
image_embedder->Embed(crop));
// Check results.
CheckMobileNetV3Result(image_result, false);
CheckMobileNetV3Result(crop_result, false);
// CheckCosineSimilarity.
MP_ASSERT_OK_AND_ASSIGN(
double similarity,
ImageEmbedder::CosineSimilarity(image_result.embeddings(0).entries(0),
crop_result.embeddings(0).entries(0)));
MP_ASSERT_OK_AND_ASSIGN(double similarity, ImageEmbedder::CosineSimilarity(
image_result.embeddings[0],
crop_result.embeddings[0]));
double expected_similarity = 0.999931;
EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy);
}
@@ -367,20 +358,19 @@ TEST_F(ImageModeTest, SucceedsWithRotation) {
image_processing_options.rotation_degrees = -90;
// Extract both embeddings.
MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& image_result,
MP_ASSERT_OK_AND_ASSIGN(const ImageEmbedderResult& image_result,
image_embedder->Embed(image));
MP_ASSERT_OK_AND_ASSIGN(
const EmbeddingResult& rotated_result,
const ImageEmbedderResult& rotated_result,
image_embedder->Embed(rotated, image_processing_options));
// Check results.
CheckMobileNetV3Result(image_result, false);
CheckMobileNetV3Result(rotated_result, false);
// CheckCosineSimilarity.
MP_ASSERT_OK_AND_ASSIGN(
double similarity,
ImageEmbedder::CosineSimilarity(image_result.embeddings(0).entries(0),
rotated_result.embeddings(0).entries(0)));
MP_ASSERT_OK_AND_ASSIGN(double similarity, ImageEmbedder::CosineSimilarity(
image_result.embeddings[0],
rotated_result.embeddings[0]));
double expected_similarity = 0.572265;
EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy);
}
@@ -403,20 +393,19 @@ TEST_F(ImageModeTest, SucceedsWithRegionOfInterestAndRotation) {
/*rotation_degrees=*/-90};
// Extract both embeddings.
MP_ASSERT_OK_AND_ASSIGN(const EmbeddingResult& crop_result,
MP_ASSERT_OK_AND_ASSIGN(const ImageEmbedderResult& crop_result,
image_embedder->Embed(crop));
MP_ASSERT_OK_AND_ASSIGN(
const EmbeddingResult& rotated_result,
const ImageEmbedderResult& rotated_result,
image_embedder->Embed(rotated, image_processing_options));
// Check results.
CheckMobileNetV3Result(crop_result, false);
CheckMobileNetV3Result(rotated_result, false);
// CheckCosineSimilarity.
MP_ASSERT_OK_AND_ASSIGN(
double similarity,
ImageEmbedder::CosineSimilarity(crop_result.embeddings(0).entries(0),
rotated_result.embeddings(0).entries(0)));
MP_ASSERT_OK_AND_ASSIGN(double similarity, ImageEmbedder::CosineSimilarity(
crop_result.embeddings[0],
rotated_result.embeddings[0]));
double expected_similarity = 0.62838;
EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy);
}
@@ -487,16 +476,16 @@ TEST_F(VideoModeTest, Succeeds) {
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageEmbedder> image_embedder,
ImageEmbedder::Create(std::move(options)));
EmbeddingResult previous_results;
ImageEmbedderResult previous_results;
for (int i = 0; i < iterations; ++i) {
MP_ASSERT_OK_AND_ASSIGN(auto results,
image_embedder->EmbedForVideo(image, i));
CheckMobileNetV3Result(results, false);
if (i > 0) {
MP_ASSERT_OK_AND_ASSIGN(double similarity,
ImageEmbedder::CosineSimilarity(
results.embeddings(0).entries(0),
previous_results.embeddings(0).entries(0)));
MP_ASSERT_OK_AND_ASSIGN(
double similarity,
ImageEmbedder::CosineSimilarity(results.embeddings[0],
previous_results.embeddings[0]));
double expected_similarity = 1.000000;
EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy);
}
@@ -515,7 +504,7 @@ TEST_F(LiveStreamModeTest, FailsWithCallingWrongMethod) {
options->base_options.model_asset_path =
JoinPath("./", kTestDataDirectory, kMobileNetV3Embedder);
options->running_mode = core::RunningMode::LIVE_STREAM;
options->result_callback = [](absl::StatusOr<EmbeddingResult>,
options->result_callback = [](absl::StatusOr<ImageEmbedderResult>,
const Image& image, int64 timestamp_ms) {};
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageEmbedder> image_embedder,
ImageEmbedder::Create(std::move(options)));
@@ -546,7 +535,7 @@ TEST_F(LiveStreamModeTest, FailsWithOutOfOrderInputTimestamps) {
options->base_options.model_asset_path =
JoinPath("./", kTestDataDirectory, kMobileNetV3Embedder);
options->running_mode = core::RunningMode::LIVE_STREAM;
options->result_callback = [](absl::StatusOr<EmbeddingResult>,
options->result_callback = [](absl::StatusOr<ImageEmbedderResult>,
const Image& image, int64 timestamp_ms) {};
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageEmbedder> image_embedder,
ImageEmbedder::Create(std::move(options)));
@@ -564,7 +553,7 @@ TEST_F(LiveStreamModeTest, FailsWithOutOfOrderInputTimestamps) {
}
struct LiveStreamModeResults {
EmbeddingResult embedding_result;
ImageEmbedderResult embedding_result;
std::pair<int, int> image_size;
int64 timestamp_ms;
};
@@ -580,7 +569,7 @@ TEST_F(LiveStreamModeTest, Succeeds) {
JoinPath("./", kTestDataDirectory, kMobileNetV3Embedder);
options->running_mode = core::RunningMode::LIVE_STREAM;
options->result_callback =
[&results](absl::StatusOr<EmbeddingResult> embedding_result,
[&results](absl::StatusOr<ImageEmbedderResult> embedding_result,
const Image& image, int64 timestamp_ms) {
MP_ASSERT_OK(embedding_result.status());
results.push_back(
@@ -612,8 +601,8 @@ TEST_F(LiveStreamModeTest, Succeeds) {
MP_ASSERT_OK_AND_ASSIGN(
double similarity,
ImageEmbedder::CosineSimilarity(
result.embedding_result.embeddings(0).entries(0),
results[i - 1].embedding_result.embeddings(0).entries(0)));
result.embedding_result.embeddings[0],
results[i - 1].embedding_result.embeddings[0]));
double expected_similarity = 1.000000;
EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy);
}
mediapipe/tasks/cc/vision/image_embedder/proto/BUILD
+1 -1
View File
@@ -24,7 +24,7 @@ mediapipe_proto_library(
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
"//mediapipe/tasks/cc/components/proto:embedder_options_proto",
"//mediapipe/tasks/cc/components/processors/proto:embedder_options_proto",
"//mediapipe/tasks/cc/core/proto:base_options_proto",
],
)
mediapipe/tasks/cc/vision/image_embedder/proto/image_embedder_graph_options.proto
+2 -2
View File
@@ -18,7 +18,7 @@ syntax = "proto2";
package mediapipe.tasks.vision.image_embedder.proto;
import "mediapipe/framework/calculator.proto";
import "mediapipe/tasks/cc/components/proto/embedder_options.proto";
import "mediapipe/tasks/cc/components/processors/proto/embedder_options.proto";
import "mediapipe/tasks/cc/core/proto/base_options.proto";
message ImageEmbedderGraphOptions {
@@ -31,5 +31,5 @@ message ImageEmbedderGraphOptions {
// Options for configuring the embedder behavior, such as normalization or
// quantization.
optional components.proto.EmbedderOptions embedder_options = 2;
optional components.processors.proto.EmbedderOptions embedder_options = 2;
}
mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_graph.cc
+3 -4
View File
@@ -287,10 +287,9 @@ class ImageSegmenterGraph : public core::ModelTaskGraph {
tensor_to_images[Output<Image>::Multiple(kSegmentationTag)][i]));
}
}
return {{
.segmented_masks = segmented_masks,
.image = preprocessing[Output<Image>(kImageTag)],
}};
return ImageSegmenterOutputs{
/*segmented_masks=*/segmented_masks,
/*image=*/preprocessing[Output<Image>(kImageTag)]};
}
};
mediapipe/tasks/java/com/google/mediapipe/tasks/mediapipe_tasks_aar.bzl
+1
View File
@@ -23,6 +23,7 @@ _CORE_TASKS_JAVA_PROTO_LITE_TARGETS = [
"//mediapipe/tasks/cc/components/containers/proto:embeddings_java_proto_lite",
"//mediapipe/tasks/cc/components/containers/proto:landmarks_detection_result_java_proto_lite",
"//mediapipe/tasks/cc/components/processors/proto:classifier_options_java_proto_lite",
"//mediapipe/tasks/cc/components/processors/proto:embedder_options_java_proto_lite",
"//mediapipe/tasks/cc/core/proto:acceleration_java_proto_lite",
"//mediapipe/tasks/cc/core/proto:base_options_java_proto_lite",
"//mediapipe/tasks/cc/core/proto:external_file_java_proto_lite",
mediapipe/tasks/python/components/containers/BUILD
+5
View File
@@ -18,6 +18,11 @@ package(default_visibility = ["//mediapipe/tasks:internal"])
licenses(["notice"])
py_library(
name = "audio_data",
srcs = ["audio_data.py"],
)
py_library(
name = "bounding_box",
srcs = ["bounding_box.py"],
mediapipe/tasks/python/components/containers/audio_data.py
+109
View File
@@ -0,0 +1,109 @@
# Copyright 2022 The MediaPipe Authors. All Rights Reserved.
#
# 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.
"""MediaPipe audio data."""
import dataclasses
from typing import Optional
import numpy as np
@dataclasses.dataclass
class AudioFormat:
"""Audio format metadata.
Attributes:
num_channels: the number of channels of the audio data.
sample_rate: the audio sample rate.
"""
num_channels: int = 1
sample_rate: Optional[float] = None
class AudioData(object):
"""MediaPipe Tasks' audio container."""
def __init__(
self, buffer_length: int,
audio_format: AudioFormat = AudioFormat()) -> None:
"""Initializes the `AudioData` object.
Args:
buffer_length: the length of the audio buffer.
audio_format: the audio format metadata.
"""
self._audio_format = audio_format
self._buffer = np.zeros([buffer_length, self._audio_format.num_channels],
dtype=np.float32)
def clear(self):
"""Clears the internal buffer and fill it with zeros."""
self._buffer.fill(0)
def load_from_array(self,
src: np.ndarray,
offset: int = 0,
size: int = -1) -> None:
"""Loads the audio data from a NumPy array.
Args:
src: A NumPy source array contains the input audio.
offset: An optional offset for loading a slice of the `src` array to the
buffer.
size: An optional size parameter denoting the number of samples to load
from the `src` array.
Raises:
ValueError: If the input array has an incorrect shape or if
`offset` + `size` exceeds the length of the `src` array.
"""
if src.shape[1] != self._audio_format.num_channels:
raise ValueError(f"Input audio contains an invalid number of channels. "
f"Expect {self._audio_format.num_channels}.")
if size < 0:
size = len(src)
if offset + size > len(src):
raise ValueError(
f"Index out of range. offset {offset} + size {size} should be <= "
f"src's length: {len(src)}")
if len(src) >= len(self._buffer):
# If the internal buffer is shorter than the load target (src), copy
# values from the end of the src array to the internal buffer.
new_offset = offset + size - len(self._buffer)
new_size = len(self._buffer)
self._buffer = src[new_offset:new_offset + new_size].copy()
else:
# Shift the internal buffer backward and add the incoming data to the end
# of the buffer.
shift = size
self._buffer = np.roll(self._buffer, -shift, axis=0)
self._buffer[-shift:, :] = src[offset:offset + size].copy()
@property
def audio_format(self) -> AudioFormat:
"""Gets the audio format of the audio."""
return self._audio_format
@property
def buffer_length(self) -> int:
"""Gets the sample count of the audio."""
return self._buffer.shape[0]
@property
def buffer(self) -> np.ndarray:
"""Gets the internal buffer."""
return self._buffer
mediapipe/tasks/python/test/BUILD
+1 -1
View File
@@ -24,7 +24,7 @@ py_library(
srcs = ["test_utils.py"],
srcs_version = "PY3",
visibility = [
"//mediapipe/model_maker/python/vision/gesture_recognizer:__pkg__",
"//mediapipe/model_maker/python:__subpackages__",
"//mediapipe/tasks:internal",
],
deps = [
mediapipe/tasks/python/test/vision/image_segmenter_test.py
+2 -2
View File
@@ -33,8 +33,8 @@ from mediapipe.tasks.python.vision.core import vision_task_running_mode
_BaseOptions = base_options_module.BaseOptions
_Image = image_module.Image
_ImageFormat = image_frame.ImageFormat
_OutputType = image_segmenter.OutputType
_Activation = image_segmenter.Activation
_OutputType = image_segmenter.ImageSegmenterOptions.OutputType
_Activation = image_segmenter.ImageSegmenterOptions.Activation
_ImageSegmenter = image_segmenter.ImageSegmenter
_ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions
_RUNNING_MODE = vision_task_running_mode.VisionTaskRunningMode
mediapipe/tasks/python/vision/__init__.py
+8
View File
@@ -15,17 +15,25 @@
"""MediaPipe Tasks Vision API."""
import mediapipe.tasks.python.vision.core
import mediapipe.tasks.python.vision.gesture_recognizer
import mediapipe.tasks.python.vision.image_classifier
import mediapipe.tasks.python.vision.image_segmenter
import mediapipe.tasks.python.vision.object_detector
GestureRecognizer = gesture_recognizer.GestureRecognizer
GestureRecognizerOptions = gesture_recognizer.GestureRecognizerOptions
ImageClassifier = image_classifier.ImageClassifier
ImageClassifierOptions = image_classifier.ImageClassifierOptions
ImageSegmenter = image_segmenter.ImageSegmenter
ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions
ObjectDetector = object_detector.ObjectDetector
ObjectDetectorOptions = object_detector.ObjectDetectorOptions
RunningMode = core.vision_task_running_mode.VisionTaskRunningMode
# Remove unnecessary modules to avoid duplication in API docs.
del core
del gesture_recognizer
del image_classifier
del image_segmenter
del object_detector
del mediapipe
mediapipe/tasks/python/vision/image_segmenter.py
+11 -12
View File
@@ -44,18 +44,6 @@ _TASK_GRAPH_NAME = 'mediapipe.tasks.vision.ImageSegmenterGraph'
_MICRO_SECONDS_PER_MILLISECOND = 1000
class OutputType(enum.Enum):
UNSPECIFIED = 0
CATEGORY_MASK = 1
CONFIDENCE_MASK = 2
class Activation(enum.Enum):
NONE = 0
SIGMOID = 1
SOFTMAX = 2
@dataclasses.dataclass
class ImageSegmenterOptions:
"""Options for the image segmenter task.
@@ -74,6 +62,17 @@ class ImageSegmenterOptions:
data. The result callback should only be specified when the running mode
is set to the live stream mode.
"""
class OutputType(enum.Enum):
UNSPECIFIED = 0
CATEGORY_MASK = 1
CONFIDENCE_MASK = 2
class Activation(enum.Enum):
NONE = 0
SIGMOID = 1
SOFTMAX = 2
base_options: _BaseOptions
running_mode: _RunningMode = _RunningMode.IMAGE
output_type: Optional[OutputType] = OutputType.CATEGORY_MASK
mediapipe/tasks/web/components/containers/BUILD
+21
View File
@@ -0,0 +1,21 @@
# This package contains options shared by all MediaPipe Tasks for Web.
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_ts_library")
package(default_visibility = ["//mediapipe/tasks:internal"])
mediapipe_ts_library(
name = "category",
srcs = ["category.d.ts"],
)
mediapipe_ts_library(
name = "classifications",
srcs = ["classifications.d.ts"],
deps = [":category"],
)
mediapipe_ts_library(
name = "landmark",
srcs = ["landmark.d.ts"],
)
mediapipe/tasks/web/components/containers/category.d.ts
Vendored
+38
View File
@@ -0,0 +1,38 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* 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.
*/
/** A classification category. */
export interface Category {
/** The probability score of this label category. */
score: number;
/** The index of the category in the corresponding label file. */
index: number;
/**
* The label of this category object. Defaults to an empty string if there is
* no category.
*/
categoryName: string;
/**
* The display name of the label, which may be translated for different
* locales. For example, a label, "apple", may be translated into Spanish for
* display purpose, so that the `display_name` is "manzana". Defaults to an
* empty string if there is no display name.
*/
displayName: string;
}
mediapipe/tasks/web/components/containers/classifications.d.ts
Vendored
+51
View File
@@ -0,0 +1,51 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* 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.
*/
import {Category} from '../../../../tasks/web/components/containers/category';
/** List of predicted categories with an optional timestamp. */
export interface ClassificationEntry {
/**
* The array of predicted categories, usually sorted by descending scores,
* e.g., from high to low probability.
*/
categories: Category[];
/**
* The optional timestamp (in milliseconds) associated to the classification
* entry. This is useful for time series use cases, e.g., audio
* classification.
*/
timestampMs?: number;
}
/** Classifications for a given classifier head. */
export interface Classifications {
/** A list of classification entries. */
entries: ClassificationEntry[];
/**
* The index of the classifier head these categories refer to. This is
* useful for multi-head models.
*/
headIndex: number;
/**
* The name of the classifier head, which is the corresponding tensor
* metadata name.
*/
headName: string;
}
mediapipe/tasks/web/components/containers/landmark.d.ts
Vendored
+35
View File
@@ -0,0 +1,35 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* 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.
*/
/**
* Landmark represents a point in 3D space with x, y, z coordinates. If
* normalized is true, the landmark coordinates is normalized respect to the
* dimension of image, and the coordinates values are in the range of [0,1].
* Otherwise, it represenet a point in world coordinates.
*/
export class Landmark {
/** The x coordinates of the landmark. */
x: number;
/** The y coordinates of the landmark. */
y: number;
/** The z coordinates of the landmark. */
z: number;
/** Whether this landmark is normalized with respect to the image size. */
normalized: boolean;
}
mediapipe/tasks/web/components/processors/BUILD
+33
View File
@@ -0,0 +1,33 @@
# This package contains options shared by all MediaPipe Tasks for Web.
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_ts_library")
package(default_visibility = ["//mediapipe/tasks:internal"])
mediapipe_ts_library(
name = "classifier_options",
srcs = ["classifier_options.ts"],
deps = [
"//mediapipe/tasks/cc/components/processors/proto:classifier_options_jspb_proto",
"//mediapipe/tasks/web/core:classifier_options",
],
)
mediapipe_ts_library(
name = "classifier_result",
srcs = ["classifier_result.ts"],
deps = [
"//mediapipe/tasks/cc/components/containers/proto:classifications_jspb_proto",
"//mediapipe/tasks/web/components/containers:classifications",
],
)
mediapipe_ts_library(
name = "base_options",
srcs = ["base_options.ts"],
deps = [
"//mediapipe/tasks/cc/core/proto:base_options_jspb_proto",
"//mediapipe/tasks/cc/core/proto:external_file_jspb_proto",
"//mediapipe/tasks/web/core",
],
)
mediapipe/tasks/web/components/processors/base_options.ts
+50
View File
@@ -0,0 +1,50 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* 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.
*/
import {BaseOptions as BaseOptionsProto} from '../../../../tasks/cc/core/proto/base_options_pb';
import {ExternalFile} from '../../../../tasks/cc/core/proto/external_file_pb';
import {BaseOptions} from '../../../../tasks/web/core/base_options';
// The OSS JS API does not support the builder pattern.
// tslint:disable:jspb-use-builder-pattern
/**
* Converts a BaseOptions API object to its Protobuf representation.
* @throws If neither a model assset path or buffer is provided
*/
export async function convertBaseOptionsToProto(baseOptions: BaseOptions):
Promise<BaseOptionsProto> {
if (baseOptions.modelAssetPath && baseOptions.modelAssetBuffer) {
throw new Error(
'Cannot set both baseOptions.modelAssetPath and baseOptions.modelAssetBuffer');
}
if (!baseOptions.modelAssetPath && !baseOptions.modelAssetBuffer) {
throw new Error(
'Either baseOptions.modelAssetPath or baseOptions.modelAssetBuffer must be set');
}
let modelAssetBuffer = baseOptions.modelAssetBuffer;
if (!modelAssetBuffer) {
const response = await fetch(baseOptions.modelAssetPath!.toString());
modelAssetBuffer = new Uint8Array(await response.arrayBuffer());
}
const proto = new BaseOptionsProto();
const externalFile = new ExternalFile();
externalFile.setFileContent(modelAssetBuffer);
proto.setModelAsset(externalFile);
return proto;
}
mediapipe/tasks/web/components/processors/classifier_options.ts
+62
View File
@@ -0,0 +1,62 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* 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.
*/
import {ClassifierOptions as ClassifierOptionsProto} from '../../../../tasks/cc/components/processors/proto/classifier_options_pb';
import {ClassifierOptions} from '../../../../tasks/web/core/classifier_options';
/**
* Converts a ClassifierOptions object to its Proto representation, optionally
* based on existing definition.
* @param options The options object to convert to a Proto. Only options that
* are expliclty provided are set.
* @param baseOptions A base object that options can be merged into.
*/
export function convertClassifierOptionsToProto(
options: ClassifierOptions,
baseOptions?: ClassifierOptionsProto): ClassifierOptionsProto {
const classifierOptions =
baseOptions ? baseOptions.clone() : new ClassifierOptionsProto();
if (options.displayNamesLocale) {
classifierOptions.setDisplayNamesLocale(options.displayNamesLocale);
} else if (options.displayNamesLocale === undefined) {
classifierOptions.clearDisplayNamesLocale();
}
if (options.maxResults) {
classifierOptions.setMaxResults(options.maxResults);
} else if ('maxResults' in options) { // Check for undefined
classifierOptions.clearMaxResults();
}
if (options.scoreThreshold) {
classifierOptions.setScoreThreshold(options.scoreThreshold);
} else if ('scoreThreshold' in options) { // Check for undefined
classifierOptions.clearScoreThreshold();
}
if (options.categoryAllowlist) {
classifierOptions.setCategoryAllowlistList(options.categoryAllowlist);
} else if ('categoryAllowlist' in options) { // Check for undefined
classifierOptions.clearCategoryAllowlistList();
}
if (options.categoryDenylist) {
classifierOptions.setCategoryDenylistList(options.categoryDenylist);
} else if ('categoryDenylist' in options) { // Check for undefined
classifierOptions.clearCategoryDenylistList();
}
return classifierOptions;
}
mediapipe/tasks/web/components/processors/classifier_result.ts
+61
View File
@@ -0,0 +1,61 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* 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.
*/
import {ClassificationEntry as ClassificationEntryProto, ClassificationResult} from '../../../../tasks/cc/components/containers/proto/classifications_pb';
import {ClassificationEntry, Classifications} from '../../../../tasks/web/components/containers/classifications';
const DEFAULT_INDEX = -1;
const DEFAULT_SCORE = 0.0;
/**
* Converts a ClassificationEntry proto to the ClassificationEntry result
* type.
*/
function convertFromClassificationEntryProto(source: ClassificationEntryProto):
ClassificationEntry {
const categories = source.getCategoriesList().map(category => {
return {
index: category.getIndex() ?? DEFAULT_INDEX,
score: category.getScore() ?? DEFAULT_SCORE,
displayName: category.getDisplayName() ?? '',
categoryName: category.getCategoryName() ?? '',
};
});
return {
categories,
timestampMs: source.getTimestampMs(),
};
}
/**
* Converts a ClassificationResult proto to a list of classifications.
*/
export function convertFromClassificationResultProto(
classificationResult: ClassificationResult) : Classifications[] {
const result: Classifications[] = [];
for (const classificationsProto of
classificationResult.getClassificationsList()) {
const classifications: Classifications = {
entries: classificationsProto.getEntriesList().map(
entry => convertFromClassificationEntryProto(entry)),
headIndex: classificationsProto.getHeadIndex() ?? DEFAULT_INDEX,
headName: classificationsProto.getHeadName() ?? '',
};
result.push(classifications);
}
return result;
}
mediapipe/tasks/web/core/BUILD
+21
View File
@@ -0,0 +1,21 @@
# This package contains options shared by all MediaPipe Tasks for Web.
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_ts_library")
package(default_visibility = ["//mediapipe/tasks:internal"])
mediapipe_ts_library(
name = "core",
srcs = [
"base_options.d.ts",
"wasm_loader_options.d.ts",
],
)
mediapipe_ts_library(
name = "classifier_options",
srcs = [
"classifier_options.d.ts",
],
deps = [":core"],
)
mediapipe/tasks/web/core/base_options.d.ts
Vendored
+31
View File
@@ -0,0 +1,31 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* 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.
*/
// Placeholder for internal dependency on trusted resource url
/** Options to configure MediaPipe Tasks in general. */
export interface BaseOptions {
/**
* The model path to the model asset file. Only one of `modelAssetPath` or
* `modelAssetBuffer` can be set.
*/
modelAssetPath?: string;
/**
* A buffer containing the model aaset. Only one of `modelAssetPath` or
* `modelAssetBuffer` can be set.
*/
modelAssetBuffer?: Uint8Array;
}
mediapipe/tasks/web/core/classifier_options.d.ts
Vendored
+52
View File
@@ -0,0 +1,52 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* 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.
*/
import {BaseOptions} from '../../../tasks/web/core/base_options';
/** Options to configure the Mediapipe Classifier Task. */
export interface ClassifierOptions {
/** Options to configure the loading of the model assets. */
baseOptions?: BaseOptions;
/**
* The locale to use for display names specified through the TFLite Model
* Metadata, if any. Defaults to English.
*/
displayNamesLocale?: string|undefined;
/** The maximum number of top-scored detection results to return. */
maxResults?: number|undefined;
/**
* Overrides the value provided in the model metadata. Results below this
* value are rejected.
*/
scoreThreshold?: number|undefined;
/**
* Allowlist of category names. If non-empty, detection results whose category
* name is not in this set will be filtered out. Duplicate or unknown category
* names are ignored. Mutually exclusive with `categoryDenylist`.
*/
categoryAllowlist?: string[]|undefined;
/**
* Denylist of category names. If non-empty, detection results whose category
* name is in this set will be filtered out. Duplicate or unknown category
* names are ignored. Mutually exclusive with `categoryAllowlist`.
*/
categoryDenylist?: string[]|undefined;
}
mediapipe/tasks/web/core/wasm_loader_options.d.ts
Vendored
+25
View File
@@ -0,0 +1,25 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* 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.
*/
// Placeholder for internal dependency on trusted resource url
/** An object containing the locations of all Wasm assets */
export interface WasmLoaderOptions {
/** The path to the Wasm loader script. */
wasmLoaderPath: string;
/** The path to the Wasm binary. */
wasmBinaryPath: string;
}