migrate geometry_pipeline_calculator

PiperOrigin-RevId: 513368520
This commit is contained in:
MediaPipe Team 2023-03-01 16:07:59 -08:00 committed by Copybara-Service
parent f15637b62c
commit 6796977a04
3 changed files with 270 additions and 0 deletions

View File

@ -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,
)

View File

@ -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

View File

@ -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;
}