cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add MultiLandmarksSmoothingCalculator

Browse Source 
PiperOrigin-RevId: 536478601
This commit is contained in:
MediaPipe Team 2023-05-30 12:34:34 -07:00 committed by Copybara-Service
parent a432559123
commit c73027926f
7 changed files with 449 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
mediapipe/calculators/util/BUILD
View File

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

31
mediapipe/calculators/util/landmarks_smoothing_calculator_utils.cc
View File

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

16
mediapipe/calculators/util/landmarks_smoothing_calculator_utils.h
View File

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

113
mediapipe/calculators/util/multi_landmarks_smoothing_calculator.cc Normal file
View File

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

81
mediapipe/calculators/util/multi_landmarks_smoothing_calculator.h Normal file
View File

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

100
mediapipe/calculators/util/multi_world_landmarks_smoothing_calculator.cc Normal file
View File

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

74
mediapipe/calculators/util/multi_world_landmarks_smoothing_calculator.h Normal file
View File

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