Initial upload: aesthetic-immich-starter
This commit is contained in:
Abderrahim-Mouhcine
@@ -0,0 +1,13 @@
|
||||
# Copyright 2022 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
+589
@@ -0,0 +1,589 @@
|
||||
# Copyright 2023 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tests for face detector."""
|
||||
|
||||
import enum
|
||||
import os
|
||||
import threading
|
||||
from unittest import mock
|
||||
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
|
||||
from google.protobuf import text_format
|
||||
from mediapipe.framework.formats import detection_pb2
|
||||
from mediapipe.tasks.python.components.containers import detections as detections_module
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.test import test_utils
|
||||
from mediapipe.tasks.python.test.vision import proto_utils
|
||||
from mediapipe.tasks.python.vision import face_detector
|
||||
from mediapipe.tasks.python.vision.core import image as image_module
|
||||
from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module
|
||||
from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module
|
||||
|
||||
FaceDetectorResult = detections_module.DetectionResult
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
_Image = image_module.Image
|
||||
_FaceDetector = face_detector.FaceDetector
|
||||
_FaceDetectorOptions = face_detector.FaceDetectorOptions
|
||||
_RUNNING_MODE = running_mode_module.VisionTaskRunningMode
|
||||
_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions
|
||||
|
||||
_SHORT_RANGE_BLAZE_FACE_MODEL = 'face_detection_short_range.tflite'
|
||||
_FULL_RANGE_BLAZE_FACE_MODEL = 'face_detection_full_range.tflite'
|
||||
_PORTRAIT_IMAGE = 'portrait.jpg'
|
||||
_PORTRAIT_EXPECTED_DETECTION = 'portrait_expected_detection.pbtxt'
|
||||
_PORTRAIT_EXPECTED_FULL_RANGE_DETECTION = (
|
||||
'portrait_expected_full_range_detection.pbtxt'
|
||||
)
|
||||
_PORTRAIT_ROTATED_IMAGE = 'portrait_rotated.jpg'
|
||||
_PORTRAIT_ROTATED_EXPECTED_DETECTION = (
|
||||
'portrait_rotated_expected_detection.pbtxt'
|
||||
)
|
||||
_CAT_IMAGE = 'cat.jpg'
|
||||
_KEYPOINT_ERROR_THRESHOLD = 1e-2
|
||||
_BOUNDING_BOX_ERROR_THRESHOLD = 5
|
||||
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision'
|
||||
|
||||
|
||||
def _get_expected_face_detector_result(file_name: str) -> FaceDetectorResult:
|
||||
face_detection_result_file_path = test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, file_name)
|
||||
)
|
||||
with open(face_detection_result_file_path, 'rb') as f:
|
||||
face_detection_proto = detection_pb2.Detection()
|
||||
text_format.Parse(f.read(), face_detection_proto)
|
||||
face_detection = proto_utils.create_detection_from_proto(face_detection_proto)
|
||||
return FaceDetectorResult(detections=[face_detection])
|
||||
|
||||
|
||||
class ModelFileType(enum.Enum):
|
||||
FILE_CONTENT = 1
|
||||
FILE_NAME = 2
|
||||
|
||||
|
||||
class FaceDetectorTest(parameterized.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _PORTRAIT_IMAGE)
|
||||
)
|
||||
)
|
||||
self.model_path = test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _SHORT_RANGE_BLAZE_FACE_MODEL)
|
||||
)
|
||||
self.full_range_model_path = test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _FULL_RANGE_BLAZE_FACE_MODEL)
|
||||
)
|
||||
|
||||
def test_create_from_file_succeeds_with_valid_model_path(self):
|
||||
# Creates with default option and valid model file successfully.
|
||||
with _FaceDetector.create_from_model_path(self.model_path) as detector:
|
||||
self.assertIsInstance(detector, _FaceDetector)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_path(self):
|
||||
# Creates with options containing model file successfully.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _FaceDetectorOptions(base_options=base_options)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
self.assertIsInstance(detector, _FaceDetector)
|
||||
|
||||
def test_create_from_options_fails_with_invalid_model_path(self):
|
||||
with self.assertRaisesRegex(
|
||||
FileNotFoundError,
|
||||
'Unable to open file at /path/to/invalid/model.tflite',
|
||||
):
|
||||
base_options = _BaseOptions(
|
||||
model_asset_path='/path/to/invalid/model.tflite'
|
||||
)
|
||||
options = _FaceDetectorOptions(base_options=base_options)
|
||||
detector = _FaceDetector.create_from_options(options)
|
||||
detector.close()
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_content(self):
|
||||
# Creates with options containing model content successfully.
|
||||
with open(self.model_path, 'rb') as f:
|
||||
base_options = _BaseOptions(model_asset_buffer=f.read())
|
||||
options = _FaceDetectorOptions(base_options=base_options)
|
||||
detector = _FaceDetector.create_from_options(options)
|
||||
self.assertIsInstance(detector, _FaceDetector)
|
||||
detector.close()
|
||||
|
||||
def _expect_keypoints_correct(self, actual_keypoints, expected_keypoints):
|
||||
self.assertLen(actual_keypoints, len(expected_keypoints))
|
||||
for i in range(len(actual_keypoints)):
|
||||
self.assertAlmostEqual(
|
||||
actual_keypoints[i].x,
|
||||
expected_keypoints[i].x,
|
||||
delta=_KEYPOINT_ERROR_THRESHOLD,
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
actual_keypoints[i].y,
|
||||
expected_keypoints[i].y,
|
||||
delta=_KEYPOINT_ERROR_THRESHOLD,
|
||||
)
|
||||
|
||||
def _expect_bounding_box_correct(self, actual_bbox, expected_bbox):
|
||||
self.assertAlmostEqual(
|
||||
actual_bbox.origin_x,
|
||||
expected_bbox.origin_x,
|
||||
delta=_BOUNDING_BOX_ERROR_THRESHOLD,
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
actual_bbox.origin_y,
|
||||
expected_bbox.origin_y,
|
||||
delta=_BOUNDING_BOX_ERROR_THRESHOLD,
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
actual_bbox.width,
|
||||
expected_bbox.width,
|
||||
delta=_BOUNDING_BOX_ERROR_THRESHOLD,
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
actual_bbox.height,
|
||||
expected_bbox.height,
|
||||
delta=_BOUNDING_BOX_ERROR_THRESHOLD,
|
||||
)
|
||||
|
||||
def _expect_face_detector_results_correct(
|
||||
self, actual_results, expected_results
|
||||
):
|
||||
self.assertLen(actual_results.detections, len(expected_results.detections))
|
||||
for i in range(len(actual_results.detections)):
|
||||
actual_bbox = actual_results.detections[i].bounding_box
|
||||
expected_bbox = expected_results.detections[i].bounding_box
|
||||
self._expect_bounding_box_correct(actual_bbox, expected_bbox)
|
||||
self.assertNotEmpty(actual_results.detections[i].keypoints)
|
||||
self._expect_keypoints_correct(
|
||||
actual_results.detections[i].keypoints,
|
||||
expected_results.detections[i].keypoints,
|
||||
)
|
||||
|
||||
@parameterized.parameters(
|
||||
(ModelFileType.FILE_NAME, _PORTRAIT_EXPECTED_DETECTION),
|
||||
(ModelFileType.FILE_CONTENT, _PORTRAIT_EXPECTED_DETECTION),
|
||||
)
|
||||
def test_detect(self, model_file_type, expected_detection_result_file):
|
||||
# Creates detector.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _FaceDetectorOptions(base_options=base_options)
|
||||
detector = _FaceDetector.create_from_options(options)
|
||||
|
||||
# Performs face detection on the input.
|
||||
detection_result = detector.detect(self.test_image)
|
||||
# Comparing results.
|
||||
expected_detection_result = _get_expected_face_detector_result(
|
||||
expected_detection_result_file
|
||||
)
|
||||
self._expect_face_detector_results_correct(
|
||||
detection_result, expected_detection_result
|
||||
)
|
||||
# Closes the detector explicitly when the detector is not used in
|
||||
# a context.
|
||||
detector.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(ModelFileType.FILE_NAME, _PORTRAIT_EXPECTED_DETECTION),
|
||||
(ModelFileType.FILE_CONTENT, _PORTRAIT_EXPECTED_DETECTION),
|
||||
)
|
||||
def test_detect_in_context(
|
||||
self, model_file_type, expected_detection_result_file
|
||||
):
|
||||
# Creates detector.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _FaceDetectorOptions(base_options=base_options)
|
||||
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
# Performs face detection on the input.
|
||||
detection_result = detector.detect(self.test_image)
|
||||
# Comparing results.
|
||||
expected_detection_result = _get_expected_face_detector_result(
|
||||
expected_detection_result_file
|
||||
)
|
||||
self._expect_face_detector_results_correct(
|
||||
detection_result, expected_detection_result
|
||||
)
|
||||
|
||||
def test_detect_succeeds_with_full_range_model(self):
|
||||
full_range_model_path = test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _FULL_RANGE_BLAZE_FACE_MODEL)
|
||||
)
|
||||
base_options = _BaseOptions(model_asset_path=full_range_model_path)
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=base_options,
|
||||
)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
# Performs face detection on the input.
|
||||
detection_result = detector.detect(self.test_image)
|
||||
|
||||
# Comparing results.
|
||||
expected_detection_result = _get_expected_face_detector_result(
|
||||
_PORTRAIT_EXPECTED_FULL_RANGE_DETECTION
|
||||
)
|
||||
self._expect_face_detector_results_correct(
|
||||
detection_result, expected_detection_result
|
||||
)
|
||||
detector.close()
|
||||
|
||||
def test_detect_succeeds_with_rotated_image(self):
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _FaceDetectorOptions(base_options=base_options)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
# Load the test image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _PORTRAIT_ROTATED_IMAGE)
|
||||
)
|
||||
)
|
||||
# Rotated input image.
|
||||
image_processing_options = _ImageProcessingOptions(rotation_degrees=-90)
|
||||
# Performs face detection on the input.
|
||||
detection_result = detector.detect(test_image, image_processing_options)
|
||||
# Comparing results.
|
||||
expected_detection_result = _get_expected_face_detector_result(
|
||||
_PORTRAIT_ROTATED_EXPECTED_DETECTION
|
||||
)
|
||||
self._expect_face_detector_results_correct(
|
||||
detection_result, expected_detection_result
|
||||
)
|
||||
|
||||
def test_empty_detection_outputs(self):
|
||||
# Load a test image with no faces.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, _CAT_IMAGE))
|
||||
)
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path)
|
||||
)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
# Performs face detection on the input.
|
||||
detection_result = detector.detect(test_image)
|
||||
self.assertEmpty(detection_result.detections)
|
||||
|
||||
def test_missing_result_callback(self):
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback must be provided'
|
||||
):
|
||||
with _FaceDetector.create_from_options(options) as unused_detector:
|
||||
pass
|
||||
|
||||
@parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO))
|
||||
def test_illegal_result_callback(self, running_mode):
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=running_mode,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback should not be provided'
|
||||
):
|
||||
with _FaceDetector.create_from_options(options) as unused_detector:
|
||||
pass
|
||||
|
||||
def test_calling_detect_for_video_in_image_mode(self):
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
with self.assertRaises(ValueError):
|
||||
detector.detect_for_video(self.test_image, 0)
|
||||
|
||||
def test_calling_detect_async_in_image_mode(self):
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
with self.assertRaises(ValueError):
|
||||
detector.detect_async(self.test_image, 0)
|
||||
|
||||
def test_calling_detect_in_video_mode(self):
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
with self.assertRaises(ValueError):
|
||||
detector.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_async_in_video_mode(self):
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
with self.assertRaises(ValueError):
|
||||
detector.detect_async(self.test_image, 0)
|
||||
|
||||
def test_detect_for_video_succeeds_with_full_range_model(self):
|
||||
base_options = _BaseOptions(model_asset_path=self.full_range_model_path)
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=base_options,
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
for timestamp in range(0, 300, 30):
|
||||
# Performs face detection on the input.
|
||||
detection_result = detector.detect_for_video(self.test_image, timestamp)
|
||||
# Comparing results.
|
||||
expected_detection_result = _get_expected_face_detector_result(
|
||||
_PORTRAIT_EXPECTED_FULL_RANGE_DETECTION
|
||||
)
|
||||
self._expect_face_detector_results_correct(
|
||||
detection_result, expected_detection_result
|
||||
)
|
||||
|
||||
def test_detect_for_video_with_out_of_order_timestamp(self):
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
unused_result = detector.detect_for_video(self.test_image, 1)
|
||||
with self.assertRaises(ValueError):
|
||||
detector.detect_for_video(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
_PORTRAIT_IMAGE,
|
||||
0,
|
||||
_get_expected_face_detector_result(_PORTRAIT_EXPECTED_DETECTION),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
_PORTRAIT_IMAGE,
|
||||
0,
|
||||
_get_expected_face_detector_result(_PORTRAIT_EXPECTED_DETECTION),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
_PORTRAIT_ROTATED_IMAGE,
|
||||
-90,
|
||||
_get_expected_face_detector_result(
|
||||
_PORTRAIT_ROTATED_EXPECTED_DETECTION
|
||||
),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
_PORTRAIT_ROTATED_IMAGE,
|
||||
-90,
|
||||
_get_expected_face_detector_result(
|
||||
_PORTRAIT_ROTATED_EXPECTED_DETECTION
|
||||
),
|
||||
),
|
||||
(ModelFileType.FILE_NAME, _CAT_IMAGE, 0, FaceDetectorResult([])),
|
||||
(ModelFileType.FILE_CONTENT, _CAT_IMAGE, 0, FaceDetectorResult([])),
|
||||
)
|
||||
def test_detect_for_video(
|
||||
self,
|
||||
model_file_type,
|
||||
test_image_file_name,
|
||||
rotation_degrees,
|
||||
expected_detection_result,
|
||||
):
|
||||
# Creates detector.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=base_options, running_mode=_RUNNING_MODE.VIDEO
|
||||
)
|
||||
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
for timestamp in range(0, 300, 30):
|
||||
# Load the test image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, test_image_file_name)
|
||||
)
|
||||
)
|
||||
# Set the image processing options.
|
||||
image_processing_options = _ImageProcessingOptions(
|
||||
rotation_degrees=rotation_degrees
|
||||
)
|
||||
# Performs face detection on the input.
|
||||
detection_result = detector.detect_for_video(
|
||||
test_image, timestamp, image_processing_options
|
||||
)
|
||||
# Comparing results.
|
||||
self._expect_face_detector_results_correct(
|
||||
detection_result, expected_detection_result
|
||||
)
|
||||
|
||||
def test_calling_detect_in_live_stream_mode(self):
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
with self.assertRaises(ValueError):
|
||||
detector.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_for_video_in_live_stream_mode(self):
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
with self.assertRaises(ValueError):
|
||||
detector.detect_for_video(self.test_image, 0)
|
||||
|
||||
def test_detect_async_calls_with_illegal_timestamp(self):
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
detector.detect_async(self.test_image, 100)
|
||||
with self.assertRaises(ValueError):
|
||||
detector.detect_async(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
_PORTRAIT_IMAGE,
|
||||
0,
|
||||
_get_expected_face_detector_result(_PORTRAIT_EXPECTED_DETECTION),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
_PORTRAIT_IMAGE,
|
||||
0,
|
||||
_get_expected_face_detector_result(_PORTRAIT_EXPECTED_DETECTION),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
_PORTRAIT_ROTATED_IMAGE,
|
||||
-90,
|
||||
_get_expected_face_detector_result(
|
||||
_PORTRAIT_ROTATED_EXPECTED_DETECTION
|
||||
),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
_PORTRAIT_ROTATED_IMAGE,
|
||||
-90,
|
||||
_get_expected_face_detector_result(
|
||||
_PORTRAIT_ROTATED_EXPECTED_DETECTION
|
||||
),
|
||||
),
|
||||
(ModelFileType.FILE_NAME, _CAT_IMAGE, 0, FaceDetectorResult([])),
|
||||
(ModelFileType.FILE_CONTENT, _CAT_IMAGE, 0, FaceDetectorResult([])),
|
||||
)
|
||||
def test_detect_async_calls(
|
||||
self,
|
||||
model_file_type,
|
||||
test_image_file_name,
|
||||
rotation_degrees,
|
||||
expected_detection_result,
|
||||
):
|
||||
# Creates detector.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
callback_event = threading.Event()
|
||||
callback_exception: None | Exception = None
|
||||
observed_timestamp_ms = -1
|
||||
|
||||
def check_result(
|
||||
result: FaceDetectorResult,
|
||||
output_image: _Image,
|
||||
timestamp_ms: int,
|
||||
):
|
||||
nonlocal callback_event, callback_exception, observed_timestamp_ms
|
||||
|
||||
try:
|
||||
self._expect_face_detector_results_correct(
|
||||
result, expected_detection_result
|
||||
)
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
self.assertEqual(output_image.width, test_image.width)
|
||||
self.assertEqual(output_image.height, test_image.height)
|
||||
observed_timestamp_ms = timestamp_ms
|
||||
except AssertionError as e:
|
||||
callback_exception = e
|
||||
finally:
|
||||
callback_event.set()
|
||||
|
||||
options = _FaceDetectorOptions(
|
||||
base_options=base_options,
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=check_result,
|
||||
)
|
||||
|
||||
# Load the test image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, test_image_file_name)
|
||||
)
|
||||
)
|
||||
|
||||
with _FaceDetector.create_from_options(options) as detector:
|
||||
for timestamp in range(0, 300, 30):
|
||||
# Set the image processing options.
|
||||
image_processing_options = _ImageProcessingOptions(
|
||||
rotation_degrees=rotation_degrees
|
||||
)
|
||||
detector.detect_async(test_image, timestamp, image_processing_options)
|
||||
callback_event.wait(timeout=3.0)
|
||||
if callback_exception is not None:
|
||||
raise callback_exception
|
||||
callback_event.clear()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
+568
@@ -0,0 +1,568 @@
|
||||
# Copyright 2023 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tests for face landmarker."""
|
||||
|
||||
import enum
|
||||
import threading
|
||||
from unittest import mock
|
||||
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import numpy as np
|
||||
|
||||
from google.protobuf import text_format
|
||||
from mediapipe.framework.formats import classification_pb2
|
||||
from mediapipe.framework.formats import landmark_pb2
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.test import test_utils
|
||||
from mediapipe.tasks.python.test.vision import proto_utils
|
||||
from mediapipe.tasks.python.vision import face_landmarker
|
||||
from mediapipe.tasks.python.vision.core import image as image_module
|
||||
from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module
|
||||
from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module
|
||||
|
||||
|
||||
FaceLandmarkerResult = face_landmarker.FaceLandmarkerResult
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
_Image = image_module.Image
|
||||
_FaceLandmarker = face_landmarker.FaceLandmarker
|
||||
_FaceLandmarkerOptions = face_landmarker.FaceLandmarkerOptions
|
||||
_RUNNING_MODE = running_mode_module.VisionTaskRunningMode
|
||||
_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions
|
||||
|
||||
_FACE_LANDMARKER_BUNDLE_ASSET_FILE = 'face_landmarker_v2.task'
|
||||
_PORTRAIT_IMAGE = 'portrait.jpg'
|
||||
_CAT_IMAGE = 'cat.jpg'
|
||||
_PORTRAIT_EXPECTED_FACE_LANDMARKS = 'portrait_expected_face_landmarks.pbtxt'
|
||||
_PORTRAIT_EXPECTED_BLENDSHAPES = 'portrait_expected_blendshapes.pbtxt'
|
||||
_LANDMARKS_MARGIN = 0.03
|
||||
_BLENDSHAPES_MARGIN = 0.13
|
||||
_FACIAL_TRANSFORMATION_MATRIX_MARGIN = 0.02
|
||||
|
||||
|
||||
def _get_expected_face_landmarks(file_path: str):
|
||||
proto_file_path = test_utils.get_test_data_path(file_path)
|
||||
face_landmarks_results = []
|
||||
with open(proto_file_path, 'rb') as f:
|
||||
proto = landmark_pb2.NormalizedLandmarkList()
|
||||
text_format.Parse(f.read(), proto)
|
||||
face_landmarks = proto_utils.create_normalized_landmark_list_from_proto(
|
||||
proto
|
||||
)
|
||||
face_landmarks_results.append(face_landmarks)
|
||||
return face_landmarks_results
|
||||
|
||||
|
||||
def _get_expected_face_blendshapes(file_path: str):
|
||||
proto_file_path = test_utils.get_test_data_path(file_path)
|
||||
face_blendshapes_results = []
|
||||
with open(proto_file_path, 'rb') as f:
|
||||
proto = classification_pb2.ClassificationList()
|
||||
text_format.Parse(f.read(), proto)
|
||||
face_blendshapes_categories = (
|
||||
proto_utils.create_classification_list_from_proto(proto)
|
||||
)
|
||||
face_blendshapes_results.append(face_blendshapes_categories)
|
||||
return face_blendshapes_results
|
||||
|
||||
|
||||
def _get_expected_facial_transformation_matrixes():
|
||||
matrix = np.array([
|
||||
[0.9995292, -0.01294756, 0.038823195, -0.3691378],
|
||||
[0.0072318087, 0.9937692, -0.1101321, 22.75809],
|
||||
[-0.03715533, 0.11070588, 0.99315894, -65.765925],
|
||||
[0, 0, 0, 1],
|
||||
])
|
||||
facial_transformation_matrixes_results = []
|
||||
facial_transformation_matrixes_results.append(matrix)
|
||||
return facial_transformation_matrixes_results
|
||||
|
||||
|
||||
class ModelFileType(enum.Enum):
|
||||
FILE_CONTENT = 1
|
||||
FILE_NAME = 2
|
||||
|
||||
|
||||
class FaceLandmarkerTest(parameterized.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(_PORTRAIT_IMAGE)
|
||||
)
|
||||
self.model_path = test_utils.get_test_data_path(
|
||||
_FACE_LANDMARKER_BUNDLE_ASSET_FILE
|
||||
)
|
||||
|
||||
def _expect_landmarks_correct(self, actual_landmarks, expected_landmarks):
|
||||
# Expects to have the same number of faces detected.
|
||||
self.assertLen(actual_landmarks, len(expected_landmarks))
|
||||
|
||||
for i, _ in enumerate(actual_landmarks):
|
||||
for j, elem in enumerate(actual_landmarks[i]):
|
||||
self.assertAlmostEqual(
|
||||
elem.x, expected_landmarks[i][j].x, delta=_LANDMARKS_MARGIN
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
elem.y, expected_landmarks[i][j].y, delta=_LANDMARKS_MARGIN
|
||||
)
|
||||
|
||||
def _expect_blendshapes_correct(
|
||||
self, actual_blendshapes, expected_blendshapes
|
||||
):
|
||||
# Expects to have the same number of blendshapes.
|
||||
self.assertLen(actual_blendshapes, len(expected_blendshapes))
|
||||
|
||||
for i, _ in enumerate(actual_blendshapes):
|
||||
for j, elem in enumerate(actual_blendshapes[i]):
|
||||
self.assertEqual(elem.index, expected_blendshapes[i][j].index)
|
||||
self.assertAlmostEqual(
|
||||
elem.score,
|
||||
expected_blendshapes[i][j].score,
|
||||
delta=_BLENDSHAPES_MARGIN,
|
||||
)
|
||||
|
||||
def _expect_facial_transformation_matrixes_correct(
|
||||
self, actual_matrix_list, expected_matrix_list
|
||||
):
|
||||
self.assertLen(actual_matrix_list, len(expected_matrix_list))
|
||||
|
||||
for i, elem in enumerate(actual_matrix_list):
|
||||
self.assertEqual(elem.shape[0], expected_matrix_list[i].shape[0])
|
||||
self.assertEqual(elem.shape[1], expected_matrix_list[i].shape[1])
|
||||
self.assertSequenceAlmostEqual(
|
||||
elem.flatten(),
|
||||
expected_matrix_list[i].flatten(),
|
||||
delta=_FACIAL_TRANSFORMATION_MATRIX_MARGIN,
|
||||
)
|
||||
|
||||
def test_create_from_file_succeeds_with_valid_model_path(self):
|
||||
# Creates with default option and valid model file successfully.
|
||||
with _FaceLandmarker.create_from_model_path(self.model_path) as landmarker:
|
||||
self.assertIsInstance(landmarker, _FaceLandmarker)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_path(self):
|
||||
# Creates with options containing model file successfully.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _FaceLandmarkerOptions(base_options=base_options)
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
self.assertIsInstance(landmarker, _FaceLandmarker)
|
||||
|
||||
def test_create_from_options_fails_with_invalid_model_path(self):
|
||||
# Invalid empty model path.
|
||||
with self.assertRaisesRegex(
|
||||
FileNotFoundError,
|
||||
'Unable to open file at /path/to/invalid/model.tflite',
|
||||
):
|
||||
base_options = _BaseOptions(
|
||||
model_asset_path='/path/to/invalid/model.tflite'
|
||||
)
|
||||
options = _FaceLandmarkerOptions(base_options=base_options)
|
||||
landmarker = _FaceLandmarker.create_from_options(options)
|
||||
landmarker.close()
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_content(self):
|
||||
# Creates with options containing model content successfully.
|
||||
with open(self.model_path, 'rb') as f:
|
||||
base_options = _BaseOptions(model_asset_buffer=f.read())
|
||||
options = _FaceLandmarkerOptions(base_options=base_options)
|
||||
landmarker = _FaceLandmarker.create_from_options(options)
|
||||
self.assertIsInstance(landmarker, _FaceLandmarker)
|
||||
landmarker.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
_FACE_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
_get_expected_face_landmarks(_PORTRAIT_EXPECTED_FACE_LANDMARKS),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
_FACE_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
_get_expected_face_landmarks(_PORTRAIT_EXPECTED_FACE_LANDMARKS),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
)
|
||||
def test_detect(
|
||||
self,
|
||||
model_file_type,
|
||||
model_name,
|
||||
expected_face_landmarks,
|
||||
expected_face_blendshapes,
|
||||
expected_facial_transformation_matrixes,
|
||||
):
|
||||
# Creates face landmarker.
|
||||
model_path = test_utils.get_test_data_path(model_name)
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=base_options,
|
||||
output_face_blendshapes=True if expected_face_blendshapes else False,
|
||||
output_facial_transformation_matrixes=True
|
||||
if expected_facial_transformation_matrixes
|
||||
else False,
|
||||
)
|
||||
landmarker = _FaceLandmarker.create_from_options(options)
|
||||
|
||||
# Performs face landmarks detection on the input.
|
||||
detection_result = landmarker.detect(self.test_image)
|
||||
# Comparing results.
|
||||
if expected_face_landmarks is not None:
|
||||
self._expect_landmarks_correct(
|
||||
detection_result.face_landmarks, expected_face_landmarks
|
||||
)
|
||||
if expected_face_blendshapes is not None:
|
||||
self._expect_blendshapes_correct(
|
||||
detection_result.face_blendshapes, expected_face_blendshapes
|
||||
)
|
||||
if expected_facial_transformation_matrixes is not None:
|
||||
self._expect_facial_transformation_matrixes_correct(
|
||||
detection_result.facial_transformation_matrixes,
|
||||
expected_facial_transformation_matrixes,
|
||||
)
|
||||
|
||||
# Closes the face landmarker explicitly when the face landmarker is not used
|
||||
# in a context.
|
||||
landmarker.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
_FACE_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
_get_expected_face_landmarks(_PORTRAIT_EXPECTED_FACE_LANDMARKS),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
_FACE_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
_get_expected_face_landmarks(_PORTRAIT_EXPECTED_FACE_LANDMARKS),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
)
|
||||
def test_detect_in_context(
|
||||
self,
|
||||
model_file_type,
|
||||
model_name,
|
||||
expected_face_landmarks,
|
||||
expected_face_blendshapes,
|
||||
expected_facial_transformation_matrixes,
|
||||
):
|
||||
# Creates face landmarker.
|
||||
model_path = test_utils.get_test_data_path(model_name)
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=base_options,
|
||||
output_face_blendshapes=True if expected_face_blendshapes else False,
|
||||
output_facial_transformation_matrixes=True
|
||||
if expected_facial_transformation_matrixes
|
||||
else False,
|
||||
)
|
||||
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
# Performs face landmarks detection on the input.
|
||||
detection_result = landmarker.detect(self.test_image)
|
||||
# Comparing results.
|
||||
if expected_face_landmarks is not None:
|
||||
self._expect_landmarks_correct(
|
||||
detection_result.face_landmarks, expected_face_landmarks
|
||||
)
|
||||
if expected_face_blendshapes is not None:
|
||||
self._expect_blendshapes_correct(
|
||||
detection_result.face_blendshapes, expected_face_blendshapes
|
||||
)
|
||||
if expected_facial_transformation_matrixes is not None:
|
||||
self._expect_facial_transformation_matrixes_correct(
|
||||
detection_result.facial_transformation_matrixes,
|
||||
expected_facial_transformation_matrixes,
|
||||
)
|
||||
|
||||
def test_empty_detection_outputs(self):
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path)
|
||||
)
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
# Load the image with no faces.
|
||||
no_faces_test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(_CAT_IMAGE)
|
||||
)
|
||||
# Performs face landmarks detection on the input.
|
||||
detection_result = landmarker.detect(no_faces_test_image)
|
||||
self.assertEmpty(detection_result.face_landmarks)
|
||||
self.assertEmpty(detection_result.face_blendshapes)
|
||||
self.assertEmpty(detection_result.facial_transformation_matrixes)
|
||||
|
||||
def test_missing_result_callback(self):
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback must be provided'
|
||||
):
|
||||
with _FaceLandmarker.create_from_options(options) as unused_landmarker:
|
||||
pass
|
||||
|
||||
@parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO))
|
||||
def test_illegal_result_callback(self, running_mode):
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=running_mode,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback should not be provided'
|
||||
):
|
||||
with _FaceLandmarker.create_from_options(options) as unused_landmarker:
|
||||
pass
|
||||
|
||||
def test_calling_detect_for_video_in_image_mode(self):
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the video mode'
|
||||
):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
# TODO: Change the errors to ValueError once we return MpStatus.
|
||||
def test_calling_detect_async_in_image_mode(self):
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the live stream mode'
|
||||
):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
def test_calling_detect_in_video_mode(self):
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the image mode'
|
||||
):
|
||||
landmarker.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_async_in_video_mode(self):
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the live stream mode'
|
||||
):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
def test_detect_for_video_with_out_of_order_timestamp(self):
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
unused_result = landmarker.detect_for_video(self.test_image, 1)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'Input timestamp must be monotonically increasing'
|
||||
):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
_FACE_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
_get_expected_face_landmarks(_PORTRAIT_EXPECTED_FACE_LANDMARKS),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
)
|
||||
def test_detect_for_video(
|
||||
self,
|
||||
model_name,
|
||||
expected_face_landmarks,
|
||||
expected_face_blendshapes,
|
||||
expected_facial_transformation_matrixes,
|
||||
):
|
||||
# Creates face landmarker.
|
||||
model_path = test_utils.get_test_data_path(model_name)
|
||||
base_options = _BaseOptions(model_asset_path=model_path)
|
||||
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=base_options,
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
output_face_blendshapes=True if expected_face_blendshapes else False,
|
||||
output_facial_transformation_matrixes=True
|
||||
if expected_facial_transformation_matrixes
|
||||
else False,
|
||||
)
|
||||
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
for timestamp in range(0, 300, 30):
|
||||
# Performs face landmarks detection on the input.
|
||||
detection_result = landmarker.detect_for_video(
|
||||
self.test_image, timestamp
|
||||
)
|
||||
# Comparing results.
|
||||
if expected_face_landmarks is not None:
|
||||
self._expect_landmarks_correct(
|
||||
detection_result.face_landmarks, expected_face_landmarks
|
||||
)
|
||||
if expected_face_blendshapes is not None:
|
||||
self._expect_blendshapes_correct(
|
||||
detection_result.face_blendshapes, expected_face_blendshapes
|
||||
)
|
||||
if expected_facial_transformation_matrixes is not None:
|
||||
self._expect_facial_transformation_matrixes_correct(
|
||||
detection_result.facial_transformation_matrixes,
|
||||
expected_facial_transformation_matrixes,
|
||||
)
|
||||
|
||||
def test_calling_detect_in_live_stream_mode(self):
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the image mode'
|
||||
):
|
||||
landmarker.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_for_video_in_live_stream_mode(self):
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the video mode'
|
||||
):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
def test_detect_async_calls_with_illegal_timestamp(self):
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
landmarker.detect_async(self.test_image, 100)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'Input timestamp must be monotonically increasing'
|
||||
):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
_PORTRAIT_IMAGE,
|
||||
_FACE_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
_get_expected_face_landmarks(_PORTRAIT_EXPECTED_FACE_LANDMARKS),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
)
|
||||
def test_detect_async_calls(
|
||||
self,
|
||||
image_path,
|
||||
model_name,
|
||||
expected_face_landmarks,
|
||||
expected_face_blendshapes,
|
||||
expected_facial_transformation_matrixes,
|
||||
):
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(image_path)
|
||||
)
|
||||
callback_event = threading.Event()
|
||||
callback_exception: None | Exception = None
|
||||
observed_timestamp_ms = -1
|
||||
|
||||
def check_result(
|
||||
result: FaceLandmarkerResult, output_image: _Image, timestamp_ms: int
|
||||
):
|
||||
nonlocal callback_event, callback_exception, observed_timestamp_ms
|
||||
|
||||
try:
|
||||
# Comparing results.
|
||||
if expected_face_landmarks is not None:
|
||||
self._expect_landmarks_correct(
|
||||
result.face_landmarks, expected_face_landmarks
|
||||
)
|
||||
if expected_face_blendshapes is not None:
|
||||
self._expect_blendshapes_correct(
|
||||
result.face_blendshapes, expected_face_blendshapes
|
||||
)
|
||||
if expected_facial_transformation_matrixes is not None:
|
||||
self._expect_facial_transformation_matrixes_correct(
|
||||
result.facial_transformation_matrixes,
|
||||
expected_facial_transformation_matrixes,
|
||||
)
|
||||
self.assertTrue(
|
||||
np.array_equal(output_image.numpy_view(), test_image.numpy_view())
|
||||
)
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
observed_timestamp_ms = timestamp_ms
|
||||
except AssertionError as e:
|
||||
callback_exception = e
|
||||
finally:
|
||||
callback_event.set()
|
||||
|
||||
model_path = test_utils.get_test_data_path(model_name)
|
||||
options = _FaceLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
output_face_blendshapes=True if expected_face_blendshapes else False,
|
||||
output_facial_transformation_matrixes=True
|
||||
if expected_facial_transformation_matrixes
|
||||
else False,
|
||||
result_callback=check_result,
|
||||
)
|
||||
with _FaceLandmarker.create_from_options(options) as landmarker:
|
||||
for timestamp in range(0, 300, 30):
|
||||
landmarker.detect_async(test_image, timestamp)
|
||||
callback_event.wait(timeout=3.0)
|
||||
if callback_exception is not None:
|
||||
raise callback_exception
|
||||
callback_event.clear()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
+455
@@ -0,0 +1,455 @@
|
||||
# Copyright 2022 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tests for hand landmarker."""
|
||||
|
||||
import enum
|
||||
import threading
|
||||
from unittest import mock
|
||||
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import numpy as np
|
||||
|
||||
from google.protobuf import text_format
|
||||
from mediapipe.tasks.cc.components.containers.proto import landmarks_detection_result_pb2
|
||||
from mediapipe.tasks.python.components.containers import landmark as landmark_module
|
||||
from mediapipe.tasks.python.components.containers import landmark_detection_result as landmark_detection_result_module
|
||||
from mediapipe.tasks.python.components.containers import rect as rect_module
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.test import test_utils
|
||||
from mediapipe.tasks.python.test.vision import proto_utils
|
||||
from mediapipe.tasks.python.vision import hand_landmarker
|
||||
from mediapipe.tasks.python.vision.core import image as image_module
|
||||
from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module
|
||||
from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module
|
||||
|
||||
_LandmarksDetectionResultProto = (
|
||||
landmarks_detection_result_pb2.LandmarksDetectionResult)
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
_RectF = rect_module.RectF
|
||||
_Landmark = landmark_module.Landmark
|
||||
_NormalizedLandmark = landmark_module.NormalizedLandmark
|
||||
_LandmarksDetectionResult = (
|
||||
landmark_detection_result_module.LandmarksDetectionResult)
|
||||
_Image = image_module.Image
|
||||
_HandLandmarker = hand_landmarker.HandLandmarker
|
||||
_HandLandmarkerOptions = hand_landmarker.HandLandmarkerOptions
|
||||
_HandLandmarkerResult = hand_landmarker.HandLandmarkerResult
|
||||
_RUNNING_MODE = running_mode_module.VisionTaskRunningMode
|
||||
_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions
|
||||
|
||||
_HAND_LANDMARKER_BUNDLE_ASSET_FILE = 'hand_landmarker.task'
|
||||
_NO_HANDS_IMAGE = 'cats_and_dogs.jpg'
|
||||
_TWO_HANDS_IMAGE = 'right_hands.jpg'
|
||||
_THUMB_UP_IMAGE = 'thumb_up.jpg'
|
||||
_THUMB_UP_LANDMARKS = 'thumb_up_landmarks.pbtxt'
|
||||
_POINTING_UP_IMAGE = 'pointing_up.jpg'
|
||||
_POINTING_UP_LANDMARKS = 'pointing_up_landmarks.pbtxt'
|
||||
_POINTING_UP_ROTATED_IMAGE = 'pointing_up_rotated.jpg'
|
||||
_POINTING_UP_ROTATED_LANDMARKS = 'pointing_up_rotated_landmarks.pbtxt'
|
||||
_LANDMARKS_MARGIN = 0.03
|
||||
_HANDEDNESS_MARGIN = 0.05
|
||||
|
||||
|
||||
def _get_expected_hand_landmarker_result(
|
||||
file_path: str) -> _HandLandmarkerResult:
|
||||
landmarks_detection_result_file_path = test_utils.get_test_data_path(
|
||||
file_path)
|
||||
with open(landmarks_detection_result_file_path, 'rb') as f:
|
||||
landmarks_detection_result_proto = _LandmarksDetectionResultProto()
|
||||
# Use this if a .pb file is available.
|
||||
# landmarks_detection_result_proto.ParseFromString(f.read())
|
||||
text_format.Parse(f.read(), landmarks_detection_result_proto)
|
||||
landmarks_detection_result = (
|
||||
proto_utils.create_landmarks_detection_result_from_proto(
|
||||
landmarks_detection_result_proto
|
||||
)
|
||||
)
|
||||
return _HandLandmarkerResult(
|
||||
handedness=[landmarks_detection_result.categories],
|
||||
hand_landmarks=[landmarks_detection_result.landmarks],
|
||||
hand_world_landmarks=[landmarks_detection_result.world_landmarks])
|
||||
|
||||
|
||||
class ModelFileType(enum.Enum):
|
||||
FILE_CONTENT = 1
|
||||
FILE_NAME = 2
|
||||
|
||||
|
||||
class HandLandmarkerTest(parameterized.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(_THUMB_UP_IMAGE))
|
||||
self.model_path = test_utils.get_test_data_path(
|
||||
_HAND_LANDMARKER_BUNDLE_ASSET_FILE)
|
||||
|
||||
def _expect_hand_landmarks_correct(
|
||||
self, actual_landmarks, expected_landmarks, margin
|
||||
):
|
||||
# Expects to have the same number of hands detected.
|
||||
self.assertLen(actual_landmarks, len(expected_landmarks))
|
||||
|
||||
for i, _ in enumerate(actual_landmarks):
|
||||
for j, elem in enumerate(actual_landmarks[i]):
|
||||
self.assertAlmostEqual(elem.x, expected_landmarks[i][j].x, delta=margin)
|
||||
self.assertAlmostEqual(elem.y, expected_landmarks[i][j].y, delta=margin)
|
||||
|
||||
def _expect_handedness_correct(
|
||||
self, actual_handedness, expected_handedness, margin
|
||||
):
|
||||
# Actual top handedness matches expected top handedness.
|
||||
actual_top_handedness = actual_handedness[0][0]
|
||||
expected_top_handedness = expected_handedness[0][0]
|
||||
self.assertEqual(actual_top_handedness.index, expected_top_handedness.index)
|
||||
self.assertEqual(actual_top_handedness.category_name,
|
||||
expected_top_handedness.category_name)
|
||||
self.assertAlmostEqual(
|
||||
actual_top_handedness.score, expected_top_handedness.score, delta=margin
|
||||
)
|
||||
|
||||
def _expect_hand_landmarker_results_correct(
|
||||
self,
|
||||
actual_result: _HandLandmarkerResult,
|
||||
expected_result: _HandLandmarkerResult,
|
||||
):
|
||||
self._expect_hand_landmarks_correct(
|
||||
actual_result.hand_landmarks,
|
||||
expected_result.hand_landmarks,
|
||||
_LANDMARKS_MARGIN,
|
||||
)
|
||||
self._expect_handedness_correct(
|
||||
actual_result.handedness, expected_result.handedness, _HANDEDNESS_MARGIN
|
||||
)
|
||||
|
||||
def test_create_from_file_succeeds_with_valid_model_path(self):
|
||||
# Creates with default option and valid model file successfully.
|
||||
with _HandLandmarker.create_from_model_path(self.model_path) as landmarker:
|
||||
self.assertIsInstance(landmarker, _HandLandmarker)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_path(self):
|
||||
# Creates with options containing model file successfully.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _HandLandmarkerOptions(base_options=base_options)
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
self.assertIsInstance(landmarker, _HandLandmarker)
|
||||
|
||||
def test_create_from_options_fails_with_invalid_model_path(self):
|
||||
# Invalid empty model path.
|
||||
with self.assertRaisesRegex(
|
||||
FileNotFoundError,
|
||||
'Unable to open file at /path/to/invalid/model.tflite',
|
||||
):
|
||||
base_options = _BaseOptions(
|
||||
model_asset_path='/path/to/invalid/model.tflite')
|
||||
options = _HandLandmarkerOptions(base_options=base_options)
|
||||
landmarker = _HandLandmarker.create_from_options(options)
|
||||
landmarker.close()
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_content(self):
|
||||
# Creates with options containing model content successfully.
|
||||
with open(self.model_path, 'rb') as f:
|
||||
base_options = _BaseOptions(model_asset_buffer=f.read())
|
||||
options = _HandLandmarkerOptions(base_options=base_options)
|
||||
landmarker = _HandLandmarker.create_from_options(options)
|
||||
self.assertIsInstance(landmarker, _HandLandmarker)
|
||||
landmarker.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(ModelFileType.FILE_NAME,
|
||||
_get_expected_hand_landmarker_result(_THUMB_UP_LANDMARKS)),
|
||||
(ModelFileType.FILE_CONTENT,
|
||||
_get_expected_hand_landmarker_result(_THUMB_UP_LANDMARKS)))
|
||||
def test_detect(self, model_file_type, expected_detection_result):
|
||||
# Creates hand landmarker.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _HandLandmarkerOptions(base_options=base_options)
|
||||
landmarker = _HandLandmarker.create_from_options(options)
|
||||
|
||||
# Performs hand landmarks detection on the input.
|
||||
detection_result = landmarker.detect(self.test_image)
|
||||
# Comparing results.
|
||||
self._expect_hand_landmarker_results_correct(
|
||||
detection_result, expected_detection_result
|
||||
)
|
||||
# Closes the hand landmarker explicitly when the hand landmarker is not used
|
||||
# in a context.
|
||||
landmarker.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(ModelFileType.FILE_NAME,
|
||||
_get_expected_hand_landmarker_result(_THUMB_UP_LANDMARKS)),
|
||||
(ModelFileType.FILE_CONTENT,
|
||||
_get_expected_hand_landmarker_result(_THUMB_UP_LANDMARKS)))
|
||||
def test_detect_in_context(self, model_file_type, expected_detection_result):
|
||||
# Creates hand landmarker.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _HandLandmarkerOptions(base_options=base_options)
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
# Performs hand landmarks detection on the input.
|
||||
detection_result = landmarker.detect(self.test_image)
|
||||
# Comparing results.
|
||||
self._expect_hand_landmarker_results_correct(
|
||||
detection_result, expected_detection_result
|
||||
)
|
||||
|
||||
def test_detect_succeeds_with_num_hands(self):
|
||||
# Creates hand landmarker.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _HandLandmarkerOptions(base_options=base_options, num_hands=2)
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
# Load the two hands image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(_TWO_HANDS_IMAGE))
|
||||
# Performs hand landmarks detection on the input.
|
||||
detection_result = landmarker.detect(test_image)
|
||||
# Comparing results.
|
||||
self.assertLen(detection_result.handedness, 2)
|
||||
|
||||
def test_detect_succeeds_with_rotation(self):
|
||||
# Creates hand landmarker.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _HandLandmarkerOptions(base_options=base_options)
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
# Load the pointing up rotated image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(_POINTING_UP_ROTATED_IMAGE))
|
||||
# Set rotation parameters using ImageProcessingOptions.
|
||||
image_processing_options = _ImageProcessingOptions(rotation_degrees=-90)
|
||||
# Performs hand landmarks detection on the input.
|
||||
detection_result = landmarker.detect(test_image, image_processing_options)
|
||||
expected_detection_result = _get_expected_hand_landmarker_result(
|
||||
_POINTING_UP_ROTATED_LANDMARKS)
|
||||
# Comparing results.
|
||||
self._expect_hand_landmarker_results_correct(
|
||||
detection_result, expected_detection_result
|
||||
)
|
||||
|
||||
def test_detect_fails_with_region_of_interest(self):
|
||||
# Creates hand landmarker.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _HandLandmarkerOptions(base_options=base_options)
|
||||
with self.assertRaises(ValueError):
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
# Set the `region_of_interest` parameter using `ImageProcessingOptions`.
|
||||
image_processing_options = _ImageProcessingOptions(
|
||||
region_of_interest=_RectF(0.0, 0.0, 1.0, 1.0)
|
||||
)
|
||||
# Attempt to perform hand landmarks detection on the cropped input.
|
||||
landmarker.detect(self.test_image, image_processing_options)
|
||||
|
||||
def test_empty_detection_outputs(self):
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path))
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
# Load the image with no hands.
|
||||
no_hands_test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(_NO_HANDS_IMAGE))
|
||||
# Performs hand landmarks detection on the input.
|
||||
detection_result = landmarker.detect(no_hands_test_image)
|
||||
self.assertEmpty(detection_result.hand_landmarks)
|
||||
self.assertEmpty(detection_result.hand_world_landmarks)
|
||||
self.assertEmpty(detection_result.handedness)
|
||||
|
||||
def test_missing_result_callback(self):
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM)
|
||||
with self.assertRaisesRegex(ValueError,
|
||||
r'result callback must be provided'):
|
||||
with _HandLandmarker.create_from_options(options) as unused_landmarker:
|
||||
pass
|
||||
|
||||
@parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO))
|
||||
def test_illegal_result_callback(self, running_mode):
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=running_mode,
|
||||
result_callback=mock.MagicMock())
|
||||
with self.assertRaisesRegex(ValueError,
|
||||
r'result callback should not be provided'):
|
||||
with _HandLandmarker.create_from_options(options) as unused_landmarker:
|
||||
pass
|
||||
|
||||
def test_calling_detect_for_video_in_image_mode(self):
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE)
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
def test_calling_detect_async_in_image_mode(self):
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE)
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
def test_calling_detect_in_video_mode(self):
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO)
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_async_in_video_mode(self):
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO)
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
def test_detect_for_video_with_out_of_order_timestamp(self):
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO)
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
unused_result = landmarker.detect_for_video(self.test_image, 1)
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(_THUMB_UP_IMAGE, 0,
|
||||
_get_expected_hand_landmarker_result(_THUMB_UP_LANDMARKS)),
|
||||
(_POINTING_UP_IMAGE, 0,
|
||||
_get_expected_hand_landmarker_result(_POINTING_UP_LANDMARKS)),
|
||||
(_POINTING_UP_ROTATED_IMAGE, -90,
|
||||
_get_expected_hand_landmarker_result(_POINTING_UP_ROTATED_LANDMARKS)),
|
||||
(_NO_HANDS_IMAGE, 0, _HandLandmarkerResult([], [], [])))
|
||||
def test_detect_for_video(self, image_path, rotation, expected_result):
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(image_path))
|
||||
# Set rotation parameters using ImageProcessingOptions.
|
||||
image_processing_options = _ImageProcessingOptions(
|
||||
rotation_degrees=rotation)
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO)
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
for timestamp in range(0, 300, 30):
|
||||
result = landmarker.detect_for_video(test_image, timestamp,
|
||||
image_processing_options)
|
||||
if (result.hand_landmarks and result.hand_world_landmarks and
|
||||
result.handedness):
|
||||
self._expect_hand_landmarker_results_correct(result, expected_result)
|
||||
else:
|
||||
self.assertEqual(result, expected_result)
|
||||
|
||||
def test_calling_detect_in_live_stream_mode(self):
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock())
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_for_video_in_live_stream_mode(self):
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock())
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
def test_detect_async_calls_with_illegal_timestamp(self):
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock())
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
landmarker.detect_async(self.test_image, 100)
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(_THUMB_UP_IMAGE, 0,
|
||||
_get_expected_hand_landmarker_result(_THUMB_UP_LANDMARKS)),
|
||||
(_POINTING_UP_IMAGE, 0,
|
||||
_get_expected_hand_landmarker_result(_POINTING_UP_LANDMARKS)),
|
||||
(_POINTING_UP_ROTATED_IMAGE, -90,
|
||||
_get_expected_hand_landmarker_result(_POINTING_UP_ROTATED_LANDMARKS)),
|
||||
(_NO_HANDS_IMAGE, 0, _HandLandmarkerResult([], [], [])))
|
||||
def test_detect_async_calls(self, image_path, rotation, expected_result):
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(image_path))
|
||||
# Set rotation parameters using ImageProcessingOptions.
|
||||
image_processing_options = _ImageProcessingOptions(
|
||||
rotation_degrees=rotation)
|
||||
callback_event = threading.Event()
|
||||
callback_exception: None | Exception = None
|
||||
observed_timestamp_ms = -1
|
||||
|
||||
def check_result(result: _HandLandmarkerResult, output_image: _Image,
|
||||
timestamp_ms: int):
|
||||
nonlocal callback_exception, observed_timestamp_ms
|
||||
|
||||
try:
|
||||
if (
|
||||
result.hand_landmarks
|
||||
and result.hand_world_landmarks
|
||||
and result.handedness
|
||||
):
|
||||
self._expect_hand_landmarker_results_correct(result, expected_result)
|
||||
else:
|
||||
self.assertEqual(result, expected_result)
|
||||
self.assertTrue(
|
||||
np.array_equal(output_image.numpy_view(), test_image.numpy_view())
|
||||
)
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
observed_timestamp_ms = timestamp_ms
|
||||
except AssertionError as e:
|
||||
callback_exception = e
|
||||
finally:
|
||||
callback_event.set()
|
||||
|
||||
options = _HandLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=check_result)
|
||||
with _HandLandmarker.create_from_options(options) as landmarker:
|
||||
for timestamp in range(0, 300, 30):
|
||||
landmarker.detect_async(test_image, timestamp, image_processing_options)
|
||||
callback_event.wait(timeout=3.0)
|
||||
if callback_exception is not None:
|
||||
raise callback_exception
|
||||
callback_event.clear()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
+594
@@ -0,0 +1,594 @@
|
||||
# Copyright 2023 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tests for holistic landmarker."""
|
||||
|
||||
import enum
|
||||
from unittest import mock
|
||||
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import numpy as np
|
||||
|
||||
from google.protobuf import text_format
|
||||
from mediapipe.tasks.cc.vision.holistic_landmarker.proto import holistic_result_pb2
|
||||
from mediapipe.tasks.python.components.containers import category as category_module
|
||||
from mediapipe.tasks.python.components.containers import landmark as landmark_module
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.test import test_utils
|
||||
from mediapipe.tasks.python.vision import holistic_landmarker
|
||||
from mediapipe.tasks.python.vision.core import image as image_lib
|
||||
from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module
|
||||
from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module
|
||||
|
||||
|
||||
HolisticLandmarkerResult = holistic_landmarker.HolisticLandmarkerResult
|
||||
_HolisticResultProto = holistic_result_pb2.HolisticResult
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
_Image = image_lib.Image
|
||||
_HolisticLandmarker = holistic_landmarker.HolisticLandmarker
|
||||
_HolisticLandmarkerOptions = holistic_landmarker.HolisticLandmarkerOptions
|
||||
_RUNNING_MODE = running_mode_module.VisionTaskRunningMode
|
||||
_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions
|
||||
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE = 'holistic_landmarker.task'
|
||||
_POSE_IMAGE = 'male_full_height_hands.jpg'
|
||||
_CAT_IMAGE = 'cat.jpg'
|
||||
_EXPECTED_HOLISTIC_RESULT = 'male_full_height_hands_result_cpu.pbtxt'
|
||||
_IMAGE_WIDTH = 638
|
||||
_IMAGE_HEIGHT = 1000
|
||||
_LANDMARKS_MARGIN = 0.03
|
||||
_BLENDSHAPES_MARGIN = 0.13
|
||||
_VIDEO_LANDMARKS_MARGIN = 0.03
|
||||
_VIDEO_BLENDSHAPES_MARGIN = 0.31
|
||||
_LIVE_STREAM_LANDMARKS_MARGIN = 0.03
|
||||
_LIVE_STREAM_BLENDSHAPES_MARGIN = 0.31
|
||||
|
||||
|
||||
def _create_landmarks_from_proto(landmark_list_proto, landmark_class):
|
||||
"""Creates a list of landmarks from a proto."""
|
||||
return [
|
||||
landmark_class(
|
||||
x=landmark.x,
|
||||
y=landmark.y,
|
||||
z=landmark.z,
|
||||
visibility=(
|
||||
landmark.visibility if landmark.HasField('visibility') else None
|
||||
),
|
||||
presence=(
|
||||
landmark.presence if landmark.HasField('presence') else None
|
||||
),
|
||||
)
|
||||
for landmark in landmark_list_proto.landmark
|
||||
]
|
||||
|
||||
|
||||
def _create_blendshapes_from_proto(blendshapes_proto):
|
||||
"""Creates a list of blendshapes from a proto."""
|
||||
return [
|
||||
category_module.Category(
|
||||
score=classification.score,
|
||||
index=classification.index,
|
||||
category_name=classification.label,
|
||||
display_name=classification.display_name,
|
||||
)
|
||||
for classification in blendshapes_proto.classification
|
||||
]
|
||||
|
||||
|
||||
def _get_expected_holistic_landmarker_result(
|
||||
file_path: str,
|
||||
) -> HolisticLandmarkerResult:
|
||||
holistic_result_file_path = test_utils.get_test_data_path(file_path)
|
||||
with open(holistic_result_file_path, 'rb') as f:
|
||||
holistic_result_proto = _HolisticResultProto()
|
||||
# Use this if a .pb file is available.
|
||||
# holistic_result_proto.ParseFromString(f.read())
|
||||
text_format.Parse(f.read(), holistic_result_proto)
|
||||
|
||||
face_landmarks = _create_landmarks_from_proto(
|
||||
holistic_result_proto.face_landmarks, landmark_module.NormalizedLandmark
|
||||
)
|
||||
pose_landmarks = _create_landmarks_from_proto(
|
||||
holistic_result_proto.pose_landmarks, landmark_module.NormalizedLandmark
|
||||
)
|
||||
pose_world_landmarks = _create_landmarks_from_proto(
|
||||
holistic_result_proto.pose_world_landmarks, landmark_module.Landmark
|
||||
)
|
||||
face_blendshapes_list = _create_blendshapes_from_proto(
|
||||
holistic_result_proto.face_blendshapes
|
||||
)
|
||||
|
||||
return HolisticLandmarkerResult(
|
||||
face_landmarks=face_landmarks,
|
||||
pose_landmarks=pose_landmarks,
|
||||
pose_world_landmarks=pose_world_landmarks,
|
||||
left_hand_landmarks=[],
|
||||
right_hand_landmarks=[],
|
||||
left_hand_world_landmarks=[],
|
||||
right_hand_world_landmarks=[],
|
||||
face_blendshapes=face_blendshapes_list,
|
||||
segmentation_mask=None,
|
||||
)
|
||||
|
||||
|
||||
class ModelFileType(enum.Enum):
|
||||
FILE_CONTENT = 1
|
||||
FILE_NAME = 2
|
||||
|
||||
|
||||
class HolisticLandmarkerTest(parameterized.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(_POSE_IMAGE)
|
||||
)
|
||||
self.model_path = test_utils.get_test_data_path(
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE
|
||||
)
|
||||
|
||||
def _expect_landmarks_correct(
|
||||
self, actual_landmarks, expected_landmarks, margin
|
||||
):
|
||||
# Expects to have the same number of landmarks detected.
|
||||
self.assertLen(actual_landmarks, len(expected_landmarks))
|
||||
|
||||
for i, elem in enumerate(actual_landmarks):
|
||||
self.assertAlmostEqual(elem.x, expected_landmarks[i].x, delta=margin)
|
||||
self.assertAlmostEqual(elem.y, expected_landmarks[i].y, delta=margin)
|
||||
|
||||
def _expect_blendshapes_correct(
|
||||
self, actual_blendshapes, expected_blendshapes, margin
|
||||
):
|
||||
# Expects to have the same number of blendshapes.
|
||||
self.assertLen(actual_blendshapes, len(expected_blendshapes))
|
||||
|
||||
for i, elem in enumerate(actual_blendshapes):
|
||||
self.assertEqual(elem.index, expected_blendshapes[i].index)
|
||||
self.assertEqual(
|
||||
elem.category_name, expected_blendshapes[i].category_name
|
||||
)
|
||||
self.assertAlmostEqual(
|
||||
elem.score,
|
||||
expected_blendshapes[i].score,
|
||||
delta=margin,
|
||||
)
|
||||
|
||||
def _expect_holistic_landmarker_results_correct(
|
||||
self,
|
||||
actual_result: HolisticLandmarkerResult,
|
||||
expected_result: HolisticLandmarkerResult,
|
||||
output_segmentation_mask: bool,
|
||||
landmarks_margin: float,
|
||||
blendshapes_margin: float,
|
||||
):
|
||||
self._expect_landmarks_correct(
|
||||
actual_result.pose_landmarks,
|
||||
expected_result.pose_landmarks,
|
||||
landmarks_margin,
|
||||
)
|
||||
self._expect_landmarks_correct(
|
||||
actual_result.face_landmarks,
|
||||
expected_result.face_landmarks,
|
||||
landmarks_margin,
|
||||
)
|
||||
self._expect_blendshapes_correct(
|
||||
actual_result.face_blendshapes,
|
||||
expected_result.face_blendshapes,
|
||||
blendshapes_margin,
|
||||
)
|
||||
if output_segmentation_mask:
|
||||
segmentation_mask = actual_result.segmentation_mask
|
||||
self.assertIsInstance(segmentation_mask, _Image)
|
||||
self.assertEqual(segmentation_mask.width, _IMAGE_WIDTH)
|
||||
self.assertEqual(segmentation_mask.height, _IMAGE_HEIGHT)
|
||||
else:
|
||||
self.assertIsNone(actual_result.segmentation_mask)
|
||||
|
||||
def test_create_from_file_succeeds_with_valid_model_path(self):
|
||||
# Creates with default option and valid model file successfully.
|
||||
with _HolisticLandmarker.create_from_model_path(
|
||||
self.model_path
|
||||
) as landmarker:
|
||||
self.assertIsInstance(landmarker, _HolisticLandmarker)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_path(self):
|
||||
# Creates with options containing model file successfully.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _HolisticLandmarkerOptions(base_options=base_options)
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
self.assertIsInstance(landmarker, _HolisticLandmarker)
|
||||
|
||||
def test_create_from_options_fails_with_invalid_model_path(self):
|
||||
# Invalid empty model path.
|
||||
with self.assertRaisesRegex(
|
||||
FileNotFoundError,
|
||||
'Unable to open file at /path/to/invalid/model.tflite',
|
||||
):
|
||||
base_options = _BaseOptions(
|
||||
model_asset_path='/path/to/invalid/model.tflite'
|
||||
)
|
||||
options = _HolisticLandmarkerOptions(base_options=base_options)
|
||||
_HolisticLandmarker.create_from_options(options)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_content(self):
|
||||
# Creates with options containing model content successfully.
|
||||
with open(self.model_path, 'rb') as f:
|
||||
base_options = _BaseOptions(model_asset_buffer=f.read())
|
||||
options = _HolisticLandmarkerOptions(base_options=base_options)
|
||||
landmarker = _HolisticLandmarker.create_from_options(options)
|
||||
self.assertIsInstance(landmarker, _HolisticLandmarker)
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
False,
|
||||
_get_expected_holistic_landmarker_result(_EXPECTED_HOLISTIC_RESULT),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
False,
|
||||
_get_expected_holistic_landmarker_result(_EXPECTED_HOLISTIC_RESULT),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
True,
|
||||
_get_expected_holistic_landmarker_result(_EXPECTED_HOLISTIC_RESULT),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
True,
|
||||
_get_expected_holistic_landmarker_result(_EXPECTED_HOLISTIC_RESULT),
|
||||
),
|
||||
)
|
||||
def test_detect(
|
||||
self,
|
||||
model_file_type,
|
||||
model_name,
|
||||
output_segmentation_mask,
|
||||
expected_holistic_landmarker_result,
|
||||
):
|
||||
# Creates holistic landmarker.
|
||||
model_path = test_utils.get_test_data_path(model_name)
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=base_options,
|
||||
output_face_blendshapes=True
|
||||
if expected_holistic_landmarker_result.face_blendshapes
|
||||
else False,
|
||||
output_segmentation_mask=output_segmentation_mask,
|
||||
)
|
||||
landmarker = _HolisticLandmarker.create_from_options(options)
|
||||
|
||||
# Performs holistic landmarks detection on the input.
|
||||
detection_result = landmarker.detect(self.test_image)
|
||||
self._expect_holistic_landmarker_results_correct(
|
||||
detection_result,
|
||||
expected_holistic_landmarker_result,
|
||||
output_segmentation_mask,
|
||||
_LANDMARKS_MARGIN,
|
||||
_BLENDSHAPES_MARGIN,
|
||||
)
|
||||
# Closes the holistic landmarker explicitly when the holistic landmarker is
|
||||
# not used in a context.
|
||||
landmarker.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
False,
|
||||
_get_expected_holistic_landmarker_result(_EXPECTED_HOLISTIC_RESULT),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
True,
|
||||
_get_expected_holistic_landmarker_result(_EXPECTED_HOLISTIC_RESULT),
|
||||
),
|
||||
)
|
||||
def test_detect_in_context(
|
||||
self,
|
||||
model_file_type,
|
||||
model_name,
|
||||
output_segmentation_mask,
|
||||
expected_holistic_landmarker_result,
|
||||
):
|
||||
# Creates holistic landmarker.
|
||||
model_path = test_utils.get_test_data_path(model_name)
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=base_options,
|
||||
output_face_blendshapes=True
|
||||
if expected_holistic_landmarker_result.face_blendshapes
|
||||
else False,
|
||||
output_segmentation_mask=output_segmentation_mask,
|
||||
)
|
||||
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
# Performs holistic landmarks detection on the input.
|
||||
detection_result = landmarker.detect(self.test_image)
|
||||
self._expect_holistic_landmarker_results_correct(
|
||||
detection_result,
|
||||
expected_holistic_landmarker_result,
|
||||
output_segmentation_mask,
|
||||
_LANDMARKS_MARGIN,
|
||||
_BLENDSHAPES_MARGIN,
|
||||
)
|
||||
|
||||
def test_empty_detection_outputs(self):
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path)
|
||||
)
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
# Load the cat image.
|
||||
cat_test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(_CAT_IMAGE)
|
||||
)
|
||||
# Performs holistic landmarks detection on the input.
|
||||
detection_result = landmarker.detect(cat_test_image)
|
||||
self.assertEmpty(detection_result.face_landmarks)
|
||||
self.assertIsNone(detection_result.face_blendshapes)
|
||||
self.assertIsNone(detection_result.segmentation_mask)
|
||||
|
||||
def test_missing_result_callback(self):
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback must be provided'
|
||||
):
|
||||
with _HolisticLandmarker.create_from_options(
|
||||
options
|
||||
) as unused_landmarker:
|
||||
pass
|
||||
|
||||
@parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO))
|
||||
def test_illegal_result_callback(self, running_mode):
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=running_mode,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback should not be provided'
|
||||
):
|
||||
with _HolisticLandmarker.create_from_options(
|
||||
options
|
||||
) as unused_landmarker:
|
||||
pass
|
||||
|
||||
def test_calling_detect_for_video_in_image_mode(self):
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the video mode'
|
||||
):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
def test_calling_detect_async_in_image_mode(self):
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the live stream mode'
|
||||
):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
def test_calling_detect_in_video_mode(self):
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the image mode'
|
||||
):
|
||||
landmarker.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_async_in_video_mode(self):
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the live stream mode'
|
||||
):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
def test_detect_for_video_with_out_of_order_timestamp(self):
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
unused_result = landmarker.detect_for_video(self.test_image, 1)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'Input timestamp must be monotonically increasing'
|
||||
):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
False,
|
||||
_get_expected_holistic_landmarker_result(_EXPECTED_HOLISTIC_RESULT),
|
||||
),
|
||||
(
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
True,
|
||||
_get_expected_holistic_landmarker_result(_EXPECTED_HOLISTIC_RESULT),
|
||||
),
|
||||
)
|
||||
def test_detect_for_video(
|
||||
self,
|
||||
model_name,
|
||||
output_segmentation_mask,
|
||||
expected_holistic_landmarker_result,
|
||||
):
|
||||
# Creates holistic landmarker.
|
||||
model_path = test_utils.get_test_data_path(model_name)
|
||||
base_options = _BaseOptions(model_asset_path=model_path)
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=base_options,
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
output_face_blendshapes=True
|
||||
if expected_holistic_landmarker_result.face_blendshapes
|
||||
else False,
|
||||
output_segmentation_mask=output_segmentation_mask,
|
||||
)
|
||||
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
for timestamp in range(0, 300, 30):
|
||||
# Performs holistic landmarks detection on the input.
|
||||
detection_result = landmarker.detect_for_video(
|
||||
self.test_image, timestamp
|
||||
)
|
||||
# Comparing results.
|
||||
self._expect_holistic_landmarker_results_correct(
|
||||
detection_result,
|
||||
expected_holistic_landmarker_result,
|
||||
output_segmentation_mask,
|
||||
_VIDEO_LANDMARKS_MARGIN,
|
||||
_VIDEO_BLENDSHAPES_MARGIN,
|
||||
)
|
||||
|
||||
def test_calling_detect_in_live_stream_mode(self):
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the image mode'
|
||||
):
|
||||
landmarker.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_for_video_in_live_stream_mode(self):
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'not initialized with the video mode'
|
||||
):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
def test_detect_async_calls_with_illegal_timestamp(self):
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
landmarker.detect_async(self.test_image, 100)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'Input timestamp must be monotonically increasing'
|
||||
):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
_POSE_IMAGE,
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
False,
|
||||
_get_expected_holistic_landmarker_result(_EXPECTED_HOLISTIC_RESULT),
|
||||
),
|
||||
(
|
||||
_POSE_IMAGE,
|
||||
_HOLISTIC_LANDMARKER_BUNDLE_ASSET_FILE,
|
||||
True,
|
||||
_get_expected_holistic_landmarker_result(_EXPECTED_HOLISTIC_RESULT),
|
||||
),
|
||||
)
|
||||
def test_detect_async_calls(
|
||||
self,
|
||||
image_path,
|
||||
model_name,
|
||||
output_segmentation_mask,
|
||||
expected_holistic_landmarker_result,
|
||||
):
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(image_path)
|
||||
)
|
||||
observed_timestamp_ms = -1
|
||||
|
||||
def check_result(
|
||||
result: HolisticLandmarkerResult,
|
||||
output_image: _Image,
|
||||
timestamp_ms: int,
|
||||
):
|
||||
# Comparing results.
|
||||
self._expect_holistic_landmarker_results_correct(
|
||||
result,
|
||||
expected_holistic_landmarker_result,
|
||||
output_segmentation_mask,
|
||||
_LIVE_STREAM_LANDMARKS_MARGIN,
|
||||
_LIVE_STREAM_BLENDSHAPES_MARGIN,
|
||||
)
|
||||
self.assertTrue(
|
||||
np.array_equal(output_image.numpy_view(), test_image.numpy_view())
|
||||
)
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
self.observed_timestamp_ms = timestamp_ms
|
||||
|
||||
model_path = test_utils.get_test_data_path(model_name)
|
||||
options = _HolisticLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
output_face_blendshapes=True
|
||||
if expected_holistic_landmarker_result.face_blendshapes
|
||||
else False,
|
||||
output_segmentation_mask=output_segmentation_mask,
|
||||
result_callback=check_result,
|
||||
)
|
||||
with _HolisticLandmarker.create_from_options(options) as landmarker:
|
||||
for timestamp in range(0, 300, 30):
|
||||
landmarker.detect_async(test_image, timestamp)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
+679
@@ -0,0 +1,679 @@
|
||||
# Copyright 2022 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tests for image classifier."""
|
||||
|
||||
import enum
|
||||
import os
|
||||
import threading
|
||||
from unittest import mock
|
||||
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import numpy as np
|
||||
|
||||
from mediapipe.tasks.python.components.containers import category as category_module
|
||||
from mediapipe.tasks.python.components.containers import classification_result as classification_result_module
|
||||
from mediapipe.tasks.python.components.containers import rect as rect_module
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.test import test_utils
|
||||
from mediapipe.tasks.python.vision import image_classifier
|
||||
from mediapipe.tasks.python.vision.core import image as image_module
|
||||
from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module
|
||||
from mediapipe.tasks.python.vision.core import vision_task_running_mode
|
||||
|
||||
ImageClassifierResult = classification_result_module.ClassificationResult
|
||||
_RectF = rect_module.RectF
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
_Category = category_module.Category
|
||||
_Classifications = classification_result_module.Classifications
|
||||
_Image = image_module.Image
|
||||
_ImageClassifier = image_classifier.ImageClassifier
|
||||
_ImageClassifierOptions = image_classifier.ImageClassifierOptions
|
||||
_RUNNING_MODE = vision_task_running_mode.VisionTaskRunningMode
|
||||
_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions
|
||||
|
||||
_MODEL_FILE = 'mobilenet_v2_1.0_224.tflite'
|
||||
_IMAGE_FILE = 'burger.jpg'
|
||||
_ALLOW_LIST = ['cheeseburger', 'guacamole']
|
||||
_DENY_LIST = ['cheeseburger']
|
||||
_SCORE_THRESHOLD = 0.5
|
||||
_MAX_RESULTS = 3
|
||||
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision'
|
||||
|
||||
|
||||
def _generate_empty_results() -> ImageClassifierResult:
|
||||
return ImageClassifierResult(
|
||||
classifications=[
|
||||
_Classifications(categories=[], head_index=0, head_name='probability')
|
||||
],
|
||||
timestamp_ms=0,
|
||||
)
|
||||
|
||||
|
||||
def _generate_burger_results(timestamp_ms=0) -> ImageClassifierResult:
|
||||
return ImageClassifierResult(
|
||||
classifications=[
|
||||
_Classifications(
|
||||
categories=[
|
||||
_Category(
|
||||
index=934,
|
||||
score=0.793959,
|
||||
display_name='',
|
||||
category_name='cheeseburger',
|
||||
),
|
||||
_Category(
|
||||
index=932,
|
||||
score=0.0273929,
|
||||
display_name='',
|
||||
category_name='bagel',
|
||||
),
|
||||
_Category(
|
||||
index=925,
|
||||
score=0.0193408,
|
||||
display_name='',
|
||||
category_name='guacamole',
|
||||
),
|
||||
_Category(
|
||||
index=963,
|
||||
score=0.00632786,
|
||||
display_name='',
|
||||
category_name='meat loaf',
|
||||
),
|
||||
],
|
||||
head_index=0,
|
||||
head_name='probability',
|
||||
)
|
||||
],
|
||||
timestamp_ms=timestamp_ms,
|
||||
)
|
||||
|
||||
|
||||
def _generate_soccer_ball_results(timestamp_ms=0) -> ImageClassifierResult:
|
||||
return ImageClassifierResult(
|
||||
classifications=[
|
||||
_Classifications(
|
||||
categories=[
|
||||
_Category(
|
||||
index=806,
|
||||
score=0.996527,
|
||||
display_name='',
|
||||
category_name='soccer ball',
|
||||
)
|
||||
],
|
||||
head_index=0,
|
||||
head_name='probability',
|
||||
)
|
||||
],
|
||||
timestamp_ms=timestamp_ms,
|
||||
)
|
||||
|
||||
|
||||
class ModelFileType(enum.Enum):
|
||||
FILE_CONTENT = 1
|
||||
FILE_NAME = 2
|
||||
|
||||
|
||||
class ImageClassifierTest(parameterized.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, _IMAGE_FILE))
|
||||
)
|
||||
self.model_path = test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _MODEL_FILE)
|
||||
)
|
||||
|
||||
def assertCategoryAlmostEqual(
|
||||
self, actual: _Category, expected: _Category, delta: float = 1e-6
|
||||
):
|
||||
self.assertEqual(actual.index, expected.index)
|
||||
self.assertAlmostEqual(actual.score, expected.score, delta=delta)
|
||||
self.assertEqual(actual.display_name or '', expected.display_name)
|
||||
self.assertEqual(actual.category_name, expected.category_name)
|
||||
|
||||
def assertClassificationsAlmostEqual(
|
||||
self, actual: _Classifications, expected: _Classifications
|
||||
):
|
||||
self.assertEqual(actual.head_index, expected.head_index)
|
||||
self.assertEqual(actual.head_name, expected.head_name)
|
||||
self.assertLen(actual.categories, len(expected.categories))
|
||||
for i, actual_category in enumerate(actual.categories):
|
||||
self.assertCategoryAlmostEqual(actual_category, expected.categories[i])
|
||||
|
||||
def assertClassificationResultCorrect(
|
||||
self, actual: ImageClassifierResult, expected: ImageClassifierResult
|
||||
):
|
||||
self.assertEqual(actual.timestamp_ms, expected.timestamp_ms)
|
||||
self.assertLen(actual.classifications, len(expected.classifications))
|
||||
for i, actual_classifications in enumerate(actual.classifications):
|
||||
self.assertClassificationsAlmostEqual(
|
||||
actual_classifications, expected.classifications[i]
|
||||
)
|
||||
|
||||
def test_create_from_file_succeeds_with_valid_model_path(self):
|
||||
# Creates with default option and valid model file successfully.
|
||||
with _ImageClassifier.create_from_model_path(self.model_path) as classifier:
|
||||
self.assertIsInstance(classifier, _ImageClassifier)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_path(self):
|
||||
# Creates with options containing model file successfully.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _ImageClassifierOptions(base_options=base_options)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
self.assertIsInstance(classifier, _ImageClassifier)
|
||||
|
||||
def test_create_from_options_fails_with_invalid_model_path(self):
|
||||
with self.assertRaisesRegex(
|
||||
FileNotFoundError,
|
||||
'Unable to open file at /path/to/invalid/model.tflite',
|
||||
):
|
||||
base_options = _BaseOptions(
|
||||
model_asset_path='/path/to/invalid/model.tflite'
|
||||
)
|
||||
options = _ImageClassifierOptions(base_options=base_options)
|
||||
_ImageClassifier.create_from_options(options)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_content(self):
|
||||
# Creates with options containing model content successfully.
|
||||
with open(self.model_path, 'rb') as f:
|
||||
base_options = _BaseOptions(model_asset_buffer=f.read())
|
||||
options = _ImageClassifierOptions(base_options=base_options)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
self.assertIsInstance(classifier, _ImageClassifier)
|
||||
|
||||
@parameterized.parameters(
|
||||
(ModelFileType.FILE_NAME, 4, _generate_burger_results()),
|
||||
(ModelFileType.FILE_CONTENT, 4, _generate_burger_results()),
|
||||
)
|
||||
def test_classify(
|
||||
self, model_file_type, max_results, expected_classification_result
|
||||
):
|
||||
# Creates classifier.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=base_options, max_results=max_results
|
||||
)
|
||||
classifier = _ImageClassifier.create_from_options(options)
|
||||
|
||||
# Performs image classification on the input.
|
||||
image_result = classifier.classify(self.test_image)
|
||||
# Comparing results.
|
||||
self.assertClassificationResultCorrect(
|
||||
image_result, expected_classification_result
|
||||
)
|
||||
# Closes the classifier explicitly when the classifier is not used in
|
||||
# a context.
|
||||
classifier.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(ModelFileType.FILE_NAME, 4, _generate_burger_results()),
|
||||
(ModelFileType.FILE_CONTENT, 4, _generate_burger_results()),
|
||||
)
|
||||
def test_classify_in_context(
|
||||
self, model_file_type, max_results, expected_classification_result
|
||||
):
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=base_options, max_results=max_results
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
# Performs image classification on the input.
|
||||
image_result = classifier.classify(self.test_image)
|
||||
# Comparing results.
|
||||
self.assertClassificationResultCorrect(
|
||||
image_result, expected_classification_result
|
||||
)
|
||||
|
||||
def test_classify_succeeds_with_region_of_interest(self):
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _ImageClassifierOptions(base_options=base_options, max_results=1)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
# Load the test image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, 'multi_objects.jpg')
|
||||
)
|
||||
)
|
||||
# Region-of-interest around the soccer ball.
|
||||
roi = _RectF(left=0.45, top=0.3075, right=0.614, bottom=0.7345)
|
||||
image_processing_options = _ImageProcessingOptions(roi)
|
||||
# Performs image classification on the input.
|
||||
image_result = classifier.classify(test_image, image_processing_options)
|
||||
# Comparing results.
|
||||
self.assertClassificationResultCorrect(
|
||||
image_result, _generate_soccer_ball_results()
|
||||
)
|
||||
|
||||
def test_classify_succeeds_with_rotation(self):
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _ImageClassifierOptions(base_options=base_options, max_results=3)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
# Load the test image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, 'burger_rotated.jpg')
|
||||
)
|
||||
)
|
||||
# Specify a 90° anti-clockwise rotation.
|
||||
image_processing_options = _ImageProcessingOptions(None, -90)
|
||||
# Performs image classification on the input.
|
||||
image_result = classifier.classify(test_image, image_processing_options)
|
||||
# Comparing results.
|
||||
expected = ImageClassifierResult(
|
||||
classifications=[
|
||||
_Classifications(
|
||||
categories=[
|
||||
_Category(
|
||||
index=934,
|
||||
score=0.754467,
|
||||
display_name='',
|
||||
category_name='cheeseburger',
|
||||
),
|
||||
_Category(
|
||||
index=925,
|
||||
score=0.0288028,
|
||||
display_name='',
|
||||
category_name='guacamole',
|
||||
),
|
||||
_Category(
|
||||
index=932,
|
||||
score=0.0286119,
|
||||
display_name='',
|
||||
category_name='bagel',
|
||||
),
|
||||
],
|
||||
head_index=0,
|
||||
head_name='probability',
|
||||
)
|
||||
],
|
||||
timestamp_ms=0,
|
||||
)
|
||||
self.assertClassificationResultCorrect(image_result, expected)
|
||||
|
||||
def test_classify_succeeds_with_region_of_interest_and_rotation(self):
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _ImageClassifierOptions(base_options=base_options, max_results=1)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
# Load the test image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, 'multi_objects_rotated.jpg')
|
||||
)
|
||||
)
|
||||
# Region-of-interest around the soccer ball, with 90° anti-clockwise
|
||||
# rotation.
|
||||
roi = _RectF(left=0.2655, top=0.45, right=0.6925, bottom=0.614)
|
||||
image_processing_options = _ImageProcessingOptions(roi, -90)
|
||||
# Performs image classification on the input.
|
||||
image_result = classifier.classify(test_image, image_processing_options)
|
||||
# Comparing results.
|
||||
expected = ImageClassifierResult(
|
||||
classifications=[
|
||||
_Classifications(
|
||||
categories=[
|
||||
_Category(
|
||||
index=806,
|
||||
score=0.997684,
|
||||
display_name='',
|
||||
category_name='soccer ball',
|
||||
),
|
||||
],
|
||||
head_index=0,
|
||||
head_name='probability',
|
||||
)
|
||||
],
|
||||
timestamp_ms=0,
|
||||
)
|
||||
self.assertClassificationResultCorrect(image_result, expected)
|
||||
|
||||
def test_score_threshold_option(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
score_threshold=_SCORE_THRESHOLD,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
# Performs image classification on the input.
|
||||
image_result = classifier.classify(self.test_image)
|
||||
classifications = image_result.classifications
|
||||
|
||||
for classification in classifications:
|
||||
for category in classification.categories:
|
||||
score = category.score
|
||||
self.assertGreaterEqual(
|
||||
score,
|
||||
_SCORE_THRESHOLD,
|
||||
(
|
||||
'Classification with score lower than threshold found. '
|
||||
f'{classification}'
|
||||
),
|
||||
)
|
||||
|
||||
def test_max_results_option(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
score_threshold=_SCORE_THRESHOLD,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
# Performs image classification on the input.
|
||||
image_result = classifier.classify(self.test_image)
|
||||
categories = image_result.classifications[0].categories
|
||||
|
||||
self.assertLessEqual(
|
||||
len(categories), _MAX_RESULTS, 'Too many results returned.'
|
||||
)
|
||||
|
||||
def test_allow_list_option(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
category_allowlist=_ALLOW_LIST,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
# Performs image classification on the input.
|
||||
image_result = classifier.classify(self.test_image)
|
||||
classifications = image_result.classifications
|
||||
|
||||
for classification in classifications:
|
||||
for category in classification.categories:
|
||||
label = category.category_name
|
||||
self.assertIn(
|
||||
label,
|
||||
_ALLOW_LIST,
|
||||
f'Label {label} found but not in label allow list',
|
||||
)
|
||||
|
||||
def test_deny_list_option(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
category_denylist=_DENY_LIST,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
# Performs image classification on the input.
|
||||
image_result = classifier.classify(self.test_image)
|
||||
classifications = image_result.classifications
|
||||
|
||||
for classification in classifications:
|
||||
for category in classification.categories:
|
||||
label = category.category_name
|
||||
self.assertNotIn(
|
||||
label, _DENY_LIST, f'Label {label} found but in deny list.'
|
||||
)
|
||||
|
||||
def test_combined_allowlist_and_denylist(self):
|
||||
# Fails with combined allowlist and denylist
|
||||
with self.assertRaises(ValueError):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
category_allowlist=['foo'],
|
||||
category_denylist=['bar'],
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as unused_classifier:
|
||||
pass
|
||||
|
||||
def test_empty_classification_outputs(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
score_threshold=1,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
# Performs image classification on the input.
|
||||
image_result = classifier.classify(self.test_image)
|
||||
self.assertEmpty(image_result.classifications[0].categories)
|
||||
|
||||
def test_missing_result_callback(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback must be provided'
|
||||
):
|
||||
with _ImageClassifier.create_from_options(options) as unused_classifier:
|
||||
pass
|
||||
|
||||
@parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO))
|
||||
def test_illegal_result_callback(self, running_mode):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=running_mode,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback should not be provided'
|
||||
):
|
||||
with _ImageClassifier.create_from_options(options) as unused_classifier:
|
||||
pass
|
||||
|
||||
def test_calling_classify_for_video_in_image_mode(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
with self.assertRaises(ValueError):
|
||||
classifier.classify_for_video(self.test_image, 0)
|
||||
|
||||
def test_calling_classify_async_in_image_mode(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
with self.assertRaises(ValueError):
|
||||
classifier.classify_async(self.test_image, 0)
|
||||
|
||||
def test_calling_classify_in_video_mode(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
with self.assertRaises(ValueError):
|
||||
classifier.classify(self.test_image)
|
||||
|
||||
def test_calling_classify_async_in_video_mode(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
with self.assertRaises(ValueError):
|
||||
classifier.classify_async(self.test_image, 0)
|
||||
|
||||
def test_classify_for_video_with_out_of_order_timestamp(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
unused_result = classifier.classify_for_video(self.test_image, 1)
|
||||
with self.assertRaises(ValueError):
|
||||
classifier.classify_for_video(self.test_image, 0)
|
||||
|
||||
def test_classify_for_video(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
max_results=4,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
for timestamp in range(0, 300, 30):
|
||||
classification_result = classifier.classify_for_video(
|
||||
self.test_image, timestamp
|
||||
)
|
||||
self.assertClassificationResultCorrect(
|
||||
classification_result, _generate_burger_results(timestamp)
|
||||
)
|
||||
|
||||
def test_classify_for_video_succeeds_with_region_of_interest(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
max_results=1,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
# Load the test image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, 'multi_objects.jpg')
|
||||
)
|
||||
)
|
||||
# Region-of-interest around the soccer ball.
|
||||
roi = _RectF(left=0.45, top=0.3075, right=0.614, bottom=0.7345)
|
||||
image_processing_options = _ImageProcessingOptions(roi)
|
||||
for timestamp in range(0, 300, 30):
|
||||
classification_result = classifier.classify_for_video(
|
||||
test_image, timestamp, image_processing_options
|
||||
)
|
||||
self.assertClassificationResultCorrect(
|
||||
classification_result, _generate_soccer_ball_results(timestamp)
|
||||
)
|
||||
|
||||
def test_calling_classify_in_live_stream_mode(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
with self.assertRaises(ValueError):
|
||||
classifier.classify(self.test_image)
|
||||
|
||||
def test_calling_classify_for_video_in_live_stream_mode(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
with self.assertRaises(ValueError):
|
||||
classifier.classify_for_video(self.test_image, 0)
|
||||
|
||||
def test_classify_async_calls_with_illegal_timestamp(self):
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
max_results=4,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
classifier.classify_async(self.test_image, 100)
|
||||
with self.assertRaises(ValueError):
|
||||
classifier.classify_async(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(0, _generate_burger_results()), (1, _generate_empty_results())
|
||||
)
|
||||
def test_classify_async_calls(self, threshold, expected_result):
|
||||
observed_timestamp_ms = -1
|
||||
callback_event = threading.Event()
|
||||
callback_exception = None
|
||||
|
||||
def check_result(
|
||||
result: ImageClassifierResult, output_image: _Image, timestamp_ms: int
|
||||
):
|
||||
nonlocal callback_exception, observed_timestamp_ms
|
||||
try:
|
||||
self.assertClassificationResultCorrect(result, expected_result)
|
||||
self.assertTrue(
|
||||
np.array_equal(
|
||||
output_image.numpy_view(), self.test_image.numpy_view()
|
||||
)
|
||||
)
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
observed_timestamp_ms = timestamp_ms
|
||||
except AssertionError as e:
|
||||
callback_exception = e
|
||||
finally:
|
||||
callback_event.set()
|
||||
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
max_results=4,
|
||||
score_threshold=threshold,
|
||||
result_callback=check_result,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
classifier.classify_async(self.test_image, 0)
|
||||
callback_event.wait(3)
|
||||
if callback_exception is not None:
|
||||
raise callback_exception
|
||||
callback_event.clear()
|
||||
|
||||
def test_classify_async_succeeds_with_region_of_interest(self):
|
||||
# Load the test image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, 'multi_objects.jpg')
|
||||
)
|
||||
)
|
||||
# Region-of-interest around the soccer ball.
|
||||
roi = _RectF(left=0.45, top=0.3075, right=0.614, bottom=0.7345)
|
||||
image_processing_options = _ImageProcessingOptions(roi)
|
||||
observed_timestamp_ms = -1
|
||||
callback_event = threading.Event()
|
||||
callback_exception = None
|
||||
|
||||
def check_result(
|
||||
result: ImageClassifierResult, output_image: _Image, timestamp_ms: int
|
||||
):
|
||||
nonlocal callback_exception, observed_timestamp_ms
|
||||
try:
|
||||
self.assertClassificationResultCorrect(
|
||||
result, _generate_soccer_ball_results(100)
|
||||
)
|
||||
self.assertEqual(output_image.width, test_image.width)
|
||||
self.assertEqual(output_image.height, test_image.height)
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
observed_timestamp_ms = timestamp_ms
|
||||
except AssertionError as e:
|
||||
callback_exception = e
|
||||
finally:
|
||||
callback_event.set()
|
||||
|
||||
options = _ImageClassifierOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
max_results=1,
|
||||
result_callback=check_result,
|
||||
)
|
||||
with _ImageClassifier.create_from_options(options) as classifier:
|
||||
classifier.classify_async(test_image, 100, image_processing_options)
|
||||
callback_event.wait(3)
|
||||
if callback_exception is not None:
|
||||
raise callback_exception
|
||||
callback_event.clear()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
+450
@@ -0,0 +1,450 @@
|
||||
# Copyright 2022 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tests for image embedder."""
|
||||
|
||||
import enum
|
||||
import os
|
||||
import threading
|
||||
from unittest import mock
|
||||
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import numpy as np
|
||||
|
||||
from mediapipe.tasks.python.components.containers import embedding_result as embedding_result_module
|
||||
from mediapipe.tasks.python.components.containers import rect as rect_module
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.test import test_utils
|
||||
from mediapipe.tasks.python.vision import image_embedder
|
||||
from mediapipe.tasks.python.vision.core import image as image_module
|
||||
from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module
|
||||
from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module
|
||||
|
||||
_RectF = rect_module.RectF
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
_Embedding = embedding_result_module.Embedding
|
||||
_Image = image_module.Image
|
||||
_ImageEmbedder = image_embedder.ImageEmbedder
|
||||
_ImageEmbedderOptions = image_embedder.ImageEmbedderOptions
|
||||
_ImageEmbedderResult = image_embedder.ImageEmbedderResult
|
||||
_RUNNING_MODE = running_mode_module.VisionTaskRunningMode
|
||||
_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions
|
||||
|
||||
_MODEL_FILE = 'mobilenet_v3_small_100_224_embedder.tflite'
|
||||
_BURGER_IMAGE_FILE = 'burger.jpg'
|
||||
_BURGER_CROPPED_IMAGE_FILE = 'burger_crop.jpg'
|
||||
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision'
|
||||
# Tolerance for embedding vector coordinate values.
|
||||
_EPSILON = 1e-4
|
||||
# Tolerance for cosine similarity evaluation.
|
||||
_SIMILARITY_TOLERANCE = 1e-6
|
||||
|
||||
|
||||
class ModelFileType(enum.Enum):
|
||||
FILE_CONTENT = 1
|
||||
FILE_NAME = 2
|
||||
|
||||
|
||||
class ImageEmbedderTest(parameterized.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _BURGER_IMAGE_FILE)))
|
||||
self.test_cropped_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _BURGER_CROPPED_IMAGE_FILE)))
|
||||
self.model_path = test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _MODEL_FILE))
|
||||
|
||||
def test_create_from_file_succeeds_with_valid_model_path(self):
|
||||
# Creates with default option and valid model file successfully.
|
||||
with _ImageEmbedder.create_from_model_path(self.model_path) as embedder:
|
||||
self.assertIsInstance(embedder, _ImageEmbedder)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_path(self):
|
||||
# Creates with options containing model file successfully.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _ImageEmbedderOptions(base_options=base_options)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
self.assertIsInstance(embedder, _ImageEmbedder)
|
||||
|
||||
def test_create_from_options_fails_with_invalid_model_path(self):
|
||||
with self.assertRaisesRegex(
|
||||
FileNotFoundError,
|
||||
'Unable to open file at /path/to/invalid/model.tflite',
|
||||
):
|
||||
base_options = _BaseOptions(
|
||||
model_asset_path='/path/to/invalid/model.tflite')
|
||||
options = _ImageEmbedderOptions(base_options=base_options)
|
||||
embedder = _ImageEmbedder.create_from_options(options)
|
||||
embedder.close()
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_content(self):
|
||||
# Creates with options containing model content successfully.
|
||||
with open(self.model_path, 'rb') as f:
|
||||
base_options = _BaseOptions(model_asset_buffer=f.read())
|
||||
options = _ImageEmbedderOptions(base_options=base_options)
|
||||
embedder = _ImageEmbedder.create_from_options(options)
|
||||
self.assertIsInstance(embedder, _ImageEmbedder)
|
||||
embedder.close()
|
||||
|
||||
def _check_embedding_value(self, result, expected_first_value):
|
||||
# Check embedding first value.
|
||||
self.assertAlmostEqual(
|
||||
result.embeddings[0].embedding[0], expected_first_value, delta=_EPSILON)
|
||||
|
||||
def _check_embedding_size(self, result, quantize, expected_embedding_size):
|
||||
# Check embedding size.
|
||||
self.assertLen(result.embeddings, 1)
|
||||
embedding_result = result.embeddings[0]
|
||||
self.assertLen(embedding_result.embedding, expected_embedding_size)
|
||||
if quantize:
|
||||
self.assertEqual(embedding_result.embedding.dtype, np.uint8)
|
||||
else:
|
||||
self.assertEqual(embedding_result.embedding.dtype, float)
|
||||
|
||||
def _check_cosine_similarity(self, result0, result1, expected_similarity):
|
||||
# Checks cosine similarity.
|
||||
similarity = _ImageEmbedder.cosine_similarity(result0.embeddings[0],
|
||||
result1.embeddings[0])
|
||||
self.assertAlmostEqual(
|
||||
similarity, expected_similarity, delta=_SIMILARITY_TOLERANCE)
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
False,
|
||||
False,
|
||||
False,
|
||||
ModelFileType.FILE_NAME,
|
||||
0.925519,
|
||||
1024,
|
||||
(-0.2101883, -0.193027),
|
||||
),
|
||||
(
|
||||
True,
|
||||
False,
|
||||
False,
|
||||
ModelFileType.FILE_NAME,
|
||||
0.925519,
|
||||
1024,
|
||||
(-0.0142344, -0.0131606),
|
||||
),
|
||||
(
|
||||
False,
|
||||
True,
|
||||
False,
|
||||
ModelFileType.FILE_NAME,
|
||||
0.926791,
|
||||
1024,
|
||||
(229, 231),
|
||||
),
|
||||
(
|
||||
False,
|
||||
False,
|
||||
True,
|
||||
ModelFileType.FILE_CONTENT,
|
||||
0.999931,
|
||||
1024,
|
||||
(-0.195062, -0.193027),
|
||||
),
|
||||
)
|
||||
def test_embed(self, l2_normalize, quantize, with_roi, model_file_type,
|
||||
expected_similarity, expected_size, expected_first_values):
|
||||
# Creates embedder.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=base_options, l2_normalize=l2_normalize, quantize=quantize)
|
||||
embedder = _ImageEmbedder.create_from_options(options)
|
||||
|
||||
image_processing_options = None
|
||||
if with_roi:
|
||||
# Region-of-interest in "burger.jpg" corresponding to "burger_crop.jpg".
|
||||
roi = _RectF(left=0.0, top=0.0, right=0.833333, bottom=1.0)
|
||||
image_processing_options = _ImageProcessingOptions(roi)
|
||||
|
||||
# Extracts both embeddings.
|
||||
image_result = embedder.embed(self.test_image, image_processing_options)
|
||||
crop_result = embedder.embed(self.test_cropped_image)
|
||||
|
||||
# Checks embeddings and cosine similarity.
|
||||
expected_result0_value, expected_result1_value = expected_first_values
|
||||
self._check_embedding_size(image_result, quantize, expected_size)
|
||||
self._check_embedding_size(crop_result, quantize, expected_size)
|
||||
self._check_embedding_value(image_result, expected_result0_value)
|
||||
self._check_embedding_value(crop_result, expected_result1_value)
|
||||
self._check_cosine_similarity(image_result, crop_result,
|
||||
expected_similarity)
|
||||
# Closes the embedder explicitly when the embedder is not used in
|
||||
# a context.
|
||||
embedder.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(False, False, ModelFileType.FILE_NAME, 0.925519),
|
||||
(False, False, ModelFileType.FILE_CONTENT, 0.925519))
|
||||
def test_embed_in_context(self, l2_normalize, quantize, model_file_type,
|
||||
expected_similarity):
|
||||
# Creates embedder.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=base_options, l2_normalize=l2_normalize, quantize=quantize)
|
||||
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
# Extracts both embeddings.
|
||||
image_result = embedder.embed(self.test_image)
|
||||
crop_result = embedder.embed(self.test_cropped_image)
|
||||
|
||||
# Checks cosine similarity.
|
||||
self._check_cosine_similarity(image_result, crop_result,
|
||||
expected_similarity)
|
||||
|
||||
def test_missing_result_callback(self):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM)
|
||||
with self.assertRaisesRegex(ValueError,
|
||||
r'result callback must be provided'):
|
||||
with _ImageEmbedder.create_from_options(options) as unused_embedder:
|
||||
pass
|
||||
|
||||
@parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO))
|
||||
def test_illegal_result_callback(self, running_mode):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=running_mode,
|
||||
result_callback=mock.MagicMock())
|
||||
with self.assertRaisesRegex(ValueError,
|
||||
r'result callback should not be provided'):
|
||||
with _ImageEmbedder.create_from_options(options) as unused_embedder:
|
||||
pass
|
||||
|
||||
def test_calling_embed_for_video_in_image_mode(self):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
with self.assertRaises(ValueError):
|
||||
embedder.embed_for_video(self.test_image, 0)
|
||||
|
||||
def test_calling_embed_async_in_image_mode(self):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
with self.assertRaises(ValueError):
|
||||
embedder.embed_async(self.test_image, 0)
|
||||
|
||||
def test_calling_embed_in_video_mode(self):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
with self.assertRaises(ValueError):
|
||||
embedder.embed(self.test_image)
|
||||
|
||||
def test_calling_embed_async_in_video_mode(self):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
with self.assertRaises(ValueError):
|
||||
embedder.embed_async(self.test_image, 0)
|
||||
|
||||
def test_embed_for_video_with_out_of_order_timestamp(self):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
unused_result = embedder.embed_for_video(self.test_image, 1)
|
||||
with self.assertRaises(ValueError):
|
||||
embedder.embed_for_video(self.test_image, 0)
|
||||
|
||||
def test_embed_for_video(self):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder0, \
|
||||
_ImageEmbedder.create_from_options(options) as embedder1:
|
||||
for timestamp in range(0, 300, 30):
|
||||
# Extracts both embeddings.
|
||||
image_result = embedder0.embed_for_video(self.test_image, timestamp)
|
||||
crop_result = embedder1.embed_for_video(self.test_cropped_image,
|
||||
timestamp)
|
||||
# Checks cosine similarity.
|
||||
self._check_cosine_similarity(
|
||||
image_result, crop_result, expected_similarity=0.925519)
|
||||
|
||||
def test_embed_for_video_succeeds_with_region_of_interest(self):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder0, \
|
||||
_ImageEmbedder.create_from_options(options) as embedder1:
|
||||
# Region-of-interest in "burger.jpg" corresponding to "burger_crop.jpg".
|
||||
roi = _RectF(left=0.0, top=0.0, right=0.833333, bottom=1.0)
|
||||
image_processing_options = _ImageProcessingOptions(roi)
|
||||
|
||||
for timestamp in range(0, 300, 30):
|
||||
# Extracts both embeddings.
|
||||
image_result = embedder0.embed_for_video(self.test_image, timestamp,
|
||||
image_processing_options)
|
||||
crop_result = embedder1.embed_for_video(self.test_cropped_image,
|
||||
timestamp)
|
||||
|
||||
# Checks cosine similarity.
|
||||
self._check_cosine_similarity(
|
||||
image_result, crop_result, expected_similarity=0.999931)
|
||||
|
||||
def test_calling_embed_in_live_stream_mode(self):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock())
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
with self.assertRaises(ValueError):
|
||||
embedder.embed(self.test_image)
|
||||
|
||||
def test_calling_embed_for_video_in_live_stream_mode(self):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock())
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
with self.assertRaises(ValueError):
|
||||
embedder.embed_for_video(self.test_image, 0)
|
||||
|
||||
def test_embed_async_calls_with_illegal_timestamp(self):
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock())
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
embedder.embed_async(self.test_image, 100)
|
||||
with self.assertRaises(ValueError):
|
||||
embedder.embed_async(self.test_image, 0)
|
||||
|
||||
def test_embed_async_calls(self):
|
||||
# Get the embedding result for the cropped image.
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
crop_result = embedder.embed(self.test_cropped_image)
|
||||
|
||||
callback_event = threading.Event()
|
||||
callback_exception: None | Exception = None
|
||||
observed_timestamp_ms = -1
|
||||
|
||||
def check_result(result: _ImageEmbedderResult, output_image: _Image,
|
||||
timestamp_ms: int):
|
||||
nonlocal callback_exception, observed_timestamp_ms
|
||||
try:
|
||||
# Checks cosine similarity.
|
||||
self._check_cosine_similarity(
|
||||
result, crop_result, expected_similarity=0.925519
|
||||
)
|
||||
self.assertTrue(
|
||||
np.array_equal(
|
||||
output_image.numpy_view(), self.test_image.numpy_view()
|
||||
)
|
||||
)
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
observed_timestamp_ms = timestamp_ms
|
||||
except AssertionError as e:
|
||||
callback_exception = e
|
||||
finally:
|
||||
callback_event.set()
|
||||
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=check_result)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
for timestamp in range(0, 300, 30):
|
||||
embedder.embed_async(self.test_image, timestamp)
|
||||
callback_event.wait()
|
||||
if callback_exception is not None:
|
||||
raise callback_exception
|
||||
callback_event.clear()
|
||||
|
||||
def test_embed_async_succeeds_with_region_of_interest(self):
|
||||
# Get the embedding result for the cropped image.
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
crop_result = embedder.embed(self.test_cropped_image)
|
||||
|
||||
# Region-of-interest in "burger.jpg" corresponding to "burger_crop.jpg".
|
||||
roi = _RectF(left=0.0, top=0.0, right=0.833333, bottom=1.0)
|
||||
image_processing_options = _ImageProcessingOptions(roi)
|
||||
callback_event = threading.Event()
|
||||
callback_exception: None | Exception = None
|
||||
observed_timestamp_ms = -1
|
||||
|
||||
def check_result(result: _ImageEmbedderResult, output_image: _Image,
|
||||
timestamp_ms: int):
|
||||
nonlocal callback_exception, observed_timestamp_ms
|
||||
try:
|
||||
# Checks cosine similarity.
|
||||
self._check_cosine_similarity(
|
||||
result, crop_result, expected_similarity=0.999931
|
||||
)
|
||||
self.assertTrue(
|
||||
np.array_equal(
|
||||
output_image.numpy_view(), self.test_image.numpy_view()
|
||||
)
|
||||
)
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
observed_timestamp_ms = timestamp_ms
|
||||
except AssertionError as e:
|
||||
callback_exception = e
|
||||
finally:
|
||||
callback_event.set()
|
||||
|
||||
options = _ImageEmbedderOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=check_result)
|
||||
with _ImageEmbedder.create_from_options(options) as embedder:
|
||||
for timestamp in range(0, 300, 30):
|
||||
embedder.embed_async(self.test_image, timestamp,
|
||||
image_processing_options)
|
||||
callback_event.wait()
|
||||
if callback_exception is not None:
|
||||
raise callback_exception
|
||||
callback_event.clear()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
+529
@@ -0,0 +1,529 @@
|
||||
# Copyright 2022 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tests for image segmenter."""
|
||||
|
||||
import enum
|
||||
import os
|
||||
import threading
|
||||
from unittest import mock
|
||||
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.test import test_utils
|
||||
from mediapipe.tasks.python.vision import image_segmenter
|
||||
from mediapipe.tasks.python.vision.core import image as image_module
|
||||
from mediapipe.tasks.python.vision.core import vision_task_running_mode
|
||||
|
||||
ImageSegmenterResult = image_segmenter.ImageSegmenterResult
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
_Image = image_module.Image
|
||||
_ImageFormat = image_module.ImageFormat
|
||||
_ImageSegmenter = image_segmenter.ImageSegmenter
|
||||
_ImageSegmenterOptions = image_segmenter.ImageSegmenterOptions
|
||||
_RUNNING_MODE = vision_task_running_mode.VisionTaskRunningMode
|
||||
|
||||
_MODEL_FILE = 'deeplabv3.tflite'
|
||||
_IMAGE_FILE = 'segmentation_input_rotation0.jpg'
|
||||
_SEGMENTATION_FILE = 'segmentation_golden_rotation0.png'
|
||||
_CAT_IMAGE = 'cat.jpg'
|
||||
_CAT_MASK = 'cat_mask.jpg'
|
||||
_MASK_MAGNIFICATION_FACTOR = 10
|
||||
_MASK_SIMILARITY_THRESHOLD = 0.98
|
||||
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision'
|
||||
_EXPECTED_LABELS = [
|
||||
'background',
|
||||
'aeroplane',
|
||||
'bicycle',
|
||||
'bird',
|
||||
'boat',
|
||||
'bottle',
|
||||
'bus',
|
||||
'car',
|
||||
'cat',
|
||||
'chair',
|
||||
'cow',
|
||||
'dining table',
|
||||
'dog',
|
||||
'horse',
|
||||
'motorbike',
|
||||
'person',
|
||||
'potted plant',
|
||||
'sheep',
|
||||
'sofa',
|
||||
'train',
|
||||
'tv',
|
||||
]
|
||||
|
||||
|
||||
def _calculate_soft_iou(m1, m2):
|
||||
intersection_sum = np.sum(m1 * m2)
|
||||
union_sum = np.sum(m1 * m1) + np.sum(m2 * m2) - intersection_sum
|
||||
|
||||
if union_sum > 0:
|
||||
return intersection_sum / union_sum
|
||||
else:
|
||||
return 0
|
||||
|
||||
|
||||
def _similar_to_float_mask(actual_mask, expected_mask, similarity_threshold):
|
||||
if not actual_mask:
|
||||
raise ValueError('Result mask used for comparison was None')
|
||||
actual_mask = actual_mask.numpy_view()
|
||||
expected_mask = expected_mask.numpy_view() / 255.0
|
||||
|
||||
return (
|
||||
actual_mask.shape == expected_mask.shape
|
||||
and _calculate_soft_iou(actual_mask, expected_mask) > similarity_threshold
|
||||
)
|
||||
|
||||
|
||||
def _similar_to_uint8_mask(actual_mask, expected_mask):
|
||||
if not actual_mask:
|
||||
raise ValueError('Result mask used for comparison was None')
|
||||
actual_mask_pixels = actual_mask.numpy_view().flatten()
|
||||
expected_mask_pixels = expected_mask.numpy_view().flatten()
|
||||
|
||||
consistent_pixels = 0
|
||||
num_pixels = len(expected_mask_pixels)
|
||||
|
||||
for index in range(num_pixels):
|
||||
consistent_pixels += (
|
||||
actual_mask_pixels[index] * _MASK_MAGNIFICATION_FACTOR
|
||||
== expected_mask_pixels[index]
|
||||
)
|
||||
|
||||
return consistent_pixels / num_pixels >= _MASK_SIMILARITY_THRESHOLD
|
||||
|
||||
|
||||
class ModelFileType(enum.Enum):
|
||||
FILE_CONTENT = 1
|
||||
FILE_NAME = 2
|
||||
|
||||
|
||||
class ImageSegmenterTest(parameterized.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
# Load the test input image.
|
||||
self.test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, _IMAGE_FILE))
|
||||
)
|
||||
# Loads ground truth segmentation file.
|
||||
gt_segmentation_data = cv2.imread(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _SEGMENTATION_FILE)
|
||||
),
|
||||
cv2.IMREAD_GRAYSCALE,
|
||||
)
|
||||
self.test_seg_image = _Image(_ImageFormat.GRAY8, gt_segmentation_data)
|
||||
self.model_path = test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _MODEL_FILE)
|
||||
)
|
||||
|
||||
def _load_segmentation_mask(self, file_path: str):
|
||||
# Loads ground truth segmentation file.
|
||||
gt_segmentation_data = cv2.imread(
|
||||
test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, file_path)),
|
||||
cv2.IMREAD_GRAYSCALE,
|
||||
)
|
||||
return _Image(_ImageFormat.GRAY8, gt_segmentation_data)
|
||||
|
||||
def test_create_from_file_succeeds_with_valid_model_path(self):
|
||||
# Creates with default option and valid model file successfully.
|
||||
with _ImageSegmenter.create_from_model_path(self.model_path) as segmenter:
|
||||
self.assertIsInstance(segmenter, _ImageSegmenter)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_path(self):
|
||||
# Creates with options containing model file successfully.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _ImageSegmenterOptions(base_options=base_options)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
self.assertIsInstance(segmenter, _ImageSegmenter)
|
||||
|
||||
def test_create_from_options_fails_with_invalid_model_path(self):
|
||||
with self.assertRaisesRegex(
|
||||
FileNotFoundError,
|
||||
'Unable to open file at /path/to/invalid/model.tflite',
|
||||
):
|
||||
base_options = _BaseOptions(
|
||||
model_asset_path='/path/to/invalid/model.tflite'
|
||||
)
|
||||
options = _ImageSegmenterOptions(base_options=base_options)
|
||||
segmenter = _ImageSegmenter.create_from_options(options)
|
||||
segmenter.close()
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_content(self):
|
||||
# Creates with options containing model content successfully.
|
||||
with open(self.model_path, 'rb') as f:
|
||||
base_options = _BaseOptions(model_asset_buffer=f.read())
|
||||
options = _ImageSegmenterOptions(base_options=base_options)
|
||||
segmenter = _ImageSegmenter.create_from_options(options)
|
||||
self.assertIsInstance(segmenter, _ImageSegmenter)
|
||||
segmenter.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(ModelFileType.FILE_NAME,), (ModelFileType.FILE_CONTENT,)
|
||||
)
|
||||
def test_segment_succeeds_with_category_mask(self, model_file_type):
|
||||
# Creates segmenter.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=base_options,
|
||||
output_category_mask=True,
|
||||
output_confidence_masks=False,
|
||||
)
|
||||
segmenter = _ImageSegmenter.create_from_options(options)
|
||||
|
||||
# Performs image segmentation on the input.
|
||||
segmentation_result = segmenter.segment(self.test_image)
|
||||
category_mask = segmentation_result.category_mask
|
||||
if not category_mask:
|
||||
raise ValueError('Result category mask was None')
|
||||
result_pixels = category_mask.numpy_view().flatten()
|
||||
|
||||
# Check if data type of `category_mask` is correct.
|
||||
self.assertEqual(result_pixels.dtype, np.uint8)
|
||||
|
||||
self.assertTrue(
|
||||
_similar_to_uint8_mask(category_mask, self.test_seg_image),
|
||||
(
|
||||
'Number of pixels in the candidate mask differing from that of the'
|
||||
f' ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.'
|
||||
),
|
||||
)
|
||||
|
||||
# Closes the segmenter explicitly when the segmenter is not used in
|
||||
# a context.
|
||||
segmenter.close()
|
||||
|
||||
def test_segment_succeeds_with_confidence_mask(self):
|
||||
# Creates segmenter.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
|
||||
# Load the cat image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, _CAT_IMAGE))
|
||||
)
|
||||
|
||||
# Run segmentation on the model in CONFIDENCE_MASK mode.
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=base_options,
|
||||
output_category_mask=False,
|
||||
output_confidence_masks=True,
|
||||
)
|
||||
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
segmentation_result = segmenter.segment(test_image)
|
||||
confidence_masks = segmentation_result.confidence_masks
|
||||
|
||||
# Check if confidence mask shape is correct.
|
||||
self.assertLen(
|
||||
confidence_masks,
|
||||
21,
|
||||
'Number of confidence masks must match with number of categories.',
|
||||
)
|
||||
|
||||
# Loads ground truth segmentation file.
|
||||
expected_mask = self._load_segmentation_mask(_CAT_MASK)
|
||||
|
||||
self.assertTrue(
|
||||
_similar_to_float_mask(
|
||||
confidence_masks[8], expected_mask, _MASK_SIMILARITY_THRESHOLD
|
||||
)
|
||||
)
|
||||
|
||||
@parameterized.parameters((True, False), (False, True))
|
||||
def test_labels_succeeds(self, output_category_mask, output_confidence_masks):
|
||||
expected_labels = _EXPECTED_LABELS
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=base_options,
|
||||
output_category_mask=output_category_mask,
|
||||
output_confidence_masks=output_confidence_masks,
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
# Performs image segmentation on the input.
|
||||
actual_labels = segmenter.labels
|
||||
self.assertListEqual(actual_labels, expected_labels)
|
||||
|
||||
def test_missing_result_callback(self):
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback must be provided'
|
||||
):
|
||||
with _ImageSegmenter.create_from_options(options) as unused_segmenter:
|
||||
pass
|
||||
|
||||
@parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO))
|
||||
def test_illegal_result_callback(self, running_mode):
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=running_mode,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback should not be provided'
|
||||
):
|
||||
with _ImageSegmenter.create_from_options(options) as unused_segmenter:
|
||||
pass
|
||||
|
||||
def test_calling_segment_for_video_in_image_mode(self):
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
with self.assertRaises(ValueError):
|
||||
segmenter.segment_for_video(self.test_image, 0)
|
||||
|
||||
def test_calling_segment_async_in_image_mode(self):
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
with self.assertRaises(ValueError):
|
||||
segmenter.segment_async(self.test_image, 0)
|
||||
|
||||
def test_calling_segment_in_video_mode(self):
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
with self.assertRaises(ValueError):
|
||||
segmenter.segment(self.test_image)
|
||||
|
||||
def test_calling_segment_async_in_video_mode(self):
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
with self.assertRaises(ValueError):
|
||||
segmenter.segment_async(self.test_image, 0)
|
||||
|
||||
def test_segment_for_video_with_out_of_order_timestamp(self):
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
unused_result = segmenter.segment_for_video(self.test_image, 1)
|
||||
with self.assertRaises(ValueError):
|
||||
segmenter.segment_for_video(self.test_image, 0)
|
||||
|
||||
def test_segment_for_video_in_category_mask_mode(self):
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
output_category_mask=True,
|
||||
output_confidence_masks=False,
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
for timestamp in range(0, 300, 30):
|
||||
segmentation_result = segmenter.segment_for_video(
|
||||
self.test_image, timestamp
|
||||
)
|
||||
category_mask = segmentation_result.category_mask
|
||||
self.assertTrue(
|
||||
_similar_to_uint8_mask(category_mask, self.test_seg_image),
|
||||
(
|
||||
'Number of pixels in the candidate mask differing from that of'
|
||||
f' the ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.'
|
||||
),
|
||||
)
|
||||
|
||||
def test_segment_for_video_in_confidence_mask_mode(self):
|
||||
# Load the cat image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, _CAT_IMAGE))
|
||||
)
|
||||
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
output_category_mask=False,
|
||||
output_confidence_masks=True,
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
for timestamp in range(0, 300, 30):
|
||||
segmentation_result = segmenter.segment_for_video(test_image, timestamp)
|
||||
confidence_masks = segmentation_result.confidence_masks
|
||||
|
||||
# Check if confidence mask shape is correct.
|
||||
self.assertLen(
|
||||
confidence_masks,
|
||||
21,
|
||||
'Number of confidence masks must match with number of categories.',
|
||||
)
|
||||
|
||||
# Loads ground truth segmentation file.
|
||||
expected_mask = self._load_segmentation_mask(_CAT_MASK)
|
||||
self.assertTrue(
|
||||
_similar_to_float_mask(
|
||||
confidence_masks[8], expected_mask, _MASK_SIMILARITY_THRESHOLD
|
||||
)
|
||||
)
|
||||
|
||||
def test_calling_segment_in_live_stream_mode(self):
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
with self.assertRaises(ValueError):
|
||||
segmenter.segment(self.test_image)
|
||||
|
||||
def test_calling_segment_for_video_in_live_stream_mode(self):
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
with self.assertRaises(ValueError):
|
||||
segmenter.segment_for_video(self.test_image, 0)
|
||||
|
||||
def test_segment_async_calls_with_illegal_timestamp(self):
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
segmenter.segment_async(self.test_image, 100)
|
||||
with self.assertRaises(ValueError):
|
||||
segmenter.segment_async(self.test_image, 0)
|
||||
|
||||
def test_segment_async_calls_in_category_mask_mode(self):
|
||||
observed_timestamp_ms = -1
|
||||
callback_event = threading.Event()
|
||||
callback_exception = None
|
||||
|
||||
def check_result(
|
||||
result: ImageSegmenterResult, output_image: _Image, timestamp_ms: int
|
||||
):
|
||||
nonlocal callback_exception, observed_timestamp_ms
|
||||
try:
|
||||
# Get the output category mask.
|
||||
category_mask = result.category_mask
|
||||
self.assertEqual(output_image.width, self.test_image.width)
|
||||
self.assertEqual(output_image.height, self.test_image.height)
|
||||
self.assertEqual(output_image.width, self.test_seg_image.width)
|
||||
self.assertEqual(output_image.height, self.test_seg_image.height)
|
||||
self.assertTrue(
|
||||
_similar_to_uint8_mask(category_mask, self.test_seg_image),
|
||||
(
|
||||
'Number of pixels in the candidate mask differing from that of'
|
||||
f' the ground truth mask exceeds {_MASK_SIMILARITY_THRESHOLD}.'
|
||||
),
|
||||
)
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
self.observed_timestamp_ms = timestamp_ms
|
||||
except AssertionError as e:
|
||||
callback_exception = e
|
||||
finally:
|
||||
callback_event.set()
|
||||
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
output_category_mask=True,
|
||||
output_confidence_masks=False,
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=check_result,
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
for timestamp in range(0, 300, 30):
|
||||
segmenter.segment_async(self.test_image, timestamp)
|
||||
callback_event.wait(3)
|
||||
if callback_exception is not None:
|
||||
raise callback_exception
|
||||
callback_event.clear()
|
||||
|
||||
def test_segment_async_calls_in_confidence_mask_mode(self):
|
||||
# Load the cat image.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, _CAT_IMAGE))
|
||||
)
|
||||
|
||||
# Loads ground truth segmentation file.
|
||||
expected_mask = self._load_segmentation_mask(_CAT_MASK)
|
||||
observed_timestamp_ms = -1
|
||||
callback_event = threading.Event()
|
||||
callback_exception = None
|
||||
|
||||
def check_result(
|
||||
result: ImageSegmenterResult, output_image: _Image, timestamp_ms: int
|
||||
):
|
||||
nonlocal callback_exception, observed_timestamp_ms
|
||||
try:
|
||||
# Get the output category mask.
|
||||
confidence_masks = result.confidence_masks
|
||||
|
||||
# Check if confidence mask shape is correct.
|
||||
self.assertLen(
|
||||
confidence_masks,
|
||||
21,
|
||||
'Number of confidence masks must match with number of categories.',
|
||||
)
|
||||
self.assertEqual(output_image.width, test_image.width)
|
||||
self.assertEqual(output_image.height, test_image.height)
|
||||
self.assertTrue(
|
||||
_similar_to_float_mask(
|
||||
confidence_masks[8], expected_mask, _MASK_SIMILARITY_THRESHOLD
|
||||
)
|
||||
)
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
self.observed_timestamp_ms = timestamp_ms
|
||||
except AssertionError as e:
|
||||
callback_exception = e
|
||||
finally:
|
||||
callback_event.set()
|
||||
|
||||
options = _ImageSegmenterOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
output_category_mask=False,
|
||||
output_confidence_masks=True,
|
||||
result_callback=check_result,
|
||||
)
|
||||
with _ImageSegmenter.create_from_options(options) as segmenter:
|
||||
for timestamp in range(0, 300, 30):
|
||||
segmenter.segment_async(test_image, timestamp)
|
||||
callback_event.wait(3)
|
||||
if callback_exception is not None:
|
||||
raise callback_exception
|
||||
callback_event.clear()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
@@ -0,0 +1,113 @@
|
||||
# Copyright 2025 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import numpy as np
|
||||
|
||||
from mediapipe.tasks.python.test import test_utils
|
||||
from mediapipe.tasks.python.vision.core import image
|
||||
|
||||
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision'
|
||||
_IMAGE_FILE = 'portrait.jpg'
|
||||
|
||||
|
||||
class ImageTest(parameterized.TestCase):
|
||||
|
||||
def test_create_from_numpy(self):
|
||||
np_array = np.zeros((3, 4, 3), dtype=np.uint8)
|
||||
img = image.Image(image.ImageFormat.SRGB, np_array)
|
||||
self.assertEqual(img.width, 4)
|
||||
self.assertEqual(img.height, 3)
|
||||
self.assertEqual(img.channels, 3)
|
||||
self.assertEqual(img.image_format, image.ImageFormat.SRGB)
|
||||
np.testing.assert_array_equal(img.numpy_view(), np_array)
|
||||
|
||||
def test_create_from_file(self):
|
||||
test_image_path = test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _IMAGE_FILE)
|
||||
)
|
||||
img = image.Image.create_from_file(test_image_path)
|
||||
# portrait.jpg is 820x1024, 3 channels (SRGB)
|
||||
self.assertEqual(img.width, 820)
|
||||
self.assertEqual(img.height, 1024)
|
||||
self.assertEqual(img.channels, 3)
|
||||
self.assertEqual(img.image_format, image.ImageFormat.SRGB)
|
||||
|
||||
def test_get_item_uint8(self):
|
||||
pixel_data = np.array(
|
||||
[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]], dtype=np.uint8
|
||||
)
|
||||
img = image.Image(image.ImageFormat.SRGB, pixel_data)
|
||||
self.assertEqual(img[0, 1, 1], 5) # row 0, col 1, channel 1
|
||||
self.assertEqual(img[1, 0, 2], 9) # row 1, col 0, channel 2
|
||||
|
||||
def test_get_item_uint8_grayscale(self):
|
||||
pixel_data = np.array([[[1], [2]], [[3], [4]]], dtype=np.uint8)
|
||||
img = image.Image(image.ImageFormat.GRAY8, pixel_data)
|
||||
self.assertEqual(img[0, 1], 2) # row 0, col 1
|
||||
self.assertEqual(img[1, 0], 3) # row 1, col 0
|
||||
|
||||
def test_get_item_uint16(self):
|
||||
pixel_data = np.array(
|
||||
[
|
||||
[[100, 200, 300], [400, 500, 600]],
|
||||
[[700, 800, 900], [1000, 1100, 1200]],
|
||||
],
|
||||
dtype=np.uint16,
|
||||
)
|
||||
img = image.Image(image.ImageFormat.SRGB48, pixel_data)
|
||||
self.assertEqual(img[0, 1, 1], 500) # row 0, col 1, channel 1
|
||||
self.assertEqual(img[1, 0, 0], 700) # row 1, col 0, channel 0
|
||||
|
||||
def test_get_item_float32(self):
|
||||
pixel_data = np.array(
|
||||
[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]], dtype=np.float32
|
||||
)
|
||||
img = image.Image(image.ImageFormat.VEC32F2, pixel_data)
|
||||
self.assertAlmostEqual(img[0, 1, 1], 4.0) # row 0, col 1, channel 1
|
||||
self.assertAlmostEqual(img[1, 0, 0], 5.0) # row 1, col 0, channel 0
|
||||
|
||||
def test_uses_gpu(self):
|
||||
pixel_data = np.array([[[1], [2]], [[3], [4]]], dtype=np.uint8)
|
||||
img = image.Image(image.ImageFormat.GRAY8, pixel_data)
|
||||
self.assertFalse(img.uses_gpu())
|
||||
|
||||
def test_is_contiguous(self):
|
||||
pixel_data = np.array([[[1], [2]], [[3], [4]]], dtype=np.uint8)
|
||||
img = image.Image(image.ImageFormat.GRAY8, pixel_data)
|
||||
self.assertFalse(img.is_contiguous())
|
||||
|
||||
def test_is_empty(self):
|
||||
pixel_data = np.array([[[1], [2]], [[3], [4]]], dtype=np.uint8)
|
||||
img = image.Image(image.ImageFormat.GRAY8, pixel_data)
|
||||
self.assertFalse(img.is_empty())
|
||||
|
||||
def test_is_aligned(self):
|
||||
pixel_data = np.array([[[1], [2]], [[3], [4]]], dtype=np.uint8)
|
||||
img = image.Image(image.ImageFormat.GRAY8, pixel_data)
|
||||
self.assertTrue(img.is_aligned(1))
|
||||
self.assertFalse(img.is_aligned(32))
|
||||
|
||||
def test_create_from_numpy_error(self):
|
||||
# SRGB expects 3 channels, but the array has 2.
|
||||
np_array = np.zeros((3, 4, 2), dtype=np.uint8)
|
||||
with self.assertRaisesRegex(ValueError, 'Pixel data size is too small'):
|
||||
image.Image(image.ImageFormat.SRGB, np_array)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
+346
@@ -0,0 +1,346 @@
|
||||
# Copyright 2023 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tests for interactive segmenter."""
|
||||
|
||||
import enum
|
||||
import os
|
||||
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from mediapipe.tasks.python.components.containers import keypoint as keypoint_module
|
||||
from mediapipe.tasks.python.components.containers import rect as rect_module
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.test import test_utils
|
||||
from mediapipe.tasks.python.vision import interactive_segmenter
|
||||
from mediapipe.tasks.python.vision.core import image as image_module
|
||||
from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module
|
||||
|
||||
InteractiveSegmenterResult = interactive_segmenter.InteractiveSegmenterResult
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
_Image = image_module.Image
|
||||
_ImageFormat = image_module.ImageFormat
|
||||
_NormalizedKeypoint = keypoint_module.NormalizedKeypoint
|
||||
_RectF = rect_module.RectF
|
||||
_InteractiveSegmenter = interactive_segmenter.InteractiveSegmenter
|
||||
_InteractiveSegmenterOptions = interactive_segmenter.InteractiveSegmenterOptions
|
||||
_RegionOfInterest = interactive_segmenter.RegionOfInterest
|
||||
_Format = interactive_segmenter.RegionOfInterest.Format
|
||||
_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions
|
||||
|
||||
_MODEL_FILE = 'ptm_512_hdt_ptm_woid.tflite'
|
||||
_CATS_AND_DOGS = 'cats_and_dogs.jpg'
|
||||
_CATS_AND_DOGS_MASK_DOG_1 = 'cats_and_dogs_mask_dog1.png'
|
||||
_CATS_AND_DOGS_MASK_DOG_2 = 'cats_and_dogs_mask_dog2.png'
|
||||
_MASK_MAGNIFICATION_FACTOR = 255
|
||||
_MASK_SIMILARITY_THRESHOLD = 0.96
|
||||
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision'
|
||||
|
||||
|
||||
def _calculate_soft_iou(m1, m2):
|
||||
intersection_sum = np.sum(m1 * m2)
|
||||
union_sum = np.sum(m1 * m1) + np.sum(m2 * m2) - intersection_sum
|
||||
|
||||
if union_sum > 0:
|
||||
return intersection_sum / union_sum
|
||||
else:
|
||||
return 0
|
||||
|
||||
|
||||
def _similar_to_float_mask(
|
||||
actual_mask: _Image, expected_mask: _Image, similarity_threshold: float
|
||||
):
|
||||
actual_mask = actual_mask.numpy_view()
|
||||
expected_mask = expected_mask.numpy_view() / 255.0
|
||||
|
||||
return (
|
||||
actual_mask.shape == expected_mask.shape
|
||||
and _calculate_soft_iou(actual_mask, expected_mask) > similarity_threshold
|
||||
)
|
||||
|
||||
|
||||
def _similar_to_uint8_mask(
|
||||
actual_mask: _Image, expected_mask: _Image, similarity_threshold: float
|
||||
):
|
||||
actual_mask_pixels = actual_mask.numpy_view().flatten()
|
||||
expected_mask_pixels = expected_mask.numpy_view().flatten()
|
||||
|
||||
consistent_pixels = 0
|
||||
num_pixels = len(expected_mask_pixels)
|
||||
|
||||
for index in range(num_pixels):
|
||||
consistent_pixels += (
|
||||
actual_mask_pixels[index] * _MASK_MAGNIFICATION_FACTOR
|
||||
== expected_mask_pixels[index]
|
||||
)
|
||||
|
||||
return consistent_pixels / num_pixels >= similarity_threshold
|
||||
|
||||
|
||||
class ModelFileType(enum.Enum):
|
||||
FILE_CONTENT = 1
|
||||
FILE_NAME = 2
|
||||
|
||||
|
||||
class InteractiveSegmenterTest(parameterized.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
# Load the test input image.
|
||||
self.test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _CATS_AND_DOGS)
|
||||
)
|
||||
)
|
||||
# Loads ground truth segmentation file.
|
||||
self.test_seg_image = self._load_segmentation_mask(
|
||||
_CATS_AND_DOGS_MASK_DOG_1
|
||||
)
|
||||
self.model_path = test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _MODEL_FILE)
|
||||
)
|
||||
|
||||
def _load_segmentation_mask(self, file_path: str):
|
||||
# Loads ground truth segmentation file.
|
||||
gt_segmentation_data = cv2.imread(
|
||||
test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, file_path)),
|
||||
cv2.IMREAD_GRAYSCALE,
|
||||
)
|
||||
return _Image(_ImageFormat.GRAY8, gt_segmentation_data)
|
||||
|
||||
def test_create_from_file_succeeds_with_valid_model_path(self):
|
||||
# Creates with default option and valid model file successfully.
|
||||
with _InteractiveSegmenter.create_from_model_path(
|
||||
self.model_path
|
||||
) as segmenter:
|
||||
self.assertIsInstance(segmenter, _InteractiveSegmenter)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_path(self):
|
||||
# Creates with options containing model file successfully.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _InteractiveSegmenterOptions(base_options=base_options)
|
||||
with _InteractiveSegmenter.create_from_options(options) as segmenter:
|
||||
self.assertIsInstance(segmenter, _InteractiveSegmenter)
|
||||
|
||||
def test_create_from_options_fails_with_invalid_model_path(self):
|
||||
with self.assertRaisesRegex(
|
||||
FileNotFoundError,
|
||||
'Unable to open file at /path/to/invalid/model.tflite',
|
||||
):
|
||||
base_options = _BaseOptions(
|
||||
model_asset_path='/path/to/invalid/model.tflite'
|
||||
)
|
||||
options = _InteractiveSegmenterOptions(base_options=base_options)
|
||||
segmenter = _InteractiveSegmenter.create_from_options(options)
|
||||
segmenter.close()
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_content(self):
|
||||
# Creates with options containing model content successfully.
|
||||
with open(self.model_path, 'rb') as f:
|
||||
base_options = _BaseOptions(model_asset_buffer=f.read())
|
||||
options = _InteractiveSegmenterOptions(base_options=base_options)
|
||||
segmenter = _InteractiveSegmenter.create_from_options(options)
|
||||
self.assertIsInstance(segmenter, _InteractiveSegmenter)
|
||||
segmenter.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
_RegionOfInterest.Format.KEYPOINT,
|
||||
_NormalizedKeypoint(0.44, 0.7),
|
||||
_CATS_AND_DOGS_MASK_DOG_1,
|
||||
0.84,
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
_RegionOfInterest.Format.KEYPOINT,
|
||||
_NormalizedKeypoint(0.44, 0.7),
|
||||
_CATS_AND_DOGS_MASK_DOG_1,
|
||||
0.84,
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
_RegionOfInterest.Format.KEYPOINT,
|
||||
_NormalizedKeypoint(0.66, 0.66),
|
||||
_CATS_AND_DOGS_MASK_DOG_2,
|
||||
_MASK_SIMILARITY_THRESHOLD,
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
_RegionOfInterest.Format.KEYPOINT,
|
||||
_NormalizedKeypoint(0.66, 0.66),
|
||||
_CATS_AND_DOGS_MASK_DOG_2,
|
||||
_MASK_SIMILARITY_THRESHOLD,
|
||||
),
|
||||
)
|
||||
def test_segment_succeeds_with_category_mask(
|
||||
self,
|
||||
model_file_type,
|
||||
roi_format,
|
||||
keypoint,
|
||||
output_mask,
|
||||
similarity_threshold,
|
||||
):
|
||||
# Creates segmenter.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _InteractiveSegmenterOptions(
|
||||
base_options=base_options,
|
||||
output_category_mask=True,
|
||||
output_confidence_masks=False,
|
||||
)
|
||||
segmenter = _InteractiveSegmenter.create_from_options(options)
|
||||
|
||||
# Performs image segmentation on the input.
|
||||
roi = _RegionOfInterest(format=roi_format, keypoint=keypoint)
|
||||
segmentation_result = segmenter.segment(self.test_image, roi)
|
||||
category_mask = segmentation_result.category_mask
|
||||
assert category_mask is not None, 'Category mask was None'
|
||||
result_pixels = category_mask.numpy_view().flatten()
|
||||
|
||||
# Check if data type of `category_mask` is correct.
|
||||
self.assertEqual(result_pixels.dtype, np.uint8)
|
||||
|
||||
# Loads ground truth segmentation file.
|
||||
test_seg_image = self._load_segmentation_mask(output_mask)
|
||||
|
||||
self.assertTrue(
|
||||
_similar_to_uint8_mask(
|
||||
category_mask, test_seg_image, similarity_threshold
|
||||
),
|
||||
(
|
||||
'Number of pixels in the candidate mask differing from that of the'
|
||||
f' ground truth mask exceeds {similarity_threshold}.'
|
||||
),
|
||||
)
|
||||
|
||||
# Closes the segmenter explicitly when the segmenter is not used in
|
||||
# a context.
|
||||
segmenter.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
_RegionOfInterest.Format.KEYPOINT,
|
||||
_NormalizedKeypoint(0.44, 0.7),
|
||||
_CATS_AND_DOGS_MASK_DOG_1,
|
||||
0.84,
|
||||
),
|
||||
(
|
||||
_RegionOfInterest.Format.KEYPOINT,
|
||||
_NormalizedKeypoint(0.66, 0.66),
|
||||
_CATS_AND_DOGS_MASK_DOG_2,
|
||||
0.84,
|
||||
),
|
||||
)
|
||||
def test_segment_succeeds_with_confidence_mask(
|
||||
self, roi_format, keypoint, output_mask, similarity_threshold
|
||||
):
|
||||
# Creates segmenter.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
roi = _RegionOfInterest(format=roi_format, keypoint=keypoint)
|
||||
|
||||
# Run segmentation on the model in CONFIDENCE_MASK mode.
|
||||
options = _InteractiveSegmenterOptions(
|
||||
base_options=base_options,
|
||||
output_category_mask=False,
|
||||
output_confidence_masks=True,
|
||||
)
|
||||
|
||||
with _InteractiveSegmenter.create_from_options(options) as segmenter:
|
||||
# Perform segmentation
|
||||
segmentation_result = segmenter.segment(self.test_image, roi)
|
||||
confidence_masks = segmentation_result.confidence_masks
|
||||
|
||||
# Check if confidence mask shape is correct.
|
||||
self.assertLen(
|
||||
confidence_masks,
|
||||
2,
|
||||
'Number of confidence masks must match with number of categories.',
|
||||
)
|
||||
|
||||
# Loads ground truth segmentation file.
|
||||
expected_mask = self._load_segmentation_mask(output_mask)
|
||||
|
||||
self.assertTrue(
|
||||
_similar_to_float_mask(
|
||||
confidence_masks[1], expected_mask, similarity_threshold
|
||||
)
|
||||
)
|
||||
|
||||
def test_segment_succeeds_with_rotation(self):
|
||||
# Creates segmenter.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
roi = _RegionOfInterest(
|
||||
format=_RegionOfInterest.Format.KEYPOINT,
|
||||
keypoint=_NormalizedKeypoint(0.66, 0.66),
|
||||
)
|
||||
|
||||
# Run segmentation on the model in CONFIDENCE_MASK mode.
|
||||
options = _InteractiveSegmenterOptions(
|
||||
base_options=base_options,
|
||||
output_category_mask=False,
|
||||
output_confidence_masks=True,
|
||||
)
|
||||
|
||||
with _InteractiveSegmenter.create_from_options(options) as segmenter:
|
||||
# Perform segmentation
|
||||
image_processing_options = _ImageProcessingOptions(rotation_degrees=-90)
|
||||
segmentation_result = segmenter.segment(
|
||||
self.test_image, roi, image_processing_options
|
||||
)
|
||||
confidence_masks = segmentation_result.confidence_masks
|
||||
|
||||
# Check if confidence mask shape is correct.
|
||||
self.assertLen(
|
||||
confidence_masks,
|
||||
2,
|
||||
'Number of confidence masks must match with number of categories.',
|
||||
)
|
||||
|
||||
def test_segment_fails_with_roi_in_image_processing_options(self):
|
||||
# Creates segmenter.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
roi = _RegionOfInterest(
|
||||
format=_RegionOfInterest.Format.KEYPOINT,
|
||||
keypoint=_NormalizedKeypoint(0.66, 0.66),
|
||||
)
|
||||
|
||||
# Run segmentation on the model in CONFIDENCE_MASK mode.
|
||||
options = _InteractiveSegmenterOptions(
|
||||
base_options=base_options,
|
||||
output_category_mask=False,
|
||||
output_confidence_masks=True,
|
||||
)
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
with _InteractiveSegmenter.create_from_options(options) as segmenter:
|
||||
# Perform segmentation
|
||||
image_processing_options = _ImageProcessingOptions(
|
||||
_RectF(left=0.1, top=0.0, right=0.9, bottom=1.0)
|
||||
)
|
||||
segmenter.segment(self.test_image, roi, image_processing_options)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
+520
@@ -0,0 +1,520 @@
|
||||
# Copyright 2022 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tests for object detector."""
|
||||
|
||||
import enum
|
||||
import os
|
||||
import threading
|
||||
from unittest import mock
|
||||
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import numpy as np
|
||||
|
||||
from mediapipe.tasks.python.components.containers import bounding_box as bounding_box_module
|
||||
from mediapipe.tasks.python.components.containers import category as category_module
|
||||
from mediapipe.tasks.python.components.containers import detections as detections_module
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.test import test_utils
|
||||
from mediapipe.tasks.python.vision import object_detector
|
||||
from mediapipe.tasks.python.vision.core import image as image_module
|
||||
from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module
|
||||
from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module
|
||||
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
_Category = category_module.Category
|
||||
_BoundingBox = bounding_box_module.BoundingBox
|
||||
_Detection = detections_module.Detection
|
||||
_DetectionResult = detections_module.DetectionResult
|
||||
_Image = image_module.Image
|
||||
_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions
|
||||
_ObjectDetector = object_detector.ObjectDetector
|
||||
_ObjectDetectorOptions = object_detector.ObjectDetectorOptions
|
||||
|
||||
_RUNNING_MODE = running_mode_module.VisionTaskRunningMode
|
||||
|
||||
_MODEL_FILE = 'coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.tflite'
|
||||
_NO_NMS_MODEL_FILE = 'efficientdet_lite0_fp16_no_nms.tflite'
|
||||
_IMAGE_FILE = 'cats_and_dogs.jpg'
|
||||
_EXPECTED_DETECTION_RESULT = _DetectionResult(
|
||||
detections=[
|
||||
_Detection(
|
||||
bounding_box=_BoundingBox(
|
||||
origin_x=608,
|
||||
origin_y=164,
|
||||
width=381,
|
||||
height=432,
|
||||
),
|
||||
categories=[
|
||||
_Category(
|
||||
index=None,
|
||||
score=0.69921875,
|
||||
display_name=None,
|
||||
category_name='cat',
|
||||
)
|
||||
],
|
||||
),
|
||||
_Detection(
|
||||
bounding_box=_BoundingBox(
|
||||
origin_x=57,
|
||||
origin_y=398,
|
||||
width=386,
|
||||
height=196,
|
||||
),
|
||||
categories=[
|
||||
_Category(
|
||||
index=None,
|
||||
score=0.65625,
|
||||
display_name=None,
|
||||
category_name='cat',
|
||||
)
|
||||
],
|
||||
),
|
||||
_Detection(
|
||||
bounding_box=_BoundingBox(
|
||||
origin_x=256,
|
||||
origin_y=394,
|
||||
width=173,
|
||||
height=202,
|
||||
),
|
||||
categories=[
|
||||
_Category(
|
||||
index=None,
|
||||
score=0.51171875,
|
||||
display_name=None,
|
||||
category_name='cat',
|
||||
)
|
||||
],
|
||||
),
|
||||
_Detection(
|
||||
bounding_box=_BoundingBox(
|
||||
origin_x=360,
|
||||
origin_y=195,
|
||||
width=330,
|
||||
height=412,
|
||||
),
|
||||
categories=[
|
||||
_Category(
|
||||
index=None,
|
||||
score=0.48828125,
|
||||
display_name=None,
|
||||
category_name='cat',
|
||||
)
|
||||
],
|
||||
),
|
||||
]
|
||||
)
|
||||
_ALLOW_LIST = ['cat', 'dog']
|
||||
_DENY_LIST = ['cat']
|
||||
_SCORE_THRESHOLD = 0.3
|
||||
_MAX_RESULTS = 3
|
||||
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/vision'
|
||||
|
||||
|
||||
class ModelFileType(enum.Enum):
|
||||
FILE_CONTENT = 1
|
||||
FILE_NAME = 2
|
||||
|
||||
|
||||
class ObjectDetectorTest(parameterized.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(os.path.join(_TEST_DATA_DIR, _IMAGE_FILE))
|
||||
)
|
||||
self.model_path = test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _MODEL_FILE)
|
||||
)
|
||||
|
||||
def test_create_from_file_succeeds_with_valid_model_path(self):
|
||||
# Creates with default option and valid model file successfully.
|
||||
with _ObjectDetector.create_from_model_path(self.model_path) as detector:
|
||||
self.assertIsInstance(detector, _ObjectDetector)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_path(self):
|
||||
# Creates with options containing model file successfully.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _ObjectDetectorOptions(base_options=base_options)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
self.assertIsInstance(detector, _ObjectDetector)
|
||||
|
||||
def test_create_from_options_fails_with_invalid_model_path(self):
|
||||
with self.assertRaisesRegex(
|
||||
FileNotFoundError,
|
||||
'Unable to open file at /path/to/invalid/model.tflite',
|
||||
):
|
||||
base_options = _BaseOptions(
|
||||
model_asset_path='/path/to/invalid/model.tflite'
|
||||
)
|
||||
options = _ObjectDetectorOptions(base_options=base_options)
|
||||
detector = _ObjectDetector.create_from_options(options)
|
||||
detector.close()
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_content(self):
|
||||
# Creates with options containing model content successfully.
|
||||
with open(self.model_path, 'rb') as f:
|
||||
base_options = _BaseOptions(model_asset_buffer=f.read())
|
||||
options = _ObjectDetectorOptions(base_options=base_options)
|
||||
detector = _ObjectDetector.create_from_options(options)
|
||||
self.assertIsInstance(detector, _ObjectDetector)
|
||||
detector.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(ModelFileType.FILE_NAME, 4, _EXPECTED_DETECTION_RESULT),
|
||||
(ModelFileType.FILE_CONTENT, 4, _EXPECTED_DETECTION_RESULT),
|
||||
)
|
||||
def test_detect(
|
||||
self, model_file_type, max_results, expected_detection_result
|
||||
):
|
||||
# Creates detector.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=base_options, max_results=max_results
|
||||
)
|
||||
detector = _ObjectDetector.create_from_options(options)
|
||||
|
||||
# Performs object detection on the input.
|
||||
detection_result = detector.detect(self.test_image)
|
||||
# Comparing results.
|
||||
self.assertEqual(detection_result, expected_detection_result)
|
||||
# Closes the detector explicitly when the detector is not used in
|
||||
# a context.
|
||||
detector.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(ModelFileType.FILE_NAME, 4, _EXPECTED_DETECTION_RESULT),
|
||||
(ModelFileType.FILE_CONTENT, 4, _EXPECTED_DETECTION_RESULT),
|
||||
)
|
||||
def test_detect_in_context(
|
||||
self, model_file_type, max_results, expected_detection_result
|
||||
):
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_contents = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_contents)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=base_options, max_results=max_results
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
# Performs object detection on the input.
|
||||
detection_result = detector.detect(self.test_image)
|
||||
# Comparing results.
|
||||
self.assertEqual(detection_result, expected_detection_result)
|
||||
|
||||
def test_score_threshold_option(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
score_threshold=_SCORE_THRESHOLD,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
# Performs object detection on the input.
|
||||
detection_result = detector.detect(self.test_image)
|
||||
detections = detection_result.detections
|
||||
|
||||
for detection in detections:
|
||||
score = detection.categories[0].score
|
||||
self.assertGreaterEqual(
|
||||
score,
|
||||
_SCORE_THRESHOLD,
|
||||
f'Detection with score lower than threshold found. {detection}',
|
||||
)
|
||||
|
||||
def test_max_results_option(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
max_results=_MAX_RESULTS,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
# Performs object detection on the input.
|
||||
detection_result = detector.detect(self.test_image)
|
||||
detections = detection_result.detections
|
||||
|
||||
self.assertLessEqual(
|
||||
len(detections), _MAX_RESULTS, 'Too many results returned.'
|
||||
)
|
||||
|
||||
def test_allow_list_option(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
category_allowlist=_ALLOW_LIST,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
# Performs object detection on the input.
|
||||
detection_result = detector.detect(self.test_image)
|
||||
detections = detection_result.detections
|
||||
|
||||
for detection in detections:
|
||||
label = detection.categories[0].category_name
|
||||
self.assertIn(
|
||||
label,
|
||||
_ALLOW_LIST,
|
||||
f'Label {label} found but not in label allow list',
|
||||
)
|
||||
|
||||
def test_deny_list_option(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
category_denylist=_DENY_LIST,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
# Performs object detection on the input.
|
||||
detection_result = detector.detect(self.test_image)
|
||||
detections = detection_result.detections
|
||||
|
||||
for detection in detections:
|
||||
label = detection.categories[0].category_name
|
||||
self.assertNotIn(
|
||||
label, _DENY_LIST, f'Label {label} found but in deny list.'
|
||||
)
|
||||
|
||||
def test_combined_allowlist_and_denylist(self):
|
||||
# Fails with combined allowlist and denylist
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
'`category_allowlist` and `category_denylist` are mutually exclusive'
|
||||
' options.',
|
||||
):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
category_allowlist=['foo'],
|
||||
category_denylist=['bar'],
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as unused_detector:
|
||||
pass
|
||||
|
||||
def test_empty_detection_outputs_with_in_model_nms(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
score_threshold=1,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
# Performs object detection on the input.
|
||||
detection_result = detector.detect(self.test_image)
|
||||
self.assertEmpty(detection_result.detections)
|
||||
|
||||
def test_empty_detection_outputs_without_in_model_nms(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(
|
||||
model_asset_path=test_utils.get_test_data_path(
|
||||
os.path.join(_TEST_DATA_DIR, _NO_NMS_MODEL_FILE)
|
||||
)
|
||||
),
|
||||
score_threshold=1,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
# Performs object detection on the input.
|
||||
detection_result = detector.detect(self.test_image)
|
||||
self.assertEmpty(detection_result.detections)
|
||||
|
||||
def test_missing_result_callback(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback must be provided'
|
||||
):
|
||||
with _ObjectDetector.create_from_options(options) as unused_detector:
|
||||
pass
|
||||
|
||||
@parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO))
|
||||
def test_illegal_result_callback(self, running_mode):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=running_mode,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback should not be provided'
|
||||
):
|
||||
with _ObjectDetector.create_from_options(options) as unused_detector:
|
||||
pass
|
||||
|
||||
# TODO: Change the errors to ValueError once we return MpStatus.
|
||||
def test_calling_detect_for_video_in_image_mode(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
'Task is not initialized with the video mode. Current running'
|
||||
' mode:image mode',
|
||||
):
|
||||
detector.detect_for_video(self.test_image, 0)
|
||||
|
||||
def test_calling_detect_async_in_image_mode(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
'Task is not initialized with the live stream mode. Current running'
|
||||
' mode:image mode',
|
||||
):
|
||||
detector.detect_async(self.test_image, 0)
|
||||
|
||||
def test_calling_detect_in_video_mode(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
'Task is not initialized with the image mode. Current running'
|
||||
' mode:video mode',
|
||||
):
|
||||
detector.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_async_in_video_mode(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
'Task is not initialized with the live stream mode. Current running'
|
||||
' mode:video mode',
|
||||
):
|
||||
detector.detect_async(self.test_image, 0)
|
||||
|
||||
def test_detect_for_video_with_out_of_order_timestamp(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
unused_result = detector.detect_for_video(self.test_image, 1)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, 'Input timestamp must be monotonically increasing.'
|
||||
):
|
||||
detector.detect_for_video(self.test_image, 0)
|
||||
|
||||
# TODO: Tests how `detect_for_video` handles the temporal data
|
||||
# with a real video.
|
||||
def test_detect_for_video(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
max_results=4,
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
for timestamp in range(0, 300, 30):
|
||||
detection_result = detector.detect_for_video(self.test_image, timestamp)
|
||||
self.assertEqual(detection_result, _EXPECTED_DETECTION_RESULT)
|
||||
|
||||
def test_calling_detect_in_live_stream_mode(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
'Task is not initialized with the image mode. Current running'
|
||||
' mode:live stream mode',
|
||||
):
|
||||
detector.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_for_video_in_live_stream_mode(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
'Task is not initialized with the video mode. Current running'
|
||||
' mode:live stream mode',
|
||||
):
|
||||
detector.detect_for_video(self.test_image, 0)
|
||||
|
||||
def test_detect_async_calls_with_illegal_timestamp(self):
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
max_results=4,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _ObjectDetector.create_from_options(options) as detector:
|
||||
detector.detect_async(self.test_image, 100)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, 'Input timestamp must be monotonically increasing.'
|
||||
):
|
||||
detector.detect_async(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(0, _EXPECTED_DETECTION_RESULT), (1, _DetectionResult(detections=[]))
|
||||
)
|
||||
def test_detect_async_calls(self, threshold, expected_result):
|
||||
callback_event = threading.Event()
|
||||
callback_exception: None | Exception = None
|
||||
observed_timestamp_ms = -1
|
||||
|
||||
def check_result(
|
||||
result: _DetectionResult, output_image: _Image, timestamp_ms: int
|
||||
):
|
||||
nonlocal callback_exception, observed_timestamp_ms
|
||||
try:
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
self.assertEqual(result, expected_result)
|
||||
self.assertTrue(
|
||||
np.array_equal(
|
||||
output_image.numpy_view(), self.test_image.numpy_view()
|
||||
)
|
||||
)
|
||||
observed_timestamp_ms = timestamp_ms
|
||||
except AssertionError as e:
|
||||
callback_exception = e
|
||||
finally:
|
||||
callback_event.set()
|
||||
|
||||
options = _ObjectDetectorOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
max_results=4,
|
||||
score_threshold=threshold,
|
||||
result_callback=check_result,
|
||||
)
|
||||
detector = _ObjectDetector.create_from_options(options)
|
||||
for timestamp in range(0, 300, 30):
|
||||
detector.detect_async(self.test_image, timestamp)
|
||||
detector.close()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
+525
@@ -0,0 +1,525 @@
|
||||
# Copyright 2023 The MediaPipe Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tests for pose landmarker."""
|
||||
|
||||
import enum
|
||||
import threading
|
||||
from typing import List
|
||||
from unittest import mock
|
||||
|
||||
from absl.testing import absltest
|
||||
from absl.testing import parameterized
|
||||
import numpy as np
|
||||
|
||||
from google.protobuf import text_format
|
||||
from mediapipe.tasks.cc.components.containers.proto import landmarks_detection_result_pb2
|
||||
from mediapipe.tasks.python.components.containers import landmark as landmark_module
|
||||
from mediapipe.tasks.python.components.containers import landmark_detection_result as landmark_detection_result_module
|
||||
from mediapipe.tasks.python.components.containers import rect as rect_module
|
||||
from mediapipe.tasks.python.core import base_options as base_options_module
|
||||
from mediapipe.tasks.python.test import test_utils
|
||||
from mediapipe.tasks.python.test.vision import proto_utils
|
||||
from mediapipe.tasks.python.vision import pose_landmarker
|
||||
from mediapipe.tasks.python.vision.core import image as image_module
|
||||
from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module
|
||||
from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module
|
||||
|
||||
PoseLandmarkerResult = pose_landmarker.PoseLandmarkerResult
|
||||
_LandmarksDetectionResultProto = (
|
||||
landmarks_detection_result_pb2.LandmarksDetectionResult
|
||||
)
|
||||
_BaseOptions = base_options_module.BaseOptions
|
||||
_RectF = rect_module.RectF
|
||||
_Landmark = landmark_module.Landmark
|
||||
_NormalizedLandmark = landmark_module.NormalizedLandmark
|
||||
_LandmarksDetectionResult = (
|
||||
landmark_detection_result_module.LandmarksDetectionResult
|
||||
)
|
||||
_Image = image_module.Image
|
||||
_PoseLandmarker = pose_landmarker.PoseLandmarker
|
||||
_PoseLandmarkerOptions = pose_landmarker.PoseLandmarkerOptions
|
||||
_RUNNING_MODE = running_mode_module.VisionTaskRunningMode
|
||||
_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions
|
||||
|
||||
_POSE_LANDMARKER_BUNDLE_ASSET_FILE = 'pose_landmarker.task'
|
||||
_BURGER_IMAGE = 'burger.jpg'
|
||||
_POSE_IMAGE = 'pose.jpg'
|
||||
_POSE_LANDMARKS = 'pose_landmarks.pbtxt'
|
||||
_LANDMARKS_MARGIN = 0.03
|
||||
|
||||
|
||||
def _get_expected_pose_landmarker_result(
|
||||
file_path: str,
|
||||
) -> PoseLandmarkerResult:
|
||||
landmarks_detection_result_file_path = test_utils.get_test_data_path(
|
||||
file_path
|
||||
)
|
||||
with open(landmarks_detection_result_file_path, 'rb') as f:
|
||||
landmarks_detection_result_proto = _LandmarksDetectionResultProto()
|
||||
# Use this if a .pb file is available.
|
||||
# landmarks_detection_result_proto.ParseFromString(f.read())
|
||||
text_format.Parse(f.read(), landmarks_detection_result_proto)
|
||||
landmarks_detection_result = (
|
||||
proto_utils.create_landmarks_detection_result_from_proto(
|
||||
landmarks_detection_result_proto
|
||||
)
|
||||
)
|
||||
return PoseLandmarkerResult(
|
||||
pose_landmarks=[landmarks_detection_result.landmarks],
|
||||
pose_world_landmarks=[],
|
||||
)
|
||||
|
||||
|
||||
class ModelFileType(enum.Enum):
|
||||
FILE_CONTENT = 1
|
||||
FILE_NAME = 2
|
||||
|
||||
|
||||
class PoseLandmarkerTest(parameterized.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(_POSE_IMAGE)
|
||||
)
|
||||
self.model_path = test_utils.get_test_data_path(
|
||||
_POSE_LANDMARKER_BUNDLE_ASSET_FILE
|
||||
)
|
||||
|
||||
def _expect_pose_landmarks_correct(
|
||||
self, actual_landmarks, expected_landmarks, margin
|
||||
):
|
||||
# Expects to have the same number of poses detected.
|
||||
self.assertLen(actual_landmarks, len(expected_landmarks))
|
||||
|
||||
for i, _ in enumerate(actual_landmarks):
|
||||
for j, elem in enumerate(actual_landmarks[i]):
|
||||
self.assertAlmostEqual(elem.x, expected_landmarks[i][j].x, delta=margin)
|
||||
self.assertAlmostEqual(elem.y, expected_landmarks[i][j].y, delta=margin)
|
||||
|
||||
def _expect_pose_landmarker_results_correct(
|
||||
self,
|
||||
actual_result: PoseLandmarkerResult,
|
||||
expected_result: PoseLandmarkerResult,
|
||||
output_segmentation_masks: bool,
|
||||
margin: float,
|
||||
):
|
||||
self._expect_pose_landmarks_correct(
|
||||
actual_result.pose_landmarks, expected_result.pose_landmarks, margin
|
||||
)
|
||||
if output_segmentation_masks:
|
||||
self.assertIsInstance(actual_result.segmentation_masks, List)
|
||||
for _, mask in enumerate(actual_result.segmentation_masks):
|
||||
self.assertIsInstance(mask, _Image)
|
||||
else:
|
||||
self.assertIsNone(actual_result.segmentation_masks)
|
||||
|
||||
def test_create_from_file_succeeds_with_valid_model_path(self):
|
||||
# Creates with default option and valid model file successfully.
|
||||
with _PoseLandmarker.create_from_model_path(self.model_path) as landmarker:
|
||||
self.assertIsInstance(landmarker, _PoseLandmarker)
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_path(self):
|
||||
# Creates with options containing model file successfully.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _PoseLandmarkerOptions(base_options=base_options)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
self.assertIsInstance(landmarker, _PoseLandmarker)
|
||||
|
||||
def test_create_from_options_fails_with_invalid_model_path(self):
|
||||
# Invalid empty model path.
|
||||
with self.assertRaisesWithLiteralMatch(
|
||||
FileNotFoundError,
|
||||
'Unable to open file at /path/to/invalid/model.tflite',
|
||||
):
|
||||
base_options = _BaseOptions(
|
||||
model_asset_path='/path/to/invalid/model.tflite'
|
||||
)
|
||||
options = _PoseLandmarkerOptions(base_options=base_options)
|
||||
landmarker = _PoseLandmarker.create_from_options(options)
|
||||
landmarker.close()
|
||||
|
||||
def test_create_from_options_succeeds_with_valid_model_content(self):
|
||||
# Creates with options containing model content successfully.
|
||||
with open(self.model_path, 'rb') as f:
|
||||
base_options = _BaseOptions(model_asset_buffer=f.read())
|
||||
options = _PoseLandmarkerOptions(base_options=base_options)
|
||||
landmarker = _PoseLandmarker.create_from_options(options)
|
||||
self.assertIsInstance(landmarker, _PoseLandmarker)
|
||||
landmarker.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
False,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
False,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
True,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
True,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
)
|
||||
def test_detect(
|
||||
self,
|
||||
model_file_type,
|
||||
output_segmentation_masks,
|
||||
expected_detection_result,
|
||||
):
|
||||
# Creates pose landmarker.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=base_options,
|
||||
output_segmentation_masks=output_segmentation_masks,
|
||||
)
|
||||
landmarker = _PoseLandmarker.create_from_options(options)
|
||||
|
||||
# Performs pose landmarks detection on the input.
|
||||
detection_result = landmarker.detect(self.test_image)
|
||||
|
||||
# Comparing results.
|
||||
self._expect_pose_landmarker_results_correct(
|
||||
detection_result,
|
||||
expected_detection_result,
|
||||
output_segmentation_masks,
|
||||
_LANDMARKS_MARGIN,
|
||||
)
|
||||
# Closes the pose landmarker explicitly when the pose landmarker is not used
|
||||
# in a context.
|
||||
landmarker.close()
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
False,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
False,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_NAME,
|
||||
True,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
(
|
||||
ModelFileType.FILE_CONTENT,
|
||||
True,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
)
|
||||
def test_detect_in_context(
|
||||
self,
|
||||
model_file_type,
|
||||
output_segmentation_masks,
|
||||
expected_detection_result,
|
||||
):
|
||||
# Creates pose landmarker.
|
||||
if model_file_type is ModelFileType.FILE_NAME:
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
elif model_file_type is ModelFileType.FILE_CONTENT:
|
||||
with open(self.model_path, 'rb') as f:
|
||||
model_content = f.read()
|
||||
base_options = _BaseOptions(model_asset_buffer=model_content)
|
||||
else:
|
||||
# Should never happen
|
||||
raise ValueError('model_file_type is invalid.')
|
||||
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=base_options,
|
||||
output_segmentation_masks=output_segmentation_masks,
|
||||
)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
# Performs pose landmarks detection on the input.
|
||||
detection_result = landmarker.detect(self.test_image)
|
||||
|
||||
# Comparing results.
|
||||
self._expect_pose_landmarker_results_correct(
|
||||
detection_result,
|
||||
expected_detection_result,
|
||||
output_segmentation_masks,
|
||||
_LANDMARKS_MARGIN,
|
||||
)
|
||||
|
||||
def test_detect_fails_with_region_of_interest(self):
|
||||
# Creates pose landmarker.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _PoseLandmarkerOptions(base_options=base_options)
|
||||
with self.assertRaisesWithLiteralMatch(
|
||||
ValueError, "This task doesn't support region-of-interest."
|
||||
):
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
# Set the `region_of_interest` parameter using `ImageProcessingOptions`.
|
||||
image_processing_options = _ImageProcessingOptions(
|
||||
region_of_interest=_RectF(0.0, 0.0, 1.0, 1.0)
|
||||
)
|
||||
# Attempt to perform pose landmarks detection on the cropped input.
|
||||
landmarker.detect(self.test_image, image_processing_options)
|
||||
|
||||
def test_empty_detection_outputs(self):
|
||||
# Creates pose landmarker.
|
||||
base_options = _BaseOptions(model_asset_path=self.model_path)
|
||||
options = _PoseLandmarkerOptions(base_options=base_options)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
# Load an image with no poses.
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(_BURGER_IMAGE)
|
||||
)
|
||||
# Performs pose landmarks detection on the input.
|
||||
detection_result = landmarker.detect(test_image)
|
||||
# Comparing results.
|
||||
self.assertEmpty(detection_result.pose_landmarks)
|
||||
self.assertEmpty(detection_result.pose_world_landmarks)
|
||||
|
||||
def test_missing_result_callback(self):
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
)
|
||||
with self.assertRaises(ValueError):
|
||||
with _PoseLandmarker.create_from_options(options) as unused_landmarker:
|
||||
pass
|
||||
|
||||
@parameterized.parameters((_RUNNING_MODE.IMAGE), (_RUNNING_MODE.VIDEO))
|
||||
def test_illegal_result_callback(self, running_mode):
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=running_mode,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with self.assertRaisesRegex(
|
||||
ValueError, r'result callback should not be provided'
|
||||
):
|
||||
with _PoseLandmarker.create_from_options(options) as unused_landmarker:
|
||||
pass
|
||||
|
||||
def test_calling_detect_for_video_in_image_mode(self):
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
def test_calling_detect_async_in_image_mode(self):
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.IMAGE,
|
||||
)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
def test_calling_detect_in_video_mode(self):
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_async_in_video_mode(self):
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
def test_detect_for_video_with_out_of_order_timestamp(self):
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
unused_result = landmarker.detect_for_video(self.test_image, 1)
|
||||
with self.assertRaises(ValueError):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
_POSE_IMAGE,
|
||||
0,
|
||||
False,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
(
|
||||
_POSE_IMAGE,
|
||||
0,
|
||||
True,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
(_BURGER_IMAGE, 0, False, PoseLandmarkerResult([], [])),
|
||||
)
|
||||
def test_detect_for_video(
|
||||
self, image_path, rotation, output_segmentation_masks, expected_result
|
||||
):
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(image_path)
|
||||
)
|
||||
# Set rotation parameters using ImageProcessingOptions.
|
||||
image_processing_options = _ImageProcessingOptions(
|
||||
rotation_degrees=rotation
|
||||
)
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
output_segmentation_masks=output_segmentation_masks,
|
||||
running_mode=_RUNNING_MODE.VIDEO,
|
||||
)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
for timestamp in range(0, 300, 30):
|
||||
result = landmarker.detect_for_video(
|
||||
test_image, timestamp, image_processing_options
|
||||
)
|
||||
if result.pose_landmarks:
|
||||
self._expect_pose_landmarker_results_correct(
|
||||
result,
|
||||
expected_result,
|
||||
output_segmentation_masks,
|
||||
_LANDMARKS_MARGIN,
|
||||
)
|
||||
else:
|
||||
self.assertEqual(result, expected_result)
|
||||
|
||||
def test_calling_detect_in_live_stream_mode(self):
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesWithLiteralMatch(
|
||||
ValueError,
|
||||
'Task is not initialized with the image mode. Current running'
|
||||
' mode:live stream mode',
|
||||
):
|
||||
landmarker.detect(self.test_image)
|
||||
|
||||
def test_calling_detect_for_video_in_live_stream_mode(self):
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
with self.assertRaisesWithLiteralMatch(
|
||||
ValueError,
|
||||
'Task is not initialized with the video mode. Current running'
|
||||
' mode:live stream mode',
|
||||
):
|
||||
landmarker.detect_for_video(self.test_image, 0)
|
||||
|
||||
def test_detect_async_calls_with_illegal_timestamp(self):
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=mock.MagicMock(),
|
||||
)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
landmarker.detect_async(self.test_image, 100)
|
||||
with self.assertRaisesWithLiteralMatch(
|
||||
ValueError, 'Input timestamp must be monotonically increasing.'
|
||||
):
|
||||
landmarker.detect_async(self.test_image, 0)
|
||||
|
||||
@parameterized.parameters(
|
||||
(
|
||||
_POSE_IMAGE,
|
||||
0,
|
||||
False,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
(
|
||||
_POSE_IMAGE,
|
||||
0,
|
||||
True,
|
||||
_get_expected_pose_landmarker_result(_POSE_LANDMARKS),
|
||||
),
|
||||
(_BURGER_IMAGE, 0, False, PoseLandmarkerResult([], [])),
|
||||
)
|
||||
def test_detect_async_calls(
|
||||
self, image_path, rotation, output_segmentation_masks, expected_result
|
||||
):
|
||||
test_image = _Image.create_from_file(
|
||||
test_utils.get_test_data_path(image_path)
|
||||
)
|
||||
callback_event = threading.Event()
|
||||
callback_exception: Exception | None
|
||||
# Set rotation parameters using ImageProcessingOptions.
|
||||
image_processing_options = _ImageProcessingOptions(
|
||||
rotation_degrees=rotation
|
||||
)
|
||||
observed_timestamp_ms = -1
|
||||
|
||||
def check_result(
|
||||
result: PoseLandmarkerResult, output_image: _Image, timestamp_ms: int
|
||||
):
|
||||
nonlocal callback_exception, observed_timestamp_ms
|
||||
try:
|
||||
if result.pose_landmarks:
|
||||
self._expect_pose_landmarker_results_correct(
|
||||
result,
|
||||
expected_result,
|
||||
output_segmentation_masks,
|
||||
_LANDMARKS_MARGIN,
|
||||
)
|
||||
else:
|
||||
self.assertEqual(result, expected_result)
|
||||
self.assertTrue(
|
||||
np.array_equal(output_image.numpy_view(), test_image.numpy_view())
|
||||
)
|
||||
self.assertLess(observed_timestamp_ms, timestamp_ms)
|
||||
self.observed_timestamp_ms = timestamp_ms
|
||||
except AssertionError as e:
|
||||
callback_exception = e
|
||||
finally:
|
||||
callback_event.set()
|
||||
|
||||
options = _PoseLandmarkerOptions(
|
||||
base_options=_BaseOptions(model_asset_path=self.model_path),
|
||||
output_segmentation_masks=output_segmentation_masks,
|
||||
running_mode=_RUNNING_MODE.LIVE_STREAM,
|
||||
result_callback=check_result,
|
||||
)
|
||||
with _PoseLandmarker.create_from_options(options) as landmarker:
|
||||
for timestamp in range(0, 300, 30):
|
||||
landmarker.detect_async(test_image, timestamp, image_processing_options)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
absltest.main()
|
||||
@@ -0,0 +1,182 @@
|
||||
# Copyright 2025 The MediaPipe Authors.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Proto loading utilities for vision tests."""
|
||||
|
||||
from typing import List
|
||||
|
||||
from mediapipe.framework.formats import classification_pb2
|
||||
from mediapipe.framework.formats import detection_pb2
|
||||
from mediapipe.framework.formats import landmark_pb2
|
||||
from mediapipe.framework.formats import location_data_pb2
|
||||
from mediapipe.tasks.cc.components.containers.proto import landmarks_detection_result_pb2
|
||||
from mediapipe.tasks.python.components.containers import bounding_box as bounding_box_module
|
||||
from mediapipe.tasks.python.components.containers import category as category_module
|
||||
from mediapipe.tasks.python.components.containers import detections as detections_module
|
||||
from mediapipe.tasks.python.components.containers import keypoint as keypoint_module
|
||||
from mediapipe.tasks.python.components.containers import landmark as landmark_module
|
||||
from mediapipe.tasks.python.components.containers import landmark_detection_result as landmark_detection_result_module
|
||||
from mediapipe.tasks.python.components.containers import rect as rect_module
|
||||
from mediapipe.tasks.python.vision import hand_landmarker
|
||||
from mediapipe.tasks.python.vision import pose_landmarker
|
||||
|
||||
PoseLandmarkerResult = pose_landmarker.PoseLandmarkerResult
|
||||
HandLandmarkerResult = hand_landmarker.HandLandmarkerResult
|
||||
_LandmarksDetectionResultProto = (
|
||||
landmarks_detection_result_pb2.LandmarksDetectionResult
|
||||
)
|
||||
_NormalizedRect = rect_module.NormalizedRect
|
||||
|
||||
|
||||
def create_bounding_box_from_proto(
|
||||
pb2_obj: location_data_pb2.LocationData.BoundingBox,
|
||||
) -> bounding_box_module.BoundingBox:
|
||||
return bounding_box_module.BoundingBox(
|
||||
origin_x=pb2_obj.xmin,
|
||||
origin_y=pb2_obj.ymin,
|
||||
width=pb2_obj.width,
|
||||
height=pb2_obj.height,
|
||||
)
|
||||
|
||||
|
||||
def create_category_from_proto(
|
||||
pb2_obj: classification_pb2.Classification,
|
||||
) -> category_module.Category:
|
||||
return category_module.Category(
|
||||
index=pb2_obj.index,
|
||||
score=pb2_obj.score,
|
||||
display_name=pb2_obj.display_name,
|
||||
category_name=pb2_obj.label,
|
||||
)
|
||||
|
||||
|
||||
def create_normalized_keypoint_from_proto(
|
||||
pb2_obj: location_data_pb2.LocationData.RelativeKeypoint,
|
||||
) -> keypoint_module.NormalizedKeypoint:
|
||||
return keypoint_module.NormalizedKeypoint(
|
||||
x=pb2_obj.x,
|
||||
y=pb2_obj.y,
|
||||
label=pb2_obj.keypoint_label,
|
||||
score=pb2_obj.score,
|
||||
)
|
||||
|
||||
|
||||
def create_detection_from_proto(
|
||||
pb2_obj: detection_pb2.Detection,
|
||||
) -> detections_module.Detection:
|
||||
"""Creates a `Detection` from the given protobuf object."""
|
||||
categories = [
|
||||
create_category_from_proto(
|
||||
classification_pb2.Classification(
|
||||
score=score,
|
||||
index=(
|
||||
pb2_obj.label_id[idx] if idx < len(pb2_obj.label_id) else None
|
||||
),
|
||||
label=pb2_obj.label[idx] if idx < len(pb2_obj.label) else None,
|
||||
display_name=(
|
||||
pb2_obj.display_name[idx]
|
||||
if idx < len(pb2_obj.display_name)
|
||||
else None
|
||||
),
|
||||
)
|
||||
)
|
||||
for idx, score in enumerate(pb2_obj.score)
|
||||
]
|
||||
|
||||
if pb2_obj.location_data.relative_keypoints:
|
||||
keypoints = [
|
||||
create_normalized_keypoint_from_proto(elem)
|
||||
for elem in pb2_obj.location_data.relative_keypoints
|
||||
]
|
||||
else:
|
||||
keypoints = []
|
||||
|
||||
return detections_module.Detection(
|
||||
bounding_box=create_bounding_box_from_proto(
|
||||
pb2_obj.location_data.bounding_box
|
||||
),
|
||||
categories=categories,
|
||||
keypoints=keypoints,
|
||||
)
|
||||
|
||||
|
||||
def create_landmark_from_proto(
|
||||
pb2_obj: landmark_pb2.Landmark,
|
||||
) -> landmark_module.Landmark:
|
||||
return landmark_module.Landmark(
|
||||
x=pb2_obj.x,
|
||||
y=pb2_obj.y,
|
||||
z=pb2_obj.z,
|
||||
visibility=pb2_obj.visibility if pb2_obj.HasField('visibility') else None,
|
||||
presence=pb2_obj.presence if pb2_obj.HasField('presence') else None,
|
||||
)
|
||||
|
||||
|
||||
def create_normalized_landmark_from_proto(
|
||||
pb2_obj: landmark_pb2.NormalizedLandmark,
|
||||
) -> landmark_module.NormalizedLandmark:
|
||||
return landmark_module.NormalizedLandmark(
|
||||
x=pb2_obj.x,
|
||||
y=pb2_obj.y,
|
||||
z=pb2_obj.z,
|
||||
visibility=pb2_obj.visibility if pb2_obj.HasField('visibility') else None,
|
||||
presence=pb2_obj.presence if pb2_obj.HasField('presence') else None,
|
||||
)
|
||||
|
||||
|
||||
def create_landmarks_detection_result_from_proto(
|
||||
pb2_obj: landmarks_detection_result_pb2.LandmarksDetectionResult,
|
||||
) -> landmark_detection_result_module.LandmarksDetectionResult:
|
||||
"""Creates a `LandmarksDetectionResult` from the given protobuf object."""
|
||||
if pb2_obj.HasField('landmarks'):
|
||||
landmarks = [
|
||||
create_normalized_landmark_from_proto(landmark_proto)
|
||||
for landmark_proto in pb2_obj.landmarks.landmark
|
||||
]
|
||||
else:
|
||||
landmarks = []
|
||||
|
||||
if pb2_obj.HasField('classifications'):
|
||||
categories = [
|
||||
create_category_from_proto(category_proto)
|
||||
for category_proto in pb2_obj.classifications.classification
|
||||
]
|
||||
else:
|
||||
categories = []
|
||||
|
||||
return landmark_detection_result_module.LandmarksDetectionResult(
|
||||
landmarks=landmarks,
|
||||
world_landmarks=[],
|
||||
categories=categories,
|
||||
rect=_NormalizedRect(x_center=0.0, y_center=0.0, width=0.0, height=0.0),
|
||||
)
|
||||
|
||||
|
||||
def create_normalized_landmark_list_from_proto(
|
||||
pb2_obj: landmark_pb2.NormalizedLandmarkList,
|
||||
) -> List[landmark_module.NormalizedLandmark]:
|
||||
"""Creates a list of `NormalizedLandmark` from the given protobuf object."""
|
||||
return [
|
||||
create_normalized_landmark_from_proto(landmark)
|
||||
for landmark in pb2_obj.landmark
|
||||
]
|
||||
|
||||
|
||||
def create_classification_list_from_proto(
|
||||
pb2_obj: classification_pb2.ClassificationList,
|
||||
) -> List[category_module.Category]:
|
||||
"""Creates a list of `Category` from the given protobuf object."""
|
||||
return [
|
||||
create_category_from_proto(classification)
|
||||
for classification in pb2_obj.classification
|
||||
]
|
||||
Reference in New Issue
Block a user