From 803210a86b460f53fa6eb6c4ef4d67131e70d457 Mon Sep 17 00:00:00 2001 From: kinaryml Date: Fri, 14 Oct 2022 03:00:29 -0700 Subject: [PATCH] Simplified async test cases to invoke the classifier in context --- .../python/test/vision/image_classifier_test.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/mediapipe/tasks/python/test/vision/image_classifier_test.py b/mediapipe/tasks/python/test/vision/image_classifier_test.py index 783718d0..a90ddd53 100644 --- a/mediapipe/tasks/python/test/vision/image_classifier_test.py +++ b/mediapipe/tasks/python/test/vision/image_classifier_test.py @@ -471,10 +471,9 @@ class ImageClassifierTest(parameterized.TestCase): running_mode=_RUNNING_MODE.LIVE_STREAM, classifier_options=classifier_options, result_callback=check_result) - classifier = _ImageClassifier.create_from_options(options) - for timestamp in range(0, 300, 30): - classifier.classify_async(self.test_image, timestamp) - classifier.close() + with _ImageClassifier.create_from_options(options) as classifier: + for timestamp in range(0, 300, 30): + classifier.classify_async(self.test_image, timestamp) def test_classify_async_succeeds_with_region_of_interest(self): # Load the test image. @@ -499,10 +498,9 @@ class ImageClassifierTest(parameterized.TestCase): running_mode=_RUNNING_MODE.LIVE_STREAM, classifier_options=classifier_options, result_callback=check_result) - classifier = _ImageClassifier.create_from_options(options) - for timestamp in range(0, 300, 30): - classifier.classify_async(test_image, timestamp, roi) - classifier.close() + with _ImageClassifier.create_from_options(options) as classifier: + for timestamp in range(0, 300, 30): + classifier.classify_async(test_image, timestamp, roi) if __name__ == '__main__':