Merge branch 'master' into interactive-segmenter-python

This commit is contained in:
Kinar R
2023-04-05 10:51:45 +05:30
committed by GitHub
370 changed files with 15587 additions and 2256 deletions
@@ -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
+4 -4
View File
@@ -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_
@@ -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];
+1
View File
@@ -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",
@@ -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);
});
}
@@ -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
@@ -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(),
@@ -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, &timestamps](
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,
&timestamps](
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];
}
@@ -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) {
-1
View File
@@ -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
+2
View File
@@ -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;
}
@@ -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(
@@ -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>
@@ -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.");
}
}
}
@@ -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();
}
@@ -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>
@@ -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();
}
}
}
@@ -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();
}
@@ -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.
@@ -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;
}
@@ -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