Add MultiLandmarksSmoothingCalculator

PiperOrigin-RevId: 536478601

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"],

mediapipe/calculators/util/landmarks_smoothing_calculator_utils.cc

@@ -318,5 +318,36 @@ absl::StatusOr<std::unique_ptr<LandmarksFilter>> InitializeLandmarksFilter(
   }
 }
 
+absl::StatusOr<LandmarksFilter*> 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<int64_t>& tracking_ids) {
+  std::vector<int64_t> 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

mediapipe/calculators/util/landmarks_smoothing_calculator_utils.h

@@ -55,6 +55,22 @@ class LandmarksFilter {
 absl::StatusOr<std::unique_ptr<LandmarksFilter>> InitializeLandmarksFilter(
     const mediapipe::LandmarksSmoothingCalculatorOptions& options);
 
+class MultiLandmarkFilters {
+ public:
+  virtual ~MultiLandmarkFilters() = default;
+
+  virtual absl::StatusOr<LandmarksFilter*> GetOrCreate(
+      const int64_t tracking_id,
+      const mediapipe::LandmarksSmoothingCalculatorOptions& options);
+
+  virtual void ClearUnused(const std::vector<int64_t>& tracking_ids);
+
+  virtual void Clear();
+
+ private:
+  std::map<int64_t, std::unique_ptr<LandmarksFilter>> filters_;
+};
+
 }  // namespace landmarks_smoothing
 }  // namespace mediapipe
 

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 <cstdint>
+#include <memory>
+#include <optional>
+#include <vector>
+
+#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<MultiLandmarksSmoothingCalculator> {
+ 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<std::vector<NormalizedRect>> 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<NormalizedLandmarkList> 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<float> 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<LandmarksSmoothingCalculatorOptions>()));
+
+      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

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<NormalizedLandmarkList> of landmarks you want
+//     to smooth.
+//   TRACKING_IDS: A std<int64_t> 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<int, int> 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<NormRect> 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<NormalizedLandmarkList> 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<std::vector<mediapipe::NormalizedLandmarkList>>
+      kInNormLandmarks{"NORM_LANDMARKS"};
+  static constexpr Input<std::vector<int64_t>> kTrackingIds{"TRACKING_IDS"};
+  static constexpr Input<std::pair<int, int>> kImageSize{"IMAGE_SIZE"};
+  static constexpr Input<std::vector<NormalizedRect>>::Optional kObjectScaleRoi{
+      "OBJECT_SCALE_ROI"};
+  static constexpr Output<std::vector<mediapipe::NormalizedLandmarkList>>
+      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_

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 <cstdint>
+#include <memory>
+#include <optional>
+#include <vector>
+
+#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<MultiWorldLandmarksSmoothingCalculator> {
+ 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<std::vector<Rect>> 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<LandmarkList> out_landmarks_vec;
+    for (int i = 0; i < tracking_ids.size(); ++i) {
+      const auto& in_landmarks = in_landmarks_vec[i];
+
+      std::optional<float> 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<LandmarksSmoothingCalculatorOptions>()));
+
+      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

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<LandmarkList> of landmarks you want to
+//     smooth.
+//   TRACKING_IDS: A std<int64_t> 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<Rect> 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<LandmarkList> 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<std::vector<mediapipe::LandmarkList>> kInLandmarks{
+      "LANDMARKS"};
+  static constexpr Input<std::vector<int64_t>> kTrackingIds{"TRACKING_IDS"};
+  static constexpr Input<std::vector<Rect>>::Optional kObjectScaleRoi{
+      "OBJECT_SCALE_ROI"};
+  static constexpr Output<std::vector<mediapipe::LandmarkList>> 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_