// 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 <algorithm>
#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_normalized(true);
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());
}
if (box_proto.has_label()) {
auto* label_annotation = render_data->add_render_annotations();
label_annotation->mutable_color()->set_r(options.box_color().r());
label_annotation->mutable_color()->set_g(options.box_color().g());
label_annotation->mutable_color()->set_b(options.box_color().b());
label_annotation->set_thickness(options.thickness());
RenderAnnotation::Text* text = label_annotation->mutable_text();
text->set_display_text(box_proto.label());
text->set_normalized(true);
constexpr float text_left_start = 0.2f;
text->set_left((1.0f - text_left_start) * box_proto.left() +
text_left_start * box_proto.right());
constexpr float text_baseline = 0.6f;
text->set_baseline(text_baseline * box_proto.bottom() +
(1.0f - text_baseline) * box_proto.top());
constexpr float text_height = 0.1f;
text->set_font_height(std::min(box_proto.bottom() - box_proto.top(),
box_proto.right() - box_proto.left()) *
text_height);
}
}
} // 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 absl::Status GetContract(CalculatorContract* cc);
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
private:
TimedBoxListToRenderDataCalculatorOptions options_;
};
REGISTER_CALCULATOR(TimedBoxListToRenderDataCalculator);
absl::Status TimedBoxListToRenderDataCalculator::GetContract(
CalculatorContract* cc) {
if (cc->Inputs().HasTag(kTimedBoxListTag)) {
cc->Inputs().Tag(kTimedBoxListTag).Set<TimedBoxProtoList>();
}
cc->Outputs().Tag(kRenderDataTag).Set<RenderData>();
return absl::OkStatus();
}
absl::Status TimedBoxListToRenderDataCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
options_ = cc->Options<TimedBoxListToRenderDataCalculatorOptions>();
return absl::OkStatus();
}
absl::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 absl::OkStatus();
}
} // namespace mediapipe
Codebase Browser
chromium