From c73027926ff09d760b4438173cc676e7971968f1 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Tue, 30 May 2023 12:34:34 -0700 Subject: [PATCH] Add MultiLandmarksSmoothingCalculator PiperOrigin-RevId: 536478601 --- mediapipe/calculators/util/BUILD | 34 ++++++ .../landmarks_smoothing_calculator_utils.cc | 31 +++++ .../landmarks_smoothing_calculator_utils.h | 16 +++ .../multi_landmarks_smoothing_calculator.cc | 113 ++++++++++++++++++ .../multi_landmarks_smoothing_calculator.h | 81 +++++++++++++ ...ti_world_landmarks_smoothing_calculator.cc | 100 ++++++++++++++++ ...lti_world_landmarks_smoothing_calculator.h | 74 ++++++++++++ 7 files changed, 449 insertions(+) create mode 100644 mediapipe/calculators/util/multi_landmarks_smoothing_calculator.cc create mode 100644 mediapipe/calculators/util/multi_landmarks_smoothing_calculator.h create mode 100644 mediapipe/calculators/util/multi_world_landmarks_smoothing_calculator.cc create mode 100644 mediapipe/calculators/util/multi_world_landmarks_smoothing_calculator.h diff --git a/mediapipe/calculators/util/BUILD b/mediapipe/calculators/util/BUILD index 7f2d7305a..ea593aa3d 100644 --- a/mediapipe/calculators/util/BUILD +++ b/mediapipe/calculators/util/BUILD @@ -929,6 +929,40 @@ cc_library( alwayslink = 1, ) +cc_library( + name = "multi_landmarks_smoothing_calculator", + srcs = ["multi_landmarks_smoothing_calculator.cc"], + hdrs = ["multi_landmarks_smoothing_calculator.h"], + deps = [ + ":landmarks_smoothing_calculator_cc_proto", + ":landmarks_smoothing_calculator_utils", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework:timestamp", + "//mediapipe/framework/api2:node", + "//mediapipe/framework/formats:landmark_cc_proto", + "//mediapipe/framework/formats:rect_cc_proto", + "//mediapipe/framework/port:ret_check", + ], + alwayslink = 1, +) + +cc_library( + name = "multi_world_landmarks_smoothing_calculator", + srcs = ["multi_world_landmarks_smoothing_calculator.cc"], + hdrs = ["multi_world_landmarks_smoothing_calculator.h"], + deps = [ + ":landmarks_smoothing_calculator_cc_proto", + ":landmarks_smoothing_calculator_utils", + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework:timestamp", + "//mediapipe/framework/api2:node", + "//mediapipe/framework/formats:landmark_cc_proto", + "//mediapipe/framework/formats:rect_cc_proto", + "//mediapipe/framework/port:ret_check", + ], + alwayslink = 1, +) + mediapipe_proto_library( name = "visibility_smoothing_calculator_proto", srcs = ["visibility_smoothing_calculator.proto"], diff --git a/mediapipe/calculators/util/landmarks_smoothing_calculator_utils.cc b/mediapipe/calculators/util/landmarks_smoothing_calculator_utils.cc index 59e773ed9..7e275807e 100644 --- a/mediapipe/calculators/util/landmarks_smoothing_calculator_utils.cc +++ b/mediapipe/calculators/util/landmarks_smoothing_calculator_utils.cc @@ -318,5 +318,36 @@ absl::StatusOr> InitializeLandmarksFilter( } } +absl::StatusOr MultiLandmarkFilters::GetOrCreate( + const int64_t tracking_id, + const mediapipe::LandmarksSmoothingCalculatorOptions& options) { + const auto it = filters_.find(tracking_id); + if (it != filters_.end()) { + return it->second.get(); + } + + ASSIGN_OR_RETURN(auto landmarks_filter, InitializeLandmarksFilter(options)); + filters_[tracking_id] = std::move(landmarks_filter); + return filters_[tracking_id].get(); +} + +void MultiLandmarkFilters::ClearUnused( + const std::vector& tracking_ids) { + std::vector unused_tracking_ids; + for (const auto& it : filters_) { + bool unused = true; + for (int64_t tracking_id : tracking_ids) { + if (tracking_id == it.first) unused = false; + } + if (unused) unused_tracking_ids.push_back(it.first); + } + + for (int64_t tracking_id : unused_tracking_ids) { + filters_.erase(tracking_id); + } +} + +void MultiLandmarkFilters::Clear() { filters_.clear(); } + } // namespace landmarks_smoothing } // namespace mediapipe diff --git a/mediapipe/calculators/util/landmarks_smoothing_calculator_utils.h b/mediapipe/calculators/util/landmarks_smoothing_calculator_utils.h index 267053907..c63db926a 100644 --- a/mediapipe/calculators/util/landmarks_smoothing_calculator_utils.h +++ b/mediapipe/calculators/util/landmarks_smoothing_calculator_utils.h @@ -55,6 +55,22 @@ class LandmarksFilter { absl::StatusOr> InitializeLandmarksFilter( const mediapipe::LandmarksSmoothingCalculatorOptions& options); +class MultiLandmarkFilters { + public: + virtual ~MultiLandmarkFilters() = default; + + virtual absl::StatusOr GetOrCreate( + const int64_t tracking_id, + const mediapipe::LandmarksSmoothingCalculatorOptions& options); + + virtual void ClearUnused(const std::vector& tracking_ids); + + virtual void Clear(); + + private: + std::map> filters_; +}; + } // namespace landmarks_smoothing } // namespace mediapipe diff --git a/mediapipe/calculators/util/multi_landmarks_smoothing_calculator.cc b/mediapipe/calculators/util/multi_landmarks_smoothing_calculator.cc new file mode 100644 index 000000000..40098935c --- /dev/null +++ b/mediapipe/calculators/util/multi_landmarks_smoothing_calculator.cc @@ -0,0 +1,113 @@ +// 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 "mediapipe/calculators/util/multi_landmarks_smoothing_calculator.h" + +#include +#include +#include +#include + +#include "mediapipe/calculators/util/landmarks_smoothing_calculator.pb.h" +#include "mediapipe/calculators/util/landmarks_smoothing_calculator_utils.h" +#include "mediapipe/framework/api2/node.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/landmark.pb.h" +#include "mediapipe/framework/formats/rect.pb.h" +#include "mediapipe/framework/port/ret_check.h" +#include "mediapipe/framework/timestamp.h" + +namespace mediapipe { +namespace api2 { + +namespace { + +using ::mediapipe::NormalizedRect; +using ::mediapipe::landmarks_smoothing::GetObjectScale; +using ::mediapipe::landmarks_smoothing::LandmarksToNormalizedLandmarks; +using ::mediapipe::landmarks_smoothing::MultiLandmarkFilters; +using ::mediapipe::landmarks_smoothing::NormalizedLandmarksToLandmarks; + +} // namespace + +class MultiLandmarksSmoothingCalculatorImpl + : public NodeImpl { + public: + absl::Status Process(CalculatorContext* cc) override { + // Check that landmarks are not empty and reset the filter if so. + // Don't emit an empty packet for this timestamp. + if (kInNormLandmarks(cc).IsEmpty()) { + multi_filters_.Clear(); + return absl::OkStatus(); + } + + const auto& timestamp = + absl::Microseconds(cc->InputTimestamp().Microseconds()); + + const auto& tracking_ids = kTrackingIds(cc).Get(); + multi_filters_.ClearUnused(tracking_ids); + + const auto& in_norm_landmarks_vec = kInNormLandmarks(cc).Get(); + RET_CHECK_EQ(in_norm_landmarks_vec.size(), tracking_ids.size()); + + int image_width; + int image_height; + std::tie(image_width, image_height) = kImageSize(cc).Get(); + + std::optional> object_scale_roi_vec; + if (kObjectScaleRoi(cc).IsConnected() && !kObjectScaleRoi(cc).IsEmpty()) { + object_scale_roi_vec = kObjectScaleRoi(cc).Get(); + RET_CHECK_EQ(object_scale_roi_vec.value().size(), tracking_ids.size()); + } + + std::vector out_norm_landmarks_vec; + for (int i = 0; i < tracking_ids.size(); ++i) { + LandmarkList in_landmarks; + NormalizedLandmarksToLandmarks(in_norm_landmarks_vec[i], image_width, + image_height, in_landmarks); + + std::optional object_scale; + if (object_scale_roi_vec) { + object_scale = GetObjectScale(object_scale_roi_vec.value()[i], + image_width, image_height); + } + + ASSIGN_OR_RETURN(auto* landmarks_filter, + multi_filters_.GetOrCreate( + tracking_ids[i], + cc->Options())); + + LandmarkList out_landmarks; + MP_RETURN_IF_ERROR(landmarks_filter->Apply(in_landmarks, timestamp, + object_scale, out_landmarks)); + + NormalizedLandmarkList out_norm_landmarks; + LandmarksToNormalizedLandmarks(out_landmarks, image_width, image_height, + out_norm_landmarks); + + out_norm_landmarks_vec.push_back(std::move(out_norm_landmarks)); + } + + kOutNormLandmarks(cc).Send(std::move(out_norm_landmarks_vec)); + + return absl::OkStatus(); + } + + private: + MultiLandmarkFilters multi_filters_; +}; +MEDIAPIPE_NODE_IMPLEMENTATION(MultiLandmarksSmoothingCalculatorImpl); + +} // namespace api2 +} // namespace mediapipe diff --git a/mediapipe/calculators/util/multi_landmarks_smoothing_calculator.h b/mediapipe/calculators/util/multi_landmarks_smoothing_calculator.h new file mode 100644 index 000000000..6c834ef56 --- /dev/null +++ b/mediapipe/calculators/util/multi_landmarks_smoothing_calculator.h @@ -0,0 +1,81 @@ +// 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. + +#ifndef MEDIAPIPE_CALCULATORS_UTIL_MULTI_LANDMARKS_SMOOTHING_CALCULATOR_H_ +#define MEDIAPIPE_CALCULATORS_UTIL_MULTI_LANDMARKS_SMOOTHING_CALCULATOR_H_ + +#include "mediapipe/framework/api2/node.h" +#include "mediapipe/framework/formats/landmark.pb.h" +#include "mediapipe/framework/formats/rect.pb.h" + +namespace mediapipe { +namespace api2 { + +// A calculator to smooth landmarks over time. +// +// Inputs: +// NORM_LANDMARKS: A std::vector of landmarks you want +// to smooth. +// TRACKING_IDS: A std vector of tracking IDs used to associate +// landmarks over time. When new ID arrives - calculator will initialize new +// filter. When tracking ID is no longer provided - calculator will forget +// smoothing state. +// IMAGE_SIZE: A std::pair represention of image width and height. +// Required to perform all computations in absolute coordinates to avoid any +// influence of normalized values. +// OBJECT_SCALE_ROI (optional): A std::vector used to determine the +// object scale for some of the filters. If not provided - object scale will +// be calculated from landmarks. +// +// Outputs: +// NORM_FILTERED_LANDMARKS: A std::vector of smoothed +// landmarks. +// +// Example config: +// node { +// calculator: "MultiLandmarksSmoothingCalculator" +// input_stream: "NORM_LANDMARKS:pose_landmarks" +// input_stream: "IMAGE_SIZE:image_size" +// input_stream: "OBJECT_SCALE_ROI:roi" +// output_stream: "NORM_FILTERED_LANDMARKS:pose_landmarks_filtered" +// options: { +// [mediapipe.LandmarksSmoothingCalculatorOptions.ext] { +// velocity_filter: { +// window_size: 5 +// velocity_scale: 10.0 +// } +// } +// } +// } +// +class MultiLandmarksSmoothingCalculator : public NodeIntf { + public: + static constexpr Input> + kInNormLandmarks{"NORM_LANDMARKS"}; + static constexpr Input> kTrackingIds{"TRACKING_IDS"}; + static constexpr Input> kImageSize{"IMAGE_SIZE"}; + static constexpr Input>::Optional kObjectScaleRoi{ + "OBJECT_SCALE_ROI"}; + static constexpr Output> + kOutNormLandmarks{"NORM_FILTERED_LANDMARKS"}; + + MEDIAPIPE_NODE_INTERFACE(MultiLandmarksSmoothingCalculator, kInNormLandmarks, + kTrackingIds, kImageSize, kObjectScaleRoi, + kOutNormLandmarks); +}; + +} // namespace api2 +} // namespace mediapipe + +#endif // MEDIAPIPE_CALCULATORS_UTIL_MULTI_LANDMARKS_SMOOTHING_CALCULATOR_H_ diff --git a/mediapipe/calculators/util/multi_world_landmarks_smoothing_calculator.cc b/mediapipe/calculators/util/multi_world_landmarks_smoothing_calculator.cc new file mode 100644 index 000000000..ddc16d296 --- /dev/null +++ b/mediapipe/calculators/util/multi_world_landmarks_smoothing_calculator.cc @@ -0,0 +1,100 @@ +// 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 "mediapipe/calculators/util/multi_world_landmarks_smoothing_calculator.h" + +#include +#include +#include +#include + +#include "mediapipe/calculators/util/landmarks_smoothing_calculator.pb.h" +#include "mediapipe/calculators/util/landmarks_smoothing_calculator_utils.h" +#include "mediapipe/framework/api2/node.h" +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/landmark.pb.h" +#include "mediapipe/framework/formats/rect.pb.h" +#include "mediapipe/framework/port/ret_check.h" +#include "mediapipe/framework/timestamp.h" + +namespace mediapipe { +namespace api2 { + +namespace { + +using ::mediapipe::Rect; +using ::mediapipe::landmarks_smoothing::GetObjectScale; +using ::mediapipe::landmarks_smoothing::MultiLandmarkFilters; + +} // namespace + +class MultiWorldLandmarksSmoothingCalculatorImpl + : public NodeImpl { + public: + absl::Status Process(CalculatorContext* cc) override { + // Check that landmarks are not empty and reset the filter if so. + // Don't emit an empty packet for this timestamp. + if (kInLandmarks(cc).IsEmpty()) { + multi_filters_.Clear(); + return absl::OkStatus(); + } + + const auto& timestamp = + absl::Microseconds(cc->InputTimestamp().Microseconds()); + + const auto& tracking_ids = kTrackingIds(cc).Get(); + multi_filters_.ClearUnused(tracking_ids); + + const auto& in_landmarks_vec = kInLandmarks(cc).Get(); + RET_CHECK_EQ(in_landmarks_vec.size(), tracking_ids.size()); + + std::optional> object_scale_roi_vec; + if (kObjectScaleRoi(cc).IsConnected() && !kObjectScaleRoi(cc).IsEmpty()) { + object_scale_roi_vec = kObjectScaleRoi(cc).Get(); + RET_CHECK_EQ(object_scale_roi_vec.value().size(), tracking_ids.size()); + } + + std::vector out_landmarks_vec; + for (int i = 0; i < tracking_ids.size(); ++i) { + const auto& in_landmarks = in_landmarks_vec[i]; + + std::optional object_scale; + if (object_scale_roi_vec) { + object_scale = GetObjectScale(object_scale_roi_vec.value()[i]); + } + + ASSIGN_OR_RETURN(auto* landmarks_filter, + multi_filters_.GetOrCreate( + tracking_ids[i], + cc->Options())); + + LandmarkList out_landmarks; + MP_RETURN_IF_ERROR(landmarks_filter->Apply(in_landmarks, timestamp, + object_scale, out_landmarks)); + + out_landmarks_vec.push_back(std::move(out_landmarks)); + } + + kOutLandmarks(cc).Send(std::move(out_landmarks_vec)); + + return absl::OkStatus(); + } + + private: + MultiLandmarkFilters multi_filters_; +}; +MEDIAPIPE_NODE_IMPLEMENTATION(MultiWorldLandmarksSmoothingCalculatorImpl); + +} // namespace api2 +} // namespace mediapipe diff --git a/mediapipe/calculators/util/multi_world_landmarks_smoothing_calculator.h b/mediapipe/calculators/util/multi_world_landmarks_smoothing_calculator.h new file mode 100644 index 000000000..2c54ae53c --- /dev/null +++ b/mediapipe/calculators/util/multi_world_landmarks_smoothing_calculator.h @@ -0,0 +1,74 @@ +// 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. + +#ifndef MEDIAPIPE_CALCULATORS_UTIL_MULTI_WORLD_LANDMARKS_SMOOTHING_CALCULATOR_H_ +#define MEDIAPIPE_CALCULATORS_UTIL_MULTI_WORLD_LANDMARKS_SMOOTHING_CALCULATOR_H_ + +#include "mediapipe/framework/api2/node.h" +#include "mediapipe/framework/formats/landmark.pb.h" +#include "mediapipe/framework/formats/rect.pb.h" + +namespace mediapipe { +namespace api2 { + +// A calculator to smooth landmarks over time. +// +// Inputs: +// LANDMARKS: A std::vector of landmarks you want to +// smooth. +// TRACKING_IDS: A std vector of tracking IDs used to associate +// landmarks over time. When new ID arrives - calculator will initialize new +// filter. When tracking ID is no longer provided - calculator will forget +// smoothing state. +// OBJECT_SCALE_ROI (optional): A std::vector used to determine the +// object scale for some of the filters. If not provided - object scale will +// be calculated from landmarks. +// +// Outputs: +// FILTERED_LANDMARKS: A std::vector of smoothed landmarks. +// +// Example config: +// node { +// calculator: "MultiWorldLandmarksSmoothingCalculator" +// input_stream: "LANDMARKS:landmarks" +// input_stream: "OBJECT_SCALE_ROI:roi" +// output_stream: "FILTERED_LANDMARKS:landmarks_filtered" +// options: { +// [mediapipe.LandmarksSmoothingCalculatorOptions.ext] { +// velocity_filter: { +// window_size: 5 +// velocity_scale: 10.0 +// } +// } +// } +// } +// +class MultiWorldLandmarksSmoothingCalculator : public NodeIntf { + public: + static constexpr Input> kInLandmarks{ + "LANDMARKS"}; + static constexpr Input> kTrackingIds{"TRACKING_IDS"}; + static constexpr Input>::Optional kObjectScaleRoi{ + "OBJECT_SCALE_ROI"}; + static constexpr Output> kOutLandmarks{ + "FILTERED_LANDMARKS"}; + + MEDIAPIPE_NODE_INTERFACE(MultiWorldLandmarksSmoothingCalculator, kInLandmarks, + kTrackingIds, kObjectScaleRoi, kOutLandmarks); +}; + +} // namespace api2 +} // namespace mediapipe + +#endif // MEDIAPIPE_CALCULATORS_UTIL_MULTI_WORLD_LANDMARKS_SMOOTHING_CALCULATOR_H_