Project import generated by Copybara.

GitOrigin-RevId: ea8d45731f5a052f79745e35bfd8240d6ac568d2
This commit is contained in:
MediaPipe Team
2020-12-16 00:05:25 -05:00
committed by chuoling
parent 38be2ec58f
commit 39309bedba
109 changed files with 5803 additions and 1500 deletions
+6
View File
@@ -59,6 +59,7 @@ cc_library(
deps = [
":inference_calculator_cc_proto",
"@com_google_absl//absl/memory",
"//mediapipe/framework/api2:node",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/formats:tensor",
"//mediapipe/util/tflite:tflite_model_loader",
@@ -234,6 +235,7 @@ cc_library(
"//mediapipe/framework/formats:detection_cc_proto",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
"//mediapipe/framework/api2:node",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:port",
"//mediapipe/framework/deps:file_path",
@@ -286,6 +288,7 @@ cc_library(
deps = [
":tensors_to_landmarks_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/formats:landmark_cc_proto",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:ret_check",
@@ -317,6 +320,7 @@ cc_library(
deps = [
":tensors_to_floats_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:ret_check",
],
@@ -355,6 +359,7 @@ cc_library(
":tensors_to_classification_calculator_cc_proto",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/formats:classification_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/formats:location",
@@ -421,6 +426,7 @@ cc_library(
":image_to_tensor_converter",
":image_to_tensor_converter_opencv",
":image_to_tensor_utils",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:rect_cc_proto",
"//mediapipe/framework/formats:tensor",
@@ -20,6 +20,7 @@
#include "mediapipe/calculators/tensor/image_to_tensor_converter.h"
#include "mediapipe/calculators/tensor/image_to_tensor_converter_opencv.h"
#include "mediapipe/calculators/tensor/image_to_tensor_utils.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/rect.pb.h"
@@ -45,16 +46,15 @@
#endif // !MEDIAPIPE_DISABLE_GPU
namespace {
constexpr char kInputCpu[] = "IMAGE";
constexpr char kInputGpu[] = "IMAGE_GPU";
constexpr char kOutputMatrix[] = "MATRIX";
constexpr char kOutput[] = "TENSORS";
constexpr char kInputNormRect[] = "NORM_RECT";
constexpr char kOutputLetterboxPadding[] = "LETTERBOX_PADDING";
} // namespace
namespace mediapipe {
namespace api2 {
#if MEDIAPIPE_DISABLE_GPU
// Just a placeholder to not have to depend on mediapipe::GpuBuffer.
using GpuBuffer = AnyType;
#else
using GpuBuffer = mediapipe::GpuBuffer;
#endif // MEDIAPIPE_DISABLE_GPU
// Converts image into Tensor, possibly with cropping, resizing and
// normalization, according to specified inputs and options.
@@ -110,9 +110,21 @@ namespace mediapipe {
// }
// }
// }
class ImageToTensorCalculator : public CalculatorBase {
class ImageToTensorCalculator : public Node {
public:
static mediapipe::Status GetContract(CalculatorContract* cc) {
static constexpr Input<mediapipe::ImageFrame>::Optional kInCpu{"IMAGE"};
static constexpr Input<GpuBuffer>::Optional kInGpu{"IMAGE_GPU"};
static constexpr Input<mediapipe::NormalizedRect>::Optional kInNormRect{
"NORM_RECT"};
static constexpr Output<std::vector<Tensor>> kOutTensors{"TENSORS"};
static constexpr Output<std::array<float, 4>>::Optional kOutLetterboxPadding{
"LETTERBOX_PADDING"};
static constexpr Output<std::array<float, 16>>::Optional kOutMatrix{"MATRIX"};
MEDIAPIPE_NODE_CONTRACT(kInCpu, kInGpu, kInNormRect, kOutTensors,
kOutLetterboxPadding, kOutMatrix);
static ::mediapipe::Status UpdateContract(CalculatorContract* cc) {
const auto& options =
cc->Options<mediapipe::ImageToTensorCalculatorOptions>();
@@ -126,24 +138,10 @@ class ImageToTensorCalculator : public CalculatorBase {
RET_CHECK_GT(options.output_tensor_height(), 0)
<< "Valid output tensor height is required.";
if (cc->Inputs().HasTag(kInputNormRect)) {
cc->Inputs().Tag(kInputNormRect).Set<mediapipe::NormalizedRect>();
}
if (cc->Outputs().HasTag(kOutputLetterboxPadding)) {
cc->Outputs().Tag(kOutputLetterboxPadding).Set<std::array<float, 4>>();
}
if (cc->Outputs().HasTag(kOutputMatrix)) {
cc->Outputs().Tag(kOutputMatrix).Set<std::array<float, 16>>();
}
RET_CHECK(kInCpu(cc).IsConnected() ^ kInGpu(cc).IsConnected())
<< "One and only one of CPU or GPU input is expected.";
const bool has_cpu_input = cc->Inputs().HasTag(kInputCpu);
const bool has_gpu_input = cc->Inputs().HasTag(kInputGpu);
RET_CHECK_EQ((has_cpu_input ? 1 : 0) + (has_gpu_input ? 1 : 0), 1)
<< "Either CPU or GPU input is expected, not both.";
if (has_cpu_input) {
cc->Inputs().Tag(kInputCpu).Set<mediapipe::ImageFrame>();
} else if (has_gpu_input) {
if (kInGpu(cc).IsConnected()) {
#if MEDIAPIPE_DISABLE_GPU
return mediapipe::UnimplementedError("GPU processing is disabled");
#else
@@ -153,25 +151,20 @@ class ImageToTensorCalculator : public CalculatorBase {
#else
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
#endif // MEDIAPIPE_METAL_ENABLED
cc->Inputs().Tag(kInputGpu).Set<mediapipe::GpuBuffer>();
#endif // MEDIAPIPE_DISABLE_GPU
}
cc->Outputs().Tag(kOutput).Set<std::vector<Tensor>>();
return mediapipe::OkStatus();
}
mediapipe::Status Open(CalculatorContext* cc) {
// Makes sure outputs' next timestamp bound update is handled automatically
// by the framework.
cc->SetOffset(TimestampDiff(0));
options_ = cc->Options<mediapipe::ImageToTensorCalculatorOptions>();
output_width_ = options_.output_tensor_width();
output_height_ = options_.output_tensor_height();
range_min_ = options_.output_tensor_float_range().min();
range_max_ = options_.output_tensor_float_range().max();
if (cc->Inputs().HasTag(kInputCpu)) {
if (kInCpu(cc).IsConnected()) {
ASSIGN_OR_RETURN(converter_, CreateOpenCvConverter(cc, GetBorderMode()));
} else {
#if MEDIAPIPE_DISABLE_GPU
@@ -196,21 +189,20 @@ class ImageToTensorCalculator : public CalculatorBase {
}
mediapipe::Status Process(CalculatorContext* cc) {
const InputStreamShard& input = cc->Inputs().Tag(
cc->Inputs().HasTag(kInputCpu) ? kInputCpu : kInputGpu);
if (input.IsEmpty()) {
const PacketBase& image_packet =
kInCpu(cc).IsConnected() ? kInCpu(cc).packet() : kInGpu(cc).packet();
if (image_packet.IsEmpty()) {
// Timestamp bound update happens automatically. (See Open().)
return mediapipe::OkStatus();
}
absl::optional<mediapipe::NormalizedRect> norm_rect;
if (cc->Inputs().HasTag(kInputNormRect)) {
if (cc->Inputs().Tag(kInputNormRect).IsEmpty()) {
if (kInNormRect(cc).IsConnected()) {
if (kInNormRect(cc).IsEmpty()) {
// Timestamp bound update happens automatically. (See Open().)
return mediapipe::OkStatus();
}
norm_rect =
cc->Inputs().Tag(kInputNormRect).Get<mediapipe::NormalizedRect>();
norm_rect = *kInNormRect(cc);
if (norm_rect->width() == 0 && norm_rect->height() == 0) {
// WORKAROUND: some existing graphs may use sentinel rects {width=0,
// height=0, ...} quite often and calculator has to handle them
@@ -223,27 +215,20 @@ class ImageToTensorCalculator : public CalculatorBase {
}
}
const Packet& image_packet = input.Value();
const Size& size = converter_->GetImageSize(image_packet);
RotatedRect roi = GetRoi(size.width, size.height, norm_rect);
ASSIGN_OR_RETURN(auto padding, PadRoi(options_.output_tensor_width(),
options_.output_tensor_height(),
options_.keep_aspect_ratio(), &roi));
if (cc->Outputs().HasTag(kOutputLetterboxPadding)) {
cc->Outputs()
.Tag(kOutputLetterboxPadding)
.AddPacket(MakePacket<std::array<float, 4>>(padding).At(
cc->InputTimestamp()));
if (kOutLetterboxPadding(cc).IsConnected()) {
kOutLetterboxPadding(cc).Send(padding);
}
if (cc->Outputs().HasTag(kOutputMatrix)) {
if (kOutMatrix(cc).IsConnected()) {
std::array<float, 16> matrix;
GetRotatedSubRectToRectTransformMatrix(roi, size.width, size.height,
/*flip_horizontaly=*/false,
&matrix);
cc->Outputs()
.Tag(kOutputMatrix)
.AddPacket(MakePacket<std::array<float, 16>>(std::move(matrix))
.At(cc->InputTimestamp()));
kOutMatrix(cc).Send(std::move(matrix));
}
ASSIGN_OR_RETURN(
@@ -251,11 +236,9 @@ class ImageToTensorCalculator : public CalculatorBase {
converter_->Convert(image_packet, roi, {output_width_, output_height_},
range_min_, range_max_));
std::vector<Tensor> result;
result.push_back(std::move(tensor));
cc->Outputs().Tag(kOutput).AddPacket(
MakePacket<std::vector<Tensor>>(std::move(result))
.At(cc->InputTimestamp()));
auto result = std::make_unique<std::vector<Tensor>>();
result->push_back(std::move(tensor));
kOutTensors(cc).Send(std::move(result));
return mediapipe::OkStatus();
}
@@ -286,6 +269,7 @@ class ImageToTensorCalculator : public CalculatorBase {
float range_max_ = 1.0f;
};
REGISTER_CALCULATOR(ImageToTensorCalculator);
MEDIAPIPE_REGISTER_NODE(ImageToTensorCalculator);
} // namespace api2
} // namespace mediapipe
@@ -19,6 +19,7 @@
#include "absl/memory/memory.h"
#include "mediapipe/calculators/tensor/inference_calculator.pb.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/ret_check.h"
@@ -88,7 +89,6 @@ bool ShouldUseGpu(const mediapipe::InferenceCalculatorOptions& options) {
}
constexpr char kTensorsTag[] = "TENSORS";
} // namespace
#if defined(MEDIAPIPE_EDGE_TPU)
#include "edgetpu.h"
@@ -112,7 +112,10 @@ std::unique_ptr<tflite::Interpreter> BuildEdgeTpuInterpreter(
}
#endif // MEDIAPIPE_EDGE_TPU
} // namespace
namespace mediapipe {
namespace api2 {
#if MEDIAPIPE_TFLITE_METAL_INFERENCE
namespace {
@@ -224,12 +227,19 @@ int GetXnnpackNumThreads(const mediapipe::InferenceCalculatorOptions& opts) {
// Tensors are assumed to be ordered correctly (sequentially added to model).
// Input tensors are assumed to be of the correct size and already normalized.
class InferenceCalculator : public CalculatorBase {
class InferenceCalculator : public Node {
public:
using TfLiteDelegatePtr =
std::unique_ptr<TfLiteDelegate, std::function<void(TfLiteDelegate*)>>;
static mediapipe::Status GetContract(CalculatorContract* cc);
static constexpr Input<std::vector<Tensor>> kInTensors{"TENSORS"};
static constexpr SideInput<tflite::ops::builtin::BuiltinOpResolver>::Optional
kSideInCustomOpResolver{"CUSTOM_OP_RESOLVER"};
static constexpr SideInput<TfLiteModelPtr>::Optional kSideInModel{"MODEL"};
static constexpr Output<std::vector<Tensor>> kOutTensors{"TENSORS"};
MEDIAPIPE_NODE_CONTRACT(kInTensors, kSideInCustomOpResolver, kSideInModel,
kOutTensors);
static mediapipe::Status UpdateContract(CalculatorContract* cc);
mediapipe::Status Open(CalculatorContext* cc) override;
mediapipe::Status Process(CalculatorContext* cc) override;
@@ -239,11 +249,12 @@ class InferenceCalculator : public CalculatorBase {
mediapipe::Status ReadKernelsFromFile();
mediapipe::Status WriteKernelsToFile();
mediapipe::Status LoadModel(CalculatorContext* cc);
mediapipe::StatusOr<Packet> GetModelAsPacket(const CalculatorContext& cc);
mediapipe::StatusOr<mediapipe::Packet> GetModelAsPacket(
const CalculatorContext& cc);
mediapipe::Status LoadDelegate(CalculatorContext* cc);
mediapipe::Status InitTFLiteGPURunner(CalculatorContext* cc);
Packet model_packet_;
mediapipe::Packet model_packet_;
std::unique_ptr<tflite::Interpreter> interpreter_;
TfLiteDelegatePtr delegate_;
@@ -277,28 +288,13 @@ class InferenceCalculator : public CalculatorBase {
std::string cached_kernel_filename_;
};
REGISTER_CALCULATOR(InferenceCalculator);
mediapipe::Status InferenceCalculator::GetContract(CalculatorContract* cc) {
RET_CHECK(cc->Inputs().HasTag(kTensorsTag));
cc->Inputs().Tag(kTensorsTag).Set<std::vector<Tensor>>();
RET_CHECK(cc->Outputs().HasTag(kTensorsTag));
cc->Outputs().Tag(kTensorsTag).Set<std::vector<Tensor>>();
MEDIAPIPE_REGISTER_NODE(InferenceCalculator);
mediapipe::Status InferenceCalculator::UpdateContract(CalculatorContract* cc) {
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
RET_CHECK(!options.model_path().empty() ^
cc->InputSidePackets().HasTag("MODEL"))
RET_CHECK(!options.model_path().empty() ^ kSideInModel(cc).IsConnected())
<< "Either model as side packet or model path in options is required.";
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
cc->InputSidePackets()
.Tag("CUSTOM_OP_RESOLVER")
.Set<tflite::ops::builtin::BuiltinOpResolver>();
}
if (cc->InputSidePackets().HasTag("MODEL")) {
cc->InputSidePackets().Tag("MODEL").Set<TfLiteModelPtr>();
}
if (ShouldUseGpu(options)) {
#if MEDIAPIPE_TFLITE_GL_INFERENCE
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
@@ -310,8 +306,6 @@ mediapipe::Status InferenceCalculator::GetContract(CalculatorContract* cc) {
}
mediapipe::Status InferenceCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
#if MEDIAPIPE_TFLITE_GL_INFERENCE || MEDIAPIPE_TFLITE_METAL_INFERENCE
const auto& options = cc->Options<::mediapipe::InferenceCalculatorOptions>();
if (ShouldUseGpu(options)) {
@@ -361,11 +355,10 @@ mediapipe::Status InferenceCalculator::Open(CalculatorContext* cc) {
}
mediapipe::Status InferenceCalculator::Process(CalculatorContext* cc) {
if (cc->Inputs().Tag(kTensorsTag).IsEmpty()) {
if (kInTensors(cc).IsEmpty()) {
return mediapipe::OkStatus();
}
const auto& input_tensors =
cc->Inputs().Tag(kTensorsTag).Get<std::vector<Tensor>>();
const auto& input_tensors = *kInTensors(cc);
RET_CHECK(!input_tensors.empty());
auto output_tensors = absl::make_unique<std::vector<Tensor>>();
#if MEDIAPIPE_TFLITE_METAL_INFERENCE
@@ -509,9 +502,7 @@ mediapipe::Status InferenceCalculator::Process(CalculatorContext* cc) {
output_tensors->back().bytes());
}
}
cc->Outputs()
.Tag(kTensorsTag)
.Add(output_tensors.release(), cc->InputTimestamp());
kOutTensors(cc).Send(std::move(output_tensors));
return mediapipe::OkStatus();
}
@@ -575,12 +566,9 @@ mediapipe::Status InferenceCalculator::InitTFLiteGPURunner(
#if MEDIAPIPE_TFLITE_GL_INFERENCE
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(*cc));
const auto& model = *model_packet_.Get<TfLiteModelPtr>();
tflite::ops::builtin::BuiltinOpResolver op_resolver;
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
op_resolver = cc->InputSidePackets()
.Tag("CUSTOM_OP_RESOLVER")
.Get<tflite::ops::builtin::BuiltinOpResolver>();
}
tflite::ops::builtin::BuiltinOpResolver op_resolver =
kSideInCustomOpResolver(cc).GetOr(
tflite::ops::builtin::BuiltinOpResolver());
// Create runner
tflite::gpu::InferenceOptions options;
@@ -629,12 +617,9 @@ mediapipe::Status InferenceCalculator::InitTFLiteGPURunner(
mediapipe::Status InferenceCalculator::LoadModel(CalculatorContext* cc) {
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(*cc));
const auto& model = *model_packet_.Get<TfLiteModelPtr>();
tflite::ops::builtin::BuiltinOpResolver op_resolver;
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
op_resolver = cc->InputSidePackets()
.Tag("CUSTOM_OP_RESOLVER")
.Get<tflite::ops::builtin::BuiltinOpResolver>();
}
tflite::ops::builtin::BuiltinOpResolver op_resolver =
kSideInCustomOpResolver(cc).GetOr(
tflite::ops::builtin::BuiltinOpResolver());
#if defined(MEDIAPIPE_EDGE_TPU)
interpreter_ =
@@ -659,7 +644,7 @@ mediapipe::Status InferenceCalculator::LoadModel(CalculatorContext* cc) {
return mediapipe::OkStatus();
}
mediapipe::StatusOr<Packet> InferenceCalculator::GetModelAsPacket(
mediapipe::StatusOr<mediapipe::Packet> InferenceCalculator::GetModelAsPacket(
const CalculatorContext& cc) {
const auto& options = cc.Options<mediapipe::InferenceCalculatorOptions>();
if (!options.model_path().empty()) {
@@ -845,4 +830,5 @@ mediapipe::Status InferenceCalculator::LoadDelegate(CalculatorContext* cc) {
return mediapipe::OkStatus();
}
} // namespace api2
} // namespace mediapipe
@@ -19,6 +19,7 @@
#include "absl/strings/str_format.h"
#include "absl/types/span.h"
#include "mediapipe/calculators/tensor/tensors_to_classification_calculator.pb.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/classification.pb.h"
#include "mediapipe/framework/formats/tensor.h"
@@ -32,6 +33,7 @@
#endif
namespace mediapipe {
namespace api2 {
// Convert result tensors from classification models into MediaPipe
// classifications.
@@ -57,9 +59,12 @@ namespace mediapipe {
// }
// }
// }
class TensorsToClassificationCalculator : public CalculatorBase {
class TensorsToClassificationCalculator : public Node {
public:
static mediapipe::Status GetContract(CalculatorContract* cc);
static constexpr Input<std::vector<Tensor>> kInTensors{"TENSORS"};
static constexpr Output<ClassificationList> kOutClassificationList{
"CLASSIFICATIONS"};
MEDIAPIPE_NODE_CONTRACT(kInTensors, kOutClassificationList);
mediapipe::Status Open(CalculatorContext* cc) override;
mediapipe::Status Process(CalculatorContext* cc) override;
@@ -71,28 +76,10 @@ class TensorsToClassificationCalculator : public CalculatorBase {
std::unordered_map<int, std::string> label_map_;
bool label_map_loaded_ = false;
};
REGISTER_CALCULATOR(TensorsToClassificationCalculator);
mediapipe::Status TensorsToClassificationCalculator::GetContract(
CalculatorContract* cc) {
RET_CHECK(!cc->Inputs().GetTags().empty());
RET_CHECK(!cc->Outputs().GetTags().empty());
if (cc->Inputs().HasTag("TENSORS")) {
cc->Inputs().Tag("TENSORS").Set<std::vector<Tensor>>();
}
if (cc->Outputs().HasTag("CLASSIFICATIONS")) {
cc->Outputs().Tag("CLASSIFICATIONS").Set<ClassificationList>();
}
return mediapipe::OkStatus();
}
MEDIAPIPE_REGISTER_NODE(TensorsToClassificationCalculator);
mediapipe::Status TensorsToClassificationCalculator::Open(
CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
options_ =
cc->Options<::mediapipe::TensorsToClassificationCalculatorOptions>();
@@ -118,9 +105,7 @@ mediapipe::Status TensorsToClassificationCalculator::Open(
mediapipe::Status TensorsToClassificationCalculator::Process(
CalculatorContext* cc) {
const auto& input_tensors =
cc->Inputs().Tag("TENSORS").Get<std::vector<Tensor>>();
const auto& input_tensors = *kInTensors(cc);
RET_CHECK_EQ(input_tensors.size(), 1);
int num_classes = input_tensors[0].shape().num_elements();
@@ -182,10 +167,7 @@ mediapipe::Status TensorsToClassificationCalculator::Process(
raw_classification_list->DeleteSubrange(
top_k_, raw_classification_list->size() - top_k_);
}
cc->Outputs()
.Tag("CLASSIFICATIONS")
.Add(classification_list.release(), cc->InputTimestamp());
kOutClassificationList(cc).Send(std::move(classification_list));
return mediapipe::OkStatus();
}
@@ -194,4 +176,5 @@ mediapipe::Status TensorsToClassificationCalculator::Close(
return mediapipe::OkStatus();
}
} // namespace api2
} // namespace mediapipe
@@ -18,6 +18,7 @@
#include "absl/strings/str_format.h"
#include "absl/types/span.h"
#include "mediapipe/calculators/tensor/tensors_to_detections_calculator.pb.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/deps/file_path.h"
#include "mediapipe/framework/formats/detection.pb.h"
@@ -47,9 +48,6 @@
namespace {
constexpr int kNumInputTensorsWithAnchors = 3;
constexpr int kNumCoordsPerBox = 4;
constexpr char kDetectionsTag[] = "DETECTIONS";
constexpr char kTensorsTag[] = "TENSORS";
constexpr char kAnchorsTag[] = "ANCHORS";
bool CanUseGpu() {
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE) || MEDIAPIPE_METAL_ENABLED
@@ -63,6 +61,7 @@ bool CanUseGpu() {
} // namespace
namespace mediapipe {
namespace api2 {
namespace {
@@ -128,9 +127,14 @@ void ConvertAnchorsToRawValues(const std::vector<Anchor>& anchors,
// }
// }
// }
class TensorsToDetectionsCalculator : public CalculatorBase {
class TensorsToDetectionsCalculator : public Node {
public:
static mediapipe::Status GetContract(CalculatorContract* cc);
static constexpr Input<std::vector<Tensor>> kInTensors{"TENSORS"};
static constexpr SideInput<std::vector<Anchor>>::Optional kInAnchors{
"ANCHORS"};
static constexpr Output<std::vector<Detection>> kOutDetections{"DETECTIONS"};
MEDIAPIPE_NODE_CONTRACT(kInTensors, kInAnchors, kOutDetections);
static mediapipe::Status UpdateContract(CalculatorContract* cc);
mediapipe::Status Open(CalculatorContext* cc) override;
mediapipe::Status Process(CalculatorContext* cc) override;
@@ -161,7 +165,6 @@ class TensorsToDetectionsCalculator : public CalculatorBase {
::mediapipe::TensorsToDetectionsCalculatorOptions options_;
std::vector<Anchor> anchors_;
bool side_packet_anchors_{};
#ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
mediapipe::GlCalculatorHelper gpu_helper_;
@@ -179,22 +182,10 @@ class TensorsToDetectionsCalculator : public CalculatorBase {
bool gpu_input_ = false;
bool anchors_init_ = false;
};
REGISTER_CALCULATOR(TensorsToDetectionsCalculator);
MEDIAPIPE_REGISTER_NODE(TensorsToDetectionsCalculator);
mediapipe::Status TensorsToDetectionsCalculator::GetContract(
mediapipe::Status TensorsToDetectionsCalculator::UpdateContract(
CalculatorContract* cc) {
RET_CHECK(cc->Inputs().HasTag(kTensorsTag));
cc->Inputs().Tag(kTensorsTag).Set<std::vector<Tensor>>();
RET_CHECK(cc->Outputs().HasTag(kDetectionsTag));
cc->Outputs().Tag(kDetectionsTag).Set<std::vector<Detection>>();
if (cc->InputSidePackets().UsesTags()) {
if (cc->InputSidePackets().HasTag(kAnchorsTag)) {
cc->InputSidePackets().Tag(kAnchorsTag).Set<std::vector<Anchor>>();
}
}
if (CanUseGpu()) {
#ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
@@ -207,8 +198,6 @@ mediapipe::Status TensorsToDetectionsCalculator::GetContract(
}
mediapipe::Status TensorsToDetectionsCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
side_packet_anchors_ = cc->InputSidePackets().HasTag(kAnchorsTag);
MP_RETURN_IF_ERROR(LoadOptions(cc));
if (CanUseGpu()) {
@@ -226,18 +215,12 @@ mediapipe::Status TensorsToDetectionsCalculator::Open(CalculatorContext* cc) {
mediapipe::Status TensorsToDetectionsCalculator::Process(
CalculatorContext* cc) {
if (cc->Inputs().Tag(kTensorsTag).IsEmpty()) {
return mediapipe::OkStatus();
}
auto output_detections = absl::make_unique<std::vector<Detection>>();
bool gpu_processing = false;
if (CanUseGpu()) {
// Use GPU processing only if at least one input tensor is already on GPU
// (to avoid CPU->GPU overhead).
for (const auto& tensor :
cc->Inputs().Tag(kTensorsTag).Get<std::vector<Tensor>>()) {
for (const auto& tensor : *kInTensors(cc)) {
if (tensor.ready_on_gpu()) {
gpu_processing = true;
break;
@@ -251,18 +234,13 @@ mediapipe::Status TensorsToDetectionsCalculator::Process(
MP_RETURN_IF_ERROR(ProcessCPU(cc, output_detections.get()));
}
// Output
cc->Outputs()
.Tag(kDetectionsTag)
.Add(output_detections.release(), cc->InputTimestamp());
kOutDetections(cc).Send(std::move(output_detections));
return mediapipe::OkStatus();
}
mediapipe::Status TensorsToDetectionsCalculator::ProcessCPU(
CalculatorContext* cc, std::vector<Detection>* output_detections) {
const auto& input_tensors =
cc->Inputs().Tag(kTensorsTag).Get<std::vector<Tensor>>();
const auto& input_tensors = *kInTensors(cc);
if (input_tensors.size() == 2 ||
input_tensors.size() == kNumInputTensorsWithAnchors) {
@@ -294,10 +272,8 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessCPU(
auto anchor_view = anchor_tensor->GetCpuReadView();
auto raw_anchors = anchor_view.buffer<float>();
ConvertRawValuesToAnchors(raw_anchors, num_boxes_, &anchors_);
} else if (side_packet_anchors_) {
CHECK(!cc->InputSidePackets().Tag("ANCHORS").IsEmpty());
anchors_ =
cc->InputSidePackets().Tag("ANCHORS").Get<std::vector<Anchor>>();
} else if (!kInAnchors(cc).IsEmpty()) {
anchors_ = *kInAnchors(cc);
} else {
return mediapipe::UnavailableError("No anchor data available.");
}
@@ -391,8 +367,7 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessCPU(
mediapipe::Status TensorsToDetectionsCalculator::ProcessGPU(
CalculatorContext* cc, std::vector<Detection>* output_detections) {
const auto& input_tensors =
cc->Inputs().Tag(kTensorsTag).Get<std::vector<Tensor>>();
const auto& input_tensors = *kInTensors(cc);
RET_CHECK_GE(input_tensors.size(), 2);
#ifndef MEDIAPIPE_DISABLE_GL_COMPUTE
@@ -400,21 +375,20 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessGPU(
&output_detections]()
-> mediapipe::Status {
if (!anchors_init_) {
if (side_packet_anchors_) {
CHECK(!cc->InputSidePackets().Tag(kAnchorsTag).IsEmpty());
const auto& anchors =
cc->InputSidePackets().Tag(kAnchorsTag).Get<std::vector<Anchor>>();
auto anchors_view = raw_anchors_buffer_->GetCpuWriteView();
auto raw_anchors = anchors_view.buffer<float>();
ConvertAnchorsToRawValues(anchors, num_boxes_, raw_anchors);
} else {
CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
if (input_tensors.size() == kNumInputTensorsWithAnchors) {
auto read_view = input_tensors[2].GetOpenGlBufferReadView();
glBindBuffer(GL_COPY_READ_BUFFER, read_view.name());
auto write_view = raw_anchors_buffer_->GetOpenGlBufferWriteView();
glBindBuffer(GL_COPY_WRITE_BUFFER, write_view.name());
glCopyBufferSubData(GL_COPY_READ_BUFFER, GL_COPY_WRITE_BUFFER, 0, 0,
input_tensors[2].bytes());
} else if (!kInAnchors(cc).IsEmpty()) {
const auto& anchors = *kInAnchors(cc);
auto anchors_view = raw_anchors_buffer_->GetCpuWriteView();
auto raw_anchors = anchors_view.buffer<float>();
ConvertAnchorsToRawValues(anchors, num_boxes_, raw_anchors);
} else {
return mediapipe::UnavailableError("No anchor data available.");
}
anchors_init_ = true;
}
@@ -464,14 +438,7 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessGPU(
#elif MEDIAPIPE_METAL_ENABLED
id<MTLDevice> device = gpu_helper_.mtlDevice;
if (!anchors_init_) {
if (side_packet_anchors_) {
CHECK(!cc->InputSidePackets().Tag(kAnchorsTag).IsEmpty());
const auto& anchors =
cc->InputSidePackets().Tag(kAnchorsTag).Get<std::vector<Anchor>>();
auto raw_anchors_view = raw_anchors_buffer_->GetCpuWriteView();
ConvertAnchorsToRawValues(anchors, num_boxes_,
raw_anchors_view.buffer<float>());
} else {
if (input_tensors.size() == kNumInputTensorsWithAnchors) {
RET_CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
auto command_buffer = [gpu_helper_ commandBuffer];
auto src_buffer = input_tensors[2].GetMtlBufferReadView(command_buffer);
@@ -486,6 +453,13 @@ mediapipe::Status TensorsToDetectionsCalculator::ProcessGPU(
size:input_tensors[2].bytes()];
[blit_command endEncoding];
[command_buffer commit];
} else if (!kInAnchors(cc).IsEmpty()) {
const auto& anchors = *kInAnchors(cc);
auto raw_anchors_view = raw_anchors_buffer_->GetCpuWriteView();
ConvertAnchorsToRawValues(anchors, num_boxes_,
raw_anchors_view.buffer<float>());
} else {
return mediapipe::UnavailableError("No anchor data available.");
}
anchors_init_ = true;
}
@@ -1157,4 +1131,5 @@ kernel void scoreKernel(
return mediapipe::OkStatus();
}
} // namespace api2
} // namespace mediapipe
@@ -13,6 +13,7 @@
// limitations under the License.
#include "mediapipe/calculators/tensor/tensors_to_floats_calculator.pb.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/ret_check.h"
@@ -43,47 +44,39 @@ inline float Sigmoid(float value) { return 1.0f / (1.0f + std::exp(-value)); }
// input_stream: "TENSORS:tensors"
// output_stream: "FLOATS:floats"
// }
class TensorsToFloatsCalculator : public CalculatorBase {
namespace api2 {
class TensorsToFloatsCalculator : public Node {
public:
static mediapipe::Status GetContract(CalculatorContract* cc);
static constexpr Input<std::vector<Tensor>> kInTensors{"TENSORS"};
static constexpr Output<float>::Optional kOutFloat{"FLOAT"};
static constexpr Output<std::vector<float>>::Optional kOutFloats{"FLOATS"};
MEDIAPIPE_NODE_INTERFACE(TensorsToFloatsCalculator, kInTensors, kOutFloat,
kOutFloats);
mediapipe::Status Open(CalculatorContext* cc) override;
mediapipe::Status Process(CalculatorContext* cc) override;
static mediapipe::Status UpdateContract(CalculatorContract* cc);
mediapipe::Status Open(CalculatorContext* cc) final;
mediapipe::Status Process(CalculatorContext* cc) final;
private:
::mediapipe::TensorsToFloatsCalculatorOptions options_;
};
REGISTER_CALCULATOR(TensorsToFloatsCalculator);
MEDIAPIPE_REGISTER_NODE(TensorsToFloatsCalculator);
mediapipe::Status TensorsToFloatsCalculator::GetContract(
mediapipe::Status TensorsToFloatsCalculator::UpdateContract(
CalculatorContract* cc) {
RET_CHECK(cc->Inputs().HasTag("TENSORS"));
RET_CHECK(cc->Outputs().HasTag("FLOATS") || cc->Outputs().HasTag("FLOAT"));
cc->Inputs().Tag("TENSORS").Set<std::vector<Tensor>>();
if (cc->Outputs().HasTag("FLOATS")) {
cc->Outputs().Tag("FLOATS").Set<std::vector<float>>();
}
if (cc->Outputs().HasTag("FLOAT")) {
cc->Outputs().Tag("FLOAT").Set<float>();
}
// Only exactly a single output allowed.
RET_CHECK(kOutFloat(cc).IsConnected() ^ kOutFloats(cc).IsConnected());
return mediapipe::OkStatus();
}
mediapipe::Status TensorsToFloatsCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
options_ = cc->Options<::mediapipe::TensorsToFloatsCalculatorOptions>();
return mediapipe::OkStatus();
}
mediapipe::Status TensorsToFloatsCalculator::Process(CalculatorContext* cc) {
RET_CHECK(!cc->Inputs().Tag("TENSORS").IsEmpty());
const auto& input_tensors =
cc->Inputs().Tag("TENSORS").Get<std::vector<Tensor>>();
const auto& input_tensors = *kInTensors(cc);
RET_CHECK(!input_tensors.empty());
// TODO: Add option to specify which tensor to take from.
auto view = input_tensors[0].GetCpuReadView();
auto raw_floats = view.buffer<float>();
@@ -100,18 +93,15 @@ mediapipe::Status TensorsToFloatsCalculator::Process(CalculatorContext* cc) {
break;
}
if (cc->Outputs().HasTag("FLOAT")) {
// TODO: Could add an index in the option to specifiy returning one
// value of a float array.
if (kOutFloat(cc).IsConnected()) {
// TODO: Could add an index in the option to specifiy returning
// one value of a float array.
RET_CHECK_EQ(num_values, 1);
cc->Outputs().Tag("FLOAT").AddPacket(
MakePacket<float>(output_floats->at(0)).At(cc->InputTimestamp()));
kOutFloat(cc).Send(output_floats->at(0));
} else {
kOutFloats(cc).Send(std::move(output_floats));
}
if (cc->Outputs().HasTag("FLOATS")) {
cc->Outputs().Tag("FLOATS").Add(output_floats.release(),
cc->InputTimestamp());
}
return mediapipe::OkStatus();
}
} // namespace api2
} // namespace mediapipe
@@ -13,12 +13,14 @@
// limitations under the License.
#include "mediapipe/calculators/tensor/tensors_to_landmarks_calculator.pb.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/landmark.pb.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/framework/port/ret_check.h"
namespace mediapipe {
namespace api2 {
namespace {
@@ -85,9 +87,18 @@ float ApplyActivation(
// }
// }
// }
class TensorsToLandmarksCalculator : public CalculatorBase {
class TensorsToLandmarksCalculator : public Node {
public:
static mediapipe::Status GetContract(CalculatorContract* cc);
static constexpr Input<std::vector<Tensor>> kInTensors{"TENSORS"};
static constexpr Input<bool>::SideFallback::Optional kFlipHorizontally{
"FLIP_HORIZONTALLY"};
static constexpr Input<bool>::SideFallback::Optional kFlipVertically{
"FLIP_VERTICALLY"};
static constexpr Output<LandmarkList>::Optional kOutLandmarkList{"LANDMARKS"};
static constexpr Output<NormalizedLandmarkList>::Optional
kOutNormalizedLandmarkList{"NORM_LANDMARKS"};
MEDIAPIPE_NODE_CONTRACT(kInTensors, kFlipHorizontally, kFlipVertically,
kOutLandmarkList, kOutNormalizedLandmarkList);
mediapipe::Status Open(CalculatorContext* cc) override;
mediapipe::Status Process(CalculatorContext* cc) override;
@@ -95,100 +106,39 @@ class TensorsToLandmarksCalculator : public CalculatorBase {
private:
mediapipe::Status LoadOptions(CalculatorContext* cc);
int num_landmarks_ = 0;
bool flip_vertically_ = false;
bool flip_horizontally_ = false;
::mediapipe::TensorsToLandmarksCalculatorOptions options_;
};
REGISTER_CALCULATOR(TensorsToLandmarksCalculator);
mediapipe::Status TensorsToLandmarksCalculator::GetContract(
CalculatorContract* cc) {
RET_CHECK(!cc->Inputs().GetTags().empty());
RET_CHECK(!cc->Outputs().GetTags().empty());
if (cc->Inputs().HasTag("TENSORS")) {
cc->Inputs().Tag("TENSORS").Set<std::vector<Tensor>>();
}
if (cc->Inputs().HasTag("FLIP_HORIZONTALLY")) {
cc->Inputs().Tag("FLIP_HORIZONTALLY").Set<bool>();
}
if (cc->Inputs().HasTag("FLIP_VERTICALLY")) {
cc->Inputs().Tag("FLIP_VERTICALLY").Set<bool>();
}
if (cc->InputSidePackets().HasTag("FLIP_HORIZONTALLY")) {
cc->InputSidePackets().Tag("FLIP_HORIZONTALLY").Set<bool>();
}
if (cc->InputSidePackets().HasTag("FLIP_VERTICALLY")) {
cc->InputSidePackets().Tag("FLIP_VERTICALLY").Set<bool>();
}
if (cc->Outputs().HasTag("LANDMARKS")) {
cc->Outputs().Tag("LANDMARKS").Set<LandmarkList>();
}
if (cc->Outputs().HasTag("NORM_LANDMARKS")) {
cc->Outputs().Tag("NORM_LANDMARKS").Set<NormalizedLandmarkList>();
}
return mediapipe::OkStatus();
}
MEDIAPIPE_REGISTER_NODE(TensorsToLandmarksCalculator);
mediapipe::Status TensorsToLandmarksCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
MP_RETURN_IF_ERROR(LoadOptions(cc));
if (cc->Outputs().HasTag("NORM_LANDMARKS")) {
if (kOutNormalizedLandmarkList(cc).IsConnected()) {
RET_CHECK(options_.has_input_image_height() &&
options_.has_input_image_width())
<< "Must provide input with/height for getting normalized landmarks.";
}
if (cc->Outputs().HasTag("LANDMARKS") &&
(options_.flip_vertically() || options_.flip_horizontally() ||
cc->InputSidePackets().HasTag("FLIP_HORIZONTALLY") ||
cc->InputSidePackets().HasTag("FLIP_VERTICALLY"))) {
if (kOutLandmarkList(cc).IsConnected() &&
(options_.flip_horizontally() || options_.flip_vertically() ||
kFlipHorizontally(cc).IsConnected() ||
kFlipVertically(cc).IsConnected())) {
RET_CHECK(options_.has_input_image_height() &&
options_.has_input_image_width())
<< "Must provide input with/height for using flip_vertically option "
"when outputing landmarks in absolute coordinates.";
<< "Must provide input with/height for using flipping when outputing "
"landmarks in absolute coordinates.";
}
flip_horizontally_ =
cc->InputSidePackets().HasTag("FLIP_HORIZONTALLY")
? cc->InputSidePackets().Tag("FLIP_HORIZONTALLY").Get<bool>()
: options_.flip_horizontally();
flip_vertically_ =
cc->InputSidePackets().HasTag("FLIP_VERTICALLY")
? cc->InputSidePackets().Tag("FLIP_VERTICALLY").Get<bool>()
: options_.flip_vertically();
return mediapipe::OkStatus();
}
mediapipe::Status TensorsToLandmarksCalculator::Process(CalculatorContext* cc) {
// Override values if specified so.
if (cc->Inputs().HasTag("FLIP_HORIZONTALLY") &&
!cc->Inputs().Tag("FLIP_HORIZONTALLY").IsEmpty()) {
flip_horizontally_ = cc->Inputs().Tag("FLIP_HORIZONTALLY").Get<bool>();
}
if (cc->Inputs().HasTag("FLIP_VERTICALLY") &&
!cc->Inputs().Tag("FLIP_VERTICALLY").IsEmpty()) {
flip_vertically_ = cc->Inputs().Tag("FLIP_VERTICALLY").Get<bool>();
}
if (cc->Inputs().Tag("TENSORS").IsEmpty()) {
if (kInTensors(cc).IsEmpty()) {
return mediapipe::OkStatus();
}
bool flip_horizontally =
kFlipHorizontally(cc).GetOr(options_.flip_horizontally());
bool flip_vertically = kFlipVertically(cc).GetOr(options_.flip_vertically());
const auto& input_tensors =
cc->Inputs().Tag("TENSORS").Get<std::vector<Tensor>>();
const auto& input_tensors = *kInTensors(cc);
int num_values = input_tensors[0].shape().num_elements();
const int num_dimensions = num_values / num_landmarks_;
CHECK_GT(num_dimensions, 0);
@@ -202,13 +152,13 @@ mediapipe::Status TensorsToLandmarksCalculator::Process(CalculatorContext* cc) {
const int offset = ld * num_dimensions;
Landmark* landmark = output_landmarks.add_landmark();
if (flip_horizontally_) {
if (flip_horizontally) {
landmark->set_x(options_.input_image_width() - raw_landmarks[offset]);
} else {
landmark->set_x(raw_landmarks[offset]);
}
if (num_dimensions > 1) {
if (flip_vertically_) {
if (flip_vertically) {
landmark->set_y(options_.input_image_height() -
raw_landmarks[offset + 1]);
} else {
@@ -229,7 +179,7 @@ mediapipe::Status TensorsToLandmarksCalculator::Process(CalculatorContext* cc) {
}
// Output normalized landmarks if required.
if (cc->Outputs().HasTag("NORM_LANDMARKS")) {
if (kOutNormalizedLandmarkList(cc).IsConnected()) {
NormalizedLandmarkList output_norm_landmarks;
for (int i = 0; i < output_landmarks.landmark_size(); ++i) {
const Landmark& landmark = output_landmarks.landmark(i);
@@ -246,18 +196,12 @@ mediapipe::Status TensorsToLandmarksCalculator::Process(CalculatorContext* cc) {
norm_landmark->set_presence(landmark.presence());
}
}
cc->Outputs()
.Tag("NORM_LANDMARKS")
.AddPacket(MakePacket<NormalizedLandmarkList>(output_norm_landmarks)
.At(cc->InputTimestamp()));
kOutNormalizedLandmarkList(cc).Send(std::move(output_norm_landmarks));
}
// Output absolute landmarks.
if (cc->Outputs().HasTag("LANDMARKS")) {
cc->Outputs()
.Tag("LANDMARKS")
.AddPacket(MakePacket<LandmarkList>(output_landmarks)
.At(cc->InputTimestamp()));
if (kOutLandmarkList(cc).IsConnected()) {
kOutLandmarkList(cc).Send(std::move(output_landmarks));
}
return mediapipe::OkStatus();
@@ -272,4 +216,5 @@ mediapipe::Status TensorsToLandmarksCalculator::LoadOptions(
return mediapipe::OkStatus();
}
} // namespace api2
} // namespace mediapipe