Merge branch 'master' into interactive-segmenter-python
This commit is contained in:
@@ -111,7 +111,7 @@ class ClassificationAggregationCalculator : public Node {
|
||||
private:
|
||||
std::vector<std::string> head_names_;
|
||||
bool time_aggregation_enabled_;
|
||||
std::unordered_map<int64, std::vector<ClassificationList>>
|
||||
std::unordered_map<int64_t, std::vector<ClassificationList>>
|
||||
cached_classifications_;
|
||||
|
||||
ClassificationResult ConvertToClassificationResult(CalculatorContext* cc);
|
||||
|
||||
@@ -83,7 +83,7 @@ class EmbeddingAggregationCalculator : public Node {
|
||||
|
||||
private:
|
||||
bool time_aggregation_enabled_;
|
||||
std::unordered_map<int64, EmbeddingResult> cached_embeddings_;
|
||||
std::unordered_map<int64_t, EmbeddingResult> cached_embeddings_;
|
||||
};
|
||||
|
||||
absl::Status EmbeddingAggregationCalculator::UpdateContract(
|
||||
|
||||
@@ -107,7 +107,7 @@ absl::Status ConfigureImageToTensorCalculator(
|
||||
options->mutable_output_tensor_float_range()->set_max((255.0f - mean) /
|
||||
std);
|
||||
}
|
||||
// TODO: need to support different GPU origin on differnt
|
||||
// TODO: need to support different GPU origin on different
|
||||
// platforms or applications.
|
||||
options->set_gpu_origin(mediapipe::GpuOrigin::TOP_LEFT);
|
||||
return absl::OkStatus();
|
||||
|
||||
@@ -30,4 +30,8 @@ message TextPreprocessingGraphOptions {
|
||||
// The maximum input sequence length for the TFLite model. Used with
|
||||
// BERT_MODEL and REGEX_MODEL.
|
||||
optional int32 max_seq_len = 2;
|
||||
|
||||
// The model's input tensors are dynamic rather than static.
|
||||
// Used with BERT_MODEL.
|
||||
optional bool has_dynamic_input_tensors = 3;
|
||||
}
|
||||
|
||||
@@ -114,6 +114,60 @@ absl::StatusOr<int> GetMaxSeqLen(const tflite::SubGraph& model_graph) {
|
||||
}
|
||||
return max_seq_len;
|
||||
}
|
||||
|
||||
// Determines whether the TFLite model for `model_graph` has input tensors with
|
||||
// dynamic shape rather than static shape or returns an error if the input
|
||||
// tensors have invalid shape signatures. This util assumes that the model has
|
||||
// the correct input tensors type and count for the BertPreprocessorCalculator.
|
||||
absl::StatusOr<bool> HasDynamicInputTensors(
|
||||
const tflite::SubGraph& model_graph) {
|
||||
const flatbuffers::Vector<int32_t>& input_indices = *model_graph.inputs();
|
||||
const flatbuffers::Vector<flatbuffers::Offset<tflite::Tensor>>&
|
||||
model_tensors = *model_graph.tensors();
|
||||
|
||||
// Static input tensors may have undefined shape signatures.
|
||||
if (absl::c_all_of(input_indices, [&model_tensors](int i) {
|
||||
return model_tensors[i]->shape_signature() == nullptr;
|
||||
})) {
|
||||
return false;
|
||||
} else if (absl::c_any_of(input_indices, [&model_tensors](int i) {
|
||||
return model_tensors[i]->shape_signature() == nullptr;
|
||||
})) {
|
||||
return CreateStatusWithPayload(absl::StatusCode::kInvalidArgument,
|
||||
"Input tensors contain a mix of defined and "
|
||||
"undefined shape signatures.");
|
||||
}
|
||||
|
||||
for (int i : input_indices) {
|
||||
const tflite::Tensor* tensor = model_tensors[i];
|
||||
if (tensor->shape_signature()->size() != 2) {
|
||||
return CreateStatusWithPayload(
|
||||
absl::StatusCode::kInvalidArgument,
|
||||
absl::Substitute(
|
||||
"Model should take 2-D shape signatures, got dimension: $0",
|
||||
tensor->shape_signature()->size()),
|
||||
MediaPipeTasksStatus::kInvalidInputTensorDimensionsError);
|
||||
}
|
||||
}
|
||||
|
||||
// For dynamic input tensors, the shape_signature entry corresponding to the
|
||||
// input size is -1.
|
||||
if (absl::c_all_of(input_indices, [&model_tensors](int i) {
|
||||
return (*model_tensors[i]->shape_signature())[1] != -1;
|
||||
})) {
|
||||
return false;
|
||||
} else if (absl::c_all_of(input_indices, [&model_tensors](int i) {
|
||||
return (*model_tensors[i]->shape_signature())[1] == -1;
|
||||
})) {
|
||||
return true;
|
||||
} else {
|
||||
return CreateStatusWithPayload(
|
||||
absl::StatusCode::kInvalidArgument,
|
||||
"Input tensors contain a mix of static and dynamic shapes.");
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
absl::Status ConfigureTextPreprocessingGraph(
|
||||
@@ -128,6 +182,8 @@ absl::Status ConfigureTextPreprocessingGraph(
|
||||
|
||||
ASSIGN_OR_RETURN(TextModelType::ModelType model_type,
|
||||
GetModelType(model_resources));
|
||||
const tflite::SubGraph& model_graph =
|
||||
*(*model_resources.GetTfLiteModel()->subgraphs())[0];
|
||||
options.set_model_type(model_type);
|
||||
switch (model_type) {
|
||||
case TextModelType::UNSPECIFIED_MODEL:
|
||||
@@ -137,13 +193,15 @@ absl::Status ConfigureTextPreprocessingGraph(
|
||||
}
|
||||
case TextModelType::BERT_MODEL:
|
||||
case TextModelType::REGEX_MODEL: {
|
||||
ASSIGN_OR_RETURN(
|
||||
int max_seq_len,
|
||||
GetMaxSeqLen(*(*model_resources.GetTfLiteModel()->subgraphs())[0]));
|
||||
ASSIGN_OR_RETURN(int max_seq_len, GetMaxSeqLen(model_graph));
|
||||
options.set_max_seq_len(max_seq_len);
|
||||
}
|
||||
}
|
||||
|
||||
if (model_type == TextModelType::BERT_MODEL) {
|
||||
ASSIGN_OR_RETURN(bool has_dynamic_input_tensors,
|
||||
HasDynamicInputTensors(model_graph));
|
||||
options.set_has_dynamic_input_tensors(has_dynamic_input_tensors);
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
@@ -200,6 +258,8 @@ class TextPreprocessingGraph : public mediapipe::Subgraph {
|
||||
case TextModelType::BERT_MODEL: {
|
||||
text_preprocessor.GetOptions<BertPreprocessorCalculatorOptions>()
|
||||
.set_bert_max_seq_len(options.max_seq_len());
|
||||
text_preprocessor.GetOptions<BertPreprocessorCalculatorOptions>()
|
||||
.set_has_dynamic_input_tensors(options.has_dynamic_input_tensors());
|
||||
metadata_extractor_in >>
|
||||
text_preprocessor.SideIn(kMetadataExtractorTag);
|
||||
break;
|
||||
|
||||
@@ -68,8 +68,8 @@ TEST(DisallowGate, VerifyConfig) {
|
||||
input_stream: "VALUE_3:__stream_3"
|
||||
)pb")));
|
||||
|
||||
CalculatorGraph calcualtor_graph;
|
||||
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||
CalculatorGraph calculator_graph;
|
||||
MP_EXPECT_OK(calculator_graph.Initialize(graph.GetConfig()));
|
||||
}
|
||||
|
||||
TEST(DisallowIf, VerifyConfig) {
|
||||
@@ -99,8 +99,8 @@ TEST(DisallowIf, VerifyConfig) {
|
||||
input_stream: "VALUE:__stream_1"
|
||||
)pb")));
|
||||
|
||||
CalculatorGraph calcualtor_graph;
|
||||
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||
CalculatorGraph calculator_graph;
|
||||
MP_EXPECT_OK(calculator_graph.Initialize(graph.GetConfig()));
|
||||
}
|
||||
|
||||
TEST(DisallowIf, VerifyConfigWithSideCondition) {
|
||||
@@ -130,8 +130,8 @@ TEST(DisallowIf, VerifyConfigWithSideCondition) {
|
||||
input_side_packet: "CONDITION:__side_packet_1"
|
||||
)pb")));
|
||||
|
||||
CalculatorGraph calcualtor_graph;
|
||||
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||
CalculatorGraph calculator_graph;
|
||||
MP_EXPECT_OK(calculator_graph.Initialize(graph.GetConfig()));
|
||||
}
|
||||
|
||||
TEST(AllowGate, VerifyConfig) {
|
||||
@@ -166,8 +166,8 @@ TEST(AllowGate, VerifyConfig) {
|
||||
input_stream: "VALUE_3:__stream_3"
|
||||
)pb")));
|
||||
|
||||
CalculatorGraph calcualtor_graph;
|
||||
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||
CalculatorGraph calculator_graph;
|
||||
MP_EXPECT_OK(calculator_graph.Initialize(graph.GetConfig()));
|
||||
}
|
||||
|
||||
TEST(AllowIf, VerifyConfig) {
|
||||
@@ -192,8 +192,8 @@ TEST(AllowIf, VerifyConfig) {
|
||||
input_stream: "VALUE:__stream_1"
|
||||
)pb")));
|
||||
|
||||
CalculatorGraph calcualtor_graph;
|
||||
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||
CalculatorGraph calculator_graph;
|
||||
MP_EXPECT_OK(calculator_graph.Initialize(graph.GetConfig()));
|
||||
}
|
||||
|
||||
TEST(AllowIf, VerifyConfigWithSideConition) {
|
||||
@@ -218,8 +218,8 @@ TEST(AllowIf, VerifyConfigWithSideConition) {
|
||||
input_side_packet: "CONDITION:__side_packet_1"
|
||||
)pb")));
|
||||
|
||||
CalculatorGraph calcualtor_graph;
|
||||
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||
CalculatorGraph calculator_graph;
|
||||
MP_EXPECT_OK(calculator_graph.Initialize(graph.GetConfig()));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
@@ -55,15 +55,15 @@ cc_library(
|
||||
srcs = ["external_file_handler.cc"],
|
||||
hdrs = ["external_file_handler.h"],
|
||||
deps = [
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/tasks/cc:common",
|
||||
"//mediapipe/tasks/cc/core/proto:external_file_cc_proto",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/status",
|
||||
"@com_google_absl//absl/status:statusor",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/strings:str_format",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/tasks/cc:common",
|
||||
"//mediapipe/tasks/cc/core/proto:external_file_cc_proto",
|
||||
] + select({
|
||||
"//mediapipe:windows": ["@bazel_tools//tools/cpp/runfiles"],
|
||||
"//conditions:default": [],
|
||||
|
||||
@@ -100,10 +100,11 @@ absl::StatusOr<std::string> PathToResourceAsFile(std::string path) {
|
||||
#ifndef _WIN32
|
||||
return path;
|
||||
#else
|
||||
if (absl::StartsWith(path, "./")) {
|
||||
path = "mediapipe" + path.substr(1);
|
||||
std::string qualified_path = path;
|
||||
if (absl::StartsWith(qualified_path, "./")) {
|
||||
qualified_path = "mediapipe" + qualified_path.substr(1);
|
||||
} else if (path[0] != '/') {
|
||||
path = "mediapipe/" + path;
|
||||
qualified_path = "mediapipe/" + qualified_path;
|
||||
}
|
||||
|
||||
std::string error;
|
||||
@@ -112,9 +113,10 @@ absl::StatusOr<std::string> PathToResourceAsFile(std::string path) {
|
||||
std::unique_ptr<::bazel::tools::cpp::runfiles::Runfiles> runfiles(
|
||||
::bazel::tools::cpp::runfiles::Runfiles::Create("", &error));
|
||||
if (!runfiles) {
|
||||
return absl::InternalError("Unable to initialize runfiles: " + error);
|
||||
// Return the original path when Runfiles is not available (e.g. for Python)
|
||||
return path;
|
||||
}
|
||||
return runfiles->Rlocation(path);
|
||||
return runfiles->Rlocation(qualified_path);
|
||||
#endif // _WIN32
|
||||
}
|
||||
|
||||
|
||||
@@ -46,9 +46,9 @@ MediaPipeBuiltinOpResolver::MediaPipeBuiltinOpResolver() {
|
||||
mediapipe::tflite_operations::RegisterLandmarksToTransformMatrixV2(),
|
||||
/*version=*/2);
|
||||
// For the LanguageDetector model.
|
||||
AddCustom("NGramHash", ::tflite::ops::custom::Register_NGRAM_HASH());
|
||||
AddCustom("NGramHash", mediapipe::tflite_operations::Register_NGRAM_HASH());
|
||||
AddCustom("KmeansEmbeddingLookup",
|
||||
::tflite::ops::custom::Register_KmeansEmbeddingLookup());
|
||||
mediapipe::tflite_operations::Register_KmeansEmbeddingLookup());
|
||||
}
|
||||
} // namespace core
|
||||
} // namespace tasks
|
||||
|
||||
@@ -23,7 +23,7 @@ limitations under the License.
|
||||
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
|
||||
#include "tensorflow/lite/kernels/kernel_util.h"
|
||||
|
||||
namespace tflite::ops::custom {
|
||||
namespace mediapipe::tflite_operations {
|
||||
namespace kmeans_embedding_lookup_op {
|
||||
|
||||
namespace {
|
||||
@@ -33,6 +33,10 @@ constexpr int kEncodingTable = 1;
|
||||
constexpr int kCodebook = 2;
|
||||
constexpr int kOutputLabel = 0;
|
||||
|
||||
using ::tflite::GetInput;
|
||||
using ::tflite::GetOutput;
|
||||
using ::tflite::GetTensorData;
|
||||
|
||||
} // namespace
|
||||
|
||||
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
||||
@@ -142,4 +146,4 @@ TfLiteRegistration* Register_KmeansEmbeddingLookup() {
|
||||
return &r;
|
||||
}
|
||||
|
||||
} // namespace tflite::ops::custom
|
||||
} // namespace mediapipe::tflite_operations
|
||||
|
||||
@@ -27,10 +27,10 @@ limitations under the License.
|
||||
|
||||
#include "tensorflow/lite/kernels/register.h"
|
||||
|
||||
namespace tflite::ops::custom {
|
||||
namespace mediapipe::tflite_operations {
|
||||
|
||||
TfLiteRegistration* Register_KmeansEmbeddingLookup();
|
||||
|
||||
} // namespace tflite::ops::custom
|
||||
} // namespace mediapipe::tflite_operations
|
||||
|
||||
#endif // MEDIAPIPE_TASKS_CC_TEXT_LANGUAGE_DETECTOR_CUSTOM_OPS_KMEANS_EMBEDDING_LOOKUP_H_
|
||||
|
||||
+7
-7
@@ -12,14 +12,14 @@
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
#include "tensorflow/lite/kernels/test_util.h"
|
||||
|
||||
namespace tflite::ops::custom {
|
||||
namespace mediapipe::tflite_operations {
|
||||
namespace {
|
||||
|
||||
using ::testing::ElementsAreArray;
|
||||
using ::tflite::ArrayFloatNear;
|
||||
|
||||
// Helper class for testing the op.
|
||||
class KmeansEmbeddingLookupModel : public SingleOpModel {
|
||||
class KmeansEmbeddingLookupModel : public tflite::SingleOpModel {
|
||||
public:
|
||||
explicit KmeansEmbeddingLookupModel(
|
||||
std::initializer_list<int> input_shape,
|
||||
@@ -27,7 +27,7 @@ class KmeansEmbeddingLookupModel : public SingleOpModel {
|
||||
std::initializer_list<int> codebook_shape,
|
||||
std::initializer_list<int> output_shape) {
|
||||
// Setup the model inputs and the interpreter.
|
||||
output_ = AddOutput({TensorType_FLOAT32, output_shape});
|
||||
output_ = AddOutput({tflite::TensorType_FLOAT32, output_shape});
|
||||
SetCustomOp("KmeansEmbeddingLookup", std::vector<uint8_t>(),
|
||||
Register_KmeansEmbeddingLookup);
|
||||
BuildInterpreter({input_shape, encoding_table_shape, codebook_shape});
|
||||
@@ -68,9 +68,9 @@ class KmeansEmbeddingLookupModel : public SingleOpModel {
|
||||
std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
|
||||
|
||||
private:
|
||||
int input_ = AddInput(TensorType_INT32);
|
||||
int encoding_table_ = AddInput(TensorType_UINT8);
|
||||
int codebook_ = AddInput(TensorType_FLOAT32);
|
||||
int input_ = AddInput(tflite::TensorType_INT32);
|
||||
int encoding_table_ = AddInput(tflite::TensorType_UINT8);
|
||||
int codebook_ = AddInput(tflite::TensorType_FLOAT32);
|
||||
int output_;
|
||||
};
|
||||
|
||||
@@ -173,4 +173,4 @@ TEST(KmeansEmbeddingLookupTest, ThrowsErrorWhenGivenInvalidInputBatchSize) {
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace tflite::ops::custom
|
||||
} // namespace mediapipe::tflite_operations
|
||||
|
||||
@@ -25,7 +25,7 @@ limitations under the License.
|
||||
#include "tensorflow/lite/kernels/kernel_util.h"
|
||||
#include "tensorflow/lite/string_util.h"
|
||||
|
||||
namespace tflite::ops::custom {
|
||||
namespace mediapipe::tflite_operations {
|
||||
|
||||
namespace ngram_op {
|
||||
|
||||
@@ -217,21 +217,21 @@ void Free(TfLiteContext* context, void* buffer) {
|
||||
}
|
||||
|
||||
TfLiteStatus Resize(TfLiteContext* context, TfLiteNode* node) {
|
||||
TfLiteTensor* output = GetOutput(context, node, kOutputLabel);
|
||||
TfLiteTensor* output = tflite::GetOutput(context, node, kOutputLabel);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
SetTensorToDynamic(output);
|
||||
tflite::SetTensorToDynamic(output);
|
||||
return kTfLiteOk;
|
||||
}
|
||||
|
||||
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
||||
NGramHashParams* params = reinterpret_cast<NGramHashParams*>(node->user_data);
|
||||
TF_LITE_ENSURE_OK(
|
||||
context,
|
||||
params->PreprocessInput(GetInput(context, node, kInputMessage), context));
|
||||
context, params->PreprocessInput(
|
||||
tflite::GetInput(context, node, kInputMessage), context));
|
||||
|
||||
TfLiteTensor* output = GetOutput(context, node, kOutputLabel);
|
||||
TfLiteTensor* output = tflite::GetOutput(context, node, kOutputLabel);
|
||||
TF_LITE_ENSURE(context, output != nullptr);
|
||||
if (IsDynamicTensor(output)) {
|
||||
if (tflite::IsDynamicTensor(output)) {
|
||||
TfLiteIntArray* output_size = TfLiteIntArrayCreate(3);
|
||||
output_size->data[0] = 1;
|
||||
output_size->data[1] = params->GetNumNGrams();
|
||||
@@ -261,4 +261,4 @@ TfLiteRegistration* Register_NGRAM_HASH() {
|
||||
return &r;
|
||||
}
|
||||
|
||||
} // namespace tflite::ops::custom
|
||||
} // namespace mediapipe::tflite_operations
|
||||
|
||||
@@ -18,10 +18,10 @@ limitations under the License.
|
||||
|
||||
#include "tensorflow/lite/kernels/register.h"
|
||||
|
||||
namespace tflite::ops::custom {
|
||||
namespace mediapipe::tflite_operations {
|
||||
|
||||
TfLiteRegistration* Register_NGRAM_HASH();
|
||||
|
||||
} // namespace tflite::ops::custom
|
||||
} // namespace mediapipe::tflite_operations
|
||||
|
||||
#endif // MEDIAPIPE_TASKS_CC_TEXT_LANGUAGE_DETECTOR_CUSTOM_OPS_NGRAM_HASH_H_
|
||||
|
||||
@@ -32,7 +32,7 @@ limitations under the License.
|
||||
#include "tensorflow/lite/model.h"
|
||||
#include "tensorflow/lite/string_util.h"
|
||||
|
||||
namespace tflite::ops::custom {
|
||||
namespace mediapipe::tflite_operations {
|
||||
namespace {
|
||||
|
||||
using ::flexbuffers::Builder;
|
||||
@@ -42,7 +42,7 @@ using ::testing::ElementsAreArray;
|
||||
using ::testing::Message;
|
||||
|
||||
// Helper class for testing the op.
|
||||
class NGramHashModel : public SingleOpModel {
|
||||
class NGramHashModel : public tflite::SingleOpModel {
|
||||
public:
|
||||
explicit NGramHashModel(const uint64_t seed,
|
||||
const std::vector<int>& ngram_lengths,
|
||||
@@ -71,7 +71,7 @@ class NGramHashModel : public SingleOpModel {
|
||||
}
|
||||
fbb.EndMap(start);
|
||||
fbb.Finish();
|
||||
output_ = AddOutput({TensorType_INT32, {}});
|
||||
output_ = AddOutput({tflite::TensorType_INT32, {}});
|
||||
SetCustomOp("NGramHash", fbb.GetBuffer(), Register_NGRAM_HASH);
|
||||
BuildInterpreter({GetShape(input_)});
|
||||
}
|
||||
@@ -100,7 +100,7 @@ class NGramHashModel : public SingleOpModel {
|
||||
std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
|
||||
|
||||
private:
|
||||
int input_ = AddInput(TensorType_STRING);
|
||||
int input_ = AddInput(tflite::TensorType_STRING);
|
||||
int output_;
|
||||
};
|
||||
|
||||
@@ -173,7 +173,7 @@ TEST(NGramHashTest, ReturnsExpectedValueWhenInputIsSane) {
|
||||
|
||||
NGramHashModel m(kSeed, ngram_lengths, vocab_sizes);
|
||||
for (int test_idx = 0; test_idx < testcase_inputs.size(); test_idx++) {
|
||||
const string& testcase_input = testcase_inputs[test_idx];
|
||||
const std::string& testcase_input = testcase_inputs[test_idx];
|
||||
m.Invoke(testcase_input);
|
||||
SCOPED_TRACE(Message() << "Where the testcases' input is: "
|
||||
<< testcase_input);
|
||||
@@ -310,4 +310,4 @@ TEST(NGramHashTest, MismatchNgramLengthsAndVocabSizes) {
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace tflite::ops::custom
|
||||
} // namespace mediapipe::tflite_operations
|
||||
|
||||
@@ -49,6 +49,7 @@ using ::testing::HasSubstr;
|
||||
using ::testing::Optional;
|
||||
|
||||
constexpr int kMaxSeqLen = 128;
|
||||
const float kPrecision = 1e-6;
|
||||
constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/text/";
|
||||
constexpr char kTestBertModelPath[] = "bert_text_classifier.tflite";
|
||||
constexpr char kInvalidModelPath[] = "i/do/not/exist.tflite";
|
||||
@@ -66,7 +67,6 @@ std::string GetFullPath(absl::string_view file_name) {
|
||||
// TODO: create shared matcher for ClassificationResult.
|
||||
void ExpectApproximatelyEqual(const TextClassifierResult& actual,
|
||||
const TextClassifierResult& expected) {
|
||||
const float kPrecision = 1e-6;
|
||||
ASSERT_EQ(actual.classifications.size(), expected.classifications.size());
|
||||
for (int i = 0; i < actual.classifications.size(); ++i) {
|
||||
const Classifications& a = actual.classifications[i];
|
||||
|
||||
@@ -39,6 +39,7 @@ cc_library(
|
||||
":running_mode",
|
||||
"//mediapipe/calculators/core:flow_limiter_calculator",
|
||||
"//mediapipe/calculators/tensor:image_to_tensor_calculator_cc_proto",
|
||||
"//mediapipe/framework/formats:image",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/tasks/cc/components/containers:rect",
|
||||
"//mediapipe/tasks/cc/core:base_task_api",
|
||||
|
||||
@@ -17,6 +17,7 @@ limitations under the License.
|
||||
#define MEDIAPIPE_TASKS_CC_VISION_CORE_BASE_VISION_TASK_API_H_
|
||||
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <string>
|
||||
@@ -26,6 +27,7 @@ limitations under the License.
|
||||
#include "absl/status/statusor.h"
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "mediapipe/calculators/tensor/image_to_tensor_calculator.pb.h"
|
||||
#include "mediapipe/framework/formats/image.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/tasks/cc/components/containers/rect.h"
|
||||
#include "mediapipe/tasks/cc/core/base_task_api.h"
|
||||
@@ -136,7 +138,8 @@ class BaseVisionTaskApi : public tasks::core::BaseTaskApi {
|
||||
// to 0. If 'roi_allowed' is false, an error will be returned if the input
|
||||
// ImageProcessingOptions has its 'region_or_interest' field set.
|
||||
static absl::StatusOr<mediapipe::NormalizedRect> ConvertToNormalizedRect(
|
||||
std::optional<ImageProcessingOptions> options, bool roi_allowed = true) {
|
||||
std::optional<ImageProcessingOptions> options,
|
||||
const mediapipe::Image& image, bool roi_allowed = true) {
|
||||
mediapipe::NormalizedRect normalized_rect;
|
||||
normalized_rect.set_rotation(0);
|
||||
normalized_rect.set_x_center(0.5);
|
||||
@@ -181,6 +184,21 @@ class BaseVisionTaskApi : public tasks::core::BaseTaskApi {
|
||||
normalized_rect.set_width(roi.right - roi.left);
|
||||
normalized_rect.set_height(roi.bottom - roi.top);
|
||||
}
|
||||
|
||||
// For 90° and 270° rotations, we need to swap width and height.
|
||||
// This is due to the internal behavior of ImageToTensorCalculator, which:
|
||||
// - first denormalizes the provided rect by multiplying the rect width or
|
||||
// height by the image width or height, repectively.
|
||||
// - then rotates this by denormalized rect by the provided rotation, and
|
||||
// uses this for cropping,
|
||||
// - then finally rotates this back.
|
||||
if (std::abs(options->rotation_degrees) % 180 != 0) {
|
||||
float w = normalized_rect.height() * image.height() / image.width();
|
||||
float h = normalized_rect.width() * image.width() / image.height();
|
||||
normalized_rect.set_width(w);
|
||||
normalized_rect.set_height(h);
|
||||
}
|
||||
|
||||
return normalized_rect;
|
||||
}
|
||||
|
||||
|
||||
@@ -136,9 +136,9 @@ absl::StatusOr<std::unique_ptr<FaceDetector>> FaceDetector::Create(
|
||||
absl::StatusOr<FaceDetectorResult> FaceDetector::Detect(
|
||||
mediapipe::Image image,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options) {
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessImageData(
|
||||
@@ -156,9 +156,9 @@ absl::StatusOr<FaceDetectorResult> FaceDetector::Detect(
|
||||
absl::StatusOr<FaceDetectorResult> FaceDetector::DetectForVideo(
|
||||
mediapipe::Image image, uint64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options) {
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessVideoData(
|
||||
@@ -179,9 +179,9 @@ absl::StatusOr<FaceDetectorResult> FaceDetector::DetectForVideo(
|
||||
absl::Status FaceDetector::DetectAsync(
|
||||
mediapipe::Image image, uint64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options) {
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
return SendLiveStreamData(
|
||||
{{kImageInStreamName,
|
||||
MakePacket<Image>(std::move(image))
|
||||
|
||||
@@ -19,6 +19,9 @@ package mediapipe.tasks.vision.face_geometry;
|
||||
import "mediapipe/framework/calculator_options.proto";
|
||||
import "mediapipe/tasks/cc/core/proto/external_file.proto";
|
||||
|
||||
option java_package = "com.google.mediapipe.tasks.vision.facegeometry.calculators.proto";
|
||||
option java_outer_classname = "FaceGeometryPipelineCalculatorOptionsProto";
|
||||
|
||||
message FaceGeometryPipelineCalculatorOptions {
|
||||
extend mediapipe.CalculatorOptions {
|
||||
optional FaceGeometryPipelineCalculatorOptions ext = 512499200;
|
||||
|
||||
@@ -19,6 +19,9 @@ package mediapipe.tasks.vision.face_geometry.proto;
|
||||
import "mediapipe/framework/calculator_options.proto";
|
||||
import "mediapipe/tasks/cc/vision/face_geometry/calculators/geometry_pipeline_calculator.proto";
|
||||
|
||||
option java_package = "com.google.mediapipe.tasks.vision.facegeometry.proto";
|
||||
option java_outer_classname = "FaceGeometryGraphOptionsProto";
|
||||
|
||||
message FaceGeometryGraphOptions {
|
||||
extend mediapipe.CalculatorOptions {
|
||||
optional FaceGeometryGraphOptions ext = 515723506;
|
||||
|
||||
@@ -17,7 +17,7 @@ syntax = "proto2";
|
||||
package mediapipe.tasks.vision.face_geometry.proto;
|
||||
|
||||
option java_package = "com.google.mediapipe.tasks.vision.facegeometry.proto";
|
||||
option java_outer_classname = "Mesh3dProto";
|
||||
option java_outer_classname = "Mesh3DProto";
|
||||
|
||||
message Mesh3d {
|
||||
enum VertexType {
|
||||
|
||||
@@ -181,9 +181,12 @@ cc_library(
|
||||
":face_landmarks_detector_graph",
|
||||
"//mediapipe/calculators/core:begin_loop_calculator",
|
||||
"//mediapipe/calculators/core:clip_vector_size_calculator_cc_proto",
|
||||
"//mediapipe/calculators/core:concatenate_vector_calculator",
|
||||
"//mediapipe/calculators/core:end_loop_calculator",
|
||||
"//mediapipe/calculators/core:gate_calculator",
|
||||
"//mediapipe/calculators/core:gate_calculator_cc_proto",
|
||||
"//mediapipe/calculators/core:get_vector_item_calculator",
|
||||
"//mediapipe/calculators/core:get_vector_item_calculator_cc_proto",
|
||||
"//mediapipe/calculators/core:pass_through_calculator",
|
||||
"//mediapipe/calculators/core:previous_loopback_calculator",
|
||||
"//mediapipe/calculators/image:image_properties_calculator",
|
||||
@@ -191,6 +194,8 @@ cc_library(
|
||||
"//mediapipe/calculators/util:association_norm_rect_calculator",
|
||||
"//mediapipe/calculators/util:collection_has_min_size_calculator",
|
||||
"//mediapipe/calculators/util:collection_has_min_size_calculator_cc_proto",
|
||||
"//mediapipe/calculators/util:landmarks_smoothing_calculator",
|
||||
"//mediapipe/calculators/util:landmarks_smoothing_calculator_cc_proto",
|
||||
"//mediapipe/framework/api2:builder",
|
||||
"//mediapipe/framework/api2:port",
|
||||
"//mediapipe/framework/formats:classification_cc_proto",
|
||||
|
||||
@@ -194,7 +194,7 @@ absl::StatusOr<FaceLandmarkerResult> FaceLandmarker::Detect(
|
||||
mediapipe::Image image,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options) {
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
@@ -212,7 +212,7 @@ absl::StatusOr<FaceLandmarkerResult> FaceLandmarker::DetectForVideo(
|
||||
mediapipe::Image image, int64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options) {
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
@@ -233,7 +233,7 @@ absl::Status FaceLandmarker::DetectAsync(
|
||||
mediapipe::Image image, int64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options) {
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
return SendLiveStreamData(
|
||||
{{kImageInStreamName,
|
||||
|
||||
@@ -20,9 +20,13 @@ limitations under the License.
|
||||
|
||||
#include "absl/strings/str_format.h"
|
||||
#include "mediapipe/calculators/core/clip_vector_size_calculator.pb.h"
|
||||
#include "mediapipe/calculators/core/concatenate_vector_calculator.h"
|
||||
#include "mediapipe/calculators/core/gate_calculator.pb.h"
|
||||
#include "mediapipe/calculators/core/get_vector_item_calculator.h"
|
||||
#include "mediapipe/calculators/core/get_vector_item_calculator.pb.h"
|
||||
#include "mediapipe/calculators/util/association_calculator.pb.h"
|
||||
#include "mediapipe/calculators/util/collection_has_min_size_calculator.pb.h"
|
||||
#include "mediapipe/calculators/util/landmarks_smoothing_calculator.pb.h"
|
||||
#include "mediapipe/framework/api2/builder.h"
|
||||
#include "mediapipe/framework/api2/port.h"
|
||||
#include "mediapipe/framework/formats/classification.pb.h"
|
||||
@@ -91,6 +95,9 @@ constexpr char kEnvironmentTag[] = "ENVIRONMENT";
|
||||
constexpr char kBlendshapesTag[] = "BLENDSHAPES";
|
||||
constexpr char kImageSizeTag[] = "IMAGE_SIZE";
|
||||
constexpr char kSizeTag[] = "SIZE";
|
||||
constexpr char kVectorTag[] = "VECTOR";
|
||||
constexpr char kItemTag[] = "ITEM";
|
||||
constexpr char kNormFilteredLandmarksTag[] = "NORM_FILTERED_LANDMARKS";
|
||||
constexpr char kFaceDetectorTFLiteName[] = "face_detector.tflite";
|
||||
constexpr char kFaceLandmarksDetectorTFLiteName[] =
|
||||
"face_landmarks_detector.tflite";
|
||||
@@ -166,6 +173,18 @@ absl::Status SetSubTaskBaseOptions(const ModelAssetBundleResources& resources,
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
void ConfigureLandmarksSmoothingCalculator(
|
||||
mediapipe::LandmarksSmoothingCalculatorOptions& options) {
|
||||
// Min cutoff 0.05 results into ~0.01 alpha in landmark EMA filter when
|
||||
// landmark is static.
|
||||
options.mutable_one_euro_filter()->set_min_cutoff(0.05f);
|
||||
// Beta 80.0 in combintation with min_cutoff 0.05 results into ~0.94
|
||||
// alpha in landmark EMA filter when landmark is moving fast.
|
||||
options.mutable_one_euro_filter()->set_beta(80.0f);
|
||||
// Derivative cutoff 1.0 results into ~0.17 alpha in landmark velocity
|
||||
// EMA filter.
|
||||
options.mutable_one_euro_filter()->set_derivate_cutoff(1.0f);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
// A "mediapipe.tasks.vision.face_landmarker.FaceLandmarkerGraph" performs face
|
||||
@@ -428,13 +447,46 @@ class FaceLandmarkerGraph : public core::ModelTaskGraph {
|
||||
image_in >> face_landmarks_detector_graph.In(kImageTag);
|
||||
clipped_face_rects >> face_landmarks_detector_graph.In(kNormRectTag);
|
||||
|
||||
// TODO: add landmarks smoothing calculators.
|
||||
auto landmarks = face_landmarks_detector_graph.Out(kNormLandmarksTag)
|
||||
std::optional<Source<std::vector<NormalizedLandmarkList>>> face_landmarks;
|
||||
face_landmarks = face_landmarks_detector_graph.Out(kNormLandmarksTag)
|
||||
.Cast<std::vector<NormalizedLandmarkList>>();
|
||||
auto face_rects_for_next_frame =
|
||||
face_landmarks_detector_graph.Out(kFaceRectsNextFrameTag)
|
||||
.Cast<std::vector<NormalizedRect>>();
|
||||
|
||||
auto& image_properties = graph.AddNode("ImagePropertiesCalculator");
|
||||
image_in >> image_properties.In(kImageTag);
|
||||
auto image_size = image_properties.Out(kSizeTag);
|
||||
|
||||
// Apply smoothing filter only on the single face landmarks, because
|
||||
// landmakrs smoothing calculator doesn't support multiple landmarks yet.
|
||||
if (tasks_options.face_detector_graph_options().num_faces() == 1) {
|
||||
// Get the single face landmarks
|
||||
auto& get_vector_item =
|
||||
graph.AddNode("GetNormalizedLandmarkListVectorItemCalculator");
|
||||
get_vector_item.GetOptions<mediapipe::GetVectorItemCalculatorOptions>()
|
||||
.set_item_index(0);
|
||||
*face_landmarks >> get_vector_item.In(kVectorTag);
|
||||
auto single_face_landmarks = get_vector_item.Out(kItemTag);
|
||||
|
||||
// Apply smoothing filter on face landmarks.
|
||||
auto& landmarks_smoothing = graph.AddNode("LandmarksSmoothingCalculator");
|
||||
ConfigureLandmarksSmoothingCalculator(
|
||||
landmarks_smoothing
|
||||
.GetOptions<mediapipe::LandmarksSmoothingCalculatorOptions>());
|
||||
single_face_landmarks >> landmarks_smoothing.In(kNormLandmarksTag);
|
||||
image_size >> landmarks_smoothing.In(kImageSizeTag);
|
||||
auto smoothed_single_face_landmarks =
|
||||
landmarks_smoothing.Out(kNormFilteredLandmarksTag);
|
||||
|
||||
// Wrap the single face landmarks into a vector of landmarks.
|
||||
auto& concatenate_vector =
|
||||
graph.AddNode("ConcatenateNormalizedLandmarkListVectorCalculator");
|
||||
smoothed_single_face_landmarks >> concatenate_vector.In("");
|
||||
face_landmarks.emplace(concatenate_vector.Out("")
|
||||
.Cast<std::vector<NormalizedLandmarkList>>());
|
||||
}
|
||||
|
||||
if (tasks_options.base_options().use_stream_mode()) {
|
||||
auto& previous_loopback = graph.AddNode("PreviousLoopbackCalculator");
|
||||
image_in >> previous_loopback.In(kMainTag);
|
||||
@@ -491,9 +543,6 @@ class FaceLandmarkerGraph : public core::ModelTaskGraph {
|
||||
// Optional face geometry output.
|
||||
std::optional<Source<std::vector<FaceGeometry>>> face_geometry;
|
||||
if (output_geometry) {
|
||||
auto& image_properties = graph.AddNode("ImagePropertiesCalculator");
|
||||
image_in >> image_properties.In(kImageTag);
|
||||
auto image_size = image_properties.Out(kSizeTag);
|
||||
auto& face_geometry_from_landmarks = graph.AddNode(
|
||||
"mediapipe.tasks.vision.face_geometry."
|
||||
"FaceGeometryFromLandmarksGraph");
|
||||
@@ -503,7 +552,7 @@ class FaceLandmarkerGraph : public core::ModelTaskGraph {
|
||||
if (environment.has_value()) {
|
||||
*environment >> face_geometry_from_landmarks.SideIn(kEnvironmentTag);
|
||||
}
|
||||
landmarks >> face_geometry_from_landmarks.In(kFaceLandmarksTag);
|
||||
*face_landmarks >> face_geometry_from_landmarks.In(kFaceLandmarksTag);
|
||||
image_size >> face_geometry_from_landmarks.In(kImageSizeTag);
|
||||
face_geometry = face_geometry_from_landmarks.Out(kFaceGeometryTag)
|
||||
.Cast<std::vector<FaceGeometry>>();
|
||||
@@ -515,7 +564,7 @@ class FaceLandmarkerGraph : public core::ModelTaskGraph {
|
||||
image_in >> pass_through.In("");
|
||||
|
||||
return {{
|
||||
/* landmark_lists= */ landmarks,
|
||||
/* landmark_lists= */ *face_landmarks,
|
||||
/* face_rects_next_frame= */
|
||||
face_rects_for_next_frame,
|
||||
/* face_rects= */
|
||||
|
||||
@@ -48,8 +48,6 @@ cc_library(
|
||||
}),
|
||||
deps = [
|
||||
":tensors_to_image_calculator_cc_proto",
|
||||
"@com_google_absl//absl/status",
|
||||
"@com_google_absl//absl/strings",
|
||||
"//mediapipe/calculators/tensor:image_to_tensor_utils",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_options_cc_proto",
|
||||
@@ -67,6 +65,8 @@ cc_library(
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:vector",
|
||||
"//mediapipe/gpu:gpu_origin_cc_proto",
|
||||
"@com_google_absl//absl/status",
|
||||
"@com_google_absl//absl/strings",
|
||||
] + select({
|
||||
"//mediapipe/gpu:disable_gpu": [],
|
||||
"//conditions:default": ["tensor_to_image_calculator_gpu_deps"],
|
||||
|
||||
@@ -138,7 +138,7 @@ absl::StatusOr<Image> FaceStylizer::Stylize(
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options));
|
||||
ConvertToNormalizedRect(image_processing_options, image));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessImageData(
|
||||
@@ -157,7 +157,7 @@ absl::StatusOr<Image> FaceStylizer::StylizeForVideo(
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options));
|
||||
ConvertToNormalizedRect(image_processing_options, image));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessVideoData(
|
||||
@@ -180,7 +180,7 @@ absl::Status FaceStylizer::StylizeAsync(
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options));
|
||||
ConvertToNormalizedRect(image_processing_options, image));
|
||||
return SendLiveStreamData(
|
||||
{{kImageInStreamName,
|
||||
MakePacket<Image>(std::move(image))
|
||||
|
||||
@@ -222,9 +222,9 @@ absl::StatusOr<GestureRecognizerResult> GestureRecognizer::Recognize(
|
||||
"GPU input images are currently not supported.",
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessImageData(
|
||||
@@ -258,9 +258,9 @@ absl::StatusOr<GestureRecognizerResult> GestureRecognizer::RecognizeForVideo(
|
||||
absl::StrCat("GPU input images are currently not supported."),
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessVideoData(
|
||||
@@ -297,9 +297,9 @@ absl::Status GestureRecognizer::RecognizeAsync(
|
||||
absl::StrCat("GPU input images are currently not supported."),
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
return SendLiveStreamData(
|
||||
{{kImageInStreamName,
|
||||
MakePacket<Image>(std::move(image))
|
||||
|
||||
@@ -185,9 +185,9 @@ absl::StatusOr<HandLandmarkerResult> HandLandmarker::Detect(
|
||||
"GPU input images are currently not supported.",
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessImageData(
|
||||
@@ -223,9 +223,9 @@ absl::StatusOr<HandLandmarkerResult> HandLandmarker::DetectForVideo(
|
||||
absl::StrCat("GPU input images are currently not supported."),
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessVideoData(
|
||||
@@ -264,9 +264,9 @@ absl::Status HandLandmarker::DetectAsync(
|
||||
absl::StrCat("GPU input images are currently not supported."),
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
return SendLiveStreamData(
|
||||
{{kImageInStreamName,
|
||||
MakePacket<Image>(std::move(image))
|
||||
|
||||
@@ -156,7 +156,7 @@ absl::StatusOr<ImageClassifierResult> ImageClassifier::Classify(
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options));
|
||||
ConvertToNormalizedRect(image_processing_options, image));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessImageData(
|
||||
@@ -167,7 +167,7 @@ absl::StatusOr<ImageClassifierResult> ImageClassifier::Classify(
|
||||
}
|
||||
|
||||
absl::StatusOr<ImageClassifierResult> ImageClassifier::ClassifyForVideo(
|
||||
Image image, int64 timestamp_ms,
|
||||
Image image, int64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options) {
|
||||
if (image.UsesGpu()) {
|
||||
return CreateStatusWithPayload(
|
||||
@@ -176,7 +176,7 @@ absl::StatusOr<ImageClassifierResult> ImageClassifier::ClassifyForVideo(
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options));
|
||||
ConvertToNormalizedRect(image_processing_options, image));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessVideoData(
|
||||
@@ -191,7 +191,7 @@ absl::StatusOr<ImageClassifierResult> ImageClassifier::ClassifyForVideo(
|
||||
}
|
||||
|
||||
absl::Status ImageClassifier::ClassifyAsync(
|
||||
Image image, int64 timestamp_ms,
|
||||
Image image, int64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options) {
|
||||
if (image.UsesGpu()) {
|
||||
return CreateStatusWithPayload(
|
||||
@@ -200,7 +200,7 @@ absl::Status ImageClassifier::ClassifyAsync(
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options));
|
||||
ConvertToNormalizedRect(image_processing_options, image));
|
||||
return SendLiveStreamData(
|
||||
{{kImageInStreamName,
|
||||
MakePacket<Image>(std::move(image))
|
||||
|
||||
@@ -233,7 +233,7 @@ TEST_F(CreateTest, FailsWithIllegalCallbackInImageOrVideoMode) {
|
||||
JoinPath("./", kTestDataDirectory, kMobileNetQuantizedWithMetadata);
|
||||
options->running_mode = running_mode;
|
||||
options->result_callback = [](absl::StatusOr<ImageClassifierResult>,
|
||||
const Image& image, int64 timestamp_ms) {};
|
||||
const Image& image, int64_t timestamp_ms) {};
|
||||
|
||||
auto image_classifier = ImageClassifier::Create(std::move(options));
|
||||
|
||||
@@ -505,11 +505,9 @@ TEST_F(ImageModeTest, SucceedsWithRotation) {
|
||||
ImageClassifierResult expected;
|
||||
expected.classifications.emplace_back(Classifications{
|
||||
/*categories=*/{
|
||||
{/*index=*/934, /*score=*/0.6371766,
|
||||
/*category_name=*/"cheeseburger"},
|
||||
{/*index=*/963, /*score=*/0.049443405, /*category_name=*/"meat loaf"},
|
||||
{/*index=*/925, /*score=*/0.047918003,
|
||||
/*category_name=*/"guacamole"}},
|
||||
{/*index=*/934, /*score=*/0.754467, /*category_name=*/"cheeseburger"},
|
||||
{/*index=*/925, /*score=*/0.0288028, /*category_name=*/"guacamole"},
|
||||
{/*index=*/932, /*score=*/0.0286119, /*category_name=*/"bagel"}},
|
||||
/*head_index=*/0,
|
||||
/*head_name=*/"probability"});
|
||||
ExpectApproximatelyEqual(results, expected);
|
||||
@@ -525,9 +523,10 @@ TEST_F(ImageModeTest, SucceedsWithRegionOfInterestAndRotation) {
|
||||
options->classifier_options.max_results = 1;
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageClassifier> image_classifier,
|
||||
ImageClassifier::Create(std::move(options)));
|
||||
// Region-of-interest around the chair, with 90° anti-clockwise rotation.
|
||||
RectF roi{/*left=*/0.006, /*top=*/0.1763, /*right=*/0.5702,
|
||||
/*bottom=*/0.3049};
|
||||
// Region-of-interest around the soccer ball, with 90° anti-clockwise
|
||||
// rotation.
|
||||
RectF roi{/*left=*/0.2655, /*top=*/0.45, /*right=*/0.6925,
|
||||
/*bottom=*/0.614};
|
||||
ImageProcessingOptions image_processing_options{roi,
|
||||
/*rotation_degrees=*/-90};
|
||||
|
||||
@@ -536,8 +535,8 @@ TEST_F(ImageModeTest, SucceedsWithRegionOfInterestAndRotation) {
|
||||
|
||||
ImageClassifierResult expected;
|
||||
expected.classifications.emplace_back(
|
||||
Classifications{/*categories=*/{{/*index=*/560, /*score=*/0.6522213,
|
||||
/*category_name=*/"folding chair"}},
|
||||
Classifications{/*categories=*/{{/*index=*/806, /*score=*/0.997684,
|
||||
/*category_name=*/"soccer ball"}},
|
||||
/*head_index=*/0,
|
||||
/*head_name=*/"probability"});
|
||||
ExpectApproximatelyEqual(results, expected);
|
||||
@@ -719,7 +718,7 @@ TEST_F(LiveStreamModeTest, FailsWithCallingWrongMethod) {
|
||||
JoinPath("./", kTestDataDirectory, kMobileNetFloatWithMetadata);
|
||||
options->running_mode = core::RunningMode::LIVE_STREAM;
|
||||
options->result_callback = [](absl::StatusOr<ImageClassifierResult>,
|
||||
const Image& image, int64 timestamp_ms) {};
|
||||
const Image& image, int64_t timestamp_ms) {};
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageClassifier> image_classifier,
|
||||
ImageClassifier::Create(std::move(options)));
|
||||
|
||||
@@ -750,7 +749,7 @@ TEST_F(LiveStreamModeTest, FailsWithOutOfOrderInputTimestamps) {
|
||||
JoinPath("./", kTestDataDirectory, kMobileNetFloatWithMetadata);
|
||||
options->running_mode = core::RunningMode::LIVE_STREAM;
|
||||
options->result_callback = [](absl::StatusOr<ImageClassifierResult>,
|
||||
const Image& image, int64 timestamp_ms) {};
|
||||
const Image& image, int64_t timestamp_ms) {};
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageClassifier> image_classifier,
|
||||
ImageClassifier::Create(std::move(options)));
|
||||
|
||||
@@ -769,7 +768,7 @@ TEST_F(LiveStreamModeTest, FailsWithOutOfOrderInputTimestamps) {
|
||||
struct LiveStreamModeResults {
|
||||
ImageClassifierResult classification_result;
|
||||
std::pair<int, int> image_size;
|
||||
int64 timestamp_ms;
|
||||
int64_t timestamp_ms;
|
||||
};
|
||||
|
||||
TEST_F(LiveStreamModeTest, Succeeds) {
|
||||
@@ -785,7 +784,7 @@ TEST_F(LiveStreamModeTest, Succeeds) {
|
||||
options->classifier_options.max_results = 3;
|
||||
options->result_callback =
|
||||
[&results](absl::StatusOr<ImageClassifierResult> classification_result,
|
||||
const Image& image, int64 timestamp_ms) {
|
||||
const Image& image, int64_t timestamp_ms) {
|
||||
MP_ASSERT_OK(classification_result.status());
|
||||
results.push_back(
|
||||
{.classification_result = std::move(classification_result).value(),
|
||||
@@ -804,7 +803,7 @@ TEST_F(LiveStreamModeTest, Succeeds) {
|
||||
// number of iterations.
|
||||
ASSERT_LE(results.size(), iterations);
|
||||
ASSERT_GT(results.size(), 0);
|
||||
int64 timestamp_ms = -1;
|
||||
int64_t timestamp_ms = -1;
|
||||
for (const auto& result : results) {
|
||||
EXPECT_GT(result.timestamp_ms, timestamp_ms);
|
||||
timestamp_ms = result.timestamp_ms;
|
||||
@@ -828,7 +827,7 @@ TEST_F(LiveStreamModeTest, SucceedsWithRegionOfInterest) {
|
||||
options->classifier_options.max_results = 1;
|
||||
options->result_callback =
|
||||
[&results](absl::StatusOr<ImageClassifierResult> classification_result,
|
||||
const Image& image, int64 timestamp_ms) {
|
||||
const Image& image, int64_t timestamp_ms) {
|
||||
MP_ASSERT_OK(classification_result.status());
|
||||
results.push_back(
|
||||
{.classification_result = std::move(classification_result).value(),
|
||||
@@ -851,7 +850,7 @@ TEST_F(LiveStreamModeTest, SucceedsWithRegionOfInterest) {
|
||||
// number of iterations.
|
||||
ASSERT_LE(results.size(), iterations);
|
||||
ASSERT_GT(results.size(), 0);
|
||||
int64 timestamp_ms = -1;
|
||||
int64_t timestamp_ms = -1;
|
||||
for (const auto& result : results) {
|
||||
EXPECT_GT(result.timestamp_ms, timestamp_ms);
|
||||
timestamp_ms = result.timestamp_ms;
|
||||
|
||||
@@ -151,7 +151,7 @@ absl::StatusOr<ImageEmbedderResult> ImageEmbedder::Embed(
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options));
|
||||
ConvertToNormalizedRect(image_processing_options, image));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessImageData(
|
||||
@@ -172,7 +172,7 @@ absl::StatusOr<ImageEmbedderResult> ImageEmbedder::EmbedForVideo(
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options));
|
||||
ConvertToNormalizedRect(image_processing_options, image));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessVideoData(
|
||||
@@ -196,7 +196,7 @@ absl::Status ImageEmbedder::EmbedAsync(
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options));
|
||||
ConvertToNormalizedRect(image_processing_options, image));
|
||||
return SendLiveStreamData(
|
||||
{{kImageInStreamName,
|
||||
MakePacket<Image>(std::move(image))
|
||||
|
||||
@@ -371,7 +371,7 @@ TEST_F(ImageModeTest, SucceedsWithRotation) {
|
||||
MP_ASSERT_OK_AND_ASSIGN(double similarity, ImageEmbedder::CosineSimilarity(
|
||||
image_result.embeddings[0],
|
||||
rotated_result.embeddings[0]));
|
||||
double expected_similarity = 0.572265;
|
||||
double expected_similarity = 0.98223;
|
||||
EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy);
|
||||
}
|
||||
|
||||
@@ -406,7 +406,7 @@ TEST_F(ImageModeTest, SucceedsWithRegionOfInterestAndRotation) {
|
||||
MP_ASSERT_OK_AND_ASSIGN(double similarity, ImageEmbedder::CosineSimilarity(
|
||||
crop_result.embeddings[0],
|
||||
rotated_result.embeddings[0]));
|
||||
double expected_similarity = 0.62838;
|
||||
double expected_similarity = 0.974683;
|
||||
EXPECT_LE(abs(similarity - expected_similarity), kSimilarityTolerancy);
|
||||
}
|
||||
|
||||
|
||||
@@ -16,6 +16,13 @@ package(default_visibility = ["//mediapipe/tasks:internal"])
|
||||
|
||||
licenses(["notice"])
|
||||
|
||||
cc_library(
|
||||
name = "image_segmenter_result",
|
||||
hdrs = ["image_segmenter_result.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework/formats:image"],
|
||||
)
|
||||
|
||||
# Docs for Mediapipe Tasks Image Segmenter
|
||||
# https://developers.google.com/mediapipe/solutions/vision/image_segmenter
|
||||
cc_library(
|
||||
@@ -25,6 +32,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":image_segmenter_graph",
|
||||
":image_segmenter_result",
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/framework/api2:builder",
|
||||
"//mediapipe/framework/formats:image",
|
||||
@@ -82,6 +90,7 @@ cc_library(
|
||||
"//mediapipe/tasks/cc/vision/utils:image_tensor_specs",
|
||||
"//mediapipe/tasks/metadata:image_segmenter_metadata_schema_cc",
|
||||
"//mediapipe/tasks/metadata:metadata_schema_cc",
|
||||
"//mediapipe/util:graph_builder_utils",
|
||||
"//mediapipe/util:label_map_cc_proto",
|
||||
"//mediapipe/util:label_map_util",
|
||||
"@com_google_absl//absl/status",
|
||||
|
||||
+248
-106
@@ -22,7 +22,21 @@ using mediapipe::kBasicVertexShader;
|
||||
using ::mediapipe::tasks::vision::Shape;
|
||||
using ::mediapipe::tasks::vision::image_segmenter::proto::SegmenterOptions;
|
||||
|
||||
// TODO: This part of the setup code is so common, we should really
|
||||
// refactor to a helper utility.
|
||||
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
|
||||
const GLint attr_location[NUM_ATTRIBUTES] = {
|
||||
ATTRIB_VERTEX,
|
||||
ATTRIB_TEXTURE_POSITION,
|
||||
};
|
||||
const GLchar* attr_name[NUM_ATTRIBUTES] = {
|
||||
"position",
|
||||
"texture_coordinate",
|
||||
};
|
||||
|
||||
// We assume ES3.0+ for some of our shaders here so we can make liberal use of
|
||||
// MRT easily.
|
||||
static constexpr char kEs30RequirementHeader[] = "#version 300 es\n";
|
||||
|
||||
static constexpr char kActivationFragmentShader[] = R"(
|
||||
DEFAULT_PRECISION(mediump, float)
|
||||
@@ -140,55 +154,93 @@ void main() {
|
||||
gl_FragColor = vec4(out_value, out_value, out_value, out_value);
|
||||
})";
|
||||
|
||||
// Quick softmax shader hardcoded to max of N=12 classes. Performs softmax
|
||||
// calculations, but renders to one chunk at a time.
|
||||
// TODO: For more efficiency, should at least use MRT to render all
|
||||
// chunks simultaneously.
|
||||
static constexpr char kSoftmaxShader[] = R"(
|
||||
// Softmax is in 3 steps:
|
||||
// - First we find max over all masks
|
||||
// - Then we transform all masks to be exp(val - maxval), and also add to
|
||||
// cumulative-sum image with MRT
|
||||
// - Then we normalize all masks by cumulative-sum image
|
||||
|
||||
// Part one: max shader
|
||||
// To start with, we just do this chunk by chunk, using GL_MAX blend mode so we
|
||||
// don't need to tap into the max-so-far texture.
|
||||
static constexpr char kMaxShader[] = R"(
|
||||
DEFAULT_PRECISION(mediump, float)
|
||||
in vec2 sample_coordinate;
|
||||
uniform sampler2D input_texture0;
|
||||
uniform sampler2D input_texture1;
|
||||
uniform sampler2D input_texture2;
|
||||
uniform int chunk_select;
|
||||
uniform sampler2D current_chunk;
|
||||
uniform int num_channels; // how many channels from current chunk to use (1-4)
|
||||
|
||||
float max4(vec4 vec) {
|
||||
return max(max(vec.x, vec.y), max(vec.z, vec.w));
|
||||
}
|
||||
|
||||
vec4 expTransform(vec4 vec, float maxval) {
|
||||
return exp(vec - maxval);
|
||||
float max3(vec4 vec) {
|
||||
return max(max(vec.x, vec.y), vec.z);
|
||||
}
|
||||
float max2(vec4 vec) {
|
||||
return max(vec.x, vec.y);
|
||||
}
|
||||
void main() {
|
||||
vec4 chunk_pixel = texture2D(current_chunk, sample_coordinate);
|
||||
float new_max;
|
||||
if (num_channels == 1) {
|
||||
new_max = chunk_pixel.x;
|
||||
} else if (num_channels == 2) {
|
||||
new_max = max2(chunk_pixel);
|
||||
} else if (num_channels == 3) {
|
||||
new_max = max3(chunk_pixel);
|
||||
} else {
|
||||
new_max = max4(chunk_pixel);
|
||||
}
|
||||
gl_FragColor = vec4(new_max, 0.0, 0.0, 1.0);
|
||||
})";
|
||||
|
||||
// Part two: transform-and-sum shader
|
||||
// We use GL blending so we can more easily render a cumulative sum texture, and
|
||||
// this only costs us a glClear for the output chunk (needed since using MRT).
|
||||
static constexpr char kTransformAndSumShader[] = R"(
|
||||
DEFAULT_PRECISION(highp, float)
|
||||
in vec2 sample_coordinate;
|
||||
uniform sampler2D max_value_texture;
|
||||
uniform sampler2D current_chunk;
|
||||
uniform int num_channels; // how many channels from current chunk to use (1-4)
|
||||
|
||||
layout(location = 0) out vec4 cumulative_sum_texture;
|
||||
layout(location = 1) out vec4 out_chunk_texture;
|
||||
|
||||
void main() {
|
||||
// Grab all vecs
|
||||
vec4 pixel0 = texture2D(input_texture0, sample_coordinate);
|
||||
vec4 pixel1 = texture2D(input_texture1, sample_coordinate);
|
||||
vec4 pixel2 = texture2D(input_texture2, sample_coordinate);
|
||||
float max_pixel = texture(max_value_texture, sample_coordinate).r;
|
||||
vec4 chunk_pixel = texture(current_chunk, sample_coordinate);
|
||||
vec4 new_chunk_pixel = exp(chunk_pixel - max_pixel);
|
||||
|
||||
// Find maxval amongst all vectors
|
||||
float max0 = max4(pixel0);
|
||||
float max1 = max4(pixel1);
|
||||
float max2 = max4(pixel2);
|
||||
float maxval = max(max(max0, max1), max2);
|
||||
float sum_so_far;
|
||||
if (num_channels == 1) {
|
||||
sum_so_far = new_chunk_pixel.x;
|
||||
} else if (num_channels == 2) {
|
||||
sum_so_far = dot(vec2(1.0, 1.0), new_chunk_pixel.xy);
|
||||
} else if (num_channels == 3) {
|
||||
sum_so_far = dot(vec3(1.0, 1.0, 1.0), new_chunk_pixel.xyz);
|
||||
} else {
|
||||
sum_so_far = dot(vec4(1.0, 1.0, 1.0, 1.0), new_chunk_pixel);
|
||||
}
|
||||
|
||||
vec4 outPixel0 = expTransform(pixel0, maxval);
|
||||
vec4 outPixel1 = expTransform(pixel1, maxval);
|
||||
vec4 outPixel2 = expTransform(pixel2, maxval);
|
||||
cumulative_sum_texture = vec4(sum_so_far, 0.0, 0.0, 1.0);
|
||||
out_chunk_texture = new_chunk_pixel;
|
||||
})";
|
||||
|
||||
// Quick hack to sum all components in vec4: dot with <1, 1, 1, 1>
|
||||
vec4 ones = vec4(1.0, 1.0, 1.0, 1.0);
|
||||
float weightSum = dot(ones, outPixel0) + dot(ones, outPixel1) + dot(ones, outPixel2);
|
||||
// Part three: normalization shader
|
||||
static constexpr char kNormalizationShader[] = R"(
|
||||
DEFAULT_PRECISION(mediump, float)
|
||||
in vec2 sample_coordinate;
|
||||
uniform sampler2D sum_texture; // cumulative summation value (to normalize by)
|
||||
uniform sampler2D current_chunk; // current chunk
|
||||
|
||||
vec4 outPixel;
|
||||
if (chunk_select == 0) {
|
||||
outPixel = outPixel0 / weightSum;
|
||||
} else if (chunk_select == 1) {
|
||||
outPixel = outPixel1 / weightSum;
|
||||
} else {
|
||||
outPixel = outPixel2 / weightSum;
|
||||
}
|
||||
gl_FragColor = outPixel;
|
||||
void main() {
|
||||
float sum_pixel = texture2D(sum_texture, sample_coordinate).r;
|
||||
vec4 chunk_pixel = texture2D(current_chunk, sample_coordinate);
|
||||
|
||||
// NOTE: We assume non-zero sum_pixel here, which is a safe assumption for
|
||||
// result of an exp transform, but not if this shader is extended to other
|
||||
// uses.
|
||||
gl_FragColor = chunk_pixel / sum_pixel;
|
||||
})";
|
||||
|
||||
} // namespace
|
||||
@@ -208,19 +260,38 @@ absl::Status SegmentationPostprocessorGl::Initialize(
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SegmentationPostprocessorGl::CreateBasicFragmentShaderProgram(
|
||||
std::string const& program_name, std::string const& fragment_shader_source,
|
||||
std::vector<std::string> const& uniform_names, GlShader* shader_struct_ptr,
|
||||
bool is_es30_only = false) {
|
||||
// Format source and create basic ES3.0+ fragment-shader-only program
|
||||
const std::string frag_shader_source =
|
||||
absl::StrCat(is_es30_only ? std::string(kEs30RequirementHeader) : "",
|
||||
std::string(mediapipe::kMediaPipeFragmentShaderPreamble),
|
||||
std::string(fragment_shader_source));
|
||||
const std::string vert_shader_source =
|
||||
absl::StrCat(is_es30_only ? std::string(kEs30RequirementHeader) : "",
|
||||
std::string(kBasicVertexShader));
|
||||
mediapipe::GlhCreateProgram(
|
||||
vert_shader_source.c_str(), frag_shader_source.c_str(), NUM_ATTRIBUTES,
|
||||
&attr_name[0], attr_location, &shader_struct_ptr->program,
|
||||
/* force_log_errors */ true);
|
||||
RET_CHECK(shader_struct_ptr->program)
|
||||
<< "Problem initializing the " << program_name << " program.";
|
||||
|
||||
// Hook up all desired uniforms
|
||||
for (const auto& uniform_name : uniform_names) {
|
||||
shader_struct_ptr->uniforms[uniform_name] =
|
||||
glGetUniformLocation(shader_struct_ptr->program, uniform_name.c_str());
|
||||
RET_CHECK(shader_struct_ptr->uniforms[uniform_name] > 0)
|
||||
<< uniform_name << " uniform not found for " << program_name
|
||||
<< " program";
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
absl::Status SegmentationPostprocessorGl::GlInit() {
|
||||
return helper_.RunInGlContext([this]() -> absl::Status {
|
||||
// TODO: This part of the setup code is so common, we should really
|
||||
// refactor to a helper utility.
|
||||
const GLint attr_location[NUM_ATTRIBUTES] = {
|
||||
ATTRIB_VERTEX,
|
||||
ATTRIB_TEXTURE_POSITION,
|
||||
};
|
||||
const GLchar* attr_name[NUM_ATTRIBUTES] = {
|
||||
"position",
|
||||
"texture_coordinate",
|
||||
};
|
||||
|
||||
// Default to passthrough/NONE
|
||||
std::string activation_fn = "vec4 out_value = in_value;";
|
||||
switch (options_.segmenter_options().activation()) {
|
||||
@@ -263,9 +334,17 @@ absl::Status SegmentationPostprocessorGl::GlInit() {
|
||||
absl::StrCat(std::string(mediapipe::kMediaPipeFragmentShaderPreamble),
|
||||
std::string(kArgmaxShader));
|
||||
|
||||
const std::string softmax_shader_source =
|
||||
absl::StrCat(std::string(mediapipe::kMediaPipeFragmentShaderPreamble),
|
||||
std::string(kSoftmaxShader));
|
||||
// Softmax shaders (Max, Transform+Sum, and Normalization)
|
||||
MP_RETURN_IF_ERROR(CreateBasicFragmentShaderProgram(
|
||||
"softmax max", kMaxShader, {"current_chunk", "num_channels"},
|
||||
&softmax_max_shader_));
|
||||
MP_RETURN_IF_ERROR(CreateBasicFragmentShaderProgram(
|
||||
"softmax transform-and-sum", kTransformAndSumShader,
|
||||
{"max_value_texture", "current_chunk", "num_channels"},
|
||||
&softmax_transform_and_sum_shader_, true /* is_es30_only */));
|
||||
MP_RETURN_IF_ERROR(CreateBasicFragmentShaderProgram(
|
||||
"softmax normalization", kNormalizationShader,
|
||||
{"sum_texture", "current_chunk"}, &softmax_normalization_shader_));
|
||||
|
||||
// Compile all our shader programs.
|
||||
// Note: we enable `force_log_errors` so that we get full debugging error
|
||||
@@ -299,12 +378,6 @@ absl::Status SegmentationPostprocessorGl::GlInit() {
|
||||
/* force_log_errors */ true);
|
||||
RET_CHECK(argmax_program_) << "Problem initializing the argmax program.";
|
||||
|
||||
mediapipe::GlhCreateProgram(kBasicVertexShader,
|
||||
softmax_shader_source.c_str(), NUM_ATTRIBUTES,
|
||||
&attr_name[0], attr_location, &softmax_program_,
|
||||
/* force_log_errors */ true);
|
||||
RET_CHECK(softmax_program_) << "Problem initializing the softmax program.";
|
||||
|
||||
// Get uniform locations.
|
||||
activation_texture_uniform_ =
|
||||
glGetUniformLocation(activation_program_, "input_texture");
|
||||
@@ -341,23 +414,6 @@ absl::Status SegmentationPostprocessorGl::GlInit() {
|
||||
RET_CHECK(argmax_texture2_uniform_ > 0)
|
||||
<< "argmax input_texture2 uniform not found.";
|
||||
|
||||
softmax_texture0_uniform_ =
|
||||
glGetUniformLocation(softmax_program_, "input_texture0");
|
||||
RET_CHECK(softmax_texture0_uniform_ > 0)
|
||||
<< "softmax input_texture0 uniform not found.";
|
||||
softmax_texture1_uniform_ =
|
||||
glGetUniformLocation(softmax_program_, "input_texture1");
|
||||
RET_CHECK(softmax_texture1_uniform_ > 0)
|
||||
<< "softmax input_texture1 uniform not found.";
|
||||
softmax_texture2_uniform_ =
|
||||
glGetUniformLocation(softmax_program_, "input_texture2");
|
||||
RET_CHECK(softmax_texture2_uniform_ > 0)
|
||||
<< "softmax input_texture2 uniform not found.";
|
||||
softmax_chunk_select_uniform_ =
|
||||
glGetUniformLocation(softmax_program_, "chunk_select");
|
||||
RET_CHECK(softmax_chunk_select_uniform_ > 0)
|
||||
<< "softmax chunk select uniform not found.";
|
||||
|
||||
// TODO: If ES3.0+ only, switch to VAO for handling attributes.
|
||||
glGenBuffers(1, &square_vertices_);
|
||||
glBindBuffer(GL_ARRAY_BUFFER, square_vertices_);
|
||||
@@ -408,6 +464,9 @@ SegmentationPostprocessorGl::GetSegmentationResultGpu(const Shape& input_shape,
|
||||
|
||||
// Uint8 pipeline and conversions are lacking, so for now we just use F32
|
||||
// textures even for category masks.
|
||||
// TODO: Also, some platforms (like certain iOS devices) do not
|
||||
// allow for rendering to RGBAF32 textures, so we should switch to using
|
||||
// F16 textures in those instances.
|
||||
const GpuBufferFormat final_output_format = GpuBufferFormat::kGrayFloat32;
|
||||
const Tensor::OpenGlTexture2dView read_view =
|
||||
tensor.GetOpenGlTexture2dReadView();
|
||||
@@ -467,7 +526,7 @@ SegmentationPostprocessorGl::GetSegmentationResultGpu(const Shape& input_shape,
|
||||
((float)i + tex_offset) / (float)(input_width));
|
||||
// Technically duplicated, but fine for now; we want this after the bind
|
||||
glBindTexture(GL_TEXTURE_2D, activated_texture.name());
|
||||
// Disable HW interpolation
|
||||
// Disable hardware GPU interpolation
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
|
||||
glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
|
||||
// Render
|
||||
@@ -477,45 +536,126 @@ SegmentationPostprocessorGl::GetSegmentationResultGpu(const Shape& input_shape,
|
||||
|
||||
std::vector<GlTexture> softmax_chunks;
|
||||
if (is_softmax) {
|
||||
// Step 2.5: For SOFTMAX, apply softmax shader with up to 3 textures to
|
||||
// create softmax-transformed chunks before channel extraction.
|
||||
RET_CHECK(num_chunks <= 3)
|
||||
<< "Cannot handle more than 12 classes in softmax shader.";
|
||||
// Step 2.5: For SOFTMAX, apply softmax shaders (max, transformAndSum, and
|
||||
// normalization) to create softmax-transformed chunks before channel
|
||||
// extraction.
|
||||
// NOTE: exp(x-C) / sum_over_x(exp(x-C)) = exp(x) / sum_over_x(exp(x)). So
|
||||
// theoretically we can skip the max shader step entirely. However,
|
||||
// applying it does bring all our values into a nice (0, 1] range, so it
|
||||
// will likely be better for precision, especially when dealing with an
|
||||
// exponential function on arbitrary values. Therefore, we keep it, but
|
||||
// this is potentially a skippable step for known "good" models, if we
|
||||
// ever want to provide that as an option.
|
||||
// TODO: For a tiny bit more efficiency, could combine channel
|
||||
// extraction into last step of this via MRT.
|
||||
|
||||
glUseProgram(softmax_program_);
|
||||
glUniform1i(softmax_texture0_uniform_, 1);
|
||||
glUniform1i(softmax_texture1_uniform_, 2);
|
||||
glUniform1i(softmax_texture2_uniform_, 3);
|
||||
// Max
|
||||
glUseProgram(softmax_max_shader_.program);
|
||||
glUniform1i(softmax_max_shader_.uniforms["current_chunk"], 1);
|
||||
|
||||
// We just need one channel, so format will match final output confidence
|
||||
// masks
|
||||
auto max_texture =
|
||||
helper_.CreateDestinationTexture(width, height, final_output_format);
|
||||
helper_.BindFramebuffer(max_texture);
|
||||
|
||||
// We clear our newly-created destination texture to a reasonable minimum.
|
||||
glClearColor(0.0, 0.0, 0.0, 0.0);
|
||||
glClear(GL_COLOR_BUFFER_BIT);
|
||||
|
||||
// We will use hardware GPU blending to apply max to all our writes.
|
||||
glEnable(GL_BLEND);
|
||||
glBlendEquation(GL_MAX);
|
||||
|
||||
glActiveTexture(GL_TEXTURE1);
|
||||
for (int i = 0; i < num_chunks; i++) {
|
||||
int num_channels = 4;
|
||||
if ((i + 1) * 4 > num_outputs) num_channels = num_outputs % 4;
|
||||
glUniform1i(softmax_max_shader_.uniforms["num_channels"], num_channels);
|
||||
glBindTexture(GL_TEXTURE_2D, chunks[i].name());
|
||||
glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
|
||||
}
|
||||
|
||||
// Transform & Sum
|
||||
std::vector<GlTexture> unnormalized_softmax_chunks;
|
||||
glUseProgram(softmax_transform_and_sum_shader_.program);
|
||||
glUniform1i(softmax_transform_and_sum_shader_.uniforms["current_chunk"],
|
||||
1);
|
||||
glUniform1i(
|
||||
softmax_transform_and_sum_shader_.uniforms["max_value_texture"], 2);
|
||||
|
||||
auto sum_texture =
|
||||
helper_.CreateDestinationTexture(width, height, final_output_format);
|
||||
helper_.BindFramebuffer(sum_texture);
|
||||
glClear(GL_COLOR_BUFFER_BIT);
|
||||
|
||||
glActiveTexture(GL_TEXTURE2);
|
||||
glBindTexture(GL_TEXTURE_2D, max_texture.name());
|
||||
|
||||
glBlendEquation(GL_FUNC_ADD);
|
||||
glBlendFunc(GL_ONE, GL_ONE);
|
||||
glActiveTexture(GL_TEXTURE1);
|
||||
|
||||
// We use glDrawBuffers to clear only the new texture, then again to
|
||||
// draw to both textures simultaneously for rendering.
|
||||
GLuint both_attachments[2] = {GL_COLOR_ATTACHMENT0, GL_COLOR_ATTACHMENT1};
|
||||
GLuint one_attachment[2] = {GL_NONE, GL_COLOR_ATTACHMENT1};
|
||||
for (int i = 0; i < num_chunks; i++) {
|
||||
int num_channels = 4;
|
||||
if ((i + 1) * 4 > num_outputs) num_channels = num_outputs % 4;
|
||||
glUniform1i(softmax_transform_and_sum_shader_.uniforms["num_channels"],
|
||||
num_channels);
|
||||
unnormalized_softmax_chunks.push_back(helper_.CreateDestinationTexture(
|
||||
width, height, chunk_output_format));
|
||||
glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT1,
|
||||
GL_TEXTURE_2D,
|
||||
unnormalized_softmax_chunks.back().name(), 0);
|
||||
|
||||
// Note that we must bind AFTER the CreateDestinationTexture, or else we
|
||||
// end up with (0, 0, 0, 1) data being read from an unbound texture
|
||||
// unit.
|
||||
glBindTexture(GL_TEXTURE_2D, chunks[i].name());
|
||||
|
||||
// Clear *only* the new chunk
|
||||
glDrawBuffers(2, one_attachment);
|
||||
glClear(GL_COLOR_BUFFER_BIT);
|
||||
|
||||
// Then draw into both
|
||||
glDrawBuffers(2, both_attachments);
|
||||
glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
|
||||
}
|
||||
|
||||
// Turn off MRT and blending, and unbind second color attachment
|
||||
glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT1,
|
||||
GL_TEXTURE_2D, 0, 0);
|
||||
glDrawBuffers(1, both_attachments);
|
||||
glDisable(GL_BLEND);
|
||||
|
||||
// Normalize each chunk into a new chunk as our final step
|
||||
glUseProgram(softmax_normalization_shader_.program);
|
||||
glUniform1i(softmax_normalization_shader_.uniforms["current_chunk"], 1);
|
||||
glUniform1i(softmax_normalization_shader_.uniforms["sum_texture"], 2);
|
||||
|
||||
glActiveTexture(GL_TEXTURE2);
|
||||
glBindTexture(GL_TEXTURE_2D, sum_texture.name());
|
||||
glActiveTexture(GL_TEXTURE1);
|
||||
|
||||
for (int i = 0; i < num_chunks; i++) {
|
||||
glUniform1i(softmax_chunk_select_uniform_, i);
|
||||
softmax_chunks.push_back(helper_.CreateDestinationTexture(
|
||||
output_width, output_height, chunk_output_format));
|
||||
width, height, chunk_output_format));
|
||||
helper_.BindFramebuffer(softmax_chunks.back());
|
||||
|
||||
// Bind however many chunks we have
|
||||
for (int j = 0; j < num_chunks; ++j) {
|
||||
glActiveTexture(GL_TEXTURE1 + j);
|
||||
glBindTexture(GL_TEXTURE_2D, chunks[j].name());
|
||||
}
|
||||
|
||||
for (int j = num_chunks; j < 3; ++j) { // 3 is hard-coded max chunks
|
||||
glActiveTexture(GL_TEXTURE1 + j);
|
||||
// If texture is unbound, sampling from it should always give zeros.
|
||||
// This is not ideal, but is ok for now for not polluting the argmax
|
||||
// shader results too much.
|
||||
glBindTexture(GL_TEXTURE_2D, 0);
|
||||
}
|
||||
|
||||
glBindTexture(GL_TEXTURE_2D, unnormalized_softmax_chunks[i].name());
|
||||
glClear(GL_COLOR_BUFFER_BIT);
|
||||
glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
|
||||
}
|
||||
|
||||
// Unbind the extra textures here.
|
||||
for (int i = 0; i < num_chunks; ++i) {
|
||||
glActiveTexture(GL_TEXTURE1 + i);
|
||||
glBindTexture(GL_TEXTURE_2D, 0);
|
||||
}
|
||||
// Unbind textures here
|
||||
glActiveTexture(GL_TEXTURE2);
|
||||
glBindTexture(GL_TEXTURE_2D, 0);
|
||||
// We make sure to switch back to texture unit 1, since our confidence
|
||||
// mask extraction code assumes that's our default.
|
||||
glActiveTexture(GL_TEXTURE1);
|
||||
glBindTexture(GL_TEXTURE_2D, 0);
|
||||
}
|
||||
|
||||
std::vector<GlTexture> outputs;
|
||||
@@ -607,17 +747,19 @@ SegmentationPostprocessorGl::~SegmentationPostprocessorGl() {
|
||||
glDeleteProgram(activation_program_);
|
||||
glDeleteProgram(argmax_program_);
|
||||
glDeleteProgram(channel_select_program_);
|
||||
glDeleteProgram(softmax_program_);
|
||||
glDeleteProgram(split_program_);
|
||||
glDeleteBuffers(1, &square_vertices_);
|
||||
glDeleteBuffers(1, &texture_vertices_);
|
||||
activation_program_ = 0;
|
||||
argmax_program_ = 0;
|
||||
channel_select_program_ = 0;
|
||||
softmax_program_ = 0;
|
||||
split_program_ = 0;
|
||||
square_vertices_ = 0;
|
||||
texture_vertices_ = 0;
|
||||
|
||||
glDeleteProgram(softmax_max_shader_.program);
|
||||
glDeleteProgram(softmax_transform_and_sum_shader_.program);
|
||||
glDeleteProgram(softmax_normalization_shader_.program);
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
+14
-5
@@ -38,7 +38,17 @@ class SegmentationPostprocessorGl {
|
||||
const Tensor& tensor);
|
||||
|
||||
private:
|
||||
struct GlShader {
|
||||
GLuint program = 0;
|
||||
absl::flat_hash_map<std::string, GLint> uniforms;
|
||||
};
|
||||
|
||||
absl::Status GlInit();
|
||||
absl::Status CreateBasicFragmentShaderProgram(
|
||||
std::string const& program_name,
|
||||
std::string const& fragment_shader_source,
|
||||
std::vector<std::string> const& uniform_names,
|
||||
GlShader* shader_struct_ptr, bool is_es30_only);
|
||||
|
||||
TensorsToSegmentationCalculatorOptions options_;
|
||||
GlCalculatorHelper helper_;
|
||||
@@ -47,7 +57,6 @@ class SegmentationPostprocessorGl {
|
||||
GLuint activation_program_ = 0;
|
||||
GLuint argmax_program_ = 0;
|
||||
GLuint channel_select_program_ = 0;
|
||||
GLuint softmax_program_ = 0;
|
||||
GLuint split_program_ = 0;
|
||||
GLuint square_vertices_ = 0;
|
||||
GLuint texture_vertices_ = 0;
|
||||
@@ -57,12 +66,12 @@ class SegmentationPostprocessorGl {
|
||||
GLint argmax_texture2_uniform_;
|
||||
GLint channel_select_texture_uniform_;
|
||||
GLint channel_select_index_uniform_;
|
||||
GLint softmax_texture0_uniform_;
|
||||
GLint softmax_texture1_uniform_;
|
||||
GLint softmax_texture2_uniform_;
|
||||
GLint softmax_chunk_select_uniform_;
|
||||
GLint split_texture_uniform_;
|
||||
GLint split_x_offset_uniform_;
|
||||
|
||||
GlShader softmax_max_shader_;
|
||||
GlShader softmax_transform_and_sum_shader_;
|
||||
GlShader softmax_normalization_shader_;
|
||||
};
|
||||
|
||||
} // namespace tasks
|
||||
|
||||
+72
-33
@@ -80,10 +80,10 @@ void Sigmoid(absl::Span<const float> values,
|
||||
[](float value) { return 1. / (1 + std::exp(-value)); });
|
||||
}
|
||||
|
||||
std::vector<Image> ProcessForCategoryMaskCpu(const Shape& input_shape,
|
||||
const Shape& output_shape,
|
||||
const SegmenterOptions& options,
|
||||
const float* tensors_buffer) {
|
||||
Image ProcessForCategoryMaskCpu(const Shape& input_shape,
|
||||
const Shape& output_shape,
|
||||
const SegmenterOptions& options,
|
||||
const float* tensors_buffer) {
|
||||
cv::Mat resized_tensors_mat;
|
||||
cv::Mat tensors_mat_view(
|
||||
input_shape.height, input_shape.width, CV_32FC(input_shape.channels),
|
||||
@@ -135,7 +135,7 @@ std::vector<Image> ProcessForCategoryMaskCpu(const Shape& input_shape,
|
||||
pixel = maximum_category_idx;
|
||||
}
|
||||
});
|
||||
return {category_mask};
|
||||
return category_mask;
|
||||
}
|
||||
|
||||
std::vector<Image> ProcessForConfidenceMaskCpu(const Shape& input_shape,
|
||||
@@ -209,7 +209,9 @@ std::vector<Image> ProcessForConfidenceMaskCpu(const Shape& input_shape,
|
||||
|
||||
} // namespace
|
||||
|
||||
// Converts Tensors from a vector of Tensor to Segmentation.
|
||||
// Converts Tensors from a vector of Tensor to Segmentation masks. The
|
||||
// calculator always output confidence masks, and an optional category mask if
|
||||
// CATEGORY_MASK is connected.
|
||||
//
|
||||
// Performs optional resizing to OUTPUT_SIZE dimension if provided,
|
||||
// otherwise the segmented masks is the same size as input tensor.
|
||||
@@ -221,7 +223,12 @@ std::vector<Image> ProcessForConfidenceMaskCpu(const Shape& input_shape,
|
||||
// the size to resize masks to.
|
||||
//
|
||||
// Output:
|
||||
// Segmentation: Segmentation proto.
|
||||
// CONFIDENCE_MASK @Multiple: Multiple masks of float image where, for each
|
||||
// mask, each pixel represents the prediction confidence, usually in the [0,
|
||||
// 1] range.
|
||||
// CATEGORY_MASK @Optional: A category mask of uint8 image where each pixel
|
||||
// represents the class which the pixel in the original image was predicted to
|
||||
// belong to.
|
||||
//
|
||||
// Options:
|
||||
// See tensors_to_segmentation_calculator.proto
|
||||
@@ -231,13 +238,13 @@ std::vector<Image> ProcessForConfidenceMaskCpu(const Shape& input_shape,
|
||||
// calculator: "TensorsToSegmentationCalculator"
|
||||
// input_stream: "TENSORS:tensors"
|
||||
// input_stream: "OUTPUT_SIZE:size"
|
||||
// output_stream: "SEGMENTATION:0:segmentation"
|
||||
// output_stream: "SEGMENTATION:1:segmentation"
|
||||
// output_stream: "CONFIDENCE_MASK:0:confidence_mask"
|
||||
// output_stream: "CONFIDENCE_MASK:1:confidence_mask"
|
||||
// output_stream: "CATEGORY_MASK:category_mask"
|
||||
// options {
|
||||
// [mediapipe.tasks.TensorsToSegmentationCalculatorOptions.ext] {
|
||||
// segmenter_options {
|
||||
// activation: SOFTMAX
|
||||
// output_type: CONFIDENCE_MASK
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
@@ -248,7 +255,11 @@ class TensorsToSegmentationCalculator : public Node {
|
||||
static constexpr Input<std::pair<int, int>>::Optional kOutputSizeIn{
|
||||
"OUTPUT_SIZE"};
|
||||
static constexpr Output<Image>::Multiple kSegmentationOut{"SEGMENTATION"};
|
||||
MEDIAPIPE_NODE_CONTRACT(kTensorsIn, kOutputSizeIn, kSegmentationOut);
|
||||
static constexpr Output<Image>::Multiple kConfidenceMaskOut{
|
||||
"CONFIDENCE_MASK"};
|
||||
static constexpr Output<Image>::Optional kCategoryMaskOut{"CATEGORY_MASK"};
|
||||
MEDIAPIPE_NODE_CONTRACT(kTensorsIn, kOutputSizeIn, kSegmentationOut,
|
||||
kConfidenceMaskOut, kCategoryMaskOut);
|
||||
|
||||
static absl::Status UpdateContract(CalculatorContract* cc);
|
||||
|
||||
@@ -279,9 +290,13 @@ absl::Status TensorsToSegmentationCalculator::UpdateContract(
|
||||
absl::Status TensorsToSegmentationCalculator::Open(
|
||||
mediapipe::CalculatorContext* cc) {
|
||||
options_ = cc->Options<TensorsToSegmentationCalculatorOptions>();
|
||||
RET_CHECK_NE(options_.segmenter_options().output_type(),
|
||||
SegmenterOptions::UNSPECIFIED)
|
||||
<< "Must specify output_type as one of [CONFIDENCE_MASK|CATEGORY_MASK].";
|
||||
// TODO: remove deprecated output type support.
|
||||
if (options_.segmenter_options().has_output_type()) {
|
||||
RET_CHECK_NE(options_.segmenter_options().output_type(),
|
||||
SegmenterOptions::UNSPECIFIED)
|
||||
<< "Must specify output_type as one of "
|
||||
"[CONFIDENCE_MASK|CATEGORY_MASK].";
|
||||
}
|
||||
#ifdef __EMSCRIPTEN__
|
||||
MP_RETURN_IF_ERROR(postprocessor_.Initialize(cc, options_));
|
||||
#endif // __EMSCRIPTEN__
|
||||
@@ -309,6 +324,10 @@ absl::Status TensorsToSegmentationCalculator::Process(
|
||||
if (cc->Inputs().HasTag("OUTPUT_SIZE")) {
|
||||
std::tie(output_width, output_height) = kOutputSizeIn(cc).Get();
|
||||
}
|
||||
|
||||
// Use GPU postprocessing on web when Tensor is there already and has <= 12
|
||||
// categories.
|
||||
#ifdef __EMSCRIPTEN__
|
||||
Shape output_shape = {
|
||||
/* height= */ output_height,
|
||||
/* width= */ output_width,
|
||||
@@ -316,33 +335,53 @@ absl::Status TensorsToSegmentationCalculator::Process(
|
||||
SegmenterOptions::CATEGORY_MASK
|
||||
? 1
|
||||
: input_shape.channels};
|
||||
|
||||
// Use GPU postprocessing on web when Tensor is there already and has <= 12
|
||||
// categories.
|
||||
#ifdef __EMSCRIPTEN__
|
||||
if (input_tensor.ready_as_opengl_texture_2d() && input_shape.channels <= 12) {
|
||||
std::vector<std::unique_ptr<Image>> segmented_masks =
|
||||
postprocessor_.GetSegmentationResultGpu(input_shape, output_shape,
|
||||
input_tensor);
|
||||
for (int i = 0; i < segmented_masks.size(); ++i) {
|
||||
// Real output on GPU.
|
||||
// kSegmentationOut(cc)[i].Send(std::move(segmented_masks[i]));
|
||||
|
||||
// Reformat as CPU for now for testing.
|
||||
// TODO: Switch to real GPU output when GPU output pipeline is
|
||||
// ready.
|
||||
Image new_image(segmented_masks[i]->GetImageFrameSharedPtr());
|
||||
kSegmentationOut(cc)[i].Send(std::move(new_image));
|
||||
kSegmentationOut(cc)[i].Send(std::move(segmented_masks[i]));
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
#endif // __EMSCRIPTEN__
|
||||
|
||||
// Otherwise, use CPU postprocessing.
|
||||
std::vector<Image> segmented_masks = GetSegmentationResultCpu(
|
||||
input_shape, output_shape, input_tensor.GetCpuReadView().buffer<float>());
|
||||
for (int i = 0; i < segmented_masks.size(); ++i) {
|
||||
kSegmentationOut(cc)[i].Send(std::move(segmented_masks[i]));
|
||||
const float* tensors_buffer = input_tensor.GetCpuReadView().buffer<float>();
|
||||
|
||||
// TODO: remove deprecated output type support.
|
||||
if (options_.segmenter_options().has_output_type()) {
|
||||
std::vector<Image> segmented_masks = GetSegmentationResultCpu(
|
||||
input_shape,
|
||||
{/* height= */ output_height,
|
||||
/* width= */ output_width,
|
||||
/* channels= */ options_.segmenter_options().output_type() ==
|
||||
SegmenterOptions::CATEGORY_MASK
|
||||
? 1
|
||||
: input_shape.channels},
|
||||
input_tensor.GetCpuReadView().buffer<float>());
|
||||
for (int i = 0; i < segmented_masks.size(); ++i) {
|
||||
kSegmentationOut(cc)[i].Send(std::move(segmented_masks[i]));
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
std::vector<Image> confidence_masks =
|
||||
ProcessForConfidenceMaskCpu(input_shape,
|
||||
{/* height= */ output_height,
|
||||
/* width= */ output_width,
|
||||
/* channels= */ input_shape.channels},
|
||||
options_.segmenter_options(), tensors_buffer);
|
||||
for (int i = 0; i < confidence_masks.size(); ++i) {
|
||||
kConfidenceMaskOut(cc)[i].Send(std::move(confidence_masks[i]));
|
||||
}
|
||||
if (cc->Outputs().HasTag("CATEGORY_MASK")) {
|
||||
kCategoryMaskOut(cc).Send(ProcessForCategoryMaskCpu(
|
||||
input_shape,
|
||||
{/* height= */ output_height,
|
||||
/* width= */ output_width,
|
||||
/* channels= */ 1},
|
||||
options_.segmenter_options(), tensors_buffer));
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
@@ -352,9 +391,9 @@ std::vector<Image> TensorsToSegmentationCalculator::GetSegmentationResultCpu(
|
||||
const float* tensors_buffer) {
|
||||
if (options_.segmenter_options().output_type() ==
|
||||
SegmenterOptions::CATEGORY_MASK) {
|
||||
return ProcessForCategoryMaskCpu(input_shape, output_shape,
|
||||
options_.segmenter_options(),
|
||||
tensors_buffer);
|
||||
return {ProcessForCategoryMaskCpu(input_shape, output_shape,
|
||||
options_.segmenter_options(),
|
||||
tensors_buffer)};
|
||||
} else {
|
||||
return ProcessForConfidenceMaskCpu(input_shape, output_shape,
|
||||
options_.segmenter_options(),
|
||||
|
||||
+37
-49
@@ -79,8 +79,9 @@ void PushTensorsToRunner(int tensor_height, int tensor_width,
|
||||
std::vector<Packet> GetPackets(const CalculatorRunner& runner) {
|
||||
std::vector<Packet> mask_packets;
|
||||
for (int i = 0; i < runner.Outputs().NumEntries(); ++i) {
|
||||
EXPECT_EQ(runner.Outputs().Get("SEGMENTATION", i).packets.size(), 1);
|
||||
mask_packets.push_back(runner.Outputs().Get("SEGMENTATION", i).packets[0]);
|
||||
EXPECT_EQ(runner.Outputs().Get("CONFIDENCE_MASK", i).packets.size(), 1);
|
||||
mask_packets.push_back(
|
||||
runner.Outputs().Get("CONFIDENCE_MASK", i).packets[0]);
|
||||
}
|
||||
return mask_packets;
|
||||
}
|
||||
@@ -118,13 +119,10 @@ TEST(TensorsToSegmentationCalculatorTest, FailsInvalidTensorDimensionOne) {
|
||||
R"pb(
|
||||
calculator: "mediapipe.tasks.TensorsToSegmentationCalculator"
|
||||
input_stream: "TENSORS:tensors"
|
||||
output_stream: "SEGMENTATION:segmentation"
|
||||
output_stream: "CONFIDENCE_MASK:segmentation"
|
||||
options {
|
||||
[mediapipe.tasks.TensorsToSegmentationCalculatorOptions.ext] {
|
||||
segmenter_options {
|
||||
activation: SOFTMAX
|
||||
output_type: CONFIDENCE_MASK
|
||||
}
|
||||
segmenter_options { activation: SOFTMAX }
|
||||
}
|
||||
}
|
||||
)pb"));
|
||||
@@ -145,13 +143,10 @@ TEST(TensorsToSegmentationCalculatorTest, FailsInvalidTensorDimensionFive) {
|
||||
R"pb(
|
||||
calculator: "mediapipe.tasks.TensorsToSegmentationCalculator"
|
||||
input_stream: "TENSORS:tensors"
|
||||
output_stream: "SEGMENTATION:segmentation"
|
||||
output_stream: "CONFIDENCE_MASK:segmentation"
|
||||
options {
|
||||
[mediapipe.tasks.TensorsToSegmentationCalculatorOptions.ext] {
|
||||
segmenter_options {
|
||||
activation: SOFTMAX
|
||||
output_type: CONFIDENCE_MASK
|
||||
}
|
||||
segmenter_options { activation: SOFTMAX }
|
||||
}
|
||||
}
|
||||
)pb"));
|
||||
@@ -173,16 +168,13 @@ TEST(TensorsToSegmentationCalculatorTest, SucceedsConfidenceMaskWithSoftmax) {
|
||||
R"pb(
|
||||
calculator: "mediapipe.tasks.TensorsToSegmentationCalculator"
|
||||
input_stream: "TENSORS:tensors"
|
||||
output_stream: "SEGMENTATION:0:segmented_mask_0"
|
||||
output_stream: "SEGMENTATION:1:segmented_mask_1"
|
||||
output_stream: "SEGMENTATION:2:segmented_mask_2"
|
||||
output_stream: "SEGMENTATION:3:segmented_mask_3"
|
||||
output_stream: "CONFIDENCE_MASK:0:segmented_mask_0"
|
||||
output_stream: "CONFIDENCE_MASK:1:segmented_mask_1"
|
||||
output_stream: "CONFIDENCE_MASK:2:segmented_mask_2"
|
||||
output_stream: "CONFIDENCE_MASK:3:segmented_mask_3"
|
||||
options {
|
||||
[mediapipe.tasks.TensorsToSegmentationCalculatorOptions.ext] {
|
||||
segmenter_options {
|
||||
activation: SOFTMAX
|
||||
output_type: CONFIDENCE_MASK
|
||||
}
|
||||
segmenter_options { activation: SOFTMAX }
|
||||
}
|
||||
}
|
||||
)pb"));
|
||||
@@ -218,16 +210,13 @@ TEST(TensorsToSegmentationCalculatorTest, SucceedsConfidenceMaskWithNone) {
|
||||
R"pb(
|
||||
calculator: "mediapipe.tasks.TensorsToSegmentationCalculator"
|
||||
input_stream: "TENSORS:tensors"
|
||||
output_stream: "SEGMENTATION:0:segmented_mask_0"
|
||||
output_stream: "SEGMENTATION:1:segmented_mask_1"
|
||||
output_stream: "SEGMENTATION:2:segmented_mask_2"
|
||||
output_stream: "SEGMENTATION:3:segmented_mask_3"
|
||||
output_stream: "CONFIDENCE_MASK:0:segmented_mask_0"
|
||||
output_stream: "CONFIDENCE_MASK:1:segmented_mask_1"
|
||||
output_stream: "CONFIDENCE_MASK:2:segmented_mask_2"
|
||||
output_stream: "CONFIDENCE_MASK:3:segmented_mask_3"
|
||||
options {
|
||||
[mediapipe.tasks.TensorsToSegmentationCalculatorOptions.ext] {
|
||||
segmenter_options {
|
||||
activation: NONE
|
||||
output_type: CONFIDENCE_MASK
|
||||
}
|
||||
segmenter_options { activation: NONE }
|
||||
}
|
||||
}
|
||||
)pb"));
|
||||
@@ -259,16 +248,13 @@ TEST(TensorsToSegmentationCalculatorTest, SucceedsConfidenceMaskWithSigmoid) {
|
||||
R"pb(
|
||||
calculator: "mediapipe.tasks.TensorsToSegmentationCalculator"
|
||||
input_stream: "TENSORS:tensors"
|
||||
output_stream: "SEGMENTATION:0:segmented_mask_0"
|
||||
output_stream: "SEGMENTATION:1:segmented_mask_1"
|
||||
output_stream: "SEGMENTATION:2:segmented_mask_2"
|
||||
output_stream: "SEGMENTATION:3:segmented_mask_3"
|
||||
output_stream: "CONFIDENCE_MASK:0:segmented_mask_0"
|
||||
output_stream: "CONFIDENCE_MASK:1:segmented_mask_1"
|
||||
output_stream: "CONFIDENCE_MASK:2:segmented_mask_2"
|
||||
output_stream: "CONFIDENCE_MASK:3:segmented_mask_3"
|
||||
options {
|
||||
[mediapipe.tasks.TensorsToSegmentationCalculatorOptions.ext] {
|
||||
segmenter_options {
|
||||
activation: SIGMOID
|
||||
output_type: CONFIDENCE_MASK
|
||||
}
|
||||
segmenter_options { activation: SIGMOID }
|
||||
}
|
||||
}
|
||||
)pb"));
|
||||
@@ -301,13 +287,14 @@ TEST(TensorsToSegmentationCalculatorTest, SucceedsCategoryMask) {
|
||||
R"pb(
|
||||
calculator: "mediapipe.tasks.TensorsToSegmentationCalculator"
|
||||
input_stream: "TENSORS:tensors"
|
||||
output_stream: "SEGMENTATION:segmentation"
|
||||
output_stream: "CONFIDENCE_MASK:0:segmented_mask_0"
|
||||
output_stream: "CONFIDENCE_MASK:1:segmented_mask_1"
|
||||
output_stream: "CONFIDENCE_MASK:2:segmented_mask_2"
|
||||
output_stream: "CONFIDENCE_MASK:3:segmented_mask_3"
|
||||
output_stream: "CATEGORY_MASK:segmentation"
|
||||
options {
|
||||
[mediapipe.tasks.TensorsToSegmentationCalculatorOptions.ext] {
|
||||
segmenter_options {
|
||||
activation: NONE
|
||||
output_type: CATEGORY_MASK
|
||||
}
|
||||
segmenter_options { activation: NONE }
|
||||
}
|
||||
}
|
||||
)pb"));
|
||||
@@ -318,11 +305,11 @@ TEST(TensorsToSegmentationCalculatorTest, SucceedsCategoryMask) {
|
||||
tensor_height, tensor_width,
|
||||
std::vector<float>(kTestValues.begin(), kTestValues.end()), &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
ASSERT_EQ(runner.Outputs().NumEntries(), 1);
|
||||
ASSERT_EQ(runner.Outputs().NumEntries(), 5);
|
||||
// Largest element index is 3.
|
||||
const int expected_index = 3;
|
||||
const std::vector<int> buffer_indices = {0};
|
||||
std::vector<Packet> packets = GetPackets(runner);
|
||||
std::vector<Packet> packets = runner.Outputs().Tag("CATEGORY_MASK").packets;
|
||||
EXPECT_THAT(packets, testing::ElementsAre(
|
||||
Uint8ImagePacket(tensor_height, tensor_width,
|
||||
expected_index, buffer_indices)));
|
||||
@@ -335,13 +322,14 @@ TEST(TensorsToSegmentationCalculatorTest, SucceedsCategoryMaskResize) {
|
||||
calculator: "mediapipe.tasks.TensorsToSegmentationCalculator"
|
||||
input_stream: "TENSORS:tensors"
|
||||
input_stream: "OUTPUT_SIZE:size"
|
||||
output_stream: "SEGMENTATION:segmentation"
|
||||
output_stream: "CONFIDENCE_MASK:0:segmented_mask_0"
|
||||
output_stream: "CONFIDENCE_MASK:1:segmented_mask_1"
|
||||
output_stream: "CONFIDENCE_MASK:2:segmented_mask_2"
|
||||
output_stream: "CONFIDENCE_MASK:3:segmented_mask_3"
|
||||
output_stream: "CATEGORY_MASK:segmentation"
|
||||
options {
|
||||
[mediapipe.tasks.TensorsToSegmentationCalculatorOptions.ext] {
|
||||
segmenter_options {
|
||||
activation: NONE
|
||||
output_type: CATEGORY_MASK
|
||||
}
|
||||
segmenter_options { activation: NONE }
|
||||
}
|
||||
}
|
||||
)pb"));
|
||||
@@ -367,7 +355,7 @@ TEST(TensorsToSegmentationCalculatorTest, SucceedsCategoryMaskResize) {
|
||||
const std::vector<int> buffer_indices = {
|
||||
0 * output_width + 0, 0 * output_width + 1, 1 * output_width + 0,
|
||||
1 * output_width + 1};
|
||||
std::vector<Packet> packets = GetPackets(runner);
|
||||
std::vector<Packet> packets = runner.Outputs().Tag("CATEGORY_MASK").packets;
|
||||
EXPECT_THAT(packets, testing::ElementsAre(
|
||||
Uint8ImagePacket(output_height, output_width,
|
||||
expected_index, buffer_indices)));
|
||||
|
||||
@@ -37,8 +37,10 @@ namespace vision {
|
||||
namespace image_segmenter {
|
||||
namespace {
|
||||
|
||||
constexpr char kSegmentationStreamName[] = "segmented_mask_out";
|
||||
constexpr char kGroupedSegmentationTag[] = "GROUPED_SEGMENTATION";
|
||||
constexpr char kConfidenceMasksTag[] = "CONFIDENCE_MASKS";
|
||||
constexpr char kConfidenceMasksStreamName[] = "confidence_masks";
|
||||
constexpr char kCategoryMaskTag[] = "CATEGORY_MASK";
|
||||
constexpr char kCategoryMaskStreamName[] = "category_mask";
|
||||
constexpr char kImageInStreamName[] = "image_in";
|
||||
constexpr char kImageOutStreamName[] = "image_out";
|
||||
constexpr char kImageTag[] = "IMAGE";
|
||||
@@ -51,7 +53,6 @@ constexpr int kMicroSecondsPerMilliSecond = 1000;
|
||||
using ::mediapipe::CalculatorGraphConfig;
|
||||
using ::mediapipe::Image;
|
||||
using ::mediapipe::NormalizedRect;
|
||||
using ::mediapipe::tasks::vision::image_segmenter::proto::SegmenterOptions;
|
||||
using ImageSegmenterGraphOptionsProto = ::mediapipe::tasks::vision::
|
||||
image_segmenter::proto::ImageSegmenterGraphOptions;
|
||||
|
||||
@@ -59,21 +60,24 @@ using ImageSegmenterGraphOptionsProto = ::mediapipe::tasks::vision::
|
||||
// "mediapipe.tasks.vision.image_segmenter.ImageSegmenterGraph".
|
||||
CalculatorGraphConfig CreateGraphConfig(
|
||||
std::unique_ptr<ImageSegmenterGraphOptionsProto> options,
|
||||
bool enable_flow_limiting) {
|
||||
bool output_category_mask, bool enable_flow_limiting) {
|
||||
api2::builder::Graph graph;
|
||||
auto& task_subgraph = graph.AddNode(kSubgraphTypeName);
|
||||
task_subgraph.GetOptions<ImageSegmenterGraphOptionsProto>().Swap(
|
||||
options.get());
|
||||
graph.In(kImageTag).SetName(kImageInStreamName);
|
||||
graph.In(kNormRectTag).SetName(kNormRectStreamName);
|
||||
task_subgraph.Out(kGroupedSegmentationTag).SetName(kSegmentationStreamName) >>
|
||||
graph.Out(kGroupedSegmentationTag);
|
||||
task_subgraph.Out(kConfidenceMasksTag).SetName(kConfidenceMasksStreamName) >>
|
||||
graph.Out(kConfidenceMasksTag);
|
||||
if (output_category_mask) {
|
||||
task_subgraph.Out(kCategoryMaskTag).SetName(kCategoryMaskStreamName) >>
|
||||
graph.Out(kCategoryMaskTag);
|
||||
}
|
||||
task_subgraph.Out(kImageTag).SetName(kImageOutStreamName) >>
|
||||
graph.Out(kImageTag);
|
||||
if (enable_flow_limiting) {
|
||||
return tasks::core::AddFlowLimiterCalculator(graph, task_subgraph,
|
||||
{kImageTag, kNormRectTag},
|
||||
kGroupedSegmentationTag);
|
||||
return tasks::core::AddFlowLimiterCalculator(
|
||||
graph, task_subgraph, {kImageTag, kNormRectTag}, kConfidenceMasksTag);
|
||||
}
|
||||
graph.In(kImageTag) >> task_subgraph.In(kImageTag);
|
||||
graph.In(kNormRectTag) >> task_subgraph.In(kNormRectTag);
|
||||
@@ -91,16 +95,6 @@ ConvertImageSegmenterOptionsToProto(ImageSegmenterOptions* options) {
|
||||
options_proto->mutable_base_options()->set_use_stream_mode(
|
||||
options->running_mode != core::RunningMode::IMAGE);
|
||||
options_proto->set_display_names_locale(options->display_names_locale);
|
||||
switch (options->output_type) {
|
||||
case ImageSegmenterOptions::OutputType::CATEGORY_MASK:
|
||||
options_proto->mutable_segmenter_options()->set_output_type(
|
||||
SegmenterOptions::CATEGORY_MASK);
|
||||
break;
|
||||
case ImageSegmenterOptions::OutputType::CONFIDENCE_MASK:
|
||||
options_proto->mutable_segmenter_options()->set_output_type(
|
||||
SegmenterOptions::CONFIDENCE_MASK);
|
||||
break;
|
||||
}
|
||||
return options_proto;
|
||||
}
|
||||
|
||||
@@ -145,6 +139,7 @@ absl::StatusOr<std::unique_ptr<ImageSegmenter>> ImageSegmenter::Create(
|
||||
tasks::core::PacketsCallback packets_callback = nullptr;
|
||||
if (options->result_callback) {
|
||||
auto result_callback = options->result_callback;
|
||||
bool output_category_mask = options->output_category_mask;
|
||||
packets_callback =
|
||||
[=](absl::StatusOr<tasks::core::PacketMap> status_or_packets) {
|
||||
if (!status_or_packets.ok()) {
|
||||
@@ -156,34 +151,41 @@ absl::StatusOr<std::unique_ptr<ImageSegmenter>> ImageSegmenter::Create(
|
||||
if (status_or_packets.value()[kImageOutStreamName].IsEmpty()) {
|
||||
return;
|
||||
}
|
||||
Packet segmented_masks =
|
||||
status_or_packets.value()[kSegmentationStreamName];
|
||||
Packet confidence_masks =
|
||||
status_or_packets.value()[kConfidenceMasksStreamName];
|
||||
std::optional<Image> category_mask;
|
||||
if (output_category_mask) {
|
||||
category_mask =
|
||||
status_or_packets.value()[kCategoryMaskStreamName].Get<Image>();
|
||||
}
|
||||
Packet image_packet = status_or_packets.value()[kImageOutStreamName];
|
||||
result_callback(segmented_masks.Get<std::vector<Image>>(),
|
||||
image_packet.Get<Image>(),
|
||||
segmented_masks.Timestamp().Value() /
|
||||
kMicroSecondsPerMilliSecond);
|
||||
result_callback(
|
||||
{{confidence_masks.Get<std::vector<Image>>(), category_mask}},
|
||||
image_packet.Get<Image>(),
|
||||
confidence_masks.Timestamp().Value() /
|
||||
kMicroSecondsPerMilliSecond);
|
||||
};
|
||||
}
|
||||
|
||||
auto image_segmenter =
|
||||
core::VisionTaskApiFactory::Create<ImageSegmenter,
|
||||
ImageSegmenterGraphOptionsProto>(
|
||||
CreateGraphConfig(
|
||||
std::move(options_proto),
|
||||
std::move(options_proto), options->output_category_mask,
|
||||
options->running_mode == core::RunningMode::LIVE_STREAM),
|
||||
std::move(options->base_options.op_resolver), options->running_mode,
|
||||
std::move(packets_callback));
|
||||
if (!image_segmenter.ok()) {
|
||||
return image_segmenter.status();
|
||||
}
|
||||
image_segmenter.value()->output_category_mask_ =
|
||||
options->output_category_mask;
|
||||
ASSIGN_OR_RETURN(
|
||||
(*image_segmenter)->labels_,
|
||||
GetLabelsFromGraphConfig((*image_segmenter)->runner_->GetGraphConfig()));
|
||||
return image_segmenter;
|
||||
}
|
||||
|
||||
absl::StatusOr<std::vector<Image>> ImageSegmenter::Segment(
|
||||
absl::StatusOr<ImageSegmenterResult> ImageSegmenter::Segment(
|
||||
mediapipe::Image image,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options) {
|
||||
if (image.UsesGpu()) {
|
||||
@@ -192,20 +194,26 @@ absl::StatusOr<std::vector<Image>> ImageSegmenter::Segment(
|
||||
absl::StrCat("GPU input images are currently not supported."),
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessImageData(
|
||||
{{kImageInStreamName, mediapipe::MakePacket<Image>(std::move(image))},
|
||||
{kNormRectStreamName,
|
||||
MakePacket<NormalizedRect>(std::move(norm_rect))}}));
|
||||
return output_packets[kSegmentationStreamName].Get<std::vector<Image>>();
|
||||
std::vector<Image> confidence_masks =
|
||||
output_packets[kConfidenceMasksStreamName].Get<std::vector<Image>>();
|
||||
std::optional<Image> category_mask;
|
||||
if (output_category_mask_) {
|
||||
category_mask = output_packets[kCategoryMaskStreamName].Get<Image>();
|
||||
}
|
||||
return {{confidence_masks, category_mask}};
|
||||
}
|
||||
|
||||
absl::StatusOr<std::vector<Image>> ImageSegmenter::SegmentForVideo(
|
||||
mediapipe::Image image, int64 timestamp_ms,
|
||||
absl::StatusOr<ImageSegmenterResult> ImageSegmenter::SegmentForVideo(
|
||||
mediapipe::Image image, int64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options) {
|
||||
if (image.UsesGpu()) {
|
||||
return CreateStatusWithPayload(
|
||||
@@ -213,9 +221,9 @@ absl::StatusOr<std::vector<Image>> ImageSegmenter::SegmentForVideo(
|
||||
absl::StrCat("GPU input images are currently not supported."),
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessVideoData(
|
||||
@@ -225,11 +233,17 @@ absl::StatusOr<std::vector<Image>> ImageSegmenter::SegmentForVideo(
|
||||
{kNormRectStreamName,
|
||||
MakePacket<NormalizedRect>(std::move(norm_rect))
|
||||
.At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}}));
|
||||
return output_packets[kSegmentationStreamName].Get<std::vector<Image>>();
|
||||
std::vector<Image> confidence_masks =
|
||||
output_packets[kConfidenceMasksStreamName].Get<std::vector<Image>>();
|
||||
std::optional<Image> category_mask;
|
||||
if (output_category_mask_) {
|
||||
category_mask = output_packets[kCategoryMaskStreamName].Get<Image>();
|
||||
}
|
||||
return {{confidence_masks, category_mask}};
|
||||
}
|
||||
|
||||
absl::Status ImageSegmenter::SegmentAsync(
|
||||
Image image, int64 timestamp_ms,
|
||||
Image image, int64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options) {
|
||||
if (image.UsesGpu()) {
|
||||
return CreateStatusWithPayload(
|
||||
@@ -237,9 +251,9 @@ absl::Status ImageSegmenter::SegmentAsync(
|
||||
absl::StrCat("GPU input images are currently not supported."),
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
return SendLiveStreamData(
|
||||
{{kImageInStreamName,
|
||||
MakePacket<Image>(std::move(image))
|
||||
|
||||
@@ -26,6 +26,7 @@ limitations under the License.
|
||||
#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"
|
||||
#include "mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_result.h"
|
||||
#include "tensorflow/lite/kernels/register.h"
|
||||
|
||||
namespace mediapipe {
|
||||
@@ -52,23 +53,14 @@ struct ImageSegmenterOptions {
|
||||
// Metadata, if any. Defaults to English.
|
||||
std::string display_names_locale = "en";
|
||||
|
||||
// The output type of segmentation results.
|
||||
enum OutputType {
|
||||
// Gives a single output mask where each pixel represents the class which
|
||||
// the pixel in the original image was predicted to belong to.
|
||||
CATEGORY_MASK = 0,
|
||||
// Gives a list of output masks where, for each mask, each pixel represents
|
||||
// the prediction confidence, usually in the [0, 1] range.
|
||||
CONFIDENCE_MASK = 1,
|
||||
};
|
||||
|
||||
OutputType output_type = OutputType::CATEGORY_MASK;
|
||||
// Whether to output category mask.
|
||||
bool output_category_mask = false;
|
||||
|
||||
// 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<std::vector<mediapipe::Image>>,
|
||||
const Image&, int64)>
|
||||
std::function<void(absl::StatusOr<ImageSegmenterResult>, const Image&,
|
||||
int64_t)>
|
||||
result_callback = nullptr;
|
||||
};
|
||||
|
||||
@@ -84,13 +76,9 @@ struct ImageSegmenterOptions {
|
||||
// 1 or 3).
|
||||
// - if type is kTfLiteFloat32, NormalizationOptions are required to be
|
||||
// attached to the metadata for input normalization.
|
||||
// Output tensors:
|
||||
// (kTfLiteUInt8/kTfLiteFloat32)
|
||||
// - list of segmented masks.
|
||||
// - if `output_type` is CATEGORY_MASK, uint8 Image, Image vector of size 1.
|
||||
// - if `output_type` is CONFIDENCE_MASK, float32 Image list of size
|
||||
// `channels`.
|
||||
// - batch is always 1
|
||||
// Output ImageSegmenterResult:
|
||||
// Provides confidence masks and an optional category mask if
|
||||
// `output_category_mask` is set true.
|
||||
// An example of such model can be found at:
|
||||
// https://tfhub.dev/tensorflow/lite-model/deeplabv3/1/metadata/2
|
||||
class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
@@ -114,12 +102,8 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
// setting its 'rotation_degrees' field. Note that specifying a
|
||||
// region-of-interest using the 'region_of_interest' field is NOT supported
|
||||
// and will result in an invalid argument error being returned.
|
||||
//
|
||||
// If the output_type is CATEGORY_MASK, the returned vector of images is
|
||||
// per-category segmented image mask.
|
||||
// If the output_type is CONFIDENCE_MASK, the returned vector of images
|
||||
// contains only one confidence image mask.
|
||||
absl::StatusOr<std::vector<mediapipe::Image>> Segment(
|
||||
|
||||
absl::StatusOr<ImageSegmenterResult> Segment(
|
||||
mediapipe::Image image,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options =
|
||||
std::nullopt);
|
||||
@@ -137,13 +121,8 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
// setting its 'rotation_degrees' field. Note that specifying a
|
||||
// region-of-interest using the 'region_of_interest' field is NOT supported
|
||||
// and will result in an invalid argument error being returned.
|
||||
//
|
||||
// If the output_type is CATEGORY_MASK, the returned vector of images is
|
||||
// per-category segmented image mask.
|
||||
// If the output_type is CONFIDENCE_MASK, the returned vector of images
|
||||
// contains only one confidence image mask.
|
||||
absl::StatusOr<std::vector<mediapipe::Image>> SegmentForVideo(
|
||||
mediapipe::Image image, int64 timestamp_ms,
|
||||
absl::StatusOr<ImageSegmenterResult> SegmentForVideo(
|
||||
mediapipe::Image image, int64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options =
|
||||
std::nullopt);
|
||||
|
||||
@@ -164,17 +143,13 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
// and will result in an invalid argument error being returned.
|
||||
//
|
||||
// The "result_callback" prvoides
|
||||
// - A vector of segmented image masks.
|
||||
// If the output_type is CATEGORY_MASK, the returned vector of images is
|
||||
// per-category segmented image mask.
|
||||
// If the output_type is CONFIDENCE_MASK, the returned vector of images
|
||||
// contains only one confidence image mask.
|
||||
// - An ImageSegmenterResult.
|
||||
// - The const reference to the corresponding input image that the image
|
||||
// segmentation runs on. Note that the const reference to the image will
|
||||
// no longer be valid when the callback returns. To access the image data
|
||||
// outside of the callback, callers need to make a copy of the image.
|
||||
// - The input timestamp in milliseconds.
|
||||
absl::Status SegmentAsync(mediapipe::Image image, int64 timestamp_ms,
|
||||
absl::Status SegmentAsync(mediapipe::Image image, int64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions>
|
||||
image_processing_options = std::nullopt);
|
||||
|
||||
@@ -182,9 +157,9 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
absl::Status Close() { return runner_->Close(); }
|
||||
|
||||
// Get the category label list of the ImageSegmenter can recognize. For
|
||||
// CATEGORY_MASK type, the index in the category mask corresponds to the
|
||||
// category in the label list. For CONFIDENCE_MASK type, the output mask list
|
||||
// at index corresponds to the category in the label list.
|
||||
// CATEGORY_MASK, the index in the category mask corresponds to the category
|
||||
// in the label list. For CONFIDENCE_MASK, the output mask list at index
|
||||
// corresponds to the category in the label list.
|
||||
//
|
||||
// If there is no labelmap provided in the model file, empty label list is
|
||||
// returned.
|
||||
@@ -192,6 +167,7 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
|
||||
private:
|
||||
std::vector<std::string> labels_;
|
||||
bool output_category_mask_;
|
||||
};
|
||||
|
||||
} // namespace image_segmenter
|
||||
|
||||
@@ -14,6 +14,7 @@ limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
#include <type_traits>
|
||||
#include <vector>
|
||||
|
||||
@@ -42,6 +43,7 @@ limitations under the License.
|
||||
#include "mediapipe/tasks/cc/vision/utils/image_tensor_specs.h"
|
||||
#include "mediapipe/tasks/metadata/image_segmenter_metadata_schema_generated.h"
|
||||
#include "mediapipe/tasks/metadata/metadata_schema_generated.h"
|
||||
#include "mediapipe/util/graph_builder_utils.h"
|
||||
#include "mediapipe/util/label_map.pb.h"
|
||||
#include "mediapipe/util/label_map_util.h"
|
||||
#include "tensorflow/lite/schema/schema_generated.h"
|
||||
@@ -65,10 +67,13 @@ using ::mediapipe::tasks::vision::image_segmenter::proto::
|
||||
ImageSegmenterGraphOptions;
|
||||
using ::mediapipe::tasks::vision::image_segmenter::proto::SegmenterOptions;
|
||||
using ::tflite::TensorMetadata;
|
||||
using LabelItems = mediapipe::proto_ns::Map<int64, ::mediapipe::LabelMapItem>;
|
||||
using LabelItems = mediapipe::proto_ns::Map<int64_t, ::mediapipe::LabelMapItem>;
|
||||
|
||||
constexpr char kSegmentationTag[] = "SEGMENTATION";
|
||||
constexpr char kGroupedSegmentationTag[] = "GROUPED_SEGMENTATION";
|
||||
constexpr char kConfidenceMaskTag[] = "CONFIDENCE_MASK";
|
||||
constexpr char kConfidenceMasksTag[] = "CONFIDENCE_MASKS";
|
||||
constexpr char kCategoryMaskTag[] = "CATEGORY_MASK";
|
||||
constexpr char kImageTag[] = "IMAGE";
|
||||
constexpr char kImageCpuTag[] = "IMAGE_CPU";
|
||||
constexpr char kImageGpuTag[] = "IMAGE_GPU";
|
||||
@@ -80,7 +85,9 @@ constexpr char kSegmentationMetadataName[] = "SEGMENTER_METADATA";
|
||||
// Struct holding the different output streams produced by the image segmenter
|
||||
// subgraph.
|
||||
struct ImageSegmenterOutputs {
|
||||
std::vector<Source<Image>> segmented_masks;
|
||||
std::optional<std::vector<Source<Image>>> segmented_masks;
|
||||
std::optional<std::vector<Source<Image>>> confidence_masks;
|
||||
std::optional<Source<Image>> category_mask;
|
||||
// The same as the input image, mainly used for live stream mode.
|
||||
Source<Image> image;
|
||||
};
|
||||
@@ -95,8 +102,10 @@ struct ImageAndTensorsOnDevice {
|
||||
} // namespace
|
||||
|
||||
absl::Status SanityCheckOptions(const ImageSegmenterGraphOptions& options) {
|
||||
if (options.segmenter_options().output_type() ==
|
||||
SegmenterOptions::UNSPECIFIED) {
|
||||
// TODO: remove deprecated output type support.
|
||||
if (options.segmenter_options().has_output_type() &&
|
||||
options.segmenter_options().output_type() ==
|
||||
SegmenterOptions::UNSPECIFIED) {
|
||||
return CreateStatusWithPayload(absl::StatusCode::kInvalidArgument,
|
||||
"`output_type` must not be UNSPECIFIED",
|
||||
MediaPipeTasksStatus::kInvalidArgumentError);
|
||||
@@ -133,9 +142,8 @@ absl::Status ConfigureTensorsToSegmentationCalculator(
|
||||
const core::ModelResources& model_resources,
|
||||
TensorsToSegmentationCalculatorOptions* options) {
|
||||
// Set default activation function NONE
|
||||
options->mutable_segmenter_options()->set_output_type(
|
||||
segmenter_option.segmenter_options().output_type());
|
||||
options->mutable_segmenter_options()->set_activation(SegmenterOptions::NONE);
|
||||
options->mutable_segmenter_options()->CopyFrom(
|
||||
segmenter_option.segmenter_options());
|
||||
// Find the custom metadata of ImageSegmenterOptions type in model metadata.
|
||||
const auto* metadata_extractor = model_resources.GetMetadataExtractor();
|
||||
bool found_activation_in_metadata = false;
|
||||
@@ -317,12 +325,14 @@ absl::StatusOr<ImageAndTensorsOnDevice> ConvertImageToTensors(
|
||||
}
|
||||
}
|
||||
|
||||
// An "mediapipe.tasks.vision.ImageSegmenterGraph" performs semantic
|
||||
// segmentation.
|
||||
// Two kinds of outputs are provided: SEGMENTATION and GROUPED_SEGMENTATION.
|
||||
// Users can retrieve segmented mask of only particular category/channel from
|
||||
// SEGMENTATION, and users can also get all segmented masks from
|
||||
// GROUPED_SEGMENTATION.
|
||||
// An "mediapipe.tasks.vision.image_segmenter.ImageSegmenterGraph" performs
|
||||
// semantic segmentation. The graph always output confidence masks, and an
|
||||
// optional category mask if CATEGORY_MASK is connected.
|
||||
//
|
||||
// Two kinds of outputs for confidence mask are provided: CONFIDENCE_MASK and
|
||||
// CONFIDENCE_MASKS. Users can retrieve segmented mask of only particular
|
||||
// category/channel from CONFIDENCE_MASK, and users can also get all segmented
|
||||
// confidence masks from CONFIDENCE_MASKS.
|
||||
// - Accepts CPU input images and outputs segmented masks on CPU.
|
||||
//
|
||||
// Inputs:
|
||||
@@ -334,11 +344,13 @@ absl::StatusOr<ImageAndTensorsOnDevice> ConvertImageToTensors(
|
||||
// @Optional: rect covering the whole image is used if not specified.
|
||||
//
|
||||
// Outputs:
|
||||
// SEGMENTATION - mediapipe::Image @Multiple
|
||||
// Segmented masks for individual category. Segmented mask of single
|
||||
// CONFIDENCE_MASK - mediapipe::Image @Multiple
|
||||
// Confidence masks for individual category. Confidence mask of single
|
||||
// category can be accessed by index based output stream.
|
||||
// GROUPED_SEGMENTATION - std::vector<mediapipe::Image>
|
||||
// The output segmented masks grouped in a vector.
|
||||
// CONFIDENCE_MASKS - std::vector<mediapipe::Image>
|
||||
// The output confidence masks grouped in a vector.
|
||||
// CATEGORY_MASK - mediapipe::Image @Optional
|
||||
// Optional Category mask.
|
||||
// IMAGE - mediapipe::Image
|
||||
// The image that image segmenter runs on.
|
||||
//
|
||||
@@ -369,23 +381,39 @@ class ImageSegmenterGraph : public core::ModelTaskGraph {
|
||||
ASSIGN_OR_RETURN(const auto* model_resources,
|
||||
CreateModelResources<ImageSegmenterGraphOptions>(sc));
|
||||
Graph graph;
|
||||
const auto& options = sc->Options<ImageSegmenterGraphOptions>();
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_streams,
|
||||
BuildSegmentationTask(
|
||||
sc->Options<ImageSegmenterGraphOptions>(), *model_resources,
|
||||
graph[Input<Image>(kImageTag)],
|
||||
graph[Input<NormalizedRect>::Optional(kNormRectTag)], graph));
|
||||
options, *model_resources, graph[Input<Image>(kImageTag)],
|
||||
graph[Input<NormalizedRect>::Optional(kNormRectTag)],
|
||||
HasOutput(sc->OriginalNode(), kCategoryMaskTag), graph));
|
||||
|
||||
auto& merge_images_to_vector =
|
||||
graph.AddNode("MergeImagesToVectorCalculator");
|
||||
for (int i = 0; i < output_streams.segmented_masks.size(); ++i) {
|
||||
output_streams.segmented_masks[i] >>
|
||||
merge_images_to_vector[Input<Image>::Multiple("")][i];
|
||||
output_streams.segmented_masks[i] >>
|
||||
graph[Output<Image>::Multiple(kSegmentationTag)][i];
|
||||
// TODO: remove deprecated output type support.
|
||||
if (options.segmenter_options().has_output_type()) {
|
||||
for (int i = 0; i < output_streams.segmented_masks->size(); ++i) {
|
||||
output_streams.segmented_masks->at(i) >>
|
||||
merge_images_to_vector[Input<Image>::Multiple("")][i];
|
||||
output_streams.segmented_masks->at(i) >>
|
||||
graph[Output<Image>::Multiple(kSegmentationTag)][i];
|
||||
}
|
||||
merge_images_to_vector.Out("") >>
|
||||
graph[Output<std::vector<Image>>(kGroupedSegmentationTag)];
|
||||
} else {
|
||||
for (int i = 0; i < output_streams.confidence_masks->size(); ++i) {
|
||||
output_streams.confidence_masks->at(i) >>
|
||||
merge_images_to_vector[Input<Image>::Multiple("")][i];
|
||||
output_streams.confidence_masks->at(i) >>
|
||||
graph[Output<Image>::Multiple(kConfidenceMaskTag)][i];
|
||||
}
|
||||
merge_images_to_vector.Out("") >>
|
||||
graph[Output<std::vector<Image>>(kConfidenceMasksTag)];
|
||||
if (output_streams.category_mask) {
|
||||
*output_streams.category_mask >> graph[Output<Image>(kCategoryMaskTag)];
|
||||
}
|
||||
}
|
||||
merge_images_to_vector.Out("") >>
|
||||
graph[Output<std::vector<Image>>(kGroupedSegmentationTag)];
|
||||
output_streams.image >> graph[Output<Image>(kImageTag)];
|
||||
return graph.GetConfig();
|
||||
}
|
||||
@@ -403,7 +431,8 @@ class ImageSegmenterGraph : public core::ModelTaskGraph {
|
||||
absl::StatusOr<ImageSegmenterOutputs> BuildSegmentationTask(
|
||||
const ImageSegmenterGraphOptions& task_options,
|
||||
const core::ModelResources& model_resources, Source<Image> image_in,
|
||||
Source<NormalizedRect> norm_rect_in, Graph& graph) {
|
||||
Source<NormalizedRect> norm_rect_in, bool output_category_mask,
|
||||
Graph& graph) {
|
||||
MP_RETURN_IF_ERROR(SanityCheckOptions(task_options));
|
||||
|
||||
// Adds preprocessing calculators and connects them to the graph input image
|
||||
@@ -435,22 +464,46 @@ class ImageSegmenterGraph : public core::ModelTaskGraph {
|
||||
image_properties.Out("SIZE") >> tensor_to_images.In(kOutputSizeTag);
|
||||
|
||||
// Exports multiple segmented masks.
|
||||
std::vector<Source<Image>> segmented_masks;
|
||||
if (task_options.segmenter_options().output_type() ==
|
||||
SegmenterOptions::CATEGORY_MASK) {
|
||||
segmented_masks.push_back(
|
||||
Source<Image>(tensor_to_images[Output<Image>(kSegmentationTag)]));
|
||||
// TODO: remove deprecated output type support.
|
||||
if (task_options.segmenter_options().has_output_type()) {
|
||||
std::vector<Source<Image>> segmented_masks;
|
||||
if (task_options.segmenter_options().output_type() ==
|
||||
SegmenterOptions::CATEGORY_MASK) {
|
||||
segmented_masks.push_back(
|
||||
Source<Image>(tensor_to_images[Output<Image>(kSegmentationTag)]));
|
||||
} else {
|
||||
ASSIGN_OR_RETURN(const tflite::Tensor* output_tensor,
|
||||
GetOutputTensor(model_resources));
|
||||
int segmentation_streams_num = *output_tensor->shape()->rbegin();
|
||||
for (int i = 0; i < segmentation_streams_num; ++i) {
|
||||
segmented_masks.push_back(Source<Image>(
|
||||
tensor_to_images[Output<Image>::Multiple(kSegmentationTag)][i]));
|
||||
}
|
||||
}
|
||||
return ImageSegmenterOutputs{/*segmented_masks=*/segmented_masks,
|
||||
/*confidence_masks=*/std::nullopt,
|
||||
/*category_mask=*/std::nullopt,
|
||||
/*image=*/image_and_tensors.image};
|
||||
} else {
|
||||
ASSIGN_OR_RETURN(const tflite::Tensor* output_tensor,
|
||||
GetOutputTensor(model_resources));
|
||||
int segmentation_streams_num = *output_tensor->shape()->rbegin();
|
||||
std::vector<Source<Image>> confidence_masks;
|
||||
confidence_masks.reserve(segmentation_streams_num);
|
||||
for (int i = 0; i < segmentation_streams_num; ++i) {
|
||||
segmented_masks.push_back(Source<Image>(
|
||||
tensor_to_images[Output<Image>::Multiple(kSegmentationTag)][i]));
|
||||
confidence_masks.push_back(Source<Image>(
|
||||
tensor_to_images[Output<Image>::Multiple(kConfidenceMaskTag)][i]));
|
||||
}
|
||||
return ImageSegmenterOutputs{
|
||||
/*segmented_masks=*/std::nullopt,
|
||||
/*confidence_masks=*/confidence_masks,
|
||||
/*category_mask=*/
|
||||
output_category_mask
|
||||
? std::make_optional(
|
||||
tensor_to_images[Output<Image>(kCategoryMaskTag)])
|
||||
: std::nullopt,
|
||||
/*image=*/image_and_tensors.image};
|
||||
}
|
||||
return ImageSegmenterOutputs{/*segmented_masks=*/segmented_masks,
|
||||
/*image=*/image_and_tensors.image};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
/* Copyright 2023 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_VISION_IMAGE_SEGMENTER_IMAGE_SEGMENTER_RESULT_H_
|
||||
#define MEDIAPIPE_TASKS_CC_VISION_IMAGE_SEGMENTER_IMAGE_SEGMENTER_RESULT_H_
|
||||
|
||||
#include <optional>
|
||||
|
||||
#include "mediapipe/framework/formats/image.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace tasks {
|
||||
namespace vision {
|
||||
namespace image_segmenter {
|
||||
|
||||
// The output result of ImageSegmenter
|
||||
struct ImageSegmenterResult {
|
||||
// Multiple masks of float image in VEC32F1 format where, for each mask, each
|
||||
// pixel represents the prediction confidence, usually in the [0, 1] range.
|
||||
std::vector<Image> confidence_masks;
|
||||
// A category mask of uint8 image in GRAY8 format where each pixel represents
|
||||
// the class which the pixel in the original image was predicted to belong to.
|
||||
std::optional<Image> category_mask;
|
||||
};
|
||||
|
||||
} // namespace image_segmenter
|
||||
} // namespace vision
|
||||
} // namespace tasks
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_TASKS_CC_VISION_IMAGE_SEGMENTER_IMAGE_SEGMENTER_RESULT_H_
|
||||
@@ -36,6 +36,7 @@ limitations under the License.
|
||||
#include "mediapipe/tasks/cc/core/proto/external_file.pb.h"
|
||||
#include "mediapipe/tasks/cc/vision/core/image_processing_options.h"
|
||||
#include "mediapipe/tasks/cc/vision/image_segmenter/calculators/tensors_to_segmentation_calculator.pb.h"
|
||||
#include "mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_result.h"
|
||||
#include "mediapipe/tasks/cc/vision/image_segmenter/proto/image_segmenter_graph_options.pb.h"
|
||||
#include "mediapipe/tasks/cc/vision/utils/image_utils.h"
|
||||
#include "tensorflow/lite/core/shims/cc/shims_test_util.h"
|
||||
@@ -256,7 +257,6 @@ TEST(GetLabelsTest, SucceedsWithLabelsInModel) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CATEGORY_MASK;
|
||||
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
@@ -278,15 +278,14 @@ TEST_F(ImageModeTest, SucceedsWithCategoryMask) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CATEGORY_MASK;
|
||||
|
||||
options->output_category_mask = true;
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto category_masks, segmenter->Segment(image));
|
||||
EXPECT_EQ(category_masks.size(), 1);
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image));
|
||||
EXPECT_TRUE(result.category_mask.has_value());
|
||||
|
||||
cv::Mat actual_mask = mediapipe::formats::MatView(
|
||||
category_masks[0].GetImageFrameSharedPtr().get());
|
||||
result.category_mask->GetImageFrameSharedPtr().get());
|
||||
|
||||
cv::Mat expected_mask = cv::imread(
|
||||
JoinPath("./", kTestDataDirectory, "segmentation_golden_rotation0.png"),
|
||||
@@ -303,12 +302,11 @@ TEST_F(ImageModeTest, SucceedsWithConfidenceMask) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CONFIDENCE_MASK;
|
||||
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto confidence_masks, segmenter->Segment(image));
|
||||
EXPECT_EQ(confidence_masks.size(), 21);
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image));
|
||||
EXPECT_EQ(result.confidence_masks.size(), 21);
|
||||
|
||||
cv::Mat expected_mask = cv::imread(
|
||||
JoinPath("./", kTestDataDirectory, "cat_mask.jpg"), cv::IMREAD_GRAYSCALE);
|
||||
@@ -317,7 +315,7 @@ TEST_F(ImageModeTest, SucceedsWithConfidenceMask) {
|
||||
|
||||
// Cat category index 8.
|
||||
cv::Mat cat_mask = mediapipe::formats::MatView(
|
||||
confidence_masks[8].GetImageFrameSharedPtr().get());
|
||||
result.confidence_masks[8].GetImageFrameSharedPtr().get());
|
||||
EXPECT_THAT(cat_mask,
|
||||
SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity));
|
||||
}
|
||||
@@ -331,15 +329,14 @@ TEST_F(ImageModeTest, DISABLED_SucceedsWithRotation) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CONFIDENCE_MASK;
|
||||
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
ImageProcessingOptions image_processing_options;
|
||||
image_processing_options.rotation_degrees = -90;
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto confidence_masks,
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto result,
|
||||
segmenter->Segment(image, image_processing_options));
|
||||
EXPECT_EQ(confidence_masks.size(), 21);
|
||||
EXPECT_EQ(result.confidence_masks.size(), 21);
|
||||
|
||||
cv::Mat expected_mask =
|
||||
cv::imread(JoinPath("./", kTestDataDirectory, "cat_rotated_mask.jpg"),
|
||||
@@ -349,7 +346,7 @@ TEST_F(ImageModeTest, DISABLED_SucceedsWithRotation) {
|
||||
|
||||
// Cat category index 8.
|
||||
cv::Mat cat_mask = mediapipe::formats::MatView(
|
||||
confidence_masks[8].GetImageFrameSharedPtr().get());
|
||||
result.confidence_masks[8].GetImageFrameSharedPtr().get());
|
||||
EXPECT_THAT(cat_mask,
|
||||
SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity));
|
||||
}
|
||||
@@ -361,7 +358,6 @@ TEST_F(ImageModeTest, FailsWithRegionOfInterest) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CONFIDENCE_MASK;
|
||||
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
@@ -384,12 +380,11 @@ TEST_F(ImageModeTest, SucceedsSelfie128x128Segmentation) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kSelfie128x128WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CONFIDENCE_MASK;
|
||||
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto confidence_masks, segmenter->Segment(image));
|
||||
EXPECT_EQ(confidence_masks.size(), 2);
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image));
|
||||
EXPECT_EQ(result.confidence_masks.size(), 2);
|
||||
|
||||
cv::Mat expected_mask =
|
||||
cv::imread(JoinPath("./", kTestDataDirectory,
|
||||
@@ -400,7 +395,7 @@ TEST_F(ImageModeTest, SucceedsSelfie128x128Segmentation) {
|
||||
|
||||
// Selfie category index 1.
|
||||
cv::Mat selfie_mask = mediapipe::formats::MatView(
|
||||
confidence_masks[1].GetImageFrameSharedPtr().get());
|
||||
result.confidence_masks[1].GetImageFrameSharedPtr().get());
|
||||
EXPECT_THAT(selfie_mask,
|
||||
SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity));
|
||||
}
|
||||
@@ -411,11 +406,10 @@ TEST_F(ImageModeTest, SucceedsSelfie144x256Segmentations) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kSelfie144x256WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CONFIDENCE_MASK;
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto confidence_masks, segmenter->Segment(image));
|
||||
EXPECT_EQ(confidence_masks.size(), 1);
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image));
|
||||
EXPECT_EQ(result.confidence_masks.size(), 1);
|
||||
|
||||
cv::Mat expected_mask =
|
||||
cv::imread(JoinPath("./", kTestDataDirectory,
|
||||
@@ -425,7 +419,7 @@ TEST_F(ImageModeTest, SucceedsSelfie144x256Segmentations) {
|
||||
expected_mask.convertTo(expected_mask_float, CV_32FC1, 1 / 255.f);
|
||||
|
||||
cv::Mat selfie_mask = mediapipe::formats::MatView(
|
||||
confidence_masks[0].GetImageFrameSharedPtr().get());
|
||||
result.confidence_masks[0].GetImageFrameSharedPtr().get());
|
||||
EXPECT_THAT(selfie_mask,
|
||||
SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity));
|
||||
}
|
||||
@@ -436,12 +430,11 @@ TEST_F(ImageModeTest, SucceedsPortraitSelfieSegmentationConfidenceMask) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kSelfieSegmentation);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CONFIDENCE_MASK;
|
||||
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto confidence_masks, segmenter->Segment(image));
|
||||
EXPECT_EQ(confidence_masks.size(), 1);
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image));
|
||||
EXPECT_EQ(result.confidence_masks.size(), 1);
|
||||
MP_ASSERT_OK(segmenter->Close());
|
||||
|
||||
cv::Mat expected_mask = cv::imread(
|
||||
@@ -452,7 +445,7 @@ TEST_F(ImageModeTest, SucceedsPortraitSelfieSegmentationConfidenceMask) {
|
||||
expected_mask.convertTo(expected_mask_float, CV_32FC1, 1 / 255.f);
|
||||
|
||||
cv::Mat selfie_mask = mediapipe::formats::MatView(
|
||||
confidence_masks[0].GetImageFrameSharedPtr().get());
|
||||
result.confidence_masks[0].GetImageFrameSharedPtr().get());
|
||||
EXPECT_THAT(selfie_mask,
|
||||
SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity));
|
||||
}
|
||||
@@ -463,16 +456,15 @@ TEST_F(ImageModeTest, SucceedsPortraitSelfieSegmentationCategoryMask) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kSelfieSegmentation);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CATEGORY_MASK;
|
||||
|
||||
options->output_category_mask = true;
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto category_mask, segmenter->Segment(image));
|
||||
EXPECT_EQ(category_mask.size(), 1);
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image));
|
||||
EXPECT_TRUE(result.category_mask.has_value());
|
||||
MP_ASSERT_OK(segmenter->Close());
|
||||
|
||||
cv::Mat selfie_mask = mediapipe::formats::MatView(
|
||||
category_mask[0].GetImageFrameSharedPtr().get());
|
||||
result.category_mask->GetImageFrameSharedPtr().get());
|
||||
cv::Mat expected_mask = cv::imread(
|
||||
JoinPath("./", kTestDataDirectory,
|
||||
"portrait_selfie_segmentation_expected_category_mask.jpg"),
|
||||
@@ -487,16 +479,15 @@ TEST_F(ImageModeTest, SucceedsPortraitSelfieSegmentationLandscapeCategoryMask) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kSelfieSegmentationLandscape);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CATEGORY_MASK;
|
||||
|
||||
options->output_category_mask = true;
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto category_mask, segmenter->Segment(image));
|
||||
EXPECT_EQ(category_mask.size(), 1);
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image));
|
||||
EXPECT_TRUE(result.category_mask.has_value());
|
||||
MP_ASSERT_OK(segmenter->Close());
|
||||
|
||||
cv::Mat selfie_mask = mediapipe::formats::MatView(
|
||||
category_mask[0].GetImageFrameSharedPtr().get());
|
||||
result.category_mask->GetImageFrameSharedPtr().get());
|
||||
cv::Mat expected_mask = cv::imread(
|
||||
JoinPath(
|
||||
"./", kTestDataDirectory,
|
||||
@@ -512,14 +503,13 @@ TEST_F(ImageModeTest, SucceedsHairSegmentation) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kHairSegmentationWithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CONFIDENCE_MASK;
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto confidence_masks, segmenter->Segment(image));
|
||||
EXPECT_EQ(confidence_masks.size(), 2);
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image));
|
||||
EXPECT_EQ(result.confidence_masks.size(), 2);
|
||||
|
||||
cv::Mat hair_mask = mediapipe::formats::MatView(
|
||||
confidence_masks[1].GetImageFrameSharedPtr().get());
|
||||
result.confidence_masks[1].GetImageFrameSharedPtr().get());
|
||||
MP_ASSERT_OK(segmenter->Close());
|
||||
cv::Mat expected_mask = cv::imread(
|
||||
JoinPath("./", kTestDataDirectory, "portrait_hair_expected_mask.jpg"),
|
||||
@@ -540,7 +530,6 @@ TEST_F(VideoModeTest, FailsWithCallingWrongMethod) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CATEGORY_MASK;
|
||||
options->running_mode = core::RunningMode::VIDEO;
|
||||
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
@@ -572,7 +561,7 @@ TEST_F(VideoModeTest, Succeeds) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CATEGORY_MASK;
|
||||
options->output_category_mask = true;
|
||||
options->running_mode = core::RunningMode::VIDEO;
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
@@ -580,11 +569,10 @@ TEST_F(VideoModeTest, Succeeds) {
|
||||
JoinPath("./", kTestDataDirectory, "segmentation_golden_rotation0.png"),
|
||||
cv::IMREAD_GRAYSCALE);
|
||||
for (int i = 0; i < iterations; ++i) {
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto category_masks,
|
||||
segmenter->SegmentForVideo(image, i));
|
||||
EXPECT_EQ(category_masks.size(), 1);
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->SegmentForVideo(image, i));
|
||||
EXPECT_TRUE(result.category_mask.has_value());
|
||||
cv::Mat actual_mask = mediapipe::formats::MatView(
|
||||
category_masks[0].GetImageFrameSharedPtr().get());
|
||||
result.category_mask->GetImageFrameSharedPtr().get());
|
||||
EXPECT_THAT(actual_mask,
|
||||
SimilarToUint8Mask(expected_mask, kGoldenMaskSimilarity,
|
||||
kGoldenMaskMagnificationFactor));
|
||||
@@ -601,11 +589,10 @@ TEST_F(LiveStreamModeTest, FailsWithCallingWrongMethod) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CATEGORY_MASK;
|
||||
options->running_mode = core::RunningMode::LIVE_STREAM;
|
||||
options->result_callback =
|
||||
[](absl::StatusOr<std::vector<Image>> segmented_masks, const Image& image,
|
||||
int64 timestamp_ms) {};
|
||||
[](absl::StatusOr<ImageSegmenterResult> segmented_masks,
|
||||
const Image& image, int64_t timestamp_ms) {};
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
|
||||
@@ -634,11 +621,9 @@ TEST_F(LiveStreamModeTest, FailsWithOutOfOrderInputTimestamps) {
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CATEGORY_MASK;
|
||||
options->running_mode = core::RunningMode::LIVE_STREAM;
|
||||
options->result_callback =
|
||||
[](absl::StatusOr<std::vector<Image>> segmented_masks, const Image& image,
|
||||
int64 timestamp_ms) {};
|
||||
options->result_callback = [](absl::StatusOr<ImageSegmenterResult> result,
|
||||
const Image& image, int64_t timestamp_ms) {};
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
MP_ASSERT_OK(segmenter->SegmentAsync(image, 1));
|
||||
@@ -660,23 +645,23 @@ TEST_F(LiveStreamModeTest, Succeeds) {
|
||||
Image image,
|
||||
DecodeImageFromFile(JoinPath("./", kTestDataDirectory,
|
||||
"segmentation_input_rotation0.jpg")));
|
||||
std::vector<std::vector<Image>> segmented_masks_results;
|
||||
std::vector<Image> segmented_masks_results;
|
||||
std::vector<std::pair<int, int>> image_sizes;
|
||||
std::vector<int64> timestamps;
|
||||
std::vector<int64_t> timestamps;
|
||||
auto options = std::make_unique<ImageSegmenterOptions>();
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata);
|
||||
options->output_type = ImageSegmenterOptions::OutputType::CATEGORY_MASK;
|
||||
options->output_category_mask = true;
|
||||
options->running_mode = core::RunningMode::LIVE_STREAM;
|
||||
options->result_callback =
|
||||
[&segmented_masks_results, &image_sizes, ×tamps](
|
||||
absl::StatusOr<std::vector<Image>> segmented_masks,
|
||||
const Image& image, int64 timestamp_ms) {
|
||||
MP_ASSERT_OK(segmented_masks.status());
|
||||
segmented_masks_results.push_back(std::move(segmented_masks).value());
|
||||
image_sizes.push_back({image.width(), image.height()});
|
||||
timestamps.push_back(timestamp_ms);
|
||||
};
|
||||
options->result_callback = [&segmented_masks_results, &image_sizes,
|
||||
×tamps](
|
||||
absl::StatusOr<ImageSegmenterResult> result,
|
||||
const Image& image, int64_t timestamp_ms) {
|
||||
MP_ASSERT_OK(result.status());
|
||||
segmented_masks_results.push_back(std::move(*result->category_mask));
|
||||
image_sizes.push_back({image.width(), image.height()});
|
||||
timestamps.push_back(timestamp_ms);
|
||||
};
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter,
|
||||
ImageSegmenter::Create(std::move(options)));
|
||||
for (int i = 0; i < iterations; ++i) {
|
||||
@@ -690,10 +675,9 @@ TEST_F(LiveStreamModeTest, Succeeds) {
|
||||
cv::Mat expected_mask = cv::imread(
|
||||
JoinPath("./", kTestDataDirectory, "segmentation_golden_rotation0.png"),
|
||||
cv::IMREAD_GRAYSCALE);
|
||||
for (const auto& segmented_masks : segmented_masks_results) {
|
||||
EXPECT_EQ(segmented_masks.size(), 1);
|
||||
for (const auto& category_mask : segmented_masks_results) {
|
||||
cv::Mat actual_mask = mediapipe::formats::MatView(
|
||||
segmented_masks[0].GetImageFrameSharedPtr().get());
|
||||
category_mask.GetImageFrameSharedPtr().get());
|
||||
EXPECT_THAT(actual_mask,
|
||||
SimilarToUint8Mask(expected_mask, kGoldenMaskSimilarity,
|
||||
kGoldenMaskMagnificationFactor));
|
||||
@@ -702,7 +686,7 @@ TEST_F(LiveStreamModeTest, Succeeds) {
|
||||
EXPECT_EQ(image_size.first, image.width());
|
||||
EXPECT_EQ(image_size.second, image.height());
|
||||
}
|
||||
int64 timestamp_ms = -1;
|
||||
int64_t timestamp_ms = -1;
|
||||
for (const auto& timestamp : timestamps) {
|
||||
EXPECT_GT(timestamp, timestamp_ms);
|
||||
timestamp_ms = timestamp;
|
||||
|
||||
@@ -33,7 +33,7 @@ message SegmenterOptions {
|
||||
CONFIDENCE_MASK = 2;
|
||||
}
|
||||
// Optional output mask type.
|
||||
optional OutputType output_type = 1 [default = CATEGORY_MASK];
|
||||
optional OutputType output_type = 1 [deprecated = true];
|
||||
|
||||
// Supported activation functions for filtering.
|
||||
enum Activation {
|
||||
|
||||
@@ -46,7 +46,6 @@ cc_library(
|
||||
name = "interactive_segmenter_graph",
|
||||
srcs = ["interactive_segmenter_graph.cc"],
|
||||
deps = [
|
||||
"@com_google_absl//absl/strings",
|
||||
"//mediapipe/calculators/image:set_alpha_calculator",
|
||||
"//mediapipe/calculators/util:annotation_overlay_calculator",
|
||||
"//mediapipe/calculators/util:flat_color_image_calculator",
|
||||
@@ -65,6 +64,7 @@ cc_library(
|
||||
"//mediapipe/util:color_cc_proto",
|
||||
"//mediapipe/util:label_map_cc_proto",
|
||||
"//mediapipe/util:render_data_cc_proto",
|
||||
"@com_google_absl//absl/strings",
|
||||
] + select({
|
||||
"//mediapipe/gpu:disable_gpu": [],
|
||||
"//conditions:default": [
|
||||
|
||||
@@ -142,9 +142,9 @@ absl::StatusOr<std::vector<Image>> InteractiveSegmenter::Segment(
|
||||
absl::StrCat("GPU input images are currently not supported."),
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(RenderData roi_as_render_data, ConvertRoiToRenderData(roi));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
|
||||
@@ -157,9 +157,9 @@ absl::StatusOr<ObjectDetectorResult> ObjectDetector::Detect(
|
||||
absl::StrCat("GPU input images are currently not supported."),
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessImageData(
|
||||
@@ -178,9 +178,9 @@ absl::StatusOr<ObjectDetectorResult> ObjectDetector::DetectForVideo(
|
||||
absl::StrCat("GPU input images are currently not supported."),
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(
|
||||
auto output_packets,
|
||||
ProcessVideoData(
|
||||
@@ -203,9 +203,9 @@ absl::Status ObjectDetector::DetectAsync(
|
||||
absl::StrCat("GPU input images are currently not supported."),
|
||||
MediaPipeTasksStatus::kRunnerUnexpectedInputError);
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, /*roi_allowed=*/false));
|
||||
ASSIGN_OR_RETURN(NormalizedRect norm_rect,
|
||||
ConvertToNormalizedRect(image_processing_options, image,
|
||||
/*roi_allowed=*/false));
|
||||
return SendLiveStreamData(
|
||||
{{kImageInStreamName,
|
||||
MakePacket<Image>(std::move(image))
|
||||
|
||||
@@ -575,7 +575,6 @@ TEST_F(ImageModeTest, SucceedsWithRotation) {
|
||||
"cats_and_dogs_rotated.jpg")));
|
||||
auto options = std::make_unique<ObjectDetectorOptions>();
|
||||
options->max_results = 1;
|
||||
options->category_allowlist.push_back("cat");
|
||||
options->base_options.model_asset_path =
|
||||
JoinPath("./", kTestDataDirectory, kMobileSsdWithMetadata);
|
||||
MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ObjectDetector> object_detector,
|
||||
@@ -589,10 +588,10 @@ TEST_F(ImageModeTest, SucceedsWithRotation) {
|
||||
results,
|
||||
ConvertToDetectionResult({ParseTextProtoOrDie<DetectionProto>(R"pb(
|
||||
label: "cat"
|
||||
score: 0.7109375
|
||||
score: 0.69921875
|
||||
location_data {
|
||||
format: BOUNDING_BOX
|
||||
bounding_box { xmin: 0 ymin: 622 width: 436 height: 276 }
|
||||
bounding_box { xmin: 0 ymin: 608 width: 439 height: 387 }
|
||||
})pb")}));
|
||||
}
|
||||
|
||||
|
||||
@@ -65,6 +65,8 @@ constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/vision/";
|
||||
constexpr char kPoseDetectionModel[] = "pose_detection.tflite";
|
||||
constexpr char kPortraitImage[] = "pose.jpg";
|
||||
constexpr char kPoseExpectedDetection[] = "pose_expected_detection.pbtxt";
|
||||
constexpr char kPoseExpectedExpandedRect[] =
|
||||
"pose_expected_expanded_rect.pbtxt";
|
||||
|
||||
constexpr char kImageTag[] = "IMAGE";
|
||||
constexpr char kImageName[] = "image";
|
||||
@@ -72,8 +74,11 @@ constexpr char kNormRectTag[] = "NORM_RECT";
|
||||
constexpr char kNormRectName[] = "norm_rect";
|
||||
constexpr char kDetectionsTag[] = "DETECTIONS";
|
||||
constexpr char kDetectionsName[] = "detections";
|
||||
constexpr char kExpandedPoseRectsTag[] = "EXPANDED_POSE_RECTS";
|
||||
constexpr char kExpandedPoseRectsName[] = "expanded_pose_rects";
|
||||
|
||||
constexpr float kPoseDetectionMaxDiff = 0.01;
|
||||
constexpr float kExpandedPoseRectMaxDiff = 0.01;
|
||||
|
||||
// Helper function to create a TaskRunner.
|
||||
absl::StatusOr<std::unique_ptr<TaskRunner>> CreateTaskRunner(
|
||||
@@ -99,6 +104,10 @@ absl::StatusOr<std::unique_ptr<TaskRunner>> CreateTaskRunner(
|
||||
pose_detector_graph.Out(kDetectionsTag).SetName(kDetectionsName) >>
|
||||
graph[Output<std::vector<Detection>>(kDetectionsTag)];
|
||||
|
||||
pose_detector_graph.Out(kExpandedPoseRectsTag)
|
||||
.SetName(kExpandedPoseRectsName) >>
|
||||
graph[Output<std::vector<NormalizedRect>>(kExpandedPoseRectsTag)];
|
||||
|
||||
return TaskRunner::Create(
|
||||
graph.GetConfig(), std::make_unique<core::MediaPipeBuiltinOpResolver>());
|
||||
}
|
||||
@@ -111,6 +120,14 @@ Detection GetExpectedPoseDetectionResult(absl::string_view file_name) {
|
||||
return detection;
|
||||
}
|
||||
|
||||
NormalizedRect GetExpectedExpandedPoseRect(absl::string_view file_name) {
|
||||
NormalizedRect expanded_rect;
|
||||
CHECK_OK(GetTextProto(file::JoinPath("./", kTestDataDirectory, file_name),
|
||||
&expanded_rect, Defaults()))
|
||||
<< "Expected expanded pose rect does not exist.";
|
||||
return expanded_rect;
|
||||
}
|
||||
|
||||
struct TestParams {
|
||||
// The name of this test, for convenience when displaying test results.
|
||||
std::string test_name;
|
||||
@@ -119,7 +136,9 @@ struct TestParams {
|
||||
// The filename of test image.
|
||||
std::string test_image_name;
|
||||
// Expected pose detection results.
|
||||
std::vector<Detection> expected_result;
|
||||
std::vector<Detection> expected_detection;
|
||||
// Expected expanded pose rects.
|
||||
std::vector<NormalizedRect> expected_expanded_pose_rect;
|
||||
};
|
||||
|
||||
class PoseDetectorGraphTest : public testing::TestWithParam<TestParams> {};
|
||||
@@ -144,16 +163,27 @@ TEST_P(PoseDetectorGraphTest, Succeed) {
|
||||
(*output_packets)[kDetectionsName].Get<std::vector<Detection>>();
|
||||
EXPECT_THAT(pose_detections, Pointwise(Approximately(Partially(EqualsProto()),
|
||||
kPoseDetectionMaxDiff),
|
||||
GetParam().expected_result));
|
||||
GetParam().expected_detection));
|
||||
|
||||
const std::vector<NormalizedRect>& expanded_pose_rects =
|
||||
(*output_packets)[kExpandedPoseRectsName]
|
||||
.Get<std::vector<NormalizedRect>>();
|
||||
EXPECT_THAT(expanded_pose_rects,
|
||||
Pointwise(Approximately(Partially(EqualsProto()),
|
||||
kExpandedPoseRectMaxDiff),
|
||||
GetParam().expected_expanded_pose_rect));
|
||||
}
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(
|
||||
PoseDetectorGraphTest, PoseDetectorGraphTest,
|
||||
Values(TestParams{.test_name = "DetectPose",
|
||||
.pose_detection_model_name = kPoseDetectionModel,
|
||||
.test_image_name = kPortraitImage,
|
||||
.expected_result = {GetExpectedPoseDetectionResult(
|
||||
kPoseExpectedDetection)}}),
|
||||
Values(TestParams{
|
||||
.test_name = "DetectPose",
|
||||
.pose_detection_model_name = kPoseDetectionModel,
|
||||
.test_image_name = kPortraitImage,
|
||||
.expected_detection = {GetExpectedPoseDetectionResult(
|
||||
kPoseExpectedDetection)},
|
||||
.expected_expanded_pose_rect = {GetExpectedExpandedPoseRect(
|
||||
kPoseExpectedExpandedRect)}}),
|
||||
[](const TestParamInfo<PoseDetectorGraphTest::ParamType>& info) {
|
||||
return info.param.test_name;
|
||||
});
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
# Copyright 2023 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.
|
||||
|
||||
package(default_visibility = [
|
||||
"//mediapipe/tasks:internal",
|
||||
])
|
||||
|
||||
licenses(["notice"])
|
||||
|
||||
cc_library(
|
||||
name = "pose_landmarks_detector_graph",
|
||||
srcs = ["pose_landmarks_detector_graph.cc"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/core:begin_loop_calculator",
|
||||
"//mediapipe/calculators/core:end_loop_calculator",
|
||||
"//mediapipe/calculators/core:gate_calculator",
|
||||
"//mediapipe/calculators/core:split_proto_list_calculator",
|
||||
"//mediapipe/calculators/core:split_vector_calculator",
|
||||
"//mediapipe/calculators/core:split_vector_calculator_cc_proto",
|
||||
"//mediapipe/calculators/tensor:inference_calculator",
|
||||
"//mediapipe/calculators/tensor:tensors_to_floats_calculator",
|
||||
"//mediapipe/calculators/tensor:tensors_to_landmarks_calculator",
|
||||
"//mediapipe/calculators/tensor:tensors_to_landmarks_calculator_cc_proto",
|
||||
"//mediapipe/calculators/tensor:tensors_to_segmentation_calculator",
|
||||
"//mediapipe/calculators/tensor:tensors_to_segmentation_calculator_cc_proto",
|
||||
"//mediapipe/calculators/util:detections_to_rects_calculator",
|
||||
"//mediapipe/calculators/util:landmarks_to_detection_calculator",
|
||||
"//mediapipe/calculators/util:rect_transformation_calculator",
|
||||
"//mediapipe/calculators/util:refine_landmarks_from_heatmap_calculator",
|
||||
"//mediapipe/calculators/util:refine_landmarks_from_heatmap_calculator_cc_proto",
|
||||
"//mediapipe/calculators/util:thresholding_calculator",
|
||||
"//mediapipe/calculators/util:thresholding_calculator_cc_proto",
|
||||
"//mediapipe/calculators/util:visibility_copy_calculator",
|
||||
"//mediapipe/calculators/util:visibility_copy_calculator_cc_proto",
|
||||
"//mediapipe/framework:subgraph",
|
||||
"//mediapipe/framework/api2:builder",
|
||||
"//mediapipe/framework/api2:port",
|
||||
"//mediapipe/framework/formats:image",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/gpu:gpu_origin_cc_proto",
|
||||
"//mediapipe/tasks/cc:common",
|
||||
"//mediapipe/tasks/cc/components/processors:image_preprocessing_graph",
|
||||
"//mediapipe/tasks/cc/core:model_resources",
|
||||
"//mediapipe/tasks/cc/core:model_task_graph",
|
||||
"//mediapipe/tasks/cc/vision/pose_landmarker/proto:pose_landmarks_detector_graph_options_cc_proto",
|
||||
"//mediapipe/tasks/cc/vision/utils:image_tensor_specs",
|
||||
"@com_google_absl//absl/status:statusor",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "pose_landmarker_graph",
|
||||
srcs = ["pose_landmarker_graph.cc"],
|
||||
deps = [
|
||||
":pose_landmarks_detector_graph",
|
||||
"//mediapipe/calculators/core:clip_vector_size_calculator",
|
||||
"//mediapipe/calculators/core:clip_vector_size_calculator_cc_proto",
|
||||
"//mediapipe/calculators/core:gate_calculator",
|
||||
"//mediapipe/calculators/core:gate_calculator_cc_proto",
|
||||
"//mediapipe/calculators/core:pass_through_calculator",
|
||||
"//mediapipe/calculators/core:previous_loopback_calculator",
|
||||
"//mediapipe/calculators/image:image_properties_calculator",
|
||||
"//mediapipe/calculators/util:association_calculator_cc_proto",
|
||||
"//mediapipe/calculators/util:association_norm_rect_calculator",
|
||||
"//mediapipe/calculators/util:collection_has_min_size_calculator",
|
||||
"//mediapipe/calculators/util:collection_has_min_size_calculator_cc_proto",
|
||||
"//mediapipe/framework/api2:builder",
|
||||
"//mediapipe/framework/api2:port",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/formats:image",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/formats:tensor",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/tasks/cc:common",
|
||||
"//mediapipe/tasks/cc/components/utils:gate",
|
||||
"//mediapipe/tasks/cc/core:model_asset_bundle_resources",
|
||||
"//mediapipe/tasks/cc/core:model_resources_cache",
|
||||
"//mediapipe/tasks/cc/core:model_task_graph",
|
||||
"//mediapipe/tasks/cc/core:utils",
|
||||
"//mediapipe/tasks/cc/metadata/utils:zip_utils",
|
||||
"//mediapipe/tasks/cc/vision/pose_detector:pose_detector_graph",
|
||||
"//mediapipe/tasks/cc/vision/pose_detector/proto:pose_detector_graph_options_cc_proto",
|
||||
"//mediapipe/tasks/cc/vision/pose_landmarker/proto:pose_landmarker_graph_options_cc_proto",
|
||||
"//mediapipe/tasks/cc/vision/pose_landmarker/proto:pose_landmarks_detector_graph_options_cc_proto",
|
||||
"//mediapipe/util:graph_builder_utils",
|
||||
"@com_google_absl//absl/strings:str_format",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -0,0 +1,384 @@
|
||||
/* Copyright 2023 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 <memory>
|
||||
#include <type_traits>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/strings/str_format.h"
|
||||
#include "mediapipe/calculators/core/clip_vector_size_calculator.pb.h"
|
||||
#include "mediapipe/calculators/core/gate_calculator.pb.h"
|
||||
#include "mediapipe/calculators/util/association_calculator.pb.h"
|
||||
#include "mediapipe/calculators/util/collection_has_min_size_calculator.pb.h"
|
||||
#include "mediapipe/framework/api2/builder.h"
|
||||
#include "mediapipe/framework/api2/port.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/image.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/formats/tensor.h"
|
||||
#include "mediapipe/framework/port/status_macros.h"
|
||||
#include "mediapipe/tasks/cc/common.h"
|
||||
#include "mediapipe/tasks/cc/components/utils/gate.h"
|
||||
#include "mediapipe/tasks/cc/core/model_asset_bundle_resources.h"
|
||||
#include "mediapipe/tasks/cc/core/model_resources_cache.h"
|
||||
#include "mediapipe/tasks/cc/core/model_task_graph.h"
|
||||
#include "mediapipe/tasks/cc/core/utils.h"
|
||||
#include "mediapipe/tasks/cc/metadata/utils/zip_utils.h"
|
||||
#include "mediapipe/tasks/cc/vision/pose_detector/proto/pose_detector_graph_options.pb.h"
|
||||
#include "mediapipe/tasks/cc/vision/pose_landmarker/proto/pose_landmarker_graph_options.pb.h"
|
||||
#include "mediapipe/tasks/cc/vision/pose_landmarker/proto/pose_landmarks_detector_graph_options.pb.h"
|
||||
#include "mediapipe/util/graph_builder_utils.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace tasks {
|
||||
namespace vision {
|
||||
namespace pose_landmarker {
|
||||
|
||||
namespace {
|
||||
|
||||
using ::mediapipe::NormalizedRect;
|
||||
using ::mediapipe::api2::Input;
|
||||
using ::mediapipe::api2::Output;
|
||||
using ::mediapipe::api2::builder::Graph;
|
||||
using ::mediapipe::api2::builder::SidePacket;
|
||||
using ::mediapipe::api2::builder::Source;
|
||||
using ::mediapipe::tasks::components::utils::DisallowIf;
|
||||
using ::mediapipe::tasks::core::ModelAssetBundleResources;
|
||||
using ::mediapipe::tasks::metadata::SetExternalFile;
|
||||
using ::mediapipe::tasks::vision::pose_detector::proto::
|
||||
PoseDetectorGraphOptions;
|
||||
using ::mediapipe::tasks::vision::pose_landmarker::proto::
|
||||
PoseLandmarkerGraphOptions;
|
||||
using ::mediapipe::tasks::vision::pose_landmarker::proto::
|
||||
PoseLandmarksDetectorGraphOptions;
|
||||
|
||||
constexpr char kImageTag[] = "IMAGE";
|
||||
constexpr char kNormRectTag[] = "NORM_RECT";
|
||||
constexpr char kNormLandmarksTag[] = "NORM_LANDMARKS";
|
||||
constexpr char kWorldLandmarksTag[] = "WORLD_LANDMARKS";
|
||||
constexpr char kAuxiliaryLandmarksTag[] = "AUXILIARY_LANDMARKS";
|
||||
constexpr char kPoseRectsNextFrameTag[] = "POSE_RECTS_NEXT_FRAME";
|
||||
constexpr char kExpandedPoseRectsTag[] = "EXPANDED_POSE_RECTS";
|
||||
constexpr char kDetectionsTag[] = "DETECTIONS";
|
||||
constexpr char kLoopTag[] = "LOOP";
|
||||
constexpr char kPrevLoopTag[] = "PREV_LOOP";
|
||||
constexpr char kMainTag[] = "MAIN";
|
||||
constexpr char kIterableTag[] = "ITERABLE";
|
||||
constexpr char kSegmentationMaskTag[] = "SEGMENTATION_MASK";
|
||||
|
||||
constexpr char kPoseDetectorTFLiteName[] = "pose_detector.tflite";
|
||||
constexpr char kPoseLandmarksDetectorTFLiteName[] =
|
||||
"pose_landmarks_detector.tflite";
|
||||
|
||||
struct PoseLandmarkerOutputs {
|
||||
Source<std::vector<NormalizedLandmarkList>> landmark_lists;
|
||||
Source<std::vector<LandmarkList>> world_landmark_lists;
|
||||
Source<std::vector<NormalizedLandmarkList>> auxiliary_landmark_lists;
|
||||
Source<std::vector<NormalizedRect>> pose_rects_next_frame;
|
||||
Source<std::vector<Detection>> pose_detections;
|
||||
Source<std::vector<Image>> segmentation_masks;
|
||||
Source<Image> image;
|
||||
};
|
||||
|
||||
// Sets the base options in the sub tasks.
|
||||
absl::Status SetSubTaskBaseOptions(const ModelAssetBundleResources& resources,
|
||||
PoseLandmarkerGraphOptions* options,
|
||||
bool is_copy) {
|
||||
auto* pose_detector_graph_options =
|
||||
options->mutable_pose_detector_graph_options();
|
||||
if (!pose_detector_graph_options->base_options().has_model_asset()) {
|
||||
ASSIGN_OR_RETURN(const auto pose_detector_file,
|
||||
resources.GetFile(kPoseDetectorTFLiteName));
|
||||
SetExternalFile(pose_detector_file,
|
||||
pose_detector_graph_options->mutable_base_options()
|
||||
->mutable_model_asset(),
|
||||
is_copy);
|
||||
}
|
||||
pose_detector_graph_options->mutable_base_options()
|
||||
->mutable_acceleration()
|
||||
->CopyFrom(options->base_options().acceleration());
|
||||
pose_detector_graph_options->mutable_base_options()->set_use_stream_mode(
|
||||
options->base_options().use_stream_mode());
|
||||
auto* pose_landmarks_detector_graph_options =
|
||||
options->mutable_pose_landmarks_detector_graph_options();
|
||||
if (!pose_landmarks_detector_graph_options->base_options()
|
||||
.has_model_asset()) {
|
||||
ASSIGN_OR_RETURN(const auto pose_landmarks_detector_file,
|
||||
resources.GetFile(kPoseLandmarksDetectorTFLiteName));
|
||||
SetExternalFile(
|
||||
pose_landmarks_detector_file,
|
||||
pose_landmarks_detector_graph_options->mutable_base_options()
|
||||
->mutable_model_asset(),
|
||||
is_copy);
|
||||
}
|
||||
pose_landmarks_detector_graph_options->mutable_base_options()
|
||||
->mutable_acceleration()
|
||||
->CopyFrom(options->base_options().acceleration());
|
||||
pose_landmarks_detector_graph_options->mutable_base_options()
|
||||
->set_use_stream_mode(options->base_options().use_stream_mode());
|
||||
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// A "mediapipe.tasks.vision.pose_landmarker.PoseLandmarkerGraph" performs pose
|
||||
// landmarks detection. The PoseLandmarkerGraph consists of two subgraphs:
|
||||
// PoseDetectorGraph, MultiplePoseLandmarksDetectorGraph
|
||||
//
|
||||
// MultiplePoseLandmarksDetectorGraph detects landmarks from bounding boxes
|
||||
// produced by PoseDetectorGraph. PoseLandmarkerGraph tracks the landmarks over
|
||||
// time, and skips the PoseDetectorGraph. If the tracking is lost or the
|
||||
// detected poses are less than configured max number poses, PoseDetectorGraph
|
||||
// would be triggered to detect poses.
|
||||
//
|
||||
//
|
||||
// Inputs:
|
||||
// IMAGE - Image
|
||||
// Image to perform pose landmarks detection on.
|
||||
// NORM_RECT - NormalizedRect @Optional
|
||||
// Describes image rotation and region of image to perform landmarks
|
||||
// detection on. If not provided, whole image is used for pose landmarks
|
||||
// detection.
|
||||
//
|
||||
//
|
||||
// Outputs:
|
||||
// NORM_LANDMARKS: - std::vector<NormalizedLandmarkList>
|
||||
// Vector of detected pose landmarks.
|
||||
// WORLD_LANDMARKS: std::vector<LandmarkList>
|
||||
// Vector of detected world pose landmarks.
|
||||
// AUXILIARY_LANDMARKS: - std::vector<NormalizedLandmarkList>
|
||||
// Vector of detected auxiliary landmarks.
|
||||
// POSE_RECTS_NEXT_FRAME - std::vector<NormalizedRect>
|
||||
// Vector of the expanded rects enclosing the whole pose RoI for landmark
|
||||
// detection on the next frame.
|
||||
// POSE_RECTS - std::vector<NormalizedRect>
|
||||
// Detected pose bounding boxes in normalized coordinates from pose
|
||||
// detection.
|
||||
// SEGMENTATION_MASK - std::vector<Image>
|
||||
// Segmentation masks.
|
||||
// IMAGE - Image
|
||||
// The input image that the pose landmarker runs on and has the pixel data
|
||||
// stored on the target storage (CPU vs GPU).
|
||||
// All returned coordinates are in the unrotated and uncropped input image
|
||||
// coordinates system.
|
||||
//
|
||||
// Example:
|
||||
// node {
|
||||
// calculator: "mediapipe.tasks.vision.pose_landmarker.PoseLandmarkerGraph"
|
||||
// input_stream: "IMAGE:image_in"
|
||||
// input_stream: "NORM_RECT:norm_rect"
|
||||
// output_stream: "NORM_LANDMARKS:pose_landmarks"
|
||||
// output_stream: "LANDMARKS:world_landmarks"
|
||||
// output_stream: "NORM_LANDMAKRS:auxiliary_landmarks"
|
||||
// output_stream: "POSE_RECTS_NEXT_FRAME:pose_rects_next_frame"
|
||||
// output_stream: "POSE_RECTS:pose_rects"
|
||||
// output_stream: "SEGMENTATION_MASK:segmentation_masks"
|
||||
// output_stream: "IMAGE:image_out"
|
||||
// options {
|
||||
// [mediapipe.tasks.vision.pose_landmarker.proto.PoseLandmarkerGraphOptions.ext]
|
||||
// {
|
||||
// base_options {
|
||||
// model_asset {
|
||||
// file_name: "pose_landmarker.task"
|
||||
// }
|
||||
// }
|
||||
// pose_detector_graph_options {
|
||||
// min_detection_confidence: 0.5
|
||||
// num_poses: 2
|
||||
// }
|
||||
// pose_landmarks_detector_graph_options {
|
||||
// min_detection_confidence: 0.5
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
class PoseLandmarkerGraph : public core::ModelTaskGraph {
|
||||
public:
|
||||
absl::StatusOr<CalculatorGraphConfig> GetConfig(
|
||||
SubgraphContext* sc) override {
|
||||
Graph graph;
|
||||
if (sc->Options<PoseLandmarkerGraphOptions>()
|
||||
.base_options()
|
||||
.has_model_asset()) {
|
||||
ASSIGN_OR_RETURN(
|
||||
const auto* model_asset_bundle_resources,
|
||||
CreateModelAssetBundleResources<PoseLandmarkerGraphOptions>(sc));
|
||||
// Copies the file content instead of passing the pointer of file in
|
||||
// memory if the subgraph model resource service is not available.
|
||||
MP_RETURN_IF_ERROR(SetSubTaskBaseOptions(
|
||||
*model_asset_bundle_resources,
|
||||
sc->MutableOptions<PoseLandmarkerGraphOptions>(),
|
||||
!sc->Service(::mediapipe::tasks::core::kModelResourcesCacheService)
|
||||
.IsAvailable()));
|
||||
}
|
||||
ASSIGN_OR_RETURN(
|
||||
auto outs,
|
||||
BuildPoseLandmarkerGraph(
|
||||
*sc->MutableOptions<PoseLandmarkerGraphOptions>(),
|
||||
graph[Input<Image>(kImageTag)],
|
||||
graph[Input<NormalizedRect>::Optional(kNormRectTag)], graph));
|
||||
outs.landmark_lists >>
|
||||
graph[Output<std::vector<NormalizedLandmarkList>>(kNormLandmarksTag)];
|
||||
outs.world_landmark_lists >>
|
||||
graph[Output<std::vector<LandmarkList>>(kWorldLandmarksTag)];
|
||||
outs.auxiliary_landmark_lists >>
|
||||
graph[Output<std::vector<NormalizedLandmarkList>>(
|
||||
kAuxiliaryLandmarksTag)];
|
||||
outs.pose_rects_next_frame >>
|
||||
graph[Output<std::vector<NormalizedRect>>(kPoseRectsNextFrameTag)];
|
||||
outs.segmentation_masks >>
|
||||
graph[Output<std::vector<Image>>(kSegmentationMaskTag)];
|
||||
outs.pose_detections >>
|
||||
graph[Output<std::vector<Detection>>(kDetectionsTag)];
|
||||
outs.image >> graph[Output<Image>(kImageTag)];
|
||||
|
||||
// TODO remove when support is fixed.
|
||||
// As mediapipe GraphBuilder currently doesn't support configuring
|
||||
// InputStreamInfo, modifying the CalculatorGraphConfig proto directly.
|
||||
CalculatorGraphConfig config = graph.GetConfig();
|
||||
for (int i = 0; i < config.node_size(); ++i) {
|
||||
if (config.node(i).calculator() == "PreviousLoopbackCalculator") {
|
||||
auto* info = config.mutable_node(i)->add_input_stream_info();
|
||||
info->set_tag_index(kLoopTag);
|
||||
info->set_back_edge(true);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return config;
|
||||
}
|
||||
|
||||
private:
|
||||
// Adds a mediapipe pose landmarker graph into the provided builder::Graph
|
||||
// instance.
|
||||
//
|
||||
// tasks_options: the mediapipe tasks module PoseLandmarkerGraphOptions.
|
||||
// image_in: (mediapipe::Image) stream to run pose landmark detection on.
|
||||
// graph: the mediapipe graph instance to be updated.
|
||||
absl::StatusOr<PoseLandmarkerOutputs> BuildPoseLandmarkerGraph(
|
||||
PoseLandmarkerGraphOptions& tasks_options, Source<Image> image_in,
|
||||
Source<NormalizedRect> norm_rect_in, Graph& graph) {
|
||||
const int max_num_poses =
|
||||
tasks_options.pose_detector_graph_options().num_poses();
|
||||
|
||||
auto& pose_detector =
|
||||
graph.AddNode("mediapipe.tasks.vision.pose_detector.PoseDetectorGraph");
|
||||
pose_detector.GetOptions<PoseDetectorGraphOptions>().Swap(
|
||||
tasks_options.mutable_pose_detector_graph_options());
|
||||
auto& clip_pose_rects =
|
||||
graph.AddNode("ClipNormalizedRectVectorSizeCalculator");
|
||||
clip_pose_rects.GetOptions<ClipVectorSizeCalculatorOptions>()
|
||||
.set_max_vec_size(max_num_poses);
|
||||
auto clipped_pose_rects = clip_pose_rects.Out("");
|
||||
|
||||
auto& pose_landmarks_detector_graph = graph.AddNode(
|
||||
"mediapipe.tasks.vision.pose_landmarker."
|
||||
"MultiplePoseLandmarksDetectorGraph");
|
||||
pose_landmarks_detector_graph
|
||||
.GetOptions<PoseLandmarksDetectorGraphOptions>()
|
||||
.Swap(tasks_options.mutable_pose_landmarks_detector_graph_options());
|
||||
image_in >> pose_landmarks_detector_graph.In(kImageTag);
|
||||
clipped_pose_rects >> pose_landmarks_detector_graph.In(kNormRectTag);
|
||||
|
||||
// TODO: Add landmarks smoothing calculators to
|
||||
// PoseLandmarkerGraph
|
||||
auto landmarks = pose_landmarks_detector_graph.Out("LANDMARKS")
|
||||
.Cast<std::vector<NormalizedLandmarkList>>();
|
||||
auto world_landmarks = pose_landmarks_detector_graph.Out(kWorldLandmarksTag)
|
||||
.Cast<std::vector<LandmarkList>>();
|
||||
auto aux_landmarks =
|
||||
pose_landmarks_detector_graph.Out(kAuxiliaryLandmarksTag)
|
||||
.Cast<std::vector<NormalizedLandmarkList>>();
|
||||
auto pose_rects_for_next_frame =
|
||||
pose_landmarks_detector_graph.Out(kPoseRectsNextFrameTag)
|
||||
.Cast<std::vector<NormalizedRect>>();
|
||||
auto segmentation_masks =
|
||||
pose_landmarks_detector_graph.Out(kSegmentationMaskTag)
|
||||
.Cast<std::vector<Image>>();
|
||||
|
||||
if (tasks_options.base_options().use_stream_mode()) {
|
||||
auto& previous_loopback = graph.AddNode("PreviousLoopbackCalculator");
|
||||
image_in >> previous_loopback.In(kMainTag);
|
||||
auto prev_pose_rects_from_landmarks =
|
||||
previous_loopback[Output<std::vector<NormalizedRect>>(kPrevLoopTag)];
|
||||
|
||||
auto& min_size_node =
|
||||
graph.AddNode("NormalizedRectVectorHasMinSizeCalculator");
|
||||
prev_pose_rects_from_landmarks >> min_size_node.In(kIterableTag);
|
||||
min_size_node.GetOptions<CollectionHasMinSizeCalculatorOptions>()
|
||||
.set_min_size(max_num_poses);
|
||||
auto has_enough_poses = min_size_node.Out("").Cast<bool>();
|
||||
|
||||
// While in stream mode, skip pose detector graph when we successfully
|
||||
// track the poses from the last frame.
|
||||
auto image_for_pose_detector =
|
||||
DisallowIf(image_in, has_enough_poses, graph);
|
||||
auto norm_rect_in_for_pose_detector =
|
||||
DisallowIf(norm_rect_in, has_enough_poses, graph);
|
||||
image_for_pose_detector >> pose_detector.In(kImageTag);
|
||||
norm_rect_in_for_pose_detector >> pose_detector.In(kNormRectTag);
|
||||
auto expanded_pose_rects_from_pose_detector =
|
||||
pose_detector.Out(kExpandedPoseRectsTag);
|
||||
auto& pose_association = graph.AddNode("AssociationNormRectCalculator");
|
||||
pose_association.GetOptions<mediapipe::AssociationCalculatorOptions>()
|
||||
.set_min_similarity_threshold(
|
||||
tasks_options.min_tracking_confidence());
|
||||
prev_pose_rects_from_landmarks >>
|
||||
pose_association[Input<std::vector<NormalizedRect>>::Multiple("")][0];
|
||||
expanded_pose_rects_from_pose_detector >>
|
||||
pose_association[Input<std::vector<NormalizedRect>>::Multiple("")][1];
|
||||
auto pose_rects = pose_association.Out("");
|
||||
pose_rects >> clip_pose_rects.In("");
|
||||
// Back edge.
|
||||
pose_rects_for_next_frame >> previous_loopback.In(kLoopTag);
|
||||
} else {
|
||||
// While not in stream mode, the input images are not guaranteed to be in
|
||||
// series, and we don't want to enable the tracking and rect associations
|
||||
// between input images. Always use the pose detector graph.
|
||||
image_in >> pose_detector.In(kImageTag);
|
||||
norm_rect_in >> pose_detector.In(kNormRectTag);
|
||||
auto pose_rects = pose_detector.Out(kExpandedPoseRectsTag);
|
||||
pose_rects >> clip_pose_rects.In("");
|
||||
}
|
||||
|
||||
// TODO: Replace PassThroughCalculator with a calculator that
|
||||
// converts the pixel data to be stored on the target storage (CPU vs GPU).
|
||||
auto& pass_through = graph.AddNode("PassThroughCalculator");
|
||||
image_in >> pass_through.In("");
|
||||
|
||||
return {{
|
||||
/* landmark_lists= */ landmarks,
|
||||
/* world_landmarks= */ world_landmarks,
|
||||
/* aux_landmarks= */ aux_landmarks,
|
||||
/* pose_rects_next_frame= */ pose_rects_for_next_frame,
|
||||
/* pose_detections */
|
||||
pose_detector.Out(kDetectionsTag).Cast<std::vector<Detection>>(),
|
||||
/* segmentation_masks= */ segmentation_masks,
|
||||
/* image= */
|
||||
pass_through[Output<Image>("")],
|
||||
}};
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_MEDIAPIPE_GRAPH(
|
||||
::mediapipe::tasks::vision::pose_landmarker::PoseLandmarkerGraph);
|
||||
|
||||
} // namespace pose_landmarker
|
||||
} // namespace vision
|
||||
} // namespace tasks
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,190 @@
|
||||
/* Copyright 2023 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 <optional>
|
||||
|
||||
#include "absl/flags/flag.h"
|
||||
#include "absl/status/statusor.h"
|
||||
#include "absl/strings/str_format.h"
|
||||
#include "absl/strings/string_view.h"
|
||||
#include "mediapipe/framework/api2/builder.h"
|
||||
#include "mediapipe/framework/api2/port.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/deps/file_path.h"
|
||||
#include "mediapipe/framework/formats/image.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/file_helpers.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/tasks/cc/core/mediapipe_builtin_op_resolver.h"
|
||||
#include "mediapipe/tasks/cc/core/proto/base_options.pb.h"
|
||||
#include "mediapipe/tasks/cc/core/proto/external_file.pb.h"
|
||||
#include "mediapipe/tasks/cc/core/task_runner.h"
|
||||
#include "mediapipe/tasks/cc/vision/pose_detector/proto/pose_detector_graph_options.pb.h"
|
||||
#include "mediapipe/tasks/cc/vision/pose_landmarker/proto/pose_landmarker_graph_options.pb.h"
|
||||
#include "mediapipe/tasks/cc/vision/pose_landmarker/proto/pose_landmarks_detector_graph_options.pb.h"
|
||||
#include "mediapipe/tasks/cc/vision/utils/image_utils.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace tasks {
|
||||
namespace vision {
|
||||
namespace pose_landmarker {
|
||||
namespace {
|
||||
|
||||
using ::file::Defaults;
|
||||
using ::file::GetTextProto;
|
||||
using ::mediapipe::api2::Input;
|
||||
using ::mediapipe::api2::Output;
|
||||
using ::mediapipe::api2::builder::Graph;
|
||||
using ::mediapipe::api2::builder::Source;
|
||||
using ::mediapipe::file::JoinPath;
|
||||
using ::mediapipe::tasks::core::TaskRunner;
|
||||
using ::mediapipe::tasks::vision::DecodeImageFromFile;
|
||||
using ::mediapipe::tasks::vision::pose_landmarker::proto::
|
||||
PoseLandmarkerGraphOptions;
|
||||
using ::testing::EqualsProto;
|
||||
using ::testing::Pointwise;
|
||||
using ::testing::TestParamInfo;
|
||||
using ::testing::TestWithParam;
|
||||
using ::testing::Values;
|
||||
using ::testing::proto::Approximately;
|
||||
using ::testing::proto::Partially;
|
||||
|
||||
constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/vision/";
|
||||
constexpr char kPoseLandmarkerModelBundleName[] = "pose_landmarker.task";
|
||||
constexpr char kPoseImageName[] = "pose.jpg";
|
||||
constexpr char kExpectedPoseLandmarksName[] =
|
||||
"expected_pose_landmarks.prototxt";
|
||||
|
||||
constexpr char kImageTag[] = "IMAGE";
|
||||
constexpr char kImageName[] = "image";
|
||||
constexpr char kNormRectTag[] = "NORM_RECT";
|
||||
constexpr char kNormRectName[] = "norm_rect";
|
||||
constexpr char kNormLandmarksTag[] = "NORM_LANDMARKS";
|
||||
constexpr char kNormLandmarksName[] = "norm_landmarks";
|
||||
|
||||
constexpr float kLiteModelFractionDiff = 0.05; // percentage
|
||||
|
||||
template <typename ProtoT>
|
||||
ProtoT GetExpectedProto(absl::string_view filename) {
|
||||
ProtoT expected_proto;
|
||||
MP_EXPECT_OK(GetTextProto(file::JoinPath("./", kTestDataDirectory, filename),
|
||||
&expected_proto, Defaults()));
|
||||
return expected_proto;
|
||||
}
|
||||
|
||||
// Struct holding the parameters for parameterized PoseLandmarkerGraphTest
|
||||
// class.
|
||||
struct PoseLandmarkerGraphTestParams {
|
||||
// The name of this test, for convenience when displaying test results.
|
||||
std::string test_name;
|
||||
// The filename of the model to test.
|
||||
std::string input_model_name;
|
||||
// The filename of the test image.
|
||||
std::string test_image_name;
|
||||
// The expected output landmarks positions.
|
||||
std::optional<std::vector<NormalizedLandmarkList>> expected_landmarks_list;
|
||||
// The max value difference between expected_positions and detected positions.
|
||||
float landmarks_diff_threshold;
|
||||
};
|
||||
|
||||
// Helper function to create a PoseLandmarkerGraph TaskRunner.
|
||||
absl::StatusOr<std::unique_ptr<TaskRunner>> CreatePoseLandmarkerGraphTaskRunner(
|
||||
absl::string_view model_name) {
|
||||
Graph graph;
|
||||
|
||||
auto& pose_landmarker = graph.AddNode(
|
||||
"mediapipe.tasks.vision.pose_landmarker."
|
||||
"PoseLandmarkerGraph");
|
||||
|
||||
auto* options = &pose_landmarker.GetOptions<PoseLandmarkerGraphOptions>();
|
||||
options->mutable_base_options()->mutable_model_asset()->set_file_name(
|
||||
JoinPath("./", kTestDataDirectory, model_name));
|
||||
options->mutable_pose_detector_graph_options()->set_num_poses(1);
|
||||
options->mutable_base_options()->set_use_stream_mode(true);
|
||||
|
||||
graph[Input<Image>(kImageTag)].SetName(kImageName) >>
|
||||
pose_landmarker.In(kImageTag);
|
||||
graph[Input<NormalizedRect>(kNormRectTag)].SetName(kNormRectName) >>
|
||||
pose_landmarker.In(kNormRectTag);
|
||||
|
||||
pose_landmarker.Out(kNormLandmarksTag).SetName(kNormLandmarksName) >>
|
||||
graph[Output<std::vector<NormalizedLandmarkList>>(kNormLandmarksTag)];
|
||||
|
||||
return TaskRunner::Create(
|
||||
graph.GetConfig(),
|
||||
absl::make_unique<tasks::core::MediaPipeBuiltinOpResolver>());
|
||||
}
|
||||
|
||||
// Helper function to construct NormalizeRect proto.
|
||||
NormalizedRect MakeNormRect(float x_center, float y_center, float width,
|
||||
float height, float rotation) {
|
||||
NormalizedRect pose_rect;
|
||||
pose_rect.set_x_center(x_center);
|
||||
pose_rect.set_y_center(y_center);
|
||||
pose_rect.set_width(width);
|
||||
pose_rect.set_height(height);
|
||||
pose_rect.set_rotation(rotation);
|
||||
return pose_rect;
|
||||
}
|
||||
|
||||
class PoseLandmarkerGraphTest
|
||||
: public testing::TestWithParam<PoseLandmarkerGraphTestParams> {};
|
||||
|
||||
TEST_P(PoseLandmarkerGraphTest, Succeeds) {
|
||||
MP_ASSERT_OK_AND_ASSIGN(
|
||||
Image image, DecodeImageFromFile(JoinPath("./", kTestDataDirectory,
|
||||
GetParam().test_image_name)));
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto task_runner, CreatePoseLandmarkerGraphTaskRunner(
|
||||
GetParam().input_model_name));
|
||||
|
||||
auto output_packets = task_runner->Process(
|
||||
{{kImageName, MakePacket<Image>(std::move(image))},
|
||||
{kNormRectName,
|
||||
MakePacket<NormalizedRect>(MakeNormRect(0.5, 0.5, 1.0, 1.0, 0))}});
|
||||
MP_ASSERT_OK(output_packets);
|
||||
|
||||
if (GetParam().expected_landmarks_list) {
|
||||
const std::vector<NormalizedLandmarkList>& landmarks_lists =
|
||||
(*output_packets)[kNormLandmarksName]
|
||||
.Get<std::vector<NormalizedLandmarkList>>();
|
||||
EXPECT_THAT(landmarks_lists,
|
||||
Pointwise(Approximately(Partially(EqualsProto()),
|
||||
GetParam().landmarks_diff_threshold),
|
||||
*GetParam().expected_landmarks_list));
|
||||
}
|
||||
}
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(
|
||||
PoseLandmarkerGraphTests, PoseLandmarkerGraphTest,
|
||||
Values(PoseLandmarkerGraphTestParams{
|
||||
/* test_name= */ "PoseLandmarkerLite",
|
||||
/* input_model_name= */ kPoseLandmarkerModelBundleName,
|
||||
/* test_image_name= */ kPoseImageName,
|
||||
/* expected_landmarks_list= */
|
||||
{{GetExpectedProto<NormalizedLandmarkList>(
|
||||
kExpectedPoseLandmarksName)}},
|
||||
/* landmarks_diff_threshold= */ kLiteModelFractionDiff}),
|
||||
[](const TestParamInfo<PoseLandmarkerGraphTest::ParamType>& info) {
|
||||
return info.param.test_name;
|
||||
});
|
||||
|
||||
} // namespace
|
||||
} // namespace pose_landmarker
|
||||
} // namespace vision
|
||||
} // namespace tasks
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,641 @@
|
||||
/* Copyright 2023 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 "absl/status/statusor.h"
|
||||
#include "mediapipe/calculators/core/split_vector_calculator.pb.h"
|
||||
#include "mediapipe/calculators/tensor/tensors_to_landmarks_calculator.pb.h"
|
||||
#include "mediapipe/calculators/tensor/tensors_to_segmentation_calculator.pb.h"
|
||||
#include "mediapipe/calculators/util/refine_landmarks_from_heatmap_calculator.pb.h"
|
||||
#include "mediapipe/calculators/util/thresholding_calculator.pb.h"
|
||||
#include "mediapipe/calculators/util/visibility_copy_calculator.pb.h"
|
||||
#include "mediapipe/framework/api2/builder.h"
|
||||
#include "mediapipe/framework/api2/port.h"
|
||||
#include "mediapipe/framework/formats/image.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/subgraph.h"
|
||||
#include "mediapipe/gpu/gpu_origin.pb.h"
|
||||
#include "mediapipe/tasks/cc/common.h"
|
||||
#include "mediapipe/tasks/cc/components/processors/image_preprocessing_graph.h"
|
||||
#include "mediapipe/tasks/cc/core/model_resources.h"
|
||||
#include "mediapipe/tasks/cc/core/model_task_graph.h"
|
||||
#include "mediapipe/tasks/cc/vision/pose_landmarker/proto/pose_landmarks_detector_graph_options.pb.h"
|
||||
#include "mediapipe/tasks/cc/vision/utils/image_tensor_specs.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace tasks {
|
||||
namespace vision {
|
||||
namespace pose_landmarker {
|
||||
|
||||
using ::mediapipe::NormalizedRect;
|
||||
using ::mediapipe::api2::Input;
|
||||
using ::mediapipe::api2::Output;
|
||||
using ::mediapipe::api2::builder::Graph;
|
||||
using ::mediapipe::api2::builder::Source;
|
||||
using ::mediapipe::tasks::core::ModelResources;
|
||||
using ::mediapipe::tasks::vision::pose_landmarker::proto::
|
||||
PoseLandmarksDetectorGraphOptions;
|
||||
|
||||
constexpr char kImageTag[] = "IMAGE";
|
||||
constexpr char kNormRectTag[] = "NORM_RECT";
|
||||
constexpr char kLandmarksTag[] = "LANDMARKS";
|
||||
constexpr char kNormLandmarksTag[] = "NORM_LANDMARKS";
|
||||
constexpr char kWorldLandmarksTag[] = "WORLD_LANDMARKS";
|
||||
constexpr char kAuxLandmarksTag[] = "AUXILIARY_LANDMARKS";
|
||||
constexpr char kPoseRectNextFrameTag[] = "POSE_RECT_NEXT_FRAME";
|
||||
constexpr char kPoseRectsNextFrameTag[] = "POSE_RECTS_NEXT_FRAME";
|
||||
constexpr char kPresenceTag[] = "PRESENCE";
|
||||
constexpr char kPresenceScoreTag[] = "PRESENCE_SCORE";
|
||||
constexpr char kSegmentationMaskTag[] = "SEGMENTATION_MASK";
|
||||
constexpr char kImageSizeTag[] = "IMAGE_SIZE";
|
||||
constexpr char kLandmarksToTag[] = "LANDMARKS_TO";
|
||||
constexpr char kTensorsTag[] = "TENSORS";
|
||||
constexpr char kFloatTag[] = "FLOAT";
|
||||
constexpr char kFlagTag[] = "FLAG";
|
||||
constexpr char kMaskTag[] = "MASK";
|
||||
constexpr char kDetectionTag[] = "DETECTION";
|
||||
constexpr char kNormLandmarksFromTag[] = "NORM_LANDMARKS_FROM";
|
||||
constexpr char kBatchEndTag[] = "BATCH_END";
|
||||
constexpr char kItemTag[] = "ITEM";
|
||||
constexpr char kIterableTag[] = "ITERABLE";
|
||||
|
||||
constexpr int kModelOutputTensorSplitNum = 5;
|
||||
constexpr int kLandmarksNum = 39;
|
||||
constexpr float kLandmarksNormalizeZ = 0.4;
|
||||
|
||||
struct SinglePoseLandmarkerOutputs {
|
||||
Source<NormalizedLandmarkList> pose_landmarks;
|
||||
Source<LandmarkList> world_pose_landmarks;
|
||||
Source<NormalizedLandmarkList> auxiliary_pose_landmarks;
|
||||
Source<NormalizedRect> pose_rect_next_frame;
|
||||
Source<bool> pose_presence;
|
||||
Source<float> pose_presence_score;
|
||||
Source<Image> segmentation_mask;
|
||||
};
|
||||
|
||||
struct PoseLandmarkerOutputs {
|
||||
Source<std::vector<NormalizedLandmarkList>> landmark_lists;
|
||||
Source<std::vector<LandmarkList>> world_landmark_lists;
|
||||
Source<std::vector<NormalizedLandmarkList>> auxiliary_landmark_lists;
|
||||
Source<std::vector<NormalizedRect>> pose_rects_next_frame;
|
||||
Source<std::vector<bool>> presences;
|
||||
Source<std::vector<float>> presence_scores;
|
||||
Source<std::vector<Image>> segmentation_masks;
|
||||
};
|
||||
|
||||
absl::Status SanityCheckOptions(
|
||||
const PoseLandmarksDetectorGraphOptions& options) {
|
||||
if (options.min_detection_confidence() < 0 ||
|
||||
options.min_detection_confidence() > 1) {
|
||||
return CreateStatusWithPayload(absl::StatusCode::kInvalidArgument,
|
||||
"Invalid `min_detection_confidence` option: "
|
||||
"value must be in the range [0.0, 1.0]",
|
||||
MediaPipeTasksStatus::kInvalidArgumentError);
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
// Split pose landmark detection model output tensor into five parts,
|
||||
// representing landmarks, presence scores, segmentation, heatmap, and world
|
||||
// landmarks respectively.
|
||||
void ConfigureSplitTensorVectorCalculator(
|
||||
mediapipe::SplitVectorCalculatorOptions* options) {
|
||||
for (int i = 0; i < kModelOutputTensorSplitNum; ++i) {
|
||||
auto* range = options->add_ranges();
|
||||
range->set_begin(i);
|
||||
range->set_end(i + 1);
|
||||
}
|
||||
}
|
||||
|
||||
void ConfigureTensorsToLandmarksCalculator(
|
||||
const ImageTensorSpecs& input_image_tensor_spec, bool normalize,
|
||||
bool sigmoid_activation,
|
||||
mediapipe::TensorsToLandmarksCalculatorOptions* options) {
|
||||
options->set_num_landmarks(kLandmarksNum);
|
||||
options->set_input_image_height(input_image_tensor_spec.image_height);
|
||||
options->set_input_image_width(input_image_tensor_spec.image_width);
|
||||
|
||||
if (normalize) {
|
||||
options->set_normalize_z(kLandmarksNormalizeZ);
|
||||
}
|
||||
|
||||
if (sigmoid_activation) {
|
||||
options->set_visibility_activation(
|
||||
mediapipe::TensorsToLandmarksCalculatorOptions_Activation_SIGMOID);
|
||||
options->set_presence_activation(
|
||||
mediapipe::TensorsToLandmarksCalculatorOptions_Activation_SIGMOID);
|
||||
}
|
||||
}
|
||||
|
||||
void ConfigureTensorsToSegmentationCalculator(
|
||||
mediapipe::TensorsToSegmentationCalculatorOptions* options) {
|
||||
options->set_activation(
|
||||
mediapipe::TensorsToSegmentationCalculatorOptions_Activation_SIGMOID);
|
||||
options->set_gpu_origin(mediapipe::GpuOrigin::TOP_LEFT);
|
||||
}
|
||||
|
||||
void ConfigureRefineLandmarksFromHeatmapCalculator(
|
||||
mediapipe::RefineLandmarksFromHeatmapCalculatorOptions* options) {
|
||||
// Derived from
|
||||
// mediapipe/modules/pose_landmark/tensors_to_pose_landmarks_and_segmentation.pbtxt.
|
||||
options->set_kernel_size(7);
|
||||
}
|
||||
|
||||
void ConfigureSplitNormalizedLandmarkListCalculator(
|
||||
mediapipe::SplitVectorCalculatorOptions* options) {
|
||||
// Derived from
|
||||
// mediapipe/modules/pose_landmark/tensors_to_pose_landmarks_and_segmentation.pbtxt
|
||||
auto* range = options->add_ranges();
|
||||
range->set_begin(0);
|
||||
range->set_end(33);
|
||||
auto* range_2 = options->add_ranges();
|
||||
range_2->set_begin(33);
|
||||
range_2->set_end(35);
|
||||
}
|
||||
|
||||
void ConfigureSplitLandmarkListCalculator(
|
||||
mediapipe::SplitVectorCalculatorOptions* options) {
|
||||
// Derived from
|
||||
// mediapipe/modules/pose_landmark/tensors_to_pose_landmarks_and_segmentation.pbtxt
|
||||
auto* range = options->add_ranges();
|
||||
range->set_begin(0);
|
||||
range->set_end(33);
|
||||
}
|
||||
|
||||
void ConfigureVisibilityCopyCalculator(
|
||||
mediapipe::VisibilityCopyCalculatorOptions* options) {
|
||||
// Derived from
|
||||
// mediapipe/modules/pose_landmark/tensors_to_pose_landmarks_and_segmentation.pbtxt
|
||||
options->set_copy_visibility(true);
|
||||
options->set_copy_presence(true);
|
||||
}
|
||||
|
||||
// A "mediapipe.tasks.vision.pose_landmarker.SinglePoseLandmarksDetectorGraph"
|
||||
// performs pose landmarks detection.
|
||||
// - Accepts CPU input images and outputs Landmark on CPU.
|
||||
//
|
||||
// Inputs:
|
||||
// IMAGE - Image
|
||||
// Image to perform detection on.
|
||||
// NORM_RECT - NormalizedRect @Optional
|
||||
// Rect enclosing the RoI to perform detection on. If not set, the detection
|
||||
// RoI is the whole image.
|
||||
//
|
||||
//
|
||||
// Outputs:
|
||||
// LANDMARKS: - NormalizedLandmarkList
|
||||
// Detected pose landmarks.
|
||||
// WORLD_LANDMARKS - LandmarkList
|
||||
// Detected pose landmarks in world coordinates.
|
||||
// AUXILIARY_LANDMARKS - NormalizedLandmarkList
|
||||
// Detected pose auxiliary landmarks.
|
||||
// POSE_RECT_NEXT_FRAME - NormalizedRect
|
||||
// The predicted Rect enclosing the pose RoI for landmark detection on the
|
||||
// next frame.
|
||||
// PRESENCE - bool
|
||||
// Boolean value indicates whether the pose is present.
|
||||
// PRESENCE_SCORE - float
|
||||
// Float value indicates the probability that the pose is present.
|
||||
// SEGMENTATION_MASK - Image
|
||||
// Segmentation mask for pose.
|
||||
//
|
||||
// Example:
|
||||
// node {
|
||||
// calculator:
|
||||
// "mediapipe.tasks.vision.pose_landmarker.SingleposeLandmarksDetectorGraph"
|
||||
// input_stream: "IMAGE:input_image"
|
||||
// input_stream: "POSE_RECT:pose_rect"
|
||||
// output_stream: "LANDMARKS:pose_landmarks"
|
||||
// output_stream: "WORLD_LANDMARKS:world_pose_landmarks"
|
||||
// output_stream: "AUXILIARY_LANDMARKS:auxiliary_landmarks"
|
||||
// output_stream: "POSE_RECT_NEXT_FRAME:pose_rect_next_frame"
|
||||
// output_stream: "PRESENCE:pose_presence"
|
||||
// output_stream: "PRESENCE_SCORE:pose_presence_score"
|
||||
// output_stream: "SEGMENTATION_MASK:segmentation_mask"
|
||||
// options {
|
||||
// [mediapipe.tasks.vision.pose_landmarker.proto.poseLandmarksDetectorGraphOptions.ext]
|
||||
// {
|
||||
// base_options {
|
||||
// model_asset {
|
||||
// file_name: "pose_landmark_lite.tflite"
|
||||
// }
|
||||
// }
|
||||
// min_detection_confidence: 0.5
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
class SinglePoseLandmarksDetectorGraph : public core::ModelTaskGraph {
|
||||
public:
|
||||
absl::StatusOr<CalculatorGraphConfig> GetConfig(
|
||||
SubgraphContext* sc) override {
|
||||
ASSIGN_OR_RETURN(
|
||||
const auto* model_resources,
|
||||
CreateModelResources<PoseLandmarksDetectorGraphOptions>(sc));
|
||||
Graph graph;
|
||||
ASSIGN_OR_RETURN(
|
||||
auto pose_landmark_detection_outs,
|
||||
BuildSinglePoseLandmarksDetectorGraph(
|
||||
sc->Options<PoseLandmarksDetectorGraphOptions>(), *model_resources,
|
||||
graph[Input<Image>(kImageTag)],
|
||||
graph[Input<NormalizedRect>::Optional(kNormRectTag)], graph));
|
||||
pose_landmark_detection_outs.pose_landmarks >>
|
||||
graph[Output<NormalizedLandmarkList>(kLandmarksTag)];
|
||||
pose_landmark_detection_outs.world_pose_landmarks >>
|
||||
graph[Output<LandmarkList>(kWorldLandmarksTag)];
|
||||
pose_landmark_detection_outs.auxiliary_pose_landmarks >>
|
||||
graph[Output<NormalizedLandmarkList>(kAuxLandmarksTag)];
|
||||
pose_landmark_detection_outs.pose_rect_next_frame >>
|
||||
graph[Output<NormalizedRect>(kPoseRectNextFrameTag)];
|
||||
pose_landmark_detection_outs.pose_presence >>
|
||||
graph[Output<bool>(kPresenceTag)];
|
||||
pose_landmark_detection_outs.pose_presence_score >>
|
||||
graph[Output<float>(kPresenceScoreTag)];
|
||||
pose_landmark_detection_outs.segmentation_mask >>
|
||||
graph[Output<Image>(kSegmentationMaskTag)];
|
||||
|
||||
return graph.GetConfig();
|
||||
}
|
||||
|
||||
private:
|
||||
absl::StatusOr<SinglePoseLandmarkerOutputs>
|
||||
BuildSinglePoseLandmarksDetectorGraph(
|
||||
const PoseLandmarksDetectorGraphOptions& subgraph_options,
|
||||
const ModelResources& model_resources, Source<Image> image_in,
|
||||
Source<NormalizedRect> pose_rect, Graph& graph) {
|
||||
MP_RETURN_IF_ERROR(SanityCheckOptions(subgraph_options));
|
||||
|
||||
auto& preprocessing = graph.AddNode(
|
||||
"mediapipe.tasks.components.processors.ImagePreprocessingGraph");
|
||||
bool use_gpu =
|
||||
components::processors::DetermineImagePreprocessingGpuBackend(
|
||||
subgraph_options.base_options().acceleration());
|
||||
MP_RETURN_IF_ERROR(components::processors::ConfigureImagePreprocessingGraph(
|
||||
model_resources, use_gpu,
|
||||
&preprocessing.GetOptions<tasks::components::processors::proto::
|
||||
ImagePreprocessingGraphOptions>()));
|
||||
image_in >> preprocessing.In(kImageTag);
|
||||
pose_rect >> preprocessing.In(kNormRectTag);
|
||||
auto image_size = preprocessing[Output<std::pair<int, int>>(kImageSizeTag)];
|
||||
|
||||
ASSIGN_OR_RETURN(auto image_tensor_specs,
|
||||
BuildInputImageTensorSpecs(model_resources));
|
||||
|
||||
auto& inference = AddInference(
|
||||
model_resources, subgraph_options.base_options().acceleration(), graph);
|
||||
preprocessing.Out(kTensorsTag) >> inference.In(kTensorsTag);
|
||||
|
||||
// Split model output tensors to multiple streams.
|
||||
auto& split_tensors_vector = graph.AddNode("SplitTensorVectorCalculator");
|
||||
ConfigureSplitTensorVectorCalculator(
|
||||
&split_tensors_vector
|
||||
.GetOptions<mediapipe::SplitVectorCalculatorOptions>());
|
||||
inference.Out(kTensorsTag) >> split_tensors_vector.In("");
|
||||
auto landmark_tensors = split_tensors_vector.Out(0);
|
||||
auto pose_flag_tensors = split_tensors_vector.Out(1);
|
||||
auto segmentation_tensors = split_tensors_vector.Out(2);
|
||||
auto heatmap_tensors = split_tensors_vector.Out(3);
|
||||
auto world_landmark_tensors = split_tensors_vector.Out(4);
|
||||
|
||||
// Converts the pose-flag tensor into a float that represents the confidence
|
||||
// score of pose presence.
|
||||
auto& tensors_to_pose_presence = graph.AddNode("TensorsToFloatsCalculator");
|
||||
pose_flag_tensors >> tensors_to_pose_presence.In(kTensorsTag);
|
||||
auto pose_presence_score =
|
||||
tensors_to_pose_presence[Output<float>(kFloatTag)];
|
||||
|
||||
// Applies a threshold to the confidence score to determine whether a
|
||||
// pose is present.
|
||||
auto& pose_presence_thresholding = graph.AddNode("ThresholdingCalculator");
|
||||
pose_presence_thresholding
|
||||
.GetOptions<mediapipe::ThresholdingCalculatorOptions>()
|
||||
.set_threshold(subgraph_options.min_detection_confidence());
|
||||
pose_presence_score >> pose_presence_thresholding.In(kFloatTag);
|
||||
auto pose_presence = pose_presence_thresholding[Output<bool>(kFlagTag)];
|
||||
|
||||
// GateCalculator for tensors.
|
||||
auto& tensors_gate = graph.AddNode("GateCalculator");
|
||||
landmark_tensors >> tensors_gate.In("")[0];
|
||||
segmentation_tensors >> tensors_gate.In("")[1];
|
||||
heatmap_tensors >> tensors_gate.In("")[2];
|
||||
world_landmark_tensors >> tensors_gate.In("")[3];
|
||||
pose_presence >> tensors_gate.In("ALLOW");
|
||||
auto ensured_landmarks_tensors = tensors_gate.Out(0);
|
||||
auto ensured_segmentation_tensors = tensors_gate.Out(1);
|
||||
auto ensured_heatmap_tensors = tensors_gate.Out(2);
|
||||
auto ensured_world_landmark_tensors = tensors_gate.Out(3);
|
||||
|
||||
// Decodes the landmark tensors into a list of landmarks, where the landmark
|
||||
// coordinates are normalized by the size of the input image to the model.
|
||||
auto& tensors_to_landmarks = graph.AddNode("TensorsToLandmarksCalculator");
|
||||
ConfigureTensorsToLandmarksCalculator(
|
||||
image_tensor_specs, /* normalize = */ false,
|
||||
/*sigmoid_activation= */ true,
|
||||
&tensors_to_landmarks
|
||||
.GetOptions<mediapipe::TensorsToLandmarksCalculatorOptions>());
|
||||
ensured_landmarks_tensors >> tensors_to_landmarks.In(kTensorsTag);
|
||||
|
||||
auto landmarks =
|
||||
tensors_to_landmarks[Output<NormalizedLandmarkList>(kNormLandmarksTag)];
|
||||
|
||||
// Decodes the segmentation tensor into a mask image with pixel values in
|
||||
// [0, 1] (1 for person and 0 for background).
|
||||
auto& tensors_to_segmentation =
|
||||
graph.AddNode("TensorsToSegmentationCalculator");
|
||||
ConfigureTensorsToSegmentationCalculator(
|
||||
&tensors_to_segmentation
|
||||
.GetOptions<mediapipe::TensorsToSegmentationCalculatorOptions>());
|
||||
ensured_segmentation_tensors >> tensors_to_segmentation.In(kTensorsTag);
|
||||
auto segmentation_mask = tensors_to_segmentation[Output<Image>(kMaskTag)];
|
||||
|
||||
// Refines landmarks with the heatmap tensor.
|
||||
auto& refine_landmarks_from_heatmap =
|
||||
graph.AddNode("RefineLandmarksFromHeatmapCalculator");
|
||||
ConfigureRefineLandmarksFromHeatmapCalculator(
|
||||
&refine_landmarks_from_heatmap.GetOptions<
|
||||
mediapipe::RefineLandmarksFromHeatmapCalculatorOptions>());
|
||||
ensured_heatmap_tensors >> refine_landmarks_from_heatmap.In(kTensorsTag);
|
||||
landmarks >> refine_landmarks_from_heatmap.In(kNormLandmarksTag);
|
||||
auto landmarks_from_heatmap =
|
||||
refine_landmarks_from_heatmap[Output<NormalizedLandmarkList>(
|
||||
kNormLandmarksTag)];
|
||||
|
||||
// Splits the landmarks into two sets: the actual pose landmarks and the
|
||||
// auxiliary landmarks.
|
||||
auto& split_normalized_landmark_list =
|
||||
graph.AddNode("SplitNormalizedLandmarkListCalculator");
|
||||
ConfigureSplitNormalizedLandmarkListCalculator(
|
||||
&split_normalized_landmark_list
|
||||
.GetOptions<mediapipe::SplitVectorCalculatorOptions>());
|
||||
landmarks_from_heatmap >> split_normalized_landmark_list.In("");
|
||||
auto normalized_landmarks = split_normalized_landmark_list.Out("")[0]
|
||||
.Cast<NormalizedLandmarkList>();
|
||||
auto normalized_auxiliary_landmarks =
|
||||
split_normalized_landmark_list.Out("")[1]
|
||||
.Cast<NormalizedLandmarkList>();
|
||||
|
||||
// Decodes the world-landmark tensors into a vector of world landmarks.
|
||||
auto& tensors_to_world_landmarks =
|
||||
graph.AddNode("TensorsToLandmarksCalculator");
|
||||
ConfigureTensorsToLandmarksCalculator(
|
||||
image_tensor_specs, /* normalize = */ false,
|
||||
/* sigmoid_activation= */ false,
|
||||
&tensors_to_world_landmarks
|
||||
.GetOptions<mediapipe::TensorsToLandmarksCalculatorOptions>());
|
||||
ensured_world_landmark_tensors >>
|
||||
tensors_to_world_landmarks.In(kTensorsTag);
|
||||
auto world_landmarks =
|
||||
tensors_to_world_landmarks[Output<LandmarkList>(kLandmarksTag)];
|
||||
|
||||
// Keeps only the actual world landmarks.
|
||||
auto& split_landmark_list = graph.AddNode("SplitLandmarkListCalculator");
|
||||
ConfigureSplitLandmarkListCalculator(
|
||||
&split_landmark_list
|
||||
.GetOptions<mediapipe::SplitVectorCalculatorOptions>());
|
||||
world_landmarks >> split_landmark_list.In("");
|
||||
auto split_landmarks = split_landmark_list.Out(0);
|
||||
|
||||
// Reuses the visibility and presence field in pose landmarks for the world
|
||||
// landmarks.
|
||||
auto& visibility_copy = graph.AddNode("VisibilityCopyCalculator");
|
||||
ConfigureVisibilityCopyCalculator(
|
||||
&visibility_copy
|
||||
.GetOptions<mediapipe::VisibilityCopyCalculatorOptions>());
|
||||
split_landmarks >> visibility_copy.In(kLandmarksToTag);
|
||||
normalized_landmarks >> visibility_copy.In(kNormLandmarksFromTag);
|
||||
auto world_landmarks_with_visibility =
|
||||
visibility_copy[Output<LandmarkList>(kLandmarksToTag)];
|
||||
|
||||
// Landmarks to Detections.
|
||||
auto& landmarks_to_detection =
|
||||
graph.AddNode("LandmarksToDetectionCalculator");
|
||||
landmarks >> landmarks_to_detection.In(kNormLandmarksTag);
|
||||
auto detection = landmarks_to_detection.Out(kDetectionTag);
|
||||
|
||||
// Detections to Rects.
|
||||
auto& detection_to_rects = graph.AddNode("DetectionsToRectsCalculator");
|
||||
image_size >> detection_to_rects.In(kImageSizeTag);
|
||||
detection >> detection_to_rects.In(kDetectionTag);
|
||||
auto norm_rect = detection_to_rects.Out(kNormRectTag);
|
||||
|
||||
// Expands the pose rectangle so that in the next video frame it's likely to
|
||||
// still contain the pose even with some motion.
|
||||
auto& pose_rect_transformation =
|
||||
graph.AddNode("RectTransformationCalculator");
|
||||
image_size >> pose_rect_transformation.In(kImageSizeTag);
|
||||
norm_rect >> pose_rect_transformation.In(kNormRectTag);
|
||||
auto pose_rect_next_frame =
|
||||
pose_rect_transformation[Output<NormalizedRect>("")];
|
||||
|
||||
return {{
|
||||
/* pose_landmarks= */ normalized_landmarks,
|
||||
/* world_pose_landmarks= */ world_landmarks_with_visibility,
|
||||
/* auxiliary_pose_landmarks= */ normalized_auxiliary_landmarks,
|
||||
/* pose_rect_next_frame= */ pose_rect_next_frame,
|
||||
/* pose_presence= */ pose_presence,
|
||||
/* pose_presence_score= */ pose_presence_score,
|
||||
/* segmentation_mask= */ segmentation_mask,
|
||||
}};
|
||||
}
|
||||
};
|
||||
|
||||
// clang-format off
|
||||
REGISTER_MEDIAPIPE_GRAPH(
|
||||
::mediapipe::tasks::vision::pose_landmarker::SinglePoseLandmarksDetectorGraph); // NOLINT
|
||||
// clang-format on
|
||||
|
||||
// A "mediapipe.tasks.vision.pose_landmarker.MultiplePoseLandmarksDetectorGraph"
|
||||
// performs multi pose landmark detection.
|
||||
// - Accepts CPU input image and a vector of pose rect RoIs to detect the
|
||||
// multiple poses landmarks enclosed by the RoIs. Output vectors of
|
||||
// pose landmarks related results, where each element in the vectors
|
||||
// corresponds to the result of the same pose.
|
||||
//
|
||||
// Inputs:
|
||||
// IMAGE - Image
|
||||
// Image to perform detection on.
|
||||
// NORM_RECT - std::vector<NormalizedRect>
|
||||
// A vector of multiple pose rects enclosing the pose RoI to perform
|
||||
// landmarks detection on.
|
||||
//
|
||||
//
|
||||
// Outputs:
|
||||
// LANDMARKS: - std::vector<NormalizedLandmarkList>
|
||||
// Vector of detected pose landmarks.
|
||||
// WORLD_LANDMARKS - std::vector<LandmarkList>
|
||||
// Vector of detected pose landmarks in world coordinates.
|
||||
// AUXILIARY_LANDMARKS - std::vector<NormalizedLandmarkList>
|
||||
// Vector of detected pose auxiliary landmarks.
|
||||
// POSE_RECT_NEXT_FRAME - std::vector<NormalizedRect>
|
||||
// Vector of the predicted rects enclosing the same pose RoI for landmark
|
||||
// detection on the next frame.
|
||||
// PRESENCE - std::vector<bool>
|
||||
// Vector of boolean value indicates whether the pose is present.
|
||||
// PRESENCE_SCORE - std::vector<float>
|
||||
// Vector of float value indicates the probability that the pose is present.
|
||||
// SEGMENTATION_MASK - std::vector<Image>
|
||||
// Vector of segmentation masks.
|
||||
//
|
||||
// Example:
|
||||
// node {
|
||||
// calculator:
|
||||
// "mediapipe.tasks.vision.pose_landmarker.MultiplePoseLandmarksDetectorGraph"
|
||||
// input_stream: "IMAGE:input_image"
|
||||
// input_stream: "POSE_RECT:pose_rect"
|
||||
// output_stream: "LANDMARKS:pose_landmarks"
|
||||
// output_stream: "WORLD_LANDMARKS:world_pose_landmarks"
|
||||
// output_stream: "AUXILIARY_LANDMARKS:auxiliary_landmarks"
|
||||
// output_stream: "POSE_RECT_NEXT_FRAME:pose_rect_next_frame"
|
||||
// output_stream: "PRESENCE:pose_presence"
|
||||
// output_stream: "PRESENCE_SCORE:pose_presence_score"
|
||||
// output_stream: "SEGMENTATION_MASK:segmentation_mask"
|
||||
// options {
|
||||
// [mediapipe.tasks.vision.pose_landmarker.proto.PoseLandmarksDetectorGraphOptions.ext]
|
||||
// {
|
||||
// base_options {
|
||||
// model_asset {
|
||||
// file_name: "pose_landmark_lite.tflite"
|
||||
// }
|
||||
// }
|
||||
// min_detection_confidence: 0.5
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
class MultiplePoseLandmarksDetectorGraph : public core::ModelTaskGraph {
|
||||
public:
|
||||
absl::StatusOr<CalculatorGraphConfig> GetConfig(
|
||||
SubgraphContext* sc) override {
|
||||
Graph graph;
|
||||
ASSIGN_OR_RETURN(
|
||||
auto pose_landmark_detection_outputs,
|
||||
BuildPoseLandmarksDetectorGraph(
|
||||
sc->Options<PoseLandmarksDetectorGraphOptions>(),
|
||||
graph[Input<Image>(kImageTag)],
|
||||
graph[Input<std::vector<NormalizedRect>>(kNormRectTag)], graph));
|
||||
pose_landmark_detection_outputs.landmark_lists >>
|
||||
graph[Output<std::vector<NormalizedLandmarkList>>(kLandmarksTag)];
|
||||
pose_landmark_detection_outputs.world_landmark_lists >>
|
||||
graph[Output<std::vector<LandmarkList>>(kWorldLandmarksTag)];
|
||||
pose_landmark_detection_outputs.auxiliary_landmark_lists >>
|
||||
graph[Output<std::vector<NormalizedLandmarkList>>(kAuxLandmarksTag)];
|
||||
pose_landmark_detection_outputs.pose_rects_next_frame >>
|
||||
graph[Output<std::vector<NormalizedRect>>(kPoseRectsNextFrameTag)];
|
||||
pose_landmark_detection_outputs.presences >>
|
||||
graph[Output<std::vector<bool>>(kPresenceTag)];
|
||||
pose_landmark_detection_outputs.presence_scores >>
|
||||
graph[Output<std::vector<float>>(kPresenceScoreTag)];
|
||||
pose_landmark_detection_outputs.segmentation_masks >>
|
||||
graph[Output<std::vector<Image>>(kSegmentationMaskTag)];
|
||||
|
||||
return graph.GetConfig();
|
||||
}
|
||||
|
||||
private:
|
||||
absl::StatusOr<PoseLandmarkerOutputs> BuildPoseLandmarksDetectorGraph(
|
||||
const PoseLandmarksDetectorGraphOptions& subgraph_options,
|
||||
Source<Image> image_in,
|
||||
Source<std::vector<NormalizedRect>> multi_pose_rects, Graph& graph) {
|
||||
auto& begin_loop_multi_pose_rects =
|
||||
graph.AddNode("BeginLoopNormalizedRectCalculator");
|
||||
image_in >> begin_loop_multi_pose_rects.In("CLONE");
|
||||
multi_pose_rects >> begin_loop_multi_pose_rects.In("ITERABLE");
|
||||
auto batch_end = begin_loop_multi_pose_rects.Out("BATCH_END");
|
||||
auto image = begin_loop_multi_pose_rects.Out("CLONE");
|
||||
auto pose_rect = begin_loop_multi_pose_rects.Out("ITEM");
|
||||
|
||||
auto& pose_landmark_subgraph = graph.AddNode(
|
||||
"mediapipe.tasks.vision.pose_landmarker."
|
||||
"SinglePoseLandmarksDetectorGraph");
|
||||
pose_landmark_subgraph.GetOptions<PoseLandmarksDetectorGraphOptions>()
|
||||
.CopyFrom(subgraph_options);
|
||||
image >> pose_landmark_subgraph.In(kImageTag);
|
||||
pose_rect >> pose_landmark_subgraph.In(kNormRectTag);
|
||||
auto landmarks = pose_landmark_subgraph.Out(kLandmarksTag);
|
||||
auto world_landmarks = pose_landmark_subgraph.Out(kWorldLandmarksTag);
|
||||
auto auxiliary_landmarks = pose_landmark_subgraph.Out(kAuxLandmarksTag);
|
||||
auto pose_rect_next_frame =
|
||||
pose_landmark_subgraph.Out(kPoseRectNextFrameTag);
|
||||
auto presence = pose_landmark_subgraph.Out(kPresenceTag);
|
||||
auto presence_score = pose_landmark_subgraph.Out(kPresenceScoreTag);
|
||||
auto segmentation_mask = pose_landmark_subgraph.Out(kSegmentationMaskTag);
|
||||
|
||||
auto& end_loop_landmarks =
|
||||
graph.AddNode("EndLoopNormalizedLandmarkListVectorCalculator");
|
||||
batch_end >> end_loop_landmarks.In(kBatchEndTag);
|
||||
landmarks >> end_loop_landmarks.In(kItemTag);
|
||||
auto landmark_lists =
|
||||
end_loop_landmarks[Output<std::vector<NormalizedLandmarkList>>(
|
||||
kIterableTag)];
|
||||
|
||||
auto& end_loop_world_landmarks =
|
||||
graph.AddNode("EndLoopLandmarkListVectorCalculator");
|
||||
batch_end >> end_loop_world_landmarks.In(kBatchEndTag);
|
||||
world_landmarks >> end_loop_world_landmarks.In(kItemTag);
|
||||
auto world_landmark_lists =
|
||||
end_loop_world_landmarks[Output<std::vector<LandmarkList>>(
|
||||
kIterableTag)];
|
||||
|
||||
auto& end_loop_auxiliary_landmarks =
|
||||
graph.AddNode("EndLoopNormalizedLandmarkListVectorCalculator");
|
||||
batch_end >> end_loop_auxiliary_landmarks.In(kBatchEndTag);
|
||||
auxiliary_landmarks >> end_loop_auxiliary_landmarks.In(kItemTag);
|
||||
auto auxiliary_landmark_lists = end_loop_auxiliary_landmarks
|
||||
[Output<std::vector<NormalizedLandmarkList>>(kIterableTag)];
|
||||
|
||||
auto& end_loop_rects_next_frame =
|
||||
graph.AddNode("EndLoopNormalizedRectCalculator");
|
||||
batch_end >> end_loop_rects_next_frame.In(kBatchEndTag);
|
||||
pose_rect_next_frame >> end_loop_rects_next_frame.In(kItemTag);
|
||||
auto pose_rects_next_frame =
|
||||
end_loop_rects_next_frame[Output<std::vector<NormalizedRect>>(
|
||||
kIterableTag)];
|
||||
|
||||
auto& end_loop_presence = graph.AddNode("EndLoopBooleanCalculator");
|
||||
batch_end >> end_loop_presence.In(kBatchEndTag);
|
||||
presence >> end_loop_presence.In(kItemTag);
|
||||
auto presences = end_loop_presence[Output<std::vector<bool>>(kIterableTag)];
|
||||
|
||||
auto& end_loop_presence_score = graph.AddNode("EndLoopFloatCalculator");
|
||||
batch_end >> end_loop_presence_score.In(kBatchEndTag);
|
||||
presence_score >> end_loop_presence_score.In(kItemTag);
|
||||
auto presence_scores =
|
||||
end_loop_presence_score[Output<std::vector<float>>(kIterableTag)];
|
||||
|
||||
auto& end_loop_segmentation_mask = graph.AddNode("EndLoopImageCalculator");
|
||||
batch_end >> end_loop_segmentation_mask.In(kBatchEndTag);
|
||||
segmentation_mask >> end_loop_segmentation_mask.In(kItemTag);
|
||||
auto segmentation_masks =
|
||||
end_loop_segmentation_mask[Output<std::vector<Image>>(kIterableTag)];
|
||||
|
||||
return {{
|
||||
/* landmark_lists= */ landmark_lists,
|
||||
/* world_landmark_lists= */ world_landmark_lists,
|
||||
/* auxiliary_landmark_lists= */ auxiliary_landmark_lists,
|
||||
/* pose_rects_next_frame= */ pose_rects_next_frame,
|
||||
/* presences= */ presences,
|
||||
/* presence_scores= */ presence_scores,
|
||||
/* segmentation_masks= */ segmentation_masks,
|
||||
}};
|
||||
}
|
||||
};
|
||||
|
||||
// clang-format off
|
||||
REGISTER_MEDIAPIPE_GRAPH(
|
||||
::mediapipe::tasks::vision::pose_landmarker::MultiplePoseLandmarksDetectorGraph); // NOLINT
|
||||
// clang-format on
|
||||
|
||||
} // namespace pose_landmarker
|
||||
} // namespace vision
|
||||
} // namespace tasks
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,371 @@
|
||||
/* Copyright 2023 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 <iostream>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
#include "absl/flags/flag.h"
|
||||
#include "absl/status/statusor.h"
|
||||
#include "absl/strings/str_format.h"
|
||||
#include "absl/strings/string_view.h"
|
||||
#include "mediapipe/framework/api2/builder.h"
|
||||
#include "mediapipe/framework/api2/port.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/deps/file_path.h"
|
||||
#include "mediapipe/framework/formats/image.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/formats/tensor.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/file_helpers.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/tasks/cc/core/model_resources.h"
|
||||
#include "mediapipe/tasks/cc/core/proto/base_options.pb.h"
|
||||
#include "mediapipe/tasks/cc/core/proto/external_file.pb.h"
|
||||
#include "mediapipe/tasks/cc/core/task_runner.h"
|
||||
#include "mediapipe/tasks/cc/vision/pose_landmarker/proto/pose_landmarks_detector_graph_options.pb.h"
|
||||
#include "mediapipe/tasks/cc/vision/utils/image_utils.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace tasks {
|
||||
namespace vision {
|
||||
namespace pose_landmarker {
|
||||
namespace {
|
||||
|
||||
using ::file::Defaults;
|
||||
using ::file::GetTextProto;
|
||||
using ::mediapipe::NormalizedRect;
|
||||
using ::mediapipe::api2::Input;
|
||||
using ::mediapipe::api2::Output;
|
||||
using ::mediapipe::api2::builder::Graph;
|
||||
using ::mediapipe::api2::builder::Source;
|
||||
using ::mediapipe::file::JoinPath;
|
||||
using ::mediapipe::tasks::core::TaskRunner;
|
||||
using ::mediapipe::tasks::vision::DecodeImageFromFile;
|
||||
using ::mediapipe::tasks::vision::pose_landmarker::proto::
|
||||
PoseLandmarksDetectorGraphOptions;
|
||||
using ::testing::ElementsAreArray;
|
||||
using ::testing::EqualsProto;
|
||||
using ::testing::Pointwise;
|
||||
using ::testing::TestParamInfo;
|
||||
using ::testing::TestWithParam;
|
||||
using ::testing::Values;
|
||||
using ::testing::proto::Approximately;
|
||||
using ::testing::proto::Partially;
|
||||
|
||||
constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/vision/";
|
||||
constexpr char kPoseLandmarkerLiteModel[] = "pose_landmark_lite.tflite";
|
||||
constexpr char kPoseImage[] = "pose.jpg";
|
||||
constexpr char kBurgerImage[] = "burger.jpg";
|
||||
|
||||
constexpr char kImageTag[] = "IMAGE";
|
||||
constexpr char kImageName[] = "image_in";
|
||||
constexpr char kNormRectTag[] = "NORM_RECT";
|
||||
|
||||
constexpr char kPoseRectName[] = "pose_rect_in";
|
||||
|
||||
constexpr char kLandmarksTag[] = "LANDMARKS";
|
||||
constexpr char kLandmarksName[] = "landmarks";
|
||||
constexpr char kWorldLandmarksTag[] = "WORLD_LANDMARKS";
|
||||
constexpr char kWorldLandmarksName[] = "world_landmarks";
|
||||
constexpr char kAuxLandmarksTag[] = "AUXILIARY_LANDMARKS";
|
||||
constexpr char kAuxLandmarksName[] = "auxiliary_landmarks";
|
||||
constexpr char kPoseRectNextFrameTag[] = "POSE_RECT_NEXT_FRAME";
|
||||
constexpr char kPoseRectNextFrameName[] = "pose_rect_next_frame";
|
||||
constexpr char kPoseRectsNextFrameTag[] = "POSE_RECTS_NEXT_FRAME";
|
||||
constexpr char kPoseRectsNextFrameName[] = "pose_rects_next_frame";
|
||||
constexpr char kPresenceTag[] = "PRESENCE";
|
||||
constexpr char kPresenceName[] = "presence";
|
||||
constexpr char kPresenceScoreTag[] = "PRESENCE_SCORE";
|
||||
constexpr char kPresenceScoreName[] = "presence_score";
|
||||
constexpr char kSegmentationMaskTag[] = "SEGMENTATION_MASK";
|
||||
constexpr char kSegmentationMaskName[] = "segmentation_mask";
|
||||
|
||||
// Expected pose landmarks positions, in text proto format.
|
||||
constexpr char kExpectedPoseLandmarksFilename[] =
|
||||
"expected_pose_landmarks.prototxt";
|
||||
|
||||
constexpr float kLiteModelFractionDiff = 0.05; // percentage
|
||||
constexpr float kAbsMargin = 0.03;
|
||||
|
||||
// Helper function to create a Single Pose Landmark TaskRunner.
|
||||
absl::StatusOr<std::unique_ptr<TaskRunner>> CreateSinglePoseTaskRunner(
|
||||
absl::string_view model_name) {
|
||||
Graph graph;
|
||||
|
||||
auto& pose_landmark_detection = graph.AddNode(
|
||||
"mediapipe.tasks.vision.pose_landmarker."
|
||||
"SinglePoseLandmarksDetectorGraph");
|
||||
|
||||
auto options = std::make_unique<PoseLandmarksDetectorGraphOptions>();
|
||||
options->mutable_base_options()->mutable_model_asset()->set_file_name(
|
||||
JoinPath("./", kTestDataDirectory, model_name));
|
||||
pose_landmark_detection.GetOptions<PoseLandmarksDetectorGraphOptions>().Swap(
|
||||
options.get());
|
||||
|
||||
graph[Input<Image>(kImageTag)].SetName(kImageName) >>
|
||||
pose_landmark_detection.In(kImageTag);
|
||||
graph[Input<NormalizedRect>(kNormRectTag)].SetName(kPoseRectName) >>
|
||||
pose_landmark_detection.In(kNormRectTag);
|
||||
|
||||
pose_landmark_detection.Out(kLandmarksTag).SetName(kLandmarksName) >>
|
||||
graph[Output<NormalizedLandmarkList>(kLandmarksTag)];
|
||||
pose_landmark_detection.Out(kWorldLandmarksTag)
|
||||
.SetName(kWorldLandmarksName) >>
|
||||
graph[Output<LandmarkList>(kWorldLandmarksTag)];
|
||||
pose_landmark_detection.Out(kAuxLandmarksTag).SetName(kAuxLandmarksName) >>
|
||||
graph[Output<LandmarkList>(kAuxLandmarksTag)];
|
||||
pose_landmark_detection.Out(kPresenceTag).SetName(kPresenceName) >>
|
||||
graph[Output<bool>(kPresenceTag)];
|
||||
pose_landmark_detection.Out(kPresenceScoreTag).SetName(kPresenceScoreName) >>
|
||||
graph[Output<float>(kPresenceScoreTag)];
|
||||
pose_landmark_detection.Out(kSegmentationMaskTag)
|
||||
.SetName(kSegmentationMaskName) >>
|
||||
graph[Output<Image>(kSegmentationMaskTag)];
|
||||
pose_landmark_detection.Out(kPoseRectNextFrameTag)
|
||||
.SetName(kPoseRectNextFrameName) >>
|
||||
graph[Output<NormalizedRect>(kPoseRectNextFrameTag)];
|
||||
|
||||
return TaskRunner::Create(
|
||||
graph.GetConfig(),
|
||||
absl::make_unique<tflite_shims::ops::builtin::BuiltinOpResolver>());
|
||||
}
|
||||
|
||||
// Helper function to create a Multi Pose Landmark TaskRunner.
|
||||
absl::StatusOr<std::unique_ptr<TaskRunner>> CreateMultiPoseTaskRunner(
|
||||
absl::string_view model_name) {
|
||||
Graph graph;
|
||||
|
||||
auto& multi_pose_landmark_detection = graph.AddNode(
|
||||
"mediapipe.tasks.vision.pose_landmarker."
|
||||
"MultiplePoseLandmarksDetectorGraph");
|
||||
|
||||
auto options = std::make_unique<PoseLandmarksDetectorGraphOptions>();
|
||||
options->mutable_base_options()->mutable_model_asset()->set_file_name(
|
||||
JoinPath("./", kTestDataDirectory, model_name));
|
||||
multi_pose_landmark_detection.GetOptions<PoseLandmarksDetectorGraphOptions>()
|
||||
.Swap(options.get());
|
||||
|
||||
graph[Input<Image>(kImageTag)].SetName(kImageName) >>
|
||||
multi_pose_landmark_detection.In(kImageTag);
|
||||
graph[Input<std::vector<NormalizedRect>>(kNormRectTag)].SetName(
|
||||
kPoseRectName) >>
|
||||
multi_pose_landmark_detection.In(kNormRectTag);
|
||||
|
||||
multi_pose_landmark_detection.Out(kLandmarksTag).SetName(kLandmarksName) >>
|
||||
graph[Output<std::vector<NormalizedLandmarkList>>(kLandmarksTag)];
|
||||
multi_pose_landmark_detection.Out(kWorldLandmarksTag)
|
||||
.SetName(kWorldLandmarksName) >>
|
||||
graph[Output<std::vector<LandmarkList>>(kWorldLandmarksTag)];
|
||||
multi_pose_landmark_detection.Out(kAuxLandmarksTag)
|
||||
.SetName(kAuxLandmarksName) >>
|
||||
graph[Output<std::vector<NormalizedLandmarkList>>(kAuxLandmarksTag)];
|
||||
multi_pose_landmark_detection.Out(kPresenceTag).SetName(kPresenceName) >>
|
||||
graph[Output<std::vector<bool>>(kPresenceTag)];
|
||||
multi_pose_landmark_detection.Out(kPresenceScoreTag)
|
||||
.SetName(kPresenceScoreName) >>
|
||||
graph[Output<std::vector<float>>(kPresenceScoreTag)];
|
||||
multi_pose_landmark_detection.Out(kSegmentationMaskTag)
|
||||
.SetName(kSegmentationMaskName) >>
|
||||
graph[Output<std::vector<Image>>(kSegmentationMaskTag)];
|
||||
multi_pose_landmark_detection.Out(kPoseRectsNextFrameTag)
|
||||
.SetName(kPoseRectsNextFrameName) >>
|
||||
graph[Output<std::vector<NormalizedRect>>(kPoseRectsNextFrameTag)];
|
||||
|
||||
return TaskRunner::Create(
|
||||
graph.GetConfig(),
|
||||
absl::make_unique<tflite_shims::ops::builtin::BuiltinOpResolver>());
|
||||
}
|
||||
|
||||
NormalizedLandmarkList GetExpectedLandmarkList(absl::string_view filename) {
|
||||
NormalizedLandmarkList expected_landmark_list;
|
||||
MP_EXPECT_OK(GetTextProto(file::JoinPath("./", kTestDataDirectory, filename),
|
||||
&expected_landmark_list, Defaults()));
|
||||
return expected_landmark_list;
|
||||
}
|
||||
|
||||
// Struct holding the parameters for parameterized PoseLandmarkerTest
|
||||
// class.
|
||||
struct SinglePoseTestParams {
|
||||
// The name of this test, for convenience when displaying test results.
|
||||
std::string test_name;
|
||||
// The filename of the model to test.
|
||||
std::string input_model_name;
|
||||
// The filename of the test image.
|
||||
std::string test_image_name;
|
||||
// RoI on image to detect pose.
|
||||
NormalizedRect pose_rect;
|
||||
// Expected pose presence value.
|
||||
bool expected_presence;
|
||||
// The expected output landmarks positions in pixels coornidates.
|
||||
std::optional<NormalizedLandmarkList> expected_landmarks;
|
||||
// The expected segmentation mask.
|
||||
Image expected_segmentation_mask;
|
||||
// The max value difference between expected_positions and detected positions.
|
||||
float landmarks_diff_threshold;
|
||||
};
|
||||
|
||||
struct MultiPoseTestParams {
|
||||
// The name of this test, for convenience when displaying test results.
|
||||
std::string test_name;
|
||||
// The filename of the model to test.
|
||||
std::string input_model_name;
|
||||
// The filename of the test image.
|
||||
std::string test_image_name;
|
||||
// RoIs on image to detect poses.
|
||||
std::vector<NormalizedRect> pose_rects;
|
||||
// Expected pose presence values.
|
||||
std::vector<bool> expected_presences;
|
||||
// The expected output landmarks positions in pixels coornidates.
|
||||
std::vector<NormalizedLandmarkList> expected_landmark_lists;
|
||||
// The expected segmentation_mask Image.
|
||||
std::vector<Image> expected_segmentation_masks;
|
||||
// The max value difference between expected_positions and detected positions.
|
||||
float landmarks_diff_threshold;
|
||||
};
|
||||
|
||||
// Helper function to construct NormalizeRect proto.
|
||||
NormalizedRect MakePoseRect(float x_center, float y_center, float width,
|
||||
float height, float rotation) {
|
||||
NormalizedRect pose_rect;
|
||||
pose_rect.set_x_center(x_center);
|
||||
pose_rect.set_y_center(y_center);
|
||||
pose_rect.set_width(width);
|
||||
pose_rect.set_height(height);
|
||||
pose_rect.set_rotation(rotation);
|
||||
return pose_rect;
|
||||
}
|
||||
|
||||
class PoseLandmarkerTest : public testing::TestWithParam<SinglePoseTestParams> {
|
||||
};
|
||||
|
||||
TEST_P(PoseLandmarkerTest, Succeeds) {
|
||||
MP_ASSERT_OK_AND_ASSIGN(
|
||||
Image image, DecodeImageFromFile(JoinPath("./", kTestDataDirectory,
|
||||
GetParam().test_image_name)));
|
||||
MP_ASSERT_OK_AND_ASSIGN(auto task_runner, CreateSinglePoseTaskRunner(
|
||||
GetParam().input_model_name));
|
||||
|
||||
auto output_packets = task_runner->Process(
|
||||
{{kImageName, MakePacket<Image>(std::move(image))},
|
||||
{kPoseRectName,
|
||||
MakePacket<NormalizedRect>(std::move(GetParam().pose_rect))}});
|
||||
MP_ASSERT_OK(output_packets);
|
||||
|
||||
const bool presence = (*output_packets)[kPresenceName].Get<bool>();
|
||||
ASSERT_EQ(presence, GetParam().expected_presence);
|
||||
|
||||
if (presence) {
|
||||
const NormalizedLandmarkList landmarks =
|
||||
(*output_packets)[kLandmarksName].Get<NormalizedLandmarkList>();
|
||||
|
||||
if (GetParam().expected_landmarks.has_value()) {
|
||||
const NormalizedLandmarkList& expected_landmarks =
|
||||
GetParam().expected_landmarks.value();
|
||||
|
||||
EXPECT_THAT(
|
||||
landmarks,
|
||||
Approximately(Partially(EqualsProto(expected_landmarks)),
|
||||
/*margin=*/kAbsMargin,
|
||||
/*fraction=*/GetParam().landmarks_diff_threshold));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
class MultiPoseLandmarkerTest
|
||||
: public testing::TestWithParam<MultiPoseTestParams> {};
|
||||
|
||||
TEST_P(MultiPoseLandmarkerTest, Succeeds) {
|
||||
MP_ASSERT_OK_AND_ASSIGN(
|
||||
Image image, DecodeImageFromFile(JoinPath("./", kTestDataDirectory,
|
||||
GetParam().test_image_name)));
|
||||
MP_ASSERT_OK_AND_ASSIGN(
|
||||
auto task_runner, CreateMultiPoseTaskRunner(GetParam().input_model_name));
|
||||
|
||||
auto output_packets = task_runner->Process(
|
||||
{{kImageName, MakePacket<Image>(std::move(image))},
|
||||
{kPoseRectName, MakePacket<std::vector<NormalizedRect>>(
|
||||
std::move(GetParam().pose_rects))}});
|
||||
MP_ASSERT_OK(output_packets);
|
||||
|
||||
const std::vector<bool>& presences =
|
||||
(*output_packets)[kPresenceName].Get<std::vector<bool>>();
|
||||
const std::vector<NormalizedLandmarkList>& landmark_lists =
|
||||
(*output_packets)[kLandmarksName]
|
||||
.Get<std::vector<NormalizedLandmarkList>>();
|
||||
|
||||
EXPECT_THAT(presences, ElementsAreArray(GetParam().expected_presences));
|
||||
|
||||
EXPECT_THAT(
|
||||
landmark_lists,
|
||||
Pointwise(Approximately(Partially(EqualsProto()),
|
||||
/*margin=*/kAbsMargin,
|
||||
/*fraction=*/GetParam().landmarks_diff_threshold),
|
||||
GetParam().expected_landmark_lists));
|
||||
}
|
||||
// TODO: Add additional tests for MP Tasks Pose Graphs.
|
||||
// PoseRects below are based on result from PoseDetectorGraph,
|
||||
// mediapipe/tasks/testdata/vision/pose_expected_expanded_rect.pbtxt.
|
||||
INSTANTIATE_TEST_SUITE_P(
|
||||
PoseLandmarkerTest, PoseLandmarkerTest,
|
||||
Values(
|
||||
SinglePoseTestParams{
|
||||
.test_name = "PoseLandmarkerLiteModel",
|
||||
.input_model_name = kPoseLandmarkerLiteModel,
|
||||
.test_image_name = kPoseImage,
|
||||
.pose_rect = MakePoseRect(0.5450622, 0.31605977, 0.5196669,
|
||||
0.77911085, 0.50149304),
|
||||
.expected_presence = true,
|
||||
.expected_landmarks =
|
||||
GetExpectedLandmarkList(kExpectedPoseLandmarksFilename),
|
||||
.landmarks_diff_threshold = kLiteModelFractionDiff},
|
||||
SinglePoseTestParams{
|
||||
.test_name = "PoseLandmarkerLiteModelNoPose",
|
||||
.input_model_name = kPoseLandmarkerLiteModel,
|
||||
.test_image_name = kBurgerImage,
|
||||
.pose_rect = MakePoseRect(0.5450622, 0.31605977, 0.5196669,
|
||||
0.77911085, 0.50149304),
|
||||
.expected_presence = false,
|
||||
.expected_landmarks = std::nullopt,
|
||||
.landmarks_diff_threshold = kLiteModelFractionDiff}),
|
||||
[](const TestParamInfo<PoseLandmarkerTest::ParamType>& info) {
|
||||
return info.param.test_name;
|
||||
});
|
||||
|
||||
INSTANTIATE_TEST_SUITE_P(
|
||||
MultiPoseLandmarkerTest, MultiPoseLandmarkerTest,
|
||||
Values(MultiPoseTestParams{
|
||||
.test_name = "MultiPoseLandmarkerLiteModel",
|
||||
.input_model_name = kPoseLandmarkerLiteModel,
|
||||
.test_image_name = kPoseImage,
|
||||
.pose_rects = {MakePoseRect(0.5450622, 0.31605977, 0.5196669,
|
||||
0.77911085, 0.50149304)},
|
||||
.expected_presences = {true},
|
||||
.expected_landmark_lists = {GetExpectedLandmarkList(
|
||||
kExpectedPoseLandmarksFilename)},
|
||||
.landmarks_diff_threshold = kLiteModelFractionDiff,
|
||||
}),
|
||||
[](const TestParamInfo<MultiPoseLandmarkerTest::ParamType>& info) {
|
||||
return info.param.test_name;
|
||||
});
|
||||
|
||||
} // namespace
|
||||
} // namespace pose_landmarker
|
||||
} // namespace vision
|
||||
} // namespace tasks
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,43 @@
|
||||
# Copyright 2023 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.
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_proto_library")
|
||||
|
||||
package(default_visibility = [
|
||||
"//mediapipe/tasks:internal",
|
||||
])
|
||||
|
||||
licenses(["notice"])
|
||||
|
||||
mediapipe_proto_library(
|
||||
name = "pose_landmarks_detector_graph_options_proto",
|
||||
srcs = ["pose_landmarks_detector_graph_options.proto"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_options_proto",
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/tasks/cc/core/proto:base_options_proto",
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_proto_library(
|
||||
name = "pose_landmarker_graph_options_proto",
|
||||
srcs = ["pose_landmarker_graph_options.proto"],
|
||||
deps = [
|
||||
":pose_landmarks_detector_graph_options_proto",
|
||||
"//mediapipe/framework:calculator_options_proto",
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/tasks/cc/core/proto:base_options_proto",
|
||||
"//mediapipe/tasks/cc/vision/pose_detector/proto:pose_detector_graph_options_proto",
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,48 @@
|
||||
/* Copyright 2023 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.
|
||||
==============================================================================*/
|
||||
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe.tasks.vision.pose_landmarker.proto;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
import "mediapipe/framework/calculator_options.proto";
|
||||
import "mediapipe/tasks/cc/core/proto/base_options.proto";
|
||||
import "mediapipe/tasks/cc/vision/pose_detector/proto/pose_detector_graph_options.proto";
|
||||
import "mediapipe/tasks/cc/vision/pose_landmarker/proto/pose_landmarks_detector_graph_options.proto";
|
||||
|
||||
option java_package = "com.google.mediapipe.tasks.vision.poselandmarker.proto";
|
||||
option java_outer_classname = "PoseLandmarkerGraphOptionsProto";
|
||||
|
||||
message PoseLandmarkerGraphOptions {
|
||||
extend mediapipe.CalculatorOptions {
|
||||
optional PoseLandmarkerGraphOptions ext = 516587230;
|
||||
}
|
||||
// Base options for configuring Task library, such as specifying the TfLite
|
||||
// model file with metadata, accelerator options, etc.
|
||||
optional core.proto.BaseOptions base_options = 1;
|
||||
|
||||
// Options for pose detector graph.
|
||||
optional pose_detector.proto.PoseDetectorGraphOptions
|
||||
pose_detector_graph_options = 2;
|
||||
|
||||
// Options for pose landmarks detector graph.
|
||||
optional PoseLandmarksDetectorGraphOptions
|
||||
pose_landmarks_detector_graph_options = 3;
|
||||
|
||||
// Minimum confidence for pose landmarks tracking to be considered
|
||||
// successfully.
|
||||
optional float min_tracking_confidence = 4 [default = 0.5];
|
||||
}
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
/* Copyright 2023 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.
|
||||
==============================================================================*/
|
||||
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe.tasks.vision.pose_landmarker.proto;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
import "mediapipe/framework/calculator_options.proto";
|
||||
import "mediapipe/tasks/cc/core/proto/base_options.proto";
|
||||
|
||||
option java_package = "com.google.mediapipe.tasks.vision.poselandmarker.proto";
|
||||
option java_outer_classname = "PoseLandmarksDetectorGraphOptionsProto";
|
||||
|
||||
message PoseLandmarksDetectorGraphOptions {
|
||||
extend mediapipe.CalculatorOptions {
|
||||
optional PoseLandmarksDetectorGraphOptions ext = 518928384;
|
||||
}
|
||||
// Base options for configuring MediaPipe Tasks, such as specifying the TfLite
|
||||
// model file with metadata, accelerator options, etc.
|
||||
optional core.proto.BaseOptions base_options = 1;
|
||||
|
||||
// Minimum confidence value ([0.0, 1.0]) for pose presence score to be
|
||||
// considered successfully detecting a pose in the image.
|
||||
optional float min_detection_confidence = 2 [default = 0.5];
|
||||
}
|
||||
@@ -191,8 +191,9 @@ absl::StatusOr<ImageTensorSpecs> BuildInputImageTensorSpecs(
|
||||
MediaPipeTasksStatus::kInvalidInputTensorDimensionsError);
|
||||
}
|
||||
|
||||
size_t byte_depth =
|
||||
tensor_type == tflite::TensorType_FLOAT32 ? sizeof(float) : sizeof(uint8);
|
||||
size_t byte_depth = tensor_type == tflite::TensorType_FLOAT32
|
||||
? sizeof(float)
|
||||
: sizeof(uint8_t);
|
||||
int bytes_size = byte_depth * batch * height * width * depth;
|
||||
// Sanity checks.
|
||||
if (tensor_type == tflite::TensorType_FLOAT32) {
|
||||
|
||||
@@ -24,7 +24,6 @@ objc_library(
|
||||
"//mediapipe/tasks/cc:common",
|
||||
"//mediapipe/tasks/ios/common:MPPCommon",
|
||||
"@com_google_absl//absl/status",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/strings:cord",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -44,3 +44,13 @@ objc_library(
|
||||
hdrs = ["sources/MPPEmbeddingResult.h"],
|
||||
deps = [":MPPEmbedding"],
|
||||
)
|
||||
|
||||
objc_library(
|
||||
name = "MPPDetection",
|
||||
srcs = ["sources/MPPDetection.m"],
|
||||
hdrs = ["sources/MPPDetection.h"],
|
||||
deps = [
|
||||
":MPPCategory",
|
||||
"//third_party/apple_frameworks:UIKit",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#import <Foundation/Foundation.h>
|
||||
#import <UIKit/UIKit.h>
|
||||
#import "mediapipe/tasks/ios/components/containers/sources/MPPCategory.h"
|
||||
|
||||
NS_ASSUME_NONNULL_BEGIN
|
||||
|
||||
/**
|
||||
* Normalized keypoint represents a point in 2D space with x, y coordinates. x and y are normalized
|
||||
* to [0.0, 1.0] by the image width and height respectively.
|
||||
*/
|
||||
NS_SWIFT_NAME(NormalizedKeypoint)
|
||||
@interface MPPNormalizedKeypoint : NSObject
|
||||
|
||||
/** The (x,y) coordinates location of the normalized keypoint. */
|
||||
@property(nonatomic, readonly) CGPoint location;
|
||||
|
||||
/** The optional label of the normalized keypoint. */
|
||||
@property(nonatomic, readonly, nullable) NSString *label;
|
||||
|
||||
/** The optional score of the normalized keypoint. If score is absent, it will be equal to 0.0. */
|
||||
@property(nonatomic, readonly) float score;
|
||||
|
||||
/**
|
||||
* Initializes a new `MPPNormalizedKeypoint` object with the given location, label and score.
|
||||
* You must pass 0.0 for `score` if it is not present.
|
||||
*
|
||||
* @param location The (x,y) coordinates location of the normalized keypoint.
|
||||
* @param label The optional label of the normalized keypoint.
|
||||
* @param score The optional score of the normalized keypoint. You must pass 0.0 for score if it
|
||||
* is not present.
|
||||
*
|
||||
* @return An instance of `MPPNormalizedKeypoint` initialized with the given given location, label
|
||||
* and score.
|
||||
*/
|
||||
- (instancetype)initWithLocation:(CGPoint)location
|
||||
label:(nullable NSString *)label
|
||||
score:(float)score NS_DESIGNATED_INITIALIZER;
|
||||
|
||||
- (instancetype)init NS_UNAVAILABLE;
|
||||
|
||||
+ (instancetype)new NS_UNAVAILABLE;
|
||||
|
||||
@end
|
||||
|
||||
/** Represents one detected object in the results of `MPPObjectDetector`. */
|
||||
NS_SWIFT_NAME(Detection)
|
||||
@interface MPPDetection : NSObject
|
||||
|
||||
/** An array of `MPPCategory` objects containing the predicted categories. */
|
||||
@property(nonatomic, readonly) NSArray<MPPCategory *> *categories;
|
||||
|
||||
/** The bounding box of the detected object. */
|
||||
@property(nonatomic, readonly) CGRect boundingBox;
|
||||
|
||||
/**
|
||||
* An optional array of `MPPNormalizedKeypoint` objects associated with the detection. Keypoints
|
||||
* represent interesting points related to the detection. For example, the keypoints represent the
|
||||
* eyes, ear and mouth from the from detection model. In template matching detection, e.g. KNIFT,
|
||||
* they can instead represent the feature points for template matching.
|
||||
*/
|
||||
@property(nonatomic, readonly, nullable) NSArray<MPPNormalizedKeypoint *> *keypoints;
|
||||
|
||||
/**
|
||||
* Initializes a new `MPPDetection` object with the given array of categories, bounding box and
|
||||
* optional array of keypoints;
|
||||
*
|
||||
* @param categories A list of `MPPCategory` objects that contain category name, display name,
|
||||
* score, and the label index.
|
||||
* @param boundingBox A `CGRect` that represents the bounding box.
|
||||
* @param keypoints: An optional array of `MPPNormalizedKeypoint` objects associated with the
|
||||
* detection. Keypoints represent interesting points related to the detection. For example, the
|
||||
* keypoints represent the eyes, ear and mouth from the face detection model. In template matching
|
||||
* detection, e.g. KNIFT, they can instead represent the feature points for template matching.
|
||||
*
|
||||
* @return An instance of `MPPDetection` initialized with the given array of categories, bounding
|
||||
* box and `nil` keypoints.
|
||||
*/
|
||||
- (instancetype)initWithCategories:(NSArray<MPPCategory *> *)categories
|
||||
boundingBox:(CGRect)boundingBox
|
||||
keypoints:(nullable NSArray<MPPNormalizedKeypoint *> *)keypoints
|
||||
NS_DESIGNATED_INITIALIZER;
|
||||
|
||||
- (instancetype)init NS_UNAVAILABLE;
|
||||
|
||||
+ (instancetype)new NS_UNAVAILABLE;
|
||||
|
||||
@end
|
||||
|
||||
NS_ASSUME_NONNULL_END
|
||||
@@ -0,0 +1,68 @@
|
||||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#import "mediapipe/tasks/ios/components/containers/sources/MPPDetection.h"
|
||||
|
||||
@implementation MPPNormalizedKeypoint
|
||||
|
||||
- (instancetype)initWithLocation:(CGPoint)location
|
||||
label:(nullable NSString *)label
|
||||
score:(float)score {
|
||||
self = [super init];
|
||||
if (self) {
|
||||
_location = location;
|
||||
_label = label;
|
||||
_score = score;
|
||||
}
|
||||
return self;
|
||||
}
|
||||
|
||||
// TODO: Implement hash
|
||||
|
||||
- (BOOL)isEqual:(nullable id)object {
|
||||
if (!object) {
|
||||
return NO;
|
||||
}
|
||||
|
||||
if (self == object) {
|
||||
return YES;
|
||||
}
|
||||
|
||||
if (![object isKindOfClass:[MPPNormalizedKeypoint class]]) {
|
||||
return NO;
|
||||
}
|
||||
|
||||
MPPNormalizedKeypoint *otherKeypoint = (MPPNormalizedKeypoint *)object;
|
||||
|
||||
return CGPointEqualToPoint(self.location, otherKeypoint.location) &&
|
||||
(self.label == otherKeypoint.label) && (self.score == otherKeypoint.score);
|
||||
}
|
||||
|
||||
@end
|
||||
|
||||
@implementation MPPDetection
|
||||
|
||||
- (instancetype)initWithCategories:(NSArray<MPPCategory *> *)categories
|
||||
boundingBox:(CGRect)boundingBox
|
||||
keypoints:(nullable NSArray<MPPNormalizedKeypoint *> *)keypoints {
|
||||
self = [super init];
|
||||
if (self) {
|
||||
_categories = categories;
|
||||
_boundingBox = boundingBox;
|
||||
_keypoints = keypoints;
|
||||
}
|
||||
return self;
|
||||
}
|
||||
|
||||
@end
|
||||
@@ -61,3 +61,15 @@ objc_library(
|
||||
"//mediapipe/tasks/ios/components/containers:MPPEmbeddingResult",
|
||||
],
|
||||
)
|
||||
|
||||
objc_library(
|
||||
name = "MPPDetectionHelpers",
|
||||
srcs = ["sources/MPPDetection+Helpers.mm"],
|
||||
hdrs = ["sources/MPPDetection+Helpers.h"],
|
||||
deps = [
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/formats:location_data_cc_proto",
|
||||
"//mediapipe/tasks/ios/common/utils:NSStringHelpers",
|
||||
"//mediapipe/tasks/ios/components/containers:MPPDetection",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#import "mediapipe/tasks/ios/components/containers/sources/MPPDetection.h"
|
||||
|
||||
NS_ASSUME_NONNULL_BEGIN
|
||||
|
||||
@interface MPPDetection (Helpers)
|
||||
|
||||
+ (MPPDetection *)detectionWithProto:(const mediapipe::Detection &)detectionProto;
|
||||
|
||||
@end
|
||||
|
||||
NS_ASSUME_NONNULL_END
|
||||
@@ -0,0 +1,83 @@
|
||||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#import "mediapipe/tasks/ios/components/containers/utils/sources/MPPDetection+Helpers.h"
|
||||
#import "mediapipe/framework/formats/location_data.pb.h"
|
||||
|
||||
#import "mediapipe/tasks/ios/common/utils/sources/NSString+Helpers.h"
|
||||
|
||||
static const NSInteger kDefaultCategoryIndex = -1;
|
||||
|
||||
namespace {
|
||||
using DetectionProto = ::mediapipe::Detection;
|
||||
using BoundingBoxProto = ::mediapipe::LocationData::BoundingBox;
|
||||
} // namespace
|
||||
|
||||
@implementation MPPDetection (Helpers)
|
||||
|
||||
+ (MPPDetection *)detectionWithProto:(const DetectionProto &)detectionProto {
|
||||
NSMutableArray<MPPCategory *> *categories =
|
||||
[NSMutableArray arrayWithCapacity:(NSUInteger)detectionProto.score_size()];
|
||||
|
||||
for (int idx = 0; idx < detectionProto.score_size(); ++idx) {
|
||||
NSInteger categoryIndex =
|
||||
detectionProto.label_id_size() > idx ? detectionProto.label_id(idx) : kDefaultCategoryIndex;
|
||||
NSString *categoryName = detectionProto.label_size() > idx
|
||||
? [NSString stringWithCppString:detectionProto.label(idx)]
|
||||
: nil;
|
||||
|
||||
NSString *displayName = detectionProto.display_name_size() > idx
|
||||
? [NSString stringWithCppString:detectionProto.display_name(idx)]
|
||||
: nil;
|
||||
|
||||
[categories addObject:[[MPPCategory alloc] initWithIndex:categoryIndex
|
||||
score:detectionProto.score(idx)
|
||||
categoryName:categoryName
|
||||
displayName:displayName]];
|
||||
}
|
||||
|
||||
CGRect boundingBox = CGRectZero;
|
||||
|
||||
if (detectionProto.location_data().has_bounding_box()) {
|
||||
const BoundingBoxProto &boundingBoxProto = detectionProto.location_data().bounding_box();
|
||||
boundingBox.origin.x = boundingBoxProto.xmin();
|
||||
boundingBox.origin.y = boundingBoxProto.ymin();
|
||||
boundingBox.size.width = boundingBoxProto.width();
|
||||
boundingBox.size.height = boundingBoxProto.height();
|
||||
}
|
||||
|
||||
NSMutableArray<MPPNormalizedKeypoint *> *normalizedKeypoints;
|
||||
|
||||
if (!detectionProto.location_data().relative_keypoints().empty()) {
|
||||
normalizedKeypoints = [NSMutableArray
|
||||
arrayWithCapacity:(NSUInteger)detectionProto.location_data().relative_keypoints_size()];
|
||||
for (const auto &keypoint : detectionProto.location_data().relative_keypoints()) {
|
||||
NSString *label = keypoint.has_keypoint_label()
|
||||
? [NSString stringWithCppString:keypoint.keypoint_label()]
|
||||
: nil;
|
||||
CGPoint location = CGPointMake(keypoint.x(), keypoint.y());
|
||||
float score = keypoint.has_score() ? keypoint.score() : 0.0f;
|
||||
|
||||
[normalizedKeypoints addObject:[[MPPNormalizedKeypoint alloc] initWithLocation:location
|
||||
label:label
|
||||
score:score]];
|
||||
}
|
||||
}
|
||||
|
||||
return [[MPPDetection alloc] initWithCategories:categories
|
||||
boundingBox:boundingBox
|
||||
keypoints:normalizedKeypoints];
|
||||
}
|
||||
|
||||
@end
|
||||
@@ -107,18 +107,18 @@ using ::mediapipe::InputStreamInfo;
|
||||
for (NSString *inputStream in self.inputStreams) {
|
||||
graphConfig.add_input_stream(inputStream.cppString);
|
||||
|
||||
NSString *strippedInputStream = [MPPTaskInfo stripTagIndex:inputStream];
|
||||
flowLimitCalculatorNode->add_input_stream(strippedInputStream.cppString);
|
||||
|
||||
NSString *taskInputStream = [MPPTaskInfo addStreamNamePrefix:inputStream];
|
||||
taskSubgraphNode->add_input_stream(taskInputStream.cppString);
|
||||
|
||||
NSString *strippedInputStream = [MPPTaskInfo stripTagIndex:inputStream];
|
||||
flowLimitCalculatorNode->add_input_stream(strippedInputStream.cppString);
|
||||
|
||||
NSString *strippedTaskInputStream = [MPPTaskInfo stripTagIndex:taskInputStream];
|
||||
flowLimitCalculatorNode->add_output_stream(strippedTaskInputStream.cppString);
|
||||
}
|
||||
|
||||
NSString *firstOutputStream = self.outputStreams[0];
|
||||
auto finishedOutputStream = "FINISHED:" + firstOutputStream.cppString;
|
||||
NSString *strippedFirstOutputStream = [MPPTaskInfo stripTagIndex:self.outputStreams[0]];
|
||||
auto finishedOutputStream = "FINISHED:" + strippedFirstOutputStream.cppString;
|
||||
flowLimitCalculatorNode->add_input_stream(finishedOutputStream);
|
||||
|
||||
return graphConfig;
|
||||
|
||||
@@ -20,6 +20,8 @@ NS_ASSUME_NONNULL_BEGIN
|
||||
@interface MPPBaseOptions (Helpers)
|
||||
|
||||
- (void)copyToProto:(mediapipe::tasks::core::proto::BaseOptions *)baseOptionsProto;
|
||||
- (void)copyToProto:(mediapipe::tasks::core::proto::BaseOptions *)baseOptionsProto
|
||||
withUseStreamMode:(BOOL)useStreamMode;
|
||||
|
||||
@end
|
||||
|
||||
|
||||
@@ -22,6 +22,11 @@ using BaseOptionsProto = ::mediapipe::tasks::core::proto::BaseOptions;
|
||||
|
||||
@implementation MPPBaseOptions (Helpers)
|
||||
|
||||
- (void)copyToProto:(BaseOptionsProto *)baseOptionsProto withUseStreamMode:(BOOL)useStreamMode {
|
||||
[self copyToProto:baseOptionsProto];
|
||||
baseOptionsProto->set_use_stream_mode(useStreamMode);
|
||||
}
|
||||
|
||||
- (void)copyToProto:(BaseOptionsProto *)baseOptionsProto {
|
||||
baseOptionsProto->Clear();
|
||||
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
load("@build_bazel_rules_apple//apple:ios.bzl", "ios_unit_test")
|
||||
load(
|
||||
"//mediapipe/tasks:ios/ios.bzl",
|
||||
"MPP_TASK_MINIMUM_OS_VERSION",
|
||||
)
|
||||
load(
|
||||
"@org_tensorflow//tensorflow/lite:special_rules.bzl",
|
||||
"tflite_ios_lab_runner",
|
||||
)
|
||||
|
||||
package(default_visibility = ["//mediapipe/tasks:internal"])
|
||||
|
||||
licenses(["notice"])
|
||||
|
||||
# Default tags for filtering iOS targets. Targets are restricted to Apple platforms.
|
||||
TFL_DEFAULT_TAGS = [
|
||||
"apple",
|
||||
]
|
||||
|
||||
# Following sanitizer tests are not supported by iOS test targets.
|
||||
TFL_DISABLED_SANITIZER_TAGS = [
|
||||
"noasan",
|
||||
"nomsan",
|
||||
"notsan",
|
||||
]
|
||||
|
||||
objc_library(
|
||||
name = "MPPImageClassifierObjcTestLibrary",
|
||||
testonly = 1,
|
||||
srcs = ["MPPImageClassifierTests.m"],
|
||||
copts = [
|
||||
"-ObjC++",
|
||||
"-std=c++17",
|
||||
"-x objective-c++",
|
||||
],
|
||||
data = [
|
||||
"//mediapipe/tasks/testdata/vision:test_images",
|
||||
"//mediapipe/tasks/testdata/vision:test_models",
|
||||
],
|
||||
deps = [
|
||||
"//mediapipe/tasks/ios/common:MPPCommon",
|
||||
"//mediapipe/tasks/ios/test/vision/utils:MPPImageTestUtils",
|
||||
"//mediapipe/tasks/ios/vision/image_classifier:MPPImageClassifier",
|
||||
],
|
||||
)
|
||||
|
||||
ios_unit_test(
|
||||
name = "MPPImageClassifierObjcTest",
|
||||
minimum_os_version = MPP_TASK_MINIMUM_OS_VERSION,
|
||||
runner = tflite_ios_lab_runner("IOS_LATEST"),
|
||||
tags = TFL_DEFAULT_TAGS + TFL_DISABLED_SANITIZER_TAGS,
|
||||
deps = [
|
||||
":MPPImageClassifierObjcTestLibrary",
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,675 @@
|
||||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#import <XCTest/XCTest.h>
|
||||
|
||||
#import "mediapipe/tasks/ios/common/sources/MPPCommon.h"
|
||||
#import "mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.h"
|
||||
#import "mediapipe/tasks/ios/vision/image_classifier/sources/MPPImageClassifier.h"
|
||||
|
||||
static NSString *kFloatModelName = @"mobilenet_v2_1.0_224";
|
||||
static NSString *const kQuantizedModelName = @"mobilenet_v1_0.25_224_quant";
|
||||
static NSDictionary *const kBurgerImage = @{@"name" : @"burger", @"type" : @"jpg"};
|
||||
static NSDictionary *const kBurgerRotatedImage = @{@"name" : @"burger_rotated", @"type" : @"jpg"};
|
||||
static NSDictionary *const kMultiObjectsImage = @{@"name" : @"multi_objects", @"type" : @"jpg"};
|
||||
static NSDictionary *const kMultiObjectsRotatedImage =
|
||||
@{@"name" : @"multi_objects_rotated", @"type" : @"jpg"};
|
||||
static const int kMobileNetCategoriesCount = 1001;
|
||||
static NSString *const kExpectedErrorDomain = @"com.google.mediapipe.tasks";
|
||||
|
||||
#define AssertEqualErrors(error, expectedError) \
|
||||
XCTAssertNotNil(error); \
|
||||
XCTAssertEqualObjects(error.domain, expectedError.domain); \
|
||||
XCTAssertEqual(error.code, expectedError.code); \
|
||||
XCTAssertNotEqual( \
|
||||
[error.localizedDescription rangeOfString:expectedError.localizedDescription].location, \
|
||||
NSNotFound)
|
||||
|
||||
#define AssertEqualCategoryArrays(categories, expectedCategories) \
|
||||
XCTAssertEqual(categories.count, expectedCategories.count); \
|
||||
for (int i = 0; i < categories.count; i++) { \
|
||||
XCTAssertEqual(categories[i].index, expectedCategories[i].index, @"index i = %d", i); \
|
||||
XCTAssertEqualWithAccuracy(categories[i].score, expectedCategories[i].score, 1e-3, \
|
||||
@"index i = %d", i); \
|
||||
XCTAssertEqualObjects(categories[i].categoryName, expectedCategories[i].categoryName, \
|
||||
@"index i = %d", i); \
|
||||
XCTAssertEqualObjects(categories[i].displayName, expectedCategories[i].displayName, \
|
||||
@"index i = %d", i); \
|
||||
}
|
||||
|
||||
#define AssertImageClassifierResultHasOneHead(imageClassifierResult) \
|
||||
XCTAssertNotNil(imageClassifierResult); \
|
||||
XCTAssertNotNil(imageClassifierResult.classificationResult); \
|
||||
XCTAssertEqual(imageClassifierResult.classificationResult.classifications.count, 1); \
|
||||
XCTAssertEqual(imageClassifierResult.classificationResult.classifications[0].headIndex, 0);
|
||||
|
||||
@interface MPPImageClassifierTests : XCTestCase
|
||||
@end
|
||||
|
||||
@implementation MPPImageClassifierTests
|
||||
|
||||
#pragma mark Results
|
||||
|
||||
+ (NSArray<MPPCategory *> *)expectedResultCategoriesForClassifyBurgerImageWithFloatModel {
|
||||
return @[
|
||||
[[MPPCategory alloc] initWithIndex:934
|
||||
score:0.786005f
|
||||
categoryName:@"cheeseburger"
|
||||
displayName:nil],
|
||||
[[MPPCategory alloc] initWithIndex:932 score:0.023508f categoryName:@"bagel" displayName:nil],
|
||||
[[MPPCategory alloc] initWithIndex:925
|
||||
score:0.021172f
|
||||
categoryName:@"guacamole"
|
||||
displayName:nil]
|
||||
];
|
||||
}
|
||||
|
||||
#pragma mark File
|
||||
|
||||
- (NSString *)filePathWithName:(NSString *)fileName extension:(NSString *)extension {
|
||||
NSString *filePath = [[NSBundle bundleForClass:self.class] pathForResource:fileName
|
||||
ofType:extension];
|
||||
return filePath;
|
||||
}
|
||||
|
||||
#pragma mark Classifier Initializers
|
||||
|
||||
- (MPPImageClassifierOptions *)imageClassifierOptionsWithModelName:(NSString *)modelName {
|
||||
NSString *modelPath = [self filePathWithName:modelName extension:@"tflite"];
|
||||
MPPImageClassifierOptions *imageClassifierOptions = [[MPPImageClassifierOptions alloc] init];
|
||||
imageClassifierOptions.baseOptions.modelAssetPath = modelPath;
|
||||
|
||||
return imageClassifierOptions;
|
||||
}
|
||||
|
||||
- (MPPImageClassifier *)imageClassifierFromModelFileWithName:(NSString *)modelName {
|
||||
NSString *modelPath = [self filePathWithName:modelName extension:@"tflite"];
|
||||
MPPImageClassifier *imageClassifier = [[MPPImageClassifier alloc] initWithModelPath:modelPath
|
||||
error:nil];
|
||||
XCTAssertNotNil(imageClassifier);
|
||||
|
||||
return imageClassifier;
|
||||
}
|
||||
|
||||
- (MPPImageClassifier *)imageClassifierWithOptionsSucceeds:
|
||||
(MPPImageClassifierOptions *)imageClassifierOptions {
|
||||
MPPImageClassifier *imageClassifier =
|
||||
[[MPPImageClassifier alloc] initWithOptions:imageClassifierOptions error:nil];
|
||||
XCTAssertNotNil(imageClassifier);
|
||||
|
||||
return imageClassifier;
|
||||
}
|
||||
|
||||
#pragma mark Assert Classify Results
|
||||
|
||||
- (MPPImage *)imageWithFileInfo:(NSDictionary *)fileInfo {
|
||||
MPPImage *image = [MPPImage imageFromBundleWithClass:[MPPImageClassifierTests class]
|
||||
fileName:fileInfo[@"name"]
|
||||
ofType:fileInfo[@"type"]];
|
||||
XCTAssertNotNil(image);
|
||||
|
||||
return image;
|
||||
}
|
||||
|
||||
- (MPPImage *)imageWithFileInfo:(NSDictionary *)fileInfo
|
||||
orientation:(UIImageOrientation)orientation {
|
||||
MPPImage *image = [MPPImage imageFromBundleWithClass:[MPPImageClassifierTests class]
|
||||
fileName:fileInfo[@"name"]
|
||||
ofType:fileInfo[@"type"]
|
||||
orientation:orientation];
|
||||
XCTAssertNotNil(image);
|
||||
|
||||
return image;
|
||||
}
|
||||
|
||||
- (void)assertCreateImageClassifierWithOptions:(MPPImageClassifierOptions *)imageClassifierOptions
|
||||
failsWithExpectedError:(NSError *)expectedError {
|
||||
NSError *error = nil;
|
||||
MPPImageClassifier *imageClassifier =
|
||||
[[MPPImageClassifier alloc] initWithOptions:imageClassifierOptions error:&error];
|
||||
|
||||
XCTAssertNil(imageClassifier);
|
||||
AssertEqualErrors(error, expectedError);
|
||||
}
|
||||
|
||||
- (void)assertImageClassifierResult:(MPPImageClassifierResult *)imageClassifierResult
|
||||
hasExpectedCategoriesCount:(NSInteger)expectedCategoriesCount
|
||||
expectedCategories:(NSArray<MPPCategory *> *)expectedCategories {
|
||||
AssertImageClassifierResultHasOneHead(imageClassifierResult);
|
||||
|
||||
NSArray<MPPCategory *> *resultCategories =
|
||||
imageClassifierResult.classificationResult.classifications[0].categories;
|
||||
XCTAssertEqual(resultCategories.count, expectedCategoriesCount);
|
||||
|
||||
NSArray<MPPCategory *> *categorySubsetToCompare;
|
||||
if (resultCategories.count > expectedCategories.count) {
|
||||
categorySubsetToCompare =
|
||||
[resultCategories subarrayWithRange:NSMakeRange(0, expectedCategories.count)];
|
||||
} else {
|
||||
categorySubsetToCompare = resultCategories;
|
||||
}
|
||||
AssertEqualCategoryArrays(categorySubsetToCompare, expectedCategories);
|
||||
}
|
||||
|
||||
- (void)assertResultsOfClassifyImage:(MPPImage *)mppImage
|
||||
usingImageClassifier:(MPPImageClassifier *)imageClassifier
|
||||
expectedCategoriesCount:(NSInteger)expectedCategoriesCount
|
||||
equalsCategories:(NSArray<MPPCategory *> *)expectedCategories {
|
||||
MPPImageClassifierResult *imageClassifierResult = [imageClassifier classifyImage:mppImage
|
||||
error:nil];
|
||||
|
||||
[self assertImageClassifierResult:imageClassifierResult
|
||||
hasExpectedCategoriesCount:expectedCategoriesCount
|
||||
expectedCategories:expectedCategories];
|
||||
}
|
||||
|
||||
- (void)assertResultsOfClassifyImageWithFileInfo:(NSDictionary *)fileInfo
|
||||
usingImageClassifier:(MPPImageClassifier *)imageClassifier
|
||||
expectedCategoriesCount:(NSInteger)expectedCategoriesCount
|
||||
equalsCategories:(NSArray<MPPCategory *> *)expectedCategories {
|
||||
MPPImage *mppImage = [self imageWithFileInfo:fileInfo];
|
||||
|
||||
[self assertResultsOfClassifyImage:mppImage
|
||||
usingImageClassifier:imageClassifier
|
||||
expectedCategoriesCount:expectedCategoriesCount
|
||||
equalsCategories:expectedCategories];
|
||||
}
|
||||
|
||||
#pragma mark General Tests
|
||||
|
||||
- (void)testCreateImageClassifierWithMissingModelPathFails {
|
||||
NSString *modelPath = [self filePathWithName:@"" extension:@""];
|
||||
|
||||
NSError *error = nil;
|
||||
MPPImageClassifier *imageClassifier = [[MPPImageClassifier alloc] initWithModelPath:modelPath
|
||||
error:&error];
|
||||
XCTAssertNil(imageClassifier);
|
||||
|
||||
NSError *expectedError = [NSError
|
||||
errorWithDomain:kExpectedErrorDomain
|
||||
code:MPPTasksErrorCodeInvalidArgumentError
|
||||
userInfo:@{
|
||||
NSLocalizedDescriptionKey :
|
||||
@"INVALID_ARGUMENT: ExternalFile must specify at least one of 'file_content', "
|
||||
@"'file_name', 'file_pointer_meta' or 'file_descriptor_meta'."
|
||||
}];
|
||||
AssertEqualErrors(error, expectedError);
|
||||
}
|
||||
|
||||
- (void)testCreateImageClassifierAllowlistAndDenylistFails {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
options.categoryAllowlist = @[ @"cheeseburger" ];
|
||||
options.categoryDenylist = @[ @"bagel" ];
|
||||
|
||||
[self assertCreateImageClassifierWithOptions:options
|
||||
failsWithExpectedError:
|
||||
[NSError
|
||||
errorWithDomain:kExpectedErrorDomain
|
||||
code:MPPTasksErrorCodeInvalidArgumentError
|
||||
userInfo:@{
|
||||
NSLocalizedDescriptionKey :
|
||||
@"INVALID_ARGUMENT: `category_allowlist` and "
|
||||
@"`category_denylist` are mutually exclusive options."
|
||||
}]];
|
||||
}
|
||||
|
||||
- (void)testClassifyWithModelPathAndFloatModelSucceeds {
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierFromModelFileWithName:kFloatModelName];
|
||||
|
||||
[self
|
||||
assertResultsOfClassifyImageWithFileInfo:kBurgerImage
|
||||
usingImageClassifier:imageClassifier
|
||||
expectedCategoriesCount:kMobileNetCategoriesCount
|
||||
equalsCategories:
|
||||
[MPPImageClassifierTests
|
||||
expectedResultCategoriesForClassifyBurgerImageWithFloatModel]];
|
||||
}
|
||||
|
||||
- (void)testClassifyWithOptionsAndFloatModelSucceeds {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
const NSInteger maxResults = 3;
|
||||
options.maxResults = maxResults;
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
[self
|
||||
assertResultsOfClassifyImageWithFileInfo:kBurgerImage
|
||||
usingImageClassifier:imageClassifier
|
||||
expectedCategoriesCount:maxResults
|
||||
equalsCategories:
|
||||
[MPPImageClassifierTests
|
||||
expectedResultCategoriesForClassifyBurgerImageWithFloatModel]];
|
||||
}
|
||||
|
||||
- (void)testClassifyWithQuantizedModelSucceeds {
|
||||
MPPImageClassifierOptions *options =
|
||||
[self imageClassifierOptionsWithModelName:kQuantizedModelName];
|
||||
|
||||
const NSInteger maxResults = 1;
|
||||
options.maxResults = maxResults;
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
NSArray<MPPCategory *> *expectedCategories = @[ [[MPPCategory alloc] initWithIndex:934
|
||||
score:0.972656f
|
||||
categoryName:@"cheeseburger"
|
||||
displayName:nil] ];
|
||||
|
||||
[self assertResultsOfClassifyImageWithFileInfo:kBurgerImage
|
||||
usingImageClassifier:imageClassifier
|
||||
expectedCategoriesCount:maxResults
|
||||
equalsCategories:expectedCategories];
|
||||
}
|
||||
|
||||
- (void)testClassifyWithScoreThresholdSucceeds {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
options.scoreThreshold = 0.25f;
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
NSArray<MPPCategory *> *expectedCategories = @[ [[MPPCategory alloc] initWithIndex:934
|
||||
score:0.786005f
|
||||
categoryName:@"cheeseburger"
|
||||
displayName:nil] ];
|
||||
|
||||
[self assertResultsOfClassifyImageWithFileInfo:kBurgerImage
|
||||
usingImageClassifier:imageClassifier
|
||||
expectedCategoriesCount:expectedCategories.count
|
||||
equalsCategories:expectedCategories];
|
||||
}
|
||||
|
||||
- (void)testClassifyWithAllowlistSucceeds {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
options.categoryAllowlist = @[ @"cheeseburger", @"guacamole", @"meat loaf" ];
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
NSArray<MPPCategory *> *expectedCategories = @[
|
||||
[[MPPCategory alloc] initWithIndex:934
|
||||
score:0.786005f
|
||||
categoryName:@"cheeseburger"
|
||||
displayName:nil],
|
||||
[[MPPCategory alloc] initWithIndex:925
|
||||
score:0.021172f
|
||||
categoryName:@"guacamole"
|
||||
displayName:nil],
|
||||
[[MPPCategory alloc] initWithIndex:963
|
||||
score:0.006279315f
|
||||
categoryName:@"meat loaf"
|
||||
displayName:nil],
|
||||
|
||||
];
|
||||
|
||||
[self assertResultsOfClassifyImageWithFileInfo:kBurgerImage
|
||||
usingImageClassifier:imageClassifier
|
||||
expectedCategoriesCount:expectedCategories.count
|
||||
equalsCategories:expectedCategories];
|
||||
}
|
||||
|
||||
- (void)testClassifyWithDenylistSucceeds {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
options.categoryDenylist = @[
|
||||
@"bagel",
|
||||
];
|
||||
options.maxResults = 3;
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
NSArray<MPPCategory *> *expectedCategories = @[
|
||||
[[MPPCategory alloc] initWithIndex:934
|
||||
score:0.786005f
|
||||
categoryName:@"cheeseburger"
|
||||
displayName:nil],
|
||||
[[MPPCategory alloc] initWithIndex:925
|
||||
score:0.021172f
|
||||
categoryName:@"guacamole"
|
||||
displayName:nil],
|
||||
[[MPPCategory alloc] initWithIndex:963
|
||||
score:0.006279315f
|
||||
categoryName:@"meat loaf"
|
||||
displayName:nil],
|
||||
|
||||
];
|
||||
|
||||
[self assertResultsOfClassifyImageWithFileInfo:kBurgerImage
|
||||
usingImageClassifier:imageClassifier
|
||||
expectedCategoriesCount:expectedCategories.count
|
||||
equalsCategories:expectedCategories];
|
||||
}
|
||||
|
||||
- (void)testClassifyWithRegionOfInterestSucceeds {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
NSInteger maxResults = 1;
|
||||
options.maxResults = maxResults;
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
NSArray<MPPCategory *> *expectedCategories = @[ [[MPPCategory alloc] initWithIndex:806
|
||||
score:0.997122f
|
||||
categoryName:@"soccer ball"
|
||||
displayName:nil] ];
|
||||
|
||||
MPPImage *image = [self imageWithFileInfo:kMultiObjectsImage];
|
||||
|
||||
// roi around soccer ball
|
||||
MPPImageClassifierResult *imageClassifierResult =
|
||||
[imageClassifier classifyImage:image
|
||||
regionOfInterest:CGRectMake(0.450f, 0.308f, 0.164f, 0.426f)
|
||||
error:nil];
|
||||
[self assertImageClassifierResult:imageClassifierResult
|
||||
hasExpectedCategoriesCount:maxResults
|
||||
expectedCategories:expectedCategories];
|
||||
}
|
||||
|
||||
- (void)testClassifyWithOrientationSucceeds {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
NSInteger maxResults = 3;
|
||||
options.maxResults = maxResults;
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
NSArray<MPPCategory *> *expectedCategories = @[
|
||||
[[MPPCategory alloc] initWithIndex:934
|
||||
score:0.622074f
|
||||
categoryName:@"cheeseburger"
|
||||
displayName:nil],
|
||||
[[MPPCategory alloc] initWithIndex:963
|
||||
score:0.051214f
|
||||
categoryName:@"meat loaf"
|
||||
displayName:nil],
|
||||
[[MPPCategory alloc] initWithIndex:925
|
||||
score:0.048719f
|
||||
categoryName:@"guacamole"
|
||||
displayName:nil]
|
||||
|
||||
];
|
||||
|
||||
MPPImage *image = [self imageWithFileInfo:kBurgerRotatedImage
|
||||
orientation:UIImageOrientationRight];
|
||||
|
||||
[self assertResultsOfClassifyImage:image
|
||||
usingImageClassifier:imageClassifier
|
||||
expectedCategoriesCount:maxResults
|
||||
equalsCategories:expectedCategories];
|
||||
}
|
||||
|
||||
- (void)testClassifyWithRegionOfInterestAndOrientationSucceeds {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
NSInteger maxResults = 1;
|
||||
options.maxResults = maxResults;
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
NSArray<MPPCategory *> *expectedCategories =
|
||||
@[ [[MPPCategory alloc] initWithIndex:560
|
||||
score:0.682305f
|
||||
categoryName:@"folding chair"
|
||||
displayName:nil] ];
|
||||
|
||||
MPPImage *image = [self imageWithFileInfo:kMultiObjectsRotatedImage
|
||||
orientation:UIImageOrientationRight];
|
||||
|
||||
// roi around folding chair
|
||||
MPPImageClassifierResult *imageClassifierResult =
|
||||
[imageClassifier classifyImage:image
|
||||
regionOfInterest:CGRectMake(0.0f, 0.1763f, 0.5642f, 0.1286f)
|
||||
error:nil];
|
||||
[self assertImageClassifierResult:imageClassifierResult
|
||||
hasExpectedCategoriesCount:maxResults
|
||||
expectedCategories:expectedCategories];
|
||||
}
|
||||
|
||||
#pragma mark Running Mode Tests
|
||||
|
||||
- (void)testCreateImageClassifierFailsWithResultListenerInNonLiveStreamMode {
|
||||
MPPRunningMode runningModesToTest[] = {MPPRunningModeImage, MPPRunningModeVideo};
|
||||
for (int i = 0; i < sizeof(runningModesToTest) / sizeof(runningModesToTest[0]); i++) {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
options.runningMode = runningModesToTest[i];
|
||||
options.completion = ^(MPPImageClassifierResult *result, NSError *error) {
|
||||
};
|
||||
|
||||
[self
|
||||
assertCreateImageClassifierWithOptions:options
|
||||
failsWithExpectedError:
|
||||
[NSError
|
||||
errorWithDomain:kExpectedErrorDomain
|
||||
code:MPPTasksErrorCodeInvalidArgumentError
|
||||
userInfo:@{
|
||||
NSLocalizedDescriptionKey :
|
||||
@"The vision task is in image or video mode, a "
|
||||
@"user-defined result callback should not be provided."
|
||||
}]];
|
||||
}
|
||||
}
|
||||
|
||||
- (void)testCreateImageClassifierFailsWithMissingResultListenerInLiveStreamMode {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
options.runningMode = MPPRunningModeLiveStream;
|
||||
|
||||
[self assertCreateImageClassifierWithOptions:options
|
||||
failsWithExpectedError:
|
||||
[NSError errorWithDomain:kExpectedErrorDomain
|
||||
code:MPPTasksErrorCodeInvalidArgumentError
|
||||
userInfo:@{
|
||||
NSLocalizedDescriptionKey :
|
||||
@"The vision task is in live stream mode, a "
|
||||
@"user-defined result callback must be provided."
|
||||
}]];
|
||||
}
|
||||
|
||||
- (void)testClassifyFailsWithCallingWrongApiInImageMode {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
MPPImage *image = [self imageWithFileInfo:kBurgerImage];
|
||||
|
||||
NSError *liveStreamApiCallError;
|
||||
XCTAssertFalse([imageClassifier classifyAsyncImage:image
|
||||
timestampMs:0
|
||||
error:&liveStreamApiCallError]);
|
||||
|
||||
NSError *expectedLiveStreamApiCallError =
|
||||
[NSError errorWithDomain:kExpectedErrorDomain
|
||||
code:MPPTasksErrorCodeInvalidArgumentError
|
||||
userInfo:@{
|
||||
NSLocalizedDescriptionKey : @"The vision task is not initialized with live "
|
||||
@"stream mode. Current Running Mode: Image"
|
||||
}];
|
||||
|
||||
AssertEqualErrors(liveStreamApiCallError, expectedLiveStreamApiCallError);
|
||||
|
||||
NSError *videoApiCallError;
|
||||
XCTAssertFalse([imageClassifier classifyVideoFrame:image timestampMs:0 error:&videoApiCallError]);
|
||||
|
||||
NSError *expectedVideoApiCallError =
|
||||
[NSError errorWithDomain:kExpectedErrorDomain
|
||||
code:MPPTasksErrorCodeInvalidArgumentError
|
||||
userInfo:@{
|
||||
NSLocalizedDescriptionKey : @"The vision task is not initialized with "
|
||||
@"video mode. Current Running Mode: Image"
|
||||
}];
|
||||
AssertEqualErrors(videoApiCallError, expectedVideoApiCallError);
|
||||
}
|
||||
|
||||
- (void)testClassifyFailsWithCallingWrongApiInVideoMode {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
options.runningMode = MPPRunningModeVideo;
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
MPPImage *image = [self imageWithFileInfo:kBurgerImage];
|
||||
|
||||
NSError *liveStreamApiCallError;
|
||||
XCTAssertFalse([imageClassifier classifyAsyncImage:image
|
||||
timestampMs:0
|
||||
error:&liveStreamApiCallError]);
|
||||
|
||||
NSError *expectedLiveStreamApiCallError =
|
||||
[NSError errorWithDomain:kExpectedErrorDomain
|
||||
code:MPPTasksErrorCodeInvalidArgumentError
|
||||
userInfo:@{
|
||||
NSLocalizedDescriptionKey : @"The vision task is not initialized with live "
|
||||
@"stream mode. Current Running Mode: Video"
|
||||
}];
|
||||
|
||||
AssertEqualErrors(liveStreamApiCallError, expectedLiveStreamApiCallError);
|
||||
|
||||
NSError *imageApiCallError;
|
||||
XCTAssertFalse([imageClassifier classifyImage:image error:&imageApiCallError]);
|
||||
|
||||
NSError *expectedImageApiCallError =
|
||||
[NSError errorWithDomain:kExpectedErrorDomain
|
||||
code:MPPTasksErrorCodeInvalidArgumentError
|
||||
userInfo:@{
|
||||
NSLocalizedDescriptionKey : @"The vision task is not initialized with "
|
||||
@"image mode. Current Running Mode: Video"
|
||||
}];
|
||||
AssertEqualErrors(imageApiCallError, expectedImageApiCallError);
|
||||
}
|
||||
|
||||
- (void)testClassifyFailsWithCallingWrongApiInLiveStreamMode {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
options.runningMode = MPPRunningModeLiveStream;
|
||||
options.completion = ^(MPPImageClassifierResult *result, NSError *error) {
|
||||
|
||||
};
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
MPPImage *image = [self imageWithFileInfo:kBurgerImage];
|
||||
|
||||
NSError *imageApiCallError;
|
||||
XCTAssertFalse([imageClassifier classifyImage:image error:&imageApiCallError]);
|
||||
|
||||
NSError *expectedImageApiCallError =
|
||||
[NSError errorWithDomain:kExpectedErrorDomain
|
||||
code:MPPTasksErrorCodeInvalidArgumentError
|
||||
userInfo:@{
|
||||
NSLocalizedDescriptionKey : @"The vision task is not initialized with "
|
||||
@"image mode. Current Running Mode: Live Stream"
|
||||
}];
|
||||
AssertEqualErrors(imageApiCallError, expectedImageApiCallError);
|
||||
|
||||
NSError *videoApiCallError;
|
||||
XCTAssertFalse([imageClassifier classifyVideoFrame:image timestampMs:0 error:&videoApiCallError]);
|
||||
|
||||
NSError *expectedVideoApiCallError =
|
||||
[NSError errorWithDomain:kExpectedErrorDomain
|
||||
code:MPPTasksErrorCodeInvalidArgumentError
|
||||
userInfo:@{
|
||||
NSLocalizedDescriptionKey : @"The vision task is not initialized with "
|
||||
@"video mode. Current Running Mode: Live Stream"
|
||||
}];
|
||||
AssertEqualErrors(videoApiCallError, expectedVideoApiCallError);
|
||||
}
|
||||
|
||||
- (void)testClassifyWithVideoModeSucceeds {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
options.runningMode = MPPRunningModeVideo;
|
||||
|
||||
NSInteger maxResults = 3;
|
||||
options.maxResults = maxResults;
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
MPPImage *image = [self imageWithFileInfo:kBurgerImage];
|
||||
|
||||
for (int i = 0; i < 3; i++) {
|
||||
MPPImageClassifierResult *imageClassifierResult = [imageClassifier classifyVideoFrame:image
|
||||
timestampMs:i
|
||||
error:nil];
|
||||
[self assertImageClassifierResult:imageClassifierResult
|
||||
hasExpectedCategoriesCount:maxResults
|
||||
expectedCategories:
|
||||
[MPPImageClassifierTests
|
||||
expectedResultCategoriesForClassifyBurgerImageWithFloatModel]];
|
||||
}
|
||||
}
|
||||
|
||||
- (void)testClassifyWithOutOfOrderTimestampsAndLiveStreamModeFails {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
NSInteger maxResults = 3;
|
||||
options.maxResults = maxResults;
|
||||
|
||||
options.runningMode = MPPRunningModeLiveStream;
|
||||
options.completion = ^(MPPImageClassifierResult *result, NSError *error) {
|
||||
[self assertImageClassifierResult:result
|
||||
hasExpectedCategoriesCount:maxResults
|
||||
expectedCategories:
|
||||
[MPPImageClassifierTests
|
||||
expectedResultCategoriesForClassifyBurgerImageWithFloatModel]];
|
||||
};
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
MPPImage *image = [self imageWithFileInfo:kBurgerImage];
|
||||
|
||||
XCTAssertTrue([imageClassifier classifyAsyncImage:image timestampMs:1 error:nil]);
|
||||
|
||||
NSError *error;
|
||||
XCTAssertFalse([imageClassifier classifyAsyncImage:image timestampMs:0 error:&error]);
|
||||
|
||||
NSError *expectedError =
|
||||
[NSError errorWithDomain:kExpectedErrorDomain
|
||||
code:MPPTasksErrorCodeInvalidArgumentError
|
||||
userInfo:@{
|
||||
NSLocalizedDescriptionKey :
|
||||
@"INVALID_ARGUMENT: Input timestamp must be monotonically increasing."
|
||||
}];
|
||||
AssertEqualErrors(error, expectedError);
|
||||
}
|
||||
|
||||
- (void)testClassifyWithLiveStreamModeSucceeds {
|
||||
MPPImageClassifierOptions *options = [self imageClassifierOptionsWithModelName:kFloatModelName];
|
||||
|
||||
NSInteger maxResults = 3;
|
||||
options.maxResults = maxResults;
|
||||
|
||||
options.runningMode = MPPRunningModeLiveStream;
|
||||
options.completion = ^(MPPImageClassifierResult *result, NSError *error) {
|
||||
[self assertImageClassifierResult:result
|
||||
hasExpectedCategoriesCount:maxResults
|
||||
expectedCategories:
|
||||
[MPPImageClassifierTests
|
||||
expectedResultCategoriesForClassifyBurgerImageWithFloatModel]];
|
||||
};
|
||||
|
||||
MPPImageClassifier *imageClassifier = [self imageClassifierWithOptionsSucceeds:options];
|
||||
|
||||
// TODO: Mimic initialization from CMSampleBuffer as live stream mode is most likely to be used
|
||||
// with the iOS camera. AVCaptureVideoDataOutput sample buffer delegates provide frames of type
|
||||
// `CMSampleBuffer`.
|
||||
MPPImage *image = [self imageWithFileInfo:kBurgerImage];
|
||||
|
||||
for (int i = 0; i < 3; i++) {
|
||||
XCTAssertTrue([imageClassifier classifyAsyncImage:image timestampMs:i error:nil]);
|
||||
}
|
||||
}
|
||||
|
||||
@end
|
||||
@@ -0,0 +1,11 @@
|
||||
package(default_visibility = ["//mediapipe/tasks:internal"])
|
||||
|
||||
licenses(["notice"])
|
||||
|
||||
objc_library(
|
||||
name = "MPPImageTestUtils",
|
||||
srcs = ["sources/MPPImage+TestUtils.m"],
|
||||
hdrs = ["sources/MPPImage+TestUtils.h"],
|
||||
module_name = "MPPImageTestUtils",
|
||||
deps = ["//mediapipe/tasks/ios/vision/core:MPPImage"],
|
||||
)
|
||||
@@ -0,0 +1,63 @@
|
||||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#import <Foundation/Foundation.h>
|
||||
|
||||
#import "mediapipe/tasks/ios/vision/core/sources/MPPImage.h"
|
||||
|
||||
NS_ASSUME_NONNULL_BEGIN
|
||||
|
||||
/**
|
||||
* Helper utility for initializing `MPPImage` for MediaPipe iOS vision library tests.
|
||||
*/
|
||||
@interface MPPImage (TestUtils)
|
||||
|
||||
/**
|
||||
* Loads an image from a file in an app bundle into a `MPPImage` object.
|
||||
*
|
||||
* @param classObject The specified class associated with the bundle containing the file to be
|
||||
* loaded.
|
||||
* @param name Name of the image file.
|
||||
* @param type Extenstion of the image file.
|
||||
*
|
||||
* @return The `MPPImage` object contains the loaded image. This method returns
|
||||
* nil if it cannot load the image.
|
||||
*/
|
||||
+ (nullable MPPImage *)imageFromBundleWithClass:(Class)classObject
|
||||
fileName:(NSString *)name
|
||||
ofType:(NSString *)type
|
||||
NS_SWIFT_NAME(imageFromBundle(class:filename:type:));
|
||||
|
||||
/**
|
||||
* Loads an image from a file in an app bundle into a `MPPImage` object with the specified
|
||||
* orientation.
|
||||
*
|
||||
* @param classObject The specified class associated with the bundle containing the file to be
|
||||
* loaded.
|
||||
* @param name Name of the image file.
|
||||
* @param type Extenstion of the image file.
|
||||
* @param orientation Orientation of the image.
|
||||
*
|
||||
* @return The `MPPImage` object contains the loaded image. This method returns
|
||||
* nil if it cannot load the image.
|
||||
*/
|
||||
+ (nullable MPPImage *)imageFromBundleWithClass:(Class)classObject
|
||||
fileName:(NSString *)name
|
||||
ofType:(NSString *)type
|
||||
orientation:(UIImageOrientation)imageOrientation
|
||||
NS_SWIFT_NAME(imageFromBundle(class:filename:type:orientation:));
|
||||
|
||||
@end
|
||||
|
||||
NS_ASSUME_NONNULL_END
|
||||
@@ -0,0 +1,57 @@
|
||||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#import "mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.h"
|
||||
|
||||
@interface UIImage (FileUtils)
|
||||
|
||||
+ (nullable UIImage *)imageFromBundleWithClass:(Class)classObject
|
||||
fileName:(NSString *)name
|
||||
ofType:(NSString *)type;
|
||||
|
||||
@end
|
||||
|
||||
@implementation UIImage (FileUtils)
|
||||
|
||||
+ (nullable UIImage *)imageFromBundleWithClass:(Class)classObject
|
||||
fileName:(NSString *)name
|
||||
ofType:(NSString *)type {
|
||||
NSString *imagePath = [[NSBundle bundleForClass:classObject] pathForResource:name ofType:type];
|
||||
if (!imagePath) return nil;
|
||||
|
||||
return [[UIImage alloc] initWithContentsOfFile:imagePath];
|
||||
}
|
||||
|
||||
@end
|
||||
|
||||
@implementation MPPImage (TestUtils)
|
||||
|
||||
+ (nullable MPPImage *)imageFromBundleWithClass:(Class)classObject
|
||||
fileName:(NSString *)name
|
||||
ofType:(NSString *)type {
|
||||
UIImage *image = [UIImage imageFromBundleWithClass:classObject fileName:name ofType:type];
|
||||
|
||||
return [[MPPImage alloc] initWithUIImage:image error:nil];
|
||||
}
|
||||
|
||||
+ (nullable MPPImage *)imageFromBundleWithClass:(Class)classObject
|
||||
fileName:(NSString *)name
|
||||
ofType:(NSString *)type
|
||||
orientation:(UIImageOrientation)imageOrientation {
|
||||
UIImage *image = [UIImage imageFromBundleWithClass:classObject fileName:name ofType:type];
|
||||
|
||||
return [[MPPImage alloc] initWithUIImage:image orientation:imageOrientation error:nil];
|
||||
}
|
||||
|
||||
@end
|
||||
@@ -54,11 +54,13 @@ objc_library(
|
||||
],
|
||||
deps = [
|
||||
":MPPRunningMode",
|
||||
"//mediapipe/calculators/core:flow_limiter_calculator",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/tasks/ios/common:MPPCommon",
|
||||
"//mediapipe/tasks/ios/common/utils:MPPCommonUtils",
|
||||
"//mediapipe/tasks/ios/core:MPPTaskRunner",
|
||||
"//third_party/apple_frameworks:UIKit",
|
||||
"@com_google_absl//absl/status:statusor",
|
||||
"@ios_opencv//:OpencvFramework",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -44,9 +44,9 @@ NS_INLINE NSString *MPPRunningModeDisplayName(MPPRunningMode runningMode) {
|
||||
}
|
||||
|
||||
NSString *displayNameMap[MPPRunningModeLiveStream + 1] = {
|
||||
[MPPRunningModeImage] = @"#MPPRunningModeImage",
|
||||
[MPPRunningModeVideo] = @ "#MPPRunningModeVideo",
|
||||
[MPPRunningModeLiveStream] = @ "#MPPRunningModeLiveStream"};
|
||||
[MPPRunningModeImage] = @"Image",
|
||||
[MPPRunningModeVideo] = @"Video",
|
||||
[MPPRunningModeLiveStream] = @"Live Stream"};
|
||||
|
||||
return displayNameMap[runningMode];
|
||||
}
|
||||
|
||||
@@ -91,7 +91,7 @@ NS_ASSUME_NONNULL_BEGIN
|
||||
* A synchronous method to invoke the C++ task runner to process single image inputs. The call
|
||||
* blocks the current thread until a failure status or a successful result is returned.
|
||||
*
|
||||
* @param packetMap A `PackeMap` containing pairs of input stream name and data packet.
|
||||
* @param packetMap A `PacketMap` containing pairs of input stream name and data packet.
|
||||
* @param error Pointer to the memory location where errors if any should be
|
||||
* saved. If @c NULL, no error will be saved.
|
||||
*
|
||||
@@ -105,7 +105,7 @@ NS_ASSUME_NONNULL_BEGIN
|
||||
* A synchronous method to invoke the C++ task runner to process continuous video frames. The call
|
||||
* blocks the current thread until a failure status or a successful result is returned.
|
||||
*
|
||||
* @param packetMap A `PackeMap` containing pairs of input stream name and data packet.
|
||||
* @param packetMap A `PacketMap` containing pairs of input stream name and data packet.
|
||||
* @param error Pointer to the memory location where errors if any should be saved. If @c NULL, no
|
||||
* error will be saved.
|
||||
*
|
||||
@@ -121,7 +121,7 @@ NS_ASSUME_NONNULL_BEGIN
|
||||
* available in the user-defined `packetsCallback` that was provided during initialization of the
|
||||
* `MPPVisionTaskRunner`.
|
||||
*
|
||||
* @param packetMap A `PackeMap` containing pairs of input stream name and data packet.
|
||||
* @param packetMap A `PacketMap` containing pairs of input stream name and data packet.
|
||||
* @param error Pointer to the memory location where errors if any should be saved. If @c NULL, no
|
||||
* error will be saved.
|
||||
*
|
||||
|
||||
@@ -28,13 +28,13 @@ using ::mediapipe::tasks::core::PacketMap;
|
||||
using ::mediapipe::tasks::core::PacketsCallback;
|
||||
} // namespace
|
||||
|
||||
/** Rotation degress for a 90 degree rotation to the right. */
|
||||
/** Rotation degrees for a 90 degree rotation to the right. */
|
||||
static const NSInteger kMPPOrientationDegreesRight = -90;
|
||||
|
||||
/** Rotation degress for a 180 degree rotation. */
|
||||
/** Rotation degrees for a 180 degree rotation. */
|
||||
static const NSInteger kMPPOrientationDegreesDown = -180;
|
||||
|
||||
/** Rotation degress for a 90 degree rotation to the left. */
|
||||
/** Rotation degrees for a 90 degree rotation to the left. */
|
||||
static const NSInteger kMPPOrientationDegreesLeft = -270;
|
||||
|
||||
@interface MPPVisionTaskRunner () {
|
||||
@@ -97,7 +97,7 @@ static const NSInteger kMPPOrientationDegreesLeft = -270;
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
CGRect calculatedRoi = CGRectEqualToRect(roi, CGRectZero) ? roi : CGRectMake(0.0, 0.0, 1.0, 1.0);
|
||||
CGRect calculatedRoi = CGRectEqualToRect(roi, CGRectZero) ? CGRectMake(0.0, 0.0, 1.0, 1.0) : roi;
|
||||
|
||||
NormalizedRect normalizedRect;
|
||||
normalizedRect.set_x_center(CGRectGetMidX(calculatedRoi));
|
||||
|
||||
@@ -131,7 +131,9 @@ using ::mediapipe::ImageFrame;
|
||||
|
||||
size_t width = CVPixelBufferGetWidth(pixelBuffer);
|
||||
size_t height = CVPixelBufferGetHeight(pixelBuffer);
|
||||
size_t stride = CVPixelBufferGetBytesPerRow(pixelBuffer);
|
||||
|
||||
size_t destinationChannelCount = 3;
|
||||
size_t destinationStride = destinationChannelCount * width;
|
||||
|
||||
uint8_t *rgbPixelData = [MPPPixelDataUtils
|
||||
rgbPixelDataFromPixelData:(uint8_t *)CVPixelBufferGetBaseAddress(pixelBuffer)
|
||||
@@ -147,9 +149,10 @@ using ::mediapipe::ImageFrame;
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
std::unique_ptr<ImageFrame> imageFrame = absl::make_unique<ImageFrame>(
|
||||
::mediapipe::ImageFormat::SRGB, width, height, stride, static_cast<uint8 *>(rgbPixelData),
|
||||
/*deleter=*/free);
|
||||
std::unique_ptr<ImageFrame> imageFrame =
|
||||
absl::make_unique<ImageFrame>(::mediapipe::ImageFormat::SRGB, width, height,
|
||||
destinationStride, static_cast<uint8 *>(rgbPixelData),
|
||||
/*deleter=*/free);
|
||||
|
||||
return imageFrame;
|
||||
}
|
||||
@@ -183,11 +186,14 @@ using ::mediapipe::ImageFrame;
|
||||
|
||||
NSInteger bitsPerComponent = 8;
|
||||
NSInteger channelCount = 4;
|
||||
size_t bytesPerRow = channelCount * width;
|
||||
|
||||
NSInteger destinationChannelCount = 3;
|
||||
size_t destinationBytesPerRow = destinationChannelCount * width;
|
||||
|
||||
UInt8 *pixelDataToReturn = NULL;
|
||||
|
||||
CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceRGB();
|
||||
size_t bytesPerRow = channelCount * width;
|
||||
|
||||
// iOS infers bytesPerRow if it is set to 0.
|
||||
// See https://developer.apple.com/documentation/coregraphics/1455939-cgbitmapcontextcreate
|
||||
// But for segmentation test image, this was not the case.
|
||||
@@ -219,10 +225,14 @@ using ::mediapipe::ImageFrame;
|
||||
|
||||
CGColorSpaceRelease(colorSpace);
|
||||
|
||||
std::unique_ptr<ImageFrame> imageFrame =
|
||||
absl::make_unique<ImageFrame>(mediapipe::ImageFormat::SRGB, (int)width, (int)height,
|
||||
(int)bytesPerRow, static_cast<uint8 *>(pixelDataToReturn),
|
||||
/*deleter=*/free);
|
||||
if (!pixelDataToReturn) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
std::unique_ptr<ImageFrame> imageFrame = absl::make_unique<ImageFrame>(
|
||||
mediapipe::ImageFormat::SRGB, (int)width, (int)height, (int)destinationBytesPerRow,
|
||||
static_cast<uint8 *>(pixelDataToReturn),
|
||||
/*deleter=*/free);
|
||||
|
||||
return imageFrame;
|
||||
}
|
||||
|
||||
@@ -36,15 +36,17 @@ static NSString *const kClassificationsTag = @"CLASSIFICATIONS";
|
||||
static NSString *const kImageInStreamName = @"image_in";
|
||||
static NSString *const kImageOutStreamName = @"image_out";
|
||||
static NSString *const kImageTag = @"IMAGE";
|
||||
static NSString *const kNormRectName = @"norm_rect_in";
|
||||
static NSString *const kNormRectStreamName = @"norm_rect_in";
|
||||
static NSString *const kNormRectTag = @"NORM_RECT";
|
||||
|
||||
static NSString *const kTaskGraphName =
|
||||
@"mediapipe.tasks.vision.image_classifier.ImageClassifierGraph";
|
||||
|
||||
#define InputPacketMap(imagePacket, normalizedRectPacket) \
|
||||
{ \
|
||||
{kImageInStreamName.cppString, imagePacket}, { kNormRectName.cppString, normalizedRectPacket } \
|
||||
#define InputPacketMap(imagePacket, normalizedRectPacket) \
|
||||
{ \
|
||||
{kImageInStreamName.cppString, imagePacket}, { \
|
||||
kNormRectStreamName.cppString, normalizedRectPacket \
|
||||
} \
|
||||
}
|
||||
|
||||
@interface MPPImageClassifier () {
|
||||
@@ -60,12 +62,17 @@ static NSString *const kTaskGraphName =
|
||||
if (self) {
|
||||
MPPTaskInfo *taskInfo = [[MPPTaskInfo alloc]
|
||||
initWithTaskGraphName:kTaskGraphName
|
||||
inputStreams:@[ [NSString
|
||||
stringWithFormat:@"%@:%@", kImageTag, kImageInStreamName] ]
|
||||
outputStreams:@[ [NSString stringWithFormat:@"%@:%@", kClassificationsTag,
|
||||
kClassificationsStreamName] ]
|
||||
inputStreams:@[
|
||||
[NSString stringWithFormat:@"%@:%@", kImageTag, kImageInStreamName],
|
||||
[NSString stringWithFormat:@"%@:%@", kNormRectTag, kNormRectStreamName]
|
||||
]
|
||||
outputStreams:@[
|
||||
[NSString
|
||||
stringWithFormat:@"%@:%@", kClassificationsTag, kClassificationsStreamName],
|
||||
[NSString stringWithFormat:@"%@:%@", kImageTag, kImageOutStreamName]
|
||||
]
|
||||
taskOptions:options
|
||||
enableFlowLimiting:NO
|
||||
enableFlowLimiting:options.runningMode == MPPRunningModeLiveStream
|
||||
error:error];
|
||||
|
||||
if (!taskInfo) {
|
||||
@@ -130,8 +137,8 @@ static NSString *const kTaskGraphName =
|
||||
|
||||
PacketMap inputPacketMap = InputPacketMap(imagePacket, normalizedRectPacket);
|
||||
|
||||
std::optional<PacketMap> outputPacketMap = [_visionTaskRunner processPacketMap:inputPacketMap
|
||||
error:error];
|
||||
std::optional<PacketMap> outputPacketMap = [_visionTaskRunner processImagePacketMap:inputPacketMap
|
||||
error:error];
|
||||
if (!outputPacketMap.has_value()) {
|
||||
return nil;
|
||||
}
|
||||
|
||||
+4
-2
@@ -21,7 +21,7 @@
|
||||
#include "mediapipe/tasks/cc/vision/image_classifier/proto/image_classifier_graph_options.pb.h"
|
||||
|
||||
namespace {
|
||||
using CalculatorOptionsProto = ::mediapipe::CalculatorOptions;
|
||||
using CalculatorOptionsProto = mediapipe::CalculatorOptions;
|
||||
using ImageClassifierGraphOptionsProto =
|
||||
::mediapipe::tasks::vision::image_classifier::proto::ImageClassifierGraphOptions;
|
||||
using ClassifierOptionsProto = ::mediapipe::tasks::components::processors::proto::ClassifierOptions;
|
||||
@@ -32,7 +32,9 @@ using ClassifierOptionsProto = ::mediapipe::tasks::components::processors::proto
|
||||
- (void)copyToProto:(CalculatorOptionsProto *)optionsProto {
|
||||
ImageClassifierGraphOptionsProto *graphOptions =
|
||||
optionsProto->MutableExtension(ImageClassifierGraphOptionsProto::ext);
|
||||
[self.baseOptions copyToProto:graphOptions->mutable_base_options()];
|
||||
|
||||
[self.baseOptions copyToProto:graphOptions->mutable_base_options()
|
||||
withUseStreamMode:self.runningMode != MPPRunningModeImage];
|
||||
|
||||
ClassifierOptionsProto *classifierOptionsProto = graphOptions->mutable_classifier_options();
|
||||
classifierOptionsProto->Clear();
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
# Copyright 2023 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.
|
||||
|
||||
package(default_visibility = ["//mediapipe/tasks:internal"])
|
||||
|
||||
licenses(["notice"])
|
||||
|
||||
objc_library(
|
||||
name = "MPPObjectDetectionResult",
|
||||
srcs = ["sources/MPPObjectDetectionResult.m"],
|
||||
hdrs = ["sources/MPPObjectDetectionResult.h"],
|
||||
deps = [
|
||||
"//mediapipe/tasks/ios/components/containers:MPPDetection",
|
||||
"//mediapipe/tasks/ios/core:MPPTaskResult",
|
||||
],
|
||||
)
|
||||
|
||||
objc_library(
|
||||
name = "MPPObjectDetectorOptions",
|
||||
srcs = ["sources/MPPObjectDetectorOptions.m"],
|
||||
hdrs = ["sources/MPPObjectDetectorOptions.h"],
|
||||
deps = [
|
||||
":MPPObjectDetectionResult",
|
||||
"//mediapipe/tasks/ios/core:MPPTaskOptions",
|
||||
"//mediapipe/tasks/ios/vision/core:MPPRunningMode",
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,49 @@
|
||||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#import <Foundation/Foundation.h>
|
||||
#import "mediapipe/tasks/ios/components/containers/sources/MPPDetection.h"
|
||||
#import "mediapipe/tasks/ios/core/sources/MPPTaskResult.h"
|
||||
|
||||
NS_ASSUME_NONNULL_BEGIN
|
||||
|
||||
/** Represents the detection results generated by `MPPObjectDetector`. */
|
||||
NS_SWIFT_NAME(ObjectDetectionResult)
|
||||
@interface MPPObjectDetectionResult : MPPTaskResult
|
||||
|
||||
/**
|
||||
* The array of `MPPDetection` objects each of which has a bounding box that is expressed in the
|
||||
* unrotated input frame of reference coordinates system, i.e. in `[0,image_width) x
|
||||
* [0,image_height)`, which are the dimensions of the underlying image data.
|
||||
*/
|
||||
@property(nonatomic, readonly) NSArray<MPPDetection *> *detections;
|
||||
|
||||
/**
|
||||
* Initializes a new `MPPObjectDetectionResult` with the given array of detections and timestamp (in
|
||||
* milliseconds).
|
||||
*
|
||||
* @param detections An array of `MPPDetection` objects each of which has a bounding box that is
|
||||
* expressed in the unrotated input frame of reference coordinates system, i.e. in `[0,image_width)
|
||||
* x [0,image_height)`, which are the dimensions of the underlying image data.
|
||||
* @param timestampMs The timestamp for this result.
|
||||
*
|
||||
* @return An instance of `MPPObjectDetectionResult` initialized with the given array of detections
|
||||
* and timestamp (in milliseconds).
|
||||
*/
|
||||
- (instancetype)initWithDetections:(NSArray<MPPDetection *> *)detections
|
||||
timestampMs:(NSInteger)timestampMs;
|
||||
|
||||
@end
|
||||
|
||||
NS_ASSUME_NONNULL_END
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#import "mediapipe/tasks/ios/vision/object_detector/sources/MPPObjectDetectionResult.h"
|
||||
|
||||
@implementation MPPObjectDetectionResult
|
||||
|
||||
- (instancetype)initWithDetections:(NSArray<MPPDetection *> *)detections
|
||||
timestampMs:(NSInteger)timestampMs {
|
||||
self = [super initWithTimestampMs:timestampMs];
|
||||
if (self) {
|
||||
_detections = detections;
|
||||
}
|
||||
return self;
|
||||
}
|
||||
|
||||
@end
|
||||
@@ -0,0 +1,71 @@
|
||||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#import <Foundation/Foundation.h>
|
||||
|
||||
#import "mediapipe/tasks/ios/core/sources/MPPTaskOptions.h"
|
||||
#import "mediapipe/tasks/ios/vision/core/sources/MPPRunningMode.h"
|
||||
#import "mediapipe/tasks/ios/vision/object_detector/sources/MPPObjectDetectionResult.h"
|
||||
|
||||
NS_ASSUME_NONNULL_BEGIN
|
||||
|
||||
/** Options for setting up a `MPPObjectDetector`. */
|
||||
NS_SWIFT_NAME(ObjectDetectorOptions)
|
||||
@interface MPPObjectDetectorOptions : MPPTaskOptions <NSCopying>
|
||||
|
||||
@property(nonatomic) MPPRunningMode runningMode;
|
||||
|
||||
/**
|
||||
* The user-defined result callback for processing live stream data. The result callback should only
|
||||
* be specified when the running mode is set to the live stream mode.
|
||||
* TODO: Add parameter `MPPImage` in the callback.
|
||||
*/
|
||||
@property(nonatomic, copy) void (^completion)
|
||||
(MPPObjectDetectionResult *result, NSInteger timestampMs, NSError *error);
|
||||
|
||||
/**
|
||||
* The locale to use for display names specified through the TFLite Model Metadata, if any. Defaults
|
||||
* to English.
|
||||
*/
|
||||
@property(nonatomic, copy) NSString *displayNamesLocale;
|
||||
|
||||
/**
|
||||
* The maximum number of top-scored classification results to return. If < 0, all available results
|
||||
* will be returned. If 0, an invalid argument error is returned.
|
||||
*/
|
||||
@property(nonatomic) NSInteger maxResults;
|
||||
|
||||
/**
|
||||
* Score threshold to override the one provided in the model metadata (if any). Results below this
|
||||
* value are rejected.
|
||||
*/
|
||||
@property(nonatomic) float scoreThreshold;
|
||||
|
||||
/**
|
||||
* The 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.
|
||||
*/
|
||||
@property(nonatomic, copy) NSArray<NSString *> *categoryAllowlist;
|
||||
|
||||
/**
|
||||
* The 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.
|
||||
*/
|
||||
@property(nonatomic, copy) NSArray<NSString *> *categoryDenylist;
|
||||
|
||||
@end
|
||||
|
||||
NS_ASSUME_NONNULL_END
|
||||
@@ -0,0 +1,41 @@
|
||||
// Copyright 2023 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#import "mediapipe/tasks/ios/vision/object_detector/sources/MPPObjectDetectorOptions.h"
|
||||
|
||||
@implementation MPPObjectDetectorOptions
|
||||
|
||||
- (instancetype)init {
|
||||
self = [super init];
|
||||
if (self) {
|
||||
_maxResults = -1;
|
||||
_scoreThreshold = 0;
|
||||
}
|
||||
return self;
|
||||
}
|
||||
|
||||
- (id)copyWithZone:(NSZone *)zone {
|
||||
MPPObjectDetectorOptions *objectDetectorOptions = [super copyWithZone:zone];
|
||||
|
||||
objectDetectorOptions.scoreThreshold = self.scoreThreshold;
|
||||
objectDetectorOptions.maxResults = self.maxResults;
|
||||
objectDetectorOptions.categoryDenylist = self.categoryDenylist;
|
||||
objectDetectorOptions.categoryAllowlist = self.categoryAllowlist;
|
||||
objectDetectorOptions.displayNamesLocale = self.displayNamesLocale;
|
||||
objectDetectorOptions.completion = self.completion;
|
||||
|
||||
return objectDetectorOptions;
|
||||
}
|
||||
|
||||
@end
|
||||
@@ -54,10 +54,10 @@ load("//mediapipe/tasks/java/com/google/mediapipe/tasks:mediapipe_tasks_aar.bzl"
|
||||
mediapipe_tasks_core_aar(
|
||||
name = "tasks_core",
|
||||
srcs = glob(["**/*.java"]) + [
|
||||
"//mediapipe/java/com/google/mediapipe/framework/image:java_src",
|
||||
"//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:java_src",
|
||||
"//mediapipe/tasks/java/com/google/mediapipe/tasks/components/processors:java_src",
|
||||
"//mediapipe/tasks/java/com/google/mediapipe/tasks/components/utils:java_src",
|
||||
"//mediapipe/java/com/google/mediapipe/framework/image:java_src",
|
||||
],
|
||||
manifest = "AndroidManifest.xml",
|
||||
)
|
||||
|
||||
@@ -33,7 +33,7 @@ public class OutputHandler<OutputT extends TaskResult, InputT> {
|
||||
|
||||
/**
|
||||
* Interface for the customizable MediaPipe task result listener that can reteive both task result
|
||||
* objects and the correpsonding input data.
|
||||
* objects and the corresponding input data.
|
||||
*/
|
||||
public interface ResultListener<OutputT extends TaskResult, InputT> {
|
||||
void run(OutputT result, InputT input);
|
||||
@@ -90,8 +90,8 @@ public class OutputHandler<OutputT extends TaskResult, InputT> {
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets whether the output handler should react to the timestamp bound changes that are reprsented
|
||||
* as empty output {@link Packet}s.
|
||||
* Sets whether the output handler should react to the timestamp bound changes that are
|
||||
* represented as empty output {@link Packet}s.
|
||||
*
|
||||
* @param handleTimestampBoundChanges A boolean value.
|
||||
*/
|
||||
|
||||
@@ -24,7 +24,7 @@ import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* {@link TaskInfo} contains all needed informaton to initialize a MediaPipe Task {@link
|
||||
* {@link TaskInfo} contains all needed information to initialize a MediaPipe Task {@link
|
||||
* com.google.mediapipe.framework.Graph}.
|
||||
*/
|
||||
@AutoValue
|
||||
|
||||
@@ -22,8 +22,8 @@ import com.google.mediapipe.framework.AndroidPacketCreator;
|
||||
import com.google.mediapipe.framework.Graph;
|
||||
import com.google.mediapipe.framework.MediaPipeException;
|
||||
import com.google.mediapipe.framework.Packet;
|
||||
import com.google.mediapipe.tasks.core.logging.TasksStatsLogger;
|
||||
import com.google.mediapipe.tasks.core.logging.TasksStatsDummyLogger;
|
||||
import com.google.mediapipe.tasks.core.logging.TasksStatsLogger;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
|
||||
|
||||
@@ -35,6 +35,13 @@ _AUDIO_TASKS_JAVA_PROTO_LITE_TARGETS = [
|
||||
|
||||
_VISION_TASKS_JAVA_PROTO_LITE_TARGETS = [
|
||||
"//mediapipe/tasks/cc/vision/face_detector/proto:face_detector_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_geometry/proto:face_geometry_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_geometry/proto:face_geometry_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_geometry/proto:mesh_3d_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_landmarker/proto:face_blendshapes_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_landmarker/proto:face_landmarker_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_landmarker/proto:face_landmarks_detector_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_stylizer/proto:face_stylizer_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_classifier_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_embedder_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/gesture_recognizer/proto:gesture_recognizer_graph_options_java_proto_lite",
|
||||
@@ -110,6 +117,11 @@ def mediapipe_tasks_core_aar(name, srcs, manifest):
|
||||
src_out = "com/google/mediapipe/tasks/TensorsToSegmentationCalculatorOptionsProto.java",
|
||||
))
|
||||
|
||||
mediapipe_tasks_java_proto_srcs.append(mediapipe_java_proto_src_extractor(
|
||||
target = "//mediapipe/tasks/cc/vision/face_geometry/calculators:geometry_pipeline_calculator_java_proto_lite",
|
||||
src_out = "com/google/mediapipe/tasks/vision/facegeometry/calculators/proto/FaceGeometryPipelineCalculatorOptionsProto.java",
|
||||
))
|
||||
|
||||
android_library(
|
||||
name = name,
|
||||
srcs = srcs + [
|
||||
@@ -308,6 +320,7 @@ def _mediapipe_tasks_aar(name, srcs, manifest, java_proto_lite_targets, native_l
|
||||
"//mediapipe/framework/formats:detection_java_proto_lite",
|
||||
"//mediapipe/framework/formats:landmark_java_proto_lite",
|
||||
"//mediapipe/framework/formats:location_data_java_proto_lite",
|
||||
"//mediapipe/framework/formats:matrix_data_java_proto_lite",
|
||||
"//mediapipe/framework/formats:rect_java_proto_lite",
|
||||
"//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:audiodata",
|
||||
"//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:detection",
|
||||
|
||||
@@ -46,6 +46,8 @@ cc_binary(
|
||||
"//mediapipe/calculators/core:flow_limiter_calculator",
|
||||
"//mediapipe/java/com/google/mediapipe/framework/jni:mediapipe_framework_jni",
|
||||
"//mediapipe/tasks/cc/vision/face_detector:face_detector_graph",
|
||||
"//mediapipe/tasks/cc/vision/face_landmarker:face_landmarker_graph",
|
||||
"//mediapipe/tasks/cc/vision/face_stylizer:face_stylizer_graph",
|
||||
"//mediapipe/tasks/cc/vision/gesture_recognizer:gesture_recognizer_graph",
|
||||
"//mediapipe/tasks/cc/vision/image_classifier:image_classifier_graph",
|
||||
"//mediapipe/tasks/cc/vision/image_embedder:image_embedder_graph",
|
||||
@@ -114,6 +116,29 @@ android_library(
|
||||
],
|
||||
)
|
||||
|
||||
android_library(
|
||||
name = "facestylizer",
|
||||
srcs = [
|
||||
"facestylizer/FaceStylizer.java",
|
||||
"facestylizer/FaceStylizerResult.java",
|
||||
],
|
||||
javacopts = [
|
||||
"-Xep:AndroidJdkLibsChecker:OFF",
|
||||
],
|
||||
manifest = "imagesegmenter/AndroidManifest.xml",
|
||||
deps = [
|
||||
":core",
|
||||
"//mediapipe/framework:calculator_options_java_proto_lite",
|
||||
"//mediapipe/java/com/google/mediapipe/framework:android_framework",
|
||||
"//mediapipe/java/com/google/mediapipe/framework/image",
|
||||
"//mediapipe/tasks/cc/core/proto:base_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_stylizer/proto:face_stylizer_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/java/com/google/mediapipe/tasks/core",
|
||||
"//third_party:autovalue",
|
||||
"@maven//:com_google_guava_guava",
|
||||
],
|
||||
)
|
||||
|
||||
android_library(
|
||||
name = "gesturerecognizer",
|
||||
srcs = [
|
||||
@@ -289,6 +314,37 @@ android_library(
|
||||
],
|
||||
)
|
||||
|
||||
android_library(
|
||||
name = "facelandmarker",
|
||||
srcs = [
|
||||
"facelandmarker/FaceLandmarker.java",
|
||||
"facelandmarker/FaceLandmarkerResult.java",
|
||||
],
|
||||
javacopts = [
|
||||
"-Xep:AndroidJdkLibsChecker:OFF",
|
||||
],
|
||||
manifest = "facedetector/AndroidManifest.xml",
|
||||
deps = [
|
||||
":core",
|
||||
"//mediapipe/framework:calculator_options_java_proto_lite",
|
||||
"//mediapipe/framework/formats:classification_java_proto_lite",
|
||||
"//mediapipe/framework/formats:landmark_java_proto_lite",
|
||||
"//mediapipe/framework/formats:matrix_data_java_proto_lite",
|
||||
"//mediapipe/java/com/google/mediapipe/framework:android_framework",
|
||||
"//mediapipe/java/com/google/mediapipe/framework/image",
|
||||
"//mediapipe/tasks/cc/core/proto:base_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_detector/proto:face_detector_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_geometry/proto:face_geometry_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_landmarker/proto:face_landmarker_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/cc/vision/face_landmarker/proto:face_landmarks_detector_graph_options_java_proto_lite",
|
||||
"//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:category",
|
||||
"//mediapipe/tasks/java/com/google/mediapipe/tasks/components/containers:normalized_landmark",
|
||||
"//mediapipe/tasks/java/com/google/mediapipe/tasks/core",
|
||||
"//third_party:autovalue",
|
||||
"@maven//:com_google_guava_guava",
|
||||
],
|
||||
)
|
||||
|
||||
load("//mediapipe/tasks/java/com/google/mediapipe/tasks:mediapipe_tasks_aar.bzl", "mediapipe_tasks_vision_aar")
|
||||
|
||||
mediapipe_tasks_vision_aar(
|
||||
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
<?xml version="1.0" encoding="utf-8"?>
|
||||
<manifest xmlns:android="http://schemas.android.com/apk/res/android"
|
||||
package="com.google.mediapipe.tasks.vision.facelandmarker">
|
||||
|
||||
<uses-sdk android:minSdkVersion="24"
|
||||
android:targetSdkVersion="30" />
|
||||
|
||||
</manifest>
|
||||
+550
@@ -0,0 +1,550 @@
|
||||
// Copyright 2023 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.
|
||||
|
||||
package com.google.mediapipe.tasks.vision.facelandmarker;
|
||||
|
||||
import android.content.Context;
|
||||
import android.os.ParcelFileDescriptor;
|
||||
import com.google.auto.value.AutoValue;
|
||||
import com.google.mediapipe.formats.proto.LandmarkProto.NormalizedLandmarkList;
|
||||
import com.google.mediapipe.proto.CalculatorOptionsProto.CalculatorOptions;
|
||||
import com.google.mediapipe.formats.proto.ClassificationProto.ClassificationList;
|
||||
import com.google.mediapipe.framework.AndroidPacketGetter;
|
||||
import com.google.mediapipe.framework.Packet;
|
||||
import com.google.mediapipe.framework.PacketGetter;
|
||||
import com.google.mediapipe.framework.image.BitmapImageBuilder;
|
||||
import com.google.mediapipe.framework.image.MPImage;
|
||||
import com.google.mediapipe.tasks.core.BaseOptions;
|
||||
import com.google.mediapipe.tasks.core.ErrorListener;
|
||||
import com.google.mediapipe.tasks.core.OutputHandler;
|
||||
import com.google.mediapipe.tasks.core.OutputHandler.ResultListener;
|
||||
import com.google.mediapipe.tasks.core.TaskInfo;
|
||||
import com.google.mediapipe.tasks.core.TaskOptions;
|
||||
import com.google.mediapipe.tasks.core.TaskRunner;
|
||||
import com.google.mediapipe.tasks.core.proto.BaseOptionsProto;
|
||||
import com.google.mediapipe.tasks.vision.core.BaseVisionTaskApi;
|
||||
import com.google.mediapipe.tasks.vision.core.ImageProcessingOptions;
|
||||
import com.google.mediapipe.tasks.vision.core.RunningMode;
|
||||
import com.google.mediapipe.tasks.vision.facedetector.proto.FaceDetectorGraphOptionsProto;
|
||||
import com.google.mediapipe.tasks.vision.facegeometry.proto.FaceGeometryProto.FaceGeometry;
|
||||
import com.google.mediapipe.tasks.vision.facelandmarker.proto.FaceLandmarkerGraphOptionsProto;
|
||||
import com.google.mediapipe.tasks.vision.facelandmarker.proto.FaceLandmarksDetectorGraphOptionsProto;
|
||||
import com.google.mediapipe.formats.proto.MatrixDataProto.MatrixData;
|
||||
import java.io.File;
|
||||
import java.io.IOException;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
|
||||
/**
|
||||
* Performs face landmarks detection on images.
|
||||
*
|
||||
* <p>This API expects a pre-trained face landmarks model asset bundle. See <TODO link
|
||||
* to the DevSite documentation page>.
|
||||
*
|
||||
* <ul>
|
||||
* <li>Input image {@link MPImage}
|
||||
* <ul>
|
||||
* <li>The image that face landmarks detection runs on.
|
||||
* </ul>
|
||||
* <li>Output {@link FaceLandmarkerResult}
|
||||
* <ul>
|
||||
* <li>A FaceLandmarkerResult containing face landmarks.
|
||||
* </ul>
|
||||
* </ul>
|
||||
*/
|
||||
public final class FaceLandmarker extends BaseVisionTaskApi {
|
||||
private static final String TAG = FaceLandmarker.class.getSimpleName();
|
||||
private static final String IMAGE_IN_STREAM_NAME = "image_in";
|
||||
private static final String NORM_RECT_IN_STREAM_NAME = "norm_rect_in";
|
||||
|
||||
@SuppressWarnings("ConstantCaseForConstants")
|
||||
private static final List<String> INPUT_STREAMS =
|
||||
Collections.unmodifiableList(
|
||||
Arrays.asList("IMAGE:" + IMAGE_IN_STREAM_NAME, "NORM_RECT:" + NORM_RECT_IN_STREAM_NAME));
|
||||
|
||||
private static final int LANDMARKS_OUT_STREAM_INDEX = 0;
|
||||
private static final int IMAGE_OUT_STREAM_INDEX = 1;
|
||||
private static int blendshapesOutStreamIndex = -1;
|
||||
private static int faceGeometryOutStreamIndex = -1;
|
||||
private static final String TASK_GRAPH_NAME =
|
||||
"mediapipe.tasks.vision.face_landmarker.FaceLandmarkerGraph";
|
||||
|
||||
/**
|
||||
* Creates a {@link FaceLandmarker} instance from a model asset bundle path and the default {@link
|
||||
* FaceLandmarkerOptions}.
|
||||
*
|
||||
* @param context an Android {@link Context}.
|
||||
* @param modelAssetPath path to the face landmarks model with metadata in the assets.
|
||||
* @throws MediaPipeException if there is an error during {@link FaceLandmarker} creation.
|
||||
*/
|
||||
public static FaceLandmarker createFromFile(Context context, String modelAssetPath) {
|
||||
BaseOptions baseOptions = BaseOptions.builder().setModelAssetPath(modelAssetPath).build();
|
||||
return createFromOptions(
|
||||
context, FaceLandmarkerOptions.builder().setBaseOptions(baseOptions).build());
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a {@link FaceLandmarker} instance from a model asset bundle file and the default {@link
|
||||
* FaceLandmarkerOptions}.
|
||||
*
|
||||
* @param context an Android {@link Context}.
|
||||
* @param modelAssetFile the face landmarks model {@link File} instance.
|
||||
* @throws IOException if an I/O error occurs when opening the tflite model file.
|
||||
* @throws MediaPipeException if there is an error during {@link FaceLandmarker} creation.
|
||||
*/
|
||||
public static FaceLandmarker createFromFile(Context context, File modelAssetFile)
|
||||
throws IOException {
|
||||
try (ParcelFileDescriptor descriptor =
|
||||
ParcelFileDescriptor.open(modelAssetFile, ParcelFileDescriptor.MODE_READ_ONLY)) {
|
||||
BaseOptions baseOptions =
|
||||
BaseOptions.builder().setModelAssetFileDescriptor(descriptor.getFd()).build();
|
||||
return createFromOptions(
|
||||
context, FaceLandmarkerOptions.builder().setBaseOptions(baseOptions).build());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a {@link FaceLandmarker} instance from a model asset bundle buffer and the default
|
||||
* {@link FaceLandmarkerOptions}.
|
||||
*
|
||||
* @param context an Android {@link Context}.
|
||||
* @param modelBuffer a direct {@link ByteBuffer} or a {@link MappedByteBuffer} of the detection
|
||||
* model.
|
||||
* @throws MediaPipeException if there is an error during {@link FaceLandmarker} creation.
|
||||
*/
|
||||
public static FaceLandmarker createFromBuffer(
|
||||
Context context, final ByteBuffer modelAssetBuffer) {
|
||||
BaseOptions baseOptions = BaseOptions.builder().setModelAssetBuffer(modelAssetBuffer).build();
|
||||
return createFromOptions(
|
||||
context, FaceLandmarkerOptions.builder().setBaseOptions(baseOptions).build());
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a {@link FaceLandmarker} instance from a {@link FaceLandmarkerOptions}.
|
||||
*
|
||||
* @param context an Android {@link Context}.
|
||||
* @param landmarkerOptions a {@link FaceLandmarkerOptions} instance.
|
||||
* @throws MediaPipeException if there is an error during {@link FaceLandmarker} creation.
|
||||
*/
|
||||
public static FaceLandmarker createFromOptions(
|
||||
Context context, FaceLandmarkerOptions landmarkerOptions) {
|
||||
List<String> outputStreams = new ArrayList<>();
|
||||
outputStreams.add("NORM_LANDMARKS:face_landmarks");
|
||||
outputStreams.add("IMAGE:image_out");
|
||||
if (landmarkerOptions.outputFaceBlendshapes()) {
|
||||
outputStreams.add("BLENDSHAPES:face_blendshapes");
|
||||
blendshapesOutStreamIndex = outputStreams.size() - 1;
|
||||
}
|
||||
if (landmarkerOptions.outputFacialTransformationMatrixes()) {
|
||||
outputStreams.add("FACE_GEOMETRY:face_geometry");
|
||||
faceGeometryOutStreamIndex = outputStreams.size() - 1;
|
||||
}
|
||||
// TODO: Consolidate OutputHandler and TaskRunner.
|
||||
OutputHandler<FaceLandmarkerResult, MPImage> handler = new OutputHandler<>();
|
||||
handler.setOutputPacketConverter(
|
||||
new OutputHandler.OutputPacketConverter<FaceLandmarkerResult, MPImage>() {
|
||||
@Override
|
||||
public FaceLandmarkerResult convertToTaskResult(List<Packet> packets) {
|
||||
// If there is no faces detected in the image, just returns empty lists.
|
||||
if (packets.get(LANDMARKS_OUT_STREAM_INDEX).isEmpty()) {
|
||||
return FaceLandmarkerResult.create(
|
||||
new ArrayList<>(),
|
||||
Optional.empty(),
|
||||
Optional.empty(),
|
||||
BaseVisionTaskApi.generateResultTimestampMs(
|
||||
landmarkerOptions.runningMode(), packets.get(LANDMARKS_OUT_STREAM_INDEX)));
|
||||
}
|
||||
|
||||
Optional<List<ClassificationList>> blendshapes = Optional.empty();
|
||||
if (landmarkerOptions.outputFaceBlendshapes()) {
|
||||
blendshapes =
|
||||
Optional.of(
|
||||
PacketGetter.getProtoVector(
|
||||
packets.get(blendshapesOutStreamIndex), ClassificationList.parser()));
|
||||
}
|
||||
|
||||
Optional<List<MatrixData>> facialTransformationMatrixes = Optional.empty();
|
||||
if (landmarkerOptions.outputFacialTransformationMatrixes()) {
|
||||
List<FaceGeometry> faceGeometryList =
|
||||
PacketGetter.getProtoVector(
|
||||
packets.get(faceGeometryOutStreamIndex), FaceGeometry.parser());
|
||||
facialTransformationMatrixes = Optional.of(new ArrayList<>());
|
||||
for (FaceGeometry faceGeometry : faceGeometryList) {
|
||||
facialTransformationMatrixes.get().add(faceGeometry.getPoseTransformMatrix());
|
||||
}
|
||||
}
|
||||
|
||||
return FaceLandmarkerResult.create(
|
||||
PacketGetter.getProtoVector(
|
||||
packets.get(LANDMARKS_OUT_STREAM_INDEX), NormalizedLandmarkList.parser()),
|
||||
blendshapes,
|
||||
facialTransformationMatrixes,
|
||||
BaseVisionTaskApi.generateResultTimestampMs(
|
||||
landmarkerOptions.runningMode(), packets.get(LANDMARKS_OUT_STREAM_INDEX)));
|
||||
}
|
||||
|
||||
@Override
|
||||
public MPImage convertToTaskInput(List<Packet> packets) {
|
||||
return new BitmapImageBuilder(
|
||||
AndroidPacketGetter.getBitmapFromRgb(packets.get(IMAGE_OUT_STREAM_INDEX)))
|
||||
.build();
|
||||
}
|
||||
});
|
||||
landmarkerOptions.resultListener().ifPresent(handler::setResultListener);
|
||||
landmarkerOptions.errorListener().ifPresent(handler::setErrorListener);
|
||||
TaskRunner runner =
|
||||
TaskRunner.create(
|
||||
context,
|
||||
TaskInfo.<FaceLandmarkerOptions>builder()
|
||||
.setTaskName(FaceLandmarker.class.getSimpleName())
|
||||
.setTaskRunningModeName(landmarkerOptions.runningMode().name())
|
||||
.setTaskGraphName(TASK_GRAPH_NAME)
|
||||
.setInputStreams(INPUT_STREAMS)
|
||||
.setOutputStreams(outputStreams)
|
||||
.setTaskOptions(landmarkerOptions)
|
||||
.setEnableFlowLimiting(landmarkerOptions.runningMode() == RunningMode.LIVE_STREAM)
|
||||
.build(),
|
||||
handler);
|
||||
return new FaceLandmarker(runner, landmarkerOptions.runningMode());
|
||||
}
|
||||
|
||||
/**
|
||||
* Constructor to initialize an {@link FaceLandmarker} from a {@link TaskRunner} and a {@link
|
||||
* RunningMode}.
|
||||
*
|
||||
* @param taskRunner a {@link TaskRunner}.
|
||||
* @param runningMode a mediapipe vision task {@link RunningMode}.
|
||||
*/
|
||||
private FaceLandmarker(TaskRunner taskRunner, RunningMode runningMode) {
|
||||
super(taskRunner, runningMode, IMAGE_IN_STREAM_NAME, NORM_RECT_IN_STREAM_NAME);
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face landmarks detection on the provided single image with default image processing
|
||||
* options, i.e. without any rotation applied. Only use this method when the {@link
|
||||
* FaceLandmarker} is created with {@link RunningMode.IMAGE}. TODO update java doc
|
||||
* for input image format.
|
||||
*
|
||||
* <p>{@link FaceLandmarker} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link Bitmap.Config.ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @throws MediaPipeException if there is an internal error.
|
||||
*/
|
||||
public FaceLandmarkerResult detect(MPImage image) {
|
||||
return detect(image, ImageProcessingOptions.builder().build());
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face landmarks detection on the provided single image. Only use this method when the
|
||||
* {@link FaceLandmarker} is created with {@link RunningMode.IMAGE}. TODO update java
|
||||
* doc for input image format.
|
||||
*
|
||||
* <p>{@link FaceLandmarker} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link Bitmap.Config.ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the
|
||||
* input image before running inference. Note that region-of-interest is <b>not</b> supported
|
||||
* by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in
|
||||
* this method throwing an IllegalArgumentException.
|
||||
* @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a
|
||||
* region-of-interest.
|
||||
* @throws MediaPipeException if there is an internal error.
|
||||
*/
|
||||
public FaceLandmarkerResult detect(MPImage image, ImageProcessingOptions imageProcessingOptions) {
|
||||
validateImageProcessingOptions(imageProcessingOptions);
|
||||
return (FaceLandmarkerResult) processImageData(image, imageProcessingOptions);
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face landmarks detection on the provided video frame with default image processing
|
||||
* options, i.e. without any rotation applied. Only use this method when the {@link
|
||||
* FaceLandmarker} is created with {@link RunningMode.VIDEO}.
|
||||
*
|
||||
* <p>It's required to provide the video frame's timestamp (in milliseconds). The input timestamps
|
||||
* must be monotonically increasing.
|
||||
*
|
||||
* <p>{@link FaceLandmarker} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link Bitmap.Config.ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws MediaPipeException if there is an internal error.
|
||||
*/
|
||||
public FaceLandmarkerResult detectForVideo(MPImage image, long timestampMs) {
|
||||
return detectForVideo(image, ImageProcessingOptions.builder().build(), timestampMs);
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face landmarks detection on the provided video frame. Only use this method when the
|
||||
* {@link FaceLandmarker} is created with {@link RunningMode.VIDEO}.
|
||||
*
|
||||
* <p>It's required to provide the video frame's timestamp (in milliseconds). The input timestamps
|
||||
* must be monotonically increasing.
|
||||
*
|
||||
* <p>{@link FaceLandmarker} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link Bitmap.Config.ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the
|
||||
* input image before running inference. Note that region-of-interest is <b>not</b> supported
|
||||
* by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in
|
||||
* this method throwing an IllegalArgumentException.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a
|
||||
* region-of-interest.
|
||||
* @throws MediaPipeException if there is an internal error.
|
||||
*/
|
||||
public FaceLandmarkerResult detectForVideo(
|
||||
MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) {
|
||||
validateImageProcessingOptions(imageProcessingOptions);
|
||||
return (FaceLandmarkerResult) processVideoData(image, imageProcessingOptions, timestampMs);
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends live image data to perform face landmarks detection with default image processing
|
||||
* options, i.e. without any rotation applied, and the results will be available via the {@link
|
||||
* ResultListener} provided in the {@link FaceLandmarkerOptions}. Only use this method when the
|
||||
* {@link FaceLandmarker } is created with {@link RunningMode.LIVE_STREAM}.
|
||||
*
|
||||
* <p>It's required to provide a timestamp (in milliseconds) to indicate when the input image is
|
||||
* sent to the face landmarker. The input timestamps must be monotonically increasing.
|
||||
*
|
||||
* <p>{@link FaceLandmarker} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link Bitmap.Config.ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws MediaPipeException if there is an internal error.
|
||||
*/
|
||||
public void detectAsync(MPImage image, long timestampMs) {
|
||||
detectAsync(image, ImageProcessingOptions.builder().build(), timestampMs);
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends live image data to perform face landmarks detection, and the results will be available
|
||||
* via the {@link ResultListener} provided in the {@link FaceLandmarkerOptions}. Only use this
|
||||
* method when the {@link FaceLandmarker} is created with {@link RunningMode.LIVE_STREAM}.
|
||||
*
|
||||
* <p>It's required to provide a timestamp (in milliseconds) to indicate when the input image is
|
||||
* sent to the face landmarker. The input timestamps must be monotonically increasing.
|
||||
*
|
||||
* <p>{@link FaceLandmarker} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link Bitmap.Config.ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the
|
||||
* input image before running inference. Note that region-of-interest is <b>not</b> supported
|
||||
* by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in
|
||||
* this method throwing an IllegalArgumentException.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a
|
||||
* region-of-interest.
|
||||
* @throws MediaPipeException if there is an internal error.
|
||||
*/
|
||||
public void detectAsync(
|
||||
MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) {
|
||||
validateImageProcessingOptions(imageProcessingOptions);
|
||||
sendLiveStreamData(image, imageProcessingOptions, timestampMs);
|
||||
}
|
||||
|
||||
/** Options for setting up an {@link FaceLandmarker}. */
|
||||
@AutoValue
|
||||
public abstract static class FaceLandmarkerOptions extends TaskOptions {
|
||||
|
||||
/** Builder for {@link FaceLandmarkerOptions}. */
|
||||
@AutoValue.Builder
|
||||
public abstract static class Builder {
|
||||
/** Sets the base options for the face landmarker task. */
|
||||
public abstract Builder setBaseOptions(BaseOptions value);
|
||||
|
||||
/**
|
||||
* Sets the running mode for the face landmarker task. Default to the image mode. Hand
|
||||
* landmarker has three modes:
|
||||
*
|
||||
* <ul>
|
||||
* <li>IMAGE: The mode for detecting face landmarks on single image inputs.
|
||||
* <li>VIDEO: The mode for detecting face landmarks on the decoded frames of a video.
|
||||
* <li>LIVE_STREAM: The mode for for detecting face landmarks on a live stream of input
|
||||
* data, such as from camera. In this mode, {@code setResultListener} must be called to
|
||||
* set up a listener to receive the detection results asynchronously.
|
||||
* </ul>
|
||||
*/
|
||||
public abstract Builder setRunningMode(RunningMode value);
|
||||
|
||||
/** Sets the maximum number of faces can be detected by the FaceLandmarker. */
|
||||
public abstract Builder setNumFaces(Integer value);
|
||||
|
||||
/** Sets minimum confidence score for the face detection to be considered successful */
|
||||
public abstract Builder setMinFaceDetectionConfidence(Float value);
|
||||
|
||||
/** Sets minimum confidence score of face presence score in the face landmark detection. */
|
||||
public abstract Builder setMinFacePresenceConfidence(Float value);
|
||||
|
||||
/** Sets the minimum confidence score for the face tracking to be considered successful. */
|
||||
public abstract Builder setMinTrackingConfidence(Float value);
|
||||
|
||||
/**
|
||||
* Whether FaceLandmarker outputs face blendshapes classification. Face blendshapes are used
|
||||
* for rendering the 3D face model.
|
||||
*/
|
||||
public abstract Builder setOutputFaceBlendshapes(Boolean value);
|
||||
|
||||
/**
|
||||
* Whether FaceLandmarker outptus facial transformation_matrix. Facial transformation matrix
|
||||
* is used to transform the face landmarks in canonical face to the detected face, so that
|
||||
* users can apply face effects on the detected landmarks.
|
||||
*/
|
||||
public abstract Builder setOutputFacialTransformationMatrixes(Boolean value);
|
||||
|
||||
/**
|
||||
* Sets the result listener to receive the detection results asynchronously when the face
|
||||
* landmarker is in the live stream mode.
|
||||
*/
|
||||
public abstract Builder setResultListener(
|
||||
ResultListener<FaceLandmarkerResult, MPImage> value);
|
||||
|
||||
/** Sets an optional error listener. */
|
||||
public abstract Builder setErrorListener(ErrorListener value);
|
||||
|
||||
abstract FaceLandmarkerOptions autoBuild();
|
||||
|
||||
/**
|
||||
* Validates and builds the {@link FaceLandmarkerOptions} instance.
|
||||
*
|
||||
* @throws IllegalArgumentException if the result listener and the running mode are not
|
||||
* properly configured. The result listener should only be set when the face landmarker is
|
||||
* in the live stream mode.
|
||||
*/
|
||||
public final FaceLandmarkerOptions build() {
|
||||
FaceLandmarkerOptions options = autoBuild();
|
||||
if (options.runningMode() == RunningMode.LIVE_STREAM) {
|
||||
if (!options.resultListener().isPresent()) {
|
||||
throw new IllegalArgumentException(
|
||||
"The face landmarker is in the live stream mode, a user-defined result listener"
|
||||
+ " must be provided in FaceLandmarkerOptions.");
|
||||
}
|
||||
} else if (options.resultListener().isPresent()) {
|
||||
throw new IllegalArgumentException(
|
||||
"The face landmarker is in the image or the video mode, a user-defined result"
|
||||
+ " listener shouldn't be provided in FaceLandmarkerOptions.");
|
||||
}
|
||||
return options;
|
||||
}
|
||||
}
|
||||
|
||||
abstract BaseOptions baseOptions();
|
||||
|
||||
abstract RunningMode runningMode();
|
||||
|
||||
abstract Optional<Integer> numFaces();
|
||||
|
||||
abstract Optional<Float> minFaceDetectionConfidence();
|
||||
|
||||
abstract Optional<Float> minFacePresenceConfidence();
|
||||
|
||||
abstract Optional<Float> minTrackingConfidence();
|
||||
|
||||
abstract Boolean outputFaceBlendshapes();
|
||||
|
||||
abstract Boolean outputFacialTransformationMatrixes();
|
||||
|
||||
abstract Optional<ResultListener<FaceLandmarkerResult, MPImage>> resultListener();
|
||||
|
||||
abstract Optional<ErrorListener> errorListener();
|
||||
|
||||
public static Builder builder() {
|
||||
return new AutoValue_FaceLandmarker_FaceLandmarkerOptions.Builder()
|
||||
.setRunningMode(RunningMode.IMAGE)
|
||||
.setNumFaces(1)
|
||||
.setMinFaceDetectionConfidence(0.5f)
|
||||
.setMinFacePresenceConfidence(0.5f)
|
||||
.setMinTrackingConfidence(0.5f)
|
||||
.setOutputFaceBlendshapes(false)
|
||||
.setOutputFacialTransformationMatrixes(false);
|
||||
}
|
||||
|
||||
/** Converts a {@link FaceLandmarkerOptions} to a {@link CalculatorOptions} protobuf message. */
|
||||
@Override
|
||||
public CalculatorOptions convertToCalculatorOptionsProto() {
|
||||
FaceLandmarkerGraphOptionsProto.FaceLandmarkerGraphOptions.Builder taskOptionsBuilder =
|
||||
FaceLandmarkerGraphOptionsProto.FaceLandmarkerGraphOptions.newBuilder()
|
||||
.setBaseOptions(
|
||||
BaseOptionsProto.BaseOptions.newBuilder()
|
||||
.setUseStreamMode(runningMode() != RunningMode.IMAGE)
|
||||
.mergeFrom(convertBaseOptionsToProto(baseOptions()))
|
||||
.build());
|
||||
|
||||
// Setup FaceDetectorGraphOptions.
|
||||
FaceDetectorGraphOptionsProto.FaceDetectorGraphOptions.Builder
|
||||
faceDetectorGraphOptionsBuilder =
|
||||
FaceDetectorGraphOptionsProto.FaceDetectorGraphOptions.newBuilder();
|
||||
numFaces().ifPresent(faceDetectorGraphOptionsBuilder::setNumFaces);
|
||||
minFaceDetectionConfidence()
|
||||
.ifPresent(faceDetectorGraphOptionsBuilder::setMinDetectionConfidence);
|
||||
|
||||
// Setup FaceLandmarkerGraphOptions.
|
||||
FaceLandmarksDetectorGraphOptionsProto.FaceLandmarksDetectorGraphOptions.Builder
|
||||
faceLandmarksDetectorGraphOptionsBuilder =
|
||||
FaceLandmarksDetectorGraphOptionsProto.FaceLandmarksDetectorGraphOptions.newBuilder();
|
||||
minFacePresenceConfidence()
|
||||
.ifPresent(faceLandmarksDetectorGraphOptionsBuilder::setMinDetectionConfidence);
|
||||
minTrackingConfidence().ifPresent(taskOptionsBuilder::setMinTrackingConfidence);
|
||||
|
||||
taskOptionsBuilder
|
||||
.setFaceDetectorGraphOptions(faceDetectorGraphOptionsBuilder.build())
|
||||
.setFaceLandmarksDetectorGraphOptions(faceLandmarksDetectorGraphOptionsBuilder.build());
|
||||
|
||||
return CalculatorOptions.newBuilder()
|
||||
.setExtension(
|
||||
FaceLandmarkerGraphOptionsProto.FaceLandmarkerGraphOptions.ext,
|
||||
taskOptionsBuilder.build())
|
||||
.build();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates that the provided {@link ImageProcessingOptions} doesn't contain a
|
||||
* region-of-interest.
|
||||
*/
|
||||
private static void validateImageProcessingOptions(
|
||||
ImageProcessingOptions imageProcessingOptions) {
|
||||
if (imageProcessingOptions.regionOfInterest().isPresent()) {
|
||||
throw new IllegalArgumentException("FaceLandmarker doesn't support region-of-interest.");
|
||||
}
|
||||
}
|
||||
}
|
||||
+115
@@ -0,0 +1,115 @@
|
||||
// Copyright 2023 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.
|
||||
|
||||
package com.google.mediapipe.tasks.vision.facelandmarker;
|
||||
|
||||
import com.google.auto.value.AutoValue;
|
||||
import com.google.mediapipe.formats.proto.LandmarkProto;
|
||||
import com.google.mediapipe.formats.proto.ClassificationProto.Classification;
|
||||
import com.google.mediapipe.formats.proto.ClassificationProto.ClassificationList;
|
||||
import com.google.mediapipe.tasks.components.containers.Category;
|
||||
import com.google.mediapipe.tasks.components.containers.NormalizedLandmark;
|
||||
import com.google.mediapipe.tasks.core.TaskResult;
|
||||
import com.google.mediapipe.formats.proto.MatrixDataProto.MatrixData;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
|
||||
/** Represents the face landmarks detection results generated by {@link FaceLandmarker}. */
|
||||
@AutoValue
|
||||
public abstract class FaceLandmarkerResult implements TaskResult {
|
||||
|
||||
/**
|
||||
* Creates a {@link FaceLandmarkerResult} instance from the list of landmarks, list of face
|
||||
* blendshapes classification, and list of facial transformation matrixes protobuf message.
|
||||
*
|
||||
* @param multiFaceLandmarksProto a List of {@link NormalizedLandmarkList}
|
||||
* @param multiFaceBendshapesProto an Optional List of {@link ClassificationList}
|
||||
* @param multiFaceTransformationMatrixesProto an Optional List of {@link MatrixData}
|
||||
* @throws IllegalArgumentException if there is error creating {@link FaceLandmarkerResult}
|
||||
*/
|
||||
static FaceLandmarkerResult create(
|
||||
List<LandmarkProto.NormalizedLandmarkList> multiFaceLandmarksProto,
|
||||
Optional<List<ClassificationList>> multiFaceBendshapesProto,
|
||||
Optional<List<MatrixData>> multiFaceTransformationMatrixesProto,
|
||||
long timestampMs) {
|
||||
List<List<NormalizedLandmark>> multiFaceLandmarks = new ArrayList<>();
|
||||
for (LandmarkProto.NormalizedLandmarkList faceLandmarksProto : multiFaceLandmarksProto) {
|
||||
List<NormalizedLandmark> faceLandmarks = new ArrayList<>();
|
||||
multiFaceLandmarks.add(faceLandmarks);
|
||||
for (LandmarkProto.NormalizedLandmark faceLandmarkProto :
|
||||
faceLandmarksProto.getLandmarkList()) {
|
||||
faceLandmarks.add(
|
||||
NormalizedLandmark.create(
|
||||
faceLandmarkProto.getX(), faceLandmarkProto.getY(), faceLandmarkProto.getZ()));
|
||||
}
|
||||
}
|
||||
Optional<List<List<Category>>> multiFaceBlendshapes = Optional.empty();
|
||||
if (multiFaceBendshapesProto.isPresent()) {
|
||||
List<List<Category>> blendshapes = new ArrayList<>();
|
||||
for (ClassificationList faceBendshapeProto : multiFaceBendshapesProto.get()) {
|
||||
List<Category> blendshape = new ArrayList<>();
|
||||
blendshapes.add(blendshape);
|
||||
for (Classification classification : faceBendshapeProto.getClassificationList()) {
|
||||
blendshape.add(
|
||||
Category.create(
|
||||
classification.getScore(),
|
||||
classification.getIndex(),
|
||||
classification.getLabel(),
|
||||
classification.getDisplayName()));
|
||||
}
|
||||
}
|
||||
multiFaceBlendshapes = Optional.of(Collections.unmodifiableList(blendshapes));
|
||||
}
|
||||
Optional<List<float[]>> multiFaceTransformationMatrixes = Optional.empty();
|
||||
if (multiFaceTransformationMatrixesProto.isPresent()) {
|
||||
List<float[]> matrixes = new ArrayList<>();
|
||||
for (MatrixData matrixProto : multiFaceTransformationMatrixesProto.get()) {
|
||||
if (matrixProto.getPackedDataCount() != 16) {
|
||||
throw new IllegalArgumentException(
|
||||
"MatrixData must contain 4x4 matrix as a size 16 float array, but get size "
|
||||
+ matrixProto.getPackedDataCount()
|
||||
+ " float array.");
|
||||
}
|
||||
float[] matrixData = new float[matrixProto.getPackedDataCount()];
|
||||
for (int i = 0; i < matrixData.length; i++) {
|
||||
matrixData[i] = matrixProto.getPackedData(i);
|
||||
}
|
||||
matrixes.add(matrixData);
|
||||
}
|
||||
multiFaceTransformationMatrixes = Optional.of(Collections.unmodifiableList(matrixes));
|
||||
}
|
||||
return new AutoValue_FaceLandmarkerResult(
|
||||
timestampMs,
|
||||
Collections.unmodifiableList(multiFaceLandmarks),
|
||||
multiFaceBlendshapes,
|
||||
multiFaceTransformationMatrixes);
|
||||
}
|
||||
|
||||
@Override
|
||||
public abstract long timestampMs();
|
||||
|
||||
/** Face landmarks of detected faces. */
|
||||
public abstract List<List<NormalizedLandmark>> faceLandmarks();
|
||||
|
||||
/** Optional face blendshapes classifications. */
|
||||
public abstract Optional<List<List<Category>>> faceBlendshapes();
|
||||
|
||||
/**
|
||||
* Optional facial transformation matrix list from canonical face to the detected face landmarks.
|
||||
* The 4x4 facial transformation matrix is represetned as a flat column-major float array.
|
||||
*/
|
||||
public abstract Optional<List<float[]>> facialTransformationMatrixes();
|
||||
}
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
<?xml version="1.0" encoding="utf-8"?>
|
||||
<manifest xmlns:android="http://schemas.android.com/apk/res/android"
|
||||
package="com.google.mediapipe.tasks.vision.facestylizer">
|
||||
|
||||
<uses-sdk android:minSdkVersion="24"
|
||||
android:targetSdkVersion="30" />
|
||||
|
||||
</manifest>
|
||||
+572
@@ -0,0 +1,572 @@
|
||||
// Copyright 2023 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.
|
||||
|
||||
package com.google.mediapipe.tasks.vision.facestylizer;
|
||||
|
||||
import android.content.Context;
|
||||
import com.google.auto.value.AutoValue;
|
||||
import com.google.mediapipe.proto.CalculatorOptionsProto.CalculatorOptions;
|
||||
import com.google.mediapipe.framework.AndroidPacketGetter;
|
||||
import com.google.mediapipe.framework.MediaPipeException;
|
||||
import com.google.mediapipe.framework.Packet;
|
||||
import com.google.mediapipe.framework.PacketGetter;
|
||||
import com.google.mediapipe.framework.image.BitmapImageBuilder;
|
||||
import com.google.mediapipe.framework.image.ByteBufferImageBuilder;
|
||||
import com.google.mediapipe.framework.image.MPImage;
|
||||
import com.google.mediapipe.tasks.core.BaseOptions;
|
||||
import com.google.mediapipe.tasks.core.ErrorListener;
|
||||
import com.google.mediapipe.tasks.core.OutputHandler;
|
||||
import com.google.mediapipe.tasks.core.OutputHandler.ResultListener;
|
||||
import com.google.mediapipe.tasks.core.TaskInfo;
|
||||
import com.google.mediapipe.tasks.core.TaskOptions;
|
||||
import com.google.mediapipe.tasks.core.TaskResult;
|
||||
import com.google.mediapipe.tasks.core.TaskRunner;
|
||||
import com.google.mediapipe.tasks.core.proto.BaseOptionsProto;
|
||||
import com.google.mediapipe.tasks.vision.core.BaseVisionTaskApi;
|
||||
import com.google.mediapipe.tasks.vision.core.ImageProcessingOptions;
|
||||
import com.google.mediapipe.tasks.vision.core.RunningMode;
|
||||
import com.google.mediapipe.tasks.vision.facestylizer.proto.FaceStylizerGraphOptionsProto;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.util.Arrays;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
|
||||
/**
|
||||
* Performs face stylization on images.
|
||||
*
|
||||
* <p>Note that, in addition to the standard stylization API, {@link #stylize} and {@link
|
||||
* #stylizeForVideo}, that take an input image and return the outputs, but involves deep copy of the
|
||||
* returns, FaceStylizer also supports the callback API, {@link #stylizeWithResultListener} and
|
||||
* {@link #stylizeForVideoWithResultListener}, which allow you to access the outputs through zero
|
||||
* copy for the duration of the result listener.
|
||||
*
|
||||
* <p>The callback API is available for all {@link RunningMode} in FaceStylizer. Set {@link
|
||||
* ResultListener} in {@link FaceStylizerOptions} properly to use the callback API.
|
||||
*
|
||||
* <p>The API expects a TFLite model with,<a
|
||||
* href="https://www.tensorflow.org/lite/convert/metadata">TFLite Model Metadata.</a>.
|
||||
*
|
||||
* <ul>
|
||||
* <li>Input image {@link MPImage}
|
||||
* <ul>
|
||||
* <li>The image that face stylizer runs on.
|
||||
* </ul>
|
||||
* <li>Output MPImage {@link MPImage}
|
||||
* <ul>
|
||||
* <li>A MPImage containing a stylized face.
|
||||
* </ul>
|
||||
* </ul>
|
||||
*/
|
||||
public final class FaceStylizer extends BaseVisionTaskApi {
|
||||
private static final String IMAGE_IN_STREAM_NAME = "image_in";
|
||||
private static final String NORM_RECT_IN_STREAM_NAME = "norm_rect_in";
|
||||
private static final String IMAGE_OUT_STREAM_NAME = "image_out";
|
||||
|
||||
@SuppressWarnings("ConstantCaseForConstants")
|
||||
private static final List<String> INPUT_STREAMS =
|
||||
Collections.unmodifiableList(
|
||||
Arrays.asList("IMAGE:" + IMAGE_IN_STREAM_NAME, "NORM_RECT:" + NORM_RECT_IN_STREAM_NAME));
|
||||
|
||||
@SuppressWarnings("ConstantCaseForConstants")
|
||||
private static final List<String> OUTPUT_STREAMS =
|
||||
Collections.singletonList("STYLIZED_IMAGE:" + IMAGE_OUT_STREAM_NAME);
|
||||
|
||||
private static final int IMAGE_OUT_STREAM_INDEX = 0;
|
||||
private static final String TASK_GRAPH_NAME =
|
||||
"mediapipe.tasks.vision.face_stylizer.FaceStylizerGraph";
|
||||
private final boolean hasResultListener;
|
||||
|
||||
/**
|
||||
* Creates an {@link FaceStylizer} instance from an {@link FaceStylizerOptions}.
|
||||
*
|
||||
* @param context an Android {@link Context}.
|
||||
* @param stylizerOptions an {@link FaceStylizerOptions} instance.
|
||||
* @throws MediaPipeException if there is an error during {@link FaceStylizer} creation.
|
||||
*/
|
||||
public static FaceStylizer createFromOptions(
|
||||
Context context, FaceStylizerOptions stylizerOptions) {
|
||||
// TODO: Consolidate OutputHandler and TaskRunner.
|
||||
OutputHandler<FaceStylizerResult, MPImage> handler = new OutputHandler<>();
|
||||
handler.setOutputPacketConverter(
|
||||
new OutputHandler.OutputPacketConverter<FaceStylizerResult, MPImage>() {
|
||||
@Override
|
||||
public FaceStylizerResult convertToTaskResult(List<Packet> packets)
|
||||
throws MediaPipeException {
|
||||
Packet packet = packets.get(IMAGE_OUT_STREAM_INDEX);
|
||||
int width = PacketGetter.getImageWidth(packet);
|
||||
int height = PacketGetter.getImageHeight(packet);
|
||||
int numChannels = PacketGetter.getImageNumChannels(packet);
|
||||
int imageFormat =
|
||||
numChannels == 3 ? MPImage.IMAGE_FORMAT_RGB : MPImage.IMAGE_FORMAT_RGBA;
|
||||
|
||||
ByteBuffer imageBuffer;
|
||||
// If resultListener is not provided, the resulted MPImage is deep copied from the
|
||||
// MediaPipe graph. If provided, the result MPImage is wrapping the MediaPipe packet
|
||||
// memory.
|
||||
if (!stylizerOptions.resultListener().isPresent()) {
|
||||
imageBuffer = ByteBuffer.allocateDirect(width * height * numChannels);
|
||||
if (!PacketGetter.getImageData(packet, imageBuffer)) {
|
||||
imageBuffer = null;
|
||||
}
|
||||
} else {
|
||||
imageBuffer = PacketGetter.getImageDataDirectly(packet);
|
||||
}
|
||||
|
||||
if (imageBuffer == null) {
|
||||
throw new MediaPipeException(
|
||||
MediaPipeException.StatusCode.INTERNAL.ordinal(),
|
||||
"There is an error getting the stylized face. It usually results from incorrect"
|
||||
+ " options of unsupported OutputType of given model.");
|
||||
}
|
||||
ByteBufferImageBuilder imageBuilder =
|
||||
new ByteBufferImageBuilder(imageBuffer, width, height, imageFormat);
|
||||
|
||||
return FaceStylizerResult.create(
|
||||
imageBuilder.build(),
|
||||
BaseVisionTaskApi.generateResultTimestampMs(
|
||||
stylizerOptions.runningMode(), packets.get(IMAGE_OUT_STREAM_INDEX)));
|
||||
}
|
||||
|
||||
@Override
|
||||
public MPImage convertToTaskInput(List<Packet> packets) {
|
||||
return new BitmapImageBuilder(
|
||||
AndroidPacketGetter.getBitmapFromRgb(packets.get(IMAGE_OUT_STREAM_INDEX)))
|
||||
.build();
|
||||
}
|
||||
});
|
||||
stylizerOptions.resultListener().ifPresent(handler::setResultListener);
|
||||
stylizerOptions.errorListener().ifPresent(handler::setErrorListener);
|
||||
TaskRunner runner =
|
||||
TaskRunner.create(
|
||||
context,
|
||||
TaskInfo.<FaceStylizerOptions>builder()
|
||||
.setTaskName(FaceStylizer.class.getSimpleName())
|
||||
.setTaskRunningModeName(stylizerOptions.runningMode().name())
|
||||
.setTaskGraphName(TASK_GRAPH_NAME)
|
||||
.setInputStreams(INPUT_STREAMS)
|
||||
.setOutputStreams(OUTPUT_STREAMS)
|
||||
.setTaskOptions(stylizerOptions)
|
||||
.setEnableFlowLimiting(stylizerOptions.runningMode() == RunningMode.LIVE_STREAM)
|
||||
.build(),
|
||||
handler);
|
||||
return new FaceStylizer(
|
||||
runner, stylizerOptions.runningMode(), stylizerOptions.resultListener().isPresent());
|
||||
}
|
||||
|
||||
/**
|
||||
* Constructor to initialize an {@link FaceStylizer} from a {@link TaskRunner} and a {@link
|
||||
* RunningMode}.
|
||||
*
|
||||
* @param taskRunner a {@link TaskRunner}.
|
||||
* @param runningMode a mediapipe vision task {@link RunningMode}.
|
||||
*/
|
||||
private FaceStylizer(TaskRunner taskRunner, RunningMode runningMode, boolean hasResultListener) {
|
||||
super(taskRunner, runningMode, IMAGE_IN_STREAM_NAME, NORM_RECT_IN_STREAM_NAME);
|
||||
this.hasResultListener = hasResultListener;
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face stylization on the provided single image with default image processing options,
|
||||
* i.e. without any rotation applied. Only use this method when the {@link FaceStylizer} is
|
||||
* created with {@link RunningMode#IMAGE}.
|
||||
*
|
||||
* <p>{@link FaceStylizer} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link android.graphics.Bitmap.Config#ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* <p>The image can be of any size. To ensure that the output image has reasonable quality, the
|
||||
* size of the stylized output is based the model output size and can be smaller than the input
|
||||
* image.
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link FaceStylizer} is created
|
||||
* with a {@link ResultListener}.
|
||||
*/
|
||||
public FaceStylizerResult stylize(MPImage image) {
|
||||
return stylize(image, ImageProcessingOptions.builder().build());
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face stylization on the provided single image. Only use this method when the {@link
|
||||
* FaceStylizer} is created with {@link RunningMode#IMAGE}.
|
||||
*
|
||||
* <p>{@link FaceStylizer} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link android.graphics.Bitmap.Config#ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* <p>The input image can be of any size, To ensure that the output image has reasonable quality,
|
||||
* the stylized output image size is the smaller of the model output size and the size of the
|
||||
* {@link ImageProcessingOptions#regionOfInterest} specified in {@code imageProcessingOptions}.
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the
|
||||
* input image before running inference. Note that region-of-interest is <b>not</b> supported
|
||||
* by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in
|
||||
* this method throwing an IllegalArgumentException.
|
||||
* @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a
|
||||
* region-of-interest.
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link FaceStylizer} is created
|
||||
* with a {@link ResultListener}.
|
||||
*/
|
||||
public FaceStylizerResult stylize(MPImage image, ImageProcessingOptions imageProcessingOptions) {
|
||||
if (hasResultListener) {
|
||||
throw new MediaPipeException(
|
||||
MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(),
|
||||
"ResultListener is provided in the FaceStylizerOptions, but this method will return an"
|
||||
+ " ImageSegmentationResult.");
|
||||
}
|
||||
return (FaceStylizerResult) processImageData(image, imageProcessingOptions);
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face stylization on the provided single image with default image processing options,
|
||||
* i.e. without any rotation applied, and provides zero-copied results via {@link ResultListener}
|
||||
* in {@link FaceStylizerOptions}. Only use this method when the {@link FaceStylizer} is created
|
||||
* with {@link RunningMode#IMAGE}.
|
||||
*
|
||||
* <p>{@link FaceStylizer} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link android.graphics.Bitmap.Config#ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* <p>The image can be of any size. To ensure that the output image has reasonable quality, the
|
||||
* size of the stylized output is based the model output size and can be smaller than the input
|
||||
* image.
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a
|
||||
* region-of-interest.
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link FaceStylizer} is not
|
||||
* created wtih {@link ResultListener} set in {@link FaceStylizerOptions}.
|
||||
*/
|
||||
public void stylizeWithResultListener(MPImage image) {
|
||||
stylizeWithResultListener(image, ImageProcessingOptions.builder().build());
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face stylization on the provided single image, and provides zero-copied results via
|
||||
* {@link ResultListener} in {@link FaceStylizerOptions}. Only use this method when the {@link
|
||||
* FaceStylizer} is created with {@link RunningMode#IMAGE}.
|
||||
*
|
||||
* <p>{@link FaceStylizer} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link android.graphics.Bitmap.Config#ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* <p>The input image can be of any size, To ensure that the output image has reasonable quality,
|
||||
* the stylized output image size is the smaller of the model output size and the size of the
|
||||
* {@link ImageProcessingOptions#regionOfInterest} specified in {@code imageProcessingOptions}.
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the
|
||||
* input image before running inference. Note that region-of-interest is <b>not</b> supported
|
||||
* by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in
|
||||
* this method throwing an IllegalArgumentException.
|
||||
* @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a
|
||||
* region-of-interest.
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link FaceStylizer} is not
|
||||
* created wtih {@link ResultListener} set in {@link FaceStylizerOptions}.
|
||||
*/
|
||||
public void stylizeWithResultListener(
|
||||
MPImage image, ImageProcessingOptions imageProcessingOptions) {
|
||||
if (!hasResultListener) {
|
||||
throw new MediaPipeException(
|
||||
MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(),
|
||||
"ResultListener is not set in the FaceStylizerOptions, but this method expects a"
|
||||
+ " ResultListener to process ImageSegmentationResult.");
|
||||
}
|
||||
TaskResult unused = processImageData(image, imageProcessingOptions);
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face stylization on the provided video frame with default image processing options,
|
||||
* i.e. without any rotation applied. Only use this method when the {@link FaceStylizer} is
|
||||
* created with {@link RunningMode#VIDEO}.
|
||||
*
|
||||
* <p>It's required to provide the video frame's timestamp (in milliseconds). The input timestamps
|
||||
* must be monotonically increasing.
|
||||
*
|
||||
* <p>{@link FaceStylizer} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link android.graphics.Bitmap.Config#ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* <p>The image can be of any size. To ensure that the output image has reasonable quality, the
|
||||
* size of the stylized output is based the model output size and can be smaller than the input
|
||||
* image.
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link FaceStylizer} is created
|
||||
* with a {@link ResultListener}.
|
||||
*/
|
||||
public FaceStylizerResult stylizeForVideo(MPImage image, long timestampMs) {
|
||||
return stylizeForVideo(image, ImageProcessingOptions.builder().build(), timestampMs);
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face stylization on the provided video frame. Only use this method when the {@link
|
||||
* FaceStylizer} is created with {@link RunningMode#VIDEO}.
|
||||
*
|
||||
* <p>It's required to provide the video frame's timestamp (in milliseconds). The input timestamps
|
||||
* must be monotonically increasing.
|
||||
*
|
||||
* <p>{@link FaceStylizer} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link android.graphics.Bitmap.Config#ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* <p>The input image can be of any size, To ensure that the output image has reasonable quality,
|
||||
* the stylized output image size is the smaller of the model output size and the size of the
|
||||
* {@link ImageProcessingOptions#regionOfInterest} specified in {@code imageProcessingOptions}.
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the
|
||||
* input image before running inference. Note that region-of-interest is <b>not</b> supported
|
||||
* by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in
|
||||
* this method throwing an IllegalArgumentException.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a
|
||||
* region-of-interest.
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link FaceStylizer} is created
|
||||
* with a {@link ResultListener}.
|
||||
*/
|
||||
public FaceStylizerResult stylizeForVideo(
|
||||
MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) {
|
||||
if (hasResultListener) {
|
||||
throw new MediaPipeException(
|
||||
MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(),
|
||||
"ResultListener is provided in the FaceStylizerOptions, but this method will return an"
|
||||
+ " ImageSegmentationResult.");
|
||||
}
|
||||
return (FaceStylizerResult) processVideoData(image, imageProcessingOptions, timestampMs);
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face stylization on the provided video frame with default image processing options,
|
||||
* i.e. without any rotation applied, and provides zero-copied results via {@link ResultListener}
|
||||
* in {@link FaceStylizerOptions}. Only use this method when the {@link FaceStylizer} is created
|
||||
* with {@link RunningMode#VIDEO}.
|
||||
*
|
||||
* <p>It's required to provide the video frame's timestamp (in milliseconds). The input timestamps
|
||||
* must be monotonically increasing.
|
||||
*
|
||||
* <p>{@link FaceStylizer} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link android.graphics.Bitmap.Config#ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* <p>The image can be of any size. To ensure that the output image has reasonable quality, the
|
||||
* size of the stylized output is based the model output size and can be smaller than the input
|
||||
* image.
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link FaceStylizer} is not
|
||||
* created wtih {@link ResultListener} set in {@link FaceStylizerOptions}.
|
||||
*/
|
||||
public void stylizeForVideoWithResultListener(MPImage image, long timestampMs) {
|
||||
stylizeForVideoWithResultListener(image, ImageProcessingOptions.builder().build(), timestampMs);
|
||||
}
|
||||
|
||||
/**
|
||||
* Performs face stylization on the provided video frame, and provides zero-copied results via
|
||||
* {@link ResultListener} in {@link FaceStylizerOptions}. Only use this method when the {@link
|
||||
* FaceStylizer} is created with {@link RunningMode#VIDEO}.
|
||||
*
|
||||
* <p>It's required to provide the video frame's timestamp (in milliseconds). The input timestamps
|
||||
* must be monotonically increasing.
|
||||
*
|
||||
* <p>{@link FaceStylizer} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link android.graphics.Bitmap.Config#ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* <p>The input image can be of any size, To ensure that the output image has reasonable quality,
|
||||
* the stylized output image size is the smaller of the model output size and the size of the
|
||||
* {@link ImageProcessingOptions#regionOfInterest} specified in {@code imageProcessingOptions}.
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link FaceStylizer} is not
|
||||
* created wtih {@link ResultListener} set in {@link FaceStylizerOptions}.
|
||||
*/
|
||||
public void stylizeForVideoWithResultListener(
|
||||
MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) {
|
||||
if (!hasResultListener) {
|
||||
throw new MediaPipeException(
|
||||
MediaPipeException.StatusCode.FAILED_PRECONDITION.ordinal(),
|
||||
"ResultListener is not set in the FaceStylizerOptions, but this method expects a"
|
||||
+ " ResultListener to process ImageSegmentationResult.");
|
||||
}
|
||||
TaskResult unused = processVideoData(image, imageProcessingOptions, timestampMs);
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends live image data to perform face stylization with default image processing options, i.e.
|
||||
* without any rotation applied, and the results will be available via the {@link ResultListener}
|
||||
* provided in the {@link FaceStylizerOptions}. Only use this method when the {@link FaceStylizer
|
||||
* } is created with {@link RunningMode#LIVE_STREAM}.
|
||||
*
|
||||
* <p>It's required to provide a timestamp (in milliseconds) to indicate when the input image is
|
||||
* sent to the face stylizer. The input timestamps must be monotonically increasing.
|
||||
*
|
||||
* <p>{@link FaceStylizer} supports the following color space types:
|
||||
*
|
||||
* <p>The image can be of any size. To ensure that the output image has reasonable quality, the
|
||||
* size of the stylized output is based the model output * size and can be smaller than the input
|
||||
* image.
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link android.graphics.Bitmap.Config#ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws MediaPipeException if there is an internal error.
|
||||
*/
|
||||
public void stylizeAsync(MPImage image, long timestampMs) {
|
||||
stylizeAsync(image, ImageProcessingOptions.builder().build(), timestampMs);
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends live image data to perform face stylization, and the results will be available via the
|
||||
* {@link ResultListener} provided in the {@link FaceStylizerOptions}. Only use this method when
|
||||
* the {@link FaceStylizer} is created with {@link RunningMode#LIVE_STREAM}.
|
||||
*
|
||||
* <p>It's required to provide a timestamp (in milliseconds) to indicate when the input image is
|
||||
* sent to the face stylizer. The input timestamps must be monotonically increasing.
|
||||
*
|
||||
* <p>{@link FaceStylizer} supports the following color space types:
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@link android.graphics.Bitmap.Config#ARGB_8888}
|
||||
* </ul>
|
||||
*
|
||||
* <p>The input image can be of any size, To ensure that the output image has reasonable quality,
|
||||
* the stylized output image size is the smaller of the model output size and the size of the
|
||||
* {@link ImageProcessingOptions#regionOfInterest} specified in {@code imageProcessingOptions}.
|
||||
*
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param imageProcessingOptions the {@link ImageProcessingOptions} specifying how to process the
|
||||
* input image before running inference. Note that region-of-interest is <b>not</b> supported
|
||||
* by this task: specifying {@link ImageProcessingOptions#regionOfInterest()} will result in
|
||||
* this method throwing an IllegalArgumentException.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a
|
||||
* region-of-interest.
|
||||
* @throws MediaPipeException if there is an internal error.
|
||||
*/
|
||||
public void stylizeAsync(
|
||||
MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) {
|
||||
sendLiveStreamData(image, imageProcessingOptions, timestampMs);
|
||||
}
|
||||
|
||||
/** Options for setting up an {@link FaceStylizer}. */
|
||||
@AutoValue
|
||||
public abstract static class FaceStylizerOptions extends TaskOptions {
|
||||
|
||||
/** Builder for {@link FaceStylizerOptions}. */
|
||||
@AutoValue.Builder
|
||||
public abstract static class Builder {
|
||||
/** Sets the base options for the face stylizer task. */
|
||||
public abstract Builder setBaseOptions(BaseOptions value);
|
||||
|
||||
/**
|
||||
* Sets the running mode for the face stylizer task. Default to the image mode. Image stylizer
|
||||
* has three modes:
|
||||
*
|
||||
* <ul>
|
||||
* <li>IMAGE: The mode for stylizeing image on single image inputs.
|
||||
* <li>VIDEO: The mode for stylizeing image on the decoded frames of a video.
|
||||
* <li>LIVE_STREAM: The mode for for stylizeing image on a live stream of input data, such
|
||||
* as from camera. In this mode, {@code setResultListener} must be called to set up a
|
||||
* listener to receive the recognition results asynchronously.
|
||||
* </ul>
|
||||
*/
|
||||
public abstract Builder setRunningMode(RunningMode value);
|
||||
|
||||
/**
|
||||
* Sets an optional {@link ResultListener} to receive the stylization results when the graph
|
||||
* pipeline is done processing an image.
|
||||
*/
|
||||
public abstract Builder setResultListener(ResultListener<FaceStylizerResult, MPImage> value);
|
||||
|
||||
/** Sets an optional {@link ErrorListener}}. */
|
||||
public abstract Builder setErrorListener(ErrorListener value);
|
||||
|
||||
abstract FaceStylizerOptions autoBuild();
|
||||
|
||||
/**
|
||||
* Validates and builds the {@link FaceStylizerOptions} instance.
|
||||
*
|
||||
* @throws IllegalArgumentException if the result listener and the running mode are not
|
||||
* properly configured. The result listener must be set when the face stylizer is in the
|
||||
* live stream mode.
|
||||
*/
|
||||
public final FaceStylizerOptions build() {
|
||||
FaceStylizerOptions options = autoBuild();
|
||||
if (options.runningMode() == RunningMode.LIVE_STREAM) {
|
||||
if (!options.resultListener().isPresent()) {
|
||||
throw new IllegalArgumentException(
|
||||
"The face stylizer is in the live stream mode, a user-defined result listener"
|
||||
+ " must be provided in FaceStylizerOptions.");
|
||||
}
|
||||
}
|
||||
return options;
|
||||
}
|
||||
}
|
||||
|
||||
abstract BaseOptions baseOptions();
|
||||
|
||||
abstract RunningMode runningMode();
|
||||
|
||||
abstract Optional<ResultListener<FaceStylizerResult, MPImage>> resultListener();
|
||||
|
||||
abstract Optional<ErrorListener> errorListener();
|
||||
|
||||
public static Builder builder() {
|
||||
return new AutoValue_FaceStylizer_FaceStylizerOptions.Builder()
|
||||
.setRunningMode(RunningMode.IMAGE);
|
||||
}
|
||||
|
||||
/** Converts an {@link FaceStylizerOptions} to a {@link CalculatorOptions} protobuf message. */
|
||||
@Override
|
||||
public CalculatorOptions convertToCalculatorOptionsProto() {
|
||||
FaceStylizerGraphOptionsProto.FaceStylizerGraphOptions taskOptions =
|
||||
FaceStylizerGraphOptionsProto.FaceStylizerGraphOptions.newBuilder()
|
||||
.setBaseOptions(
|
||||
BaseOptionsProto.BaseOptions.newBuilder()
|
||||
.setUseStreamMode(runningMode() != RunningMode.IMAGE)
|
||||
.mergeFrom(convertBaseOptionsToProto(baseOptions()))
|
||||
.build())
|
||||
.build();
|
||||
|
||||
return CalculatorOptions.newBuilder()
|
||||
.setExtension(FaceStylizerGraphOptionsProto.FaceStylizerGraphOptions.ext, taskOptions)
|
||||
.build();
|
||||
}
|
||||
}
|
||||
}
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
// Copyright 2023 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.
|
||||
|
||||
package com.google.mediapipe.tasks.vision.facestylizer;
|
||||
|
||||
import com.google.auto.value.AutoValue;
|
||||
import com.google.mediapipe.framework.image.MPImage;
|
||||
import com.google.mediapipe.tasks.core.TaskResult;
|
||||
|
||||
/** Represents the stylized image generated by {@link FaceStylizer}. */
|
||||
@AutoValue
|
||||
public abstract class FaceStylizerResult implements TaskResult {
|
||||
|
||||
/**
|
||||
* Creates an {@link FaceStylizerResult} instance from a MPImage.
|
||||
*
|
||||
* @param stylizedImage an MPImage representing the stylized face.
|
||||
* @param timestampMs a timestamp for this result.
|
||||
*/
|
||||
public static FaceStylizerResult create(MPImage stylizedImage, long timestampMs) {
|
||||
return new AutoValue_FaceStylizerResult(stylizedImage, timestampMs);
|
||||
}
|
||||
|
||||
public abstract MPImage stylizedImage();
|
||||
|
||||
@Override
|
||||
public abstract long timestampMs();
|
||||
}
|
||||
+6
-6
@@ -403,10 +403,10 @@ public final class GestureRecognizer extends BaseVisionTaskApi {
|
||||
public abstract Builder setMinTrackingConfidence(Float value);
|
||||
|
||||
/**
|
||||
* Sets the optional {@link ClassifierOptions} controling the canned gestures classifier, such
|
||||
* as score threshold, allow list and deny list of gestures. The categories for canned gesture
|
||||
* classifiers are: ["None", "Closed_Fist", "Open_Palm", "Pointing_Up", "Thumb_Down",
|
||||
* "Thumb_Up", "Victory", "ILoveYou"]
|
||||
* Sets the optional {@link ClassifierOptions} controlling the canned gestures classifier,
|
||||
* such as score threshold, allow list and deny list of gestures. The categories
|
||||
* for canned gesture classifiers are: ["None", "Closed_Fist", "Open_Palm",
|
||||
* "Pointing_Up", "Thumb_Down", "Thumb_Up", "Victory", "ILoveYou"]
|
||||
*
|
||||
* <p>TODO Note this option is subject to change, after scoring merging
|
||||
* calculator is implemented.
|
||||
@@ -415,8 +415,8 @@ public final class GestureRecognizer extends BaseVisionTaskApi {
|
||||
ClassifierOptions classifierOptions);
|
||||
|
||||
/**
|
||||
* Sets the optional {@link ClassifierOptions} controling the custom gestures classifier, such
|
||||
* as score threshold, allow list and deny list of gestures.
|
||||
* Sets the optional {@link ClassifierOptions} controlling the custom gestures classifier,
|
||||
* such as score threshold, allow list and deny list of gestures.
|
||||
*
|
||||
* <p>TODO Note this option is subject to change, after scoring merging
|
||||
* calculator is implemented.
|
||||
|
||||
+5
-5
@@ -302,7 +302,7 @@ public final class ImageSegmenter extends BaseVisionTaskApi {
|
||||
* @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a
|
||||
* region-of-interest.
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link ImageSegmenter} is not
|
||||
* created wtih {@link ResultListener} set in {@link ImageSegmenterOptions}.
|
||||
* created with {@link ResultListener} set in {@link ImageSegmenterOptions}.
|
||||
*/
|
||||
public void segmentWithResultListener(MPImage image) {
|
||||
segmentWithResultListener(image, ImageProcessingOptions.builder().build());
|
||||
@@ -329,7 +329,7 @@ public final class ImageSegmenter extends BaseVisionTaskApi {
|
||||
* @throws IllegalArgumentException if the {@link ImageProcessingOptions} specify a
|
||||
* region-of-interest.
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link ImageSegmenter} is not
|
||||
* created wtih {@link ResultListener} set in {@link ImageSegmenterOptions}.
|
||||
* created with {@link ResultListener} set in {@link ImageSegmenterOptions}.
|
||||
*/
|
||||
public void segmentWithResultListener(
|
||||
MPImage image, ImageProcessingOptions imageProcessingOptions) {
|
||||
@@ -421,7 +421,7 @@ public final class ImageSegmenter extends BaseVisionTaskApi {
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link ImageSegmenter} is not
|
||||
* created wtih {@link ResultListener} set in {@link ImageSegmenterOptions}.
|
||||
* created with {@link ResultListener} set in {@link ImageSegmenterOptions}.
|
||||
*/
|
||||
public void segmentForVideoWithResultListener(MPImage image, long timestampMs) {
|
||||
segmentForVideoWithResultListener(image, ImageProcessingOptions.builder().build(), timestampMs);
|
||||
@@ -444,7 +444,7 @@ public final class ImageSegmenter extends BaseVisionTaskApi {
|
||||
* @param image a MediaPipe {@link MPImage} object for processing.
|
||||
* @param timestampMs the input timestamp (in milliseconds).
|
||||
* @throws MediaPipeException if there is an internal error. Or if {@link ImageSegmenter} is not
|
||||
* created wtih {@link ResultListener} set in {@link ImageSegmenterOptions}.
|
||||
* created with {@link ResultListener} set in {@link ImageSegmenterOptions}.
|
||||
*/
|
||||
public void segmentForVideoWithResultListener(
|
||||
MPImage image, ImageProcessingOptions imageProcessingOptions, long timestampMs) {
|
||||
@@ -519,7 +519,7 @@ public final class ImageSegmenter extends BaseVisionTaskApi {
|
||||
*
|
||||
* <p>If there is no labelmap provided in the model file, empty label list is returned.
|
||||
*/
|
||||
List<String> getLabels() {
|
||||
public List<String> getLabels() {
|
||||
return labels;
|
||||
}
|
||||
|
||||
|
||||
+1
-1
@@ -29,7 +29,7 @@ public abstract class ImageSegmenterResult implements TaskResult {
|
||||
*
|
||||
* @param segmentations a {@link List} of MPImage representing the segmented masks. If OutputType
|
||||
* is CATEGORY_MASK, the masks will be in IMAGE_FORMAT_ALPHA format. If OutputType is
|
||||
* CONFIDENCE_MASK, the masks will be in IMAGE_FORMAT_ALPHA format.
|
||||
* CONFIDENCE_MASK, the masks will be in IMAGE_FORMAT_VEC32F1 format.
|
||||
* @param timestampMs a timestamp for this result.
|
||||
*/
|
||||
// TODO: consolidate output formats across platforms.
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user