From d9080c0d3819deef26249650ac067bfa036e43da Mon Sep 17 00:00:00 2001 From: Kinar Date: Tue, 7 Nov 2023 07:02:08 -0800 Subject: [PATCH] Updated the Image Embedder C API and added tests for cosine similarity --- .../containers/embedding_result_converter.cc | 23 +++ .../containers/embedding_result_converter.h | 4 + .../c/text/text_embedder/text_embedder.cc | 27 +++ .../c/text/text_embedder/text_embedder.h | 9 + .../text/text_embedder/text_embedder_test.cc | 36 +++- .../c/vision/image_embedder/image_embedder.cc | 169 ++++++++++++++++- .../c/vision/image_embedder/image_embedder.h | 34 +++- .../image_embedder/image_embedder_test.cc | 176 +++++++++++++++++- 8 files changed, 459 insertions(+), 19 deletions(-) diff --git a/mediapipe/tasks/c/components/containers/embedding_result_converter.cc b/mediapipe/tasks/c/components/containers/embedding_result_converter.cc index ba72c0994..cd4850c18 100644 --- a/mediapipe/tasks/c/components/containers/embedding_result_converter.cc +++ b/mediapipe/tasks/c/components/containers/embedding_result_converter.cc @@ -66,6 +66,29 @@ void CppConvertToEmbeddingResult( } } +void ConvertToCppEmbedding( + const Embedding& in, // C struct as input + mediapipe::tasks::components::containers::Embedding* out) { + // Handle float embeddings + if (in.float_embedding != nullptr) { + out->float_embedding.assign(in.float_embedding, + in.float_embedding + in.values_count); + } + + // Handle quantized embeddings + if (in.quantized_embedding != nullptr) { + out->quantized_embedding.assign(in.quantized_embedding, + in.quantized_embedding + in.values_count); + } + + out->head_index = in.head_index; + + // Copy head_name if it is present. + if (in.head_name) { + out->head_name = std::make_optional(std::string(in.head_name)); + } +} + void CppCloseEmbeddingResult(EmbeddingResult* in) { for (uint32_t i = 0; i < in->embeddings_count; ++i) { auto embedding_in = in->embeddings[i]; diff --git a/mediapipe/tasks/c/components/containers/embedding_result_converter.h b/mediapipe/tasks/c/components/containers/embedding_result_converter.h index 15bcdbdd0..5ba4e4e2b 100644 --- a/mediapipe/tasks/c/components/containers/embedding_result_converter.h +++ b/mediapipe/tasks/c/components/containers/embedding_result_converter.h @@ -29,6 +29,10 @@ void CppConvertToEmbeddingResult( const mediapipe::tasks::components::containers::EmbeddingResult& in, EmbeddingResult* out); +void ConvertToCppEmbedding( + const Embedding& in, + mediapipe::tasks::components::containers::Embedding* out); + void CppCloseEmbedding(Embedding* in); void CppCloseEmbeddingResult(EmbeddingResult* in); diff --git a/mediapipe/tasks/c/text/text_embedder/text_embedder.cc b/mediapipe/tasks/c/text/text_embedder/text_embedder.cc index c98b958f5..85b15e923 100644 --- a/mediapipe/tasks/c/text/text_embedder/text_embedder.cc +++ b/mediapipe/tasks/c/text/text_embedder/text_embedder.cc @@ -29,6 +29,8 @@ namespace mediapipe::tasks::c::text::text_embedder { namespace { + +using ::mediapipe::tasks::c::components::containers::ConvertToCppEmbedding; using ::mediapipe::tasks::c::components::containers::CppCloseEmbeddingResult; using ::mediapipe::tasks::c::components::containers:: CppConvertToEmbeddingResult; @@ -36,6 +38,7 @@ using ::mediapipe::tasks::c::components::processors:: CppConvertToEmbedderOptions; using ::mediapipe::tasks::c::core::CppConvertToBaseOptions; using ::mediapipe::tasks::text::text_embedder::TextEmbedder; +typedef ::mediapipe::tasks::components::containers::Embedding CppEmbedding; int CppProcessError(absl::Status status, char** error_msg) { if (error_msg) { @@ -91,6 +94,24 @@ int CppTextEmbedderClose(void* embedder, char** error_msg) { return 0; } +int CppTextEmbedderCosineSimilarity(const Embedding& u, const Embedding& v, + double* similarity, char** error_msg) { + CppEmbedding cpp_u; + ConvertToCppEmbedding(u, &cpp_u); + CppEmbedding cpp_v; + ConvertToCppEmbedding(v, &cpp_v); + auto status_or_similarity = + mediapipe::tasks::text::text_embedder::TextEmbedder::CosineSimilarity( + cpp_u, cpp_v); + if (status_or_similarity.ok()) { + *similarity = status_or_similarity.value(); + } else { + ABSL_LOG(ERROR) << "Cannot computer cosine similarity."; + return CppProcessError(status_or_similarity.status(), error_msg); + } + return 0; +} + } // namespace mediapipe::tasks::c::text::text_embedder extern "C" { @@ -116,4 +137,10 @@ int text_embedder_close(void* embedder, char** error_ms) { embedder, error_ms); } +int cosine_similarity(const Embedding& u, const Embedding& v, + double* similarity, char** error_msg) { + return mediapipe::tasks::c::text::text_embedder:: + CppTextEmbedderCosineSimilarity(u, v, similarity, error_msg); +} + } // extern "C" diff --git a/mediapipe/tasks/c/text/text_embedder/text_embedder.h b/mediapipe/tasks/c/text/text_embedder/text_embedder.h index c9ccf816b..7fee3c737 100644 --- a/mediapipe/tasks/c/text/text_embedder/text_embedder.h +++ b/mediapipe/tasks/c/text/text_embedder/text_embedder.h @@ -67,6 +67,15 @@ MP_EXPORT void text_embedder_close_result(TextEmbedderResult* result); // allocated for the error message. MP_EXPORT int text_embedder_close(void* embedder, char** error_msg = nullptr); +// Utility function to compute cosine similarity [1] between two embeddings. +// May return an InvalidArgumentError if e.g. the embeddings are of different +// types (quantized vs. float), have different sizes, or have a an L2-norm of +// 0. +// +// [1]: https://en.wikipedia.org/wiki/Cosine_similarity +MP_EXPORT int cosine_similarity(const Embedding& u, const Embedding& v, + double* similarity, char** error_msg = nullptr); + #ifdef __cplusplus } // extern C #endif diff --git a/mediapipe/tasks/c/text/text_embedder/text_embedder_test.cc b/mediapipe/tasks/c/text/text_embedder/text_embedder_test.cc index c823e01b4..c6b19d9f5 100644 --- a/mediapipe/tasks/c/text/text_embedder/text_embedder_test.cc +++ b/mediapipe/tasks/c/text/text_embedder/text_embedder_test.cc @@ -32,7 +32,12 @@ using testing::HasSubstr; constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/text/"; constexpr char kTestBertModelPath[] = "mobilebert_embedding_with_metadata.tflite"; -constexpr char kTestString[] = "It's beautiful outside."; +constexpr char kTestString0[] = + "When you go to this restaurant, they hold the pancake upside-down " + "before they hand it to you. It's a great gimmick."; +constexpr char kTestString1[] = + "Let's make a plan to steal the declaration of independence."; +constexpr float kPrecision = 1e-3; std::string GetFullPath(absl::string_view file_name) { return JoinPath("./", kTestDataDirectory, file_name); @@ -51,7 +56,7 @@ TEST(TextEmbedderTest, SmokeTest) { EXPECT_NE(embedder, nullptr); TextEmbedderResult result; - text_embedder_embed(embedder, kTestString, &result); + text_embedder_embed(embedder, kTestString0, &result); EXPECT_EQ(result.embeddings_count, 1); EXPECT_EQ(result.embeddings[0].values_count, 512); @@ -59,6 +64,33 @@ TEST(TextEmbedderTest, SmokeTest) { text_embedder_close(embedder); } +TEST(TextEmbedderTest, SucceedsWithCosineSimilarity) { + std::string model_path = GetFullPath(kTestBertModelPath); + TextEmbedderOptions options = { + /* base_options= */ {/* model_asset_buffer= */ nullptr, + /* model_asset_path= */ model_path.c_str()}, + /* embedder_options= */ + {/* l2_normalize= */ false, + /* quantize= */ false}}; + + void* embedder = text_embedder_create(&options); + EXPECT_NE(embedder, nullptr); + + // Extract both embeddings. + TextEmbedderResult result0; + text_embedder_embed(embedder, kTestString0, &result0); + TextEmbedderResult result1; + text_embedder_embed(embedder, kTestString1, &result1); + + // Check cosine similarity. + double similarity; + cosine_similarity(result0.embeddings[0], result1.embeddings[0], + &similarity); + double expected_similarity = 0.98077; + EXPECT_LE(abs(similarity - expected_similarity), kPrecision); + text_embedder_close(embedder); +} + TEST(TextEmbedderTest, ErrorHandling) { // It is an error to set neither the asset buffer nor the path. TextEmbedderOptions options = { diff --git a/mediapipe/tasks/c/vision/image_embedder/image_embedder.cc b/mediapipe/tasks/c/vision/image_embedder/image_embedder.cc index 4b42034cb..00a571e90 100644 --- a/mediapipe/tasks/c/vision/image_embedder/image_embedder.cc +++ b/mediapipe/tasks/c/vision/image_embedder/image_embedder.cc @@ -15,6 +15,8 @@ limitations under the License. #include "mediapipe/tasks/c/vision/image_embedder/image_embedder.h" +#include +#include #include #include @@ -26,6 +28,7 @@ limitations under the License. #include "mediapipe/tasks/c/components/containers/embedding_result_converter.h" #include "mediapipe/tasks/c/components/processors/embedder_options_converter.h" #include "mediapipe/tasks/c/core/base_options_converter.h" +#include "mediapipe/tasks/cc/vision/core/running_mode.h" #include "mediapipe/tasks/cc/vision/image_embedder/image_embedder.h" #include "mediapipe/tasks/cc/vision/utils/image_utils.h" @@ -33,6 +36,7 @@ namespace mediapipe::tasks::c::vision::image_embedder { namespace { +using ::mediapipe::tasks::c::components::containers::ConvertToCppEmbedding; using ::mediapipe::tasks::c::components::containers::CppCloseEmbeddingResult; using ::mediapipe::tasks::c::components::containers:: CppConvertToEmbeddingResult; @@ -40,7 +44,11 @@ using ::mediapipe::tasks::c::components::processors:: CppConvertToEmbedderOptions; using ::mediapipe::tasks::c::core::CppConvertToBaseOptions; using ::mediapipe::tasks::vision::CreateImageFromBuffer; +using ::mediapipe::tasks::vision::core::RunningMode; using ::mediapipe::tasks::vision::image_embedder::ImageEmbedder; +typedef ::mediapipe::tasks::components::containers::Embedding CppEmbedding; +typedef ::mediapipe::tasks::vision::image_embedder::ImageEmbedderResult + CppImageEmbedderResult; int CppProcessError(absl::Status status, char** error_msg) { if (error_msg) { @@ -59,6 +67,54 @@ ImageEmbedder* CppImageEmbedderCreate(const ImageEmbedderOptions& options, CppConvertToBaseOptions(options.base_options, &cpp_options->base_options); CppConvertToEmbedderOptions(options.embedder_options, &cpp_options->embedder_options); + cpp_options->running_mode = static_cast(options.running_mode); + + // Enable callback for processing live stream data when the running mode is + // set to RunningMode::LIVE_STREAM. + if (cpp_options->running_mode == RunningMode::LIVE_STREAM) { + if (options.result_callback == nullptr) { + const absl::Status status = absl::InvalidArgumentError( + "Provided null pointer to callback function."); + ABSL_LOG(ERROR) << "Failed to create ImageEmbedder: " << status; + CppProcessError(status, error_msg); + return nullptr; + } + + ImageEmbedderOptions::result_callback_fn result_callback = + options.result_callback; + cpp_options->result_callback = + [result_callback](absl::StatusOr cpp_result, + const Image& image, int64_t timestamp) { + char* error_msg = nullptr; + + if (!cpp_result.ok()) { + ABSL_LOG(ERROR) + << "Embedding extraction failed: " << cpp_result.status(); + CppProcessError(cpp_result.status(), &error_msg); + result_callback(nullptr, MpImage(), timestamp, error_msg); + free(error_msg); + return; + } + + // Result is valid for the lifetime of the callback function. + ImageEmbedderResult result; + CppConvertToEmbeddingResult(*cpp_result, &result); + + const auto& image_frame = image.GetImageFrameSharedPtr(); + const MpImage mp_image = { + .type = MpImage::IMAGE_FRAME, + .image_frame = { + .format = static_cast<::ImageFormat>(image_frame->Format()), + .image_buffer = image_frame->PixelData(), + .width = image_frame->Width(), + .height = image_frame->Height()}}; + + result_callback(&result, mp_image, timestamp, + /* error_msg= */ nullptr); + + CppCloseEmbeddingResult(&result); + }; + } auto embedder = ImageEmbedder::Create(std::move(cpp_options)); if (!embedder.ok()) { @@ -72,10 +128,10 @@ ImageEmbedder* CppImageEmbedderCreate(const ImageEmbedderOptions& options, int CppImageEmbedderEmbed(void* embedder, const MpImage* image, ImageEmbedderResult* result, char** error_msg) { if (image->type == MpImage::GPU_BUFFER) { - absl::Status status = - absl::InvalidArgumentError("gpu buffer not supported yet"); + const absl::Status status = + absl::InvalidArgumentError("GPU Buffer not supported yet."); - ABSL_LOG(ERROR) << "Classification failed: " << status.message(); + ABSL_LOG(ERROR) << "Embedding extraction failed: " << status.message(); return CppProcessError(status, error_msg); } @@ -92,13 +148,75 @@ int CppImageEmbedderEmbed(void* embedder, const MpImage* image, auto cpp_embedder = static_cast(embedder); auto cpp_result = cpp_embedder->Embed(*img); if (!cpp_result.ok()) { - ABSL_LOG(ERROR) << "Classification failed: " << cpp_result.status(); + ABSL_LOG(ERROR) << "Embedding extraction failed: " << cpp_result.status(); return CppProcessError(cpp_result.status(), error_msg); } CppConvertToEmbeddingResult(*cpp_result, result); return 0; } +int CppImageEmbedderEmbedForVideo(void* embedder, const MpImage* image, + int64_t timestamp_ms, + ImageEmbedderResult* result, + char** error_msg) { + if (image->type == MpImage::GPU_BUFFER) { + absl::Status status = + absl::InvalidArgumentError("GPU Buffer not supported yet"); + + ABSL_LOG(ERROR) << "Embedding extraction failed: " << status.message(); + return CppProcessError(status, error_msg); + } + + const auto img = CreateImageFromBuffer( + static_cast(image->image_frame.format), + image->image_frame.image_buffer, image->image_frame.width, + image->image_frame.height); + + if (!img.ok()) { + ABSL_LOG(ERROR) << "Failed to create Image: " << img.status(); + return CppProcessError(img.status(), error_msg); + } + + auto cpp_embedder = static_cast(embedder); + auto cpp_result = cpp_embedder->EmbedForVideo(*img, timestamp_ms); + if (!cpp_result.ok()) { + ABSL_LOG(ERROR) << "Embedding extraction failed: " << cpp_result.status(); + return CppProcessError(cpp_result.status(), error_msg); + } + CppConvertToEmbeddingResult(*cpp_result, result); + return 0; +} + +int CppImageEmbedderEmbedAsync(void* embedder, const MpImage* image, + int64_t timestamp_ms, char** error_msg) { + if (image->type == MpImage::GPU_BUFFER) { + absl::Status status = + absl::InvalidArgumentError("GPU Buffer not supported yet"); + + ABSL_LOG(ERROR) << "Embedding extraction failed: " << status.message(); + return CppProcessError(status, error_msg); + } + + const auto img = CreateImageFromBuffer( + static_cast(image->image_frame.format), + image->image_frame.image_buffer, image->image_frame.width, + image->image_frame.height); + + if (!img.ok()) { + ABSL_LOG(ERROR) << "Failed to create Image: " << img.status(); + return CppProcessError(img.status(), error_msg); + } + + auto cpp_embedder = static_cast(embedder); + auto cpp_result = cpp_embedder->EmbedAsync(*img, timestamp_ms); + if (!cpp_result.ok()) { + ABSL_LOG(ERROR) << "Data preparation for the embedding extraction failed: " + << cpp_result; + return CppProcessError(cpp_result, error_msg); + } + return 0; +} + void CppImageEmbedderCloseResult(ImageEmbedderResult* result) { CppCloseEmbeddingResult(result); } @@ -114,6 +232,24 @@ int CppImageEmbedderClose(void* embedder, char** error_msg) { return 0; } +int CppImageEmbedderCosineSimilarity(const Embedding& u, const Embedding& v, + double* similarity, char** error_msg) { + CppEmbedding cpp_u; + ConvertToCppEmbedding(u, &cpp_u); + CppEmbedding cpp_v; + ConvertToCppEmbedding(v, &cpp_v); + auto status_or_similarity = + mediapipe::tasks::vision::image_embedder::ImageEmbedder::CosineSimilarity( + cpp_u, cpp_v); + if (status_or_similarity.ok()) { + *similarity = status_or_similarity.value(); + } else { + ABSL_LOG(ERROR) << "Cannot computer cosine similarity."; + return CppProcessError(status_or_similarity.status(), error_msg); + } + return 0; +} + } // namespace mediapipe::tasks::c::vision::image_embedder extern "C" { @@ -130,14 +266,35 @@ int image_embedder_embed_image(void* embedder, const MpImage* image, embedder, image, result, error_msg); } +int image_embedder_embed_for_video(void* embedder, const MpImage* image, + int64_t timestamp_ms, + ImageEmbedderResult* result, + char** error_msg) { + return mediapipe::tasks::c::vision::image_embedder:: + CppImageEmbedderEmbedForVideo(embedder, image, timestamp_ms, result, + error_msg); +} + +int image_embedder_embed_async(void* embedder, const MpImage* image, + int64_t timestamp_ms, char** error_msg) { + return mediapipe::tasks::c::vision::image_embedder:: + CppImageEmbedderEmbedAsync(embedder, image, timestamp_ms, error_msg); +} + void image_embedder_close_result(ImageEmbedderResult* result) { mediapipe::tasks::c::vision::image_embedder::CppImageEmbedderCloseResult( result); } -int image_embedder_close(void* embedder, char** error_ms) { +int image_embedder_close(void* embedder, char** error_msg) { return mediapipe::tasks::c::vision::image_embedder::CppImageEmbedderClose( - embedder, error_ms); + embedder, error_msg); +} + +int cosine_similarity(const Embedding& u, const Embedding& v, + double* similarity, char** error_msg) { + return mediapipe::tasks::c::vision::image_embedder:: + CppImageEmbedderCosineSimilarity(u, v, similarity, error_msg); } } // extern "C" diff --git a/mediapipe/tasks/c/vision/image_embedder/image_embedder.h b/mediapipe/tasks/c/vision/image_embedder/image_embedder.h index cefe97b1b..550e30781 100644 --- a/mediapipe/tasks/c/vision/image_embedder/image_embedder.h +++ b/mediapipe/tasks/c/vision/image_embedder/image_embedder.h @@ -92,9 +92,16 @@ struct ImageEmbedderOptions { // 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. - typedef void (*result_callback_fn)(ImageEmbedderResult*, const MpImage*, - int64_t); + // to RunningMode::LIVE_STREAM. Arguments of the callback function include: + // the pointer to embedding result, the image that result was obtained + // on, the timestamp relevant to embedding extraction results and pointer to + // error message in case of any failure. The validity of the passed arguments + // is true for the lifetime of the callback function. + // + // A caller is responsible for closing image embedder result. + typedef void (*result_callback_fn)(ImageEmbedderResult* result, + const MpImage image, int64_t timestamp_ms, + char* error_msg); result_callback_fn result_callback; }; @@ -110,12 +117,20 @@ MP_EXPORT void* image_embedder_create(struct ImageEmbedderOptions* options, // If an error occurs, returns an error code and sets the error parameter to an // an error message (if `error_msg` is not nullptr). You must free the memory // allocated for the error message. -// -// TODO: Add API for video and live stream processing. MP_EXPORT int image_embedder_embed_image(void* embedder, const MpImage* image, ImageEmbedderResult* result, char** error_msg = nullptr); +MP_EXPORT int image_embedder_embed_for_video(void* embedder, + const MpImage* image, + int64_t timestamp_ms, + ImageEmbedderResult* result, + char** error_msg = nullptr); + +MP_EXPORT int image_embedder_embed_async(void* embedder, const MpImage* image, + int64_t timestamp_ms, + char** error_msg = nullptr); + // Frees the memory allocated inside a ImageEmbedderResult result. // Does not free the result pointer itself. MP_EXPORT void image_embedder_close_result(ImageEmbedderResult* result); @@ -126,6 +141,15 @@ MP_EXPORT void image_embedder_close_result(ImageEmbedderResult* result); // allocated for the error message. MP_EXPORT int image_embedder_close(void* embedder, char** error_msg = nullptr); +// Utility function to compute cosine similarity [1] between two embeddings. +// May return an InvalidArgumentError if e.g. the embeddings are of different +// types (quantized vs. float), have different sizes, or have a an L2-norm of +// 0. +// +// [1]: https://en.wikipedia.org/wiki/Cosine_similarity +MP_EXPORT int cosine_similarity(const Embedding& u, const Embedding& v, + double* similarity, char** error_msg = nullptr); + #ifdef __cplusplus } // extern C #endif diff --git a/mediapipe/tasks/c/vision/image_embedder/image_embedder_test.cc b/mediapipe/tasks/c/vision/image_embedder/image_embedder_test.cc index fb6b0b628..d31b8b736 100644 --- a/mediapipe/tasks/c/vision/image_embedder/image_embedder_test.cc +++ b/mediapipe/tasks/c/vision/image_embedder/image_embedder_test.cc @@ -37,13 +37,27 @@ constexpr char kTestDataDirectory[] = "/mediapipe/tasks/testdata/vision/"; constexpr char kModelName[] = "mobilenet_v3_small_100_224_embedder.tflite"; constexpr char kImageFile[] = "burger.jpg"; constexpr float kPrecision = 1e-6; +constexpr int kIterations = 100; std::string GetFullPath(absl::string_view file_name) { return JoinPath("./", kTestDataDirectory, file_name); } -TEST(ImageEmbedderTest, SmokeTest) { - const auto image = DecodeImageFromFile(GetFullPath("burger.jpg")); +// Utility function to check the sizes, head_index and head_names of a result +// produced by kMobileNetV3Embedder. +void CheckMobileNetV3Result(const ImageEmbedderResult& result, bool quantized) { + EXPECT_EQ(result.embeddings_count, 1); + EXPECT_EQ(result.embeddings[0].head_index, 0); + EXPECT_EQ(std::string{result.embeddings[0].head_name}, "feature"); + if (quantized) { + EXPECT_EQ(result.embeddings[0].values_count, 1024); + } else { + EXPECT_EQ(result.embeddings[0].values_count, 1024); + } +} + +TEST(ImageEmbedderTest, ImageModeTest) { + const auto image = DecodeImageFromFile(GetFullPath(kImageFile)); ASSERT_TRUE(image.ok()); const std::string model_path = GetFullPath(kModelName); @@ -53,8 +67,7 @@ TEST(ImageEmbedderTest, SmokeTest) { /* running_mode= */ RunningMode::IMAGE, /* embedder_options= */ {/* l2_normalize= */ true, - /* quantize= */ false}, - }; + /* quantize= */ false}}; void* embedder = image_embedder_create(&options); EXPECT_NE(embedder, nullptr); @@ -70,12 +83,163 @@ TEST(ImageEmbedderTest, SmokeTest) { ImageEmbedderResult result; image_embedder_embed_image(embedder, &mp_image, &result); - EXPECT_EQ(result.embeddings_count, 1); + CheckMobileNetV3Result(result, false); EXPECT_NEAR(result.embeddings[0].float_embedding[0], -0.0142344, kPrecision); image_embedder_close_result(&result); image_embedder_close(embedder); } +TEST(ImageEmbedderTest, SucceedsWithCosineSimilarity) { + const auto image = DecodeImageFromFile(GetFullPath("burger.jpg")); + ASSERT_TRUE(image.ok()); + const auto crop = DecodeImageFromFile(GetFullPath("burger_crop.jpg")); + ASSERT_TRUE(crop.ok()); + + const std::string model_path = GetFullPath(kModelName); + ImageEmbedderOptions options = { + /* base_options= */ {/* model_asset_buffer= */ nullptr, + /* model_asset_path= */ model_path.c_str()}, + /* running_mode= */ RunningMode::IMAGE, + /* embedder_options= */ + {/* l2_normalize= */ true, + /* quantize= */ false}}; + + void* embedder = image_embedder_create(&options); + EXPECT_NE(embedder, nullptr); + + const MpImage mp_image = { + .type = MpImage::IMAGE_FRAME, + .image_frame = { + .format = static_cast( + image->GetImageFrameSharedPtr()->Format()), + .image_buffer = image->GetImageFrameSharedPtr()->PixelData(), + .width = image->GetImageFrameSharedPtr()->Width(), + .height = image->GetImageFrameSharedPtr()->Height()}}; + + const MpImage mp_crop = { + .type = MpImage::IMAGE_FRAME, + .image_frame = { + .format = static_cast( + crop->GetImageFrameSharedPtr()->Format()), + .image_buffer = crop->GetImageFrameSharedPtr()->PixelData(), + .width = crop->GetImageFrameSharedPtr()->Width(), + .height = crop->GetImageFrameSharedPtr()->Height()}}; + + // Extract both embeddings. + ImageEmbedderResult image_result; + image_embedder_embed_image(embedder, &mp_image, &image_result); + ImageEmbedderResult crop_result; + image_embedder_embed_image(embedder, &mp_crop, &crop_result); + + // Check results. + CheckMobileNetV3Result(image_result, false); + CheckMobileNetV3Result(crop_result, false); + // Check cosine similarity. + double similarity; + cosine_similarity(image_result.embeddings[0], crop_result.embeddings[0], + &similarity); + double expected_similarity = 0.925519; + EXPECT_LE(abs(similarity - expected_similarity), kPrecision); + image_embedder_close_result(&image_result); + image_embedder_close_result(&crop_result); + image_embedder_close(embedder); +} + +TEST(ImageEmbedderTest, VideoModeTest) { + const auto image = DecodeImageFromFile(GetFullPath(kImageFile)); + ASSERT_TRUE(image.ok()); + + const std::string model_path = GetFullPath(kModelName); + ImageEmbedderOptions options = { + /* base_options= */ {/* model_asset_buffer= */ nullptr, + /* model_asset_path= */ model_path.c_str()}, + /* running_mode= */ RunningMode::VIDEO, + /* embedder_options= */ + {/* l2_normalize= */ true, + /* quantize= */ false}}; + + void* embedder = image_embedder_create(&options); + EXPECT_NE(embedder, nullptr); + + const auto& image_frame = image->GetImageFrameSharedPtr(); + const MpImage mp_image = { + .type = MpImage::IMAGE_FRAME, + .image_frame = {.format = static_cast(image_frame->Format()), + .image_buffer = image_frame->PixelData(), + .width = image_frame->Width(), + .height = image_frame->Height()}}; + + for (int i = 0; i < kIterations; ++i) { + ImageEmbedderResult result; + image_embedder_embed_for_video(embedder, &mp_image, i, &result); + CheckMobileNetV3Result(result, false); + EXPECT_NEAR(result.embeddings[0].float_embedding[0], -0.0142344, + kPrecision); + image_embedder_close_result(&result); + } + image_embedder_close(embedder); +} + +// A structure to support LiveStreamModeTest below. This structure holds a +// static method `Fn` for a callback function of C API. A `static` qualifier +// allows to take an address of the method to follow API style. Another static +// struct member is `last_timestamp` that is used to verify that current +// timestamp is greater than the previous one. +struct LiveStreamModeCallback { + static int64_t last_timestamp; + static void Fn(ImageEmbedderResult* embedder_result, const MpImage image, + int64_t timestamp, char* error_msg) { + ASSERT_NE(embedder_result, nullptr); + ASSERT_EQ(error_msg, nullptr); + CheckMobileNetV3Result(*embedder_result, false); + EXPECT_NEAR(embedder_result->embeddings[0].float_embedding[0], -0.0142344, + kPrecision); + EXPECT_GT(image.image_frame.width, 0); + EXPECT_GT(image.image_frame.height, 0); + EXPECT_GT(timestamp, last_timestamp); + last_timestamp++; + } +}; +int64_t LiveStreamModeCallback::last_timestamp = -1; + +TEST(ImageEmbedderTest, LiveStreamModeTest) { + const auto image = DecodeImageFromFile(GetFullPath(kImageFile)); + ASSERT_TRUE(image.ok()); + + const std::string model_path = GetFullPath(kModelName); + + ImageEmbedderOptions options = { + /* base_options= */ {/* model_asset_buffer= */ nullptr, + /* model_asset_path= */ model_path.c_str()}, + /* running_mode= */ RunningMode::LIVE_STREAM, + /* embedder_options= */ + {/* l2_normalize= */ true, + /* quantize= */ false}, + /* result_callback= */ LiveStreamModeCallback::Fn, + }; + + void* embedder = image_embedder_create(&options); + EXPECT_NE(embedder, nullptr); + + const auto& image_frame = image->GetImageFrameSharedPtr(); + const MpImage mp_image = { + .type = MpImage::IMAGE_FRAME, + .image_frame = {.format = static_cast(image_frame->Format()), + .image_buffer = image_frame->PixelData(), + .width = image_frame->Width(), + .height = image_frame->Height()}}; + + for (int i = 0; i < kIterations; ++i) { + EXPECT_GE(image_embedder_embed_async(embedder, &mp_image, i), 0); + } + image_embedder_close(embedder); + + // Due to the flow limiter, the total of outputs might be smaller than the + // number of iterations. + EXPECT_LE(LiveStreamModeCallback::last_timestamp, kIterations); + EXPECT_GT(LiveStreamModeCallback::last_timestamp, 0); +} + TEST(ImageEmbedderTest, InvalidArgumentHandling) { // It is an error to set neither the asset buffer nor the path. ImageEmbedderOptions options = { @@ -111,7 +275,7 @@ TEST(ImageEmbedderTest, FailedEmbeddingHandling) { ImageEmbedderResult result; char* error_msg; image_embedder_embed_image(embedder, &mp_image, &result, &error_msg); - EXPECT_THAT(error_msg, HasSubstr("gpu buffer not supported yet")); + EXPECT_THAT(error_msg, HasSubstr("GPU Buffer not supported yet.")); free(error_msg); image_embedder_close(embedder); }