From 8156da341833e9d6c8042a0cef1494770458c8a0 Mon Sep 17 00:00:00 2001 From: Jiuqiang Tang Date: Thu, 12 Jan 2023 13:52:27 -0800 Subject: [PATCH] ClassificationAggregationCalculator should fill in the `timestamp_ms` field of the classification results in the stream mode. Per user feedback, the consistency between the packet timestamp and the timestamp field of the classification result helps reducing the confusion. PiperOrigin-RevId: 501657922 --- .../audio_classifier/audio_classifier_test.cc | 3 +- .../classification_aggregation_calculator.cc | 1 + ...ssification_aggregation_calculator_test.cc | 7 ++- ...lassification_postprocessing_graph_test.cc | 4 ++ .../test/vision/image_classifier_test.py | 57 ++++++++++++------- 5 files changed, 46 insertions(+), 26 deletions(-) diff --git a/mediapipe/tasks/cc/audio/audio_classifier/audio_classifier_test.cc b/mediapipe/tasks/cc/audio/audio_classifier/audio_classifier_test.cc index 596b910f..2d5b221a 100644 --- a/mediapipe/tasks/cc/audio/audio_classifier/audio_classifier_test.cc +++ b/mediapipe/tasks/cc/audio/audio_classifier/audio_classifier_test.cc @@ -143,8 +143,9 @@ void CheckStreamingModeResults(std::vector outputs) { EXPECT_EQ(outputs.size(), 5); // Ignore last result, which operates on a too small chunk to return relevant // results. + std::vector timestamps_ms = {0, 975, 1950, 2925}; for (int i = 0; i < outputs.size() - 1; i++) { - EXPECT_FALSE(outputs[i].timestamp_ms.has_value()); + EXPECT_EQ(outputs[i].timestamp_ms.value(), timestamps_ms[i]); EXPECT_EQ(outputs[i].classifications.size(), 1); EXPECT_EQ(outputs[i].classifications[0].head_index, 0); EXPECT_EQ(outputs[i].classifications[0].head_name, "scores"); diff --git a/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.cc b/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.cc index ad2c668c..145076cd 100644 --- a/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.cc +++ b/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator.cc @@ -188,6 +188,7 @@ ClassificationAggregationCalculator::ConvertToClassificationResult( *classifications->mutable_classification_list() = std::move(classification_lists[i]); } + result.set_timestamp_ms(cc->InputTimestamp().Value() / 1000); cached_classifications_.erase(cc->InputTimestamp().Value()); return result; } diff --git a/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator_test.cc b/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator_test.cc index 1bc8cafd..811d7054 100644 --- a/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator_test.cc +++ b/mediapipe/tasks/cc/components/calculators/classification_aggregation_calculator_test.cc @@ -150,14 +150,15 @@ class ClassificationAggregationCalculatorTest CalculatorGraph calculator_graph_; }; -TEST_F(ClassificationAggregationCalculatorTest, SucceedsWithoutTimestamps) { +TEST_F(ClassificationAggregationCalculatorTest, SucceedsWithoutAggregation) { MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph()); MP_ASSERT_OK(Send({MakeClassificationList(0), MakeClassificationList(1)})); MP_ASSERT_OK_AND_ASSIGN(auto result, GetResult(poller)); EXPECT_THAT(result, EqualsProto(ParseTextProtoOrDie( - R"pb(classifications { + R"pb(timestamp_ms: 0, + classifications { head_index: 0 head_name: "foo" classification_list { classification { index: 0 } } @@ -169,7 +170,7 @@ TEST_F(ClassificationAggregationCalculatorTest, SucceedsWithoutTimestamps) { })pb"))); } -TEST_F(ClassificationAggregationCalculatorTest, SucceedsWithTimestamps) { +TEST_F(ClassificationAggregationCalculatorTest, SucceedsWithAggregation) { MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph(/*connect_timestamps=*/true)); MP_ASSERT_OK(Send({MakeClassificationList(0), MakeClassificationList(1)})); MP_ASSERT_OK(Send( diff --git a/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph_test.cc b/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph_test.cc index 8eb6f3c3..a11bad71 100644 --- a/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph_test.cc +++ b/mediapipe/tasks/cc/components/processors/classification_postprocessing_graph_test.cc @@ -534,6 +534,7 @@ TEST_F(PostprocessingTest, SucceedsWithoutMetadata) { // Validate results. EXPECT_THAT(results, EqualsProto(ParseTextProtoOrDie(R"pb( + timestamp_ms: 0, classifications { head_index: 0 classification_list { @@ -567,6 +568,7 @@ TEST_F(PostprocessingTest, SucceedsWithMetadata) { // Validate results. EXPECT_THAT( results, EqualsProto(ParseTextProtoOrDie(R"pb( + timestamp_ms: 0, classifications { head_index: 0 head_name: "probability" @@ -603,6 +605,7 @@ TEST_F(PostprocessingTest, SucceedsWithScoreCalibration) { // Validate results. EXPECT_THAT( results, EqualsProto(ParseTextProtoOrDie(R"pb( + timestamp_ms: 0, classifications { head_index: 0 head_name: "probability" @@ -646,6 +649,7 @@ TEST_F(PostprocessingTest, SucceedsWithMultipleHeads) { // Validate results. EXPECT_THAT( results, EqualsProto(ParseTextProtoOrDie(R"pb( + timestamp_ms: 0, classifications { head_index: 0 head_name: "yamnet_classification" diff --git a/mediapipe/tasks/python/test/vision/image_classifier_test.py b/mediapipe/tasks/python/test/vision/image_classifier_test.py index cbeaf36b..b47efb32 100644 --- a/mediapipe/tasks/python/test/vision/image_classifier_test.py +++ b/mediapipe/tasks/python/test/vision/image_classifier_test.py @@ -61,7 +61,7 @@ def _generate_empty_results() -> ImageClassifierResult: timestamp_ms=0) -def _generate_burger_results() -> ImageClassifierResult: +def _generate_burger_results(timestamp_ms=0) -> ImageClassifierResult: return ImageClassifierResult( classifications=[ _Classifications( @@ -70,30 +70,36 @@ def _generate_burger_results() -> ImageClassifierResult: index=934, score=0.793959, display_name='', - category_name='cheeseburger'), + category_name='cheeseburger', + ), _Category( index=932, score=0.0273929, display_name='', - category_name='bagel'), + category_name='bagel', + ), _Category( index=925, score=0.0193408, display_name='', - category_name='guacamole'), + category_name='guacamole', + ), _Category( index=963, score=0.00632786, display_name='', - category_name='meat loaf') + category_name='meat loaf', + ), ], head_index=0, - head_name='probability') + head_name='probability', + ) ], - timestamp_ms=0) + timestamp_ms=timestamp_ms, + ) -def _generate_soccer_ball_results() -> ImageClassifierResult: +def _generate_soccer_ball_results(timestamp_ms=0) -> ImageClassifierResult: return ImageClassifierResult( classifications=[ _Classifications( @@ -102,12 +108,15 @@ def _generate_soccer_ball_results() -> ImageClassifierResult: index=806, score=0.996527, display_name='', - category_name='soccer ball') + category_name='soccer ball', + ) ], head_index=0, - head_name='probability') + head_name='probability', + ) ], - timestamp_ms=0) + timestamp_ms=timestamp_ms, + ) class ModelFileType(enum.Enum): @@ -379,8 +388,11 @@ class ImageClassifierTest(parameterized.TestCase): for timestamp in range(0, 300, 30): classification_result = classifier.classify_for_video( self.test_image, timestamp) - test_utils.assert_proto_equals(self, classification_result.to_pb2(), - _generate_burger_results().to_pb2()) + test_utils.assert_proto_equals( + self, + classification_result.to_pb2(), + _generate_burger_results(timestamp).to_pb2(), + ) def test_classify_for_video_succeeds_with_region_of_interest(self): options = _ImageClassifierOptions( @@ -398,8 +410,11 @@ class ImageClassifierTest(parameterized.TestCase): for timestamp in range(0, 300, 30): classification_result = classifier.classify_for_video( test_image, timestamp, image_processing_options) - test_utils.assert_proto_equals(self, classification_result.to_pb2(), - _generate_soccer_ball_results().to_pb2()) + test_utils.assert_proto_equals( + self, + classification_result.to_pb2(), + _generate_soccer_ball_results(timestamp).to_pb2(), + ) def test_calling_classify_in_live_stream_mode(self): options = _ImageClassifierOptions( @@ -455,8 +470,7 @@ class ImageClassifierTest(parameterized.TestCase): score_threshold=threshold, result_callback=check_result) with _ImageClassifier.create_from_options(options) as classifier: - for timestamp in range(0, 300, 30): - classifier.classify_async(self.test_image, timestamp) + classifier.classify_async(self.test_image, 0) def test_classify_async_succeeds_with_region_of_interest(self): # Load the test image. @@ -470,8 +484,9 @@ class ImageClassifierTest(parameterized.TestCase): def check_result(result: ImageClassifierResult, output_image: _Image, timestamp_ms: int): - test_utils.assert_proto_equals(self, result.to_pb2(), - _generate_soccer_ball_results().to_pb2()) + test_utils.assert_proto_equals( + self, result.to_pb2(), _generate_soccer_ball_results(100).to_pb2() + ) self.assertEqual(output_image.width, test_image.width) self.assertEqual(output_image.height, test_image.height) self.assertLess(observed_timestamp_ms, timestamp_ms) @@ -483,9 +498,7 @@ class ImageClassifierTest(parameterized.TestCase): max_results=1, result_callback=check_result) with _ImageClassifier.create_from_options(options) as classifier: - for timestamp in range(0, 300, 30): - classifier.classify_async(test_image, timestamp, - image_processing_options) + classifier.classify_async(test_image, 100, image_processing_options) if __name__ == '__main__':