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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 6e4aff1cc351be3ae4537b677f36d139ee50ce09
This commit is contained in:
MediaPipe Team
2021-03-25 22:09:18 -04:00
committed by chuoling
parent a92cff7a60
commit 7c331ad58b
175 changed files with 4804 additions and 1325 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/python/solution_base.py
+4 -1
View File
@@ -11,7 +11,6 @@
# 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.
"""MediaPipe SolutionBase module.
MediaPipe SolutionBase is the common base class for the high-level MediaPipe
@@ -123,6 +122,8 @@ NAME_TO_TYPE: Mapping[str, '_PacketDataType'] = {
_PacketDataType.PROTO,
'::mediapipe::Landmark':
_PacketDataType.PROTO,
'::mediapipe::LandmarkList':
_PacketDataType.PROTO,
'::mediapipe::NormalizedLandmark':
_PacketDataType.PROTO,
'::mediapipe::FrameAnnotation':
@@ -145,6 +146,8 @@ NAME_TO_TYPE: Mapping[str, '_PacketDataType'] = {
_PacketDataType.PROTO_LIST,
'::std::vector<::mediapipe::Landmark>':
_PacketDataType.PROTO_LIST,
'::std::vector<::mediapipe::LandmarkList>':
_PacketDataType.PROTO_LIST,
'::std::vector<::mediapipe::NormalizedLandmark>':
_PacketDataType.PROTO_LIST,
'::std::vector<::mediapipe::NormalizedLandmarkList>':
mediapipe/python/solutions/drawing_utils.py
+13 -17
View File
@@ -28,6 +28,8 @@ from mediapipe.framework.formats import landmark_pb2
PRESENCE_THRESHOLD = 0.5
RGB_CHANNELS = 3
RED_COLOR = (0, 0, 255)
GREEN_COLOR = (0, 128, 0)
BLUE_COLOR = (255, 0, 0)
VISIBILITY_THRESHOLD = 0.5
@@ -178,9 +180,7 @@ def draw_axis(
focal_length: Tuple[float, float] = (1.0, 1.0),
principal_point: Tuple[float, float] = (0.0, 0.0),
axis_length: float = 0.1,
x_axis_drawing_spec: DrawingSpec = DrawingSpec(color=(0, 0, 255)),
y_axis_drawing_spec: DrawingSpec = DrawingSpec(color=(0, 128, 0)),
z_axis_drawing_spec: DrawingSpec = DrawingSpec(color=(255, 0, 0))):
axis_drawing_spec: DrawingSpec = DrawingSpec()):
"""Draws the 3D axis on the image.
Args:
@@ -190,12 +190,8 @@ def draw_axis(
focal_length: camera focal length along x and y directions.
principal_point: camera principal point in x and y.
axis_length: length of the axis in the drawing.
x_axis_drawing_spec: A DrawingSpec object that specifies the x axis
drawing settings such as color, line thickness.
y_axis_drawing_spec: A DrawingSpec object that specifies the y axis
drawing settings such as color, line thickness.
z_axis_drawing_spec: A DrawingSpec object that specifies the z axis
drawing settings such as color, line thickness.
axis_drawing_spec: A DrawingSpec object that specifies the xyz axis
drawing settings such as line thickness.
Raises:
ValueError: If one of the followings:
@@ -213,8 +209,8 @@ def draw_axis(
# Project 3D points to NDC space.
fx, fy = focal_length
px, py = principal_point
x_ndc = -fx * x / z + px
y_ndc = -fy * y / z + py
x_ndc = np.clip(-fx * x / (z + 1e-5) + px, -1., 1.)
y_ndc = np.clip(-fy * y / (z + 1e-5) + py, -1., 1.)
# Convert from NDC space to image space.
x_im = np.int32((1 + x_ndc) * 0.5 * image_cols)
y_im = np.int32((1 - y_ndc) * 0.5 * image_rows)
@@ -223,9 +219,9 @@ def draw_axis(
x_axis = (x_im[1], y_im[1])
y_axis = (x_im[2], y_im[2])
z_axis = (x_im[3], y_im[3])
image = cv2.arrowedLine(image, origin, x_axis, x_axis_drawing_spec.color,
x_axis_drawing_spec.thickness)
image = cv2.arrowedLine(image, origin, y_axis, y_axis_drawing_spec.color,
y_axis_drawing_spec.thickness)
image = cv2.arrowedLine(image, origin, z_axis, z_axis_drawing_spec.color,
z_axis_drawing_spec.thickness)
cv2.arrowedLine(image, origin, x_axis, RED_COLOR,
axis_drawing_spec.thickness)
cv2.arrowedLine(image, origin, y_axis, GREEN_COLOR,
axis_drawing_spec.thickness)
cv2.arrowedLine(image, origin, z_axis, BLUE_COLOR,
axis_drawing_spec.thickness)
mediapipe/python/solutions/drawing_utils_test.py
+43
View File
@@ -28,6 +28,7 @@ from mediapipe.python.solutions import drawing_utils
DEFAULT_BBOX_DRAWING_SPEC = drawing_utils.DrawingSpec()
DEFAULT_CONNECTION_DRAWING_SPEC = drawing_utils.DrawingSpec()
DEFAULT_CIRCLE_DRAWING_SPEC = drawing_utils.DrawingSpec(color=(0, 0, 255))
DEFAULT_AXIS_DRAWING_SPEC = drawing_utils.DrawingSpec()
class DrawingUtilTest(parameterized.TestCase):
@@ -40,6 +41,11 @@ class DrawingUtilTest(parameterized.TestCase):
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
drawing_utils.draw_detection(image, detection_pb2.Detection())
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
rotation = np.eye(3, dtype=np.float32)
translation = np.array([0., 0., 1.])
drawing_utils.draw_axis(image, rotation, translation)
def test_invalid_connection(self):
landmark_list = text_format.Parse(
@@ -133,6 +139,43 @@ class DrawingUtilTest(parameterized.TestCase):
image=image, landmark_list=landmark_list, connections=[(0, 1)])
np.testing.assert_array_equal(image, expected_result)
def test_draw_axis(self):
image = np.zeros((100, 100, 3), np.uint8)
expected_result = np.copy(image)
origin = (50, 50)
x_axis = (75, 50)
y_axis = (50, 22)
z_axis = (50, 77)
cv2.arrowedLine(expected_result, origin, x_axis, drawing_utils.RED_COLOR,
DEFAULT_AXIS_DRAWING_SPEC.thickness)
cv2.arrowedLine(expected_result, origin, y_axis, drawing_utils.GREEN_COLOR,
DEFAULT_AXIS_DRAWING_SPEC.thickness)
cv2.arrowedLine(expected_result, origin, z_axis, drawing_utils.BLUE_COLOR,
DEFAULT_AXIS_DRAWING_SPEC.thickness)
r = np.sqrt(2.) / 2.
rotation = np.array([[1., 0., 0.], [0., r, -r], [0., r, r]])
translation = np.array([0, 0, -0.2])
drawing_utils.draw_axis(image, rotation, translation)
np.testing.assert_array_equal(image, expected_result)
def test_draw_axis_zero_translation(self):
image = np.zeros((100, 100, 3), np.uint8)
expected_result = np.copy(image)
origin = (50, 50)
x_axis = (0, 50)
y_axis = (50, 100)
z_axis = (50, 50)
cv2.arrowedLine(expected_result, origin, x_axis, drawing_utils.RED_COLOR,
DEFAULT_AXIS_DRAWING_SPEC.thickness)
cv2.arrowedLine(expected_result, origin, y_axis, drawing_utils.GREEN_COLOR,
DEFAULT_AXIS_DRAWING_SPEC.thickness)
cv2.arrowedLine(expected_result, origin, z_axis, drawing_utils.BLUE_COLOR,
DEFAULT_AXIS_DRAWING_SPEC.thickness)
rotation = np.eye(3, dtype=np.float32)
translation = np.zeros((3,), dtype=np.float32)
drawing_utils.draw_axis(image, rotation, translation)
np.testing.assert_array_equal(image, expected_result)
def test_min_and_max_coordinate_values(self):
landmark_list = text_format.Parse(
'landmark {x: 0.0 y: 1.0}'
mediapipe/python/solutions/objectron.py
+27 -3
View File
@@ -15,7 +15,10 @@
"""MediaPipe Objectron."""
import enum
import os
import shutil
from typing import List, Tuple, NamedTuple, Optional
import urllib.request
import attr
import numpy as np
@@ -89,6 +92,23 @@ BOX_CONNECTIONS = frozenset([
(BoxLandmark.FRONT_BOTTOM_RIGHT, BoxLandmark.FRONT_TOP_RIGHT),
(BoxLandmark.BACK_TOP_RIGHT, BoxLandmark.FRONT_TOP_RIGHT),
])
_OSS_URL_PREFIX = 'https://github.com/google/mediapipe/raw/master/'
def _download_oss_model(model_path: str):
"""Download the objectron oss model from GitHub if it doesn't exist in the package."""
mp_root_path = os.sep.join(os.path.abspath(__file__).split(os.sep)[:-4])
model_abspath = os.path.join(mp_root_path, model_path)
if os.path.exists(model_abspath):
return
model_url = _OSS_URL_PREFIX + model_path
with urllib.request.urlopen(model_url) as response, open(model_abspath,
'wb') as out_file:
if response.code != 200:
raise ConnectionError('Cannot download ' + model_path +
' from the MediaPipe Github repo.')
shutil.copyfileobj(response, out_file)
@attr.s(auto_attribs=True)
@@ -132,9 +152,10 @@ _MODEL_DICT = {
}
def GetModelByName(name: str) -> ObjectronModel:
def get_model_by_name(name: str) -> ObjectronModel:
if name not in _MODEL_DICT:
raise ValueError(f'{name} is not a valid model name for Objectron.')
_download_oss_model(_MODEL_DICT[name].model_path)
return _MODEL_DICT[name]
@@ -186,6 +207,10 @@ class Objectron(SolutionBase):
conversions inside the API.
image_size (Optional): size (image_width, image_height) of the input image
, ONLY needed when use focal_length and principal_point in pixel space.
Raises:
ConnectionError: If the objectron open source model can't be downloaded
from the MediaPipe Github repo.
"""
# Get Camera parameters.
fx, fy = focal_length
@@ -199,7 +224,7 @@ class Objectron(SolutionBase):
py = - (py - half_height) / half_height
# Create and init model.
model = GetModelByName(model_name)
model = get_model_by_name(model_name)
super().__init__(
binary_graph_path=BINARYPB_FILE_PATH,
side_inputs={
@@ -275,4 +300,3 @@ class Objectron(SolutionBase):
new_outputs.append(ObjectronOutputs(landmarks_2d, landmarks_3d,
rotation, translation, scale=scale))
return new_outputs
mediapipe/python/solutions/pose.py
+1
View File
@@ -36,6 +36,7 @@ from mediapipe.calculators.util import logic_calculator_pb2
from mediapipe.calculators.util import non_max_suppression_calculator_pb2
from mediapipe.calculators.util import rect_transformation_calculator_pb2
from mediapipe.calculators.util import thresholding_calculator_pb2
from mediapipe.calculators.util import visibility_smoothing_calculator_pb2
# pylint: enable=unused-import
from mediapipe.python.solution_base import SolutionBase
mediapipe/python/solutions/pose_test.py
+66 -7
View File
@@ -13,7 +13,9 @@
# limitations under the License.
"""Tests for mediapipe.python.solutions.pose."""
import json
import os
import tempfile
from absl.testing import absltest
from absl.testing import parameterized
@@ -52,7 +54,7 @@ class PoseTest(parameterized.TestCase):
def _landmarks_list_to_array(self, landmark_list, image_shape):
rows, cols, _ = image_shape
return np.asarray([(lmk.x * cols, lmk.y * rows)
return np.asarray([(lmk.x * cols, lmk.y * rows, lmk.z * cols)
for lmk in landmark_list.landmark])
def _assert_diff_less(self, array1, array2, threshold):
@@ -81,9 +83,9 @@ class PoseTest(parameterized.TestCase):
for _ in range(num_frames):
results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._assert_diff_less(
self._landmarks_list_to_array(results.pose_landmarks, image.shape),
EXPECTED_UPPER_BODY_LANDMARKS,
DIFF_THRESHOLD)
self._landmarks_list_to_array(results.pose_landmarks,
image.shape)[:, :2],
EXPECTED_UPPER_BODY_LANDMARKS, DIFF_THRESHOLD)
@parameterized.named_parameters(('static_image_mode', True, 3),
('video_mode', False, 3))
@@ -95,9 +97,66 @@ class PoseTest(parameterized.TestCase):
for _ in range(num_frames):
results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._assert_diff_less(
self._landmarks_list_to_array(results.pose_landmarks, image.shape),
EXPECTED_FULL_BODY_LANDMARKS,
DIFF_THRESHOLD)
self._landmarks_list_to_array(results.pose_landmarks,
image.shape)[:, :2],
EXPECTED_FULL_BODY_LANDMARKS, DIFF_THRESHOLD)
@parameterized.named_parameters(
('full_body', False, 'pose_squats.full_body.npz'),
('upper_body', True, 'pose_squats.upper_body.npz'))
def test_on_video(self, upper_body_only, expected_name):
"""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
video_path = os.path.join(os.path.dirname(__file__),
'testdata/pose_squats.mp4')
expected_path = os.path.join(os.path.dirname(__file__),
'testdata/{}'.format(expected_name))
# Predict pose landmarks for each frame.
video_cap = cv2.VideoCapture(video_path)
actual_per_frame = []
with mp_pose.Pose(
static_image_mode=False, upper_body_only=upper_body_only) as pose:
while True:
# Get next frame of the video.
success, input_frame = video_cap.read()
if not success:
break
# Run pose tracker.
input_frame = cv2.cvtColor(input_frame, cv2.COLOR_BGR2RGB)
result = pose.process(image=input_frame)
pose_landmarks = self._landmarks_list_to_array(result.pose_landmarks,
input_frame.shape)
actual_per_frame.append(pose_landmarks)
actual = np.asarray(actual_per_frame)
if dump_predictions:
# Dump .npz
with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
np.savez(tmp_file, predictions=np.array(actual))
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()}
fl.write(json.dumps(dump_data, indent=2, separators=(',', ': ')))
print('Predictions saved as JSON to {}'.format(tmp_file.name))
# Validate actual vs. expected predictions.
expected = np.load(expected_path)['predictions']
assert actual.shape == expected.shape, (
'Unexpected shape of predictions: {} instead of {}'.format(
actual.shape, expected.shape))
self._assert_diff_less(
actual[..., :2], expected[..., :2], threshold=diff_threshold)
if __name__ == '__main__':