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__':