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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: b695dda274aa3ac3c7d054e150bd9eb5c1285b19
This commit is contained in:
MediaPipe Team
2020-01-17 15:49:22 -08:00
committed by chris
parent 66b377c825
commit dd02df1dbe
39 changed files with 692 additions and 203 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/calculators/core/clip_vector_size_calculator.h
+11
View File
@@ -38,6 +38,8 @@ namespace mediapipe {
// }
// }
// }
// Optionally, you can pass in a side packet that will override `max_vec_size`
// that is specified in the options.
template <typename T>
class ClipVectorSizeCalculator : public CalculatorBase {
public:
@@ -53,6 +55,10 @@ class ClipVectorSizeCalculator : public CalculatorBase {
cc->Inputs().Index(0).Set<std::vector<T>>();
cc->Outputs().Index(0).Set<std::vector<T>>();
// Optional input side packet that determines `max_vec_size`.
if (cc->InputSidePackets().NumEntries() > 0) {
cc->InputSidePackets().Index(0).Set<int>();
}
return ::mediapipe::OkStatus();
}
@@ -61,6 +67,11 @@ class ClipVectorSizeCalculator : public CalculatorBase {
cc->SetOffset(TimestampDiff(0));
max_vec_size_ = cc->Options<::mediapipe::ClipVectorSizeCalculatorOptions>()
.max_vec_size();
// Override `max_vec_size` if passed as side packet.
if (cc->InputSidePackets().NumEntries() > 0 &&
!cc->InputSidePackets().Index(0).IsEmpty()) {
max_vec_size_ = cc->InputSidePackets().Index(0).Get<int>();
}
return ::mediapipe::OkStatus();
}
mediapipe/calculators/core/clip_vector_size_calculator_test.cc
+27
View File
@@ -176,4 +176,31 @@ TEST(TestClipUniqueIntPtrVectorSizeCalculatorTest, ConsumeOneTimestamp) {
}
}
TEST(TestClipIntVectorSizeCalculatorTest, SidePacket) {
CalculatorGraphConfig::Node node_config =
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
calculator: "TestClipIntVectorSizeCalculator"
input_stream: "input_vector"
input_side_packet: "max_vec_size"
output_stream: "output_vector"
options {
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 1 }
}
)");
CalculatorRunner runner(node_config);
// This should override the default of 1 set in the options.
runner.MutableSidePackets()->Index(0) = Adopt(new int(2));
std::vector<int> input = {0, 1, 2, 3};
AddInputVector(input, /*timestamp=*/1, &runner);
MP_ASSERT_OK(runner.Run());
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
EXPECT_EQ(1, outputs.size());
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
const std::vector<int>& output = outputs[0].Get<std::vector<int>>();
EXPECT_EQ(2, output.size());
std::vector<int> expected_vector = {0, 1};
EXPECT_EQ(expected_vector, output);
}
} // namespace mediapipe
mediapipe/calculators/core/concatenate_vector_calculator.cc
+10
View File
@@ -35,6 +35,16 @@ namespace mediapipe {
typedef ConcatenateVectorCalculator<float> ConcatenateFloatVectorCalculator;
REGISTER_CALCULATOR(ConcatenateFloatVectorCalculator);
// Example config:
// node {
// calculator: "ConcatenateInt32VectorCalculator"
// input_stream: "int32_vector_1"
// input_stream: "int32_vector_2"
// output_stream: "concatenated_int32_vector"
// }
typedef ConcatenateVectorCalculator<int32> ConcatenateInt32VectorCalculator;
REGISTER_CALCULATOR(ConcatenateInt32VectorCalculator);
// Example config:
// node {
// calculator: "ConcatenateTfLiteTensorVectorCalculator"
mediapipe/calculators/image/image_transformation_calculator.cc
+20 -6
View File
@@ -138,8 +138,7 @@ mediapipe::ScaleMode_Mode ParseScaleMode(
// Note: To enable horizontal or vertical flipping, specify them in the
// calculator options. Flipping is applied after rotation.
//
// Note: Only scale mode STRETCH is currently supported on CPU,
// and flipping is not yet supported either.
// Note: Only scale mode STRETCH is currently supported on CPU.
//
class ImageTransformationCalculator : public CalculatorBase {
public:
@@ -316,6 +315,11 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
cv::Mat input_mat = formats::MatView(&input_img);
cv::Mat scaled_mat;
if (!output_height_ || !output_width_) {
output_height_ = input_height;
output_width_ = input_width;
}
if (scale_mode_ == mediapipe::ScaleMode_Mode_STRETCH) {
cv::resize(input_mat, scaled_mat, cv::Size(output_width_, output_height_));
} else {
@@ -367,10 +371,21 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
cv::Mat rotation_mat = cv::getRotationMatrix2D(src_center, angle, 1.0);
cv::warpAffine(scaled_mat, rotated_mat, rotation_mat, scaled_mat.size());
cv::Mat flipped_mat;
if (options_.flip_horizontally() || options_.flip_vertically()) {
const int flip_code =
options_.flip_horizontally() && options_.flip_vertically()
? -1
: options_.flip_horizontally();
cv::flip(rotated_mat, flipped_mat, flip_code);
} else {
flipped_mat = rotated_mat;
}
std::unique_ptr<ImageFrame> output_frame(
new ImageFrame(input_img.Format(), output_width, output_height));
cv::Mat output_mat = formats::MatView(output_frame.get());
rotated_mat.copyTo(output_mat);
flipped_mat.copyTo(output_mat);
cc->Outputs().Tag("IMAGE").Add(output_frame.release(), cc->InputTimestamp());
return ::mediapipe::OkStatus();
@@ -440,9 +455,8 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
cc->InputSidePackets().Tag("ROTATION_DEGREES").Get<int>());
}
static mediapipe::FrameScaleMode scale_mode =
mediapipe::FrameScaleModeFromProto(scale_mode_,
mediapipe::FrameScaleMode::kStretch);
mediapipe::FrameScaleMode scale_mode = mediapipe::FrameScaleModeFromProto(
scale_mode_, mediapipe::FrameScaleMode::kStretch);
mediapipe::FrameRotation rotation =
mediapipe::FrameRotationFromDegrees(RotationModeToDegrees(rotation_));
mediapipe/calculators/tensorflow/pack_media_sequence_calculator.cc
+15
View File
@@ -34,6 +34,7 @@ namespace mediapipe {
const char kSequenceExampleTag[] = "SEQUENCE_EXAMPLE";
const char kImageTag[] = "IMAGE";
const char kFloatContextFeaturePrefixTag[] = "FLOAT_CONTEXT_FEATURE_";
const char kFloatFeaturePrefixTag[] = "FLOAT_FEATURE_";
const char kForwardFlowEncodedTag[] = "FORWARD_FLOW_ENCODED";
const char kBBoxTag[] = "BBOX";
@@ -145,6 +146,9 @@ class PackMediaSequenceCalculator : public CalculatorBase {
}
cc->Inputs().Tag(tag).Set<std::vector<Detection>>();
}
if (absl::StartsWith(tag, kFloatContextFeaturePrefixTag)) {
cc->Inputs().Tag(tag).Set<std::vector<float>>();
}
if (absl::StartsWith(tag, kFloatFeaturePrefixTag)) {
cc->Inputs().Tag(tag).Set<std::vector<float>>();
}
@@ -344,6 +348,17 @@ class PackMediaSequenceCalculator : public CalculatorBase {
sequence_.get());
}
}
if (absl::StartsWith(tag, kFloatContextFeaturePrefixTag) &&
!cc->Inputs().Tag(tag).IsEmpty()) {
std::string key =
tag.substr(sizeof(kFloatContextFeaturePrefixTag) /
sizeof(*kFloatContextFeaturePrefixTag) -
1);
RET_CHECK_EQ(cc->InputTimestamp(), Timestamp::PostStream());
mpms::SetContextFeatureFloats(
key, cc->Inputs().Tag(tag).Get<std::vector<float>>(),
sequence_.get());
}
if (absl::StartsWith(tag, kFloatFeaturePrefixTag) &&
!cc->Inputs().Tag(tag).IsEmpty()) {
std::string key = tag.substr(sizeof(kFloatFeaturePrefixTag) /
mediapipe/calculators/tensorflow/pack_media_sequence_calculator_test.cc
+32
View File
@@ -194,6 +194,38 @@ TEST_F(PackMediaSequenceCalculatorTest, PacksTwoFloatLists) {
}
}
TEST_F(PackMediaSequenceCalculatorTest, PacksTwoContextFloatLists) {
SetUpCalculator(
{"FLOAT_CONTEXT_FEATURE_TEST:test", "FLOAT_CONTEXT_FEATURE_OTHER:test2"},
{}, false, true);
auto input_sequence = absl::make_unique<tf::SequenceExample>();
auto vf_ptr = absl::make_unique<std::vector<float>>(2, 3);
runner_->MutableInputs()
->Tag("FLOAT_CONTEXT_FEATURE_TEST")
.packets.push_back(Adopt(vf_ptr.release()).At(Timestamp::PostStream()));
vf_ptr = absl::make_unique<std::vector<float>>(2, 4);
runner_->MutableInputs()
->Tag("FLOAT_CONTEXT_FEATURE_OTHER")
.packets.push_back(Adopt(vf_ptr.release()).At(Timestamp::PostStream()));
runner_->MutableSidePackets()->Tag("SEQUENCE_EXAMPLE") =
Adopt(input_sequence.release());
MP_ASSERT_OK(runner_->Run());
const std::vector<Packet>& output_packets =
runner_->Outputs().Tag("SEQUENCE_EXAMPLE").packets;
ASSERT_EQ(1, output_packets.size());
const tf::SequenceExample& output_sequence =
output_packets[0].Get<tf::SequenceExample>();
ASSERT_THAT(mpms::GetContextFeatureFloats("TEST", output_sequence),
testing::ElementsAre(3, 3));
ASSERT_THAT(mpms::GetContextFeatureFloats("OTHER", output_sequence),
testing::ElementsAre(4, 4));
}
TEST_F(PackMediaSequenceCalculatorTest, PacksAdditionalContext) {
tf::Features context;
(*context.mutable_feature())["TEST"].mutable_bytes_list()->add_value("YES");
mediapipe/calculators/util/BUILD
+42
View File
@@ -508,6 +508,17 @@ proto_library(
],
)
proto_library(
name = "timed_box_list_to_render_data_calculator_proto",
srcs = ["timed_box_list_to_render_data_calculator.proto"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_proto",
"//mediapipe/util:color_proto",
"//mediapipe/util:render_data_proto",
],
)
proto_library(
name = "labels_to_render_data_calculator_proto",
srcs = ["labels_to_render_data_calculator.proto"],
@@ -651,6 +662,37 @@ cc_library(
alwayslink = 1,
)
mediapipe_cc_proto_library(
name = "timed_box_list_to_render_data_calculator_cc_proto",
srcs = ["timed_box_list_to_render_data_calculator.proto"],
cc_deps = [
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/util:color_cc_proto",
"//mediapipe/util:render_data_cc_proto",
],
visibility = ["//visibility:public"],
deps = [":timed_box_list_to_render_data_calculator_proto"],
)
cc_library(
name = "timed_box_list_to_render_data_calculator",
srcs = ["timed_box_list_to_render_data_calculator.cc"],
visibility = ["//visibility:public"],
deps = [
":timed_box_list_to_render_data_calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_options_cc_proto",
"//mediapipe/framework/port:ret_check",
"//mediapipe/util:color_cc_proto",
"//mediapipe/util:render_data_cc_proto",
"//mediapipe/util/tracking:box_tracker_cc_proto",
"//mediapipe/util/tracking:tracking_cc_proto",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
],
alwayslink = 1,
)
cc_library(
name = "labels_to_render_data_calculator",
srcs = ["labels_to_render_data_calculator.cc"],
mediapipe/calculators/util/collection_has_min_size_calculator.h
+11
View File
@@ -37,6 +37,8 @@ namespace mediapipe {
// }
// }
// }
// Optionally, uses a side packet to override `min_size` specified in the
// calculator options.
template <typename IterableT>
class CollectionHasMinSizeCalculator : public CalculatorBase {
public:
@@ -54,6 +56,10 @@ class CollectionHasMinSizeCalculator : public CalculatorBase {
cc->Inputs().Tag("ITERABLE").Set<IterableT>();
cc->Outputs().Index(0).Set<bool>();
// Optional input side packet that determines `min_size_`.
if (cc->InputSidePackets().NumEntries() > 0) {
cc->InputSidePackets().Index(0).Set<int>();
}
return ::mediapipe::OkStatus();
}
@@ -62,6 +68,11 @@ class CollectionHasMinSizeCalculator : public CalculatorBase {
min_size_ =
cc->Options<::mediapipe::CollectionHasMinSizeCalculatorOptions>()
.min_size();
// Override `min_size` if passed as side packet.
if (cc->InputSidePackets().NumEntries() > 0 &&
!cc->InputSidePackets().Index(0).IsEmpty()) {
min_size_ = cc->InputSidePackets().Index(0).Get<int>();
}
return ::mediapipe::OkStatus();
}
mediapipe/calculators/util/timed_box_list_to_render_data_calculator.cc
+146
View File
@@ -0,0 +1,146 @@
// Copyright 2019 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 "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "mediapipe/calculators/util/timed_box_list_to_render_data_calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/calculator_options.pb.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/util/color.pb.h"
#include "mediapipe/util/render_data.pb.h"
#include "mediapipe/util/tracking/box_tracker.pb.h"
#include "mediapipe/util/tracking/tracking.pb.h"
namespace mediapipe {
namespace {
constexpr char kTimedBoxListTag[] = "BOX_LIST";
constexpr char kRenderDataTag[] = "RENDER_DATA";
void AddTimedBoxProtoToRenderData(
const TimedBoxProto& box_proto,
const TimedBoxListToRenderDataCalculatorOptions& options,
RenderData* render_data) {
if (box_proto.has_quad() && box_proto.quad().vertices_size() > 0 &&
box_proto.quad().vertices_size() % 2 == 0) {
const int num_corners = box_proto.quad().vertices_size() / 2;
for (int i = 0; i < num_corners; ++i) {
const int next_corner = (i + 1) % num_corners;
auto* line_annotation = render_data->add_render_annotations();
line_annotation->mutable_color()->set_r(options.box_color().r());
line_annotation->mutable_color()->set_g(options.box_color().g());
line_annotation->mutable_color()->set_b(options.box_color().b());
line_annotation->set_thickness(options.thickness());
RenderAnnotation::Line* line = line_annotation->mutable_line();
line->set_x_start(box_proto.quad().vertices(i * 2));
line->set_y_start(box_proto.quad().vertices(i * 2 + 1));
line->set_x_end(box_proto.quad().vertices(next_corner * 2));
line->set_y_end(box_proto.quad().vertices(next_corner * 2 + 1));
}
} else {
auto* rect_annotation = render_data->add_render_annotations();
rect_annotation->mutable_color()->set_r(options.box_color().r());
rect_annotation->mutable_color()->set_g(options.box_color().g());
rect_annotation->mutable_color()->set_b(options.box_color().b());
rect_annotation->set_thickness(options.thickness());
RenderAnnotation::Rectangle* rect = rect_annotation->mutable_rectangle();
rect->set_normalized(true);
rect->set_left(box_proto.left());
rect->set_right(box_proto.right());
rect->set_top(box_proto.top());
rect->set_bottom(box_proto.bottom());
rect->set_rotation(box_proto.rotation());
}
}
} // namespace
// A calculator that converts TimedBoxProtoList proto to RenderData proto for
// visualization. If the input TimedBoxProto contains `quad` field, this
// calculator will draw a quadrilateral based on it. Otherwise this calculator
// will draw a rotated rectangle based on `top`, `bottom`, `left`, `right` and
// `rotation` fields
//
// Example config:
// node {
// calculator: "TimedBoxListToRenderDataCalculator"
// input_stream: "BOX_LIST:landmarks"
// output_stream: "RENDER_DATA:render_data"
// options {
// [TimedBoxListToRenderDataCalculatorOptions.ext] {
// box_color { r: 0 g: 255 b: 0 }
// thickness: 4.0
// }
// }
// }
class TimedBoxListToRenderDataCalculator : public CalculatorBase {
public:
TimedBoxListToRenderDataCalculator() {}
~TimedBoxListToRenderDataCalculator() override {}
TimedBoxListToRenderDataCalculator(
const TimedBoxListToRenderDataCalculator&) = delete;
TimedBoxListToRenderDataCalculator& operator=(
const TimedBoxListToRenderDataCalculator&) = delete;
static ::mediapipe::Status GetContract(CalculatorContract* cc);
::mediapipe::Status Open(CalculatorContext* cc) override;
::mediapipe::Status Process(CalculatorContext* cc) override;
private:
TimedBoxListToRenderDataCalculatorOptions options_;
};
REGISTER_CALCULATOR(TimedBoxListToRenderDataCalculator);
::mediapipe::Status TimedBoxListToRenderDataCalculator::GetContract(
CalculatorContract* cc) {
if (cc->Inputs().HasTag(kTimedBoxListTag)) {
cc->Inputs().Tag(kTimedBoxListTag).Set<TimedBoxProtoList>();
}
cc->Outputs().Tag(kRenderDataTag).Set<RenderData>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status TimedBoxListToRenderDataCalculator::Open(
CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
options_ = cc->Options<TimedBoxListToRenderDataCalculatorOptions>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status TimedBoxListToRenderDataCalculator::Process(
CalculatorContext* cc) {
auto render_data = absl::make_unique<RenderData>();
if (cc->Inputs().HasTag(kTimedBoxListTag)) {
const auto& box_list =
cc->Inputs().Tag(kTimedBoxListTag).Get<TimedBoxProtoList>();
for (const auto& box : box_list.box()) {
AddTimedBoxProtoToRenderData(box, options_, render_data.get());
}
}
cc->Outputs()
.Tag(kRenderDataTag)
.Add(render_data.release(), cc->InputTimestamp());
return ::mediapipe::OkStatus();
}
} // namespace mediapipe
mediapipe/calculators/util/timed_box_list_to_render_data_calculator.proto
+32
View File
@@ -0,0 +1,32 @@
// Copyright 2019 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.
syntax = "proto2";
package mediapipe;
import "mediapipe/framework/calculator.proto";
import "mediapipe/util/color.proto";
message TimedBoxListToRenderDataCalculatorOptions {
extend CalculatorOptions {
optional TimedBoxListToRenderDataCalculatorOptions ext = 289899854;
}
// Color of boxes.
optional Color box_color = 1;
// Thickness of the drawing of boxes.
optional double thickness = 2 [default = 1.0];
}