migrate geometry_pipeline_calculator

PiperOrigin-RevId: 513368520

mediapipe/tasks/cc/vision/face_geometry/calculators/BUILD

@@ -0,0 +1,49 @@
+# Copyright 2023 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.
+
+load("//mediapipe/framework/port:build_config.bzl", "mediapipe_proto_library")
+
+licenses(["notice"])
+
+package(default_visibility = ["//mediapipe/tasks:internal"])
+
+mediapipe_proto_library(
+    name = "geometry_pipeline_calculator_proto",
+    srcs = ["geometry_pipeline_calculator.proto"],
+    deps = [
+        "//mediapipe/framework:calculator_options_proto",
+    ],
+)
+
+cc_library(
+    name = "geometry_pipeline_calculator",
+    srcs = ["geometry_pipeline_calculator.cc"],
+    deps = [
+        ":geometry_pipeline_calculator_cc_proto",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/formats:landmark_cc_proto",
+        "//mediapipe/framework/port:logging",
+        "//mediapipe/framework/port:ret_check",
+        "//mediapipe/framework/port:status",
+        "//mediapipe/framework/port:statusor",
+        "//mediapipe/tasks/cc/vision/face_geometry/libs:geometry_pipeline",
+        "//mediapipe/tasks/cc/vision/face_geometry/libs:validation_utils",
+        "//mediapipe/tasks/cc/vision/face_geometry/proto:environment_cc_proto",
+        "//mediapipe/tasks/cc/vision/face_geometry/proto:face_geometry_cc_proto",
+        "//mediapipe/tasks/cc/vision/face_geometry/proto:geometry_pipeline_metadata_cc_proto",
+        "//mediapipe/util:resource_util",
+        "@com_google_absl//absl/memory",
+    ],
+    alwayslink = 1,
+)

mediapipe/tasks/cc/vision/face_geometry/calculators/geometry_pipeline_calculator.cc

