From 5c74ed2ae58eeb6b6f9b18aa47edf52e08a0eccb Mon Sep 17 00:00:00 2001
From: Jiuqiang Tang <jqtang@google.com>
Date: Thu, 12 Jan 2023 08:27:57 -0800
Subject: [PATCH] EmbeddingAggregationCalculator should fill in the
 `timestamp_ms` field of the embedding results in the stream mode.

Per user feedback, the consistency between the packet timestamp and the timestamp field of the embedding result helps reducing the confusion.

PiperOrigin-RevId: 501572379
---
 .../calculators/embedding_aggregation_calculator.cc       | 4 +++-
 .../calculators/embedding_aggregation_calculator_test.cc  | 8 +++++---
 .../processors/embedding_postprocessing_graph_test.cc     | 7 ++++---
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator.cc b/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator.cc
index bae926b7..6e06c4e3 100644
--- a/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator.cc
+++ b/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator.cc
@@ -120,7 +120,9 @@ absl::Status EmbeddingAggregationCalculator::Process(CalculatorContext* cc) {
     }
     kTimestampedEmbeddingsOut(cc).Send(std::move(results));
   } else {
-    kEmbeddingsOut(cc).Send(kEmbeddingsIn(cc));
+    auto result = kEmbeddingsIn(cc).Get();
+    result.set_timestamp_ms(cc->InputTimestamp().Value() / 1000);
+    kEmbeddingsOut(cc).Send(result);
   }
   RET_CHECK(cached_embeddings_.empty());
   return absl::OkStatus();
diff --git a/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator_test.cc b/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator_test.cc
index ebb4d888..f2b2fa1d 100644
--- a/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator_test.cc
+++ b/mediapipe/tasks/cc/components/calculators/embedding_aggregation_calculator_test.cc
@@ -120,7 +120,7 @@ class EmbeddingAggregationCalculatorTest : public tflite_shims::testing::Test {
   CalculatorGraph calculator_graph_;
 };
 
-TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutTimestamps) {
+TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutAggregation) {
   EmbeddingResult embedding = ParseTextProtoOrDie<EmbeddingResult>(
       R"pb(embeddings { head_index: 0 })pb");
 
@@ -129,10 +129,12 @@ TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutTimestamps) {
   MP_ASSERT_OK(Send(embedding));
   MP_ASSERT_OK_AND_ASSIGN(auto result, GetResult<EmbeddingResult>(poller));
 
-  EXPECT_THAT(result, EqualsProto(embedding));
+  EXPECT_THAT(result, EqualsProto(ParseTextProtoOrDie<EmbeddingResult>(
+                          R"pb(timestamp_ms: 0
+                               embeddings { head_index: 0 })pb")));
 }
 
-TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithTimestamps) {
+TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithAggregation) {
   MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph(/*connect_timestamps=*/true));
   MP_ASSERT_OK(Send(ParseTextProtoOrDie<EmbeddingResult>(R"pb(embeddings {
                                                                 head_index: 0
diff --git a/mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph_test.cc b/mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph_test.cc
index 163e46ee..809268a6 100644
--- a/mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph_test.cc
+++ b/mediapipe/tasks/cc/components/processors/embedding_postprocessing_graph_test.cc
@@ -246,7 +246,7 @@ class PostprocessingTest : public tflite_shims::testing::Test {
       absl::make_unique<std::vector<Tensor>>();
 };
 
-TEST_F(PostprocessingTest, SucceedsWithoutTimestamps) {
+TEST_F(PostprocessingTest, SucceedsWithoutAggregation) {
   // Build graph.
   proto::EmbedderOptions options;
   MP_ASSERT_OK_AND_ASSIGN(auto poller,
@@ -261,7 +261,8 @@ TEST_F(PostprocessingTest, SucceedsWithoutTimestamps) {
   MP_ASSERT_OK_AND_ASSIGN(auto results, GetResult<EmbeddingResult>(poller));
 
   // Validate results.
-  EXPECT_FALSE(results.has_timestamp_ms());
+  EXPECT_TRUE(results.has_timestamp_ms());
+  EXPECT_EQ(results.timestamp_ms(), 0);
   EXPECT_EQ(results.embeddings_size(), 1);
   EXPECT_EQ(results.embeddings(0).head_index(), 0);
   EXPECT_EQ(results.embeddings(0).head_name(), "feature");
@@ -273,7 +274,7 @@ TEST_F(PostprocessingTest, SucceedsWithoutTimestamps) {
   }
 }
 
-TEST_F(PostprocessingTest, SucceedsWithTimestamps) {
+TEST_F(PostprocessingTest, SucceedsWithAggregation) {
   // Build graph.
   proto::EmbedderOptions options;
   MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph(kMobileNetV3Embedder, options,