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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 6e5aa035cd1f6a9333962df5d3ab97a05bd5744e
This commit is contained in:
MediaPipe Team
2022-06-28 12:11:05 +00:00
committed by Sebastian Schmidt
parent 4a20e9909d
commit c688862570
144 changed files with 5772 additions and 2118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/python/pybind/packet_getter.cc
+16 -1
View File
@@ -247,7 +247,22 @@ void PublicPacketGetters(pybind11::module* m) {
)doc");
m->def(
"get_float_list", &GetContent<std::vector<float>>,
"get_float_list",
[](const Packet& packet) {
if (packet.ValidateAsType<std::vector<float>>().ok()) {
return packet.Get<std::vector<float>>();
} else if (packet.ValidateAsType<std::array<float, 16>>().ok()) {
auto float_array = packet.Get<std::array<float, 16>>();
return std::vector<float>(float_array.begin(), float_array.end());
} else if (packet.ValidateAsType<std::array<float, 4>>().ok()) {
auto float_array = packet.Get<std::array<float, 4>>();
return std::vector<float>(float_array.begin(), float_array.end());
} else {
throw RaisePyError(PyExc_ValueError,
"Packet doesn't contain std::vector<float> or "
"std::array<float, 4 / 16> containers.");
}
},
R"doc(Get the content of a MediaPipe float vector Packet as a float list.
Args:
mediapipe/python/solution_base.py
+51
View File
@@ -28,6 +28,7 @@ from typing import Any, Iterable, List, Mapping, NamedTuple, Optional, Union
import numpy as np
from google.protobuf.internal import containers
from google.protobuf import descriptor
from google.protobuf import message
# resources dependency
@@ -216,6 +217,7 @@ class SolutionBase:
binary_graph_path: Optional[str] = None,
graph_config: Optional[calculator_pb2.CalculatorGraphConfig] = None,
calculator_params: Optional[Mapping[str, Any]] = None,
graph_options: Optional[message.Message] = None,
side_inputs: Optional[Mapping[str, Any]] = None,
outputs: Optional[List[str]] = None,
stream_type_hints: Optional[Mapping[str, PacketDataType]] = None):
@@ -227,6 +229,7 @@ class SolutionBase:
format.
calculator_params: A mapping from the
{calculator_name}.{options_field_name} str to the field value.
graph_options: The graph options protobuf for the mediapipe graph.
side_inputs: A mapping from the side packet name to the packet raw data.
outputs: A list of the graph output stream names to observe. If the list
is empty, all the output streams listed in the graph config will be
@@ -267,6 +270,10 @@ class SolutionBase:
if calculator_params:
self._modify_calculator_options(canonical_graph_config_proto,
calculator_params)
if graph_options:
self._set_extension(canonical_graph_config_proto.graph_options,
graph_options)
self._graph = calculator_graph.CalculatorGraph(
graph_config=canonical_graph_config_proto)
self._simulated_timestamp = 0
@@ -530,6 +537,50 @@ class SolutionBase:
if num_modified < len(nested_calculator_params):
raise ValueError('Not all calculator params are valid.')
def create_graph_options(self, options_message: message.Message,
values: Mapping[str, Any]) -> message.Message:
"""Sets protobuf field values.
Args:
options_message: the options protobuf message.
values: field value pairs, where each field may be a "." separated path.
Returns:
the options protobuf message.
"""
if hasattr(values, 'items'):
values = values.items()
for pair in values:
(field, value) = pair
fields = field.split('.')
m = options_message
while len(fields) > 1:
m = getattr(m, fields[0])
del fields[0]
v = getattr(m, fields[0])
if hasattr(v, 'append'):
del v[:]
v.extend(value)
elif hasattr(v, 'CopyFrom'):
v.CopyFrom(value)
else:
setattr(m, fields[0], value)
return options_message
def _set_extension(self,
extension_list: containers.RepeatedCompositeFieldContainer,
extension_value: message.Message) -> None:
"""Sets one value in a repeated protobuf.Any extension field."""
for extension_any in extension_list:
if extension_any.Is(extension_value.DESCRIPTOR):
v = type(extension_value)()
extension_any.Unpack(v)
v.MergeFrom(extension_value)
extension_any.Pack(v)
return
extension_list.add().Pack(extension_value)
def _make_packet(self, packet_data_type: PacketDataType,
data: Any) -> packet.Packet:
if (packet_data_type == PacketDataType.IMAGE_FRAME or
mediapipe/python/solutions/drawing_utils.py
+13 -13
View File
@@ -28,7 +28,7 @@ from mediapipe.framework.formats import landmark_pb2
_PRESENCE_THRESHOLD = 0.5
_VISIBILITY_THRESHOLD = 0.5
_RGB_CHANNELS = 3
_BGR_CHANNELS = 3
WHITE_COLOR = (224, 224, 224)
BLACK_COLOR = (0, 0, 0)
@@ -74,7 +74,7 @@ def draw_detection(
"""Draws the detction bounding box and keypoints on the image.
Args:
image: A three channel RGB image represented as numpy ndarray.
image: A three channel BGR image represented as numpy ndarray.
detection: A detection proto message to be annotated on the image.
keypoint_drawing_spec: A DrawingSpec object that specifies the keypoints'
drawing settings such as color, line thickness, and circle radius.
@@ -83,13 +83,13 @@ def draw_detection(
Raises:
ValueError: If one of the followings:
a) If the input image is not three channel RGB.
a) If the input image is not three channel BGR.
b) If the location data is not relative data.
"""
if not detection.location_data:
return
if image.shape[2] != _RGB_CHANNELS:
raise ValueError('Input image must contain three channel rgb data.')
if image.shape[2] != _BGR_CHANNELS:
raise ValueError('Input image must contain three channel bgr data.')
image_rows, image_cols, _ = image.shape
location = detection.location_data
@@ -130,7 +130,7 @@ def draw_landmarks(
"""Draws the landmarks and the connections on the image.
Args:
image: A three channel RGB image represented as numpy ndarray.
image: A three channel BGR image represented as numpy ndarray.
landmark_list: A normalized landmark list proto message to be annotated on
the image.
connections: A list of landmark index tuples that specifies how landmarks to
@@ -147,13 +147,13 @@ def draw_landmarks(
Raises:
ValueError: If one of the followings:
a) If the input image is not three channel RGB.
a) If the input image is not three channel BGR.
b) If any connetions contain invalid landmark index.
"""
if not landmark_list:
return
if image.shape[2] != _RGB_CHANNELS:
raise ValueError('Input image must contain three channel rgb data.')
if image.shape[2] != _BGR_CHANNELS:
raise ValueError('Input image must contain three channel bgr data.')
image_rows, image_cols, _ = image.shape
idx_to_coordinates = {}
for idx, landmark in enumerate(landmark_list.landmark):
@@ -208,7 +208,7 @@ def draw_axis(
"""Draws the 3D axis on the image.
Args:
image: A three channel RGB image represented as numpy ndarray.
image: A three channel BGR image represented as numpy ndarray.
rotation: Rotation matrix from object to camera coordinate frame.
translation: Translation vector from object to camera coordinate frame.
focal_length: camera focal length along x and y directions.
@@ -219,10 +219,10 @@ def draw_axis(
Raises:
ValueError: If one of the followings:
a) If the input image is not three channel RGB.
a) If the input image is not three channel BGR.
"""
if image.shape[2] != _RGB_CHANNELS:
raise ValueError('Input image must contain three channel rgb data.')
if image.shape[2] != _BGR_CHANNELS:
raise ValueError('Input image must contain three channel bgr data.')
image_rows, image_cols, _ = image.shape
# Create axis points in camera coordinate frame.
axis_world = np.float32([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]])
mediapipe/python/solutions/drawing_utils_test.py
+5 -4
View File
@@ -27,7 +27,8 @@ 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_CIRCLE_DRAWING_SPEC = drawing_utils.DrawingSpec(
color=drawing_utils.RED_COLOR)
DEFAULT_AXIS_DRAWING_SPEC = drawing_utils.DrawingSpec()
DEFAULT_CYCLE_BORDER_COLOR = (224, 224, 224)
@@ -37,13 +38,13 @@ class DrawingUtilTest(parameterized.TestCase):
def test_invalid_input_image(self):
image = np.arange(18, dtype=np.uint8).reshape(3, 3, 2)
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
ValueError, 'Input image must contain three channel bgr data.'):
drawing_utils.draw_landmarks(image, landmark_pb2.NormalizedLandmarkList())
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
ValueError, 'Input image must contain three channel bgr data.'):
drawing_utils.draw_detection(image, detection_pb2.Detection())
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
ValueError, 'Input image must contain three channel bgr data.'):
rotation = np.eye(3, dtype=np.float32)
translation = np.array([0., 0., 1.])
drawing_utils.draw_axis(image, rotation, translation)
mediapipe/python/solutions/face_detection.py
+5 -12
View File
@@ -19,13 +19,7 @@ from typing import NamedTuple, Union
import numpy as np
from mediapipe.framework.formats import detection_pb2
from mediapipe.framework.formats import location_data_pb2
# pylint: disable=unused-import
from mediapipe.calculators.tensor import image_to_tensor_calculator_pb2
from mediapipe.calculators.tensor import inference_calculator_pb2
from mediapipe.calculators.tensor import tensors_to_detections_calculator_pb2
from mediapipe.calculators.tflite import ssd_anchors_calculator_pb2
from mediapipe.calculators.util import non_max_suppression_calculator_pb2
# pylint: enable=unused-import
from mediapipe.modules.face_detection import face_detection_pb2
from mediapipe.python.solution_base import SolutionBase
_SHORT_RANGE_GRAPH_FILE_PATH = 'mediapipe/modules/face_detection/face_detection_short_range_cpu.binarypb'
@@ -84,14 +78,13 @@ class FaceDetection(SolutionBase):
"""
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=binary_graph_path,
calculator_params={
subgraph_name + '__TensorsToDetectionsCalculator.min_score_thresh':
min_detection_confidence,
},
graph_options=self.create_graph_options(
face_detection_pb2.FaceDetectionOptions(), {
'min_score_thresh': min_detection_confidence,
}),
outputs=['detections'])
def process(self, image: np.ndarray) -> NamedTuple:
mediapipe/python/solutions/face_mesh.py
+1 -1
View File
@@ -99,7 +99,7 @@ class FaceMesh(SolutionBase):
'use_prev_landmarks': not static_image_mode,
},
calculator_params={
'facedetectionshortrangecpu__facedetectionshortrangecommon__TensorsToDetectionsCalculator.min_score_thresh':
'facedetectionshortrangecpu__facedetectionshortrange__facedetection__TensorsToDetectionsCalculator.min_score_thresh':
min_detection_confidence,
'facelandmarkcpu__ThresholdingCalculator.threshold':
min_tracking_confidence,