// 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 <cmath>
#include <memory>
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/landmark.pb.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/graphs/iris_tracking/calculators/iris_to_render_data_calculator.pb.h"
#include "mediapipe/util/color.pb.h"
#include "mediapipe/util/render_data.pb.h"
namespace mediapipe {
namespace {
constexpr char kIrisTag[] = "IRIS";
constexpr char kRenderDataTag[] = "RENDER_DATA";
constexpr char kImageSizeTag[] = "IMAGE_SIZE";
constexpr char kLeftIrisDepthTag[] = "LEFT_IRIS_DEPTH_MM";
constexpr char kRightIrisDepthTag[] = "RIGHT_IRIS_DEPTH_MM";
constexpr char kOvalLabel[] = "OVAL";
constexpr float kFontHeightScale = 1.5f;
constexpr int kNumIrisLandmarksPerEye = 5;
// TODO: Source.
constexpr float kIrisSizeInMM = 11.8;
inline void SetColor(RenderAnnotation* annotation, const Color& color) {
annotation->mutable_color()->set_r(color.r());
annotation->mutable_color()->set_g(color.g());
annotation->mutable_color()->set_b(color.b());
}
inline float GetDepth(float x0, float y0, float x1, float y1) {
return std::sqrt((x0 - x1) * (x0 - x1) + (y0 - y1) * (y0 - y1));
}
inline float GetLandmarkDepth(const NormalizedLandmark& ld0,
const NormalizedLandmark& ld1,
const std::pair<int, int>& image_size) {
return GetDepth(ld0.x() * image_size.first, ld0.y() * image_size.second,
ld1.x() * image_size.first, ld1.y() * image_size.second);
}
float CalculateIrisDiameter(const NormalizedLandmarkList& landmarks,
const std::pair<int, int>& image_size) {
const float dist_vert = GetLandmarkDepth(landmarks.landmark(1),
landmarks.landmark(2), image_size);
const float dist_hori = GetLandmarkDepth(landmarks.landmark(3),
landmarks.landmark(4), image_size);
return (dist_hori + dist_vert) / 2.0f;
}
float CalculateDepth(const NormalizedLandmark& center, float focal_length,
float iris_size, float img_w, float img_h) {
std::pair<float, float> origin{img_w / 2.f, img_h / 2.f};
const auto y = GetDepth(origin.first, origin.second, center.x() * img_w,
center.y() * img_h);
const auto x = std::sqrt(focal_length * focal_length + y * y);
const auto depth = kIrisSizeInMM * x / iris_size;
return depth;
}
} // namespace
// Converts iris landmarks to render data and estimates depth from the camera if
// focal length and image size. The depth will be rendered as part of the render
// data on the frame.
//
// Usage example:
// node {
// calculator: "IrisToRenderDataCalculator"
// input_stream: "IRIS:iris_landmarks"
// input_stream: "IMAGE_SIZE:image_size"
// # Note: Only one of FOCAL_LENGTH or IMAGE_FILE_PROPERTIES is necessary
// # to get focal length in pixels. Sending focal length in pixels to
// # this calculator is optional.
// input_side_packet: "FOCAL_LENGTH:focal_length_pixel"
// # OR
// input_side_packet: "IMAGE_FILE_PROPERTIES:image_file_properties"
// output_stream: "RENDER_DATA:iris_render_data"
// output_stream: "LEFT_IRIS_DEPTH_MM:left_iris_depth_mm"
// output_stream: "RIGHT_IRIS_DEPTH_MM:right_iris_depth_mm"
// node_options: {
// [type.googleapis.com/mediapipe.IrisToRenderDataCalculatorOptions] {
// color { r: 255 g: 255 b: 255 }
// thickness: 2.0
// font_height_px: 50
// horizontal_offset_px: 200
// vertical_offset_px: 200
// location: TOP_LEFT
// }
// }
// }
class IrisToRenderDataCalculator : public CalculatorBase {
public:
static absl::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Tag(kIrisTag).Set<NormalizedLandmarkList>();
cc->Outputs().Tag(kRenderDataTag).Set<RenderData>();
cc->Inputs().Tag(kImageSizeTag).Set<std::pair<int, int>>();
if (cc->Inputs().HasTag(kLeftIrisDepthTag)) {
cc->Inputs().Tag(kLeftIrisDepthTag).Set<float>();
}
if (cc->Inputs().HasTag(kRightIrisDepthTag)) {
cc->Inputs().Tag(kRightIrisDepthTag).Set<float>();
}
return absl::OkStatus();
}
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
private:
void RenderIris(const NormalizedLandmarkList& iris_landmarks,
const IrisToRenderDataCalculatorOptions& options,
const std::pair<int, int>& image_size, float iris_size,
RenderData* render_data);
void GetLeftIris(const NormalizedLandmarkList& lds,
NormalizedLandmarkList* iris);
void GetRightIris(const NormalizedLandmarkList& lds,
NormalizedLandmarkList* iris);
void AddTextRenderData(const IrisToRenderDataCalculatorOptions& options,
const std::pair<int, int>& image_size,
const std::vector<std::string>& lines,
RenderData* render_data);
static RenderAnnotation* AddOvalRenderData(
const IrisToRenderDataCalculatorOptions& options,
RenderData* render_data);
static RenderAnnotation* AddPointRenderData(
const IrisToRenderDataCalculatorOptions& options,
RenderData* render_data);
};
REGISTER_CALCULATOR(IrisToRenderDataCalculator);
absl::Status IrisToRenderDataCalculator::Open(CalculatorContext* cc) {
cc->SetOffset(TimestampDiff(0));
return absl::OkStatus();
}
absl::Status IrisToRenderDataCalculator::Process(CalculatorContext* cc) {
// Only process if there's input landmarks.
if (cc->Inputs().Tag(kIrisTag).IsEmpty()) {
return absl::OkStatus();
}
const auto& options =
cc->Options<::mediapipe::IrisToRenderDataCalculatorOptions>();
const auto& iris_landmarks =
cc->Inputs().Tag(kIrisTag).Get<NormalizedLandmarkList>();
RET_CHECK_EQ(iris_landmarks.landmark_size(), kNumIrisLandmarksPerEye * 2)
<< "Wrong number of iris landmarks";
std::pair<int, int> image_size;
RET_CHECK(!cc->Inputs().Tag(kImageSizeTag).IsEmpty());
image_size = cc->Inputs().Tag(kImageSizeTag).Get<std::pair<int, int>>();
auto render_data = absl::make_unique<RenderData>();
auto left_iris = absl::make_unique<NormalizedLandmarkList>();
auto right_iris = absl::make_unique<NormalizedLandmarkList>();
GetLeftIris(iris_landmarks, left_iris.get());
GetRightIris(iris_landmarks, right_iris.get());
const auto left_iris_size = CalculateIrisDiameter(*left_iris, image_size);
const auto right_iris_size = CalculateIrisDiameter(*right_iris, image_size);
RenderIris(*left_iris, options, image_size, left_iris_size,
render_data.get());
RenderIris(*right_iris, options, image_size, right_iris_size,
render_data.get());
std::vector<std::string> lines;
std::string line;
if (cc->Inputs().HasTag(kLeftIrisDepthTag) &&
!cc->Inputs().Tag(kLeftIrisDepthTag).IsEmpty()) {
const float left_iris_depth =
cc->Inputs().Tag(kLeftIrisDepthTag).Get<float>();
if (!std::isinf(left_iris_depth)) {
line = "Left : ";
absl::StrAppend(&line, ":", std::round(left_iris_depth / 10), " cm");
lines.emplace_back(line);
}
}
if (cc->Inputs().HasTag(kRightIrisDepthTag) &&
!cc->Inputs().Tag(kRightIrisDepthTag).IsEmpty()) {
const float right_iris_depth =
cc->Inputs().Tag(kRightIrisDepthTag).Get<float>();
if (!std::isinf(right_iris_depth)) {
line = "Right : ";
absl::StrAppend(&line, ":", std::round(right_iris_depth / 10), " cm");
lines.emplace_back(line);
}
}
AddTextRenderData(options, image_size, lines, render_data.get());
cc->Outputs()
.Tag(kRenderDataTag)
.Add(render_data.release(), cc->InputTimestamp());
return absl::OkStatus();
}
void IrisToRenderDataCalculator::AddTextRenderData(
const IrisToRenderDataCalculatorOptions& options,
const std::pair<int, int>& image_size,
const std::vector<std::string>& lines, RenderData* render_data) {
int label_baseline_px = options.vertical_offset_px();
float label_height_px =
std::ceil(options.font_height_px() * kFontHeightScale);
if (options.location() == IrisToRenderDataCalculatorOptions::TOP_LEFT) {
label_baseline_px += label_height_px;
} else if (options.location() ==
IrisToRenderDataCalculatorOptions::BOTTOM_LEFT) {
label_baseline_px += image_size.second - label_height_px * lines.size();
}
const auto label_left_px = options.horizontal_offset_px();
for (int i = 0; i < lines.size(); ++i) {
auto* label_annotation = render_data->add_render_annotations();
label_annotation->set_thickness(5);
label_annotation->mutable_color()->set_r(255);
label_annotation->mutable_color()->set_g(0);
label_annotation->mutable_color()->set_b(0);
//
auto* text = label_annotation->mutable_text();
text->set_display_text(lines[i]);
text->set_font_height(options.font_height_px());
text->set_left(label_left_px);
text->set_baseline(label_baseline_px + i * label_height_px);
text->set_font_face(options.font_face());
}
}
void IrisToRenderDataCalculator::RenderIris(
const NormalizedLandmarkList& iris_landmarks,
const IrisToRenderDataCalculatorOptions& options,
const std::pair<int, int>& image_size, float iris_size,
RenderData* render_data) {
auto* oval_data_render = AddOvalRenderData(options, render_data);
auto* oval_data = oval_data_render->mutable_oval();
const float iris_radius = iris_size / 2.f;
const auto& iris_center = iris_landmarks.landmark(0);
oval_data->mutable_rectangle()->set_top(iris_center.y() -
iris_radius / image_size.second);
oval_data->mutable_rectangle()->set_bottom(iris_center.y() +
iris_radius / image_size.second);
oval_data->mutable_rectangle()->set_left(iris_center.x() -
iris_radius / image_size.first);
oval_data->mutable_rectangle()->set_right(iris_center.x() +
iris_radius / image_size.first);
oval_data->mutable_rectangle()->set_normalized(true);
for (int i = 0; i < iris_landmarks.landmark_size(); ++i) {
const NormalizedLandmark& landmark = iris_landmarks.landmark(i);
auto* landmark_data_render = AddPointRenderData(options, render_data);
auto* landmark_data = landmark_data_render->mutable_point();
landmark_data->set_normalized(true);
landmark_data->set_x(landmark.x());
landmark_data->set_y(landmark.y());
}
}
void IrisToRenderDataCalculator::GetLeftIris(const NormalizedLandmarkList& lds,
NormalizedLandmarkList* iris) {
// Center, top, bottom, left, right
*iris->add_landmark() = lds.landmark(0);
*iris->add_landmark() = lds.landmark(2);
*iris->add_landmark() = lds.landmark(4);
*iris->add_landmark() = lds.landmark(3);
*iris->add_landmark() = lds.landmark(1);
}
void IrisToRenderDataCalculator::GetRightIris(const NormalizedLandmarkList& lds,
NormalizedLandmarkList* iris) {
// Center, top, bottom, left, right
*iris->add_landmark() = lds.landmark(5);
*iris->add_landmark() = lds.landmark(7);
*iris->add_landmark() = lds.landmark(9);
*iris->add_landmark() = lds.landmark(6);
*iris->add_landmark() = lds.landmark(8);
}
RenderAnnotation* IrisToRenderDataCalculator::AddOvalRenderData(
const IrisToRenderDataCalculatorOptions& options, RenderData* render_data) {
auto* oval_data_annotation = render_data->add_render_annotations();
oval_data_annotation->set_scene_tag(kOvalLabel);
SetColor(oval_data_annotation, options.oval_color());
oval_data_annotation->set_thickness(options.oval_thickness());
return oval_data_annotation;
}
RenderAnnotation* IrisToRenderDataCalculator::AddPointRenderData(
const IrisToRenderDataCalculatorOptions& options, RenderData* render_data) {
auto* landmark_data_annotation = render_data->add_render_annotations();
SetColor(landmark_data_annotation, options.landmark_color());
landmark_data_annotation->set_thickness(options.landmark_thickness());
return landmark_data_annotation;
}
} // namespace mediapipe