migrate geometry_pipeline_calculator
PiperOrigin-RevId: 513368520
This commit is contained in:
parent
f15637b62c
commit
6796977a04
49
mediapipe/tasks/cc/vision/face_geometry/calculators/BUILD
Normal file
49
mediapipe/tasks/cc/vision/face_geometry/calculators/BUILD
Normal 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,
|
||||
)
|
|
@ -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
|
|
@ -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;
|
||||
}
|
Loading…
Reference in New Issue
Block a user