cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Removed an unnecessary check and updated tests to check if the masks are generated or not

Browse Source
This commit is contained in:
kinaryml 2023-04-24 11:23:27 -07:00
parent 9032bce577
commit b511822815
2 changed files with 55 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
mediapipe/tasks/python/test/vision/pose_landmarker_test.py
View File

@ -104,12 +104,19 @@ class PoseLandmarkerTest(parameterized.TestCase):
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 i, 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.
@ -141,12 +148,17 @@ class PoseLandmarkerTest(parameterized.TestCase):
self.assertIsInstance(landmarker, _PoseLandmarker)
@parameterized.parameters(
(ModelFileType.FILE_NAME,
(ModelFileType.FILE_NAME, False,
_get_expected_pose_landmarker_result(_POSE_LANDMARKS)),
(ModelFileType.FILE_CONTENT,
(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, expected_detection_result):
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)
@ -158,7 +170,10 @@ class PoseLandmarkerTest(parameterized.TestCase):
# Should never happen
raise ValueError('model_file_type is invalid.')
options = _PoseLandmarkerOptions(base_options=base_options)
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.
@ -166,19 +181,27 @@ class PoseLandmarkerTest(parameterized.TestCase):
# Comparing results.
self._expect_pose_landmarker_results_correct(
detection_result, expected_detection_result, _LANDMARKS_MARGIN
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,
(ModelFileType.FILE_NAME, False,
_get_expected_pose_landmarker_result(_POSE_LANDMARKS)),
(ModelFileType.FILE_CONTENT,
(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, expected_detection_result):
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)
@ -190,14 +213,18 @@ class PoseLandmarkerTest(parameterized.TestCase):
# Should never happen
raise ValueError('model_file_type is invalid.')
options = _PoseLandmarkerOptions(base_options=base_options)
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, _LANDMARKS_MARGIN
detection_result, expected_detection_result,
output_segmentation_masks, _LANDMARKS_MARGIN
)
def test_detect_fails_with_region_of_interest(self):
@ -295,12 +322,15 @@ class PoseLandmarkerTest(parameterized.TestCase):
landmarker.detect_for_video(self.test_image, 0)
@parameterized.parameters(
(_POSE_IMAGE, 0,
(_POSE_IMAGE, 0, False,
_get_expected_pose_landmarker_result(_POSE_LANDMARKS)),
(_BURGER_IMAGE, 0,
(_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, expected_result):
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.
@ -308,6 +338,7 @@ class PoseLandmarkerTest(parameterized.TestCase):
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):
@ -315,7 +346,8 @@ class PoseLandmarkerTest(parameterized.TestCase):
image_processing_options)
if result.pose_landmarks:
self._expect_pose_landmarker_results_correct(
result, expected_result, _LANDMARKS_MARGIN
result, expected_result, output_segmentation_masks,
_LANDMARKS_MARGIN
)
else:
self.assertEqual(result, expected_result)
@ -352,12 +384,15 @@ class PoseLandmarkerTest(parameterized.TestCase):
landmarker.detect_async(self.test_image, 0)
@parameterized.parameters(
(_POSE_IMAGE, 0,
(_POSE_IMAGE, 0, False,
_get_expected_pose_landmarker_result(_POSE_LANDMARKS)),
(_BURGER_IMAGE, 0,
(_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, expected_result):
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))
# Set rotation parameters using ImageProcessingOptions.
@ -369,7 +404,8 @@ class PoseLandmarkerTest(parameterized.TestCase):
timestamp_ms: int):
if result.pose_landmarks:
self._expect_pose_landmarker_results_correct(
result, expected_result, _LANDMARKS_MARGIN
result, expected_result, output_segmentation_masks,
_LANDMARKS_MARGIN
)
else:
self.assertEqual(result, expected_result)
@ -380,6 +416,7 @@ class PoseLandmarkerTest(parameterized.TestCase):
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:

1
mediapipe/tasks/python/vision/pose_landmarker.py
View File

@ -263,7 +263,6 @@ class PoseLandmarker(base_vision_task_api.BaseVisionTaskApi):
)
output_streams = [
':'.join([_SEGMENTATION_MASK_TAG, _SEGMENTATION_MASK_STREAM_NAME]),
':'.join([_NORM_LANDMARKS_TAG, _NORM_LANDMARKS_STREAM_NAME]),
':'.join([
_POSE_WORLD_LANDMARKS_TAG, _POSE_WORLD_LANDMARKS_STREAM_NAME