// Copyright 2020 The MediaPipe Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
#include <cmath>
#include <vector>
#include "Eigen/Dense"
#include "absl/memory/memory.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/rect.pb.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/modules/objectron/calculators/annotation_data.pb.h"
#include "mediapipe/modules/objectron/calculators/frame_annotation_to_rect_calculator.pb.h"
namespace mediapipe {
using Matrix3fRM = Eigen::Matrix<float, 3, 3, Eigen::RowMajor>;
using Eigen::Vector2f;
using Eigen::Vector3f;
namespace {
constexpr char kInputFrameAnnotationTag[] = "FRAME_ANNOTATION";
constexpr char kOutputNormRectsTag[] = "NORM_RECTS";
using ::mediapipe::NormalizedRect;
} // namespace
// A calculator that converts FrameAnnotation proto to NormalizedRect.
// The rotation angle of the NormalizedRect is derived from object's 3d pose.
// The angle is calculated such that after rotation the 2d projection of y-axis.
// on the image plane is always vertical.
class FrameAnnotationToRectCalculator : public CalculatorBase {
public:
enum ViewStatus {
TOP_VIEW_ON,
TOP_VIEW_OFF,
};
static absl::Status GetContract(CalculatorContract* cc);
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
private:
void AddAnnotationToRect(const ObjectAnnotation& annotation,
std::vector<NormalizedRect>* rect);
float RotationAngleFromAnnotation(const ObjectAnnotation& annotation);
float RotationAngleFromPose(const Matrix3fRM& rotation,
const Vector3f& translation, const Vector3f& vec);
ViewStatus status_;
float off_threshold_;
float on_threshold_;
};
REGISTER_CALCULATOR(FrameAnnotationToRectCalculator);
absl::Status FrameAnnotationToRectCalculator::GetContract(
CalculatorContract* cc) {
RET_CHECK(!cc->Inputs().GetTags().empty());
RET_CHECK(!cc->Outputs().GetTags().empty());
if (cc->Inputs().HasTag(kInputFrameAnnotationTag)) {
cc->Inputs().Tag(kInputFrameAnnotationTag).Set<FrameAnnotation>();
}
if (cc->Outputs().HasTag(kOutputNormRectsTag)) {
cc->Outputs().Tag(kOutputNormRectsTag).Set<std::vector<NormalizedRect>>();
}
return absl::OkStatus();
}
absl::Status FrameAnnotationToRectCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
status_ = TOP_VIEW_OFF;
const auto& options = cc->Options<FrameAnnotationToRectCalculatorOptions>();
off_threshold_ = options.off_threshold();
on_threshold_ = options.on_threshold();
RET_CHECK(off_threshold_ <= on_threshold_);
return absl::OkStatus();
}
absl::Status FrameAnnotationToRectCalculator::Process(CalculatorContext* cc) {
if (cc->Inputs().Tag(kInputFrameAnnotationTag).IsEmpty()) {
return absl::OkStatus();
}
auto output_rects = absl::make_unique<std::vector<NormalizedRect>>();
const auto& frame_annotation =
cc->Inputs().Tag(kInputFrameAnnotationTag).Get<FrameAnnotation>();
for (const auto& object_annotation : frame_annotation.annotations()) {
AddAnnotationToRect(object_annotation, output_rects.get());
}
// Output.
cc->Outputs()
.Tag(kOutputNormRectsTag)
.Add(output_rects.release(), cc->InputTimestamp());
return absl::OkStatus();
}
void FrameAnnotationToRectCalculator::AddAnnotationToRect(
const ObjectAnnotation& annotation, std::vector<NormalizedRect>* rects) {
float x_min = std::numeric_limits<float>::max();
float x_max = std::numeric_limits<float>::min();
float y_min = std::numeric_limits<float>::max();
float y_max = std::numeric_limits<float>::min();
for (const auto& keypoint : annotation.keypoints()) {
const auto& point_2d = keypoint.point_2d();
x_min = std::min(x_min, point_2d.x());
x_max = std::max(x_max, point_2d.x());
y_min = std::min(y_min, point_2d.y());
y_max = std::max(y_max, point_2d.y());
}
NormalizedRect new_rect;
new_rect.set_x_center((x_min + x_max) / 2);
new_rect.set_y_center((y_min + y_max) / 2);
new_rect.set_width(x_max - x_min);
new_rect.set_height(y_max - y_min);
new_rect.set_rotation(RotationAngleFromAnnotation(annotation));
rects->push_back(new_rect);
}
float FrameAnnotationToRectCalculator::RotationAngleFromAnnotation(
const ObjectAnnotation& annotation) {
// Get box rotation and translation from annotation.
const auto box_rotation =
Eigen::Map<const Matrix3fRM>(annotation.rotation().data());
const auto box_translation =
Eigen::Map<const Vector3f>(annotation.translation().data());
// Rotation angle to use when top-view is on(top-view on),
// Which will make z-axis upright after the rotation.
const float angle_on =
RotationAngleFromPose(box_rotation, box_translation, Vector3f::UnitZ());
// Rotation angle to use when side-view is on(top-view off),
// Which will make y-axis upright after the rotation.
const float angle_off =
RotationAngleFromPose(box_rotation, box_translation, Vector3f::UnitY());
// Calculate angle between z-axis and viewing ray in degrees.
const float view_to_z_angle = std::acos(box_rotation(2, 1)) * 180 / M_PI;
// Determine threshold based on current status,
// on_threshold_ is used for TOP_VIEW_ON -> TOP_VIEW_OFF transition,
// off_threshold_ is used for TOP_VIEW_OFF -> TOP_VIEW_ON transition.
const float thresh =
(status_ == TOP_VIEW_ON) ? on_threshold_ : off_threshold_;
// If view_to_z_angle is smaller than threshold, then top-view is on;
// Otherwise top-view is off.
status_ = (view_to_z_angle < thresh) ? TOP_VIEW_ON : TOP_VIEW_OFF;
// Determine which angle to used based on current status_.
float angle_to_rotate = (status_ == TOP_VIEW_ON) ? angle_on : angle_off;
return angle_to_rotate;
}
float FrameAnnotationToRectCalculator::RotationAngleFromPose(
const Matrix3fRM& rotation, const Vector3f& translation,
const Vector3f& vec) {
auto p1 = rotation * vec + translation;
auto p2 = -rotation * vec + translation;
const float dy = p2[2] * p2[1] - p1[2] * p1[1];
const float dx = p2[2] * p2[0] - p1[2] * p1[0];
return M_PI / 2 - std::atan2(dy, dx);
}
} // namespace mediapipe
Codebase Browser
chromium