@@ -0,0 +1,194 @@
+// Copyright 2023 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 <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/memory/memory.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/framework/port/status_macros.h"
+#include "mediapipe/framework/port/statusor.h"
+#include "mediapipe/tasks/cc/vision/face_geometry/calculators/geometry_pipeline_calculator.pb.h"
+#include "mediapipe/tasks/cc/vision/face_geometry/libs/geometry_pipeline.h"
+#include "mediapipe/tasks/cc/vision/face_geometry/libs/validation_utils.h"
+#include "mediapipe/tasks/cc/vision/face_geometry/proto/environment.pb.h"
+#include "mediapipe/tasks/cc/vision/face_geometry/proto/face_geometry.pb.h"
+#include "mediapipe/tasks/cc/vision/face_geometry/proto/geometry_pipeline_metadata.pb.h"
+#include "mediapipe/util/resource_util.h"
+
+namespace mediapipe::tasks::vision::face_geometry {
+namespace {
+
+static constexpr char kEnvironmentTag[] = "ENVIRONMENT";
+static constexpr char kImageSizeTag[] = "IMAGE_SIZE";
+static constexpr char kMultiFaceGeometryTag[] = "MULTI_FACE_GEOMETRY";
+static constexpr char kMultiFaceLandmarksTag[] = "MULTI_FACE_LANDMARKS";
+
+using ::mediapipe::tasks::vision::face_geometry::proto::Environment;
+using ::mediapipe::tasks::vision::face_geometry::proto::FaceGeometry;
+using ::mediapipe::tasks::vision::face_geometry::proto::
+    GeometryPipelineMetadata;
+
+// A calculator that renders a visual effect for multiple faces.
+//
+// Inputs:
+//   IMAGE_SIZE (`std::pair<int, int>`, required):
+//     The size of the current frame. The first element of the pair is the frame
+//     width; the other one is the frame height.
+//
+//     The face landmarks should have been detected on a frame with the same
+//     ratio. If used as-is, the resulting face geometry visualization should be
+//     happening on a frame with the same ratio as well.
+//
+//   MULTI_FACE_LANDMARKS (`std::vector<NormalizedLandmarkList>`, required):
+//     A vector of face landmark lists.
+//
+// Input side packets:
+//   ENVIRONMENT (`proto::Environment`, required)
+//     Describes an environment; includes the camera frame origin point location
+//     as well as virtual camera parameters.
+//
+// Output:
+//   MULTI_FACE_GEOMETRY (`std::vector<FaceGeometry>`, required):
+//     A vector of face geometry data.
+//
+// Options:
+//   metadata_path (`string`, optional):
+//     Defines a path for the geometry pipeline metadata file.
+//
+//     The geometry pipeline metadata file format must be the binary
+//     `GeometryPipelineMetadata` proto.
+//
+class GeometryPipelineCalculator : public CalculatorBase {
+ public:
+  static absl::Status GetContract(CalculatorContract* cc) {
+    cc->InputSidePackets().Tag(kEnvironmentTag).Set<Environment>();
+    cc->Inputs().Tag(kImageSizeTag).Set<std::pair<int, int>>();
+    cc->Inputs()
+        .Tag(kMultiFaceLandmarksTag)
+        .Set<std::vector<mediapipe::NormalizedLandmarkList>>();
+    cc->Outputs().Tag(kMultiFaceGeometryTag).Set<std::vector<FaceGeometry>>();
+
+    return absl::OkStatus();
+  }
+
+  absl::Status Open(CalculatorContext* cc) override {
+    cc->SetOffset(mediapipe::TimestampDiff(0));
+
+    const auto& options = cc->Options<FaceGeometryPipelineCalculatorOptions>();
+
+    ASSIGN_OR_RETURN(
+        GeometryPipelineMetadata metadata,
+        ReadMetadataFromFile(options.metadata_path()),
+        _ << "Failed to read the geometry pipeline metadata from file!");
+
+    MP_RETURN_IF_ERROR(ValidateGeometryPipelineMetadata(metadata))
+        << "Invalid geometry pipeline metadata!";
+
+    const Environment& environment =
+        cc->InputSidePackets().Tag(kEnvironmentTag).Get<Environment>();
+
+    MP_RETURN_IF_ERROR(ValidateEnvironment(environment))
+        << "Invalid environment!";
+
+    ASSIGN_OR_RETURN(geometry_pipeline_,
+                     CreateGeometryPipeline(environment, metadata),
+                     _ << "Failed to create a geometry pipeline!");
+
+    return absl::OkStatus();
+  }
+
+  absl::Status Process(CalculatorContext* cc) override {
+    // Both the `IMAGE_SIZE` and the `MULTI_FACE_LANDMARKS` streams are required
+    // to have a non-empty packet. In case this requirement is not met, there's
+    // nothing to be processed at the current timestamp.
+    if (cc->Inputs().Tag(kImageSizeTag).IsEmpty() ||
+        cc->Inputs().Tag(kMultiFaceLandmarksTag).IsEmpty()) {
+      return absl::OkStatus();
+    }
+
+    const auto& image_size =
+        cc->Inputs().Tag(kImageSizeTag).Get<std::pair<int, int>>();
+    const auto& multi_face_landmarks =
+        cc->Inputs()
+            .Tag(kMultiFaceLandmarksTag)
+            .Get<std::vector<mediapipe::NormalizedLandmarkList>>();
+
+    auto multi_face_geometry = absl::make_unique<std::vector<FaceGeometry>>();
+
+    ASSIGN_OR_RETURN(
+        *multi_face_geometry,
+        geometry_pipeline_->EstimateFaceGeometry(
+            multi_face_landmarks,  //
+            /*frame_width*/ image_size.first,
+            /*frame_height*/ image_size.second),
+        _ << "Failed to estimate face geometry for multiple faces!");
+
+    cc->Outputs()
+        .Tag(kMultiFaceGeometryTag)
+        .AddPacket(mediapipe::Adopt<std::vector<FaceGeometry>>(
+                       multi_face_geometry.release())
+                       .At(cc->InputTimestamp()));
+
+    return absl::OkStatus();
+  }
+
+  absl::Status Close(CalculatorContext* cc) override {
+    return absl::OkStatus();
+  }
+
+ private:
+  static absl::StatusOr<GeometryPipelineMetadata> ReadMetadataFromFile(
+      const std::string& metadata_path) {
+    ASSIGN_OR_RETURN(std::string metadata_blob,
+                     ReadContentBlobFromFile(metadata_path),
+                     _ << "Failed to read a metadata blob from file!");
+
+    GeometryPipelineMetadata metadata;
+    RET_CHECK(metadata.ParseFromString(metadata_blob))
+        << "Failed to parse a metadata proto from a binary blob!";
+
+    return metadata;
+  }
+
+  static absl::StatusOr<std::string> ReadContentBlobFromFile(
+      const std::string& unresolved_path) {
+    ASSIGN_OR_RETURN(std::string resolved_path,
+                     mediapipe::PathToResourceAsFile(unresolved_path),
+                     _ << "Failed to resolve path! Path = " << unresolved_path);
+
+    std::string content_blob;
+    MP_RETURN_IF_ERROR(
+        mediapipe::GetResourceContents(resolved_path, &content_blob))
+        << "Failed to read content blob! Resolved path = " << resolved_path;
+
+    return content_blob;
+  }
+
+  std::unique_ptr<GeometryPipeline> geometry_pipeline_;
+};
+
+}  // namespace
+
+using FaceGeometryPipelineCalculator = GeometryPipelineCalculator;
+
+REGISTER_CALCULATOR(
+    ::mediapipe::tasks::vision::face_geometry::FaceGeometryPipelineCalculator);
+
+}  // namespace mediapipe::tasks::vision::face_geometry

mediapipe/tasks/cc/vision/face_geometry/calculators/geometry_pipeline_calculator.proto

@@ -0,0 +1,27 @@
+// Copyright 2023 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.tasks.vision.face_geometry;
+
+import "mediapipe/framework/calculator_options.proto";
+
+message FaceGeometryPipelineCalculatorOptions {
+  extend mediapipe.CalculatorOptions {
+    optional FaceGeometryPipelineCalculatorOptions ext = 512499200;
+  }
+
+  optional string metadata_path = 1;
+}