Project import generated by Copybara.
GitOrigin-RevId: ea8d45731f5a052f79745e35bfd8240d6ac568d2
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user