mediapipe/mediapipe2/python/pybind/packet_creator.cc
2021-06-10 23:01:19 +00:00

685 lines
20 KiB
C++

// Copyright 2020 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.
#include "mediapipe/python/pybind/packet_creator.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/matrix.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/timestamp.h"
#include "mediapipe/python/pybind/image_frame_util.h"
#include "mediapipe/python/pybind/util.h"
#include "pybind11/eigen.h"
#include "pybind11/pybind11.h"
#include "pybind11/stl.h"
namespace mediapipe {
namespace python {
namespace {
Packet CreateImageFramePacket(mediapipe::ImageFormat::Format format,
const py::array& data, bool copy) {
if (format == mediapipe::ImageFormat::SRGB ||
format == mediapipe::ImageFormat::SRGBA ||
format == mediapipe::ImageFormat::GRAY8) {
return Adopt(CreateImageFrame<uint8>(format, data, copy).release());
} else if (format == mediapipe::ImageFormat::GRAY16 ||
format == mediapipe::ImageFormat::SRGB48 ||
format == mediapipe::ImageFormat::SRGBA64) {
return Adopt(CreateImageFrame<uint16>(format, data, copy).release());
} else if (format == mediapipe::ImageFormat::VEC32F1 ||
format == mediapipe::ImageFormat::VEC32F2) {
return Adopt(CreateImageFrame<float>(format, data, copy).release());
}
throw RaisePyError(PyExc_RuntimeError,
absl::StrCat("Unsupported ImageFormat: ", format).c_str());
return Packet();
}
Packet CreateImagePacket(mediapipe::ImageFormat::Format format,
const py::array& data, bool copy) {
if (format == mediapipe::ImageFormat::SRGB ||
format == mediapipe::ImageFormat::SRGBA ||
format == mediapipe::ImageFormat::GRAY8) {
return MakePacket<Image>(std::make_shared<ImageFrame>(
std::move(*CreateImageFrame<uint8>(format, data, copy).release())));
} else if (format == mediapipe::ImageFormat::GRAY16 ||
format == mediapipe::ImageFormat::SRGB48 ||
format == mediapipe::ImageFormat::SRGBA64) {
return MakePacket<Image>(std::make_shared<ImageFrame>(
std::move(*CreateImageFrame<uint16>(format, data, copy).release())));
} else if (format == mediapipe::ImageFormat::VEC32F1 ||
format == mediapipe::ImageFormat::VEC32F2) {
return MakePacket<Image>(std::make_shared<ImageFrame>(
std::move(*CreateImageFrame<float>(format, data, copy).release())));
}
throw RaisePyError(PyExc_RuntimeError,
absl::StrCat("Unsupported ImageFormat: ", format).c_str());
return Packet();
}
} // namespace
namespace py = pybind11;
void PublicPacketCreators(pybind11::module* m) {
m->def(
"create_string",
[](const std::string& data) { return MakePacket<std::string>(data); },
R"doc(Create a MediaPipe std::string Packet from a str.
Args:
data: A str.
Returns:
A MediaPipe std::string Packet.
Raises:
TypeError: If the input is not a str.
Examples:
packet = mp.packet_creator.create_string('abc')
data = mp.packet_getter.get_string(packet)
)doc",
py::return_value_policy::move);
m->def(
"create_string",
[](const py::bytes& data) { return MakePacket<std::string>(data); },
R"doc(Create a MediaPipe std::string Packet from a bytes object.
Args:
data: A bytes object.
Returns:
A MediaPipe std::string Packet.
Raises:
TypeError: If the input is not a bytes object.
Examples:
packet = mp.packet_creator.create_string(b'\xd0\xd0\xd0')
data = mp.packet_getter.get_bytes(packet)
)doc",
py::return_value_policy::move);
m->def(
"create_bool", [](bool data) { return MakePacket<bool>(data); },
R"doc(Create a MediaPipe bool Packet from a boolean object.
Args:
data: A boolean object.
Returns:
A MediaPipe bool Packet.
Raises:
TypeError: If the input is not a boolean object.
Examples:
packet = mp.packet_creator.create_bool(True)
data = mp.packet_getter.get_bool(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_int",
[](int64 data) {
RaisePyErrorIfOverflow(data, INT_MIN, INT_MAX);
return MakePacket<int>(data);
},
R"doc(Create a MediaPipe int Packet from an integer.
Args:
data: An integer or a np.intc.
Returns:
A MediaPipe int Packet.
Raises:
OverflowError: If the input integer overflows.
TypeError: If the input is not an integer.
Examples:
packet = mp.packet_creator.create_int(0)
data = mp.packet_getter.get_int(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_int8",
[](int64 data) {
RaisePyErrorIfOverflow(data, INT8_MIN, INT8_MAX);
return MakePacket<int8>(data);
},
R"doc(Create a MediaPipe int8 Packet from an integer.
Args:
data: An integer or a np.int8.
Returns:
A MediaPipe int8 Packet.
Raises:
OverflowError: If the input integer overflows.
TypeError: If the input is neither an integer nor a np.int8.
Examples:
packet = mp.packet_creator.create_int8(2**7 - 1)
data = mp.packet_getter.get_int(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_int16",
[](int64 data) {
RaisePyErrorIfOverflow(data, INT16_MIN, INT16_MAX);
return MakePacket<int16>(data);
},
R"doc(Create a MediaPipe int16 Packet from an integer.
Args:
data: An integer or a np.int16.
Returns:
A MediaPipe int16 Packet.
Raises:
OverflowError: If the input integer overflows.
TypeError: If the input is neither an integer nor a np.int16.
Examples:
packet = mp.packet_creator.create_int16(2**15 - 1)
data = mp.packet_getter.get_int(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_int32",
[](int64 data) {
RaisePyErrorIfOverflow(data, INT32_MIN, INT32_MAX);
return MakePacket<int32>(data);
},
R"doc(Create a MediaPipe int32 Packet from an integer.
Args:
data: An integer or a np.int32.
Returns:
A MediaPipe int32 Packet.
Raises:
OverflowError: If the input integer overflows.
TypeError: If the input is neither an integer nor a np.int32.
Examples:
packet = mp.packet_creator.create_int32(2**31 - 1)
data = mp.packet_getter.get_int(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_int64", [](int64 data) { return MakePacket<int64>(data); },
R"doc(Create a MediaPipe int64 Packet from an integer.
Args:
data: An integer or a np.int64.
Returns:
A MediaPipe int64 Packet.
Raises:
TypeError: If the input is neither an integer nor a np.int64.
Examples:
packet = mp.packet_creator.create_int64(2**63 - 1)
data = mp.packet_getter.get_int(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_uint8",
[](int64 data) {
RaisePyErrorIfOverflow(data, 0, UINT8_MAX);
return MakePacket<uint8>(data);
},
R"doc(Create a MediaPipe uint8 Packet from an integer.
Args:
data: An integer or a np.uint8.
Returns:
A MediaPipe uint8 Packet.
Raises:
OverflowError: If the input integer overflows.
TypeError: If the input is neither an integer nor a np.uint8.
Examples:
packet = mp.packet_creator.create_uint8(2**8 - 1)
data = mp.packet_getter.get_uint(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_uint16",
[](int64 data) {
RaisePyErrorIfOverflow(data, 0, UINT16_MAX);
return MakePacket<uint16>(data);
},
R"doc(Create a MediaPipe uint16 Packet from an integer.
Args:
data: An integer or a np.uint16.
Returns:
A MediaPipe uint16 Packet.
Raises:
OverflowError: If the input integer overflows.
TypeError: If the input is neither an integer nor a np.uint16.
Examples:
packet = mp.packet_creator.create_uint16(2**16 - 1)
data = mp.packet_getter.get_uint(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_uint32",
[](int64 data) {
RaisePyErrorIfOverflow(data, 0, UINT32_MAX);
return MakePacket<uint32>(data);
},
R"doc(Create a MediaPipe uint32 Packet from an integer.
Args:
data: An integer or a np.uint32.
Returns:
A MediaPipe uint32 Packet.
Raises:
OverflowError: If the input integer overflows.
TypeError: If the input is neither an integer nor a np.uint32.
Examples:
packet = mp.packet_creator.create_uint32(2**32 - 1)
data = mp.packet_getter.get_uint(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_uint64", [](uint64 data) { return MakePacket<uint64>(data); },
R"doc(Create a MediaPipe uint64 Packet from an integer.
Args:
data: An integer or a np.uint64.
Returns:
A MediaPipe uint64 Packet.
Raises:
TypeError: If the input is neither an integer nor a np.uint64.
Examples:
packet = mp.packet_creator.create_uint64(2**64 - 1)
data = mp.packet_getter.get_uint(packet)
)doc",
// py::arg().noconvert() won't allow this to accept np.uint64 data type.
py::arg(), py::return_value_policy::move);
m->def(
"create_float", [](float data) { return MakePacket<float>(data); },
R"doc(Create a MediaPipe float Packet from a float.
Args:
data: A float or a np.float.
Returns:
A MediaPipe float Packet.
Raises:
TypeError: If the input is neither a float nor a np.float.
Examples:
packet = mp.packet_creator.create_float(0.1)
data = mp.packet_getter.get_float(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_double", [](double data) { return MakePacket<double>(data); },
R"doc(Create a MediaPipe double Packet from a float.
Args:
data: A float or a np.double.
Returns:
A MediaPipe double Packet.
Raises:
TypeError: If the input is neither a float nore a np.double.
Examples:
packet = mp.packet_creator.create_double(0.1)
data = mp.packet_getter.get_float(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_int_array",
[](const std::vector<int>& data) {
int* ints = new int[data.size()];
std::copy(data.begin(), data.end(), ints);
return Adopt(reinterpret_cast<int(*)[]>(ints));
},
R"doc(Create a MediaPipe int array Packet from a list of integers.
Args:
data: A list of integers.
Returns:
A MediaPipe int array Packet.
Raises:
TypeError: If the input is not a list of integers.
Examples:
packet = mp.packet_creator.create_int_array([1, 2, 3])
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_float_array",
[](const std::vector<float>& data) {
float* floats = new float[data.size()];
std::copy(data.begin(), data.end(), floats);
return Adopt(reinterpret_cast<float(*)[]>(floats));
},
R"doc(Create a MediaPipe float array Packet from a list of floats.
Args:
data: A list of floats.
Returns:
A MediaPipe float array Packet.
Raises:
TypeError: If the input is not a list of floats.
Examples:
packet = mp.packet_creator.create_float_array([0.1, 0.2, 0.3])
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_int_vector",
[](const std::vector<int>& data) {
return MakePacket<std::vector<int>>(data);
},
R"doc(Create a MediaPipe int vector Packet from a list of integers.
Args:
data: A list of integers.
Returns:
A MediaPipe int vector Packet.
Raises:
TypeError: If the input is not a list of integers.
Examples:
packet = mp.packet_creator.create_int_vector([1, 2, 3])
data = mp.packet_getter.get_int_vector(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_bool_vector",
[](const std::vector<bool>& data) {
return MakePacket<std::vector<bool>>(data);
},
R"doc(Create a MediaPipe bool vector Packet from a list of booleans.
Args:
data: A list of booleans.
Returns:
A MediaPipe bool vector Packet.
Raises:
TypeError: If the input is not a list of booleans.
Examples:
packet = mp.packet_creator.create_bool_vector([True, True, False])
data = mp.packet_getter.get_bool_vector(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_float_vector",
[](const std::vector<float>& data) {
return MakePacket<std::vector<float>>(data);
},
R"doc(Create a MediaPipe float vector Packet from a list of floats.
Args:
data: A list of floats
Returns:
A MediaPipe float vector Packet.
Raises:
TypeError: If the input is not a list of floats.
Examples:
packet = mp.packet_creator.create_float_vector([0.1, 0.2, 0.3])
data = mp.packet_getter.get_float_list(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_string_vector",
[](const std::vector<std::string>& data) {
return MakePacket<std::vector<std::string>>(data);
},
R"doc(Create a MediaPipe std::string vector Packet from a list of str.
Args:
data: A list of str.
Returns:
A MediaPipe std::string vector Packet.
Raises:
TypeError: If the input is not a list of str.
Examples:
packet = mp.packet_creator.create_string_vector(['a', 'b', 'c'])
data = mp.packet_getter.get_str_list(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_packet_vector",
[](const std::vector<Packet>& data) {
return MakePacket<std::vector<Packet>>(data);
},
R"doc(Create a MediaPipe Packet holds a vector of packets.
Args:
data: A list of packets.
Returns:
A MediaPipe Packet holds a vector of packets.
Raises:
TypeError: If the input is not a list of packets.
Examples:
packet = mp.packet_creator.create_packet_vector([
mp.packet_creator.create_float(0.1),
mp.packet_creator.create_int(1),
mp.packet_creator.create_string('1')
])
data = mp.packet_getter.get_packet_vector(packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_string_to_packet_map",
[](const std::map<std::string, Packet>& data) {
return MakePacket<std::map<std::string, Packet>>(data);
},
R"doc(Create a MediaPipe std::string to packet map Packet from a dictionary.
Args:
data: A dictionary that has (str, Packet) pairs.
Returns:
A MediaPipe Packet holds std::map<std::string, Packet>.
Raises:
TypeError: If the input is not a dictionary from str to packet.
Examples:
dict_packet = mp.packet_creator.create_string_to_packet_map({
'float': mp.packet_creator.create_float(0.1),
'int': mp.packet_creator.create_int(1),
'std::string': mp.packet_creator.create_string('1')
data = mp.packet_getter.get_str_to_packet_dict(dict_packet)
)doc",
py::arg().noconvert(), py::return_value_policy::move);
m->def(
"create_matrix",
// Eigen Map class
// (https://eigen.tuxfamily.org/dox/group__TutorialMapClass.html) is the
// way to reuse the external memory as an Eigen type. However, when
// creating an Eigen::MatrixXf from an Eigen Map object, the data copy
// still happens. We can make a packet of an Eigen Map type for reusing
// external memory. However,the packet data type is no longer
// Eigen::MatrixXf.
// TODO: Should take "const Eigen::Ref<const Eigen::MatrixXf>&"
// as the input argument. Investigate why bazel non-optimized mode
// triggers a memory allocation bug in Eigen::internal::aligned_free().
[](const Eigen::MatrixXf& matrix) {
// MakePacket copies the data.
return MakePacket<Matrix>(matrix);
},
R"doc(Create a MediaPipe Matrix Packet from a 2d numpy float ndarray.
The method copies data from the input MatrixXf and the returned packet owns
a MatrixXf object.
Args:
matrix: A 2d numpy float ndarray.
Returns:
A MediaPipe Matrix Packet.
Raises:
TypeError: If the input is not a 2d numpy float ndarray.
Examples:
packet = mp.packet_creator.create_matrix(
np.array([[.1, .2, .3], [.4, .5, .6]])
matrix = mp.packet_getter.get_matrix(packet)
)doc",
py::return_value_policy::move);
}
void InternalPacketCreators(pybind11::module* m) {
m->def("_create_image_frame_from_pixel_data", &CreateImageFramePacket,
py::arg("format"), py::arg("data").noconvert(), py::arg("copy"),
py::return_value_policy::move);
m->def("_create_image_from_pixel_data", &CreateImagePacket, py::arg("format"),
py::arg("data").noconvert(), py::arg("copy"),
py::return_value_policy::move);
m->def(
"_create_image_frame_from_image_frame",
[](ImageFrame& image_frame) {
auto image_frame_copy = absl::make_unique<ImageFrame>();
// Set alignment_boundary to kGlDefaultAlignmentBoundary so that
// both GPU and CPU can process it.
image_frame_copy->CopyFrom(image_frame,
ImageFrame::kGlDefaultAlignmentBoundary);
return Adopt(image_frame_copy.release());
},
py::arg("image_frame").noconvert(), py::return_value_policy::move);
m->def(
"_create_image_from_image",
[](Image& image) {
auto image_frame_copy = absl::make_unique<ImageFrame>();
// Set alignment_boundary to kGlDefaultAlignmentBoundary so that
// both GPU and CPU can process it.
image_frame_copy->CopyFrom(*image.GetImageFrameSharedPtr(),
ImageFrame::kGlDefaultAlignmentBoundary);
return MakePacket<Image>(std::make_shared<ImageFrame>(
std::move(*image_frame_copy.release())));
},
py::arg("image").noconvert(), py::return_value_policy::move);
m->def(
"_create_proto",
[](const std::string& type_name, const py::bytes& serialized_proto) {
using packet_internal::HolderBase;
absl::StatusOr<std::unique_ptr<HolderBase>> maybe_holder =
packet_internal::MessageHolderRegistry::CreateByName(type_name);
if (!maybe_holder.ok()) {
throw RaisePyError(
PyExc_RuntimeError,
absl::StrCat("Unregistered proto message type: ", type_name)
.c_str());
}
// Creates a Packet with the concrete C++ payload type.
std::unique_ptr<HolderBase> message_holder =
std::move(maybe_holder).value();
auto* copy = const_cast<proto_ns::MessageLite*>(
message_holder->GetProtoMessageLite());
copy->ParseFromString(std::string(serialized_proto));
return packet_internal::Create(message_holder.release());
},
py::return_value_policy::move);
m->def(
"_create_proto_vector",
[](const std::string& type_name,
const std::vector<py::bytes>& serialized_proto_vector) {
// TODO: Implement this.
throw RaisePyError(PyExc_NotImplementedError,
"Creating a packet from a vector of proto messages "
"is not supproted yet.");
return Packet();
},
py::return_value_policy::move);
}
void PacketCreatorSubmodule(pybind11::module* module) {
py::module m = module->def_submodule(
"_packet_creator", "MediaPipe internal packet creator module.");
PublicPacketCreators(&m);
InternalPacketCreators(&m);
}
} // namespace python
} // namespace mediapipe