talksik/mediapipe
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 33adfdf31f3a5cbf9edc07ee1ea583e95080bdc5
This commit is contained in:
MediaPipe Team
2021-06-24 17:55:26 -04:00
committed by chuoling
parent b544a314b3
commit 139237092f
151 changed files with 4023 additions and 1001 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/python/BUILD
+3 -2
View File
@@ -59,11 +59,12 @@ cc_library(
"//mediapipe/calculators/core:gate_calculator",
"//mediapipe/calculators/core:pass_through_calculator",
"//mediapipe/calculators/core:side_packet_to_stream_calculator",
"//mediapipe/calculators/core:split_normalized_landmark_list_calculator",
"//mediapipe/calculators/core:split_landmarks_calculator",
"//mediapipe/calculators/core:string_to_int_calculator",
"//mediapipe/calculators/image:image_transformation_calculator",
"//mediapipe/calculators/util:detection_unique_id_calculator",
"//mediapipe/modules/face_detection:face_detection_front_cpu",
"//mediapipe/modules/face_detection:face_detection_full_range_cpu",
"//mediapipe/modules/face_detection:face_detection_short_range_cpu",
"//mediapipe/modules/face_landmark:face_landmark_front_cpu",
"//mediapipe/modules/hand_landmark:hand_landmark_tracking_cpu",
"//mediapipe/modules/holistic_landmark:holistic_landmark_cpu",
mediapipe/python/pybind/calculator_graph.cc
+6 -1
View File
@@ -399,7 +399,12 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
stream_name,
[callback_fn, stream_name](const Packet& packet) {
absl::MutexLock lock(&callback_mutex);
callback_fn(stream_name, packet);
py::gil_scoped_release gil_release;
{
// Acquires GIL before calling Python callback.
py::gil_scoped_acquire gil_acquire;
callback_fn(stream_name, packet);
}
return absl::OkStatus();
},
observe_timestamp_bounds));
mediapipe/python/pybind/packet_getter.cc
+1 -1
View File
@@ -397,7 +397,7 @@ void InternalPacketGetters(pybind11::module* m) {
"_get_serialized_proto",
[](const Packet& packet) {
// By default, py::bytes is an extra copy of the original std::string
// object: https://github.com/pybind/pybind11/issues/1236 Howeover, when
// object: https://github.com/pybind/pybind11/issues/1236 However, when
// Pybind11 performs the C++ to Python transition, it only increases the
// py::bytes object's ref count. See the implmentation at line 1583 in
// "pybind11/cast.h".
mediapipe/python/solutions/drawing_utils.py
+70
View File
@@ -19,6 +19,7 @@ from typing import List, Optional, Tuple, Union
import cv2
import dataclasses
import matplotlib.pyplot as plt
import numpy as np
from mediapipe.framework.formats import detection_pb2
@@ -27,6 +28,7 @@ from mediapipe.framework.formats import landmark_pb2
PRESENCE_THRESHOLD = 0.5
RGB_CHANNELS = 3
BLACK_COLOR = (0, 0, 0)
RED_COLOR = (0, 0, 255)
GREEN_COLOR = (0, 128, 0)
BLUE_COLOR = (255, 0, 0)
@@ -225,3 +227,71 @@ def draw_axis(
axis_drawing_spec.thickness)
cv2.arrowedLine(image, origin, z_axis, BLUE_COLOR,
axis_drawing_spec.thickness)
def _normalize_color(color):
return tuple(v / 255. for v in color)
def plot_landmarks(landmark_list: landmark_pb2.NormalizedLandmarkList,
connections: Optional[List[Tuple[int, int]]] = None,
landmark_drawing_spec: DrawingSpec = DrawingSpec(
color=RED_COLOR, thickness=5),
connection_drawing_spec: DrawingSpec = DrawingSpec(
color=BLACK_COLOR, thickness=5),
elevation: int = 10,
azimuth: int = 10):
"""Plot the landmarks and the connections in matplotlib 3d.
Args:
landmark_list: A normalized landmark list proto message to be plotted.
connections: A list of landmark index tuples that specifies how landmarks to
be connected.
landmark_drawing_spec: A DrawingSpec object that specifies the landmarks'
drawing settings such as color and line thickness.
connection_drawing_spec: A DrawingSpec object that specifies the
connections' drawing settings such as color and line thickness.
elevation: The elevation from which to view the plot.
azimuth: the azimuth angle to rotate the plot.
Raises:
ValueError: If any connetions contain invalid landmark index.
"""
if not landmark_list:
return
plt.figure(figsize=(10, 10))
ax = plt.axes(projection='3d')
ax.view_init(elev=elevation, azim=azimuth)
plotted_landmarks = {}
for idx, landmark in enumerate(landmark_list.landmark):
if ((landmark.HasField('visibility') and
landmark.visibility < VISIBILITY_THRESHOLD) or
(landmark.HasField('presence') and
landmark.presence < PRESENCE_THRESHOLD)):
continue
ax.scatter3D(
xs=[-landmark.z],
ys=[landmark.x],
zs=[-landmark.y],
color=_normalize_color(landmark_drawing_spec.color[::-1]),
linewidth=landmark_drawing_spec.thickness)
plotted_landmarks[idx] = (-landmark.z, landmark.x, -landmark.y)
if connections:
num_landmarks = len(landmark_list.landmark)
# Draws the connections if the start and end landmarks are both visible.
for connection in connections:
start_idx = connection[0]
end_idx = connection[1]
if not (0 <= start_idx < num_landmarks and 0 <= end_idx < num_landmarks):
raise ValueError(f'Landmark index is out of range. Invalid connection '
f'from landmark #{start_idx} to landmark #{end_idx}.')
if start_idx in plotted_landmarks and end_idx in plotted_landmarks:
landmark_pair = [
plotted_landmarks[start_idx], plotted_landmarks[end_idx]
]
ax.plot3D(
xs=[landmark_pair[0][0], landmark_pair[1][0]],
ys=[landmark_pair[0][1], landmark_pair[1][1]],
zs=[landmark_pair[0][2], landmark_pair[1][2]],
color=_normalize_color(connection_drawing_spec.color[::-1]),
linewidth=connection_drawing_spec.thickness)
plt.show()
mediapipe/python/solutions/face_detection.py
+13 -4
View File
@@ -28,7 +28,8 @@ from mediapipe.calculators.util import non_max_suppression_calculator_pb2
# pylint: enable=unused-import
from mediapipe.python.solution_base import SolutionBase
BINARYPB_FILE_PATH = 'mediapipe/modules/face_detection/face_detection_front_cpu.binarypb'
SHORT_RANGE_GRAPH_FILE_PATH = 'mediapipe/modules/face_detection/face_detection_short_range_cpu.binarypb'
FULL_RANGE_GRAPH_FILE_PATH = 'mediapipe/modules/face_detection/face_detection_full_range_cpu.binarypb'
def get_key_point(
@@ -69,18 +70,26 @@ class FaceDetection(SolutionBase):
for usage examples.
"""
def __init__(self, min_detection_confidence=0.5):
def __init__(self, min_detection_confidence=0.5, model_selection=0):
"""Initializes a MediaPipe Face Detection object.
Args:
min_detection_confidence: Minimum confidence value ([0.0, 1.0]) for face
detection to be considered successful. See details in
https://solutions.mediapipe.dev/face_detection#min_detection_confidence.
model_selection: 0 or 1. 0 to select a short-range model that works
best for faces within 2 meters from the camera, and 1 for a full-range
model best for faces within 5 meters. See details in
https://solutions.mediapipe.dev/face_detection#model_selection.
"""
binary_graph_path = FULL_RANGE_GRAPH_FILE_PATH if model_selection == 1 else SHORT_RANGE_GRAPH_FILE_PATH
subgraph_name = 'facedetectionfullrangecommon' if model_selection == 1 else 'facedetectionshortrangecommon'
super().__init__(
binary_graph_path=BINARYPB_FILE_PATH,
binary_graph_path=binary_graph_path,
calculator_params={
'facedetectionfrontcommon__TensorsToDetectionsCalculator.min_score_thresh':
subgraph_name + '__TensorsToDetectionsCalculator.min_score_thresh':
min_detection_confidence,
},
outputs=['detections'])
mediapipe/python/solutions/face_detection_test.py
+23 -10
View File
@@ -18,6 +18,7 @@ import tempfile # pylint: disable=unused-import
from typing import NamedTuple
from absl.testing import absltest
from absl.testing import parameterized
import cv2
import numpy as np
import numpy.testing as npt
@@ -28,12 +29,14 @@ from mediapipe.python.solutions import drawing_utils as mp_drawing
from mediapipe.python.solutions import face_detection as mp_faces
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
EXPECTED_FACE_KEY_POINTS = [[182, 363], [186, 460], [241, 420], [284, 417],
[199, 295], [198, 502]]
SHORT_RANGE_EXPECTED_FACE_KEY_POINTS = [[363, 182], [460, 186], [420, 241],
[417, 284], [295, 199], [502, 198]]
FULL_RANGE_EXPECTED_FACE_KEY_POINTS = [[363, 181], [455, 181], [413, 233],
[411, 278], [306, 204], [499, 207]]
DIFF_THRESHOLD = 5 # pixels
class FaceDetectionTest(absltest.TestCase):
class FaceDetectionTest(parameterized.TestCase):
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
for detection in results.detections:
@@ -55,20 +58,30 @@ class FaceDetectionTest(absltest.TestCase):
results = faces.process(image)
self.assertIsNone(results.detections)
def test_face(self):
@parameterized.named_parameters(('short_range_model', 0),
('full_range_model', 1))
def test_face(self, model_selection):
image_path = os.path.join(os.path.dirname(__file__),
'testdata/portrait.jpg')
image = cv2.imread(image_path)
with mp_faces.FaceDetection(min_detection_confidence=0.5) as faces:
rows, cols, _ = image.shape
with mp_faces.FaceDetection(
min_detection_confidence=0.5, model_selection=model_selection) as faces:
for idx in range(5):
results = faces.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._annotate(image.copy(), results, idx)
location_data = results.detections[0].location_data
x = [keypoint.x for keypoint in location_data.relative_keypoints]
y = [keypoint.y for keypoint in location_data.relative_keypoints]
face_keypoints = np.transpose(np.stack((y, x))) * image.shape[0:2]
prediction_error = np.abs(
np.asarray(face_keypoints) - np.asarray(EXPECTED_FACE_KEY_POINTS))
x = [keypoint.x * cols for keypoint in location_data.relative_keypoints]
y = [keypoint.y * rows for keypoint in location_data.relative_keypoints]
face_keypoints = np.column_stack((x, y))
if model_selection == 0:
prediction_error = np.abs(
np.asarray(face_keypoints) -
np.asarray(SHORT_RANGE_EXPECTED_FACE_KEY_POINTS))
else:
prediction_error = np.abs(
np.asarray(face_keypoints) -
np.asarray(FULL_RANGE_EXPECTED_FACE_KEY_POINTS))
self.assertLen(results.detections, 1)
self.assertLen(location_data.relative_keypoints, 6)
mediapipe/python/solutions/face_mesh.py
+1 -1
View File
@@ -213,7 +213,7 @@ class FaceMesh(SolutionBase):
.ConstantSidePacketCalculatorOptions.ConstantSidePacket(
bool_value=not static_image_mode)
],
'facedetectionfrontcpu__TensorsToDetectionsCalculator.min_score_thresh':
'facedetectionshortrangecpu__TensorsToDetectionsCalculator.min_score_thresh':
min_detection_confidence,
'facelandmarkcpu__ThresholdingCalculator.threshold':
min_tracking_confidence,
mediapipe/python/solutions/face_mesh_test.py
+36 -35
View File
@@ -32,38 +32,38 @@ from mediapipe.python.solutions import face_mesh as mp_faces
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
DIFF_THRESHOLD = 5 # pixels
EYE_INDICES_TO_LANDMARKS = {
33: [178, 345],
7: [179, 348],
163: [178, 352],
144: [179, 357],
145: [179, 365],
153: [179, 371],
154: [178, 378],
155: [177, 381],
133: [177, 383],
246: [175, 347],
161: [174, 350],
160: [172, 355],
159: [170, 362],
158: [171, 368],
157: [172, 375],
173: [175, 380],
263: [176, 467],
249: [177, 464],
390: [177, 460],
373: [178, 455],
374: [179, 448],
380: [179, 441],
381: [178, 435],
382: [177, 432],
362: [177, 430],
466: [175, 465],
388: [173, 462],
387: [171, 457],
386: [170, 450],
385: [171, 444],
384: [172, 437],
398: [175, 432]
33: [345, 178],
7: [348, 179],
163: [352, 178],
144: [357, 179],
145: [365, 179],
153: [371, 179],
154: [378, 178],
155: [381, 177],
133: [383, 177],
246: [347, 175],
161: [350, 174],
160: [355, 172],
159: [362, 170],
158: [368, 171],
157: [375, 172],
173: [380, 175],
263: [467, 176],
249: [464, 177],
390: [460, 177],
373: [455, 178],
374: [448, 179],
380: [441, 179],
381: [435, 178],
382: [432, 177],
362: [430, 177],
466: [465, 175],
388: [462, 173],
387: [457, 171],
386: [450, 170],
385: [444, 171],
384: [437, 172],
398: [432, 175]
}
@@ -99,6 +99,7 @@ class FaceMeshTest(parameterized.TestCase):
image_path = os.path.join(os.path.dirname(__file__),
'testdata/portrait.jpg')
image = cv2.imread(image_path)
rows, cols, _ = image.shape
with mp_faces.FaceMesh(
static_image_mode=static_image_mode,
min_detection_confidence=0.5) as faces:
@@ -108,9 +109,9 @@ class FaceMeshTest(parameterized.TestCase):
multi_face_landmarks = []
for landmarks in results.multi_face_landmarks:
self.assertLen(landmarks.landmark, 468)
x = [landmark.x for landmark in landmarks.landmark]
y = [landmark.y for landmark in landmarks.landmark]
face_landmarks = np.transpose(np.stack((y, x))) * image.shape[0:2]
x = [landmark.x * cols for landmark in landmarks.landmark]
y = [landmark.y * rows for landmark in landmarks.landmark]
face_landmarks = np.column_stack((x, y))
multi_face_landmarks.append(face_landmarks)
self.assertLen(multi_face_landmarks, 1)
# Verify the eye landmarks are correct as sanity check.
mediapipe/python/solutions/hands_test.py
+18 -17
View File
@@ -31,20 +31,20 @@ from mediapipe.python.solutions import hands as mp_hands
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
DIFF_THRESHOLD = 15 # pixels
EXPECTED_HAND_COORDINATES_PREDICTION = [[[345, 144], [323, 211], [286, 257],
[237, 289], [203, 322], [216, 219],
[138, 238], [90, 249], [51, 253],
[204, 177], [115, 184], [60, 187],
[19, 185], [208, 138], [127, 131],
[77, 124], [36, 117], [222, 106],
[159, 92], [124, 79], [93, 68]],
[[40, 577], [56, 504], [94, 459],
[146, 429], [182, 397], [167, 496],
[245, 479], [292, 469], [330, 464],
[177, 540], [265, 534], [319, 533],
[360, 536], [172, 581], [252, 587],
[304, 593], [346, 599], [157, 615],
[223, 628], [258, 638], [288, 648]]]
EXPECTED_HAND_COORDINATES_PREDICTION = [[[144, 345], [211, 323], [257, 286],
[289, 237], [322, 203], [219, 216],
[238, 138], [249, 90], [253, 51],
[177, 204], [184, 115], [187, 60],
[185, 19], [138, 208], [131, 127],
[124, 77], [117, 36], [106, 222],
[92, 159], [79, 124], [68, 93]],
[[577, 40], [504, 56], [459, 94],
[429, 146], [397, 182], [496, 167],
[479, 245], [469, 292], [464, 330],
[540, 177], [534, 265], [533, 319],
[536, 360], [581, 172], [587, 252],
[593, 304], [599, 346], [615, 157],
[628, 223], [638, 258], [648, 288]]]
class HandsTest(parameterized.TestCase):
@@ -88,11 +88,12 @@ class HandsTest(parameterized.TestCase):
for handedness in results.multi_handedness
]
multi_hand_coordinates = []
rows, cols, _ = image.shape
for landmarks in results.multi_hand_landmarks:
self.assertLen(landmarks.landmark, 21)
x = [landmark.x for landmark in landmarks.landmark]
y = [landmark.y for landmark in landmarks.landmark]
hand_coordinates = np.transpose(np.stack((y, x))) * image.shape[0:2]
x = [landmark.x * cols for landmark in landmarks.landmark]
y = [landmark.y * rows for landmark in landmarks.landmark]
hand_coordinates = np.column_stack((x, y))
multi_hand_coordinates.append(hand_coordinates)
self.assertLen(handedness, 2)
self.assertLen(multi_hand_coordinates, 2)
mediapipe/python/solutions/holistic.py
+14 -9
View File
@@ -116,8 +116,8 @@ class Holistic(SolutionBase):
min_tracking_confidence,
},
outputs=[
'pose_landmarks', 'left_hand_landmarks', 'right_hand_landmarks',
'face_landmarks'
'pose_landmarks', 'pose_world_landmarks', 'left_hand_landmarks',
'right_hand_landmarks', 'face_landmarks'
])
def process(self, image: np.ndarray) -> NamedTuple:
@@ -131,17 +131,22 @@ class Holistic(SolutionBase):
ValueError: If the input image is not three channel RGB.
Returns:
A NamedTuple that has four fields:
1) "pose_landmarks" field that contains the pose landmarks on the most
prominent person detected.
2) "left_hand_landmarks" and "right_hand_landmarks" fields that contain
the left and right hand landmarks of the most prominent person detected.
3) "face_landmarks" field that contains the face landmarks of the most
prominent person detected.
A NamedTuple that has five fields describing the landmarks on the most
prominate person detected:
1) "pose_landmarks" field that contains the pose landmarks.
2) "pose_world_landmarks" field that contains the pose landmarks in
real-world 3D coordinates that are in meters with the origin at the
center between hips.
3) "left_hand_landmarks" field that contains the left-hand landmarks.
4) "right_hand_landmarks" field that contains the right-hand landmarks.
5) "face_landmarks" field that contains the face landmarks.
"""
results = super().process(input_data={'image': image})
if results.pose_landmarks:
for landmark in results.pose_landmarks.landmark:
landmark.ClearField('presence')
if results.pose_world_landmarks:
for landmark in results.pose_world_landmarks.landmark:
landmark.ClearField('presence')
return results
mediapipe/python/solutions/pose.py
+10 -3
View File
@@ -185,7 +185,7 @@ class Pose(SolutionBase):
'poselandmarkcpu__poselandmarkbyroicpu__ThresholdingCalculator.threshold':
min_tracking_confidence,
},
outputs=['pose_landmarks'])
outputs=['pose_landmarks', 'pose_world_landmarks'])
def process(self, image: np.ndarray) -> NamedTuple:
"""Processes an RGB image and returns the pose landmarks on the most prominent person detected.
@@ -198,12 +198,19 @@ class Pose(SolutionBase):
ValueError: If the input image is not three channel RGB.
Returns:
A NamedTuple object with a "pose_landmarks" field that contains the pose
landmarks on the most prominent person detected.
A NamedTuple that has two fields describing the landmarks on the most
prominate person detected:
1) "pose_landmarks" field that contains the pose landmarks.
2) "pose_world_landmarks" field that contains the pose landmarks in
real-world 3D coordinates that are in meters with the origin at the
center between hips.
"""
results = super().process(input_data={'image': image})
if results.pose_landmarks:
for landmark in results.pose_landmarks.landmark:
landmark.ClearField('presence')
if results.pose_world_landmarks:
for landmark in results.pose_world_landmarks.landmark:
landmark.ClearField('presence')
return results
mediapipe/python/solutions/pose_test.py
+44 -6
View File
@@ -42,6 +42,20 @@ EXPECTED_POSE_LANDMARKS = np.array([[460, 283], [467, 273], [471, 273],
[467, 471], [612, 550], [358, 490],
[701, 613], [349, 611], [709, 624],
[363, 630], [730, 633], [303, 628]])
WORLD_DIFF_THRESHOLD = 0.2 # meters
EXPECTED_POSE_WORLD_LANDMARKS = np.array([
[-0.11, -0.59, -0.15], [-0.09, -0.64, -0.16], [-0.09, -0.64, -0.16],
[-0.09, -0.64, -0.16], [-0.11, -0.64, -0.14], [-0.11, -0.64, -0.14],
[-0.11, -0.64, -0.14], [0.01, -0.65, -0.15], [-0.06, -0.64, -0.05],
[-0.07, -0.57, -0.15], [-0.09, -0.57, -0.12], [0.18, -0.49, -0.09],
[-0.14, -0.5, -0.03], [0.41, -0.48, -0.11], [-0.42, -0.5, -0.02],
[0.64, -0.49, -0.17], [-0.63, -0.51, -0.13], [0.7, -0.5, -0.19],
[-0.71, -0.53, -0.15], [0.72, -0.51, -0.23], [-0.69, -0.54, -0.19],
[0.66, -0.49, -0.19], [-0.64, -0.52, -0.15], [0.09, 0., -0.04],
[-0.09, -0., 0.03], [0.41, 0.23, -0.09], [-0.43, 0.1, -0.11],
[0.69, 0.49, -0.04], [-0.48, 0.47, -0.02], [0.72, 0.52, -0.04],
[-0.48, 0.51, -0.02], [0.8, 0.5, -0.14], [-0.59, 0.52, -0.11],
])
class PoseTest(parameterized.TestCase):
@@ -51,6 +65,10 @@ class PoseTest(parameterized.TestCase):
return np.asarray([(lmk.x * cols, lmk.y * rows, lmk.z * cols)
for lmk in landmark_list.landmark])
def _world_landmarks_list_to_array(self, landmark_list):
return np.asarray([(lmk.x, lmk.y, lmk.z)
for lmk in landmark_list.landmark])
def _assert_diff_less(self, array1, array2, threshold):
npt.assert_array_less(np.abs(array1 - array2), threshold)
@@ -87,11 +105,15 @@ class PoseTest(parameterized.TestCase):
model_complexity=model_complexity) as pose:
for idx in range(num_frames):
results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
# TODO: Add rendering of world 3D when supported.
self._annotate(image.copy(), results, idx)
self._assert_diff_less(
self._landmarks_list_to_array(results.pose_landmarks,
image.shape)[:, :2],
EXPECTED_POSE_LANDMARKS, DIFF_THRESHOLD)
self._assert_diff_less(
self._world_landmarks_list_to_array(results.pose_world_landmarks),
EXPECTED_POSE_WORLD_LANDMARKS, WORLD_DIFF_THRESHOLD)
@parameterized.named_parameters(
('full', 1, 'pose_squats.full.npz'))
@@ -99,9 +121,9 @@ class PoseTest(parameterized.TestCase):
"""Tests pose models on a video."""
# If set to `True` will dump actual predictions to .npz and JSON files.
dump_predictions = False
# Set threshold for comparing actual and expected predictions in pixels.
diff_threshold = 50
diff_threshold = 15
world_diff_threshold = 0.1
video_path = os.path.join(os.path.dirname(__file__),
'testdata/pose_squats.mp4')
@@ -111,6 +133,7 @@ class PoseTest(parameterized.TestCase):
# Predict pose landmarks for each frame.
video_cap = cv2.VideoCapture(video_path)
actual_per_frame = []
actual_world_per_frame = []
frame_idx = 0
with mp_pose.Pose(static_image_mode=False,
model_complexity=model_complexity) as pose:
@@ -125,28 +148,35 @@ class PoseTest(parameterized.TestCase):
result = pose.process(image=input_frame)
pose_landmarks = self._landmarks_list_to_array(result.pose_landmarks,
input_frame.shape)
pose_world_landmarks = self._world_landmarks_list_to_array(
result.pose_world_landmarks)
actual_per_frame.append(pose_landmarks)
actual_world_per_frame.append(pose_world_landmarks)
input_frame = cv2.cvtColor(input_frame, cv2.COLOR_RGB2BGR)
self._annotate(input_frame, result, frame_idx)
frame_idx += 1
actual = np.asarray(actual_per_frame)
actual = np.array(actual_per_frame)
actual_world = np.array(actual_world_per_frame)
if dump_predictions:
# Dump .npz
with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
np.savez(tmp_file, predictions=np.array(actual))
np.savez(tmp_file, predictions=actual, predictions_world=actual_world)
print('Predictions saved as .npz to {}'.format(tmp_file.name))
# Dump JSON
with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
with open(tmp_file.name, 'w') as fl:
dump_data = {'predictions': np.around(actual, 3).tolist()}
dump_data = {
'predictions': np.around(actual, 3).tolist(),
'predictions_world': np.around(actual_world, 3).tolist()
}
fl.write(json.dumps(dump_data, indent=2, separators=(',', ': ')))
print('Predictions saved as JSON to {}'.format(tmp_file.name))
# Validate actual vs. expected predictions.
# Validate actual vs. expected landmarks.
expected = np.load(expected_path)['predictions']
assert actual.shape == expected.shape, (
'Unexpected shape of predictions: {} instead of {}'.format(
@@ -154,6 +184,14 @@ class PoseTest(parameterized.TestCase):
self._assert_diff_less(
actual[..., :2], expected[..., :2], threshold=diff_threshold)
# Validate actual vs. expected world landmarks.
expected_world = np.load(expected_path)['predictions_world']
assert actual_world.shape == expected_world.shape, (
'Unexpected shape of world predictions: {} instead of {}'.format(
actual_world.shape, expected_world.shape))
self._assert_diff_less(
actual_world, expected_world, threshold=world_diff_threshold)
if __name__ == '__main__':
absltest.main()