Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ae6be10afe | ||
|
|
38ee2603a7 | ||
|
|
86b3283b2f | ||
|
|
7d470a1335 | ||
|
|
d16cc3be5b | ||
|
|
137867d088 | ||
|
|
446d7cf6b6 | ||
|
|
90f72bd851 | ||
|
|
4285aeddfc | ||
|
|
37287925b0 | ||
|
|
48bcbb115f | ||
|
|
9437483827 | ||
|
|
d030c13931 | ||
|
|
fce372d153 | ||
|
|
c6fea4c9d9 | ||
|
|
259b48e082 |
@@ -3,7 +3,7 @@
|
||||
# Basic build settings
|
||||
build --jobs 128
|
||||
build --define='absl=1'
|
||||
build --cxxopt='-std=c++11'
|
||||
build --cxxopt='-std=c++14'
|
||||
build --copt='-Wno-sign-compare'
|
||||
build --copt='-Wno-unused-function'
|
||||
build --copt='-Wno-uninitialized'
|
||||
@@ -12,13 +12,16 @@ build --copt='-Wno-comment'
|
||||
build --copt='-Wno-return-type'
|
||||
build --copt='-Wno-unused-local-typedefs'
|
||||
build --copt='-Wno-ignored-attributes'
|
||||
# Temporarily set the incompatiblity flag for Bazel 0.27.0 and above
|
||||
# Temporarily set the incompatibility flag for Bazel 0.27.0 and above
|
||||
build --incompatible_disable_deprecated_attr_params=false
|
||||
build --incompatible_depset_is_not_iterable=false
|
||||
|
||||
# Sets the default Apple platform to macOS.
|
||||
build --apple_platform_type=macos
|
||||
|
||||
# Allow debugging with XCODE
|
||||
build --apple_generate_dsym
|
||||
|
||||
# Android configs.
|
||||
build:android --crosstool_top=//external:android/crosstool
|
||||
build:android --host_crosstool_top=@bazel_tools//tools/cpp:toolchain
|
||||
|
||||
+5
-1
@@ -30,10 +30,13 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
unzip \
|
||||
python \
|
||||
python-pip \
|
||||
python3-pip \
|
||||
libopencv-core-dev \
|
||||
libopencv-highgui-dev \
|
||||
libopencv-imgproc-dev \
|
||||
libopencv-video-dev \
|
||||
libopencv-calib3d-dev \
|
||||
libopencv-features2d-dev \
|
||||
software-properties-common && \
|
||||
add-apt-repository -y ppa:openjdk-r/ppa && \
|
||||
apt-get update && apt-get install -y openjdk-8-jdk && \
|
||||
@@ -42,9 +45,10 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
|
||||
RUN pip install --upgrade setuptools
|
||||
RUN pip install future
|
||||
RUN pip3 install six
|
||||
|
||||
# Install bazel
|
||||
ARG BAZEL_VERSION=0.26.1
|
||||
ARG BAZEL_VERSION=1.1.0
|
||||
RUN mkdir /bazel && \
|
||||
wget --no-check-certificate -O /bazel/installer.sh "https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/b\
|
||||
azel-${BAZEL_VERSION}-installer-linux-x86_64.sh" && \
|
||||
|
||||
@@ -10,20 +10,24 @@
|
||||
## ML Solutions in MediaPipe
|
||||
|
||||
* [Hand Tracking](mediapipe/docs/hand_tracking_mobile_gpu.md)
|
||||
* [Multi-hand Tracking](mediapipe/docs/multi_hand_tracking_mobile_gpu.md)
|
||||
* [Face Detection](mediapipe/docs/face_detection_mobile_gpu.md)
|
||||
* [Hair Segmentation](mediapipe/docs/hair_segmentation_mobile_gpu.md)
|
||||
* [Object Detection](mediapipe/docs/object_detection_mobile_gpu.md)
|
||||
* [Object Detection and Tracking](mediapipe/docs/object_tracking_mobile_gpu.md)
|
||||
* [AutoFlip](mediapipe/docs/autoflip.md)
|
||||
|
||||

|
||||

|
||||

|
||||

|
||||

|
||||

|
||||
|
||||
## Installation
|
||||
Follow these [instructions](mediapipe/docs/install.md).
|
||||
|
||||
## Getting started
|
||||
See mobile and desktop [examples](mediapipe/docs/examples.md).
|
||||
See mobile, desktop and Google Coral [examples](mediapipe/docs/examples.md).
|
||||
|
||||
## Documentation
|
||||
[MediaPipe Read-the-Docs](https://mediapipe.readthedocs.io/) or [docs.mediapipe.dev](https://docs.mediapipe.dev)
|
||||
@@ -37,10 +41,19 @@ A web-based visualizer is hosted on [viz.mediapipe.dev](https://viz.mediapipe.de
|
||||
* [Discuss](https://groups.google.com/forum/#!forum/mediapipe) - General community discussion around MediaPipe
|
||||
|
||||
## Publications
|
||||
* [On-Device, Real-Time Hand Tracking with MediaPipe](https://ai.googleblog.com/2019/08/on-device-real-time-hand-tracking-with.html)
|
||||
* [MediaPipe: A Framework for Building Perception Pipelines](https://arxiv.org/abs/1906.08172)
|
||||
|
||||
## Events
|
||||
[Open sourced at CVPR 2019](https://sites.google.com/corp/view/perception-cv4arvr/mediapipe) on June 17~20 in Long Beach, CA
|
||||
* [AI Nextcon 2020, 12-16 Feb 2020, Seattle](http://aisea20.xnextcon.com/)
|
||||
* [MediaPipe Madrid Meetup, 16 Dec 2019](https://www.meetup.com/Madrid-AI-Developers-Group/events/266329088/)
|
||||
* [MediaPipe London Meetup, Google 123 Building, 12 Dec 2019](https://www.meetup.com/London-AI-Tech-Talk/events/266329038)
|
||||
* [ML Conference, Berlin, 11 Dec 2019](https://mlconference.ai/machine-learning-advanced-development/mediapipe-building-real-time-cross-platform-mobile-web-edge-desktop-video-audio-ml-pipelines/)
|
||||
* [MediaPipe Berlin Meetup, Google Berlin, 11 Dec 2019](https://www.meetup.com/Berlin-AI-Tech-Talk/events/266328794/)
|
||||
* [The 3rd Workshop on YouTube-8M Large Scale Video Understanding Workshop](https://research.google.com/youtube8m/workshop2019/index.html) Seoul, Korea ICCV 2019
|
||||
* [AI DevWorld 2019](https://aidevworld.com) on Oct 10 in San Jose, California
|
||||
* [Google Industry Workshop at ICIP 2019](http://2019.ieeeicip.org/?action=page4&id=14#Google) [Presentation](https://docs.google.com/presentation/d/e/2PACX-1vRIBBbO_LO9v2YmvbHHEt1cwyqH6EjDxiILjuT0foXy1E7g6uyh4CesB2DkkEwlRDO9_lWfuKMZx98T/pub?start=false&loop=false&delayms=3000&slide=id.g556cc1a659_0_5) on Sept 24 in Taipei, Taiwan
|
||||
* [Open sourced at CVPR 2019](https://sites.google.com/corp/view/perception-cv4arvr/mediapipe) on June 17~20 in Long Beach, CA
|
||||
|
||||
## Alpha Disclaimer
|
||||
MediaPipe is currently in alpha for v0.6. We are still making breaking API changes and expect to get to stable API by v1.0.
|
||||
|
||||
@@ -10,19 +10,25 @@ http_archive(
|
||||
sha256 = "2ef429f5d7ce7111263289644d233707dba35e39696377ebab8b0bc701f7818e",
|
||||
)
|
||||
load("@bazel_skylib//lib:versions.bzl", "versions")
|
||||
versions.check(minimum_bazel_version = "0.24.1")
|
||||
versions.check(minimum_bazel_version = "0.24.1",
|
||||
maximum_bazel_version = "1.2.1")
|
||||
|
||||
# ABSL cpp library.
|
||||
|
||||
# ABSL cpp library lts_2019_08_08.
|
||||
http_archive(
|
||||
name = "com_google_absl",
|
||||
# Head commit on 2019-04-12.
|
||||
# TODO: Switch to the latest absl version when the problem gets
|
||||
# fixed.
|
||||
urls = [
|
||||
"https://github.com/abseil/abseil-cpp/archive/a02f62f456f2c4a7ecf2be3104fe0c6e16fbad9a.tar.gz",
|
||||
"https://github.com/abseil/abseil-cpp/archive/20190808.tar.gz",
|
||||
],
|
||||
sha256 = "d437920d1434c766d22e85773b899c77c672b8b4865d5dc2cd61a29fdff3cf03",
|
||||
strip_prefix = "abseil-cpp-a02f62f456f2c4a7ecf2be3104fe0c6e16fbad9a",
|
||||
# Remove after https://github.com/abseil/abseil-cpp/issues/326 is solved.
|
||||
patches = [
|
||||
"@//third_party:com_google_absl_f863b622fe13612433fdf43f76547d5edda0c93001.diff"
|
||||
],
|
||||
patch_args = [
|
||||
"-p1",
|
||||
],
|
||||
strip_prefix = "abseil-cpp-20190808",
|
||||
sha256 = "8100085dada279bf3ee00cd064d43b5f55e5d913be0dfe2906f06f8f28d5b37e"
|
||||
)
|
||||
|
||||
http_archive(
|
||||
@@ -103,31 +109,45 @@ http_archive(
|
||||
],
|
||||
)
|
||||
|
||||
# 2019-08-15
|
||||
_TENSORFLOW_GIT_COMMIT = "67def62936e28f97c16182dfcc467d8d1cae02b4"
|
||||
_TENSORFLOW_SHA256= "ddd4e3c056e7c0ff2ef29133b30fa62781dfbf8a903e99efb91a02d292fa9562"
|
||||
# 2019-11-21
|
||||
_TENSORFLOW_GIT_COMMIT = "f482488b481a799ca07e7e2d153cf47b8e91a60c"
|
||||
_TENSORFLOW_SHA256= "8d9118c2ce186c7e1403f04b96982fe72c184060c7f7a93e30a28dca358694f0"
|
||||
http_archive(
|
||||
name = "org_tensorflow",
|
||||
urls = [
|
||||
"https://mirror.bazel.build/github.com/tensorflow/tensorflow/archive/%s.tar.gz" % _TENSORFLOW_GIT_COMMIT,
|
||||
"https://github.com/tensorflow/tensorflow/archive/%s.tar.gz" % _TENSORFLOW_GIT_COMMIT,
|
||||
],
|
||||
strip_prefix = "tensorflow-%s" % _TENSORFLOW_GIT_COMMIT,
|
||||
sha256 = _TENSORFLOW_SHA256,
|
||||
# Patch https://github.com/tensorflow/tensorflow/commit/e3a7bdbebb99352351a19e2e403136166aa52934
|
||||
patches = [
|
||||
"@//third_party:tensorflow_065c20bf79253257c87bd4614bb9a7fdef015cbb.diff",
|
||||
"@//third_party:tensorflow_f67fcbefce906cd419e4657f0d41e21019b71abd.diff",
|
||||
"@//third_party:org_tensorflow_e3a7bdbebb99352351a19e2e403136166aa52934.diff"
|
||||
],
|
||||
patch_args = [
|
||||
"-p1",
|
||||
],
|
||||
strip_prefix = "tensorflow-%s" % _TENSORFLOW_GIT_COMMIT,
|
||||
sha256 = _TENSORFLOW_SHA256,
|
||||
)
|
||||
|
||||
load("@org_tensorflow//tensorflow:workspace.bzl", "tf_workspace")
|
||||
tf_workspace(tf_repo_name = "org_tensorflow")
|
||||
|
||||
http_archive(
|
||||
name = "ceres_solver",
|
||||
url = "https://github.com/ceres-solver/ceres-solver/archive/1.14.0.zip",
|
||||
patches = [
|
||||
"@//third_party:ceres_solver_9bf9588988236279e1262f75d7f4d85711dfa172.diff"
|
||||
],
|
||||
patch_args = [
|
||||
"-p1",
|
||||
],
|
||||
strip_prefix = "ceres-solver-1.14.0",
|
||||
sha256 = "5ba6d0db4e784621fda44a50c58bb23b0892684692f0c623e2063f9c19f192f1"
|
||||
)
|
||||
|
||||
# Please run
|
||||
# $ sudo apt-get install libopencv-core-dev libopencv-highgui-dev \
|
||||
# libopencv-calib3d-dev libopencv-features2d-dev \
|
||||
# libopencv-imgproc-dev libopencv-video-dev
|
||||
new_local_repository(
|
||||
name = "linux_opencv",
|
||||
@@ -156,11 +176,10 @@ new_local_repository(
|
||||
|
||||
http_archive(
|
||||
name = "android_opencv",
|
||||
sha256 = "056b849842e4fa8751d09edbb64530cfa7a63c84ccd232d0ace330e27ba55d0b",
|
||||
build_file = "@//third_party:opencv_android.BUILD",
|
||||
strip_prefix = "OpenCV-android-sdk",
|
||||
type = "zip",
|
||||
url = "https://github.com/opencv/opencv/releases/download/4.1.0/opencv-4.1.0-android-sdk.zip",
|
||||
url = "https://github.com/opencv/opencv/releases/download/3.4.3/opencv-3.4.3-android-sdk.zip",
|
||||
)
|
||||
|
||||
# After OpenCV 3.2.0, the pre-compiled opencv2.framework has google protobuf symbols, which will
|
||||
@@ -191,13 +210,18 @@ maven_install(
|
||||
artifacts = [
|
||||
"androidx.annotation:annotation:aar:1.1.0",
|
||||
"androidx.appcompat:appcompat:aar:1.1.0-rc01",
|
||||
"androidx.camera:camera-core:aar:1.0.0-alpha06",
|
||||
"androidx.camera:camera-camera2:aar:1.0.0-alpha06",
|
||||
"androidx.constraintlayout:constraintlayout:aar:1.1.3",
|
||||
"androidx.core:core:aar:1.1.0-rc03",
|
||||
"androidx.legacy:legacy-support-v4:aar:1.0.0",
|
||||
"androidx.recyclerview:recyclerview:aar:1.1.0-beta02",
|
||||
"com.google.android.material:material:aar:1.0.0-rc01",
|
||||
],
|
||||
repositories = ["https://dl.google.com/dl/android/maven2"],
|
||||
repositories = [
|
||||
"https://dl.google.com/dl/android/maven2",
|
||||
"https://repo1.maven.org/maven2",
|
||||
],
|
||||
)
|
||||
|
||||
maven_server(
|
||||
@@ -213,10 +237,10 @@ maven_jar(
|
||||
)
|
||||
|
||||
maven_jar(
|
||||
name = "androidx_concurrent_futures",
|
||||
artifact = "androidx.concurrent:concurrent-futures:1.0.0-alpha03",
|
||||
sha1 = "b528df95c7e2fefa2210c0c742bf3e491c1818ae",
|
||||
server = "google_server",
|
||||
name = "androidx_concurrent_futures",
|
||||
artifact = "androidx.concurrent:concurrent-futures:1.0.0-alpha03",
|
||||
sha1 = "b528df95c7e2fefa2210c0c742bf3e491c1818ae",
|
||||
server = "google_server",
|
||||
)
|
||||
|
||||
maven_jar(
|
||||
@@ -254,18 +278,11 @@ android_sdk_repository(
|
||||
|
||||
# iOS basic build deps.
|
||||
|
||||
load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
|
||||
|
||||
git_repository(
|
||||
http_archive(
|
||||
name = "build_bazel_rules_apple",
|
||||
remote = "https://github.com/bazelbuild/rules_apple.git",
|
||||
tag = "0.18.0",
|
||||
patches = [
|
||||
"@//third_party:rules_apple_c0863d0596ae6b769a29fa3fb72ff036444fd249.diff",
|
||||
],
|
||||
patch_args = [
|
||||
"-p1",
|
||||
],
|
||||
sha256 = "bdc8e66e70b8a75da23b79f1f8c6207356df07d041d96d2189add7ee0780cf4e",
|
||||
strip_prefix = "rules_apple-b869b0d3868d78a1d4ffd866ccb304fb68aa12c3",
|
||||
url = "https://github.com/bazelbuild/rules_apple/archive/b869b0d3868d78a1d4ffd866ccb304fb68aa12c3.tar.gz",
|
||||
)
|
||||
|
||||
load(
|
||||
@@ -298,3 +315,4 @@ http_archive(
|
||||
strip_prefix = "google-toolbox-for-mac-2.2.1",
|
||||
build_file = "@//third_party:google_toolbox_for_mac.BUILD",
|
||||
)
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
"mediapipe/examples/ios/facedetectiongpu/BUILD",
|
||||
"mediapipe/examples/ios/handdetectiongpu/BUILD",
|
||||
"mediapipe/examples/ios/handtrackinggpu/BUILD",
|
||||
"mediapipe/examples/ios/multihandtrackinggpu/BUILD",
|
||||
"mediapipe/examples/ios/objectdetectioncpu/BUILD",
|
||||
"mediapipe/examples/ios/objectdetectiongpu/BUILD"
|
||||
],
|
||||
@@ -18,6 +19,7 @@
|
||||
"//mediapipe/examples/ios/facedetectiongpu:FaceDetectionGpuApp",
|
||||
"//mediapipe/examples/ios/handdetectiongpu:HandDetectionGpuApp",
|
||||
"//mediapipe/examples/ios/handtrackinggpu:HandTrackingGpuApp",
|
||||
"//mediapipe/examples/ios/multihandtrackinggpu:MultiHandTrackingGpuApp",
|
||||
"//mediapipe/examples/ios/objectdetectioncpu:ObjectDetectionCpuApp",
|
||||
"//mediapipe/examples/ios/objectdetectiongpu:ObjectDetectionGpuApp",
|
||||
"//mediapipe/objc:mediapipe_framework_ios"
|
||||
@@ -84,6 +86,8 @@
|
||||
"mediapipe/examples/ios/handdetectiongpu/Base.lproj",
|
||||
"mediapipe/examples/ios/handtrackinggpu",
|
||||
"mediapipe/examples/ios/handtrackinggpu/Base.lproj",
|
||||
"mediapipe/examples/ios/multihandtrackinggpu",
|
||||
"mediapipe/examples/ios/multihandtrackinggpu/Base.lproj",
|
||||
"mediapipe/examples/ios/objectdetectioncpu",
|
||||
"mediapipe/examples/ios/objectdetectioncpu/Base.lproj",
|
||||
"mediapipe/examples/ios/objectdetectiongpu",
|
||||
|
||||
@@ -16,6 +16,7 @@
|
||||
"mediapipe/examples/ios/facedetectiongpu",
|
||||
"mediapipe/examples/ios/handdetectiongpu",
|
||||
"mediapipe/examples/ios/handtrackinggpu",
|
||||
"mediapipe/examples/ios/multihandtrackinggpu",
|
||||
"mediapipe/examples/ios/objectdetectioncpu",
|
||||
"mediapipe/examples/ios/objectdetectiongpu"
|
||||
],
|
||||
|
||||
@@ -113,8 +113,15 @@ class SpectrogramCalculator : public CalculatorBase {
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
Timestamp CurrentOutputTimestamp() {
|
||||
// Current output timestamp is the *center* of the next frame to be
|
||||
Timestamp CurrentOutputTimestamp(CalculatorContext* cc) {
|
||||
if (use_local_timestamp_) {
|
||||
return cc->InputTimestamp();
|
||||
}
|
||||
return CumulativeOutputTimestamp();
|
||||
}
|
||||
|
||||
Timestamp CumulativeOutputTimestamp() {
|
||||
// Cumulative output timestamp is the *center* of the next frame to be
|
||||
// emitted, hence delayed by half a window duration compared to relevant
|
||||
// input timestamp.
|
||||
return initial_input_timestamp_ +
|
||||
@@ -141,6 +148,7 @@ class SpectrogramCalculator : public CalculatorBase {
|
||||
const OutputMatrixType postprocess_output_fn(const OutputMatrixType&),
|
||||
CalculatorContext* cc);
|
||||
|
||||
bool use_local_timestamp_;
|
||||
double input_sample_rate_;
|
||||
bool pad_final_packet_;
|
||||
int frame_duration_samples_;
|
||||
@@ -173,6 +181,8 @@ const float SpectrogramCalculator::kLnPowerToDb = 4.342944819032518;
|
||||
SpectrogramCalculatorOptions spectrogram_options =
|
||||
cc->Options<SpectrogramCalculatorOptions>();
|
||||
|
||||
use_local_timestamp_ = spectrogram_options.use_local_timestamp();
|
||||
|
||||
if (spectrogram_options.frame_duration_seconds() <= 0.0) {
|
||||
::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Invalid or missing frame_duration_seconds.\n"
|
||||
@@ -351,11 +361,11 @@ template <class OutputMatrixType>
|
||||
<< "Inconsistent number of spectrogram channels.";
|
||||
if (allow_multichannel_input_) {
|
||||
cc->Outputs().Index(0).Add(spectrogram_matrices.release(),
|
||||
CurrentOutputTimestamp());
|
||||
CurrentOutputTimestamp(cc));
|
||||
} else {
|
||||
cc->Outputs().Index(0).Add(
|
||||
new OutputMatrixType(spectrogram_matrices->at(0)),
|
||||
CurrentOutputTimestamp());
|
||||
CurrentOutputTimestamp(cc));
|
||||
}
|
||||
cumulative_completed_frames_ += output_vectors.size();
|
||||
}
|
||||
|
||||
@@ -66,4 +66,11 @@ message SpectrogramCalculatorOptions {
|
||||
// uniformly regardless of output type (i.e., even dBs are multiplied, not
|
||||
// offset).
|
||||
optional double output_scale = 7 [default = 1.0];
|
||||
|
||||
// If use_local_timestamp is true, the output packet's timestamp is based on
|
||||
// the last sample of the packet and it's inferred from the latest input
|
||||
// packet's timestamp. If false, the output packet's timestamp is based on
|
||||
// the cumulative timestamping, which is inferred from the intial input
|
||||
// timestamp and the cumulative number of samples.
|
||||
optional bool use_local_timestamp = 8 [default = false];
|
||||
}
|
||||
|
||||
@@ -13,12 +13,12 @@
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "concatenate_vector_calculator_proto",
|
||||
srcs = ["concatenate_vector_calculator.proto"],
|
||||
@@ -26,6 +26,13 @@ proto_library(
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "dequantize_byte_array_calculator_proto",
|
||||
srcs = ["dequantize_byte_array_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "packet_cloner_calculator_proto",
|
||||
srcs = ["packet_cloner_calculator.proto"],
|
||||
@@ -40,6 +47,13 @@ proto_library(
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "packet_thinner_calculator_proto",
|
||||
srcs = ["packet_thinner_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "split_vector_calculator_proto",
|
||||
srcs = ["split_vector_calculator.proto"],
|
||||
@@ -72,6 +86,13 @@ proto_library(
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "clip_vector_size_calculator_proto",
|
||||
srcs = ["clip_vector_size_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "packet_cloner_calculator_cc_proto",
|
||||
srcs = ["packet_cloner_calculator.proto"],
|
||||
@@ -88,6 +109,14 @@ mediapipe_cc_proto_library(
|
||||
deps = [":packet_resampler_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "packet_thinner_calculator_cc_proto",
|
||||
srcs = ["packet_thinner_calculator.proto"],
|
||||
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":packet_thinner_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "split_vector_calculator_cc_proto",
|
||||
srcs = ["split_vector_calculator.proto"],
|
||||
@@ -104,6 +133,22 @@ mediapipe_cc_proto_library(
|
||||
deps = [":concatenate_vector_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "clip_vector_size_calculator_cc_proto",
|
||||
srcs = ["clip_vector_size_calculator.proto"],
|
||||
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":clip_vector_size_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "dequantize_byte_array_calculator_cc_proto",
|
||||
srcs = ["dequantize_byte_array_calculator.proto"],
|
||||
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":dequantize_byte_array_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "quantize_float_vector_calculator_cc_proto",
|
||||
srcs = ["quantize_float_vector_calculator.proto"],
|
||||
@@ -135,6 +180,7 @@ cc_library(
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -154,6 +200,66 @@ cc_test(
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "begin_loop_calculator",
|
||||
srcs = ["begin_loop_calculator.cc"],
|
||||
hdrs = ["begin_loop_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_contract",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/memory",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "end_loop_calculator",
|
||||
srcs = ["end_loop_calculator.cc"],
|
||||
hdrs = ["end_loop_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_contract",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/util:render_data_cc_proto",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "begin_end_loop_calculator_graph_test",
|
||||
srcs = ["begin_end_loop_calculator_graph_test.cc"],
|
||||
deps = [
|
||||
":begin_loop_calculator",
|
||||
":end_loop_calculator",
|
||||
"//mediapipe/calculators/core:packet_cloner_calculator",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_contract",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/memory",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "concatenate_vector_calculator",
|
||||
srcs = ["concatenate_vector_calculator.cc"],
|
||||
@@ -193,6 +299,49 @@ cc_test(
|
||||
srcs = ["concatenate_vector_calculator_test.cc"],
|
||||
deps = [
|
||||
":concatenate_vector_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "clip_vector_size_calculator",
|
||||
srcs = ["clip_vector_size_calculator.cc"],
|
||||
hdrs = ["clip_vector_size_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":clip_vector_size_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@org_tensorflow//tensorflow/lite:framework",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "clip_detection_vector_size_calculator",
|
||||
srcs = ["clip_detection_vector_size_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":clip_vector_size_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "clip_vector_size_calculator_test",
|
||||
srcs = ["clip_vector_size_calculator_test.cc"],
|
||||
deps = [
|
||||
":clip_vector_size_calculator",
|
||||
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
@@ -285,7 +434,7 @@ cc_library(
|
||||
"//visibility:public",
|
||||
],
|
||||
deps = [
|
||||
"//mediapipe/calculators/core:packet_cloner_calculator_cc_proto",
|
||||
":packet_cloner_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
@@ -316,6 +465,37 @@ cc_test(
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "packet_thinner_calculator",
|
||||
srcs = ["packet_thinner_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/core:packet_thinner_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "packet_thinner_calculator_test",
|
||||
srcs = ["packet_thinner_calculator_test.cc"],
|
||||
deps = [
|
||||
":packet_thinner_calculator",
|
||||
"//mediapipe/calculators/core:packet_thinner_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "pass_through_calculator",
|
||||
srcs = ["pass_through_calculator.cc"],
|
||||
@@ -387,6 +567,32 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "string_to_int_calculator",
|
||||
srcs = ["string_to_int_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "side_packet_to_stream_calculator",
|
||||
srcs = ["side_packet_to_stream_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "immediate_mux_calculator_test",
|
||||
srcs = ["immediate_mux_calculator_test.cc"],
|
||||
@@ -411,6 +617,7 @@ cc_test(
|
||||
cc_library(
|
||||
name = "packet_resampler_calculator",
|
||||
srcs = ["packet_resampler_calculator.cc"],
|
||||
hdrs = ["packet_resampler_calculator.h"],
|
||||
visibility = [
|
||||
"//visibility:public",
|
||||
],
|
||||
@@ -434,17 +641,17 @@ cc_library(
|
||||
cc_test(
|
||||
name = "packet_resampler_calculator_test",
|
||||
timeout = "short",
|
||||
srcs = ["packet_resampler_calculator_test.cc"],
|
||||
srcs = [
|
||||
"packet_resampler_calculator_test.cc",
|
||||
],
|
||||
deps = [
|
||||
":packet_resampler_calculator",
|
||||
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
)
|
||||
@@ -531,12 +738,19 @@ cc_library(
|
||||
":split_vector_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/util:resource_util",
|
||||
"@org_tensorflow//tensorflow/lite:framework",
|
||||
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
||||
],
|
||||
] + select({
|
||||
"//mediapipe/gpu:disable_gpu": [],
|
||||
"//mediapipe:ios": [],
|
||||
"//conditions:default": [
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
|
||||
],
|
||||
}),
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
@@ -558,6 +772,32 @@ cc_test(
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "dequantize_byte_array_calculator",
|
||||
srcs = ["dequantize_byte_array_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":dequantize_byte_array_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "dequantize_byte_array_calculator_test",
|
||||
srcs = ["dequantize_byte_array_calculator_test.cc"],
|
||||
deps = [
|
||||
":dequantize_byte_array_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "quantize_float_vector_calculator",
|
||||
srcs = ["quantize_float_vector_calculator.cc"],
|
||||
@@ -694,3 +934,29 @@ cc_test(
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "stream_to_side_packet_calculator",
|
||||
srcs = ["stream_to_side_packet_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "stream_to_side_packet_calculator_test",
|
||||
srcs = ["stream_to_side_packet_calculator_test.cc"],
|
||||
deps = [
|
||||
":stream_to_side_packet_calculator",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/memory",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -13,11 +13,12 @@
|
||||
// limitations under the License.
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Attach the header from one stream to another stream.
|
||||
// Attach the header from a stream or side input to another stream.
|
||||
//
|
||||
// The header stream (tag HEADER) must not have any packets in it.
|
||||
//
|
||||
@@ -25,17 +26,53 @@ namespace mediapipe {
|
||||
// calculator to not need a header or to accept a separate stream with
|
||||
// a header, that would be more future proof.
|
||||
//
|
||||
// Example usage 1:
|
||||
// node {
|
||||
// calculator: "AddHeaderCalculator"
|
||||
// input_stream: "DATA:audio"
|
||||
// input_stream: "HEADER:audio_header"
|
||||
// output_stream: "audio_with_header"
|
||||
// }
|
||||
//
|
||||
// Example usage 2:
|
||||
// node {
|
||||
// calculator: "AddHeaderCalculator"
|
||||
// input_stream: "DATA:audio"
|
||||
// input_side_packet: "HEADER:audio_header"
|
||||
// output_stream: "audio_with_header"
|
||||
// }
|
||||
//
|
||||
class AddHeaderCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Tag("HEADER").SetNone();
|
||||
bool has_side_input = false;
|
||||
bool has_header_stream = false;
|
||||
if (cc->InputSidePackets().HasTag("HEADER")) {
|
||||
cc->InputSidePackets().Tag("HEADER").SetAny();
|
||||
has_side_input = true;
|
||||
}
|
||||
if (cc->Inputs().HasTag("HEADER")) {
|
||||
cc->Inputs().Tag("HEADER").SetNone();
|
||||
has_header_stream = true;
|
||||
}
|
||||
if (has_side_input == has_header_stream) {
|
||||
return mediapipe::InvalidArgumentError(
|
||||
"Header must be provided via exactly one of side input and input "
|
||||
"stream");
|
||||
}
|
||||
cc->Inputs().Tag("DATA").SetAny();
|
||||
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Tag("DATA"));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
const Packet& header = cc->Inputs().Tag("HEADER").Header();
|
||||
Packet header;
|
||||
if (cc->InputSidePackets().HasTag("HEADER")) {
|
||||
header = cc->InputSidePackets().Tag("HEADER");
|
||||
}
|
||||
if (cc->Inputs().HasTag("HEADER")) {
|
||||
header = cc->Inputs().Tag("HEADER").Header();
|
||||
}
|
||||
if (!header.IsEmpty()) {
|
||||
cc->Outputs().Index(0).SetHeader(header);
|
||||
}
|
||||
|
||||
@@ -14,8 +14,10 @@
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
#include "mediapipe/framework/tool/validate_type.h"
|
||||
@@ -24,7 +26,7 @@ namespace mediapipe {
|
||||
|
||||
class AddHeaderCalculatorTest : public ::testing::Test {};
|
||||
|
||||
TEST_F(AddHeaderCalculatorTest, Works) {
|
||||
TEST_F(AddHeaderCalculatorTest, HeaderStream) {
|
||||
CalculatorGraphConfig::Node node;
|
||||
node.set_calculator("AddHeaderCalculator");
|
||||
node.add_input_stream("HEADER:header_stream");
|
||||
@@ -96,4 +98,62 @@ TEST_F(AddHeaderCalculatorTest, NoPacketsOnHeaderStream) {
|
||||
ASSERT_FALSE(runner.Run().ok());
|
||||
}
|
||||
|
||||
TEST_F(AddHeaderCalculatorTest, InputSidePacket) {
|
||||
CalculatorGraphConfig::Node node;
|
||||
node.set_calculator("AddHeaderCalculator");
|
||||
node.add_input_stream("DATA:data_stream");
|
||||
node.add_output_stream("merged_stream");
|
||||
node.add_input_side_packet("HEADER:header");
|
||||
|
||||
CalculatorRunner runner(node);
|
||||
|
||||
// Set header and add 5 packets.
|
||||
runner.MutableSidePackets()->Tag("HEADER") =
|
||||
Adopt(new std::string("my_header"));
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
Packet packet = Adopt(new int(i)).At(Timestamp(i * 1000));
|
||||
runner.MutableInputs()->Tag("DATA").packets.push_back(packet);
|
||||
}
|
||||
|
||||
// Run calculator.
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
ASSERT_EQ(1, runner.Outputs().NumEntries());
|
||||
|
||||
// Test output.
|
||||
EXPECT_EQ(std::string("my_header"),
|
||||
runner.Outputs().Index(0).header.Get<std::string>());
|
||||
const std::vector<Packet>& output_packets = runner.Outputs().Index(0).packets;
|
||||
ASSERT_EQ(5, output_packets.size());
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
const int val = output_packets[i].Get<int>();
|
||||
EXPECT_EQ(i, val);
|
||||
EXPECT_EQ(Timestamp(i * 1000), output_packets[i].Timestamp());
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(AddHeaderCalculatorTest, UsingBothSideInputAndStream) {
|
||||
CalculatorGraphConfig::Node node;
|
||||
node.set_calculator("AddHeaderCalculator");
|
||||
node.add_input_stream("HEADER:header_stream");
|
||||
node.add_input_stream("DATA:data_stream");
|
||||
node.add_output_stream("merged_stream");
|
||||
node.add_input_side_packet("HEADER:header");
|
||||
|
||||
CalculatorRunner runner(node);
|
||||
|
||||
// Set both headers and add 5 packets.
|
||||
runner.MutableSidePackets()->Tag("HEADER") =
|
||||
Adopt(new std::string("my_header"));
|
||||
runner.MutableSidePackets()->Tag("HEADER") =
|
||||
Adopt(new std::string("my_header"));
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
Packet packet = Adopt(new int(i)).At(Timestamp(i * 1000));
|
||||
runner.MutableInputs()->Tag("DATA").packets.push_back(packet);
|
||||
}
|
||||
|
||||
// Run should fail because header can only be provided one way.
|
||||
EXPECT_EQ(runner.Run().code(), ::mediapipe::InvalidArgumentError("").code());
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,335 @@
|
||||
// Copyright 2019 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 <string>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "mediapipe/calculators/core/begin_loop_calculator.h"
|
||||
#include "mediapipe/calculators/core/end_loop_calculator.h"
|
||||
#include "mediapipe/framework/calculator_contract.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h" // NOLINT
|
||||
|
||||
namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
typedef BeginLoopCalculator<std::vector<int>> BeginLoopIntegerCalculator;
|
||||
REGISTER_CALCULATOR(BeginLoopIntegerCalculator);
|
||||
|
||||
class IncrementCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Index(0).Set<int>();
|
||||
cc->Outputs().Index(0).Set<int>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
const int& input_int = cc->Inputs().Index(0).Get<int>();
|
||||
auto output_int = absl::make_unique<int>(input_int + 1);
|
||||
cc->Outputs().Index(0).Add(output_int.release(), cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(IncrementCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<int>> EndLoopIntegersCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopIntegersCalculator);
|
||||
|
||||
class BeginEndLoopCalculatorGraphTest : public ::testing::Test {
|
||||
protected:
|
||||
BeginEndLoopCalculatorGraphTest() {
|
||||
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
num_threads: 4
|
||||
input_stream: "ints"
|
||||
node {
|
||||
calculator: "BeginLoopIntegerCalculator"
|
||||
input_stream: "ITERABLE:ints"
|
||||
output_stream: "ITEM:int"
|
||||
output_stream: "BATCH_END:timestamp"
|
||||
}
|
||||
node {
|
||||
calculator: "IncrementCalculator"
|
||||
input_stream: "int"
|
||||
output_stream: "int_plus_one"
|
||||
}
|
||||
node {
|
||||
calculator: "EndLoopIntegersCalculator"
|
||||
input_stream: "ITEM:int_plus_one"
|
||||
input_stream: "BATCH_END:timestamp"
|
||||
output_stream: "ITERABLE:ints_plus_one"
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("ints_plus_one", &graph_config_, &output_packets_);
|
||||
}
|
||||
|
||||
CalculatorGraphConfig graph_config_;
|
||||
std::vector<Packet> output_packets_;
|
||||
};
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphTest, SingleEmptyVector) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||
Timestamp input_timestamp = Timestamp(0);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
// EndLoopCalc will forward the timestamp bound because there are no elements
|
||||
// in collection to output.
|
||||
ASSERT_EQ(0, output_packets_.size());
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphTest, SingleNonEmptyVector) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||
input_vector->emplace_back(0);
|
||||
input_vector->emplace_back(1);
|
||||
input_vector->emplace_back(2);
|
||||
Timestamp input_timestamp = Timestamp(0);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
ASSERT_EQ(1, output_packets_.size());
|
||||
EXPECT_EQ(input_timestamp, output_packets_[0].Timestamp());
|
||||
std::vector<int> expected_output_vector = {1, 2, 3};
|
||||
EXPECT_EQ(expected_output_vector, output_packets_[0].Get<std::vector<int>>());
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphTest, MultipleVectors) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
|
||||
auto input_vector0 = absl::make_unique<std::vector<int>>();
|
||||
input_vector0->emplace_back(0);
|
||||
input_vector0->emplace_back(1);
|
||||
Timestamp input_timestamp0 = Timestamp(0);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector0.release()).At(input_timestamp0)));
|
||||
|
||||
auto input_vector1 = absl::make_unique<std::vector<int>>();
|
||||
Timestamp input_timestamp1 = Timestamp(1);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector1.release()).At(input_timestamp1)));
|
||||
|
||||
auto input_vector2 = absl::make_unique<std::vector<int>>();
|
||||
input_vector2->emplace_back(2);
|
||||
input_vector2->emplace_back(3);
|
||||
Timestamp input_timestamp2 = Timestamp(2);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector2.release()).At(input_timestamp2)));
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
|
||||
ASSERT_EQ(2, output_packets_.size());
|
||||
|
||||
EXPECT_EQ(input_timestamp0, output_packets_[0].Timestamp());
|
||||
std::vector<int> expected_output_vector0 = {1, 2};
|
||||
EXPECT_EQ(expected_output_vector0,
|
||||
output_packets_[0].Get<std::vector<int>>());
|
||||
|
||||
// At input_timestamp1, EndLoopCalc will forward timestamp bound as there are
|
||||
// no elements in vector to process.
|
||||
|
||||
EXPECT_EQ(input_timestamp2, output_packets_[1].Timestamp());
|
||||
std::vector<int> expected_output_vector2 = {3, 4};
|
||||
EXPECT_EQ(expected_output_vector2,
|
||||
output_packets_[1].Get<std::vector<int>>());
|
||||
}
|
||||
|
||||
class MultiplierCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Index(0).Set<int>();
|
||||
cc->Inputs().Index(1).Set<int>();
|
||||
cc->Outputs().Index(0).Set<int>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
const int& input_int = cc->Inputs().Index(0).Get<int>();
|
||||
const int& multiplier_int = cc->Inputs().Index(1).Get<int>();
|
||||
auto output_int = absl::make_unique<int>(input_int * multiplier_int);
|
||||
cc->Outputs().Index(0).Add(output_int.release(), cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(MultiplierCalculator);
|
||||
|
||||
class BeginEndLoopCalculatorGraphWithClonedInputsTest : public ::testing::Test {
|
||||
protected:
|
||||
BeginEndLoopCalculatorGraphWithClonedInputsTest() {
|
||||
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
num_threads: 4
|
||||
input_stream: "ints"
|
||||
input_stream: "multiplier"
|
||||
node {
|
||||
calculator: "BeginLoopIntegerCalculator"
|
||||
input_stream: "ITERABLE:ints"
|
||||
input_stream: "CLONE:multiplier"
|
||||
output_stream: "ITEM:int_at_loop"
|
||||
output_stream: "CLONE:multiplier_cloned_at_loop"
|
||||
output_stream: "BATCH_END:timestamp"
|
||||
}
|
||||
node {
|
||||
calculator: "MultiplierCalculator"
|
||||
input_stream: "int_at_loop"
|
||||
input_stream: "multiplier_cloned_at_loop"
|
||||
output_stream: "multiplied_int_at_loop"
|
||||
}
|
||||
node {
|
||||
calculator: "EndLoopIntegersCalculator"
|
||||
input_stream: "ITEM:multiplied_int_at_loop"
|
||||
input_stream: "BATCH_END:timestamp"
|
||||
output_stream: "ITERABLE:multiplied_ints"
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("multiplied_ints", &graph_config_, &output_packets_);
|
||||
}
|
||||
|
||||
CalculatorGraphConfig graph_config_;
|
||||
std::vector<Packet> output_packets_;
|
||||
};
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, SingleEmptyVector) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||
Timestamp input_timestamp = Timestamp(42);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||
auto multiplier = absl::make_unique<int>(2);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"multiplier", Adopt(multiplier.release()).At(input_timestamp)));
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
// EndLoopCalc will forward the timestamp bound because there are no elements
|
||||
// in collection to output.
|
||||
ASSERT_EQ(0, output_packets_.size());
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, SingleNonEmptyVector) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||
input_vector->emplace_back(0);
|
||||
input_vector->emplace_back(1);
|
||||
input_vector->emplace_back(2);
|
||||
Timestamp input_timestamp = Timestamp(42);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||
auto multiplier = absl::make_unique<int>(2);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"multiplier", Adopt(multiplier.release()).At(input_timestamp)));
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
ASSERT_EQ(1, output_packets_.size());
|
||||
EXPECT_EQ(input_timestamp, output_packets_[0].Timestamp());
|
||||
std::vector<int> expected_output_vector = {0, 2, 4};
|
||||
EXPECT_EQ(expected_output_vector, output_packets_[0].Get<std::vector<int>>());
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, MultipleVectors) {
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
|
||||
auto input_vector0 = absl::make_unique<std::vector<int>>();
|
||||
input_vector0->emplace_back(0);
|
||||
input_vector0->emplace_back(1);
|
||||
Timestamp input_timestamp0 = Timestamp(42);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector0.release()).At(input_timestamp0)));
|
||||
auto multiplier0 = absl::make_unique<int>(2);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"multiplier", Adopt(multiplier0.release()).At(input_timestamp0)));
|
||||
|
||||
auto input_vector1 = absl::make_unique<std::vector<int>>();
|
||||
Timestamp input_timestamp1 = Timestamp(43);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector1.release()).At(input_timestamp1)));
|
||||
auto multiplier1 = absl::make_unique<int>(2);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"multiplier", Adopt(multiplier1.release()).At(input_timestamp1)));
|
||||
|
||||
auto input_vector2 = absl::make_unique<std::vector<int>>();
|
||||
input_vector2->emplace_back(2);
|
||||
input_vector2->emplace_back(3);
|
||||
Timestamp input_timestamp2 = Timestamp(44);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"ints", Adopt(input_vector2.release()).At(input_timestamp2)));
|
||||
auto multiplier2 = absl::make_unique<int>(3);
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"multiplier", Adopt(multiplier2.release()).At(input_timestamp2)));
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
|
||||
ASSERT_EQ(2, output_packets_.size());
|
||||
|
||||
EXPECT_EQ(input_timestamp0, output_packets_[0].Timestamp());
|
||||
std::vector<int> expected_output_vector0 = {0, 2};
|
||||
EXPECT_EQ(expected_output_vector0,
|
||||
output_packets_[0].Get<std::vector<int>>());
|
||||
|
||||
// At input_timestamp1, EndLoopCalc will forward timestamp bound as there are
|
||||
// no elements in vector to process.
|
||||
|
||||
EXPECT_EQ(input_timestamp2, output_packets_[1].Timestamp());
|
||||
std::vector<int> expected_output_vector2 = {6, 9};
|
||||
EXPECT_EQ(expected_output_vector2,
|
||||
output_packets_[1].Get<std::vector<int>>());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,34 @@
|
||||
// Copyright 2019 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/core/begin_loop_calculator.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator to process std::vector<NormalizedLandmarkList>.
|
||||
typedef BeginLoopCalculator<std::vector<::mediapipe::NormalizedLandmarkList>>
|
||||
BeginLoopNormalizedLandmarkListVectorCalculator;
|
||||
REGISTER_CALCULATOR(BeginLoopNormalizedLandmarkListVectorCalculator);
|
||||
|
||||
// A calculator to process std::vector<NormalizedRect>.
|
||||
typedef BeginLoopCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||
BeginLoopNormalizedRectCalculator;
|
||||
REGISTER_CALCULATOR(BeginLoopNormalizedRectCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,157 @@
|
||||
// Copyright 2019 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_CORE_BEGIN_LOOP_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_CORE_BEGIN_LOOP_CALCULATOR_H_
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_contract.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Calculator for implementing loops on iterable collections inside a MediaPipe
|
||||
// graph.
|
||||
//
|
||||
// It is designed to be used like:
|
||||
//
|
||||
// node {
|
||||
// calculator: "BeginLoopWithIterableCalculator"
|
||||
// input_stream: "ITERABLE:input_iterable" # IterableT @ext_ts
|
||||
// output_stream: "ITEM:input_element" # ItemT @loop_internal_ts
|
||||
// output_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||
// }
|
||||
//
|
||||
// node {
|
||||
// calculator: "ElementToBlaConverterSubgraph"
|
||||
// input_stream: "ITEM:input_to_loop_body" # ItemT @loop_internal_ts
|
||||
// output_stream: "BLA:output_of_loop_body" # ItemU @loop_internal_ts
|
||||
// }
|
||||
//
|
||||
// node {
|
||||
// calculator: "EndLoopWithOutputCalculator"
|
||||
// input_stream: "ITEM:output_of_loop_body" # ItemU @loop_internal_ts
|
||||
// input_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||
// output_stream: "OUTPUT:aggregated_result" # IterableU @ext_ts
|
||||
// }
|
||||
//
|
||||
// BeginLoopCalculator accepts an optional input stream tagged with "TICK"
|
||||
// which if non-empty, wakes up the calculator and calls
|
||||
// BeginLoopCalculator::Process(). Input streams tagged with "CLONE" are cloned
|
||||
// to the corresponding output streams at loop timestamps. This ensures that a
|
||||
// MediaPipe graph or sub-graph can run multiple times, once per element in the
|
||||
// "ITERABLE" for each pakcet clone of the packets in the "CLONE" input streams.
|
||||
template <typename IterableT>
|
||||
class BeginLoopCalculator : public CalculatorBase {
|
||||
using ItemT = typename IterableT::value_type;
|
||||
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
// A non-empty packet in the optional "TICK" input stream wakes up the
|
||||
// calculator.
|
||||
if (cc->Inputs().HasTag("TICK")) {
|
||||
cc->Inputs().Tag("TICK").SetAny();
|
||||
}
|
||||
|
||||
// An iterable collection in the input stream.
|
||||
RET_CHECK(cc->Inputs().HasTag("ITERABLE"));
|
||||
cc->Inputs().Tag("ITERABLE").Set<IterableT>();
|
||||
|
||||
// An element from the collection.
|
||||
RET_CHECK(cc->Outputs().HasTag("ITEM"));
|
||||
cc->Outputs().Tag("ITEM").Set<ItemT>();
|
||||
|
||||
RET_CHECK(cc->Outputs().HasTag("BATCH_END"));
|
||||
cc->Outputs()
|
||||
.Tag("BATCH_END")
|
||||
.Set<Timestamp>(
|
||||
// A flush signal to the corresponding EndLoopCalculator for it to
|
||||
// emit the aggregated result with the timestamp contained in this
|
||||
// flush signal packet.
|
||||
);
|
||||
|
||||
// Input streams tagged with "CLONE" are cloned to the corresponding
|
||||
// "CLONE" output streams at loop timestamps.
|
||||
RET_CHECK(cc->Inputs().NumEntries("CLONE") ==
|
||||
cc->Outputs().NumEntries("CLONE"));
|
||||
if (cc->Inputs().NumEntries("CLONE") > 0) {
|
||||
for (int i = 0; i < cc->Inputs().NumEntries("CLONE"); ++i) {
|
||||
cc->Inputs().Get("CLONE", i).SetAny();
|
||||
cc->Outputs().Get("CLONE", i).SetSameAs(&cc->Inputs().Get("CLONE", i));
|
||||
}
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||
Timestamp last_timestamp = loop_internal_timestamp_;
|
||||
if (!cc->Inputs().Tag("ITERABLE").IsEmpty()) {
|
||||
const IterableT& collection =
|
||||
cc->Inputs().Tag("ITERABLE").template Get<IterableT>();
|
||||
for (const auto& item : collection) {
|
||||
cc->Outputs().Tag("ITEM").AddPacket(
|
||||
MakePacket<ItemT>(item).At(loop_internal_timestamp_));
|
||||
ForwardClonePackets(cc, loop_internal_timestamp_);
|
||||
++loop_internal_timestamp_;
|
||||
}
|
||||
}
|
||||
|
||||
// The collection was empty and nothing was processed.
|
||||
if (last_timestamp == loop_internal_timestamp_) {
|
||||
// Increment loop_internal_timestamp_ because it is used up now.
|
||||
++loop_internal_timestamp_;
|
||||
for (auto it = cc->Outputs().begin(); it < cc->Outputs().end(); ++it) {
|
||||
it->SetNextTimestampBound(loop_internal_timestamp_);
|
||||
}
|
||||
}
|
||||
|
||||
// The for loop processing the input collection already incremented
|
||||
// loop_internal_timestamp_. To emit BATCH_END packet along the last
|
||||
// non-BATCH_END packet, decrement by one.
|
||||
cc->Outputs()
|
||||
.Tag("BATCH_END")
|
||||
.AddPacket(MakePacket<Timestamp>(cc->InputTimestamp())
|
||||
.At(Timestamp(loop_internal_timestamp_ - 1)));
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
void ForwardClonePackets(CalculatorContext* cc, Timestamp output_timestamp) {
|
||||
if (cc->Inputs().NumEntries("CLONE") > 0) {
|
||||
for (int i = 0; i < cc->Inputs().NumEntries("CLONE"); ++i) {
|
||||
if (!cc->Inputs().Get("CLONE", i).IsEmpty()) {
|
||||
auto input_packet = cc->Inputs().Get("CLONE", i).Value();
|
||||
cc->Outputs()
|
||||
.Get("CLONE", i)
|
||||
.AddPacket(std::move(input_packet).At(output_timestamp));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fake timestamps generated per element in collection.
|
||||
Timestamp loop_internal_timestamp_ = Timestamp(0);
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_CORE_BEGIN_LOOP_CALCULATOR_H_
|
||||
@@ -0,0 +1,26 @@
|
||||
// Copyright 2019 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 <vector>
|
||||
|
||||
#include "mediapipe/calculators/core/clip_vector_size_calculator.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef ClipVectorSizeCalculator<::mediapipe::Detection>
|
||||
ClipDetectionVectorSizeCalculator;
|
||||
REGISTER_CALCULATOR(ClipDetectionVectorSizeCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright 2019 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/core/clip_vector_size_calculator.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef ClipVectorSizeCalculator<::mediapipe::NormalizedRect>
|
||||
ClipNormalizedRectVectorSizeCalculator;
|
||||
REGISTER_CALCULATOR(ClipNormalizedRectVectorSizeCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,137 @@
|
||||
// Copyright 2019 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_CORE_CLIP_VECTOR_SIZE_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_CORE_CLIP_VECTOR_SIZE_CALCULATOR_H_
|
||||
|
||||
#include <type_traits>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/calculators/core/clip_vector_size_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Clips the size of the input vector of type T to a specified max_vec_size.
|
||||
// In a graph it will be used as:
|
||||
// node {
|
||||
// calculator: "ClipIntVectorSizeCalculator"
|
||||
// input_stream: "input_vector"
|
||||
// output_stream: "output_vector"
|
||||
// options {
|
||||
// [mediapipe.ClipIntVectorSizeCalculatorOptions.ext] {
|
||||
// max_vec_size: 5
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
template <typename T>
|
||||
class ClipVectorSizeCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().NumEntries() == 1);
|
||||
RET_CHECK(cc->Outputs().NumEntries() == 1);
|
||||
|
||||
if (cc->Options<::mediapipe::ClipVectorSizeCalculatorOptions>()
|
||||
.max_vec_size() < 1) {
|
||||
return ::mediapipe::InternalError(
|
||||
"max_vec_size should be greater than or equal to 1.");
|
||||
}
|
||||
|
||||
cc->Inputs().Index(0).Set<std::vector<T>>();
|
||||
cc->Outputs().Index(0).Set<std::vector<T>>();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
max_vec_size_ = cc->Options<::mediapipe::ClipVectorSizeCalculatorOptions>()
|
||||
.max_vec_size();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
if (max_vec_size_ < 1) {
|
||||
return ::mediapipe::InternalError(
|
||||
"max_vec_size should be greater than or equal to 1.");
|
||||
}
|
||||
if (cc->Inputs().Index(0).IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
return ClipVectorSize<T>(std::is_copy_constructible<T>(), cc);
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
::mediapipe::Status ClipVectorSize(std::true_type, CalculatorContext* cc) {
|
||||
auto output = absl::make_unique<std::vector<U>>();
|
||||
const std::vector<U>& input_vector =
|
||||
cc->Inputs().Index(0).Get<std::vector<U>>();
|
||||
if (max_vec_size_ >= input_vector.size()) {
|
||||
output->insert(output->end(), input_vector.begin(), input_vector.end());
|
||||
} else {
|
||||
for (int i = 0; i < max_vec_size_; ++i) {
|
||||
output->push_back(input_vector[i]);
|
||||
}
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
::mediapipe::Status ClipVectorSize(std::false_type, CalculatorContext* cc) {
|
||||
return ConsumeAndClipVectorSize<T>(std::is_move_constructible<U>(), cc);
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
::mediapipe::Status ConsumeAndClipVectorSize(std::true_type,
|
||||
CalculatorContext* cc) {
|
||||
auto output = absl::make_unique<std::vector<U>>();
|
||||
::mediapipe::StatusOr<std::unique_ptr<std::vector<U>>> input_status =
|
||||
cc->Inputs().Index(0).Value().Consume<std::vector<U>>();
|
||||
|
||||
if (input_status.ok()) {
|
||||
std::unique_ptr<std::vector<U>> input_vector =
|
||||
std::move(input_status).ValueOrDie();
|
||||
auto begin_it = input_vector->begin();
|
||||
auto end_it = input_vector->end();
|
||||
if (max_vec_size_ < input_vector->size()) {
|
||||
end_it = input_vector->begin() + max_vec_size_;
|
||||
}
|
||||
output->insert(output->end(), std::make_move_iterator(begin_it),
|
||||
std::make_move_iterator(end_it));
|
||||
} else {
|
||||
return input_status.status();
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
template <typename U>
|
||||
::mediapipe::Status ConsumeAndClipVectorSize(std::false_type,
|
||||
CalculatorContext* cc) {
|
||||
return ::mediapipe::InternalError(
|
||||
"Cannot copy or move input vectors and clip their size.");
|
||||
}
|
||||
|
||||
private:
|
||||
int max_vec_size_ = 0;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_CORE_CLIP_VECTOR_SIZE_CALCULATOR_H_
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright 2019 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;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
|
||||
message ClipVectorSizeCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional ClipVectorSizeCalculatorOptions ext = 274674998;
|
||||
}
|
||||
|
||||
// Maximum size of output vector.
|
||||
optional int32 max_vec_size = 1 [default = 1];
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
// Copyright 2019 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/core/clip_vector_size_calculator.h"
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h" // NOLINT
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef ClipVectorSizeCalculator<int> TestClipIntVectorSizeCalculator;
|
||||
REGISTER_CALCULATOR(TestClipIntVectorSizeCalculator);
|
||||
|
||||
void AddInputVector(const std::vector<int>& input, int64 timestamp,
|
||||
CalculatorRunner* runner) {
|
||||
runner->MutableInputs()->Index(0).packets.push_back(
|
||||
MakePacket<std::vector<int>>(input).At(Timestamp(timestamp)));
|
||||
}
|
||||
|
||||
TEST(TestClipIntVectorSizeCalculatorTest, EmptyVectorInput) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "TestClipIntVectorSizeCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "output_vector"
|
||||
options {
|
||||
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 1 }
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
|
||||
std::vector<int> input = {};
|
||||
AddInputVector(input, /*timestamp=*/1, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||
EXPECT_TRUE(outputs[0].Get<std::vector<int>>().empty());
|
||||
}
|
||||
|
||||
TEST(TestClipIntVectorSizeCalculatorTest, OneTimestamp) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "TestClipIntVectorSizeCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "output_vector"
|
||||
options {
|
||||
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 2 }
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
|
||||
std::vector<int> input = {0, 1, 2, 3};
|
||||
AddInputVector(input, /*timestamp=*/1, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||
const std::vector<int>& output = outputs[0].Get<std::vector<int>>();
|
||||
EXPECT_EQ(2, output.size());
|
||||
std::vector<int> expected_vector = {0, 1};
|
||||
EXPECT_EQ(expected_vector, output);
|
||||
}
|
||||
|
||||
TEST(TestClipIntVectorSizeCalculatorTest, TwoInputsAtTwoTimestamps) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "TestClipIntVectorSizeCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "output_vector"
|
||||
options {
|
||||
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 3 }
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
|
||||
{
|
||||
std::vector<int> input = {0, 1, 2, 3};
|
||||
AddInputVector(input, /*timestamp=*/1, &runner);
|
||||
}
|
||||
{
|
||||
std::vector<int> input = {2, 3, 4, 5};
|
||||
AddInputVector(input, /*timestamp=*/2, &runner);
|
||||
}
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(2, outputs.size());
|
||||
{
|
||||
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||
const std::vector<int>& output = outputs[0].Get<std::vector<int>>();
|
||||
EXPECT_EQ(3, output.size());
|
||||
std::vector<int> expected_vector = {0, 1, 2};
|
||||
EXPECT_EQ(expected_vector, output);
|
||||
}
|
||||
{
|
||||
EXPECT_EQ(Timestamp(2), outputs[1].Timestamp());
|
||||
const std::vector<int>& output = outputs[1].Get<std::vector<int>>();
|
||||
EXPECT_EQ(3, output.size());
|
||||
std::vector<int> expected_vector = {2, 3, 4};
|
||||
EXPECT_EQ(expected_vector, output);
|
||||
}
|
||||
}
|
||||
|
||||
typedef ClipVectorSizeCalculator<std::unique_ptr<int>>
|
||||
TestClipUniqueIntPtrVectorSizeCalculator;
|
||||
REGISTER_CALCULATOR(TestClipUniqueIntPtrVectorSizeCalculator);
|
||||
|
||||
TEST(TestClipUniqueIntPtrVectorSizeCalculatorTest, ConsumeOneTimestamp) {
|
||||
/* Note: We don't use CalculatorRunner for this test because it keeps copies
|
||||
* of input packets, so packets sent to the graph don't have sole ownership.
|
||||
* The test needs to send packets that own the data.
|
||||
*/
|
||||
CalculatorGraphConfig graph_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: "input_vector"
|
||||
node {
|
||||
calculator: "TestClipUniqueIntPtrVectorSizeCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "output_vector"
|
||||
options {
|
||||
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 3 }
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
std::vector<Packet> outputs;
|
||||
tool::AddVectorSink("output_vector", &graph_config, &outputs);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MP_EXPECT_OK(graph.Initialize(graph_config));
|
||||
MP_EXPECT_OK(graph.StartRun({}));
|
||||
|
||||
// input1 : {0, 1, 2, 3, 4, 5}
|
||||
auto input_vector = absl::make_unique<std::vector<std::unique_ptr<int>>>(6);
|
||||
for (int i = 0; i < 6; ++i) {
|
||||
input_vector->at(i) = absl::make_unique<int>(i);
|
||||
}
|
||||
|
||||
MP_EXPECT_OK(graph.AddPacketToInputStream(
|
||||
"input_vector", Adopt(input_vector.release()).At(Timestamp(1))));
|
||||
|
||||
MP_EXPECT_OK(graph.WaitUntilIdle());
|
||||
MP_EXPECT_OK(graph.CloseAllPacketSources());
|
||||
MP_EXPECT_OK(graph.WaitUntilDone());
|
||||
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||
const std::vector<std::unique_ptr<int>>& result =
|
||||
outputs[0].Get<std::vector<std::unique_ptr<int>>>();
|
||||
EXPECT_EQ(3, result.size());
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
const std::unique_ptr<int>& v = result[i];
|
||||
EXPECT_EQ(i, *v);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -19,7 +19,7 @@
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
@@ -50,7 +50,7 @@ typedef ConcatenateVectorCalculator<::mediapipe::NormalizedLandmark>
|
||||
ConcatenateLandmarkVectorCalculator;
|
||||
REGISTER_CALCULATOR(ConcatenateLandmarkVectorCalculator);
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
typedef ConcatenateVectorCalculator<::tflite::gpu::gl::GlBuffer>
|
||||
ConcatenateGlBufferVectorCalculator;
|
||||
REGISTER_CALCULATOR(ConcatenateGlBufferVectorCalculator);
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
// Copyright 2019 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 <cfloat>
|
||||
|
||||
#include "mediapipe/calculators/core/dequantize_byte_array_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
// Dequantizes a byte array to a vector of floats.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "DequantizeByteArrayCalculator"
|
||||
// input_stream: "ENCODED:encoded"
|
||||
// output_stream: "FLOAT_VECTOR:float_vector"
|
||||
// options {
|
||||
// [mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||
// max_quantized_value: 2
|
||||
// min_quantized_value: -2
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
namespace mediapipe {
|
||||
|
||||
class DequantizeByteArrayCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Tag("ENCODED").Set<std::string>();
|
||||
cc->Outputs().Tag("FLOAT_VECTOR").Set<std::vector<float>>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) final {
|
||||
const auto options =
|
||||
cc->Options<::mediapipe::DequantizeByteArrayCalculatorOptions>();
|
||||
if (!options.has_max_quantized_value() ||
|
||||
!options.has_min_quantized_value()) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Both max_quantized_value and min_quantized_value must be provided "
|
||||
"in DequantizeByteArrayCalculatorOptions.");
|
||||
}
|
||||
float max_quantized_value = options.max_quantized_value();
|
||||
float min_quantized_value = options.min_quantized_value();
|
||||
if (max_quantized_value < min_quantized_value + FLT_EPSILON) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"max_quantized_value must be greater than min_quantized_value.");
|
||||
}
|
||||
float range = max_quantized_value - min_quantized_value;
|
||||
scalar_ = range / 255.0;
|
||||
bias_ = (range / 512.0) + min_quantized_value;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||
const std::string& encoded =
|
||||
cc->Inputs().Tag("ENCODED").Value().Get<std::string>();
|
||||
std::vector<float> float_vector;
|
||||
float_vector.reserve(encoded.length());
|
||||
for (int i = 0; i < encoded.length(); ++i) {
|
||||
float_vector.push_back(
|
||||
static_cast<unsigned char>(encoded.at(i)) * scalar_ + bias_);
|
||||
}
|
||||
cc->Outputs()
|
||||
.Tag("FLOAT_VECTOR")
|
||||
.AddPacket(MakePacket<std::vector<float>>(float_vector)
|
||||
.At(cc->InputTimestamp()));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
float scalar_;
|
||||
float bias_;
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(DequantizeByteArrayCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright 2019 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;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
|
||||
message DequantizeByteArrayCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional DequantizeByteArrayCalculatorOptions ext = 272316343;
|
||||
}
|
||||
|
||||
optional float max_quantized_value = 1;
|
||||
optional float min_quantized_value = 2;
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
// Copyright 2019 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 <string>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h" // NOLINT
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
TEST(QuantizeFloatVectorCalculatorTest, WrongConfig) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "DequantizeByteArrayCalculator"
|
||||
input_stream: "ENCODED:encoded"
|
||||
output_stream: "FLOAT_VECTOR:float_vector"
|
||||
options {
|
||||
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||
max_quantized_value: 2
|
||||
}
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
std::string empty_string;
|
||||
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||
MakePacket<std::string>(empty_string).At(Timestamp(0)));
|
||||
auto status = runner.Run();
|
||||
EXPECT_FALSE(status.ok());
|
||||
EXPECT_THAT(
|
||||
status.message(),
|
||||
testing::HasSubstr(
|
||||
"Both max_quantized_value and min_quantized_value must be provided"));
|
||||
}
|
||||
|
||||
TEST(QuantizeFloatVectorCalculatorTest, WrongConfig2) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "DequantizeByteArrayCalculator"
|
||||
input_stream: "ENCODED:encoded"
|
||||
output_stream: "FLOAT_VECTOR:float_vector"
|
||||
options {
|
||||
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||
max_quantized_value: -2
|
||||
min_quantized_value: 2
|
||||
}
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
std::string empty_string;
|
||||
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||
MakePacket<std::string>(empty_string).At(Timestamp(0)));
|
||||
auto status = runner.Run();
|
||||
EXPECT_FALSE(status.ok());
|
||||
EXPECT_THAT(
|
||||
status.message(),
|
||||
testing::HasSubstr(
|
||||
"max_quantized_value must be greater than min_quantized_value"));
|
||||
}
|
||||
|
||||
TEST(QuantizeFloatVectorCalculatorTest, WrongConfig3) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "DequantizeByteArrayCalculator"
|
||||
input_stream: "ENCODED:encoded"
|
||||
output_stream: "FLOAT_VECTOR:float_vector"
|
||||
options {
|
||||
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||
max_quantized_value: 1
|
||||
min_quantized_value: 1
|
||||
}
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
std::string empty_string;
|
||||
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||
MakePacket<std::string>(empty_string).At(Timestamp(0)));
|
||||
auto status = runner.Run();
|
||||
EXPECT_FALSE(status.ok());
|
||||
EXPECT_THAT(
|
||||
status.message(),
|
||||
testing::HasSubstr(
|
||||
"max_quantized_value must be greater than min_quantized_value"));
|
||||
}
|
||||
|
||||
TEST(DequantizeByteArrayCalculatorTest, TestDequantization) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "DequantizeByteArrayCalculator"
|
||||
input_stream: "ENCODED:encoded"
|
||||
output_stream: "FLOAT_VECTOR:float_vector"
|
||||
options {
|
||||
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||
max_quantized_value: 2
|
||||
min_quantized_value: -2
|
||||
}
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
unsigned char input[4] = {0x7F, 0xFF, 0x00, 0x01};
|
||||
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||
MakePacket<std::string>(
|
||||
std::string(reinterpret_cast<char const*>(input), 4))
|
||||
.At(Timestamp(0)));
|
||||
auto status = runner.Run();
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
const std::vector<Packet>& outputs =
|
||||
runner.Outputs().Tag("FLOAT_VECTOR").packets;
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
const std::vector<float>& result = outputs[0].Get<std::vector<float>>();
|
||||
ASSERT_FALSE(result.empty());
|
||||
EXPECT_EQ(4, result.size());
|
||||
EXPECT_NEAR(0, result[0], 0.01);
|
||||
EXPECT_NEAR(2, result[1], 0.01);
|
||||
EXPECT_NEAR(-2, result[2], 0.01);
|
||||
EXPECT_NEAR(-1.976, result[3], 0.01);
|
||||
|
||||
EXPECT_EQ(Timestamp(0), outputs[0].Timestamp());
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,40 @@
|
||||
// Copyright 2019 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/core/end_loop_calculator.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/util/render_data.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef EndLoopCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||
EndLoopNormalizedRectCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopNormalizedRectCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<::mediapipe::NormalizedLandmarkList>>
|
||||
EndLoopNormalizedLandmarkListVectorCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopNormalizedLandmarkListVectorCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<bool>> EndLoopBooleanCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopBooleanCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<::mediapipe::RenderData>>
|
||||
EndLoopRenderDataCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopRenderDataCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,106 @@
|
||||
// Copyright 2019 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_CORE_END_LOOP_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_CORE_END_LOOP_CALCULATOR_H_
|
||||
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_contract.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Calculator for completing the processing of loops on iterable collections
|
||||
// inside a MediaPipe graph. The EndLoopCalculator collects all input packets
|
||||
// from ITEM input_stream into a collection and upon receiving the flush signal
|
||||
// from the "BATCH_END" tagged input stream, it emits the aggregated results
|
||||
// at the original timestamp contained in the "BATCH_END" input stream.
|
||||
//
|
||||
// It is designed to be used like:
|
||||
//
|
||||
// node {
|
||||
// calculator: "BeginLoopWithIterableCalculator"
|
||||
// input_stream: "ITERABLE:input_iterable" # IterableT @ext_ts
|
||||
// output_stream: "ITEM:input_element" # ItemT @loop_internal_ts
|
||||
// output_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||
// }
|
||||
//
|
||||
// node {
|
||||
// calculator: "ElementToBlaConverterSubgraph"
|
||||
// input_stream: "ITEM:input_to_loop_body" # ItemT @loop_internal_ts
|
||||
// output_stream: "BLA:output_of_loop_body" # ItemU @loop_internal_ts
|
||||
// }
|
||||
//
|
||||
// node {
|
||||
// calculator: "EndLoopWithOutputCalculator"
|
||||
// input_stream: "ITEM:output_of_loop_body" # ItemU @loop_internal_ts
|
||||
// input_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||
// output_stream: "OUTPUT:aggregated_result" # IterableU @ext_ts
|
||||
// }
|
||||
template <typename IterableT>
|
||||
class EndLoopCalculator : public CalculatorBase {
|
||||
using ItemT = typename IterableT::value_type;
|
||||
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag("BATCH_END"))
|
||||
<< "Missing BATCH_END tagged input_stream.";
|
||||
cc->Inputs().Tag("BATCH_END").Set<Timestamp>();
|
||||
|
||||
RET_CHECK(cc->Inputs().HasTag("ITEM"));
|
||||
cc->Inputs().Tag("ITEM").Set<ItemT>();
|
||||
|
||||
RET_CHECK(cc->Outputs().HasTag("ITERABLE"));
|
||||
cc->Outputs().Tag("ITERABLE").Set<IterableT>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
if (!cc->Inputs().Tag("ITEM").IsEmpty()) {
|
||||
if (!input_stream_collection_) {
|
||||
input_stream_collection_.reset(new IterableT);
|
||||
}
|
||||
input_stream_collection_->push_back(
|
||||
cc->Inputs().Tag("ITEM").template Get<ItemT>());
|
||||
}
|
||||
|
||||
if (!cc->Inputs().Tag("BATCH_END").Value().IsEmpty()) { // flush signal
|
||||
Timestamp loop_control_ts =
|
||||
cc->Inputs().Tag("BATCH_END").template Get<Timestamp>();
|
||||
if (input_stream_collection_) {
|
||||
cc->Outputs()
|
||||
.Tag("ITERABLE")
|
||||
.Add(input_stream_collection_.release(), loop_control_ts);
|
||||
} else {
|
||||
// Since there is no collection, inform downstream calculators to not
|
||||
// expect any packet by updating the timestamp bounds.
|
||||
cc->Outputs()
|
||||
.Tag("ITERABLE")
|
||||
.SetNextTimestampBound(Timestamp(loop_control_ts.Value() + 1));
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
std::unique_ptr<IterableT> input_stream_collection_;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_CORE_END_LOOP_CALCULATOR_H_
|
||||
@@ -12,23 +12,9 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <cstdlib>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include "mediapipe/calculators/core/packet_resampler_calculator.h"
|
||||
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "mediapipe/calculators/core/packet_resampler_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/deps/mathutil.h"
|
||||
#include "mediapipe/framework/deps/random_base.h"
|
||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/status_macros.h"
|
||||
#include "mediapipe/framework/tool/options_util.h"
|
||||
#include <memory>
|
||||
|
||||
namespace {
|
||||
|
||||
@@ -45,114 +31,7 @@ std::unique_ptr<RandomBase> CreateSecureRandom(const std::string& seed) {
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// This calculator is used to normalize the frequency of the packets
|
||||
// out of a stream. Given a desired frame rate, packets are going to be
|
||||
// removed or added to achieve it.
|
||||
//
|
||||
// The jitter feature is disabled by default. To enable it, you need to
|
||||
// implement CreateSecureRandom(const std::string&).
|
||||
//
|
||||
// The data stream may be either specified as the only stream (by index)
|
||||
// or as the stream with tag "DATA".
|
||||
//
|
||||
// The input and output streams may be accompanied by a VIDEO_HEADER
|
||||
// stream. This stream includes a VideoHeader at Timestamp::PreStream().
|
||||
// The input VideoHeader on the VIDEO_HEADER stream will always be updated
|
||||
// with the resampler frame rate no matter what the options value for
|
||||
// output_header is before being output on the output VIDEO_HEADER stream.
|
||||
// If the input VideoHeader is not available, then only the frame rate
|
||||
// value will be set in the output.
|
||||
//
|
||||
// Related:
|
||||
// packet_downsampler_calculator.cc: skips packets regardless of timestamps.
|
||||
class PacketResamplerCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
// Logic for Process() when jitter_ != 0.0.
|
||||
::mediapipe::Status ProcessWithJitter(CalculatorContext* cc);
|
||||
|
||||
// Logic for Process() when jitter_ == 0.0.
|
||||
::mediapipe::Status ProcessWithoutJitter(CalculatorContext* cc);
|
||||
|
||||
// Given the current count of periods that have passed, this returns
|
||||
// the next valid timestamp of the middle point of the next period:
|
||||
// if count is 0, it returns the first_timestamp_.
|
||||
// if count is 1, it returns the first_timestamp_ + period (corresponding
|
||||
// to the first tick using exact fps)
|
||||
// e.g. for frame_rate=30 and first_timestamp_=0:
|
||||
// 0: 0
|
||||
// 1: 33333
|
||||
// 2: 66667
|
||||
// 3: 100000
|
||||
//
|
||||
// Can only be used if jitter_ equals zero.
|
||||
Timestamp PeriodIndexToTimestamp(int64 index) const;
|
||||
|
||||
// Given a Timestamp, finds the closest sync Timestamp based on
|
||||
// first_timestamp_ and the desired fps.
|
||||
//
|
||||
// Can only be used if jitter_ equals zero.
|
||||
int64 TimestampToPeriodIndex(Timestamp timestamp) const;
|
||||
|
||||
// Outputs a packet if it is in range (start_time_, end_time_).
|
||||
void OutputWithinLimits(CalculatorContext* cc, const Packet& packet) const;
|
||||
|
||||
// The timestamp of the first packet received.
|
||||
Timestamp first_timestamp_;
|
||||
|
||||
// Number of frames per second (desired output frequency).
|
||||
double frame_rate_;
|
||||
|
||||
// Inverse of frame_rate_.
|
||||
int64 frame_time_usec_;
|
||||
|
||||
// Number of periods that have passed (= #packets sent to the output).
|
||||
//
|
||||
// Can only be used if jitter_ equals zero.
|
||||
int64 period_count_;
|
||||
|
||||
// The last packet that was received.
|
||||
Packet last_packet_;
|
||||
|
||||
VideoHeader video_header_;
|
||||
// The "DATA" input stream.
|
||||
CollectionItemId input_data_id_;
|
||||
// The "DATA" output stream.
|
||||
CollectionItemId output_data_id_;
|
||||
|
||||
// Indicator whether to flush last packet even if its timestamp is greater
|
||||
// than the final stream timestamp. Set to false when jitter_ is non-zero.
|
||||
bool flush_last_packet_;
|
||||
|
||||
// Jitter-related variables.
|
||||
std::unique_ptr<RandomBase> random_;
|
||||
double jitter_ = 0.0;
|
||||
Timestamp next_output_timestamp_;
|
||||
|
||||
// If specified, output timestamps are aligned with base_timestamp.
|
||||
// Otherwise, they are aligned with the first input timestamp.
|
||||
Timestamp base_timestamp_;
|
||||
|
||||
// If specified, only outputs at/after start_time are included.
|
||||
Timestamp start_time_;
|
||||
|
||||
// If specified, only outputs before end_time are included.
|
||||
Timestamp end_time_;
|
||||
|
||||
// If set, the output timestamps nearest to start_time and end_time
|
||||
// are included in the output, even if the nearest timestamp is not
|
||||
// between start_time and end_time.
|
||||
bool round_limits_;
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(PacketResamplerCalculator);
|
||||
|
||||
namespace {
|
||||
// Returns a TimestampDiff (assuming microseconds) corresponding to the
|
||||
// given time in seconds.
|
||||
@@ -233,6 +112,7 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
||||
<< Timestamp::kTimestampUnitsPerSecond;
|
||||
|
||||
frame_time_usec_ = static_cast<int64>(1000000.0 / frame_rate_);
|
||||
|
||||
video_header_.frame_rate = frame_rate_;
|
||||
|
||||
if (resampler_options.output_header() !=
|
||||
@@ -272,7 +152,10 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
||||
"SecureRandom is not available. With \"jitter\" specified, "
|
||||
"PacketResamplerCalculator processing cannot proceed.");
|
||||
}
|
||||
packet_reservoir_random_ = CreateSecureRandom(seed);
|
||||
}
|
||||
packet_reservoir_ =
|
||||
std::make_unique<PacketReservoir>(packet_reservoir_random_.get());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -287,6 +170,14 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
||||
}
|
||||
}
|
||||
if (jitter_ != 0.0 && random_ != nullptr) {
|
||||
// Packet reservior is used to make sure there's an output for every period,
|
||||
// e.g. partial period at the end of the stream.
|
||||
if (packet_reservoir_->IsEnabled() &&
|
||||
(first_timestamp_ == Timestamp::Unset() ||
|
||||
(cc->InputTimestamp() - next_output_timestamp_min_).Value() >= 0)) {
|
||||
auto curr_packet = cc->Inputs().Get(input_data_id_).Value();
|
||||
packet_reservoir_->AddSample(curr_packet);
|
||||
}
|
||||
MP_RETURN_IF_ERROR(ProcessWithJitter(cc));
|
||||
} else {
|
||||
MP_RETURN_IF_ERROR(ProcessWithoutJitter(cc));
|
||||
@@ -295,6 +186,20 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
void PacketResamplerCalculator::InitializeNextOutputTimestampWithJitter() {
|
||||
next_output_timestamp_min_ = first_timestamp_;
|
||||
next_output_timestamp_ =
|
||||
first_timestamp_ + frame_time_usec_ * random_->RandFloat();
|
||||
}
|
||||
|
||||
void PacketResamplerCalculator::UpdateNextOutputTimestampWithJitter() {
|
||||
packet_reservoir_->Clear();
|
||||
packet_reservoir_->Disable();
|
||||
next_output_timestamp_ +=
|
||||
frame_time_usec_ *
|
||||
((1.0 - jitter_) + 2.0 * jitter_ * random_->RandFloat());
|
||||
}
|
||||
|
||||
::mediapipe::Status PacketResamplerCalculator::ProcessWithJitter(
|
||||
CalculatorContext* cc) {
|
||||
RET_CHECK_GT(cc->InputTimestamp(), Timestamp::PreStream());
|
||||
@@ -302,29 +207,37 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
||||
|
||||
if (first_timestamp_ == Timestamp::Unset()) {
|
||||
first_timestamp_ = cc->InputTimestamp();
|
||||
next_output_timestamp_ =
|
||||
first_timestamp_ + frame_time_usec_ * random_->RandFloat();
|
||||
InitializeNextOutputTimestampWithJitter();
|
||||
if (first_timestamp_ == next_output_timestamp_) {
|
||||
OutputWithinLimits(
|
||||
cc,
|
||||
cc->Inputs().Get(input_data_id_).Value().At(next_output_timestamp_));
|
||||
UpdateNextOutputTimestampWithJitter();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
LOG_IF(WARNING, frame_time_usec_ <
|
||||
(cc->InputTimestamp() - last_packet_.Timestamp()).Value())
|
||||
<< "Adding jitter is meaningless when upsampling.";
|
||||
if (frame_time_usec_ <
|
||||
(cc->InputTimestamp() - last_packet_.Timestamp()).Value()) {
|
||||
LOG_FIRST_N(WARNING, 2)
|
||||
<< "Adding jitter is not very useful when upsampling.";
|
||||
}
|
||||
|
||||
const int64 curr_diff =
|
||||
(next_output_timestamp_ - cc->InputTimestamp()).Value();
|
||||
const int64 last_diff =
|
||||
(next_output_timestamp_ - last_packet_.Timestamp()).Value();
|
||||
if (curr_diff * last_diff > 0) {
|
||||
return ::mediapipe::OkStatus();
|
||||
while (true) {
|
||||
const int64 last_diff =
|
||||
(next_output_timestamp_ - last_packet_.Timestamp()).Value();
|
||||
RET_CHECK_GT(last_diff, 0);
|
||||
const int64 curr_diff =
|
||||
(next_output_timestamp_ - cc->InputTimestamp()).Value();
|
||||
if (curr_diff > 0) {
|
||||
break;
|
||||
}
|
||||
OutputWithinLimits(cc, (std::abs(curr_diff) > last_diff
|
||||
? last_packet_
|
||||
: cc->Inputs().Get(input_data_id_).Value())
|
||||
.At(next_output_timestamp_));
|
||||
UpdateNextOutputTimestampWithJitter();
|
||||
}
|
||||
OutputWithinLimits(cc, (std::abs(curr_diff) > std::abs(last_diff)
|
||||
? last_packet_
|
||||
: cc->Inputs().Get(input_data_id_).Value())
|
||||
.At(next_output_timestamp_));
|
||||
next_output_timestamp_ +=
|
||||
frame_time_usec_ *
|
||||
((1.0 - jitter_) + 2.0 * jitter_ * random_->RandFloat());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -405,6 +318,9 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
||||
OutputWithinLimits(cc,
|
||||
last_packet_.At(PeriodIndexToTimestamp(period_count_)));
|
||||
}
|
||||
if (!packet_reservoir_->IsEmpty()) {
|
||||
OutputWithinLimits(cc, packet_reservoir_->GetSample());
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,168 @@
|
||||
#ifndef MEDIAPIPE_CALCULATORS_CORE_PACKET_RESAMPLER_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_CORE_PACKET_RESAMPLER_CALCULATOR_H_
|
||||
|
||||
#include <cstdlib>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "mediapipe/calculators/core/packet_resampler_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/deps/mathutil.h"
|
||||
#include "mediapipe/framework/deps/random_base.h"
|
||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/status_macros.h"
|
||||
#include "mediapipe/framework/tool/options_util.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
class PacketReservoir {
|
||||
public:
|
||||
PacketReservoir(RandomBase* rng) : rng_(rng) {}
|
||||
// Replace candidate with current packet with 1/count_ probability.
|
||||
void AddSample(Packet sample) {
|
||||
if (rng_->UnbiasedUniform(++count_) == 0) {
|
||||
reservoir_ = sample;
|
||||
}
|
||||
}
|
||||
bool IsEnabled() { return rng_ && enabled_; }
|
||||
void Disable() {
|
||||
if (enabled_) enabled_ = false;
|
||||
}
|
||||
void Clear() { count_ = 0; }
|
||||
bool IsEmpty() { return count_ == 0; }
|
||||
Packet GetSample() { return reservoir_; }
|
||||
|
||||
private:
|
||||
RandomBase* rng_;
|
||||
bool enabled_ = true;
|
||||
int32 count_ = 0;
|
||||
Packet reservoir_;
|
||||
};
|
||||
|
||||
// This calculator is used to normalize the frequency of the packets
|
||||
// out of a stream. Given a desired frame rate, packets are going to be
|
||||
// removed or added to achieve it.
|
||||
//
|
||||
// The jitter feature is disabled by default. To enable it, you need to
|
||||
// implement CreateSecureRandom(const std::string&).
|
||||
//
|
||||
// The data stream may be either specified as the only stream (by index)
|
||||
// or as the stream with tag "DATA".
|
||||
//
|
||||
// The input and output streams may be accompanied by a VIDEO_HEADER
|
||||
// stream. This stream includes a VideoHeader at Timestamp::PreStream().
|
||||
// The input VideoHeader on the VIDEO_HEADER stream will always be updated
|
||||
// with the resampler frame rate no matter what the options value for
|
||||
// output_header is before being output on the output VIDEO_HEADER stream.
|
||||
// If the input VideoHeader is not available, then only the frame rate
|
||||
// value will be set in the output.
|
||||
//
|
||||
// Related:
|
||||
// packet_downsampler_calculator.cc: skips packets regardless of timestamps.
|
||||
class PacketResamplerCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
// Calculates the first sampled timestamp that incorporates a jittering
|
||||
// offset.
|
||||
void InitializeNextOutputTimestampWithJitter();
|
||||
// Calculates the next sampled timestamp that incorporates a jittering offset.
|
||||
void UpdateNextOutputTimestampWithJitter();
|
||||
|
||||
// Logic for Process() when jitter_ != 0.0.
|
||||
::mediapipe::Status ProcessWithJitter(CalculatorContext* cc);
|
||||
|
||||
// Logic for Process() when jitter_ == 0.0.
|
||||
::mediapipe::Status ProcessWithoutJitter(CalculatorContext* cc);
|
||||
|
||||
// Given the current count of periods that have passed, this returns
|
||||
// the next valid timestamp of the middle point of the next period:
|
||||
// if count is 0, it returns the first_timestamp_.
|
||||
// if count is 1, it returns the first_timestamp_ + period (corresponding
|
||||
// to the first tick using exact fps)
|
||||
// e.g. for frame_rate=30 and first_timestamp_=0:
|
||||
// 0: 0
|
||||
// 1: 33333
|
||||
// 2: 66667
|
||||
// 3: 100000
|
||||
//
|
||||
// Can only be used if jitter_ equals zero.
|
||||
Timestamp PeriodIndexToTimestamp(int64 index) const;
|
||||
|
||||
// Given a Timestamp, finds the closest sync Timestamp based on
|
||||
// first_timestamp_ and the desired fps.
|
||||
//
|
||||
// Can only be used if jitter_ equals zero.
|
||||
int64 TimestampToPeriodIndex(Timestamp timestamp) const;
|
||||
|
||||
// Outputs a packet if it is in range (start_time_, end_time_).
|
||||
void OutputWithinLimits(CalculatorContext* cc, const Packet& packet) const;
|
||||
|
||||
// The timestamp of the first packet received.
|
||||
Timestamp first_timestamp_;
|
||||
|
||||
// Number of frames per second (desired output frequency).
|
||||
double frame_rate_;
|
||||
|
||||
// Inverse of frame_rate_.
|
||||
int64 frame_time_usec_;
|
||||
|
||||
// Number of periods that have passed (= #packets sent to the output).
|
||||
//
|
||||
// Can only be used if jitter_ equals zero.
|
||||
int64 period_count_;
|
||||
|
||||
// The last packet that was received.
|
||||
Packet last_packet_;
|
||||
|
||||
VideoHeader video_header_;
|
||||
// The "DATA" input stream.
|
||||
CollectionItemId input_data_id_;
|
||||
// The "DATA" output stream.
|
||||
CollectionItemId output_data_id_;
|
||||
|
||||
// Indicator whether to flush last packet even if its timestamp is greater
|
||||
// than the final stream timestamp. Set to false when jitter_ is non-zero.
|
||||
bool flush_last_packet_;
|
||||
|
||||
// Jitter-related variables.
|
||||
std::unique_ptr<RandomBase> random_;
|
||||
double jitter_ = 0.0;
|
||||
Timestamp next_output_timestamp_;
|
||||
Timestamp next_output_timestamp_min_;
|
||||
|
||||
// If specified, output timestamps are aligned with base_timestamp.
|
||||
// Otherwise, they are aligned with the first input timestamp.
|
||||
Timestamp base_timestamp_;
|
||||
|
||||
// If specified, only outputs at/after start_time are included.
|
||||
Timestamp start_time_;
|
||||
|
||||
// If specified, only outputs before end_time are included.
|
||||
Timestamp end_time_;
|
||||
|
||||
// If set, the output timestamps nearest to start_time and end_time
|
||||
// are included in the output, even if the nearest timestamp is not
|
||||
// between start_time and end_time.W
|
||||
bool round_limits_;
|
||||
|
||||
// packet reservior used for sampling random packet out of partial
|
||||
// period when jitter is enabled
|
||||
std::unique_ptr<PacketReservoir> packet_reservoir_;
|
||||
// random number generator used in packet_reservior_.
|
||||
std::unique_ptr<RandomBase> packet_reservoir_random_;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
#endif // MEDIAPIPE_CALCULATORS_CORE_PACKET_RESAMPLER_CALCULATOR_H_
|
||||
@@ -12,6 +12,8 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "mediapipe/calculators/core/packet_resampler_calculator.h"
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
@@ -29,7 +31,6 @@
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
// A simple version of CalculatorRunner with built-in convenience
|
||||
// methods for setting inputs from a vector and checking outputs
|
||||
// against expected outputs (both timestamps and contents).
|
||||
|
||||
@@ -0,0 +1,304 @@
|
||||
// Copyright 2019 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.
|
||||
//
|
||||
// Declaration of PacketThinnerCalculator.
|
||||
|
||||
#include <cmath> // for ceil
|
||||
#include <memory>
|
||||
|
||||
#include "mediapipe/calculators/core/packet_thinner_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
const double kTimebaseUs = 1000000; // Microseconds.
|
||||
const char* const kPeriodTag = "PERIOD";
|
||||
} // namespace
|
||||
|
||||
// This calculator is used to thin an input stream of Packets.
|
||||
// An example application would be to sample decoded frames of video
|
||||
// at a coarser temporal resolution. Unless otherwise stated, all
|
||||
// timestamps are in units of microseconds.
|
||||
//
|
||||
// Thinning can be accomplished in one of two ways:
|
||||
// 1) asynchronous thinning (known below as async):
|
||||
// Algorithm does not rely on a master clock and is parameterized only
|
||||
// by a single option -- the period. Once a packet is emitted, the
|
||||
// thinner will discard subsequent packets for the duration of the period
|
||||
// [Analogous to a refractory period during which packet emission is
|
||||
// suppressed.]
|
||||
// Packets arriving before start_time are discarded, as are packets
|
||||
// arriving at or after end_time.
|
||||
// 2) synchronous thinning (known below as sync):
|
||||
// There are two variants of this algorithm, both parameterized by a
|
||||
// start_time and a period. As in (1), packets arriving before start_time
|
||||
// or at/after end_time are discarded. Otherwise, at most one packet is
|
||||
// emitted during a period, centered at timestamps generated by the
|
||||
// expression:
|
||||
// start_time + i * period [where i is a non-negative integer]
|
||||
// During each period, the packet closest to the generated timestamp is
|
||||
// emitted (latest in the case of ties). In the first variant
|
||||
// (sync_output_timestamps = true), the emitted packet is output at the
|
||||
// generated timestamp. In the second variant, the packet is output at
|
||||
// its original timestamp. Both variants emit exactly the same packets,
|
||||
// but at different timestamps.
|
||||
//
|
||||
// Thinning period can be provided in the calculator options or via a
|
||||
// side packet with the tag "PERIOD".
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "PacketThinnerCalculator"
|
||||
// input_stream: "signal"
|
||||
// output_stream: "output"
|
||||
// options {
|
||||
// [mediapipe.PacketThinnerCalculatorOptions.ext] {
|
||||
// thinner_type: SYNC
|
||||
// period: 10
|
||||
// sync_output_timestamps: true
|
||||
// update_frame_rate: false
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
class PacketThinnerCalculator : public CalculatorBase {
|
||||
public:
|
||||
PacketThinnerCalculator() {}
|
||||
~PacketThinnerCalculator() override {}
|
||||
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Index(0).SetAny();
|
||||
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
|
||||
if (cc->InputSidePackets().HasTag(kPeriodTag)) {
|
||||
cc->InputSidePackets().Tag(kPeriodTag).Set<int64>();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
if (cc->InputTimestamp() < start_time_) {
|
||||
return ::mediapipe::OkStatus(); // Drop packets before start_time_.
|
||||
} else if (cc->InputTimestamp() >= end_time_) {
|
||||
if (!cc->Outputs().Index(0).IsClosed()) {
|
||||
cc->Outputs()
|
||||
.Index(0)
|
||||
.Close(); // No more Packets will be output after end_time_.
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
} else {
|
||||
return thinner_type_ == PacketThinnerCalculatorOptions::ASYNC
|
||||
? AsyncThinnerProcess(cc)
|
||||
: SyncThinnerProcess(cc);
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
// Implementation of ASYNC and SYNC versions of thinner algorithm.
|
||||
::mediapipe::Status AsyncThinnerProcess(CalculatorContext* cc);
|
||||
::mediapipe::Status SyncThinnerProcess(CalculatorContext* cc);
|
||||
|
||||
// Cached option.
|
||||
PacketThinnerCalculatorOptions::ThinnerType thinner_type_;
|
||||
|
||||
// Given a Timestamp, finds the closest sync Timestamp
|
||||
// based on start_time_ and period_. This can be earlier or
|
||||
// later than given Timestamp, but is guaranteed to be within
|
||||
// half a period_.
|
||||
Timestamp NearestSyncTimestamp(Timestamp now) const;
|
||||
|
||||
// Cached option used by both async and sync thinners.
|
||||
TimestampDiff period_; // Interval during which only one packet is emitted.
|
||||
Timestamp start_time_; // Cached option - default Timestamp::Min()
|
||||
Timestamp end_time_; // Cached option - default Timestamp::Max()
|
||||
|
||||
// Only used by async thinner:
|
||||
Timestamp next_valid_timestamp_; // Suppress packets until this timestamp.
|
||||
|
||||
// Only used by sync thinner:
|
||||
Packet saved_packet_; // Best packet not yet emitted.
|
||||
bool sync_output_timestamps_; // Cached option.
|
||||
};
|
||||
REGISTER_CALCULATOR(PacketThinnerCalculator);
|
||||
|
||||
namespace {
|
||||
TimestampDiff abs(TimestampDiff t) { return t < 0 ? -t : t; }
|
||||
} // namespace
|
||||
|
||||
::mediapipe::Status PacketThinnerCalculator::Open(CalculatorContext* cc) {
|
||||
auto& options = cc->Options<PacketThinnerCalculatorOptions>();
|
||||
|
||||
thinner_type_ = options.thinner_type();
|
||||
// This check enables us to assume only two thinner types exist in Process()
|
||||
CHECK(thinner_type_ == PacketThinnerCalculatorOptions::ASYNC ||
|
||||
thinner_type_ == PacketThinnerCalculatorOptions::SYNC)
|
||||
<< "Unsupported thinner type.";
|
||||
|
||||
if (thinner_type_ == PacketThinnerCalculatorOptions::ASYNC) {
|
||||
// ASYNC thinner outputs packets with the same timestamp as their input so
|
||||
// its safe to SetOffset(0). SYNC thinner manipulates timestamps of its
|
||||
// output so we don't do this for that case.
|
||||
cc->SetOffset(0);
|
||||
}
|
||||
|
||||
if (cc->InputSidePackets().HasTag(kPeriodTag)) {
|
||||
period_ =
|
||||
TimestampDiff(cc->InputSidePackets().Tag(kPeriodTag).Get<int64>());
|
||||
} else {
|
||||
period_ = TimestampDiff(options.period());
|
||||
}
|
||||
CHECK_LT(TimestampDiff(0), period_) << "Specified period must be positive.";
|
||||
|
||||
if (options.has_start_time()) {
|
||||
start_time_ = Timestamp(options.start_time());
|
||||
} else if (thinner_type_ == PacketThinnerCalculatorOptions::ASYNC) {
|
||||
start_time_ = Timestamp::Min();
|
||||
} else {
|
||||
start_time_ = Timestamp(0);
|
||||
}
|
||||
|
||||
end_time_ =
|
||||
options.has_end_time() ? Timestamp(options.end_time()) : Timestamp::Max();
|
||||
CHECK_LT(start_time_, end_time_)
|
||||
<< "Invalid PacketThinner: start_time must be earlier than end_time";
|
||||
|
||||
sync_output_timestamps_ = options.sync_output_timestamps();
|
||||
|
||||
next_valid_timestamp_ = start_time_;
|
||||
// Drop packets until this time.
|
||||
cc->Outputs().Index(0).SetNextTimestampBound(start_time_);
|
||||
|
||||
if (!cc->Inputs().Index(0).Header().IsEmpty()) {
|
||||
if (options.update_frame_rate()) {
|
||||
const VideoHeader& video_header =
|
||||
cc->Inputs().Index(0).Header().Get<VideoHeader>();
|
||||
double new_frame_rate;
|
||||
if (thinner_type_ == PacketThinnerCalculatorOptions::ASYNC) {
|
||||
new_frame_rate =
|
||||
video_header.frame_rate /
|
||||
ceil(video_header.frame_rate * options.period() / kTimebaseUs);
|
||||
} else {
|
||||
const double sampling_rate = kTimebaseUs / options.period();
|
||||
new_frame_rate = video_header.frame_rate < sampling_rate
|
||||
? video_header.frame_rate
|
||||
: sampling_rate;
|
||||
}
|
||||
std::unique_ptr<VideoHeader> header(new VideoHeader);
|
||||
header->format = video_header.format;
|
||||
header->width = video_header.width;
|
||||
header->height = video_header.height;
|
||||
header->frame_rate = new_frame_rate;
|
||||
cc->Outputs().Index(0).SetHeader(Adopt(header.release()));
|
||||
} else {
|
||||
cc->Outputs().Index(0).SetHeader(cc->Inputs().Index(0).Header());
|
||||
}
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status PacketThinnerCalculator::Close(CalculatorContext* cc) {
|
||||
// Emit any saved packets before quitting.
|
||||
if (!saved_packet_.IsEmpty()) {
|
||||
// Only sync thinner should have saved packets.
|
||||
CHECK_EQ(PacketThinnerCalculatorOptions::SYNC, thinner_type_);
|
||||
if (sync_output_timestamps_) {
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
saved_packet_.At(NearestSyncTimestamp(saved_packet_.Timestamp())));
|
||||
} else {
|
||||
cc->Outputs().Index(0).AddPacket(saved_packet_);
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status PacketThinnerCalculator::AsyncThinnerProcess(
|
||||
CalculatorContext* cc) {
|
||||
if (cc->InputTimestamp() >= next_valid_timestamp_) {
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
cc->Inputs().Index(0).Value()); // Emit current packet.
|
||||
next_valid_timestamp_ = cc->InputTimestamp() + period_;
|
||||
// Guaranteed not to emit packets seen during refractory period.
|
||||
cc->Outputs().Index(0).SetNextTimestampBound(next_valid_timestamp_);
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status PacketThinnerCalculator::SyncThinnerProcess(
|
||||
CalculatorContext* cc) {
|
||||
if (saved_packet_.IsEmpty()) {
|
||||
// If no packet has been saved, store the current packet.
|
||||
saved_packet_ = cc->Inputs().Index(0).Value();
|
||||
cc->Outputs().Index(0).SetNextTimestampBound(
|
||||
sync_output_timestamps_ ? NearestSyncTimestamp(cc->InputTimestamp())
|
||||
: cc->InputTimestamp());
|
||||
} else {
|
||||
// Saved packet exists -- update or emit.
|
||||
const Timestamp saved = saved_packet_.Timestamp();
|
||||
const Timestamp saved_sync = NearestSyncTimestamp(saved);
|
||||
const Timestamp now = cc->InputTimestamp();
|
||||
const Timestamp now_sync = NearestSyncTimestamp(now);
|
||||
CHECK_LE(saved_sync, now_sync);
|
||||
if (saved_sync == now_sync) {
|
||||
// Saved Packet is in same interval as current packet.
|
||||
// Replace saved packet with current if it is at least as
|
||||
// central as the saved packet wrt temporal interval.
|
||||
// [We break ties in favor of fresher packets]
|
||||
if (abs(now - now_sync) <= abs(saved - saved_sync)) {
|
||||
saved_packet_ = cc->Inputs().Index(0).Value();
|
||||
}
|
||||
} else {
|
||||
// Saved packet is the best packet from earlier interval: emit!
|
||||
if (sync_output_timestamps_) {
|
||||
cc->Outputs().Index(0).AddPacket(saved_packet_.At(saved_sync));
|
||||
cc->Outputs().Index(0).SetNextTimestampBound(now_sync);
|
||||
} else {
|
||||
cc->Outputs().Index(0).AddPacket(saved_packet_);
|
||||
cc->Outputs().Index(0).SetNextTimestampBound(now);
|
||||
}
|
||||
// Current packet is the first one we've seen from new interval -- save!
|
||||
saved_packet_ = cc->Inputs().Index(0).Value();
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
Timestamp PacketThinnerCalculator::NearestSyncTimestamp(Timestamp now) const {
|
||||
CHECK_NE(start_time_, Timestamp::Unset())
|
||||
<< "Method only valid for sync thinner calculator.";
|
||||
|
||||
// Computation is done using int64 arithmetic. No easy way to avoid
|
||||
// since Timestamps don't support div and multiply.
|
||||
const int64 now64 = now.Value();
|
||||
const int64 start64 = start_time_.Value();
|
||||
const int64 period64 = period_.Value();
|
||||
CHECK_LE(0, period64);
|
||||
|
||||
// Round now64 to its closest interval (units of period64).
|
||||
int64 sync64 =
|
||||
(now64 - start64 + period64 / 2) / period64 * period64 + start64;
|
||||
CHECK_LE(abs(now64 - sync64), period64 / 2)
|
||||
<< "start64: " << start64 << "; now64: " << now64
|
||||
<< "; sync64: " << sync64;
|
||||
|
||||
return Timestamp(sync64);
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,66 @@
|
||||
// Copyright 2018 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;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
|
||||
message PacketThinnerCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional PacketThinnerCalculatorOptions ext = 288533508;
|
||||
}
|
||||
|
||||
enum ThinnerType {
|
||||
ASYNC = 1; // Asynchronous thinner, described below [default].
|
||||
SYNC = 2; // Synchronous thinner, also described below.
|
||||
}
|
||||
optional ThinnerType thinner_type = 1 [default = ASYNC];
|
||||
|
||||
// The period (in microsecond) specifies the temporal interval during which
|
||||
// only a single packet is emitted in the output stream. Has subtly different
|
||||
// semantics depending on the thinner type, as follows.
|
||||
//
|
||||
// Async thinner: this option is a refractory period -- once a packet is
|
||||
// emitted, we guarantee that no packets will be emitted for period ticks.
|
||||
//
|
||||
// Sync thinner: the period specifies a temporal interval during which
|
||||
// only one packet is emitted. The emitted packet is guaranteed to be
|
||||
// the one closest to the center of the temporal interval (no guarantee on
|
||||
// how ties are broken). More specifically,
|
||||
// intervals are centered at start_time + i * period
|
||||
// (for non-negative integers i).
|
||||
// Thus, each interval extends period/2 ticks before and after its center.
|
||||
// Additionally, in the sync thinner any packets earlier than start_time
|
||||
// are discarded and the thinner calls Close() once timestamp equals or
|
||||
// exceeds end_time.
|
||||
optional int64 period = 2 [default = 1];
|
||||
|
||||
// Packets before start_time and at/after end_time are discarded.
|
||||
// Additionally, for a sync thinner, start time specifies the center of
|
||||
// time invervals as described above and therefore should be set explicitly.
|
||||
optional int64 start_time = 3; // If not specified, set to 0 for SYNC type,
|
||||
// and set to Timestamp::Min() for ASYNC type.
|
||||
optional int64 end_time = 4; // Set to Timestamp::Max() if not specified.
|
||||
|
||||
// Whether the timestamps of packets emitted by sync thinner should
|
||||
// correspond to the center of their corresponding temporal interval.
|
||||
// If false, packets emitted using original timestamp (as in async thinner).
|
||||
optional bool sync_output_timestamps = 5 [default = true];
|
||||
|
||||
// If true, update the frame rate in the header, if it's available, to an
|
||||
// estimated frame rate due to the sampling.
|
||||
optional bool update_frame_rate = 6 [default = false];
|
||||
}
|
||||
@@ -0,0 +1,357 @@
|
||||
// Copyright 2019 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 <vector>
|
||||
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "mediapipe/calculators/core/packet_thinner_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
// A simple version of CalculatorRunner with built-in convenience methods for
|
||||
// setting inputs from a vector and checking outputs against a vector of
|
||||
// expected outputs.
|
||||
class SimpleRunner : public CalculatorRunner {
|
||||
public:
|
||||
explicit SimpleRunner(const CalculatorOptions& options)
|
||||
: CalculatorRunner("PacketThinnerCalculator", options) {
|
||||
SetNumInputs(1);
|
||||
SetNumOutputs(1);
|
||||
SetNumInputSidePackets(0);
|
||||
}
|
||||
|
||||
explicit SimpleRunner(const CalculatorGraphConfig::Node& node)
|
||||
: CalculatorRunner(node) {}
|
||||
|
||||
void SetInput(const std::vector<int>& timestamp_list) {
|
||||
MutableInputs()->Index(0).packets.clear();
|
||||
for (const int ts : timestamp_list) {
|
||||
MutableInputs()->Index(0).packets.push_back(
|
||||
MakePacket<std::string>(absl::StrCat("Frame #", ts))
|
||||
.At(Timestamp(ts)));
|
||||
}
|
||||
}
|
||||
|
||||
void SetFrameRate(const double frame_rate) {
|
||||
auto video_header = absl::make_unique<VideoHeader>();
|
||||
video_header->frame_rate = frame_rate;
|
||||
MutableInputs()->Index(0).header = Adopt(video_header.release());
|
||||
}
|
||||
|
||||
std::vector<int64> GetOutputTimestamps() const {
|
||||
std::vector<int64> timestamps;
|
||||
for (const Packet& packet : Outputs().Index(0).packets) {
|
||||
timestamps.emplace_back(packet.Timestamp().Value());
|
||||
}
|
||||
return timestamps;
|
||||
}
|
||||
|
||||
double GetFrameRate() const {
|
||||
CHECK(!Outputs().Index(0).header.IsEmpty());
|
||||
return Outputs().Index(0).header.Get<VideoHeader>().frame_rate;
|
||||
}
|
||||
};
|
||||
|
||||
// Check that thinner respects start_time and end_time options.
|
||||
// We only test with one thinner because the logic for start & end time
|
||||
// handling is shared across both types of thinner in Process().
|
||||
TEST(PacketThinnerCalculatorTest, StartAndEndTimeTest) {
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||
extension->set_period(5);
|
||||
extension->set_start_time(4);
|
||||
extension->set_end_time(12);
|
||||
SimpleRunner runner(options);
|
||||
runner.SetInput({2, 3, 5, 7, 11, 13, 17, 19, 23, 29});
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {5, 11};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, AsyncUniformStreamThinningTest) {
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||
extension->set_period(5);
|
||||
SimpleRunner runner(options);
|
||||
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {2, 8, 14};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, ASyncUniformStreamThinningTestBySidePacket) {
|
||||
// Note: sync runner but outputting *original* timestamps.
|
||||
CalculatorGraphConfig::Node node;
|
||||
node.set_calculator("PacketThinnerCalculator");
|
||||
node.add_input_side_packet("PERIOD:period");
|
||||
node.add_input_stream("input_stream");
|
||||
node.add_output_stream("output_stream");
|
||||
auto* extension = node.mutable_options()->MutableExtension(
|
||||
PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||
extension->set_start_time(0);
|
||||
extension->set_sync_output_timestamps(false);
|
||||
|
||||
SimpleRunner runner(node);
|
||||
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||
runner.MutableSidePackets()->Tag("PERIOD") = MakePacket<int64>(5);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {2, 8, 14};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, SyncUniformStreamThinningTest1) {
|
||||
// Note: sync runner but outputting *original* timestamps.
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||
extension->set_start_time(0);
|
||||
extension->set_period(5);
|
||||
extension->set_sync_output_timestamps(false);
|
||||
SimpleRunner runner(options);
|
||||
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {2, 6, 10, 14};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, SyncUniformStreamThinningTestBySidePacket1) {
|
||||
// Note: sync runner but outputting *original* timestamps.
|
||||
CalculatorGraphConfig::Node node;
|
||||
node.set_calculator("PacketThinnerCalculator");
|
||||
node.add_input_side_packet("PERIOD:period");
|
||||
node.add_input_stream("input_stream");
|
||||
node.add_output_stream("output_stream");
|
||||
auto* extension = node.mutable_options()->MutableExtension(
|
||||
PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||
extension->set_start_time(0);
|
||||
extension->set_sync_output_timestamps(false);
|
||||
|
||||
SimpleRunner runner(node);
|
||||
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||
runner.MutableSidePackets()->Tag("PERIOD") = MakePacket<int64>(5);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {2, 6, 10, 14};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, SyncUniformStreamThinningTest2) {
|
||||
// Same test but now with synced timestamps.
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||
extension->set_start_time(0);
|
||||
extension->set_period(5);
|
||||
extension->set_sync_output_timestamps(true);
|
||||
SimpleRunner runner(options);
|
||||
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {0, 5, 10, 15};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
}
|
||||
|
||||
// Test: Given a stream with timestamps corresponding to first ten prime numbers
|
||||
// and period of 5, confirm whether timestamps of thinner stream matches
|
||||
// expectations.
|
||||
TEST(PacketThinnerCalculatorTest, PrimeStreamThinningTest1) {
|
||||
// ASYNC thinner.
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||
extension->set_period(5);
|
||||
SimpleRunner runner(options);
|
||||
runner.SetInput({2, 3, 5, 7, 11, 13, 17, 19, 23, 29});
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {2, 7, 13, 19, 29};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, PrimeStreamThinningTest2) {
|
||||
// SYNC with original timestamps.
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||
extension->set_start_time(0);
|
||||
extension->set_period(5);
|
||||
extension->set_sync_output_timestamps(false);
|
||||
SimpleRunner runner(options);
|
||||
runner.SetInput({2, 3, 5, 7, 11, 13, 17, 19, 23, 29});
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {2, 5, 11, 17, 19, 23, 29};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
}
|
||||
|
||||
// Confirm that Calculator correctly handles boundary cases.
|
||||
TEST(PacketThinnerCalculatorTest, BoundaryTimestampTest1) {
|
||||
// Odd period, negative start_time
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||
extension->set_start_time(-10);
|
||||
extension->set_period(5);
|
||||
extension->set_sync_output_timestamps(true);
|
||||
SimpleRunner runner(options);
|
||||
// Two timestamps falling on either side of a period boundary.
|
||||
runner.SetInput({2, 3});
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {0, 5};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, BoundaryTimestampTest2) {
|
||||
// Even period, negative start_time, negative packet timestamps.
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||
extension->set_start_time(-144);
|
||||
extension->set_period(6);
|
||||
extension->set_sync_output_timestamps(true);
|
||||
SimpleRunner runner(options);
|
||||
// Two timestamps falling on either side of a period boundary.
|
||||
runner.SetInput({-4, -3, 8, 9});
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {-6, 0, 6, 12};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, FrameRateTest1) {
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||
extension->set_period(5);
|
||||
extension->set_update_frame_rate(true);
|
||||
SimpleRunner runner(options);
|
||||
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||
runner.SetFrameRate(1000000.0 / 2);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {2, 8, 14};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
// The true sampling period is 6.
|
||||
EXPECT_DOUBLE_EQ(1000000.0 / 6, runner.GetFrameRate());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, FrameRateTest2) {
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||
extension->set_period(5);
|
||||
extension->set_update_frame_rate(true);
|
||||
SimpleRunner runner(options);
|
||||
runner.SetInput({8, 16, 24, 32, 40, 48, 56});
|
||||
runner.SetFrameRate(1000000.0 / 8);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
const std::vector<int64> expected_timestamps = {8, 16, 24, 32, 40, 48, 56};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
// The true sampling period is still 8.
|
||||
EXPECT_DOUBLE_EQ(1000000.0 / 8, runner.GetFrameRate());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, FrameRateTest3) {
|
||||
// Note: sync runner but outputting *original* timestamps.
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||
extension->set_start_time(0);
|
||||
extension->set_period(5);
|
||||
extension->set_sync_output_timestamps(false);
|
||||
extension->set_update_frame_rate(true);
|
||||
SimpleRunner runner(options);
|
||||
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||
runner.SetFrameRate(1000000.0 / 2);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {2, 6, 10, 14};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
// The true (long-run) sampling period is 5.
|
||||
EXPECT_DOUBLE_EQ(1000000.0 / 5, runner.GetFrameRate());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, FrameRateTest4) {
|
||||
// Same test but now with synced timestamps.
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||
extension->set_start_time(0);
|
||||
extension->set_period(5);
|
||||
extension->set_sync_output_timestamps(true);
|
||||
extension->set_update_frame_rate(true);
|
||||
SimpleRunner runner(options);
|
||||
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||
runner.SetFrameRate(1000000.0 / 2);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {0, 5, 10, 15};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
// The true (long-run) sampling period is 5.
|
||||
EXPECT_DOUBLE_EQ(1000000.0 / 5, runner.GetFrameRate());
|
||||
}
|
||||
|
||||
TEST(PacketThinnerCalculatorTest, FrameRateTest5) {
|
||||
CalculatorOptions options;
|
||||
auto* extension =
|
||||
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||
extension->set_start_time(0);
|
||||
extension->set_period(5);
|
||||
extension->set_sync_output_timestamps(true);
|
||||
extension->set_update_frame_rate(true);
|
||||
SimpleRunner runner(options);
|
||||
runner.SetInput({8, 16, 24, 32, 40, 48, 56});
|
||||
runner.SetFrameRate(1000000.0 / 8);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
const std::vector<int64> expected_timestamps = {10, 15, 25, 30, 40, 50, 55};
|
||||
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||
// The true (long-run) sampling period is 8.
|
||||
EXPECT_DOUBLE_EQ(1000000.0 / 8, runner.GetFrameRate());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -17,6 +17,7 @@
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
@@ -86,6 +87,7 @@ class PreviousLoopbackCalculator : public CalculatorBase {
|
||||
main_ts_.pop_front();
|
||||
}
|
||||
}
|
||||
auto& loop_out = cc->Outputs().Get(loop_out_id_);
|
||||
|
||||
while (!main_ts_.empty() && !loopback_packets_.empty()) {
|
||||
Timestamp main_timestamp = main_ts_.front();
|
||||
@@ -95,13 +97,32 @@ class PreviousLoopbackCalculator : public CalculatorBase {
|
||||
|
||||
if (previous_loopback.IsEmpty()) {
|
||||
// TODO: SetCompleteTimestampBound would be more useful.
|
||||
cc->Outputs()
|
||||
.Get(loop_out_id_)
|
||||
.SetNextTimestampBound(main_timestamp + 1);
|
||||
loop_out.SetNextTimestampBound(main_timestamp + 1);
|
||||
} else {
|
||||
cc->Outputs().Get(loop_out_id_).AddPacket(std::move(previous_loopback));
|
||||
loop_out.AddPacket(std::move(previous_loopback));
|
||||
}
|
||||
}
|
||||
|
||||
// In case of an empty loopback input, the next timestamp bound for
|
||||
// loopback input is the loopback timestamp + 1. The next timestamp bound
|
||||
// for output is set and the main_ts_ vector is truncated accordingly.
|
||||
if (loopback_packet.IsEmpty() &&
|
||||
loopback_packet.Timestamp() != Timestamp::Unstarted()) {
|
||||
Timestamp loopback_bound =
|
||||
loopback_packet.Timestamp().NextAllowedInStream();
|
||||
while (!main_ts_.empty() && main_ts_.front() <= loopback_bound) {
|
||||
main_ts_.pop_front();
|
||||
}
|
||||
if (main_ts_.empty()) {
|
||||
loop_out.SetNextTimestampBound(loopback_bound.NextAllowedInStream());
|
||||
}
|
||||
}
|
||||
if (!main_ts_.empty()) {
|
||||
loop_out.SetNextTimestampBound(main_ts_.front());
|
||||
}
|
||||
if (cc->Inputs().Get(main_id_).IsDone() && main_ts_.empty()) {
|
||||
loop_out.Close();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
|
||||
@@ -93,19 +93,178 @@ TEST(PreviousLoopbackCalculator, CorrectTimestamps) {
|
||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1}));
|
||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(1, -1));
|
||||
|
||||
send_packet("in", 2);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 2}));
|
||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(2, 1));
|
||||
|
||||
send_packet("in", 5);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 5}));
|
||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(5, 1));
|
||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 2, 5}));
|
||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(5, 2));
|
||||
|
||||
send_packet("in", 15);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 5, 15}));
|
||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 2, 5, 15}));
|
||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(15, 5));
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
// A Calculator that outputs a summary packet in CalculatorBase::Close().
|
||||
class PacketOnCloseCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Index(0).Set<int>();
|
||||
cc->Outputs().Index(0).Set<int>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) final {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||
sum_ += cc->Inputs().Index(0).Value().Get<int>();
|
||||
cc->Outputs().Index(0).AddPacket(cc->Inputs().Index(0).Value());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Close(CalculatorContext* cc) final {
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
MakePacket<int>(sum_).At(Timestamp::Max()));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
int sum_ = 0;
|
||||
};
|
||||
REGISTER_CALCULATOR(PacketOnCloseCalculator);
|
||||
|
||||
// Demonstrates that all ouput and input streams in PreviousLoopbackCalculator
|
||||
// will close as expected when all graph input streams are closed.
|
||||
TEST(PreviousLoopbackCalculator, ClosesCorrectly) {
|
||||
std::vector<Packet> outputs;
|
||||
CalculatorGraphConfig graph_config_ =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: 'in'
|
||||
node {
|
||||
calculator: 'PreviousLoopbackCalculator'
|
||||
input_stream: 'MAIN:in'
|
||||
input_stream: 'LOOP:out'
|
||||
input_stream_info: { tag_index: 'LOOP' back_edge: true }
|
||||
output_stream: 'PREV_LOOP:previous'
|
||||
}
|
||||
# This calculator synchronizes its inputs as normal, so it is used
|
||||
# to check that both "in" and "previous" are ready.
|
||||
node {
|
||||
calculator: 'PassThroughCalculator'
|
||||
input_stream: 'in'
|
||||
input_stream: 'previous'
|
||||
output_stream: 'out'
|
||||
output_stream: 'previous2'
|
||||
}
|
||||
node {
|
||||
calculator: 'PacketOnCloseCalculator'
|
||||
input_stream: 'out'
|
||||
output_stream: 'close_out'
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("close_out", &graph_config_, &outputs);
|
||||
|
||||
CalculatorGraph graph_;
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config_, {}));
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
|
||||
auto send_packet = [&graph_](const std::string& input_name, int n) {
|
||||
MP_EXPECT_OK(graph_.AddPacketToInputStream(
|
||||
input_name, MakePacket<int>(n).At(Timestamp(n))));
|
||||
};
|
||||
|
||||
send_packet("in", 1);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1}));
|
||||
|
||||
send_packet("in", 2);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1, 2}));
|
||||
|
||||
send_packet("in", 5);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1, 2, 5}));
|
||||
|
||||
send_packet("in", 15);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1, 2, 5, 15}));
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs),
|
||||
(std::vector<int64>{1, 2, 5, 15, Timestamp::Max().Value()}));
|
||||
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
// Demonstrates that downstream calculators won't be blocked by
|
||||
// always-empty-LOOP-stream.
|
||||
TEST(PreviousLoopbackCalculator, EmptyLoopForever) {
|
||||
std::vector<Packet> outputs;
|
||||
CalculatorGraphConfig graph_config_ =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: 'in'
|
||||
node {
|
||||
calculator: 'PreviousLoopbackCalculator'
|
||||
input_stream: 'MAIN:in'
|
||||
input_stream: 'LOOP:previous'
|
||||
input_stream_info: { tag_index: 'LOOP' back_edge: true }
|
||||
output_stream: 'PREV_LOOP:previous'
|
||||
}
|
||||
# This calculator synchronizes its inputs as normal, so it is used
|
||||
# to check that both "in" and "previous" are ready.
|
||||
node {
|
||||
calculator: 'PassThroughCalculator'
|
||||
input_stream: 'in'
|
||||
input_stream: 'previous'
|
||||
output_stream: 'out'
|
||||
output_stream: 'previous2'
|
||||
}
|
||||
node {
|
||||
calculator: 'PacketOnCloseCalculator'
|
||||
input_stream: 'out'
|
||||
output_stream: 'close_out'
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("close_out", &graph_config_, &outputs);
|
||||
|
||||
CalculatorGraph graph_;
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config_, {}));
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
|
||||
auto send_packet = [&graph_](const std::string& input_name, int n) {
|
||||
MP_EXPECT_OK(graph_.AddPacketToInputStream(
|
||||
input_name, MakePacket<int>(n).At(Timestamp(n))));
|
||||
};
|
||||
|
||||
send_packet("in", 0);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{0}));
|
||||
|
||||
for (int main_ts = 1; main_ts < 50; ++main_ts) {
|
||||
send_packet("in", main_ts);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
std::vector<int64> ts_values = TimestampValues(outputs);
|
||||
EXPECT_EQ(ts_values.size(), main_ts + 1);
|
||||
for (int j = 0; j < main_ts; ++j) {
|
||||
EXPECT_EQ(ts_values[j], j);
|
||||
}
|
||||
}
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
// Copyright 2019 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 <map>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <string>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
using mediapipe::PacketTypeSet;
|
||||
using mediapipe::Timestamp;
|
||||
|
||||
namespace {
|
||||
static std::map<std::string, Timestamp>* kTimestampMap = []() {
|
||||
auto* res = new std::map<std::string, Timestamp>();
|
||||
res->emplace("AT_PRESTREAM", Timestamp::PreStream());
|
||||
res->emplace("AT_POSTSTREAM", Timestamp::PostStream());
|
||||
res->emplace("AT_ZERO", Timestamp(0));
|
||||
return res;
|
||||
}();
|
||||
|
||||
} // namespace
|
||||
|
||||
// Outputs the single input_side_packet at the timestamp specified in the
|
||||
// output_stream tag. Valid tags are AT_PRESTREAM, AT_POSTSTREAM and AT_ZERO.
|
||||
class SidePacketToStreamCalculator : public CalculatorBase {
|
||||
public:
|
||||
SidePacketToStreamCalculator() = default;
|
||||
~SidePacketToStreamCalculator() override = default;
|
||||
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
};
|
||||
REGISTER_CALCULATOR(SidePacketToStreamCalculator);
|
||||
|
||||
::mediapipe::Status SidePacketToStreamCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
cc->InputSidePackets().Index(0).SetAny();
|
||||
|
||||
std::set<std::string> tags = cc->Outputs().GetTags();
|
||||
RET_CHECK_EQ(tags.size(), 1);
|
||||
|
||||
RET_CHECK_EQ(kTimestampMap->count(*tags.begin()), 1);
|
||||
cc->Outputs().Tag(*tags.begin()).SetAny();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status SidePacketToStreamCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
return mediapipe::tool::StatusStop();
|
||||
}
|
||||
|
||||
::mediapipe::Status SidePacketToStreamCalculator::Close(CalculatorContext* cc) {
|
||||
std::set<std::string> tags = cc->Outputs().GetTags();
|
||||
RET_CHECK_EQ(tags.size(), 1);
|
||||
const std::string& tag = *tags.begin();
|
||||
RET_CHECK_EQ(kTimestampMap->count(tag), 1);
|
||||
cc->Outputs().Tag(tag).AddPacket(
|
||||
cc->InputSidePackets().Index(0).At(kTimestampMap->at(tag)));
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -17,8 +17,13 @@
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Example config:
|
||||
@@ -35,10 +40,21 @@ namespace mediapipe {
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
typedef SplitVectorCalculator<TfLiteTensor> SplitTfLiteTensorVectorCalculator;
|
||||
typedef SplitVectorCalculator<TfLiteTensor, false>
|
||||
SplitTfLiteTensorVectorCalculator;
|
||||
REGISTER_CALCULATOR(SplitTfLiteTensorVectorCalculator);
|
||||
|
||||
typedef SplitVectorCalculator<::mediapipe::NormalizedLandmark>
|
||||
typedef SplitVectorCalculator<::mediapipe::NormalizedLandmark, false>
|
||||
SplitLandmarkVectorCalculator;
|
||||
REGISTER_CALCULATOR(SplitLandmarkVectorCalculator);
|
||||
|
||||
typedef SplitVectorCalculator<::mediapipe::NormalizedRect, false>
|
||||
SplitNormalizedRectVectorCalculator;
|
||||
REGISTER_CALCULATOR(SplitNormalizedRectVectorCalculator);
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
typedef SplitVectorCalculator<::tflite::gpu::gl::GlBuffer, true>
|
||||
MovableSplitGlBufferVectorCalculator;
|
||||
REGISTER_CALCULATOR(MovableSplitGlBufferVectorCalculator);
|
||||
#endif
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -15,12 +15,14 @@
|
||||
#ifndef MEDIAPIPE_CALCULATORS_CORE_SPLIT_VECTOR_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_CORE_SPLIT_VECTOR_CALCULATOR_H_
|
||||
|
||||
#include <type_traits>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/calculators/core/split_vector_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
#include "tensorflow/lite/error_reporter.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
@@ -29,15 +31,31 @@
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
template <typename T>
|
||||
using IsCopyable = std::enable_if_t<std::is_copy_constructible<T>::value, bool>;
|
||||
|
||||
template <typename T>
|
||||
using IsNotCopyable =
|
||||
std::enable_if_t<!std::is_copy_constructible<T>::value, bool>;
|
||||
|
||||
template <typename T>
|
||||
using IsMovable = std::enable_if_t<std::is_move_constructible<T>::value, bool>;
|
||||
|
||||
template <typename T>
|
||||
using IsNotMovable =
|
||||
std::enable_if_t<!std::is_move_constructible<T>::value, bool>;
|
||||
|
||||
// Splits an input packet with std::vector<T> into multiple std::vector<T>
|
||||
// output packets using the [begin, end) ranges specified in
|
||||
// SplitVectorCalculatorOptions. If the option "element_only" is set to true,
|
||||
// all ranges should be of size 1 and all outputs will be elements of type T. If
|
||||
// "element_only" is false, ranges can be non-zero in size and all outputs will
|
||||
// be of type std::vector<T>.
|
||||
// be of type std::vector<T>. If the option "combine_outputs" is set to true,
|
||||
// only one output stream can be specified and all ranges of elements will be
|
||||
// combined into one vector.
|
||||
// To use this class for a particular type T, register a calculator using
|
||||
// SplitVectorCalculator<T>.
|
||||
template <typename T>
|
||||
template <typename T, bool move_elements>
|
||||
class SplitVectorCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
@@ -49,28 +67,40 @@ class SplitVectorCalculator : public CalculatorBase {
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
|
||||
|
||||
if (cc->Outputs().NumEntries() != options.ranges_size()) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"The number of output streams should match the number of ranges "
|
||||
"specified in the CalculatorOptions.");
|
||||
if (!std::is_copy_constructible<T>::value || move_elements) {
|
||||
// Ranges of elements shouldn't overlap when the vector contains
|
||||
// non-copyable elements.
|
||||
RET_CHECK_OK(checkRangesDontOverlap(options));
|
||||
}
|
||||
|
||||
// Set the output types for each output stream.
|
||||
for (int i = 0; i < cc->Outputs().NumEntries(); ++i) {
|
||||
if (options.ranges(i).begin() < 0 || options.ranges(i).end() < 0 ||
|
||||
options.ranges(i).begin() >= options.ranges(i).end()) {
|
||||
if (options.combine_outputs()) {
|
||||
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1);
|
||||
cc->Outputs().Index(0).Set<std::vector<T>>();
|
||||
RET_CHECK_OK(checkRangesDontOverlap(options));
|
||||
} else {
|
||||
if (cc->Outputs().NumEntries() != options.ranges_size()) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Indices should be non-negative and begin index should be less "
|
||||
"than the end index.");
|
||||
"The number of output streams should match the number of ranges "
|
||||
"specified in the CalculatorOptions.");
|
||||
}
|
||||
if (options.element_only()) {
|
||||
if (options.ranges(i).end() - options.ranges(i).begin() != 1) {
|
||||
|
||||
// Set the output types for each output stream.
|
||||
for (int i = 0; i < cc->Outputs().NumEntries(); ++i) {
|
||||
if (options.ranges(i).begin() < 0 || options.ranges(i).end() < 0 ||
|
||||
options.ranges(i).begin() >= options.ranges(i).end()) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Since element_only is true, all ranges should be of size 1.");
|
||||
"Indices should be non-negative and begin index should be less "
|
||||
"than the end index.");
|
||||
}
|
||||
if (options.element_only()) {
|
||||
if (options.ranges(i).end() - options.ranges(i).begin() != 1) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Since element_only is true, all ranges should be of size 1.");
|
||||
}
|
||||
cc->Outputs().Index(i).Set<T>();
|
||||
} else {
|
||||
cc->Outputs().Index(i).Set<std::vector<T>>();
|
||||
}
|
||||
cc->Outputs().Index(i).Set<T>();
|
||||
} else {
|
||||
cc->Outputs().Index(i).Set<std::vector<T>>();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -83,41 +113,140 @@ class SplitVectorCalculator : public CalculatorBase {
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
|
||||
|
||||
element_only_ = options.element_only();
|
||||
combine_outputs_ = options.combine_outputs();
|
||||
|
||||
for (const auto& range : options.ranges()) {
|
||||
ranges_.push_back({range.begin(), range.end()});
|
||||
max_range_end_ = std::max(max_range_end_, range.end());
|
||||
total_elements_ += range.end() - range.begin();
|
||||
}
|
||||
|
||||
element_only_ = options.element_only();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
const auto& input = cc->Inputs().Index(0).Get<std::vector<T>>();
|
||||
RET_CHECK_GE(input.size(), max_range_end_);
|
||||
if (cc->Inputs().Index(0).IsEmpty()) return ::mediapipe::OkStatus();
|
||||
|
||||
if (element_only_) {
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
cc->Outputs().Index(i).AddPacket(
|
||||
MakePacket<T>(input[ranges_[i].first]).At(cc->InputTimestamp()));
|
||||
}
|
||||
if (move_elements) {
|
||||
return ProcessMovableElements<T>(cc);
|
||||
} else {
|
||||
return ProcessCopyableElements<T>(cc);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename U, IsCopyable<U> = true>
|
||||
::mediapipe::Status ProcessCopyableElements(CalculatorContext* cc) {
|
||||
// static_assert(std::is_copy_constructible<U>::value,
|
||||
// "Cannot copy non-copyable elements");
|
||||
const auto& input = cc->Inputs().Index(0).Get<std::vector<U>>();
|
||||
RET_CHECK_GE(input.size(), max_range_end_);
|
||||
if (combine_outputs_) {
|
||||
auto output = absl::make_unique<std::vector<U>>();
|
||||
output->reserve(total_elements_);
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
auto output = absl::make_unique<std::vector<T>>(
|
||||
auto elements = absl::make_unique<std::vector<U>>(
|
||||
input.begin() + ranges_[i].first,
|
||||
input.begin() + ranges_[i].second);
|
||||
cc->Outputs().Index(i).Add(output.release(), cc->InputTimestamp());
|
||||
output->insert(output->end(), elements->begin(), elements->end());
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||
} else {
|
||||
if (element_only_) {
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
cc->Outputs().Index(i).AddPacket(
|
||||
MakePacket<U>(input[ranges_[i].first]).At(cc->InputTimestamp()));
|
||||
}
|
||||
} else {
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
auto output = absl::make_unique<std::vector<T>>(
|
||||
input.begin() + ranges_[i].first,
|
||||
input.begin() + ranges_[i].second);
|
||||
cc->Outputs().Index(i).Add(output.release(), cc->InputTimestamp());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
template <typename U, IsNotCopyable<U> = true>
|
||||
::mediapipe::Status ProcessCopyableElements(CalculatorContext* cc) {
|
||||
return ::mediapipe::InternalError("Cannot copy non-copyable elements.");
|
||||
}
|
||||
|
||||
template <typename U, IsMovable<U> = true>
|
||||
::mediapipe::Status ProcessMovableElements(CalculatorContext* cc) {
|
||||
::mediapipe::StatusOr<std::unique_ptr<std::vector<U>>> input_status =
|
||||
cc->Inputs().Index(0).Value().Consume<std::vector<U>>();
|
||||
if (!input_status.ok()) return input_status.status();
|
||||
std::unique_ptr<std::vector<U>> input_vector =
|
||||
std::move(input_status).ValueOrDie();
|
||||
RET_CHECK_GE(input_vector->size(), max_range_end_);
|
||||
|
||||
if (combine_outputs_) {
|
||||
auto output = absl::make_unique<std::vector<U>>();
|
||||
output->reserve(total_elements_);
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
output->insert(
|
||||
output->end(),
|
||||
std::make_move_iterator(input_vector->begin() + ranges_[i].first),
|
||||
std::make_move_iterator(input_vector->begin() + ranges_[i].second));
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||
} else {
|
||||
if (element_only_) {
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
cc->Outputs().Index(i).AddPacket(
|
||||
MakePacket<U>(std::move(input_vector->at(ranges_[i].first)))
|
||||
.At(cc->InputTimestamp()));
|
||||
}
|
||||
} else {
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
auto output = absl::make_unique<std::vector<T>>();
|
||||
output->insert(
|
||||
output->end(),
|
||||
std::make_move_iterator(input_vector->begin() + ranges_[i].first),
|
||||
std::make_move_iterator(input_vector->begin() +
|
||||
ranges_[i].second));
|
||||
cc->Outputs().Index(i).Add(output.release(), cc->InputTimestamp());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
template <typename U, IsNotMovable<U> = true>
|
||||
::mediapipe::Status ProcessMovableElements(CalculatorContext* cc) {
|
||||
return ::mediapipe::InternalError("Cannot move non-movable elements.");
|
||||
}
|
||||
|
||||
private:
|
||||
static ::mediapipe::Status checkRangesDontOverlap(
|
||||
const ::mediapipe::SplitVectorCalculatorOptions& options) {
|
||||
for (int i = 0; i < options.ranges_size() - 1; ++i) {
|
||||
for (int j = i + 1; j < options.ranges_size(); ++j) {
|
||||
const auto& range_0 = options.ranges(i);
|
||||
const auto& range_1 = options.ranges(j);
|
||||
if ((range_0.begin() >= range_1.begin() &&
|
||||
range_0.begin() < range_1.end()) ||
|
||||
(range_1.begin() >= range_0.begin() &&
|
||||
range_1.begin() < range_0.end())) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Ranges must be non-overlapping when using combine_outputs "
|
||||
"option.");
|
||||
}
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
std::vector<std::pair<int32, int32>> ranges_;
|
||||
int32 max_range_end_ = -1;
|
||||
int32 total_elements_ = 0;
|
||||
bool element_only_ = false;
|
||||
bool combine_outputs_ = false;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -37,4 +37,7 @@ message SplitVectorCalculatorOptions {
|
||||
// just element of type T. By default, if a range specifies only one element,
|
||||
// it is outputted as an std::vector<T>.
|
||||
optional bool element_only = 2 [default = false];
|
||||
|
||||
// Combines output elements to one vector.
|
||||
optional bool combine_outputs = 3 [default = false];
|
||||
}
|
||||
|
||||
@@ -105,6 +105,34 @@ class SplitTfLiteTensorVectorCalculatorTest : public ::testing::Test {
|
||||
}
|
||||
}
|
||||
|
||||
void ValidateCombinedVectorOutput(std::vector<Packet>& output_packets,
|
||||
int expected_elements,
|
||||
std::vector<int>& input_begin_indices,
|
||||
std::vector<int>& input_end_indices) {
|
||||
ASSERT_EQ(1, output_packets.size());
|
||||
ASSERT_EQ(input_begin_indices.size(), input_end_indices.size());
|
||||
const std::vector<TfLiteTensor>& output_vec =
|
||||
output_packets[0].Get<std::vector<TfLiteTensor>>();
|
||||
ASSERT_EQ(expected_elements, output_vec.size());
|
||||
const int num_ranges = input_begin_indices.size();
|
||||
|
||||
int element_id = 0;
|
||||
for (int range_id = 0; range_id < num_ranges; ++range_id) {
|
||||
for (int i = input_begin_indices[range_id];
|
||||
i < input_end_indices[range_id]; ++i) {
|
||||
const int expected_value = i;
|
||||
const TfLiteTensor* result = &output_vec[element_id];
|
||||
float* result_buffer = result->data.f;
|
||||
ASSERT_NE(result_buffer, nullptr);
|
||||
ASSERT_EQ(result_buffer, input_buffers_[i]);
|
||||
for (int j = 0; j < width * height * channels; ++j) {
|
||||
ASSERT_EQ(expected_value, result_buffer[j]);
|
||||
}
|
||||
element_id++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void ValidateElementOutput(std::vector<Packet>& output_packets,
|
||||
int input_begin_index) {
|
||||
ASSERT_EQ(1, output_packets.size());
|
||||
@@ -234,6 +262,65 @@ TEST_F(SplitTfLiteTensorVectorCalculatorTest, InvalidOutputStreamCountTest) {
|
||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||
}
|
||||
|
||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest,
|
||||
InvalidCombineOutputsMultipleOutputsTest) {
|
||||
ASSERT_NE(interpreter_, nullptr);
|
||||
|
||||
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "tensor_in"
|
||||
node {
|
||||
calculator: "SplitTfLiteTensorVectorCalculator"
|
||||
input_stream: "tensor_in"
|
||||
output_stream: "range_0"
|
||||
output_stream: "range_1"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 1 }
|
||||
ranges: { begin: 2 end: 3 }
|
||||
combine_outputs: true
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
// The graph should fail running because the number of output streams does not
|
||||
// match the number of range elements in the options.
|
||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||
}
|
||||
|
||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest, InvalidOverlappingRangesTest) {
|
||||
ASSERT_NE(interpreter_, nullptr);
|
||||
|
||||
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "tensor_in"
|
||||
node {
|
||||
calculator: "SplitTfLiteTensorVectorCalculator"
|
||||
input_stream: "tensor_in"
|
||||
output_stream: "range_0"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 3 }
|
||||
ranges: { begin: 1 end: 4 }
|
||||
combine_outputs: true
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
// The graph should fail running because there are overlapping ranges.
|
||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||
}
|
||||
|
||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest, SmokeTestElementOnly) {
|
||||
ASSERT_NE(interpreter_, nullptr);
|
||||
|
||||
@@ -289,6 +376,53 @@ TEST_F(SplitTfLiteTensorVectorCalculatorTest, SmokeTestElementOnly) {
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest, SmokeTestCombiningOutputs) {
|
||||
ASSERT_NE(interpreter_, nullptr);
|
||||
|
||||
PrepareTfLiteTensorVector(/*vector_size=*/5);
|
||||
ASSERT_NE(input_vec_, nullptr);
|
||||
|
||||
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "tensor_in"
|
||||
node {
|
||||
calculator: "SplitTfLiteTensorVectorCalculator"
|
||||
input_stream: "tensor_in"
|
||||
output_stream: "range_0"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 1 }
|
||||
ranges: { begin: 2 end: 3 }
|
||||
ranges: { begin: 4 end: 5 }
|
||||
combine_outputs: true
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
std::vector<Packet> range_0_packets;
|
||||
tool::AddVectorSink("range_0", &graph_config, &range_0_packets);
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"tensor_in", Adopt(input_vec_.release()).At(Timestamp(0))));
|
||||
// Wait until the calculator finishes processing.
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
std::vector<int> input_begin_indices = {0, 2, 4};
|
||||
std::vector<int> input_end_indices = {1, 3, 5};
|
||||
ValidateCombinedVectorOutput(range_0_packets, /*expected_elements=*/3,
|
||||
input_begin_indices, input_end_indices);
|
||||
|
||||
// Fully close the graph at the end.
|
||||
MP_ASSERT_OK(graph.CloseInputStream("tensor_in"));
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest,
|
||||
ElementOnlyDisablesVectorOutputs) {
|
||||
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||
@@ -318,4 +452,243 @@ TEST_F(SplitTfLiteTensorVectorCalculatorTest,
|
||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||
}
|
||||
|
||||
typedef SplitVectorCalculator<std::unique_ptr<int>, true>
|
||||
MovableSplitUniqueIntPtrCalculator;
|
||||
REGISTER_CALCULATOR(MovableSplitUniqueIntPtrCalculator);
|
||||
|
||||
class MovableSplitUniqueIntPtrCalculatorTest : public ::testing::Test {
|
||||
protected:
|
||||
void ValidateVectorOutput(std::vector<Packet>& output_packets,
|
||||
int expected_elements, int input_begin_index) {
|
||||
ASSERT_EQ(1, output_packets.size());
|
||||
const std::vector<std::unique_ptr<int>>& output_vec =
|
||||
output_packets[0].Get<std::vector<std::unique_ptr<int>>>();
|
||||
ASSERT_EQ(expected_elements, output_vec.size());
|
||||
|
||||
for (int i = 0; i < expected_elements; ++i) {
|
||||
const int expected_value = input_begin_index + i;
|
||||
const std::unique_ptr<int>& result = output_vec[i];
|
||||
ASSERT_NE(result, nullptr);
|
||||
ASSERT_EQ(expected_value, *result);
|
||||
}
|
||||
}
|
||||
|
||||
void ValidateElementOutput(std::vector<Packet>& output_packets,
|
||||
int expected_value) {
|
||||
ASSERT_EQ(1, output_packets.size());
|
||||
const std::unique_ptr<int>& result =
|
||||
output_packets[0].Get<std::unique_ptr<int>>();
|
||||
ASSERT_NE(result, nullptr);
|
||||
ASSERT_EQ(expected_value, *result);
|
||||
}
|
||||
|
||||
void ValidateCombinedVectorOutput(std::vector<Packet>& output_packets,
|
||||
int expected_elements,
|
||||
std::vector<int>& input_begin_indices,
|
||||
std::vector<int>& input_end_indices) {
|
||||
ASSERT_EQ(1, output_packets.size());
|
||||
ASSERT_EQ(input_begin_indices.size(), input_end_indices.size());
|
||||
const std::vector<std::unique_ptr<int>>& output_vector =
|
||||
output_packets[0].Get<std::vector<std::unique_ptr<int>>>();
|
||||
ASSERT_EQ(expected_elements, output_vector.size());
|
||||
const int num_ranges = input_begin_indices.size();
|
||||
|
||||
int element_id = 0;
|
||||
for (int range_id = 0; range_id < num_ranges; ++range_id) {
|
||||
for (int i = input_begin_indices[range_id];
|
||||
i < input_end_indices[range_id]; ++i) {
|
||||
const int expected_value = i;
|
||||
const std::unique_ptr<int>& result = output_vector[element_id];
|
||||
ASSERT_NE(result, nullptr);
|
||||
ASSERT_EQ(expected_value, *result);
|
||||
++element_id;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, InvalidOverlappingRangesTest) {
|
||||
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "input_vector"
|
||||
node {
|
||||
calculator: "MovableSplitUniqueIntPtrCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "range_0"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 3 }
|
||||
ranges: { begin: 1 end: 4 }
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
// The graph should fail running because there are overlapping ranges.
|
||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||
}
|
||||
|
||||
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, SmokeTest) {
|
||||
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "input_vector"
|
||||
node {
|
||||
calculator: "MovableSplitUniqueIntPtrCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "range_0"
|
||||
output_stream: "range_1"
|
||||
output_stream: "range_2"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 1 }
|
||||
ranges: { begin: 1 end: 4 }
|
||||
ranges: { begin: 4 end: 5 }
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
std::vector<Packet> range_0_packets;
|
||||
tool::AddVectorSink("range_0", &graph_config, &range_0_packets);
|
||||
std::vector<Packet> range_1_packets;
|
||||
tool::AddVectorSink("range_1", &graph_config, &range_1_packets);
|
||||
std::vector<Packet> range_2_packets;
|
||||
tool::AddVectorSink("range_2", &graph_config, &range_2_packets);
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
// input_vector : {0, 1, 2, 3, 4, 5}
|
||||
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_vector =
|
||||
absl::make_unique<std::vector<std::unique_ptr<int>>>(6);
|
||||
for (int i = 0; i < 6; ++i) {
|
||||
input_vector->at(i) = absl::make_unique<int>(i);
|
||||
}
|
||||
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_vector", Adopt(input_vector.release()).At(Timestamp(1))));
|
||||
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
|
||||
ValidateVectorOutput(range_0_packets, /*expected_elements=*/1,
|
||||
/*input_begin_index=*/0);
|
||||
ValidateVectorOutput(range_1_packets, /*expected_elements=*/3,
|
||||
/*input_begin_index=*/1);
|
||||
ValidateVectorOutput(range_2_packets, /*expected_elements=*/1,
|
||||
/*input_begin_index=*/4);
|
||||
}
|
||||
|
||||
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, SmokeTestElementOnly) {
|
||||
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "input_vector"
|
||||
node {
|
||||
calculator: "MovableSplitUniqueIntPtrCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "range_0"
|
||||
output_stream: "range_1"
|
||||
output_stream: "range_2"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 1 }
|
||||
ranges: { begin: 2 end: 3 }
|
||||
ranges: { begin: 4 end: 5 }
|
||||
element_only: true
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
std::vector<Packet> range_0_packets;
|
||||
tool::AddVectorSink("range_0", &graph_config, &range_0_packets);
|
||||
std::vector<Packet> range_1_packets;
|
||||
tool::AddVectorSink("range_1", &graph_config, &range_1_packets);
|
||||
std::vector<Packet> range_2_packets;
|
||||
tool::AddVectorSink("range_2", &graph_config, &range_2_packets);
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
// input_vector : {0, 1, 2, 3, 4, 5}
|
||||
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_vector =
|
||||
absl::make_unique<std::vector<std::unique_ptr<int>>>(6);
|
||||
for (int i = 0; i < 6; ++i) {
|
||||
input_vector->at(i) = absl::make_unique<int>(i);
|
||||
}
|
||||
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_vector", Adopt(input_vector.release()).At(Timestamp(1))));
|
||||
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
|
||||
ValidateElementOutput(range_0_packets, /*expected_value=*/0);
|
||||
ValidateElementOutput(range_1_packets, /*expected_value=*/2);
|
||||
ValidateElementOutput(range_2_packets, /*expected_value=*/4);
|
||||
}
|
||||
|
||||
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, SmokeTestCombiningOutputs) {
|
||||
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "input_vector"
|
||||
node {
|
||||
calculator: "MovableSplitUniqueIntPtrCalculator"
|
||||
input_stream: "input_vector"
|
||||
output_stream: "range_0"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 1 }
|
||||
ranges: { begin: 2 end: 3 }
|
||||
ranges: { begin: 4 end: 5 }
|
||||
combine_outputs: true
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
std::vector<Packet> range_0_packets;
|
||||
tool::AddVectorSink("range_0", &graph_config, &range_0_packets);
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
// input_vector : {0, 1, 2, 3, 4, 5}
|
||||
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_vector =
|
||||
absl::make_unique<std::vector<std::unique_ptr<int>>>(6);
|
||||
for (int i = 0; i < 6; ++i) {
|
||||
input_vector->at(i) = absl::make_unique<int>(i);
|
||||
}
|
||||
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_vector", Adopt(input_vector.release()).At(Timestamp(1))));
|
||||
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
|
||||
std::vector<int> input_begin_indices = {0, 2, 4};
|
||||
std::vector<int> input_end_indices = {1, 3, 5};
|
||||
ValidateCombinedVectorOutput(range_0_packets, /*expected_elements=*/3,
|
||||
input_begin_indices, input_end_indices);
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
// Copyright 2019 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/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator that takes a packet of an input stream and converts it to an
|
||||
// output side packet. This calculator only works under the assumption that the
|
||||
// input stream only has a single packet passing through.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "StreamToSidePacketCalculator"
|
||||
// input_stream: "stream"
|
||||
// output_side_packet: "side_packet"
|
||||
// }
|
||||
class StreamToSidePacketCalculator : public mediapipe::CalculatorBase {
|
||||
public:
|
||||
static mediapipe::Status GetContract(mediapipe::CalculatorContract* cc) {
|
||||
cc->Inputs().Index(0).SetAny();
|
||||
cc->OutputSidePackets().Index(0).SetAny();
|
||||
return mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
mediapipe::Status Process(mediapipe::CalculatorContext* cc) override {
|
||||
mediapipe::Packet& packet = cc->Inputs().Index(0).Value();
|
||||
cc->OutputSidePackets().Index(0).Set(
|
||||
packet.At(mediapipe::Timestamp::Unset()));
|
||||
return mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
REGISTER_CALCULATOR(StreamToSidePacketCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,67 @@
|
||||
// Copyright 2019 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 "absl/memory/memory.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
using ::testing::Test;
|
||||
|
||||
class StreamToSidePacketCalculatorTest : public Test {
|
||||
protected:
|
||||
StreamToSidePacketCalculatorTest() {
|
||||
const char kConfig[] = R"(
|
||||
calculator: "StreamToSidePacketCalculator"
|
||||
input_stream: "stream"
|
||||
output_side_packet: "side_packet"
|
||||
)";
|
||||
runner_ = absl::make_unique<CalculatorRunner>(kConfig);
|
||||
}
|
||||
|
||||
std::unique_ptr<CalculatorRunner> runner_;
|
||||
};
|
||||
|
||||
TEST_F(StreamToSidePacketCalculatorTest,
|
||||
StreamToSidePacketCalculatorWithEmptyStreamFails) {
|
||||
EXPECT_EQ(runner_->Run().code(), mediapipe::StatusCode::kUnavailable);
|
||||
}
|
||||
|
||||
TEST_F(StreamToSidePacketCalculatorTest,
|
||||
StreamToSidePacketCalculatorWithSinglePacketCreatesSidePacket) {
|
||||
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(new std::string("test")).At(Timestamp(1)));
|
||||
MP_ASSERT_OK(runner_->Run());
|
||||
EXPECT_EQ(runner_->OutputSidePackets().Index(0).Get<std::string>(), "test");
|
||||
}
|
||||
|
||||
TEST_F(StreamToSidePacketCalculatorTest,
|
||||
StreamToSidePacketCalculatorWithMultiplePacketsFails) {
|
||||
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(new std::string("test1")).At(Timestamp(1)));
|
||||
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(new std::string("test2")).At(Timestamp(2)));
|
||||
EXPECT_EQ(runner_->Run().code(), mediapipe::StatusCode::kAlreadyExists);
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,79 @@
|
||||
// Copyright 2019 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 <sys/types.h>
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#include "absl/strings/numbers.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Calculator that converts a std::string into an integer type, or fails if the
|
||||
// conversion is not possible.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "StringToIntCalculator"
|
||||
// input_side_packet: "string"
|
||||
// output_side_packet: "index"
|
||||
// }
|
||||
template <typename IntType>
|
||||
class StringToIntCalculatorTemplate : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->InputSidePackets().Index(0).Set<std::string>();
|
||||
cc->OutputSidePackets().Index(0).Set<IntType>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
IntType number;
|
||||
if (!absl::SimpleAtoi(cc->InputSidePackets().Index(0).Get<std::string>(),
|
||||
&number)) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"The std::string could not be parsed as an integer.");
|
||||
}
|
||||
cc->OutputSidePackets().Index(0).Set(MakePacket<IntType>(number));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
using StringToIntCalculator = StringToIntCalculatorTemplate<int>;
|
||||
REGISTER_CALCULATOR(StringToIntCalculator);
|
||||
|
||||
using StringToUintCalculator = StringToIntCalculatorTemplate<uint>;
|
||||
REGISTER_CALCULATOR(StringToUintCalculator);
|
||||
|
||||
using StringToInt32Calculator = StringToIntCalculatorTemplate<int32>;
|
||||
REGISTER_CALCULATOR(StringToInt32Calculator);
|
||||
|
||||
using StringToUint32Calculator = StringToIntCalculatorTemplate<uint32>;
|
||||
REGISTER_CALCULATOR(StringToUint32Calculator);
|
||||
|
||||
using StringToInt64Calculator = StringToIntCalculatorTemplate<int64>;
|
||||
REGISTER_CALCULATOR(StringToInt64Calculator);
|
||||
|
||||
using StringToUint64Calculator = StringToIntCalculatorTemplate<uint64>;
|
||||
REGISTER_CALCULATOR(StringToUint64Calculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -12,14 +12,14 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
|
||||
exports_files(["LICENSE"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "opencv_image_encoder_calculator_proto",
|
||||
srcs = ["opencv_image_encoder_calculator.proto"],
|
||||
@@ -80,7 +80,9 @@ mediapipe_cc_proto_library(
|
||||
name = "opencv_image_encoder_calculator_cc_proto",
|
||||
srcs = ["opencv_image_encoder_calculator.proto"],
|
||||
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
visibility = [
|
||||
"//visibility:public",
|
||||
],
|
||||
deps = [":opencv_image_encoder_calculator_proto"],
|
||||
)
|
||||
|
||||
@@ -356,13 +358,13 @@ cc_library(
|
||||
"//mediapipe/framework/port:opencv_imgproc",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/gpu:gpu_buffer",
|
||||
] + select({
|
||||
"//mediapipe/gpu:disable_gpu": [],
|
||||
"//conditions:default": [
|
||||
"//mediapipe/gpu:gl_calculator_helper",
|
||||
"//mediapipe/gpu:gl_simple_shaders",
|
||||
"//mediapipe/gpu:gl_quad_renderer",
|
||||
"//mediapipe/gpu:gpu_buffer",
|
||||
"//mediapipe/gpu:shader_util",
|
||||
],
|
||||
}),
|
||||
|
||||
@@ -501,8 +501,11 @@ void ImageCroppingCalculator::GetOutputDimensions(CalculatorContext* cc,
|
||||
row_max = std::max(row_max, transformed_points_[i * 2 + 1]);
|
||||
}
|
||||
|
||||
*dst_width = std::round((col_max - col_min) * src_width);
|
||||
*dst_height = std::round((row_max - row_min) * src_height);
|
||||
int width = static_cast<int>(std::round((col_max - col_min) * src_width));
|
||||
int height = static_cast<int>(std::round((row_max - row_min) * src_height));
|
||||
// Minimum output dimension 1x1 prevents creation of textures with 0x0.
|
||||
*dst_width = std::max(1, width);
|
||||
*dst_height = std::max(1, height);
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -400,7 +400,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
QuadRenderer* renderer = nullptr;
|
||||
GlTexture src1;
|
||||
|
||||
#if defined(__APPLE__) && !TARGET_OS_OSX
|
||||
#if defined(MEDIAPIPE_IOS)
|
||||
if (input.format() == GpuBufferFormat::kBiPlanar420YpCbCr8VideoRange ||
|
||||
input.format() == GpuBufferFormat::kBiPlanar420YpCbCr8FullRange) {
|
||||
if (!yuv_renderer_) {
|
||||
|
||||
@@ -474,13 +474,20 @@ ScaleImageCalculator::~ScaleImageCalculator() {}
|
||||
input_width_, "x", input_height_));
|
||||
}
|
||||
if (input_format_ != image_frame.Format()) {
|
||||
std::string image_frame_format_desc, input_format_desc;
|
||||
#ifdef MEDIAPIPE_MOBILE
|
||||
image_frame_format_desc = std::to_string(image_frame.Format());
|
||||
input_format_desc = std::to_string(input_format_);
|
||||
#else
|
||||
const proto_ns::EnumDescriptor* desc = ImageFormat::Format_descriptor();
|
||||
image_frame_format_desc =
|
||||
desc->FindValueByNumber(image_frame.Format())->DebugString();
|
||||
input_format_desc = desc->FindValueByNumber(input_format_)->DebugString();
|
||||
#endif // MEDIAPIPE_MOBILE
|
||||
return tool::StatusFail(absl::StrCat(
|
||||
"If a header specifies a format, then image frames on "
|
||||
"the stream must have that format. Actual format ",
|
||||
desc->FindValueByNumber(image_frame.Format())->DebugString(),
|
||||
" but expected ",
|
||||
desc->FindValueByNumber(input_format_)->DebugString()));
|
||||
image_frame_format_desc, " but expected ", input_format_desc));
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
|
||||
@@ -36,7 +36,8 @@ message ScaleImageCalculatorOptions {
|
||||
|
||||
// If ratio is positive, crop the image to this minimum and maximum
|
||||
// aspect ratio (preserving the center of the frame). This is done
|
||||
// before scaling.
|
||||
// before scaling. The string must contain "/", so to disable cropping,
|
||||
// set both to "0/1".
|
||||
// For example, for a min_aspect_ratio of "9/16" and max of "16/9" the
|
||||
// following cropping will occur:
|
||||
// 1920x1080 (which is 16:9) is not cropped
|
||||
|
||||
@@ -13,12 +13,12 @@
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "graph_tensors_packet_generator_proto",
|
||||
srcs = ["graph_tensors_packet_generator.proto"],
|
||||
@@ -104,6 +104,17 @@ proto_library(
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "unpack_media_sequence_calculator_proto",
|
||||
srcs = ["unpack_media_sequence_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/core:packet_resampler_calculator_proto",
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/util:audio_decoder_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "vector_float_to_tensor_calculator_options_proto",
|
||||
srcs = ["vector_float_to_tensor_calculator_options.proto"],
|
||||
@@ -127,7 +138,7 @@ mediapipe_cc_proto_library(
|
||||
srcs = ["image_frame_to_tensor_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":image_frame_to_tensor_calculator_proto"],
|
||||
@@ -162,7 +173,7 @@ mediapipe_cc_proto_library(
|
||||
srcs = ["pack_media_sequence_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":pack_media_sequence_calculator_proto"],
|
||||
@@ -181,7 +192,7 @@ mediapipe_cc_proto_library(
|
||||
srcs = ["tensorflow_session_from_frozen_graph_generator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:packet_generator_cc_proto",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":tensorflow_session_from_frozen_graph_generator_proto"],
|
||||
@@ -192,7 +203,7 @@ mediapipe_cc_proto_library(
|
||||
srcs = ["tensorflow_session_from_frozen_graph_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":tensorflow_session_from_frozen_graph_calculator_proto"],
|
||||
@@ -261,6 +272,17 @@ mediapipe_cc_proto_library(
|
||||
deps = [":unpack_media_sequence_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "vector_int_to_tensor_calculator_options_cc_proto",
|
||||
srcs = ["vector_int_to_tensor_calculator_options.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":vector_int_to_tensor_calculator_options_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "vector_float_to_tensor_calculator_options_cc_proto",
|
||||
srcs = ["vector_float_to_tensor_calculator_options.proto"],
|
||||
@@ -274,7 +296,7 @@ cc_library(
|
||||
srcs = ["graph_tensors_packet_generator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:graph_tensors_packet_generator_cc_proto",
|
||||
":graph_tensors_packet_generator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -289,7 +311,7 @@ cc_library(
|
||||
srcs = ["image_frame_to_tensor_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:image_frame_to_tensor_calculator_cc_proto",
|
||||
":image_frame_to_tensor_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
@@ -311,7 +333,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
||||
"//mediapipe/calculators/tensorflow:matrix_to_tensor_calculator_options_cc_proto",
|
||||
":matrix_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -332,7 +354,7 @@ cc_library(
|
||||
srcs = ["lapped_tensor_buffer_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:lapped_tensor_buffer_calculator_cc_proto",
|
||||
":lapped_tensor_buffer_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -386,7 +408,7 @@ cc_library(
|
||||
"//mediapipe/util/sequence:media_sequence",
|
||||
"//mediapipe/util/sequence:media_sequence_util",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -401,7 +423,7 @@ cc_library(
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -414,7 +436,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":tensorflow_session",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_inference_calculator_cc_proto",
|
||||
":tensorflow_inference_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/tool:status_util",
|
||||
"@com_google_absl//absl/strings",
|
||||
@@ -492,7 +514,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":tensorflow_session",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
":tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/tool:status_util",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -551,7 +573,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":tensorflow_session",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_saved_model_generator_cc_proto",
|
||||
":tensorflow_session_from_saved_model_generator_cc_proto",
|
||||
"//mediapipe/framework:packet_generator",
|
||||
"//mediapipe/framework:packet_type",
|
||||
"//mediapipe/framework/tool:status_util",
|
||||
@@ -575,7 +597,7 @@ cc_library(
|
||||
srcs = ["tensor_squeeze_dimensions_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:tensor_squeeze_dimensions_calculator_cc_proto",
|
||||
":tensor_squeeze_dimensions_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -589,7 +611,7 @@ cc_library(
|
||||
srcs = ["tensor_to_image_frame_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_image_frame_calculator_cc_proto",
|
||||
":tensor_to_image_frame_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
@@ -605,7 +627,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_matrix_calculator_cc_proto",
|
||||
":tensor_to_matrix_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -621,6 +643,22 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "tfrecord_reader_calculator",
|
||||
srcs = ["tfrecord_reader_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@org_tensorflow//tensorflow/core:lib",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "tensor_to_vector_float_calculator",
|
||||
srcs = ["tensor_to_vector_float_calculator.cc"],
|
||||
@@ -629,7 +667,7 @@ cc_library(
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_vector_float_calculator_options_cc_proto",
|
||||
":tensor_to_vector_float_calculator_options_cc_proto",
|
||||
] + select({
|
||||
"//conditions:default": [
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
@@ -657,7 +695,21 @@ cc_library(
|
||||
"//mediapipe/util:audio_decoder_cc_proto",
|
||||
"//mediapipe/util/sequence:media_sequence",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "vector_int_to_tensor_calculator",
|
||||
srcs = ["vector_int_to_tensor_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":vector_int_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -667,7 +719,7 @@ cc_library(
|
||||
srcs = ["vector_float_to_tensor_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/tensorflow:vector_float_to_tensor_calculator_options_cc_proto",
|
||||
":vector_float_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -676,12 +728,26 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "unpack_yt8m_sequence_example_calculator",
|
||||
srcs = ["unpack_yt8m_sequence_example_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":lapped_tensor_buffer_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "graph_tensors_packet_generator_test",
|
||||
srcs = ["graph_tensors_packet_generator_test.cc"],
|
||||
deps = [
|
||||
":graph_tensors_packet_generator",
|
||||
"//mediapipe/calculators/tensorflow:graph_tensors_packet_generator_cc_proto",
|
||||
":graph_tensors_packet_generator_cc_proto",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework:packet_generator_cc_proto",
|
||||
"//mediapipe/framework:packet_set",
|
||||
@@ -713,7 +779,7 @@ cc_test(
|
||||
srcs = ["matrix_to_tensor_calculator_test.cc"],
|
||||
deps = [
|
||||
":matrix_to_tensor_calculator",
|
||||
"//mediapipe/calculators/tensorflow:matrix_to_tensor_calculator_options_cc_proto",
|
||||
":matrix_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
@@ -729,13 +795,13 @@ cc_test(
|
||||
srcs = ["lapped_tensor_buffer_calculator_test.cc"],
|
||||
deps = [
|
||||
":lapped_tensor_buffer_calculator",
|
||||
"//mediapipe/calculators/tensorflow:lapped_tensor_buffer_calculator_cc_proto",
|
||||
":lapped_tensor_buffer_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -774,7 +840,7 @@ cc_test(
|
||||
"//mediapipe/util/sequence:media_sequence",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -801,7 +867,7 @@ cc_test(
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:direct_session",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
"@org_tensorflow//tensorflow/core:testlib",
|
||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
||||
"@org_tensorflow//tensorflow/core/kernels:math",
|
||||
@@ -817,7 +883,7 @@ cc_test(
|
||||
":tensorflow_inference_calculator",
|
||||
":tensorflow_session",
|
||||
":tensorflow_session_from_frozen_graph_generator",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
":tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework:packet_generator_cc_proto",
|
||||
@@ -831,7 +897,7 @@ cc_test(
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:direct_session",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
"@org_tensorflow//tensorflow/core:testlib",
|
||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
||||
"@org_tensorflow//tensorflow/core/kernels:math",
|
||||
@@ -847,7 +913,7 @@ cc_test(
|
||||
":tensorflow_inference_calculator",
|
||||
":tensorflow_session",
|
||||
":tensorflow_session_from_saved_model_generator",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_saved_model_generator_cc_proto",
|
||||
":tensorflow_session_from_saved_model_generator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework:packet_generator_cc_proto",
|
||||
@@ -857,14 +923,8 @@ cc_test(
|
||||
"//mediapipe/framework/tool:tag_map_helper",
|
||||
"//mediapipe/framework/tool:validate_type",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:all_kernels",
|
||||
"@org_tensorflow//tensorflow/core:direct_session",
|
||||
"@org_tensorflow//tensorflow/core/kernels:array",
|
||||
"@org_tensorflow//tensorflow/core/kernels:bitcast_op",
|
||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
||||
"@org_tensorflow//tensorflow/core/kernels:io",
|
||||
"@org_tensorflow//tensorflow/core/kernels:state",
|
||||
"@org_tensorflow//tensorflow/core/kernels:string",
|
||||
"@org_tensorflow//tensorflow/core/kernels/data:tensor_dataset_op",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -888,14 +948,8 @@ cc_test(
|
||||
"//mediapipe/framework/tool:tag_map_helper",
|
||||
"//mediapipe/framework/tool:validate_type",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:all_kernels",
|
||||
"@org_tensorflow//tensorflow/core:direct_session",
|
||||
"@org_tensorflow//tensorflow/core/kernels:array",
|
||||
"@org_tensorflow//tensorflow/core/kernels:bitcast_op",
|
||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
||||
"@org_tensorflow//tensorflow/core/kernels:io",
|
||||
"@org_tensorflow//tensorflow/core/kernels:state",
|
||||
"@org_tensorflow//tensorflow/core/kernels:string",
|
||||
"@org_tensorflow//tensorflow/core/kernels/data:tensor_dataset_op",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -904,12 +958,12 @@ cc_test(
|
||||
srcs = ["tensor_squeeze_dimensions_calculator_test.cc"],
|
||||
deps = [
|
||||
":tensor_squeeze_dimensions_calculator",
|
||||
"//mediapipe/calculators/tensorflow:tensor_squeeze_dimensions_calculator_cc_proto",
|
||||
":tensor_squeeze_dimensions_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -919,13 +973,13 @@ cc_test(
|
||||
srcs = ["tensor_to_image_frame_calculator_test.cc"],
|
||||
deps = [
|
||||
":tensor_to_image_frame_calculator",
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_image_frame_calculator_cc_proto",
|
||||
":tensor_to_image_frame_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -935,14 +989,14 @@ cc_test(
|
||||
srcs = ["tensor_to_matrix_calculator_test.cc"],
|
||||
deps = [
|
||||
":tensor_to_matrix_calculator",
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_matrix_calculator_cc_proto",
|
||||
":tensor_to_matrix_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -951,12 +1005,12 @@ cc_test(
|
||||
srcs = ["tensor_to_vector_float_calculator_test.cc"],
|
||||
deps = [
|
||||
":tensor_to_vector_float_calculator",
|
||||
"//mediapipe/calculators/tensorflow:tensor_to_vector_float_calculator_options_cc_proto",
|
||||
":tensor_to_vector_float_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -976,7 +1030,21 @@ cc_test(
|
||||
"//mediapipe/util/sequence:media_sequence",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "vector_int_to_tensor_calculator_test",
|
||||
srcs = ["vector_int_to_tensor_calculator_test.cc"],
|
||||
deps = [
|
||||
":vector_int_to_tensor_calculator",
|
||||
":vector_int_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -985,12 +1053,12 @@ cc_test(
|
||||
srcs = ["vector_float_to_tensor_calculator_test.cc"],
|
||||
deps = [
|
||||
":vector_float_to_tensor_calculator",
|
||||
"//mediapipe/calculators/tensorflow:vector_float_to_tensor_calculator_options_cc_proto",
|
||||
":vector_float_to_tensor_calculator_options_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"@org_tensorflow//tensorflow/core:framework",
|
||||
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||
"@org_tensorflow//tensorflow/core:protos_all",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -1014,7 +1082,7 @@ cc_test(
|
||||
":tensorflow_session",
|
||||
":tensorflow_inference_calculator",
|
||||
":tensorflow_session_from_frozen_graph_generator",
|
||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
":tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||
"//mediapipe/framework/deps:file_path",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
|
||||
@@ -29,6 +29,11 @@
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
const char kBufferSize[] = "BUFFER_SIZE";
|
||||
const char kOverlap[] = "OVERLAP";
|
||||
const char kTimestampOffset[] = "TIMESTAMP_OFFSET";
|
||||
const char kCalculatorOptions[] = "CALCULATOR_OPTIONS";
|
||||
|
||||
namespace tf = tensorflow;
|
||||
|
||||
// Given an input stream of tensors, concatenates the tensors over timesteps.
|
||||
@@ -72,6 +77,9 @@ class LappedTensorBufferCalculator : public CalculatorBase {
|
||||
::mediapipe::Status AddBatchDimension(tf::Tensor* input_tensor);
|
||||
|
||||
int steps_until_output_;
|
||||
int buffer_size_;
|
||||
int overlap_;
|
||||
int timestamp_offset_;
|
||||
std::unique_ptr<CircularBuffer<Timestamp>> timestamp_buffer_;
|
||||
std::unique_ptr<CircularBuffer<tf::Tensor>> buffer_;
|
||||
LappedTensorBufferCalculatorOptions options_;
|
||||
@@ -87,6 +95,21 @@ REGISTER_CALCULATOR(LappedTensorBufferCalculator);
|
||||
);
|
||||
RET_CHECK_EQ(cc->Inputs().NumEntries(), 1)
|
||||
<< "Only one output stream is supported.";
|
||||
|
||||
if (cc->InputSidePackets().HasTag(kBufferSize)) {
|
||||
cc->InputSidePackets().Tag(kBufferSize).Set<int>();
|
||||
}
|
||||
if (cc->InputSidePackets().HasTag(kOverlap)) {
|
||||
cc->InputSidePackets().Tag(kOverlap).Set<int>();
|
||||
}
|
||||
if (cc->InputSidePackets().HasTag(kTimestampOffset)) {
|
||||
cc->InputSidePackets().Tag(kTimestampOffset).Set<int>();
|
||||
}
|
||||
if (cc->InputSidePackets().HasTag(kCalculatorOptions)) {
|
||||
cc->InputSidePackets()
|
||||
.Tag(kCalculatorOptions)
|
||||
.Set<LappedTensorBufferCalculatorOptions>();
|
||||
}
|
||||
cc->Outputs().Index(0).Set<tf::Tensor>(
|
||||
// Output tensorflow::Tensor stream with possibly overlapping steps.
|
||||
);
|
||||
@@ -95,16 +118,33 @@ REGISTER_CALCULATOR(LappedTensorBufferCalculator);
|
||||
|
||||
::mediapipe::Status LappedTensorBufferCalculator::Open(CalculatorContext* cc) {
|
||||
options_ = cc->Options<LappedTensorBufferCalculatorOptions>();
|
||||
RET_CHECK_LT(options_.overlap(), options_.buffer_size());
|
||||
RET_CHECK_GE(options_.timestamp_offset(), 0)
|
||||
if (cc->InputSidePackets().HasTag(kCalculatorOptions)) {
|
||||
options_ = cc->InputSidePackets()
|
||||
.Tag(kCalculatorOptions)
|
||||
.Get<LappedTensorBufferCalculatorOptions>();
|
||||
}
|
||||
buffer_size_ = options_.buffer_size();
|
||||
if (cc->InputSidePackets().HasTag(kBufferSize)) {
|
||||
buffer_size_ = cc->InputSidePackets().Tag(kBufferSize).Get<int>();
|
||||
}
|
||||
overlap_ = options_.overlap();
|
||||
if (cc->InputSidePackets().HasTag(kOverlap)) {
|
||||
overlap_ = cc->InputSidePackets().Tag(kOverlap).Get<int>();
|
||||
}
|
||||
timestamp_offset_ = options_.timestamp_offset();
|
||||
if (cc->InputSidePackets().HasTag(kTimestampOffset)) {
|
||||
timestamp_offset_ = cc->InputSidePackets().Tag(kTimestampOffset).Get<int>();
|
||||
}
|
||||
|
||||
RET_CHECK_LT(overlap_, buffer_size_);
|
||||
RET_CHECK_GE(timestamp_offset_, 0)
|
||||
<< "Negative timestamp_offset is not allowed.";
|
||||
RET_CHECK_LT(options_.timestamp_offset(), options_.buffer_size())
|
||||
RET_CHECK_LT(timestamp_offset_, buffer_size_)
|
||||
<< "output_frame_num_offset has to be less than buffer_size.";
|
||||
timestamp_buffer_ =
|
||||
absl::make_unique<CircularBuffer<Timestamp>>(options_.buffer_size());
|
||||
buffer_ =
|
||||
absl::make_unique<CircularBuffer<tf::Tensor>>(options_.buffer_size());
|
||||
steps_until_output_ = options_.buffer_size();
|
||||
absl::make_unique<CircularBuffer<Timestamp>>(buffer_size_);
|
||||
buffer_ = absl::make_unique<CircularBuffer<tf::Tensor>>(buffer_size_);
|
||||
steps_until_output_ = buffer_size_;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -128,11 +168,10 @@ REGISTER_CALCULATOR(LappedTensorBufferCalculator);
|
||||
concatenated.get());
|
||||
RET_CHECK(concat_status.ok()) << concat_status.ToString();
|
||||
|
||||
cc->Outputs().Index(0).Add(
|
||||
concatenated.release(),
|
||||
timestamp_buffer_->Get(options_.timestamp_offset()));
|
||||
cc->Outputs().Index(0).Add(concatenated.release(),
|
||||
timestamp_buffer_->Get(timestamp_offset_));
|
||||
|
||||
steps_until_output_ = options_.buffer_size() - options_.overlap();
|
||||
steps_until_output_ = buffer_size_ - overlap_;
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -264,7 +264,7 @@ class PackMediaSequenceCalculator : public CalculatorBase {
|
||||
if (options.output_only_if_all_present()) {
|
||||
::mediapipe::Status status = VerifySequence();
|
||||
if (!status.ok()) {
|
||||
cc->GetCounter(status.error_message())->Increment();
|
||||
cc->GetCounter(status.ToString())->Increment();
|
||||
return status;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -34,7 +34,7 @@
|
||||
#include "tensorflow/core/framework/tensor_shape.h"
|
||||
#include "tensorflow/core/framework/tensor_util.h"
|
||||
|
||||
#if !defined(__ANDROID__) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_MOBILE) && !defined(__APPLE__)
|
||||
#include "tensorflow/core/profiler/lib/traceme.h"
|
||||
#endif
|
||||
|
||||
@@ -441,7 +441,7 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
|
||||
const int64 run_start_time = absl::ToUnixMicros(clock_->TimeNow());
|
||||
tf::Status tf_status;
|
||||
{
|
||||
#if !defined(__ANDROID__) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_MOBILE) && !defined(__APPLE__)
|
||||
tensorflow::profiler::TraceMe trace(absl::string_view(cc->NodeName()));
|
||||
#endif
|
||||
tf_status = session_->Run(input_tensors, output_tensor_names,
|
||||
@@ -454,7 +454,7 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
|
||||
|
||||
// RET_CHECK on the tf::Status object itself in order to print an
|
||||
// informative error message.
|
||||
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.error_message();
|
||||
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.ToString();
|
||||
|
||||
const int64 run_end_time = absl::ToUnixMicros(clock_->TimeNow());
|
||||
cc->GetCounter(kTotalSessionRunsTimeUsecsCounterSuffix)
|
||||
|
||||
@@ -31,8 +31,7 @@
|
||||
#include "mediapipe/framework/tool/status_util.h"
|
||||
#include "tensorflow/core/public/session_options.h"
|
||||
|
||||
#if defined(MEDIAPIPE_LITE) || defined(__ANDROID__) || \
|
||||
defined(__APPLE__) && !TARGET_OS_OSX
|
||||
#if defined(MEDIAPIPE_MOBILE)
|
||||
#include "mediapipe/util/android/file/base/helpers.h"
|
||||
#else
|
||||
#include "mediapipe/framework/port/file_helpers.h"
|
||||
@@ -110,7 +109,7 @@ class TensorFlowSessionFromFrozenGraphCalculator : public CalculatorBase {
|
||||
|
||||
RET_CHECK(graph_def.ParseFromString(graph_def_serialized));
|
||||
const tf::Status tf_status = session->session->Create(graph_def);
|
||||
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.error_message();
|
||||
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.ToString();
|
||||
|
||||
for (const auto& key_value : options.tag_to_tensor_names()) {
|
||||
session->tag_to_tensor_map[key_value.first] = key_value.second;
|
||||
@@ -120,7 +119,7 @@ class TensorFlowSessionFromFrozenGraphCalculator : public CalculatorBase {
|
||||
session->session->Run({}, {}, initialization_op_names, {});
|
||||
// RET_CHECK on the tf::Status object itself in order to print an
|
||||
// informative error message.
|
||||
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.error_message();
|
||||
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.ToString();
|
||||
}
|
||||
|
||||
cc->OutputSidePackets().Tag("SESSION").Set(Adopt(session.release()));
|
||||
|
||||
@@ -109,7 +109,7 @@ class TensorFlowSessionFromFrozenGraphGenerator : public PacketGenerator {
|
||||
|
||||
RET_CHECK(graph_def.ParseFromString(graph_def_serialized));
|
||||
const tf::Status tf_status = session->session->Create(graph_def);
|
||||
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.error_message();
|
||||
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.ToString();
|
||||
|
||||
for (const auto& key_value : options.tag_to_tensor_names()) {
|
||||
session->tag_to_tensor_map[key_value.first] = key_value.second;
|
||||
@@ -119,7 +119,7 @@ class TensorFlowSessionFromFrozenGraphGenerator : public PacketGenerator {
|
||||
session->session->Run({}, {}, initialization_op_names, {});
|
||||
// RET_CHECK on the tf::Status object itself in order to print an
|
||||
// informative error message.
|
||||
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.error_message();
|
||||
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.ToString();
|
||||
}
|
||||
|
||||
output_side_packets->Tag("SESSION") = Adopt(session.release());
|
||||
|
||||
@@ -140,7 +140,7 @@ class TensorFlowSessionFromSavedModelCalculator : public CalculatorBase {
|
||||
if (!status.ok()) {
|
||||
return ::mediapipe::Status(
|
||||
static_cast<::mediapipe::StatusCode>(status.code()),
|
||||
status.error_message());
|
||||
status.ToString());
|
||||
}
|
||||
|
||||
auto session = absl::make_unique<TensorFlowSession>();
|
||||
|
||||
@@ -135,7 +135,7 @@ class TensorFlowSessionFromSavedModelGenerator : public PacketGenerator {
|
||||
if (!status.ok()) {
|
||||
return ::mediapipe::Status(
|
||||
static_cast<::mediapipe::StatusCode>(status.code()),
|
||||
status.error_message());
|
||||
status.ToString());
|
||||
}
|
||||
|
||||
auto session = absl::make_unique<TensorFlowSession>();
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
// Copyright 2019 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 "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "tensorflow/core/example/example.pb.h"
|
||||
#include "tensorflow/core/lib/core/status.h"
|
||||
#include "tensorflow/core/lib/io/record_reader.h"
|
||||
#include "tensorflow/core/platform/env.h"
|
||||
#include "tensorflow/core/platform/file_system.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
const char kTFRecordPath[] = "TFRECORD_PATH";
|
||||
const char kRecordIndex[] = "RECORD_INDEX";
|
||||
const char kExampleTag[] = "EXAMPLE";
|
||||
const char kSequenceExampleTag[] = "SEQUENCE_EXAMPLE";
|
||||
|
||||
// Reads a tensorflow example/sequence example from a tfrecord file.
|
||||
// If the "RECORD_INDEX" input side packet is provided, the calculator is going
|
||||
// to fetch the example/sequence example of the tfrecord file at the target
|
||||
// record index. Otherwise, the reader always reads the first example/sequence
|
||||
// example of the tfrecord file.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "TFRecordReaderCalculator"
|
||||
// input_side_packet: "TFRECORD_PATH:tfrecord_path"
|
||||
// input_side_packet: "RECORD_INDEX:record_index"
|
||||
// output_side_packet: "SEQUENCE_EXAMPLE:sequence_example"
|
||||
// }
|
||||
class TFRecordReaderCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
};
|
||||
|
||||
::mediapipe::Status TFRecordReaderCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
cc->InputSidePackets().Tag(kTFRecordPath).Set<std::string>();
|
||||
if (cc->InputSidePackets().HasTag(kRecordIndex)) {
|
||||
cc->InputSidePackets().Tag(kRecordIndex).Set<int>();
|
||||
}
|
||||
|
||||
RET_CHECK(cc->OutputSidePackets().HasTag(kExampleTag) ||
|
||||
cc->OutputSidePackets().HasTag(kSequenceExampleTag))
|
||||
<< "TFRecordReaderCalculator must output either Tensorflow example or "
|
||||
"sequence example.";
|
||||
if (cc->OutputSidePackets().HasTag(kExampleTag)) {
|
||||
cc->OutputSidePackets().Tag(kExampleTag).Set<tensorflow::Example>();
|
||||
} else {
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kSequenceExampleTag)
|
||||
.Set<tensorflow::SequenceExample>();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status TFRecordReaderCalculator::Open(CalculatorContext* cc) {
|
||||
std::unique_ptr<tensorflow::RandomAccessFile> file;
|
||||
auto tf_status = tensorflow::Env::Default()->NewRandomAccessFile(
|
||||
cc->InputSidePackets().Tag(kTFRecordPath).Get<std::string>(), &file);
|
||||
RET_CHECK(tf_status.ok())
|
||||
<< "Failed to open tfrecord file: " << tf_status.ToString();
|
||||
tensorflow::io::RecordReader reader(file.get(),
|
||||
tensorflow::io::RecordReaderOptions());
|
||||
tensorflow::uint64 offset = 0;
|
||||
tensorflow::tstring example_str;
|
||||
const int target_idx =
|
||||
cc->InputSidePackets().HasTag(kRecordIndex)
|
||||
? cc->InputSidePackets().Tag(kRecordIndex).Get<int>()
|
||||
: 0;
|
||||
int current_idx = 0;
|
||||
while (current_idx <= target_idx) {
|
||||
tf_status = reader.ReadRecord(&offset, &example_str);
|
||||
RET_CHECK(tf_status.ok())
|
||||
<< "Failed to read tfrecord: " << tf_status.ToString();
|
||||
if (current_idx == target_idx) {
|
||||
if (cc->OutputSidePackets().HasTag(kExampleTag)) {
|
||||
tensorflow::Example tf_example;
|
||||
tf_example.ParseFromArray(example_str.data(), example_str.size());
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kExampleTag)
|
||||
.Set(MakePacket<tensorflow::Example>(std::move(tf_example)));
|
||||
} else {
|
||||
tensorflow::SequenceExample tf_sequence_example;
|
||||
tf_sequence_example.ParseFromString(example_str);
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kSequenceExampleTag)
|
||||
.Set(MakePacket<tensorflow::SequenceExample>(
|
||||
std::move(tf_sequence_example)));
|
||||
}
|
||||
}
|
||||
++current_idx;
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status TFRecordReaderCalculator::Process(CalculatorContext* cc) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
REGISTER_CALCULATOR(TFRecordReaderCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,192 @@
|
||||
// Copyright 2019 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 <iterator>
|
||||
|
||||
#include "mediapipe/calculators/tensorflow/lapped_tensor_buffer_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "tensorflow/core/example/example.pb.h"
|
||||
#include "tensorflow/core/example/feature.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
const char kId[] = "id";
|
||||
const char kRgb[] = "rgb";
|
||||
const char kAudio[] = "audio";
|
||||
const char kDesiredSegmentSize[] = "DESIRED_SEGMENT_SIZE";
|
||||
const char kYt8mId[] = "YT8M_ID";
|
||||
const char kYt8mSequenceExample[] = "YT8M_SEQUENCE_EXAMPLE";
|
||||
const char kQuantizedRgbFeature[] = "QUANTIZED_RGB_FEATURE";
|
||||
const char kQuantizedAudioFeature[] = "QUANTIZED_AUDIO_FEATURE";
|
||||
const char kSegmentSize[] = "SEGMENT_SIZE";
|
||||
const char kLappedTensorBufferCalculatorOptions[] =
|
||||
"LAPPED_TENSOR_BUFFER_CALCULATOR_OPTIONS";
|
||||
|
||||
std::string GetQuantizedFeature(
|
||||
const tensorflow::SequenceExample& sequence_example, const std::string& key,
|
||||
int index) {
|
||||
const auto& bytes_list = sequence_example.feature_lists()
|
||||
.feature_list()
|
||||
.at(key)
|
||||
.feature()
|
||||
.Get(index)
|
||||
.bytes_list()
|
||||
.value();
|
||||
CHECK_EQ(1, bytes_list.size());
|
||||
return bytes_list.Get(0);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
// Unpacks YT8M Sequence Example. Note that the audio feature and rgb feature
|
||||
// output are quantized. DequantizeByteArrayCalculator can do the dequantization
|
||||
// for you.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "UnpackYt8mSequenceExampleCalculator"
|
||||
// input_side_packet: "YT8M_SEQUENCE_EXAMPLE:yt8m_sequence_example"
|
||||
// output_stream: "QUANTIZED_RGB_FEATURE:quantized_rgb_feature"
|
||||
// output_stream: "QUANTIZED_AUDIO_FEATURE:quantized_audio_feature"
|
||||
// }
|
||||
class UnpackYt8mSequenceExampleCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->InputSidePackets()
|
||||
.Tag(kYt8mSequenceExample)
|
||||
.Set<tensorflow::SequenceExample>();
|
||||
if (cc->InputSidePackets().HasTag(kDesiredSegmentSize)) {
|
||||
cc->InputSidePackets().Tag(kDesiredSegmentSize).Set<int>();
|
||||
}
|
||||
cc->Outputs().Tag(kQuantizedRgbFeature).Set<std::string>();
|
||||
cc->Outputs().Tag(kQuantizedAudioFeature).Set<std::string>();
|
||||
if (cc->OutputSidePackets().HasTag(kYt8mId)) {
|
||||
cc->OutputSidePackets().Tag(kYt8mId).Set<std::string>();
|
||||
}
|
||||
if (cc->OutputSidePackets().HasTag(kLappedTensorBufferCalculatorOptions)) {
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kLappedTensorBufferCalculatorOptions)
|
||||
.Set<::mediapipe::LappedTensorBufferCalculatorOptions>();
|
||||
}
|
||||
if (cc->OutputSidePackets().HasTag(kSegmentSize)) {
|
||||
cc->OutputSidePackets().Tag(kSegmentSize).Set<int>();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
const tensorflow::SequenceExample& sequence_example =
|
||||
cc->InputSidePackets()
|
||||
.Tag(kYt8mSequenceExample)
|
||||
.Get<tensorflow::SequenceExample>();
|
||||
const std::string& yt8m_id =
|
||||
sequence_example.context().feature().at(kId).bytes_list().value().Get(
|
||||
0);
|
||||
if (cc->OutputSidePackets().HasTag(kYt8mId)) {
|
||||
cc->OutputSidePackets().Tag(kYt8mId).Set(
|
||||
MakePacket<std::string>(yt8m_id));
|
||||
}
|
||||
|
||||
int rgb_feature_list_length =
|
||||
sequence_example.feature_lists().feature_list().at(kRgb).feature_size();
|
||||
int audio_feature_list_length = sequence_example.feature_lists()
|
||||
.feature_list()
|
||||
.at(kAudio)
|
||||
.feature_size();
|
||||
|
||||
if (rgb_feature_list_length != audio_feature_list_length) {
|
||||
return ::mediapipe::FailedPreconditionError(absl::StrCat(
|
||||
"Data corruption: the length of audio features and rgb features are "
|
||||
"not equal. Please check the sequence example that contains yt8m "
|
||||
"id: ",
|
||||
yt8m_id));
|
||||
}
|
||||
feature_list_length_ = rgb_feature_list_length;
|
||||
if (cc->OutputSidePackets().HasTag(kLappedTensorBufferCalculatorOptions) ||
|
||||
cc->OutputSidePackets().HasTag(kSegmentSize)) {
|
||||
// If the desired segment size is specified, take the min of the length of
|
||||
// the feature list and the desired size to be the output segment size.
|
||||
int segment_size = feature_list_length_;
|
||||
if (cc->InputSidePackets().HasTag(kDesiredSegmentSize)) {
|
||||
int desired_segment_size =
|
||||
cc->InputSidePackets().Tag(kDesiredSegmentSize).Get<int>();
|
||||
RET_CHECK(desired_segment_size > 0)
|
||||
<< "The desired segment size must be greater than zero.";
|
||||
segment_size = std::min(
|
||||
feature_list_length_,
|
||||
cc->InputSidePackets().Tag(kDesiredSegmentSize).Get<int>());
|
||||
}
|
||||
if (cc->OutputSidePackets().HasTag(
|
||||
kLappedTensorBufferCalculatorOptions)) {
|
||||
auto lapped_tensor_buffer_calculator_options = absl::make_unique<
|
||||
::mediapipe::LappedTensorBufferCalculatorOptions>();
|
||||
lapped_tensor_buffer_calculator_options->set_add_batch_dim_to_tensors(
|
||||
true);
|
||||
lapped_tensor_buffer_calculator_options->set_buffer_size(segment_size);
|
||||
lapped_tensor_buffer_calculator_options->set_overlap(segment_size - 1);
|
||||
lapped_tensor_buffer_calculator_options->set_timestamp_offset(
|
||||
segment_size - 1);
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kLappedTensorBufferCalculatorOptions)
|
||||
.Set(Adopt(lapped_tensor_buffer_calculator_options.release()));
|
||||
}
|
||||
if (cc->OutputSidePackets().HasTag(kSegmentSize)) {
|
||||
cc->OutputSidePackets()
|
||||
.Tag(kSegmentSize)
|
||||
.Set(MakePacket<int>(segment_size));
|
||||
}
|
||||
}
|
||||
LOG(INFO) << "Reading the sequence example that contains yt8m id: "
|
||||
<< yt8m_id << ". Feature list length: " << feature_list_length_;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
if (current_index_ >= feature_list_length_) {
|
||||
return ::mediapipe::tool::StatusStop();
|
||||
}
|
||||
const tensorflow::SequenceExample& sequence_example =
|
||||
cc->InputSidePackets()
|
||||
.Tag(kYt8mSequenceExample)
|
||||
.Get<tensorflow::SequenceExample>();
|
||||
|
||||
// Uses microsecond as the unit of time. In the YT8M dataset, each feature
|
||||
// represents a second.
|
||||
const Timestamp timestamp = Timestamp(current_index_ * 1000000);
|
||||
cc->Outputs()
|
||||
.Tag(kQuantizedRgbFeature)
|
||||
.AddPacket(
|
||||
MakePacket<std::string>(
|
||||
GetQuantizedFeature(sequence_example, kRgb, current_index_))
|
||||
.At(timestamp));
|
||||
cc->Outputs()
|
||||
.Tag(kQuantizedAudioFeature)
|
||||
.AddPacket(
|
||||
MakePacket<std::string>(
|
||||
GetQuantizedFeature(sequence_example, kAudio, current_index_))
|
||||
.At(timestamp));
|
||||
++current_index_;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
int current_index_ = 0;
|
||||
int feature_list_length_ = 0;
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(UnpackYt8mSequenceExampleCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -23,10 +23,12 @@
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace tf = ::tensorflow;
|
||||
|
||||
namespace {
|
||||
auto& INPUT_1D = VectorFloatToTensorCalculatorOptions::INPUT_1D;
|
||||
auto& INPUT_2D = VectorFloatToTensorCalculatorOptions::INPUT_2D;
|
||||
} // namespace
|
||||
|
||||
namespace tf = ::tensorflow;
|
||||
|
||||
// The calculator expects one input (a packet containing a vector<float> or
|
||||
// vector<vector<float>>) and generates one output (a packet containing a
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
// Copyright 2019 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.
|
||||
//
|
||||
// Converts a single int or vector<int> or vector<vector<int>> to 1D (or 2D)
|
||||
// tf::Tensor.
|
||||
|
||||
#include "mediapipe/calculators/tensorflow/vector_int_to_tensor_calculator_options.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "tensorflow/core/framework/tensor.h"
|
||||
#include "tensorflow/core/framework/types.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
const char kVectorInt[] = "VECTOR_INT";
|
||||
const char kSingleInt[] = "SINGLE_INT";
|
||||
const char kTensorOut[] = "TENSOR_OUT";
|
||||
|
||||
namespace {
|
||||
auto& INPUT_1D = VectorIntToTensorCalculatorOptions::INPUT_1D;
|
||||
auto& INPUT_2D = VectorIntToTensorCalculatorOptions::INPUT_2D;
|
||||
} // namespace
|
||||
|
||||
namespace tf = ::tensorflow;
|
||||
|
||||
template <typename TensorType>
|
||||
void AssignMatrixValue(int r, int c, int value, tf::Tensor* output_tensor) {
|
||||
output_tensor->tensor<TensorType, 2>()(r, c) = value;
|
||||
}
|
||||
|
||||
// The calculator expects one input (a packet containing a single int or
|
||||
// vector<int> or vector<vector<int>>) and generates one output (a packet
|
||||
// containing a tf::Tensor containing the same data). The output tensor will be
|
||||
// either 1D or 2D with dimensions corresponding to the input vector int. It
|
||||
// will hold DT_INT32 or DT_UINT8 or DT_INT64 values.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "VectorIntToTensorCalculator"
|
||||
// input_stream: "SINGLE_INT:segment_size_int_stream"
|
||||
// output_stream: "TENSOR_OUT:segment_size_tensor"
|
||||
// }
|
||||
//
|
||||
// or
|
||||
//
|
||||
// node {
|
||||
// calculator: "VectorIntToTensorCalculator"
|
||||
// input_stream: "VECTOR_INT:vector_int_features"
|
||||
// output_stream: "TENSOR_OUT:tensor_features"
|
||||
// }
|
||||
class VectorIntToTensorCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
VectorIntToTensorCalculatorOptions options_;
|
||||
};
|
||||
REGISTER_CALCULATOR(VectorIntToTensorCalculator);
|
||||
|
||||
::mediapipe::Status VectorIntToTensorCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
const auto& options = cc->Options<VectorIntToTensorCalculatorOptions>();
|
||||
// Start with only one input packet.
|
||||
RET_CHECK_EQ(cc->Inputs().NumEntries(), 1)
|
||||
<< "Only one input stream is supported.";
|
||||
if (options.input_size() == INPUT_2D) {
|
||||
cc->Inputs().Tag(kVectorInt).Set<std::vector<std::vector<int>>>();
|
||||
} else if (options.input_size() == INPUT_1D) {
|
||||
if (cc->Inputs().HasTag(kSingleInt)) {
|
||||
cc->Inputs().Tag(kSingleInt).Set<int>();
|
||||
} else {
|
||||
cc->Inputs().Tag(kVectorInt).Set<std::vector<int>>();
|
||||
}
|
||||
} else {
|
||||
LOG(FATAL) << "input size not supported";
|
||||
}
|
||||
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1)
|
||||
<< "Only one output stream is supported.";
|
||||
cc->Outputs().Tag(kTensorOut).Set<tf::Tensor>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status VectorIntToTensorCalculator::Open(CalculatorContext* cc) {
|
||||
options_ = cc->Options<VectorIntToTensorCalculatorOptions>();
|
||||
RET_CHECK(options_.tensor_data_type() == tf::DT_UINT8 ||
|
||||
options_.tensor_data_type() == tf::DT_INT32 ||
|
||||
options_.tensor_data_type() == tf::DT_INT64)
|
||||
<< "Output tensor data type is not supported.";
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status VectorIntToTensorCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
tf::TensorShape tensor_shape;
|
||||
if (options_.input_size() == INPUT_2D) {
|
||||
const std::vector<std::vector<int>>& input =
|
||||
cc->Inputs()
|
||||
.Tag(kVectorInt)
|
||||
.Value()
|
||||
.Get<std::vector<std::vector<int>>>();
|
||||
|
||||
const int32 rows = input.size();
|
||||
CHECK_GE(rows, 1);
|
||||
const int32 cols = input[0].size();
|
||||
CHECK_GE(cols, 1);
|
||||
for (int i = 1; i < rows; ++i) {
|
||||
CHECK_EQ(input[i].size(), cols);
|
||||
}
|
||||
if (options_.transpose()) {
|
||||
tensor_shape = tf::TensorShape({cols, rows});
|
||||
} else {
|
||||
tensor_shape = tf::TensorShape({rows, cols});
|
||||
}
|
||||
auto output = ::absl::make_unique<tf::Tensor>(options_.tensor_data_type(),
|
||||
tensor_shape);
|
||||
if (options_.transpose()) {
|
||||
for (int r = 0; r < rows; ++r) {
|
||||
for (int c = 0; c < cols; ++c) {
|
||||
switch (options_.tensor_data_type()) {
|
||||
case tf::DT_INT64:
|
||||
AssignMatrixValue<tf::int64>(c, r, input[r][c], output.get());
|
||||
break;
|
||||
case tf::DT_UINT8:
|
||||
AssignMatrixValue<uint8>(c, r, input[r][c], output.get());
|
||||
break;
|
||||
case tf::DT_INT32:
|
||||
AssignMatrixValue<int>(c, r, input[r][c], output.get());
|
||||
break;
|
||||
default:
|
||||
LOG(FATAL) << "tensor data type is not supported.";
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for (int r = 0; r < rows; ++r) {
|
||||
for (int c = 0; c < cols; ++c) {
|
||||
switch (options_.tensor_data_type()) {
|
||||
case tf::DT_INT64:
|
||||
AssignMatrixValue<tf::int64>(r, c, input[r][c], output.get());
|
||||
break;
|
||||
case tf::DT_UINT8:
|
||||
AssignMatrixValue<uint8>(r, c, input[r][c], output.get());
|
||||
break;
|
||||
case tf::DT_INT32:
|
||||
AssignMatrixValue<int>(r, c, input[r][c], output.get());
|
||||
break;
|
||||
default:
|
||||
LOG(FATAL) << "tensor data type is not supported.";
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
cc->Outputs().Tag(kTensorOut).Add(output.release(), cc->InputTimestamp());
|
||||
} else if (options_.input_size() == INPUT_1D) {
|
||||
std::vector<int> input;
|
||||
if (cc->Inputs().HasTag(kSingleInt)) {
|
||||
input.push_back(cc->Inputs().Tag(kSingleInt).Get<int>());
|
||||
} else {
|
||||
input = cc->Inputs().Tag(kVectorInt).Value().Get<std::vector<int>>();
|
||||
}
|
||||
CHECK_GE(input.size(), 1);
|
||||
const int32 length = input.size();
|
||||
tensor_shape = tf::TensorShape({length});
|
||||
auto output = ::absl::make_unique<tf::Tensor>(options_.tensor_data_type(),
|
||||
tensor_shape);
|
||||
for (int i = 0; i < length; ++i) {
|
||||
switch (options_.tensor_data_type()) {
|
||||
case tf::DT_INT64:
|
||||
output->tensor<tf::int64, 1>()(i) = input.at(i);
|
||||
break;
|
||||
case tf::DT_UINT8:
|
||||
output->tensor<uint8, 1>()(i) = input.at(i);
|
||||
break;
|
||||
case tf::DT_INT32:
|
||||
output->tensor<int, 1>()(i) = input.at(i);
|
||||
break;
|
||||
default:
|
||||
LOG(FATAL) << "tensor data type is not supported.";
|
||||
}
|
||||
}
|
||||
cc->Outputs().Tag(kTensorOut).Add(output.release(), cc->InputTimestamp());
|
||||
} else {
|
||||
LOG(FATAL) << "input size not supported";
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,43 @@
|
||||
// Copyright 2019 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;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
import "tensorflow/core/framework/types.proto";
|
||||
|
||||
message VectorIntToTensorCalculatorOptions {
|
||||
extend mediapipe.CalculatorOptions {
|
||||
optional VectorIntToTensorCalculatorOptions ext = 275364184;
|
||||
}
|
||||
enum InputSize {
|
||||
UNKNOWN = 0;
|
||||
INPUT_1D = 1;
|
||||
INPUT_2D = 2;
|
||||
}
|
||||
|
||||
// If input_size is INPUT_2D, unpack a vector<vector<int>> to a
|
||||
// 2d tensor (matrix). If INPUT_1D, convert a single int or vector<int>
|
||||
// into a 1d tensor (vector).
|
||||
optional InputSize input_size = 1 [default = INPUT_1D];
|
||||
|
||||
// If true, the output tensor is transposed.
|
||||
// Otherwise, the output tensor is not transposed.
|
||||
// It will be ignored if tensor_is_2d is INPUT_1D.
|
||||
optional bool transpose = 2 [default = false];
|
||||
|
||||
optional tensorflow.DataType tensor_data_type = 3 [default = DT_INT32];
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
// Copyright 2018 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/tensorflow/vector_int_to_tensor_calculator_options.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "tensorflow/core/framework/tensor.h"
|
||||
#include "tensorflow/core/framework/types.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
namespace tf = ::tensorflow;
|
||||
|
||||
class VectorIntToTensorCalculatorTest : public ::testing::Test {
|
||||
protected:
|
||||
void SetUpRunner(
|
||||
const VectorIntToTensorCalculatorOptions::InputSize input_size,
|
||||
const tensorflow::DataType tensor_data_type, const bool transpose,
|
||||
const bool single_value) {
|
||||
CalculatorGraphConfig::Node config;
|
||||
config.set_calculator("VectorIntToTensorCalculator");
|
||||
if (single_value) {
|
||||
config.add_input_stream("SINGLE_INT:input_int");
|
||||
} else {
|
||||
config.add_input_stream("VECTOR_INT:input_int");
|
||||
}
|
||||
config.add_output_stream("TENSOR_OUT:output_tensor");
|
||||
auto options = config.mutable_options()->MutableExtension(
|
||||
VectorIntToTensorCalculatorOptions::ext);
|
||||
options->set_input_size(input_size);
|
||||
options->set_transpose(transpose);
|
||||
options->set_tensor_data_type(tensor_data_type);
|
||||
runner_ = ::absl::make_unique<CalculatorRunner>(config);
|
||||
}
|
||||
|
||||
void TestConvertFromVectoVectorInt(const bool transpose) {
|
||||
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_2D,
|
||||
tensorflow::DT_INT32, transpose, false);
|
||||
auto input = ::absl::make_unique<std::vector<std::vector<int>>>(
|
||||
2, std::vector<int>(2));
|
||||
for (int i = 0; i < 2; ++i) {
|
||||
for (int j = 0; j < 2; ++j) {
|
||||
input->at(i).at(j) = i * 2 + j;
|
||||
}
|
||||
}
|
||||
|
||||
const int64 time = 1234;
|
||||
runner_->MutableInputs()
|
||||
->Tag("VECTOR_INT")
|
||||
.packets.push_back(Adopt(input.release()).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||
|
||||
EXPECT_EQ(2, output_tensor.dims());
|
||||
EXPECT_EQ(tf::DT_INT32, output_tensor.dtype());
|
||||
const auto matrix = output_tensor.matrix<int>();
|
||||
|
||||
for (int i = 0; i < 2; ++i) {
|
||||
for (int j = 0; j < 2; ++j) {
|
||||
if (!transpose) {
|
||||
EXPECT_EQ(i * 2 + j, matrix(i, j));
|
||||
} else {
|
||||
EXPECT_EQ(j * 2 + i, matrix(i, j));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::unique_ptr<CalculatorRunner> runner_;
|
||||
};
|
||||
|
||||
TEST_F(VectorIntToTensorCalculatorTest, TestSingleValue) {
|
||||
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||
tensorflow::DT_INT32, false, true);
|
||||
const int64 time = 1234;
|
||||
runner_->MutableInputs()
|
||||
->Tag("SINGLE_INT")
|
||||
.packets.push_back(MakePacket<int>(1).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||
|
||||
EXPECT_EQ(1, output_tensor.dims());
|
||||
EXPECT_EQ(tf::DT_INT32, output_tensor.dtype());
|
||||
const auto vec = output_tensor.vec<int32>();
|
||||
EXPECT_EQ(1, vec(0));
|
||||
}
|
||||
|
||||
TEST_F(VectorIntToTensorCalculatorTest, TesOneDim) {
|
||||
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||
tensorflow::DT_INT32, false, false);
|
||||
auto input = ::absl::make_unique<std::vector<int>>(5);
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
input->at(i) = i;
|
||||
}
|
||||
const int64 time = 1234;
|
||||
runner_->MutableInputs()
|
||||
->Tag("VECTOR_INT")
|
||||
.packets.push_back(Adopt(input.release()).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||
|
||||
EXPECT_EQ(1, output_tensor.dims());
|
||||
EXPECT_EQ(tf::DT_INT32, output_tensor.dtype());
|
||||
const auto vec = output_tensor.vec<int32>();
|
||||
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
EXPECT_EQ(i, vec(i));
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(VectorIntToTensorCalculatorTest, TestTwoDims) {
|
||||
for (bool transpose : {false, true}) {
|
||||
TestConvertFromVectoVectorInt(transpose);
|
||||
}
|
||||
}
|
||||
|
||||
TEST_F(VectorIntToTensorCalculatorTest, TestInt64) {
|
||||
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||
tensorflow::DT_INT64, false, true);
|
||||
const int64 time = 1234;
|
||||
runner_->MutableInputs()
|
||||
->Tag("SINGLE_INT")
|
||||
.packets.push_back(MakePacket<int>(2 ^ 31).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||
|
||||
EXPECT_EQ(1, output_tensor.dims());
|
||||
EXPECT_EQ(tf::DT_INT64, output_tensor.dtype());
|
||||
const auto vec = output_tensor.vec<tf::int64>();
|
||||
EXPECT_EQ(2 ^ 31, vec(0));
|
||||
}
|
||||
|
||||
TEST_F(VectorIntToTensorCalculatorTest, TestUint8) {
|
||||
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||
tensorflow::DT_UINT8, false, false);
|
||||
auto input = ::absl::make_unique<std::vector<int>>(5);
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
input->at(i) = i;
|
||||
}
|
||||
const int64 time = 1234;
|
||||
runner_->MutableInputs()
|
||||
->Tag("VECTOR_INT")
|
||||
.packets.push_back(Adopt(input.release()).At(Timestamp(time)));
|
||||
|
||||
EXPECT_TRUE(runner_->Run().ok());
|
||||
|
||||
const std::vector<Packet>& output_packets =
|
||||
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||
EXPECT_EQ(1, output_packets.size());
|
||||
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||
|
||||
EXPECT_EQ(1, output_tensor.dims());
|
||||
EXPECT_EQ(tf::DT_UINT8, output_tensor.dtype());
|
||||
const auto vec = output_tensor.vec<uint8>();
|
||||
|
||||
for (int i = 0; i < 5; ++i) {
|
||||
EXPECT_EQ(i, vec(i));
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -13,12 +13,12 @@
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "ssd_anchors_calculator_proto",
|
||||
srcs = ["ssd_anchors_calculator.proto"],
|
||||
@@ -238,6 +238,7 @@ cc_library(
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/metal:buffer_convert",
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate_internal",
|
||||
],
|
||||
"//conditions:default": [
|
||||
"//mediapipe/gpu:gl_calculator_helper",
|
||||
@@ -248,6 +249,11 @@ cc_library(
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
|
||||
],
|
||||
}) + select({
|
||||
"//conditions:default": [],
|
||||
"//mediapipe:android": [
|
||||
"@org_tensorflow//tensorflow/lite/delegates/nnapi:nnapi_delegate",
|
||||
],
|
||||
}),
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
@@ -25,7 +25,7 @@
|
||||
#include "tensorflow/lite/error_reporter.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||
#include "mediapipe/gpu/gpu_buffer.h"
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
||||
@@ -34,7 +34,7 @@
|
||||
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#if defined(MEDIAPIPE_IOS)
|
||||
#import <CoreVideo/CoreVideo.h>
|
||||
#import <Metal/Metal.h>
|
||||
#import <MetalKit/MetalKit.h>
|
||||
@@ -45,9 +45,9 @@
|
||||
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
|
||||
#endif // iOS
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
typedef id<MTLBuffer> GpuTensor;
|
||||
#endif
|
||||
|
||||
@@ -67,7 +67,7 @@ typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
||||
using ::tflite::gpu::gl::GlProgram;
|
||||
using ::tflite::gpu::gl::GlShader;
|
||||
@@ -77,7 +77,7 @@ struct GPUData {
|
||||
GlShader shader;
|
||||
GlProgram program;
|
||||
};
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
struct GPUData {
|
||||
int elements = 1;
|
||||
GpuTensor buffer;
|
||||
@@ -146,10 +146,10 @@ class TfLiteConverterCalculator : public CalculatorBase {
|
||||
|
||||
std::unique_ptr<tflite::Interpreter> interpreter_ = nullptr;
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||
std::unique_ptr<GPUData> gpu_data_out_;
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
MPPMetalHelper* gpu_helper_ = nullptr;
|
||||
std::unique_ptr<GPUData> gpu_data_out_;
|
||||
#endif
|
||||
@@ -181,7 +181,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
|
||||
if (cc->Inputs().HasTag("IMAGE")) cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
||||
if (cc->Inputs().HasTag("MATRIX")) cc->Inputs().Tag("MATRIX").Set<Matrix>();
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||
cc->Inputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
|
||||
use_gpu |= true;
|
||||
@@ -190,7 +190,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
|
||||
if (cc->Outputs().HasTag("TENSORS"))
|
||||
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
if (cc->Outputs().HasTag("TENSORS_GPU")) {
|
||||
cc->Outputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||
use_gpu |= true;
|
||||
@@ -198,9 +198,9 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
if (use_gpu) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
|
||||
#endif
|
||||
}
|
||||
@@ -218,7 +218,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU") ||
|
||||
cc->Outputs().HasTag("IMAGE_OUT_GPU")) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
use_gpu_ = true;
|
||||
#else
|
||||
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
||||
@@ -231,9 +231,9 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
cc->Outputs().HasTag("TENSORS_GPU"));
|
||||
// Cannot use quantization.
|
||||
use_quantized_tensors_ = false;
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
|
||||
RET_CHECK(gpu_helper_);
|
||||
#endif
|
||||
@@ -264,10 +264,10 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
}
|
||||
|
||||
::mediapipe::Status TfLiteConverterCalculator::Close(CalculatorContext* cc) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
gpu_helper_.RunInGlContext([this] { gpu_data_out_.reset(); });
|
||||
#endif
|
||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#if defined(MEDIAPIPE_IOS)
|
||||
gpu_data_out_.reset();
|
||||
#endif
|
||||
return ::mediapipe::OkStatus();
|
||||
@@ -294,11 +294,15 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
if (use_quantized_tensors_) {
|
||||
RET_CHECK(image_frame.Format() != mediapipe::ImageFormat::VEC32F1)
|
||||
<< "Only 8-bit input images are supported for quantization.";
|
||||
quant.type = kTfLiteAffineQuantization;
|
||||
quant.params = nullptr;
|
||||
// Optional: Set 'quant' quantization params here if needed.
|
||||
interpreter_->SetTensorParametersReadWrite(0, kTfLiteUInt8, "",
|
||||
{channels_preserved}, quant);
|
||||
} else {
|
||||
// Default TfLiteQuantization used for no quantization.
|
||||
// Initialize structure for no quantization.
|
||||
quant.type = kTfLiteNoQuantization;
|
||||
quant.params = nullptr;
|
||||
interpreter_->SetTensorParametersReadWrite(0, kTfLiteFloat32, "",
|
||||
{channels_preserved}, quant);
|
||||
}
|
||||
@@ -383,7 +387,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteConverterCalculator::ProcessGPU(
|
||||
CalculatorContext* cc) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
// GpuBuffer to tflite::gpu::GlBuffer conversion.
|
||||
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
||||
MP_RETURN_IF_ERROR(
|
||||
@@ -419,43 +423,38 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
cc->Outputs()
|
||||
.Tag("TENSORS_GPU")
|
||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
// GpuBuffer to id<MTLBuffer> conversion.
|
||||
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
||||
{
|
||||
id<MTLTexture> src_texture = [gpu_helper_ metalTextureWithGpuBuffer:input];
|
||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||
command_buffer.label = @"TfLiteConverterCalculatorConvert";
|
||||
id<MTLComputeCommandEncoder> compute_encoder =
|
||||
[command_buffer computeCommandEncoder];
|
||||
[compute_encoder setComputePipelineState:gpu_data_out_->pipeline_state];
|
||||
[compute_encoder setTexture:src_texture atIndex:0];
|
||||
[compute_encoder setBuffer:gpu_data_out_->buffer offset:0 atIndex:1];
|
||||
MTLSize threads_per_group = MTLSizeMake(kWorkgroupSize, kWorkgroupSize, 1);
|
||||
MTLSize threadgroups =
|
||||
MTLSizeMake(NumGroups(input.width(), kWorkgroupSize),
|
||||
NumGroups(input.height(), kWorkgroupSize), 1);
|
||||
[compute_encoder dispatchThreadgroups:threadgroups
|
||||
threadsPerThreadgroup:threads_per_group];
|
||||
[compute_encoder endEncoding];
|
||||
[command_buffer commit];
|
||||
[command_buffer waitUntilCompleted];
|
||||
}
|
||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||
|
||||
id<MTLTexture> src_texture = [gpu_helper_ metalTextureWithGpuBuffer:input];
|
||||
command_buffer.label = @"TfLiteConverterCalculatorConvertAndBlit";
|
||||
id<MTLComputeCommandEncoder> compute_encoder =
|
||||
[command_buffer computeCommandEncoder];
|
||||
[compute_encoder setComputePipelineState:gpu_data_out_->pipeline_state];
|
||||
[compute_encoder setTexture:src_texture atIndex:0];
|
||||
[compute_encoder setBuffer:gpu_data_out_->buffer offset:0 atIndex:1];
|
||||
MTLSize threads_per_group = MTLSizeMake(kWorkgroupSize, kWorkgroupSize, 1);
|
||||
MTLSize threadgroups =
|
||||
MTLSizeMake(NumGroups(input.width(), kWorkgroupSize),
|
||||
NumGroups(input.height(), kWorkgroupSize), 1);
|
||||
[compute_encoder dispatchThreadgroups:threadgroups
|
||||
threadsPerThreadgroup:threads_per_group];
|
||||
[compute_encoder endEncoding];
|
||||
|
||||
// Copy into outputs.
|
||||
// TODO Avoid this copy.
|
||||
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
||||
output_tensors->resize(1);
|
||||
{
|
||||
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
||||
output_tensors->at(0) =
|
||||
[device newBufferWithLength:gpu_data_out_->elements * sizeof(float)
|
||||
options:MTLResourceStorageModeShared];
|
||||
[MPPMetalUtil blitMetalBufferTo:output_tensors->at(0)
|
||||
from:gpu_data_out_->buffer
|
||||
blocking:true
|
||||
commandBuffer:[gpu_helper_ commandBuffer]];
|
||||
}
|
||||
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
||||
output_tensors->at(0) =
|
||||
[device newBufferWithLength:gpu_data_out_->elements * sizeof(float)
|
||||
options:MTLResourceStorageModeShared];
|
||||
[MPPMetalUtil blitMetalBufferTo:output_tensors->at(0)
|
||||
from:gpu_data_out_->buffer
|
||||
blocking:false
|
||||
commandBuffer:command_buffer];
|
||||
|
||||
cc->Outputs()
|
||||
.Tag("TENSORS_GPU")
|
||||
@@ -468,7 +467,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
}
|
||||
|
||||
::mediapipe::Status TfLiteConverterCalculator::InitGpu(CalculatorContext* cc) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
// Get input image sizes.
|
||||
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
||||
mediapipe::ImageFormat::Format format =
|
||||
@@ -485,7 +484,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
RET_CHECK_FAIL() << "Num input channels is less than desired output.";
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||
[this, &include_alpha, &input, &single_channel]() -> ::mediapipe::Status {
|
||||
// Device memory.
|
||||
@@ -529,7 +528,9 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
&gpu_data_out_->program));
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
|
||||
RET_CHECK(include_alpha)
|
||||
<< "iOS GPU inference currently accepts only RGBA input.";
|
||||
|
||||
@@ -610,7 +611,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
CHECK_GE(max_num_channels_, 1);
|
||||
CHECK_LE(max_num_channels_, 4);
|
||||
CHECK_NE(max_num_channels_, 2);
|
||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#if defined(MEDIAPIPE_IOS)
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU"))
|
||||
// Currently on iOS, tflite gpu input tensor must be 4 channels,
|
||||
// so input image must be 4 channels also (checked in InitGpu).
|
||||
|
||||
@@ -27,7 +27,7 @@
|
||||
#include "tensorflow/lite/kernels/register.h"
|
||||
#include "tensorflow/lite/model.h"
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||
#include "mediapipe/gpu/gpu_buffer.h"
|
||||
#include "tensorflow/lite/delegates/gpu/common/shape.h"
|
||||
@@ -35,9 +35,9 @@
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
|
||||
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
#endif // !MEDIAPIPE_DISABLE_GL_COMPUTE
|
||||
|
||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#if defined(MEDIAPIPE_IOS)
|
||||
#import <CoreVideo/CoreVideo.h>
|
||||
#import <Metal/Metal.h>
|
||||
#import <MetalKit/MetalKit.h>
|
||||
@@ -48,13 +48,22 @@
|
||||
#include "tensorflow/lite/delegates/gpu/common/shape.h"
|
||||
#include "tensorflow/lite/delegates/gpu/metal/buffer_convert.h"
|
||||
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
|
||||
#include "tensorflow/lite/delegates/gpu/metal_delegate_internal.h"
|
||||
#endif // iOS
|
||||
|
||||
namespace {
|
||||
#if defined(MEDIAPIPE_ANDROID)
|
||||
#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
|
||||
#endif // ANDROID
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
namespace {
|
||||
// Commonly used to compute the number of blocks to launch in a kernel.
|
||||
int NumGroups(const int size, const int group_size) { // NOLINT
|
||||
return (size + group_size - 1) / group_size;
|
||||
}
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
typedef id<MTLBuffer> GpuTensor;
|
||||
#endif
|
||||
|
||||
@@ -62,19 +71,41 @@ typedef id<MTLBuffer> GpuTensor;
|
||||
size_t RoundUp(size_t n, size_t m) { return ((n + m - 1) / m) * m; } // NOLINT
|
||||
} // namespace
|
||||
|
||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||
#include "edgetpu.h"
|
||||
|
||||
// Creates and returns an Edge TPU interpreter to run the given edgetpu model.
|
||||
std::unique_ptr<tflite::Interpreter> BuildEdgeTpuInterpreter(
|
||||
const tflite::FlatBufferModel& model,
|
||||
tflite::ops::builtin::BuiltinOpResolver* resolver,
|
||||
edgetpu::EdgeTpuContext* edgetpu_context) {
|
||||
resolver->AddCustom(edgetpu::kCustomOp, edgetpu::RegisterCustomOp());
|
||||
std::unique_ptr<tflite::Interpreter> interpreter;
|
||||
if (tflite::InterpreterBuilder(model, *resolver)(&interpreter) != kTfLiteOk) {
|
||||
std::cerr << "Failed to build edge TPU interpreter." << std::endl;
|
||||
}
|
||||
interpreter->SetExternalContext(kTfLiteEdgeTpuContext, edgetpu_context);
|
||||
interpreter->SetNumThreads(1);
|
||||
if (interpreter->AllocateTensors() != kTfLiteOk) {
|
||||
std::cerr << "Failed to allocate edge TPU tensors." << std::endl;
|
||||
}
|
||||
return interpreter;
|
||||
}
|
||||
#endif // MEDIAPIPE_EDGE_TPU
|
||||
|
||||
// TfLiteInferenceCalculator File Layout:
|
||||
// * Header
|
||||
// * Core
|
||||
// * Aux
|
||||
namespace mediapipe {
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
using ::tflite::gpu::gl::CopyBuffer;
|
||||
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
||||
using ::tflite::gpu::gl::GlBuffer;
|
||||
#endif
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
struct GPUData {
|
||||
int elements = 1;
|
||||
GpuTensor buffer;
|
||||
@@ -147,17 +178,23 @@ class TfLiteInferenceCalculator : public CalculatorBase {
|
||||
std::unique_ptr<tflite::FlatBufferModel> model_;
|
||||
TfLiteDelegate* delegate_ = nullptr;
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||
std::unique_ptr<GPUData> gpu_data_in_;
|
||||
std::vector<std::unique_ptr<GPUData>> gpu_data_in_;
|
||||
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
MPPMetalHelper* gpu_helper_ = nullptr;
|
||||
std::unique_ptr<GPUData> gpu_data_in_;
|
||||
std::vector<std::unique_ptr<GPUData>> gpu_data_in_;
|
||||
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
|
||||
id<MTLComputePipelineState> fp32_to_fp16_program_;
|
||||
TFLBufferConvert* converter_from_BPHWC4_ = nil;
|
||||
#endif
|
||||
|
||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||
std::shared_ptr<edgetpu::EdgeTpuContext> edgetpu_context_ =
|
||||
edgetpu::EdgeTpuManager::GetSingleton()->OpenDevice();
|
||||
#endif
|
||||
|
||||
std::string model_path_ = "";
|
||||
bool gpu_inference_ = false;
|
||||
bool gpu_input_ = false;
|
||||
@@ -179,7 +216,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
if (cc->Inputs().HasTag("TENSORS"))
|
||||
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||
use_gpu |= true;
|
||||
@@ -188,7 +225,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
if (cc->Outputs().HasTag("TENSORS"))
|
||||
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
if (cc->Outputs().HasTag("TENSORS_GPU")) {
|
||||
cc->Outputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||
use_gpu |= true;
|
||||
@@ -206,9 +243,9 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
use_gpu |= options.use_gpu();
|
||||
|
||||
if (use_gpu) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
|
||||
#endif
|
||||
}
|
||||
@@ -225,7 +262,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
MP_RETURN_IF_ERROR(LoadOptions(cc));
|
||||
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
gpu_input_ = true;
|
||||
gpu_inference_ = true; // Inference must be on GPU also.
|
||||
#else
|
||||
@@ -235,7 +272,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("TENSORS_GPU")) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
gpu_output_ = true;
|
||||
RET_CHECK(cc->Inputs().HasTag("TENSORS_GPU"))
|
||||
<< "GPU output must also have GPU Input.";
|
||||
@@ -248,20 +285,24 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
MP_RETURN_IF_ERROR(LoadModel(cc));
|
||||
|
||||
if (gpu_inference_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
|
||||
RET_CHECK(gpu_helper_);
|
||||
#endif
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||
[this, &cc]() -> ::mediapipe::Status { return LoadDelegate(cc); }));
|
||||
#else
|
||||
MP_RETURN_IF_ERROR(LoadDelegate(cc));
|
||||
#endif
|
||||
} else {
|
||||
#if defined(__EMSCRIPTEN__) || defined(MEDIAPIPE_ANDROID)
|
||||
MP_RETURN_IF_ERROR(LoadDelegate(cc));
|
||||
#endif // __EMSCRIPTEN__ || ANDROID
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -269,25 +310,44 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
// 1. Receive pre-processed tensor inputs.
|
||||
if (gpu_input_) {
|
||||
// Read GPU input into SSBO.
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||
RET_CHECK_EQ(input_tensors.size(), 1);
|
||||
RET_CHECK_GT(input_tensors.size(), 0);
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||
[this, &input_tensors]() -> ::mediapipe::Status {
|
||||
// Explicit copy input.
|
||||
RET_CHECK_CALL(CopyBuffer(input_tensors[0], gpu_data_in_->buffer));
|
||||
gpu_data_in_.resize(input_tensors.size());
|
||||
for (int i = 0; i < input_tensors.size(); ++i) {
|
||||
RET_CHECK_CALL(
|
||||
CopyBuffer(input_tensors[i], gpu_data_in_[i]->buffer));
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||
RET_CHECK_EQ(input_tensors.size(), 1);
|
||||
// Explicit copy input.
|
||||
[MPPMetalUtil blitMetalBufferTo:gpu_data_in_->buffer
|
||||
from:input_tensors[0]
|
||||
blocking:true
|
||||
commandBuffer:[gpu_helper_ commandBuffer]];
|
||||
RET_CHECK_GT(input_tensors.size(), 0);
|
||||
// Explicit copy input with conversion float 32 bits to 16 bits.
|
||||
gpu_data_in_.resize(input_tensors.size());
|
||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||
command_buffer.label = @"TfLiteInferenceCalculatorConvert";
|
||||
id<MTLComputeCommandEncoder> compute_encoder =
|
||||
[command_buffer computeCommandEncoder];
|
||||
[compute_encoder setComputePipelineState:fp32_to_fp16_program_];
|
||||
for (int i = 0; i < input_tensors.size(); ++i) {
|
||||
[compute_encoder setBuffer:input_tensors[i] offset:0 atIndex:0];
|
||||
[compute_encoder setBuffer:gpu_data_in_[i]->buffer offset:0 atIndex:1];
|
||||
constexpr int kWorkgroupSize = 64; // Block size for GPU shader.
|
||||
MTLSize threads_per_group = MTLSizeMake(kWorkgroupSize, 1, 1);
|
||||
const int threadgroups =
|
||||
NumGroups(gpu_data_in_[i]->elements, kWorkgroupSize);
|
||||
[compute_encoder dispatchThreadgroups:MTLSizeMake(threadgroups, 1, 1)
|
||||
threadsPerThreadgroup:threads_per_group];
|
||||
}
|
||||
[compute_encoder endEncoding];
|
||||
[command_buffer commit];
|
||||
#else
|
||||
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
||||
#endif
|
||||
@@ -315,13 +375,13 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
// 2. Run inference.
|
||||
if (gpu_inference_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(
|
||||
gpu_helper_.RunInGlContext([this]() -> ::mediapipe::Status {
|
||||
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
|
||||
#endif
|
||||
} else {
|
||||
@@ -330,7 +390,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
// 3. Output processed tensors.
|
||||
if (gpu_output_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
// Output result tensors (GPU).
|
||||
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||
@@ -347,7 +407,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
cc->Outputs()
|
||||
.Tag("TENSORS_GPU")
|
||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
// Output result tensors (GPU).
|
||||
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
||||
output_tensors->resize(gpu_data_out_.size());
|
||||
@@ -368,7 +428,6 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
}
|
||||
[convert_command endEncoding];
|
||||
[command_buffer commit];
|
||||
[command_buffer waitUntilCompleted];
|
||||
cc->Outputs()
|
||||
.Tag("TENSORS_GPU")
|
||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||
@@ -392,24 +451,33 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteInferenceCalculator::Close(CalculatorContext* cc) {
|
||||
if (delegate_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
|
||||
TfLiteGpuDelegateDelete(delegate_);
|
||||
gpu_data_in_.reset();
|
||||
if (gpu_inference_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
|
||||
TfLiteGpuDelegateDelete(delegate_);
|
||||
for (int i = 0; i < gpu_data_in_.size(); ++i) {
|
||||
gpu_data_in_[i].reset();
|
||||
}
|
||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||
gpu_data_out_[i].reset();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
TFLGpuDelegateDelete(delegate_);
|
||||
for (int i = 0; i < gpu_data_in_.size(); ++i) {
|
||||
gpu_data_in_[i].reset();
|
||||
}
|
||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||
gpu_data_out_[i].reset();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
TFLGpuDelegateDelete(delegate_);
|
||||
gpu_data_in_.reset();
|
||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||
gpu_data_out_[i].reset();
|
||||
}
|
||||
#endif
|
||||
}
|
||||
delegate_ = nullptr;
|
||||
}
|
||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||
edgetpu_context_.reset();
|
||||
#endif
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -423,7 +491,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
// Get model name.
|
||||
if (!options.model_path().empty()) {
|
||||
auto model_path = options.model_path();
|
||||
std::string model_path = options.model_path();
|
||||
|
||||
ASSIGN_OR_RETURN(model_path_, mediapipe::PathToResourceAsFile(model_path));
|
||||
} else {
|
||||
@@ -443,19 +511,25 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
model_ = tflite::FlatBufferModel::BuildFromFile(model_path_.c_str());
|
||||
RET_CHECK(model_);
|
||||
|
||||
tflite::ops::builtin::BuiltinOpResolver op_resolver;
|
||||
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
|
||||
const auto& op_resolver =
|
||||
cc->InputSidePackets()
|
||||
.Tag("CUSTOM_OP_RESOLVER")
|
||||
.Get<tflite::ops::builtin::BuiltinOpResolver>();
|
||||
tflite::InterpreterBuilder(*model_, op_resolver)(&interpreter_);
|
||||
} else {
|
||||
const tflite::ops::builtin::BuiltinOpResolver op_resolver;
|
||||
tflite::InterpreterBuilder(*model_, op_resolver)(&interpreter_);
|
||||
op_resolver = cc->InputSidePackets()
|
||||
.Tag("CUSTOM_OP_RESOLVER")
|
||||
.Get<tflite::ops::builtin::BuiltinOpResolver>();
|
||||
}
|
||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||
interpreter_ =
|
||||
BuildEdgeTpuInterpreter(*model_, &op_resolver, edgetpu_context_.get());
|
||||
#else
|
||||
tflite::InterpreterBuilder(*model_, op_resolver)(&interpreter_);
|
||||
#endif // MEDIAPIPE_EDGE_TPU
|
||||
|
||||
RET_CHECK(interpreter_);
|
||||
|
||||
#if defined(__EMSCRIPTEN__)
|
||||
interpreter_->SetNumThreads(1);
|
||||
#endif // __EMSCRIPTEN__
|
||||
|
||||
if (gpu_output_) {
|
||||
use_quantized_tensors_ = false;
|
||||
} else {
|
||||
@@ -471,7 +545,22 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteInferenceCalculator::LoadDelegate(
|
||||
CalculatorContext* cc) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if defined(MEDIAPIPE_ANDROID)
|
||||
if (!gpu_inference_) {
|
||||
if (cc->Options<mediapipe::TfLiteInferenceCalculatorOptions>()
|
||||
.use_nnapi()) {
|
||||
// Attempt to use NNAPI.
|
||||
// If not supported, the default CPU delegate will be created and used.
|
||||
interpreter_->SetAllowFp16PrecisionForFp32(1);
|
||||
delegate_ = tflite::NnApiDelegate();
|
||||
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
|
||||
}
|
||||
// Return, no need for GPU delegate below.
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
#endif // ANDROID
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
// Configure and create the delegate.
|
||||
TfLiteGpuDelegateOptions options = TfLiteGpuDelegateOptionsDefault();
|
||||
options.compile_options.precision_loss_allowed = 1;
|
||||
@@ -483,24 +572,24 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
if (gpu_input_) {
|
||||
// Get input image sizes.
|
||||
gpu_data_in_ = absl::make_unique<GPUData>();
|
||||
const auto& input_indices = interpreter_->inputs();
|
||||
RET_CHECK_EQ(input_indices.size(), 1); // TODO accept > 1.
|
||||
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[0]);
|
||||
gpu_data_in_->elements = 1;
|
||||
for (int d = 0; d < tensor->dims->size; ++d) {
|
||||
gpu_data_in_->elements *= tensor->dims->data[d];
|
||||
gpu_data_in_.resize(input_indices.size());
|
||||
for (int i = 0; i < input_indices.size(); ++i) {
|
||||
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[0]);
|
||||
gpu_data_in_[i] = absl::make_unique<GPUData>();
|
||||
gpu_data_in_[i]->elements = 1;
|
||||
for (int d = 0; d < tensor->dims->size; ++d) {
|
||||
gpu_data_in_[i]->elements *= tensor->dims->data[d];
|
||||
}
|
||||
// Create and bind input buffer.
|
||||
RET_CHECK_CALL(
|
||||
::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
|
||||
gpu_data_in_[i]->elements, &gpu_data_in_[i]->buffer));
|
||||
RET_CHECK_EQ(TfLiteGpuDelegateBindBufferToTensor(
|
||||
delegate_, gpu_data_in_[i]->buffer.id(),
|
||||
interpreter_->inputs()[i]),
|
||||
kTfLiteOk);
|
||||
}
|
||||
CHECK_GE(tensor->dims->data[3], 1);
|
||||
CHECK_LE(tensor->dims->data[3], 4);
|
||||
CHECK_NE(tensor->dims->data[3], 2);
|
||||
// Create and bind input buffer.
|
||||
RET_CHECK_CALL(::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
|
||||
gpu_data_in_->elements, &gpu_data_in_->buffer));
|
||||
RET_CHECK_EQ(TfLiteGpuDelegateBindBufferToTensor(
|
||||
delegate_, gpu_data_in_->buffer.id(),
|
||||
interpreter_->inputs()[0]), // First tensor only
|
||||
kTfLiteOk);
|
||||
}
|
||||
if (gpu_output_) {
|
||||
// Get output image sizes.
|
||||
@@ -531,42 +620,69 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
|
||||
#endif // OpenGL
|
||||
|
||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#if defined(MEDIAPIPE_IOS)
|
||||
const int kHalfSize = 2; // sizeof(half)
|
||||
// Configure and create the delegate.
|
||||
GpuDelegateOptions options;
|
||||
options.allow_precision_loss = false; // Must match converter, F=float/T=half
|
||||
options.wait_type = GpuDelegateOptions::WaitType::kPassive;
|
||||
TFLGpuDelegateOptions options;
|
||||
options.allow_precision_loss = true;
|
||||
options.wait_type = TFLGpuDelegateWaitType::TFLGpuDelegateWaitTypePassive;
|
||||
if (!delegate_) delegate_ = TFLGpuDelegateCreate(&options);
|
||||
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
||||
|
||||
if (gpu_input_) {
|
||||
// Get input image sizes.
|
||||
gpu_data_in_ = absl::make_unique<GPUData>();
|
||||
const auto& input_indices = interpreter_->inputs();
|
||||
RET_CHECK_EQ(input_indices.size(), 1);
|
||||
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[0]);
|
||||
gpu_data_in_->elements = 1;
|
||||
// On iOS GPU, input must be 4 channels, regardless of what model expects.
|
||||
{
|
||||
gpu_data_in_->elements *= tensor->dims->data[0]; // batch
|
||||
gpu_data_in_->elements *= tensor->dims->data[1]; // height
|
||||
gpu_data_in_->elements *= tensor->dims->data[2]; // width
|
||||
gpu_data_in_->elements *= 4; // channels
|
||||
gpu_data_in_.resize(input_indices.size());
|
||||
for (int i = 0; i < input_indices.size(); ++i) {
|
||||
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[i]);
|
||||
gpu_data_in_[i] = absl::make_unique<GPUData>();
|
||||
gpu_data_in_[i]->shape.b = tensor->dims->data[0];
|
||||
gpu_data_in_[i]->shape.h = tensor->dims->data[1];
|
||||
gpu_data_in_[i]->shape.w = tensor->dims->data[2];
|
||||
// On iOS GPU, input must be 4 channels, regardless of what model expects.
|
||||
gpu_data_in_[i]->shape.c = 4;
|
||||
gpu_data_in_[i]->elements =
|
||||
gpu_data_in_[i]->shape.b * gpu_data_in_[i]->shape.h *
|
||||
gpu_data_in_[i]->shape.w * gpu_data_in_[i]->shape.c;
|
||||
// Input to model can be RGBA only.
|
||||
if (tensor->dims->data[3] != 4) {
|
||||
LOG(WARNING) << "Please ensure input GPU tensor is 4 channels.";
|
||||
}
|
||||
const std::string shader_source =
|
||||
absl::Substitute(R"(#include <metal_stdlib>
|
||||
using namespace metal;
|
||||
kernel void convertKernel(device float4* const input_buffer [[buffer(0)]],
|
||||
device half4* output_buffer [[buffer(1)]],
|
||||
uint gid [[thread_position_in_grid]]) {
|
||||
if (gid >= $0) return;
|
||||
output_buffer[gid] = half4(input_buffer[gid]);
|
||||
})",
|
||||
gpu_data_in_[i]->elements / 4);
|
||||
NSString* library_source =
|
||||
[NSString stringWithUTF8String:shader_source.c_str()];
|
||||
NSError* error = nil;
|
||||
id<MTLLibrary> library =
|
||||
[device newLibraryWithSource:library_source options:nil error:&error];
|
||||
RET_CHECK(library != nil) << "Couldn't create shader library "
|
||||
<< [[error localizedDescription] UTF8String];
|
||||
id<MTLFunction> kernel_func = nil;
|
||||
kernel_func = [library newFunctionWithName:@"convertKernel"];
|
||||
RET_CHECK(kernel_func != nil) << "Couldn't create kernel function.";
|
||||
fp32_to_fp16_program_ =
|
||||
[device newComputePipelineStateWithFunction:kernel_func error:&error];
|
||||
RET_CHECK(fp32_to_fp16_program_ != nil)
|
||||
<< "Couldn't create pipeline state "
|
||||
<< [[error localizedDescription] UTF8String];
|
||||
|
||||
// Create and bind input buffer.
|
||||
gpu_data_in_[i]->buffer =
|
||||
[device newBufferWithLength:gpu_data_in_[i]->elements * kHalfSize
|
||||
options:MTLResourceStorageModeShared];
|
||||
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
|
||||
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
|
||||
delegate_, input_indices[i], gpu_data_in_[i]->buffer),
|
||||
true);
|
||||
}
|
||||
// Input to model can be RGBA only.
|
||||
if (tensor->dims->data[3] != 4) {
|
||||
LOG(WARNING) << "Please ensure input GPU tensor is 4 channels.";
|
||||
}
|
||||
// Create and bind input buffer.
|
||||
gpu_data_in_->buffer =
|
||||
[device newBufferWithLength:gpu_data_in_->elements * sizeof(float)
|
||||
options:MTLResourceStorageModeShared];
|
||||
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
|
||||
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
|
||||
delegate_,
|
||||
input_indices[0], // First tensor only
|
||||
gpu_data_in_->buffer),
|
||||
true);
|
||||
}
|
||||
if (gpu_output_) {
|
||||
// Get output image sizes.
|
||||
@@ -607,15 +723,16 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
interpreter_->SetAllowBufferHandleOutput(true);
|
||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||
gpu_data_out_[i]->buffer =
|
||||
[device newBufferWithLength:gpu_data_out_[i]->elements * sizeof(float)
|
||||
[device newBufferWithLength:gpu_data_out_[i]->elements * kHalfSize
|
||||
options:MTLResourceStorageModeShared];
|
||||
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
|
||||
delegate_, output_indices[i], gpu_data_out_[i]->buffer),
|
||||
true);
|
||||
}
|
||||
|
||||
// Create converter for GPU output.
|
||||
converter_from_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
|
||||
isFloat16:false
|
||||
isFloat16:true
|
||||
convertToPBHWC4:false];
|
||||
if (converter_from_BPHWC4_ == nil) {
|
||||
return mediapipe::InternalError(
|
||||
|
||||
@@ -45,4 +45,9 @@ message TfLiteInferenceCalculatorOptions {
|
||||
// input tensors are on CPU. For input tensors on GPU, GPU backend is always
|
||||
// used.
|
||||
optional bool use_gpu = 2 [default = false];
|
||||
|
||||
// Android only. When true, an NNAPI delegate will be used for inference.
|
||||
// If NNAPI is not available, then the default CPU delegate will be used
|
||||
// automatically.
|
||||
optional bool use_nnapi = 3 [default = false];
|
||||
}
|
||||
|
||||
@@ -24,7 +24,7 @@
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
||||
#if defined(MEDIAPIPE_MOBILE)
|
||||
#include "mediapipe/util/android/file/base/file.h"
|
||||
#include "mediapipe/util/android/file/base/helpers.h"
|
||||
#else
|
||||
@@ -66,8 +66,8 @@ class TfLiteTensorsToClassificationCalculator : public CalculatorBase {
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
::mediapipe::TfLiteTensorsToClassificationCalculatorOptions options_;
|
||||
int top_k_ = 0;
|
||||
double min_score_threshold_ = 0;
|
||||
std::unordered_map<int, std::string> label_map_;
|
||||
bool label_map_loaded_ = false;
|
||||
};
|
||||
@@ -93,15 +93,14 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
||||
CalculatorContext* cc) {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
auto options = cc->Options<
|
||||
options_ = cc->Options<
|
||||
::mediapipe::TfLiteTensorsToClassificationCalculatorOptions>();
|
||||
|
||||
top_k_ = options.top_k();
|
||||
min_score_threshold_ = options.min_score_threshold();
|
||||
if (options.has_label_map_path()) {
|
||||
top_k_ = options_.top_k();
|
||||
if (options_.has_label_map_path()) {
|
||||
std::string string_path;
|
||||
ASSIGN_OR_RETURN(string_path,
|
||||
PathToResourceAsFile(options.label_map_path()));
|
||||
PathToResourceAsFile(options_.label_map_path()));
|
||||
std::string label_map_string;
|
||||
MP_RETURN_IF_ERROR(file::GetContents(string_path, &label_map_string));
|
||||
|
||||
@@ -125,9 +124,11 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
||||
RET_CHECK_EQ(input_tensors.size(), 1);
|
||||
|
||||
const TfLiteTensor* raw_score_tensor = &input_tensors[0];
|
||||
RET_CHECK_EQ(raw_score_tensor->dims->size, 2);
|
||||
RET_CHECK_EQ(raw_score_tensor->dims->data[0], 1);
|
||||
int num_classes = raw_score_tensor->dims->data[1];
|
||||
int num_classes = 1;
|
||||
for (int i = 0; i < raw_score_tensor->dims->size; ++i) {
|
||||
num_classes *= raw_score_tensor->dims->data[i];
|
||||
}
|
||||
|
||||
if (label_map_loaded_) {
|
||||
RET_CHECK_EQ(num_classes, label_map_.size());
|
||||
}
|
||||
@@ -135,7 +136,8 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
||||
|
||||
auto classification_list = absl::make_unique<ClassificationList>();
|
||||
for (int i = 0; i < num_classes; ++i) {
|
||||
if (raw_scores[i] < min_score_threshold_) {
|
||||
if (options_.has_min_score_threshold() &&
|
||||
raw_scores[i] < options_.min_score_threshold()) {
|
||||
continue;
|
||||
}
|
||||
Classification* classification = classification_list->add_classification();
|
||||
@@ -148,6 +150,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
||||
|
||||
// Note that partial_sort will raise error when top_k_ >
|
||||
// classification_list->classification_size().
|
||||
CHECK_GE(classification_list->classification_size(), top_k_);
|
||||
auto raw_classification_list = classification_list->mutable_classification();
|
||||
if (top_k_ > 0 && classification_list->classification_size() >= top_k_) {
|
||||
std::partial_sort(raw_classification_list->begin(),
|
||||
|
||||
@@ -27,7 +27,7 @@
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
|
||||
@@ -35,7 +35,7 @@
|
||||
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#if defined(MEDIAPIPE_IOS)
|
||||
#import <CoreVideo/CoreVideo.h>
|
||||
#import <Metal/Metal.h>
|
||||
#import <MetalKit/MetalKit.h>
|
||||
@@ -55,22 +55,22 @@ constexpr int kNumCoordsPerBox = 4;
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
||||
using ::tflite::gpu::gl::GlShader;
|
||||
#endif
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
|
||||
typedef ::tflite::gpu::gl::GlProgram GpuProgram;
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
typedef id<MTLBuffer> GpuTensor;
|
||||
typedef id<MTLComputePipelineState> GpuProgram;
|
||||
#endif
|
||||
|
||||
namespace {
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
struct GPUData {
|
||||
GpuProgram decode_program;
|
||||
GpuProgram score_program;
|
||||
@@ -180,10 +180,10 @@ class TfLiteTensorsToDetectionsCalculator : public CalculatorBase {
|
||||
std::vector<Anchor> anchors_;
|
||||
bool side_packet_anchors_{};
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||
std::unique_ptr<GPUData> gpu_data_;
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
MPPMetalHelper* gpu_helper_ = nullptr;
|
||||
std::unique_ptr<GPUData> gpu_data_;
|
||||
#endif
|
||||
@@ -204,7 +204,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||
}
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||
use_gpu |= true;
|
||||
@@ -222,9 +222,9 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
}
|
||||
|
||||
if (use_gpu) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
|
||||
#endif
|
||||
}
|
||||
@@ -238,9 +238,9 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
gpu_input_ = true;
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
|
||||
RET_CHECK(gpu_helper_);
|
||||
#endif
|
||||
@@ -400,7 +400,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
}
|
||||
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::ProcessGPU(
|
||||
CalculatorContext* cc, std::vector<Detection>* output_detections) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||
RET_CHECK_GE(input_tensors.size(), 2);
|
||||
@@ -463,7 +463,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||
@@ -472,11 +472,11 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
// Copy inputs.
|
||||
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_boxes_buffer
|
||||
from:input_tensors[0]
|
||||
blocking:true
|
||||
blocking:false
|
||||
commandBuffer:[gpu_helper_ commandBuffer]];
|
||||
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_scores_buffer
|
||||
from:input_tensors[1]
|
||||
blocking:true
|
||||
blocking:false
|
||||
commandBuffer:[gpu_helper_ commandBuffer]];
|
||||
if (!anchors_init_) {
|
||||
if (side_packet_anchors_) {
|
||||
@@ -491,48 +491,37 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
RET_CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
|
||||
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_anchors_buffer
|
||||
from:input_tensors[2]
|
||||
blocking:true
|
||||
blocking:false
|
||||
commandBuffer:[gpu_helper_ commandBuffer]];
|
||||
}
|
||||
anchors_init_ = true;
|
||||
}
|
||||
|
||||
// Run shaders.
|
||||
{
|
||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||
command_buffer.label = @"TfLiteDecodeBoxes";
|
||||
id<MTLComputeCommandEncoder> decode_command =
|
||||
[command_buffer computeCommandEncoder];
|
||||
[decode_command setComputePipelineState:gpu_data_->decode_program];
|
||||
[decode_command setBuffer:gpu_data_->decoded_boxes_buffer
|
||||
offset:0
|
||||
atIndex:0];
|
||||
[decode_command setBuffer:gpu_data_->raw_boxes_buffer offset:0 atIndex:1];
|
||||
[decode_command setBuffer:gpu_data_->raw_anchors_buffer offset:0 atIndex:2];
|
||||
MTLSize decode_threads_per_group = MTLSizeMake(1, 1, 1);
|
||||
MTLSize decode_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
|
||||
[decode_command dispatchThreadgroups:decode_threadgroups
|
||||
threadsPerThreadgroup:decode_threads_per_group];
|
||||
[decode_command endEncoding];
|
||||
[command_buffer commit];
|
||||
[command_buffer waitUntilCompleted];
|
||||
}
|
||||
{
|
||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||
command_buffer.label = @"TfLiteScoreBoxes";
|
||||
id<MTLComputeCommandEncoder> score_command =
|
||||
[command_buffer computeCommandEncoder];
|
||||
[score_command setComputePipelineState:gpu_data_->score_program];
|
||||
[score_command setBuffer:gpu_data_->scored_boxes_buffer offset:0 atIndex:0];
|
||||
[score_command setBuffer:gpu_data_->raw_scores_buffer offset:0 atIndex:1];
|
||||
MTLSize score_threads_per_group = MTLSizeMake(1, num_classes_, 1);
|
||||
MTLSize score_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
|
||||
[score_command dispatchThreadgroups:score_threadgroups
|
||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||
command_buffer.label = @"TfLiteDecodeAndScoreBoxes";
|
||||
id<MTLComputeCommandEncoder> command_encoder =
|
||||
[command_buffer computeCommandEncoder];
|
||||
[command_encoder setComputePipelineState:gpu_data_->decode_program];
|
||||
[command_encoder setBuffer:gpu_data_->decoded_boxes_buffer
|
||||
offset:0
|
||||
atIndex:0];
|
||||
[command_encoder setBuffer:gpu_data_->raw_boxes_buffer offset:0 atIndex:1];
|
||||
[command_encoder setBuffer:gpu_data_->raw_anchors_buffer offset:0 atIndex:2];
|
||||
MTLSize decode_threads_per_group = MTLSizeMake(1, 1, 1);
|
||||
MTLSize decode_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
|
||||
[command_encoder dispatchThreadgroups:decode_threadgroups
|
||||
threadsPerThreadgroup:decode_threads_per_group];
|
||||
|
||||
[command_encoder setComputePipelineState:gpu_data_->score_program];
|
||||
[command_encoder setBuffer:gpu_data_->scored_boxes_buffer offset:0 atIndex:0];
|
||||
[command_encoder setBuffer:gpu_data_->raw_scores_buffer offset:0 atIndex:1];
|
||||
MTLSize score_threads_per_group = MTLSizeMake(1, num_classes_, 1);
|
||||
MTLSize score_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
|
||||
[command_encoder dispatchThreadgroups:score_threadgroups
|
||||
threadsPerThreadgroup:score_threads_per_group];
|
||||
[score_command endEncoding];
|
||||
[command_buffer commit];
|
||||
[command_buffer waitUntilCompleted];
|
||||
}
|
||||
[command_encoder endEncoding];
|
||||
[MPPMetalUtil commitCommandBufferAndWait:command_buffer];
|
||||
|
||||
// Copy decoded boxes from GPU to CPU.
|
||||
std::vector<float> boxes(num_boxes_ * num_coords_);
|
||||
@@ -562,11 +551,11 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::Close(
|
||||
CalculatorContext* cc) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
gpu_helper_.RunInGlContext([this] { gpu_data_.reset(); });
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
gpu_data_.reset();
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
#endif
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -715,7 +704,7 @@ Detection TfLiteTensorsToDetectionsCalculator::ConvertToDetection(
|
||||
|
||||
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::GpuInit(
|
||||
CalculatorContext* cc) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]()
|
||||
-> ::mediapipe::Status {
|
||||
gpu_data_ = absl::make_unique<GPUData>();
|
||||
@@ -928,8 +917,7 @@ void main() {
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
|
||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
||||
// TODO consolidate Metal and OpenGL shaders via vulkan.
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
|
||||
gpu_data_ = absl::make_unique<GPUData>();
|
||||
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
||||
@@ -1159,7 +1147,7 @@ kernel void scoreKernel(
|
||||
CHECK_LT(num_classes_, max_wg_size) << "# classes must be <" << max_wg_size;
|
||||
}
|
||||
|
||||
#endif // __ANDROID__ or iOS
|
||||
#endif // !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -21,7 +21,8 @@
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator for converting TFLite tensors from regression models into
|
||||
// landmarks.
|
||||
// landmarks. Note that if the landmarks in the tensor has more than 3
|
||||
// dimensions, only the first 3 dimensions will be converted to x,y,z.
|
||||
//
|
||||
// Input:
|
||||
// TENSORS - Vector of TfLiteTensor of type kTfLiteFloat32. Only the first
|
||||
@@ -75,11 +76,11 @@ REGISTER_CALCULATOR(TfLiteTensorsToLandmarksCalculator);
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("LANDMARKS")) {
|
||||
cc->Outputs().Tag("LANDMARKS").Set<std::vector<Landmark>>();
|
||||
cc->Outputs().Tag("LANDMARKS").Set<LandmarkList>();
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("NORM_LANDMARKS")) {
|
||||
cc->Outputs().Tag("NORM_LANDMARKS").Set<std::vector<NormalizedLandmark>>();
|
||||
cc->Outputs().Tag("NORM_LANDMARKS").Set<NormalizedLandmarkList>();
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
@@ -122,61 +123,59 @@ REGISTER_CALCULATOR(TfLiteTensorsToLandmarksCalculator);
|
||||
num_values *= raw_tensor->dims->data[i];
|
||||
}
|
||||
const int num_dimensions = num_values / num_landmarks_;
|
||||
// Landmarks must have less than 3 dimensions. Otherwise please consider
|
||||
// using matrix.
|
||||
CHECK_LE(num_dimensions, 3);
|
||||
CHECK_GT(num_dimensions, 0);
|
||||
|
||||
const float* raw_landmarks = raw_tensor->data.f;
|
||||
|
||||
auto output_landmarks = absl::make_unique<std::vector<Landmark>>();
|
||||
LandmarkList output_landmarks;
|
||||
|
||||
for (int ld = 0; ld < num_landmarks_; ++ld) {
|
||||
const int offset = ld * num_dimensions;
|
||||
Landmark landmark;
|
||||
Landmark* landmark = output_landmarks.add_landmark();
|
||||
|
||||
if (options_.flip_horizontally()) {
|
||||
landmark.set_x(options_.input_image_width() - raw_landmarks[offset]);
|
||||
landmark->set_x(options_.input_image_width() - raw_landmarks[offset]);
|
||||
} else {
|
||||
landmark.set_x(raw_landmarks[offset]);
|
||||
landmark->set_x(raw_landmarks[offset]);
|
||||
}
|
||||
if (num_dimensions > 1) {
|
||||
if (options_.flip_vertically()) {
|
||||
landmark.set_y(options_.input_image_height() -
|
||||
raw_landmarks[offset + 1]);
|
||||
landmark->set_y(options_.input_image_height() -
|
||||
raw_landmarks[offset + 1]);
|
||||
} else {
|
||||
landmark.set_y(raw_landmarks[offset + 1]);
|
||||
landmark->set_y(raw_landmarks[offset + 1]);
|
||||
}
|
||||
}
|
||||
if (num_dimensions > 2) {
|
||||
landmark.set_z(raw_landmarks[offset + 2]);
|
||||
landmark->set_z(raw_landmarks[offset + 2]);
|
||||
}
|
||||
output_landmarks->push_back(landmark);
|
||||
}
|
||||
|
||||
// Output normalized landmarks if required.
|
||||
if (cc->Outputs().HasTag("NORM_LANDMARKS")) {
|
||||
auto output_norm_landmarks =
|
||||
absl::make_unique<std::vector<NormalizedLandmark>>();
|
||||
for (const auto& landmark : *output_landmarks) {
|
||||
NormalizedLandmark norm_landmark;
|
||||
norm_landmark.set_x(static_cast<float>(landmark.x()) /
|
||||
options_.input_image_width());
|
||||
norm_landmark.set_y(static_cast<float>(landmark.y()) /
|
||||
options_.input_image_height());
|
||||
norm_landmark.set_z(landmark.z() / options_.normalize_z());
|
||||
|
||||
output_norm_landmarks->push_back(norm_landmark);
|
||||
NormalizedLandmarkList output_norm_landmarks;
|
||||
// for (const auto& landmark : output_landmarks) {
|
||||
for (int i = 0; i < output_landmarks.landmark_size(); ++i) {
|
||||
const Landmark& landmark = output_landmarks.landmark(i);
|
||||
NormalizedLandmark* norm_landmark = output_norm_landmarks.add_landmark();
|
||||
norm_landmark->set_x(static_cast<float>(landmark.x()) /
|
||||
options_.input_image_width());
|
||||
norm_landmark->set_y(static_cast<float>(landmark.y()) /
|
||||
options_.input_image_height());
|
||||
norm_landmark->set_z(landmark.z() / options_.normalize_z());
|
||||
}
|
||||
cc->Outputs()
|
||||
.Tag("NORM_LANDMARKS")
|
||||
.Add(output_norm_landmarks.release(), cc->InputTimestamp());
|
||||
.AddPacket(MakePacket<NormalizedLandmarkList>(output_norm_landmarks)
|
||||
.At(cc->InputTimestamp()));
|
||||
}
|
||||
|
||||
// Output absolute landmarks.
|
||||
if (cc->Outputs().HasTag("LANDMARKS")) {
|
||||
cc->Outputs()
|
||||
.Tag("LANDMARKS")
|
||||
.Add(output_landmarks.release(), cc->InputTimestamp());
|
||||
.AddPacket(MakePacket<LandmarkList>(output_landmarks)
|
||||
.At(cc->InputTimestamp()));
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
|
||||
@@ -28,7 +28,7 @@
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||
#include "mediapipe/gpu/gl_simple_shaders.h"
|
||||
#include "mediapipe/gpu/shader_util.h"
|
||||
@@ -53,7 +53,7 @@ float Clamp(float val, float min, float max) {
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
using ::tflite::gpu::gl::CopyBuffer;
|
||||
using ::tflite::gpu::gl::CreateReadWriteRgbaImageTexture;
|
||||
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
||||
@@ -129,7 +129,7 @@ class TfLiteTensorsToSegmentationCalculator : public CalculatorBase {
|
||||
int tensor_channels_ = 0;
|
||||
|
||||
bool use_gpu_ = false;
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||
std::unique_ptr<GlProgram> mask_program_with_prev_;
|
||||
std::unique_ptr<GlProgram> mask_program_no_prev_;
|
||||
@@ -159,7 +159,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
}
|
||||
|
||||
// Inputs GPU.
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GlBuffer>>();
|
||||
use_gpu |= true;
|
||||
@@ -178,7 +178,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
if (cc->Outputs().HasTag("MASK")) {
|
||||
cc->Outputs().Tag("MASK").Set<ImageFrame>();
|
||||
}
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
if (cc->Outputs().HasTag("MASK_GPU")) {
|
||||
cc->Outputs().Tag("MASK_GPU").Set<mediapipe::GpuBuffer>();
|
||||
use_gpu |= true;
|
||||
@@ -186,7 +186,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
if (use_gpu) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
}
|
||||
@@ -199,7 +199,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
use_gpu_ = true;
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
}
|
||||
@@ -207,7 +207,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
MP_RETURN_IF_ERROR(LoadOptions(cc));
|
||||
|
||||
if (use_gpu_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(
|
||||
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
|
||||
MP_RETURN_IF_ERROR(InitGpu(cc));
|
||||
@@ -224,7 +224,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
if (use_gpu_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(
|
||||
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
|
||||
MP_RETURN_IF_ERROR(ProcessGpu(cc));
|
||||
@@ -240,7 +240,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::Close(
|
||||
CalculatorContext* cc) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
gpu_helper_.RunInGlContext([this] {
|
||||
if (upsample_program_) glDeleteProgram(upsample_program_);
|
||||
upsample_program_ = 0;
|
||||
@@ -367,7 +367,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
if (cc->Inputs().Tag("TENSORS_GPU").IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
// Get input streams.
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GlBuffer>>();
|
||||
@@ -453,7 +453,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
}
|
||||
|
||||
void TfLiteTensorsToSegmentationCalculator::GlRender() {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
static const GLfloat square_vertices[] = {
|
||||
-1.0f, -1.0f, // bottom left
|
||||
1.0f, -1.0f, // bottom right
|
||||
@@ -525,7 +525,7 @@ void TfLiteTensorsToSegmentationCalculator::GlRender() {
|
||||
|
||||
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::InitGpu(
|
||||
CalculatorContext* cc) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__APPLE__)
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]()
|
||||
-> ::mediapipe::Status {
|
||||
// A shader to process a segmentation tensor into an output mask,
|
||||
|
||||
@@ -12,14 +12,14 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
package(default_visibility = ["//visibility:public"])
|
||||
|
||||
exports_files(["LICENSE"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "annotation_overlay_calculator_proto",
|
||||
srcs = ["annotation_overlay_calculator.proto"],
|
||||
@@ -72,6 +72,24 @@ proto_library(
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "collection_has_min_size_calculator_proto",
|
||||
srcs = ["collection_has_min_size_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "association_calculator_proto",
|
||||
srcs = ["association_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "annotation_overlay_calculator_cc_proto",
|
||||
srcs = ["annotation_overlay_calculator.proto"],
|
||||
@@ -141,6 +159,26 @@ mediapipe_cc_proto_library(
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "collection_has_min_size_calculator_cc_proto",
|
||||
srcs = ["collection_has_min_size_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
],
|
||||
visibility = ["//mediapipe:__subpackages__"],
|
||||
deps = [":collection_has_min_size_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "association_calculator_cc_proto",
|
||||
srcs = ["association_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
],
|
||||
visibility = ["//mediapipe:__subpackages__"],
|
||||
deps = [":association_calculator_proto"],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "packet_frequency_calculator",
|
||||
srcs = ["packet_frequency_calculator.cc"],
|
||||
@@ -234,6 +272,7 @@ cc_library(
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:vector",
|
||||
"//mediapipe/util:annotation_renderer",
|
||||
"//mediapipe/util:render_data_cc_proto",
|
||||
] + select({
|
||||
"//mediapipe/gpu:disable_gpu": [],
|
||||
"//conditions:default": [
|
||||
@@ -360,6 +399,16 @@ mediapipe_cc_proto_library(
|
||||
deps = [":landmark_projection_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "landmarks_to_floats_calculator_cc_proto",
|
||||
srcs = ["landmarks_to_floats_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":landmarks_to_floats_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "rect_transformation_calculator_cc_proto",
|
||||
srcs = ["rect_transformation_calculator.proto"],
|
||||
@@ -372,7 +421,12 @@ mediapipe_cc_proto_library(
|
||||
|
||||
cc_library(
|
||||
name = "detections_to_rects_calculator",
|
||||
srcs = ["detections_to_rects_calculator.cc"],
|
||||
srcs = [
|
||||
"detections_to_rects_calculator.cc",
|
||||
],
|
||||
hdrs = [
|
||||
"detections_to_rects_calculator.h",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":detections_to_rects_calculator_cc_proto",
|
||||
@@ -454,6 +508,17 @@ proto_library(
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "labels_to_render_data_calculator_proto",
|
||||
srcs = ["labels_to_render_data_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/util:color_proto",
|
||||
"//mediapipe/util:render_data_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "thresholding_calculator_proto",
|
||||
srcs = ["thresholding_calculator.proto"],
|
||||
@@ -483,6 +548,15 @@ proto_library(
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "landmarks_to_floats_calculator_proto",
|
||||
srcs = ["landmarks_to_floats_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "rect_transformation_calculator_proto",
|
||||
srcs = ["rect_transformation_calculator.proto"],
|
||||
@@ -577,6 +651,26 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "labels_to_render_data_calculator",
|
||||
srcs = ["labels_to_render_data_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":labels_to_render_data_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_options_cc_proto",
|
||||
"//mediapipe/framework/formats:classification_cc_proto",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:statusor",
|
||||
"//mediapipe/util:color_cc_proto",
|
||||
"//mediapipe/util:render_data_cc_proto",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "rect_to_render_data_calculator",
|
||||
srcs = ["rect_to_render_data_calculator.cc"],
|
||||
@@ -658,6 +752,22 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "landmarks_to_floats_calculator",
|
||||
srcs = ["landmarks_to_floats_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":landmarks_to_floats_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@eigen_archive//:eigen",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "detection_letterbox_removal_calculator_test",
|
||||
srcs = ["detection_letterbox_removal_calculator_test.cc"],
|
||||
@@ -714,6 +824,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":top_k_scores_calculator_cc_proto",
|
||||
"//mediapipe/framework/formats:classification_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:statusor",
|
||||
@@ -750,3 +861,153 @@ cc_test(
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "labels_to_render_data_calculator_cc_proto",
|
||||
srcs = ["labels_to_render_data_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/util:color_cc_proto",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":labels_to_render_data_calculator_proto"],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "local_file_contents_calculator",
|
||||
srcs = ["local_file_contents_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:file_helpers",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "filter_collection_calculator",
|
||||
srcs = ["filter_collection_calculator.cc"],
|
||||
hdrs = ["filter_collection_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "collection_has_min_size_calculator",
|
||||
srcs = ["collection_has_min_size_calculator.cc"],
|
||||
hdrs = ["collection_has_min_size_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":collection_has_min_size_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "association_calculator",
|
||||
hdrs = ["association_calculator.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":association_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework/port:rectangle",
|
||||
"//mediapipe/framework/port:status",
|
||||
"@com_google_absl//absl/memory",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "association_norm_rect_calculator",
|
||||
srcs = ["association_norm_rect_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":association_calculator",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:rectangle",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "association_detection_calculator",
|
||||
srcs = ["association_detection_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":association_calculator",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/formats:location",
|
||||
"//mediapipe/framework/port:rectangle",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "association_calculator_test",
|
||||
srcs = ["association_calculator_test.cc"],
|
||||
deps = [
|
||||
":association_detection_calculator",
|
||||
":association_norm_rect_calculator",
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/deps:message_matchers",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/formats:location_data_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "detections_to_timed_box_list_calculator",
|
||||
srcs = ["detections_to_timed_box_list_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/formats:location_data_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/util/tracking:box_tracker",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "detection_unique_id_calculator",
|
||||
srcs = ["detection_unique_id_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
@@ -26,6 +26,7 @@
|
||||
#include "mediapipe/framework/port/vector.h"
|
||||
#include "mediapipe/util/annotation_renderer.h"
|
||||
#include "mediapipe/util/color.pb.h"
|
||||
#include "mediapipe/util/render_data.pb.h"
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||
@@ -41,6 +42,8 @@ namespace {
|
||||
constexpr char kInputFrameTag[] = "INPUT_FRAME";
|
||||
constexpr char kOutputFrameTag[] = "OUTPUT_FRAME";
|
||||
|
||||
constexpr char kInputVectorTag[] = "VECTOR";
|
||||
|
||||
constexpr char kInputFrameTagGpu[] = "INPUT_FRAME_GPU";
|
||||
constexpr char kOutputFrameTagGpu[] = "OUTPUT_FRAME_GPU";
|
||||
|
||||
@@ -65,6 +68,9 @@ constexpr int kAnnotationBackgroundColor[] = {100, 101, 102};
|
||||
// 2. RenderData proto on variable number of input streams. All the RenderData
|
||||
// at a particular timestamp is drawn on the image in the order of their
|
||||
// input streams. No tags required.
|
||||
// 3. std::vector<RenderData> on variable number of input streams. RenderData
|
||||
// objects at a particular timestamp are drawn on the image in order of the
|
||||
// input vector items. These input streams are tagged with "VECTOR".
|
||||
//
|
||||
// Output:
|
||||
// 1. OUTPUT_FRAME or OUTPUT_FRAME_GPU: A rendered ImageFrame (or GpuBuffer).
|
||||
@@ -85,6 +91,8 @@ constexpr int kAnnotationBackgroundColor[] = {100, 101, 102};
|
||||
// input_stream: "render_data_1"
|
||||
// input_stream: "render_data_2"
|
||||
// input_stream: "render_data_3"
|
||||
// input_stream: "VECTOR:0:render_data_vec_0"
|
||||
// input_stream: "VECTOR:1:render_data_vec_1"
|
||||
// output_stream: "OUTPUT_FRAME:decorated_frames"
|
||||
// options {
|
||||
// [mediapipe.AnnotationOverlayCalculatorOptions.ext] {
|
||||
@@ -99,6 +107,8 @@ constexpr int kAnnotationBackgroundColor[] = {100, 101, 102};
|
||||
// input_stream: "render_data_1"
|
||||
// input_stream: "render_data_2"
|
||||
// input_stream: "render_data_3"
|
||||
// input_stream: "VECTOR:0:render_data_vec_0"
|
||||
// input_stream: "VECTOR:1:render_data_vec_1"
|
||||
// output_stream: "OUTPUT_FRAME_GPU:decorated_frames"
|
||||
// options {
|
||||
// [mediapipe.AnnotationOverlayCalculatorOptions.ext] {
|
||||
@@ -138,9 +148,6 @@ class AnnotationOverlayCalculator : public CalculatorBase {
|
||||
// Underlying helper renderer library.
|
||||
std::unique_ptr<AnnotationRenderer> renderer_;
|
||||
|
||||
// Number of input streams with render data.
|
||||
int num_render_streams_;
|
||||
|
||||
// Indicates if image frame is available as input.
|
||||
bool image_frame_available_ = false;
|
||||
|
||||
@@ -171,25 +178,28 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
|
||||
return ::mediapipe::InternalError("GPU output must have GPU input.");
|
||||
}
|
||||
|
||||
// Assume all inputs are render streams; adjust below.
|
||||
int num_render_streams = cc->Inputs().NumEntries();
|
||||
|
||||
// Input image to render onto copy of.
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
if (cc->Inputs().HasTag(kInputFrameTagGpu)) {
|
||||
cc->Inputs().Tag(kInputFrameTagGpu).Set<mediapipe::GpuBuffer>();
|
||||
num_render_streams = cc->Inputs().NumEntries() - 1;
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
if (cc->Inputs().HasTag(kInputFrameTag)) {
|
||||
cc->Inputs().Tag(kInputFrameTag).Set<ImageFrame>();
|
||||
num_render_streams = cc->Inputs().NumEntries() - 1;
|
||||
}
|
||||
|
||||
// Data streams to render.
|
||||
for (int i = 0; i < num_render_streams; ++i) {
|
||||
cc->Inputs().Index(i).Set<RenderData>();
|
||||
for (CollectionItemId id = cc->Inputs().BeginId(); id < cc->Inputs().EndId();
|
||||
++id) {
|
||||
auto tag_and_index = cc->Inputs().TagAndIndexFromId(id);
|
||||
std::string tag = tag_and_index.first;
|
||||
if (tag == kInputVectorTag) {
|
||||
cc->Inputs().Get(id).Set<std::vector<RenderData>>();
|
||||
} else if (tag.empty()) {
|
||||
// Empty tag defaults to accepting a single object of RenderData type.
|
||||
cc->Inputs().Get(id).Set<RenderData>();
|
||||
}
|
||||
}
|
||||
|
||||
// Rendered image.
|
||||
@@ -228,12 +238,10 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
|
||||
if (cc->Inputs().HasTag(kInputFrameTagGpu) ||
|
||||
cc->Inputs().HasTag(kInputFrameTag)) {
|
||||
image_frame_available_ = true;
|
||||
num_render_streams_ = cc->Inputs().NumEntries() - 1;
|
||||
} else {
|
||||
image_frame_available_ = false;
|
||||
RET_CHECK(options_.has_canvas_width_px());
|
||||
RET_CHECK(options_.has_canvas_height_px());
|
||||
num_render_streams_ = cc->Inputs().NumEntries();
|
||||
}
|
||||
|
||||
// Initialize the helper renderer library.
|
||||
@@ -285,12 +293,28 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
|
||||
renderer_->AdoptImage(image_mat.get());
|
||||
|
||||
// Render streams onto render target.
|
||||
for (int i = 0; i < num_render_streams_; ++i) {
|
||||
if (cc->Inputs().Index(i).IsEmpty()) {
|
||||
for (CollectionItemId id = cc->Inputs().BeginId(); id < cc->Inputs().EndId();
|
||||
++id) {
|
||||
auto tag_and_index = cc->Inputs().TagAndIndexFromId(id);
|
||||
std::string tag = tag_and_index.first;
|
||||
if (!tag.empty() && tag != kInputVectorTag) {
|
||||
continue;
|
||||
}
|
||||
const RenderData& render_data = cc->Inputs().Index(i).Get<RenderData>();
|
||||
renderer_->RenderDataOnImage(render_data);
|
||||
if (cc->Inputs().Get(id).IsEmpty()) {
|
||||
continue;
|
||||
}
|
||||
if (tag.empty()) {
|
||||
// Empty tag defaults to accepting a single object of RenderData type.
|
||||
const RenderData& render_data = cc->Inputs().Get(id).Get<RenderData>();
|
||||
renderer_->RenderDataOnImage(render_data);
|
||||
} else {
|
||||
RET_CHECK_EQ(kInputVectorTag, tag);
|
||||
const std::vector<RenderData>& render_data_vec =
|
||||
cc->Inputs().Get(id).Get<std::vector<RenderData>>();
|
||||
for (const RenderData& render_data : render_data_vec) {
|
||||
renderer_->RenderDataOnImage(render_data);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (use_gpu_) {
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
// Copyright 2019 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_ASSOCIATION_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_UTIL_ASSOCIATION_CALCULATOR_H_
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "mediapipe/calculators/util/association_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/rectangle.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Computes the overlap similarity based on Intersection over Union (IoU) of
|
||||
// two rectangles.
|
||||
inline float OverlapSimilarity(const Rectangle_f& rect1,
|
||||
const Rectangle_f& rect2) {
|
||||
if (!rect1.Intersects(rect2)) return 0.0f;
|
||||
// Compute IoU similarity score.
|
||||
const float intersection_area = Rectangle_f(rect1).Intersect(rect2).Area();
|
||||
const float normalization = rect1.Area() + rect2.Area() - intersection_area;
|
||||
return normalization > 0.0f ? intersection_area / normalization : 0.0f;
|
||||
}
|
||||
|
||||
// AssocationCalculator<T> accepts multiple inputs of vectors of type T that can
|
||||
// be converted to Rectangle_f. The output is a vector of type T that contains
|
||||
// elements from the input vectors that don't overlap with each other. When
|
||||
// two elements overlap, the element that comes in from a later input stream
|
||||
// is kept in the output. This association operation is useful for multiple
|
||||
// instance inference pipelines in MediaPipe.
|
||||
// If an input stream is tagged with "PREV" tag, IDs of overlapping elements
|
||||
// from "PREV" input stream are propagated to the output. Elements in the "PREV"
|
||||
// input stream that don't overlap with other elements are not added to the
|
||||
// output. This stream is designed to take detections from previous timestamp,
|
||||
// e.g. output of PreviousLoopbackCalculator to provide temporal association.
|
||||
// See AssociationDetectionCalculator and AssociationNormRectCalculator for
|
||||
// example uses.
|
||||
template <typename T>
|
||||
class AssociationCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
// Atmost one input stream can be tagged with "PREV".
|
||||
RET_CHECK_LE(cc->Inputs().NumEntries("PREV"), 1);
|
||||
|
||||
if (cc->Inputs().HasTag("PREV")) {
|
||||
RET_CHECK_GE(cc->Inputs().NumEntries(), 2);
|
||||
}
|
||||
|
||||
for (CollectionItemId id = cc->Inputs().BeginId();
|
||||
id < cc->Inputs().EndId(); ++id) {
|
||||
cc->Inputs().Get(id).Set<std::vector<T>>();
|
||||
}
|
||||
|
||||
cc->Outputs().Index(0).Set<std::vector<T>>();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
has_prev_input_stream_ = cc->Inputs().HasTag("PREV");
|
||||
if (has_prev_input_stream_) {
|
||||
prev_input_stream_id_ = cc->Inputs().GetId("PREV", 0);
|
||||
}
|
||||
options_ = cc->Options<::mediapipe::AssociationCalculatorOptions>();
|
||||
CHECK_GE(options_.min_similarity_threshold(), 0);
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
auto get_non_overlapping_elements = GetNonOverlappingElements(cc);
|
||||
if (!get_non_overlapping_elements.ok()) {
|
||||
return get_non_overlapping_elements.status();
|
||||
}
|
||||
std::list<T> result = get_non_overlapping_elements.ValueOrDie();
|
||||
|
||||
if (has_prev_input_stream_ &&
|
||||
!cc->Inputs().Get(prev_input_stream_id_).IsEmpty()) {
|
||||
// Processed all regular input streams. Now compare the result list
|
||||
// elements with those in the PREV input stream, and propagate IDs from
|
||||
// PREV input stream as appropriate.
|
||||
const std::vector<T>& prev_input_vec =
|
||||
cc->Inputs()
|
||||
.Get(prev_input_stream_id_)
|
||||
.template Get<std::vector<T>>();
|
||||
|
||||
MP_RETURN_IF_ERROR(
|
||||
PropagateIdsFromPreviousToCurrent(prev_input_vec, &result));
|
||||
}
|
||||
|
||||
auto output = absl::make_unique<std::vector<T>>();
|
||||
for (auto it = result.begin(); it != result.end(); ++it) {
|
||||
output->push_back(*it);
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
protected:
|
||||
::mediapipe::AssociationCalculatorOptions options_;
|
||||
|
||||
bool has_prev_input_stream_;
|
||||
CollectionItemId prev_input_stream_id_;
|
||||
|
||||
virtual ::mediapipe::StatusOr<Rectangle_f> GetRectangle(const T& input) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
virtual std::pair<bool, int> GetId(const T& input) { return {false, -1}; }
|
||||
|
||||
virtual void SetId(T* input, int id) {}
|
||||
|
||||
private:
|
||||
// Get a list of non-overlapping elements from all input streams, with
|
||||
// increasing order of priority based on input stream index.
|
||||
mediapipe::StatusOr<std::list<T>> GetNonOverlappingElements(
|
||||
CalculatorContext* cc) {
|
||||
std::list<T> result;
|
||||
|
||||
// Initialize result with the first non-empty input vector.
|
||||
CollectionItemId non_empty_id = cc->Inputs().BeginId();
|
||||
for (CollectionItemId id = cc->Inputs().BeginId();
|
||||
id < cc->Inputs().EndId(); ++id) {
|
||||
if (id == prev_input_stream_id_ || cc->Inputs().Get(id).IsEmpty()) {
|
||||
continue;
|
||||
}
|
||||
const std::vector<T>& input_vec =
|
||||
cc->Inputs().Get(id).Get<std::vector<T>>();
|
||||
if (!input_vec.empty()) {
|
||||
non_empty_id = id;
|
||||
result.push_back(input_vec[0]);
|
||||
for (int j = 1; j < input_vec.size(); ++j) {
|
||||
MP_RETURN_IF_ERROR(AddElementToList(input_vec[j], &result));
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Compare remaining input vectors with the non-empty result vector,
|
||||
// remove lower-priority overlapping elements from the result vector and
|
||||
// had corresponding higher-priority elements as necessary.
|
||||
for (CollectionItemId id = non_empty_id + 1; id < cc->Inputs().EndId();
|
||||
++id) {
|
||||
if (id == prev_input_stream_id_ || cc->Inputs().Get(id).IsEmpty()) {
|
||||
continue;
|
||||
}
|
||||
const std::vector<T>& input_vec =
|
||||
cc->Inputs().Get(id).Get<std::vector<T>>();
|
||||
|
||||
for (int vi = 0; vi < input_vec.size(); ++vi) {
|
||||
MP_RETURN_IF_ERROR(AddElementToList(input_vec[vi], &result));
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
::mediapipe::Status AddElementToList(T element, std::list<T>* current) {
|
||||
// Compare this element with elements of the input collection. If this
|
||||
// element has high overlap with elements of the collection, remove
|
||||
// those elements from the collection and add this element.
|
||||
ASSIGN_OR_RETURN(auto cur_rect, GetRectangle(element));
|
||||
|
||||
bool change_id = false;
|
||||
int new_elem_id = -1;
|
||||
|
||||
for (auto uit = current->begin(); uit != current->end();) {
|
||||
ASSIGN_OR_RETURN(auto prev_rect, GetRectangle(*uit));
|
||||
if (OverlapSimilarity(cur_rect, prev_rect) >
|
||||
options_.min_similarity_threshold()) {
|
||||
std::pair<bool, int> prev_id = GetId(*uit);
|
||||
// If prev_id.first is false when some element doesn't have an ID,
|
||||
// change_id and new_elem_id will not be updated.
|
||||
if (prev_id.first) {
|
||||
change_id = prev_id.first;
|
||||
new_elem_id = prev_id.second;
|
||||
}
|
||||
uit = current->erase(uit);
|
||||
} else {
|
||||
++uit;
|
||||
}
|
||||
}
|
||||
|
||||
if (change_id) {
|
||||
SetId(&element, new_elem_id);
|
||||
}
|
||||
current->push_back(element);
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
// Compare elements of the current list with elements in from the collection
|
||||
// of elements from the previous input stream, and propagate IDs from the
|
||||
// previous input stream as appropriate.
|
||||
::mediapipe::Status PropagateIdsFromPreviousToCurrent(
|
||||
const std::vector<T>& prev_input_vec, std::list<T>* current) {
|
||||
for (auto vit = current->begin(); vit != current->end(); ++vit) {
|
||||
auto get_cur_rectangle = GetRectangle(*vit);
|
||||
if (!get_cur_rectangle.ok()) {
|
||||
return get_cur_rectangle.status();
|
||||
}
|
||||
const Rectangle_f& cur_rect = get_cur_rectangle.ValueOrDie();
|
||||
|
||||
bool change_id = false;
|
||||
int id_for_vi = -1;
|
||||
|
||||
for (int ui = 0; ui < prev_input_vec.size(); ++ui) {
|
||||
auto get_prev_rectangle = GetRectangle(prev_input_vec[ui]);
|
||||
if (!get_prev_rectangle.ok()) {
|
||||
return get_prev_rectangle.status();
|
||||
}
|
||||
const Rectangle_f& prev_rect = get_prev_rectangle.ValueOrDie();
|
||||
|
||||
if (OverlapSimilarity(cur_rect, prev_rect) >
|
||||
options_.min_similarity_threshold()) {
|
||||
std::pair<bool, int> prev_id = GetId(prev_input_vec[ui]);
|
||||
// If prev_id.first is false when some element doesn't have an ID,
|
||||
// change_id and id_for_vi will not be updated.
|
||||
if (prev_id.first) {
|
||||
change_id = prev_id.first;
|
||||
id_for_vi = prev_id.second;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (change_id) {
|
||||
T element = *vit;
|
||||
SetId(&element, id_for_vi);
|
||||
*vit = element;
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_UTIL_ASSOCIATION_CALCULATOR_H_
|
||||
@@ -0,0 +1,27 @@
|
||||
// Copyright 2019 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;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
|
||||
message AssociationCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional AssociationCalculatorOptions ext = 275124847;
|
||||
}
|
||||
|
||||
optional float min_similarity_threshold = 1 [default = 1.0];
|
||||
}
|
||||
@@ -0,0 +1,476 @@
|
||||
// Copyright 2019 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/framework/calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/deps/message_matchers.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/location_data.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
::mediapipe::Detection DetectionWithRelativeLocationData(double xmin,
|
||||
double ymin,
|
||||
double width,
|
||||
double height) {
|
||||
::mediapipe::Detection detection;
|
||||
::mediapipe::LocationData* location_data = detection.mutable_location_data();
|
||||
location_data->set_format(::mediapipe::LocationData::RELATIVE_BOUNDING_BOX);
|
||||
location_data->mutable_relative_bounding_box()->set_xmin(xmin);
|
||||
location_data->mutable_relative_bounding_box()->set_ymin(ymin);
|
||||
location_data->mutable_relative_bounding_box()->set_width(width);
|
||||
location_data->mutable_relative_bounding_box()->set_height(height);
|
||||
return detection;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
class AssociationDetectionCalculatorTest : public ::testing::Test {
|
||||
protected:
|
||||
AssociationDetectionCalculatorTest() {
|
||||
// 0.4 ================
|
||||
// | | | |
|
||||
// 0.3 ===================== | DET2 | |
|
||||
// | | | DET1 | | | DET4 |
|
||||
// 0.2 | DET0 | =========== ================
|
||||
// | | | | | |
|
||||
// 0.1 =====|=============== |
|
||||
// | DET3 | | |
|
||||
// 0.0 ================ |
|
||||
// | DET5 |
|
||||
// -0.1 ===========
|
||||
// 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 1.1 1.2
|
||||
|
||||
// Detection det_0.
|
||||
det_0 = DetectionWithRelativeLocationData(/*xmin=*/0.1, /*ymin=*/0.1,
|
||||
/*width=*/0.2, /*height=*/0.2);
|
||||
det_0.set_detection_id(0);
|
||||
|
||||
// Detection det_1.
|
||||
det_1 = DetectionWithRelativeLocationData(/*xmin=*/0.3, /*ymin=*/0.1,
|
||||
/*width=*/0.2, /*height=*/0.2);
|
||||
det_1.set_detection_id(1);
|
||||
|
||||
// Detection det_2.
|
||||
det_2 = DetectionWithRelativeLocationData(/*xmin=*/0.9, /*ymin=*/0.2,
|
||||
/*width=*/0.2, /*height=*/0.2);
|
||||
det_2.set_detection_id(2);
|
||||
|
||||
// Detection det_3.
|
||||
det_3 = DetectionWithRelativeLocationData(/*xmin=*/0.2, /*ymin=*/0.0,
|
||||
/*width=*/0.3, /*height=*/0.3);
|
||||
det_3.set_detection_id(3);
|
||||
|
||||
// Detection det_4.
|
||||
det_4 = DetectionWithRelativeLocationData(/*xmin=*/1.0, /*ymin=*/0.2,
|
||||
/*width=*/0.2, /*height=*/0.2);
|
||||
det_4.set_detection_id(4);
|
||||
|
||||
// Detection det_5.
|
||||
det_5 = DetectionWithRelativeLocationData(/*xmin=*/0.3, /*ymin=*/-0.1,
|
||||
/*width=*/0.3, /*height=*/0.3);
|
||||
det_5.set_detection_id(5);
|
||||
}
|
||||
|
||||
::mediapipe::Detection det_0, det_1, det_2, det_3, det_4, det_5;
|
||||
};
|
||||
|
||||
TEST_F(AssociationDetectionCalculatorTest, DetectionAssocTest) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationDetectionCalculator"
|
||||
input_stream: "input_vec_0"
|
||||
input_stream: "input_vec_1"
|
||||
input_stream: "input_vec_2"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream 0: det_0, det_1, det_2.
|
||||
auto input_vec_0 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_0->push_back(det_0);
|
||||
input_vec_0->push_back(det_1);
|
||||
input_vec_0->push_back(det_2);
|
||||
runner.MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(input_vec_0.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 1: det_3, det_4.
|
||||
auto input_vec_1 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_1->push_back(det_3);
|
||||
input_vec_1->push_back(det_4);
|
||||
runner.MutableInputs()->Index(1).packets.push_back(
|
||||
Adopt(input_vec_1.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 2: det_5.
|
||||
auto input_vec_2 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_2->push_back(det_5);
|
||||
runner.MutableInputs()->Index(2).packets.push_back(
|
||||
Adopt(input_vec_2.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::Detection>>();
|
||||
|
||||
// det_3 overlaps with det_0, det_1 and det_5 overlaps with det_3. Since det_5
|
||||
// is in the highest priority, we remove other rects. det_4 overlaps with
|
||||
// det_2, and det_4 is higher priority, so we keep it. The final output
|
||||
// therefore contains 2 elements.
|
||||
EXPECT_EQ(2, assoc_rects.size());
|
||||
// Outputs are in order of inputs, so det_4 is before det_5 in output vector.
|
||||
|
||||
// det_4 overlaps with det_2, so new id for det_4 is 2.
|
||||
EXPECT_TRUE(assoc_rects[0].has_detection_id());
|
||||
EXPECT_EQ(2, assoc_rects[0].detection_id());
|
||||
det_4.set_detection_id(2);
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(det_4));
|
||||
|
||||
// det_3 overlaps with det_0, so new id for det_3 is 0.
|
||||
// det_3 overlaps with det_1, so new id for det_3 is 1.
|
||||
// det_5 overlaps with det_3, so new id for det_5 is 1.
|
||||
EXPECT_TRUE(assoc_rects[1].has_detection_id());
|
||||
EXPECT_EQ(1, assoc_rects[1].detection_id());
|
||||
det_5.set_detection_id(1);
|
||||
EXPECT_THAT(assoc_rects[1], EqualsProto(det_5));
|
||||
}
|
||||
|
||||
TEST_F(AssociationDetectionCalculatorTest, DetectionAssocTestWithPrev) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationDetectionCalculator"
|
||||
input_stream: "PREV:input_vec_0"
|
||||
input_stream: "input_vec_1"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream 0: det_3, det_4.
|
||||
auto input_vec_0 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_0->push_back(det_3);
|
||||
input_vec_0->push_back(det_4);
|
||||
CollectionItemId prev_input_stream_id =
|
||||
runner.MutableInputs()->GetId("PREV", 0);
|
||||
runner.MutableInputs()
|
||||
->Get(prev_input_stream_id)
|
||||
.packets.push_back(Adopt(input_vec_0.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 1: det_5.
|
||||
auto input_vec_1 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_1->push_back(det_5);
|
||||
CollectionItemId input_stream_id = runner.MutableInputs()->GetId("", 0);
|
||||
runner.MutableInputs()
|
||||
->Get(input_stream_id)
|
||||
.packets.push_back(Adopt(input_vec_1.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::Detection>>();
|
||||
|
||||
// det_5 overlaps with det_3 and doesn't overlap with det_4. Since det_4 is
|
||||
// in the PREV input stream, it doesn't get copied to the output, so the final
|
||||
// output contains 1 element.
|
||||
EXPECT_EQ(1, assoc_rects.size());
|
||||
|
||||
// det_5 overlaps with det_3, det_3 is in PREV, so new id for det_5 is 3.
|
||||
EXPECT_TRUE(assoc_rects[0].has_detection_id());
|
||||
EXPECT_EQ(3, assoc_rects[0].detection_id());
|
||||
det_5.set_detection_id(3);
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(det_5));
|
||||
}
|
||||
|
||||
TEST_F(AssociationDetectionCalculatorTest, DetectionAssocTestReverse) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationDetectionCalculator"
|
||||
input_stream: "input_vec_0"
|
||||
input_stream: "input_vec_1"
|
||||
input_stream: "input_vec_2"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream 0: det_5.
|
||||
auto input_vec_0 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_0->push_back(det_5);
|
||||
runner.MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(input_vec_0.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 1: det_3, det_4.
|
||||
auto input_vec_1 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_1->push_back(det_3);
|
||||
input_vec_1->push_back(det_4);
|
||||
runner.MutableInputs()->Index(1).packets.push_back(
|
||||
Adopt(input_vec_1.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 2: det_0, det_1, det_2.
|
||||
auto input_vec_2 = absl::make_unique<std::vector<::mediapipe::Detection>>();
|
||||
input_vec_2->push_back(det_0);
|
||||
input_vec_2->push_back(det_1);
|
||||
input_vec_2->push_back(det_2);
|
||||
runner.MutableInputs()->Index(2).packets.push_back(
|
||||
Adopt(input_vec_2.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::Detection>>();
|
||||
|
||||
// det_3 overlaps with det_5, so det_5 is removed. det_0 overlaps with det_3,
|
||||
// so det_3 is removed as det_0 is in higher priority for keeping. det_2
|
||||
// overlaps with det_4 so det_4 is removed as det_2 is higher priority for
|
||||
// keeping. The final output therefore contains 3 elements.
|
||||
EXPECT_EQ(3, assoc_rects.size());
|
||||
// Outputs are in same order as inputs.
|
||||
|
||||
// det_3 overlaps with det_5, so new id for det_3 is 5.
|
||||
// det_0 overlaps with det_3, so new id for det_0 is 5.
|
||||
EXPECT_TRUE(assoc_rects[0].has_detection_id());
|
||||
EXPECT_EQ(5, assoc_rects[0].detection_id());
|
||||
det_0.set_detection_id(5);
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(det_0));
|
||||
|
||||
// det_1 stays with id 1.
|
||||
EXPECT_TRUE(assoc_rects[1].has_detection_id());
|
||||
EXPECT_EQ(1, assoc_rects[1].detection_id());
|
||||
EXPECT_THAT(assoc_rects[1], EqualsProto(det_1));
|
||||
|
||||
// det_2 overlaps with det_4, so new id for det_2 is 4.
|
||||
EXPECT_TRUE(assoc_rects[2].has_detection_id());
|
||||
EXPECT_EQ(4, assoc_rects[2].detection_id());
|
||||
det_2.set_detection_id(4);
|
||||
EXPECT_THAT(assoc_rects[2], EqualsProto(det_2));
|
||||
}
|
||||
|
||||
class AssociationNormRectCalculatorTest : public ::testing::Test {
|
||||
protected:
|
||||
AssociationNormRectCalculatorTest() {
|
||||
// 0.4 ================
|
||||
// | | | |
|
||||
// 0.3 ===================== | NR2 | |
|
||||
// | | | NR1 | | | NR4 |
|
||||
// 0.2 | NR0 | =========== ================
|
||||
// | | | | | |
|
||||
// 0.1 =====|=============== |
|
||||
// | NR3 | | |
|
||||
// 0.0 ================ |
|
||||
// | NR5 |
|
||||
// -0.1 ===========
|
||||
// 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0 1.1 1.2
|
||||
|
||||
// NormalizedRect nr_0.
|
||||
nr_0.set_x_center(0.2);
|
||||
nr_0.set_y_center(0.2);
|
||||
nr_0.set_width(0.2);
|
||||
nr_0.set_height(0.2);
|
||||
|
||||
// NormalizedRect nr_1.
|
||||
nr_1.set_x_center(0.4);
|
||||
nr_1.set_y_center(0.2);
|
||||
nr_1.set_width(0.2);
|
||||
nr_1.set_height(0.2);
|
||||
|
||||
// NormalizedRect nr_2.
|
||||
nr_2.set_x_center(1.0);
|
||||
nr_2.set_y_center(0.3);
|
||||
nr_2.set_width(0.2);
|
||||
nr_2.set_height(0.2);
|
||||
|
||||
// NormalizedRect nr_3.
|
||||
nr_3.set_x_center(0.35);
|
||||
nr_3.set_y_center(0.15);
|
||||
nr_3.set_width(0.3);
|
||||
nr_3.set_height(0.3);
|
||||
|
||||
// NormalizedRect nr_4.
|
||||
nr_4.set_x_center(1.1);
|
||||
nr_4.set_y_center(0.3);
|
||||
nr_4.set_width(0.2);
|
||||
nr_4.set_height(0.2);
|
||||
|
||||
// NormalizedRect nr_5.
|
||||
nr_5.set_x_center(0.45);
|
||||
nr_5.set_y_center(0.05);
|
||||
nr_5.set_width(0.3);
|
||||
nr_5.set_height(0.3);
|
||||
}
|
||||
|
||||
::mediapipe::NormalizedRect nr_0, nr_1, nr_2, nr_3, nr_4, nr_5;
|
||||
};
|
||||
|
||||
TEST_F(AssociationNormRectCalculatorTest, NormRectAssocTest) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationNormRectCalculator"
|
||||
input_stream: "input_vec_0"
|
||||
input_stream: "input_vec_1"
|
||||
input_stream: "input_vec_2"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream 0: nr_0, nr_1, nr_2.
|
||||
auto input_vec_0 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_0->push_back(nr_0);
|
||||
input_vec_0->push_back(nr_1);
|
||||
input_vec_0->push_back(nr_2);
|
||||
runner.MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(input_vec_0.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 1: nr_3, nr_4.
|
||||
auto input_vec_1 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_1->push_back(nr_3);
|
||||
input_vec_1->push_back(nr_4);
|
||||
runner.MutableInputs()->Index(1).packets.push_back(
|
||||
Adopt(input_vec_1.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 2: nr_5.
|
||||
auto input_vec_2 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_2->push_back(nr_5);
|
||||
runner.MutableInputs()->Index(2).packets.push_back(
|
||||
Adopt(input_vec_2.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::NormalizedRect>>();
|
||||
|
||||
// nr_3 overlaps with nr_0, nr_1 and nr_5 overlaps with nr_3. Since nr_5 is
|
||||
// in the highest priority, we remove other rects.
|
||||
// nr_4 overlaps with nr_2, and nr_4 is higher priority, so we keep it.
|
||||
// The final output therefore contains 2 elements.
|
||||
EXPECT_EQ(2, assoc_rects.size());
|
||||
// Outputs are in order of inputs, so nr_4 is before nr_5 in output vector.
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(nr_4));
|
||||
EXPECT_THAT(assoc_rects[1], EqualsProto(nr_5));
|
||||
}
|
||||
|
||||
TEST_F(AssociationNormRectCalculatorTest, NormRectAssocTestReverse) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationNormRectCalculator"
|
||||
input_stream: "input_vec_0"
|
||||
input_stream: "input_vec_1"
|
||||
input_stream: "input_vec_2"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream 0: nr_5.
|
||||
auto input_vec_0 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_0->push_back(nr_5);
|
||||
runner.MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(input_vec_0.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 1: nr_3, nr_4.
|
||||
auto input_vec_1 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_1->push_back(nr_3);
|
||||
input_vec_1->push_back(nr_4);
|
||||
runner.MutableInputs()->Index(1).packets.push_back(
|
||||
Adopt(input_vec_1.release()).At(Timestamp(1)));
|
||||
|
||||
// Input Stream 2: nr_0, nr_1, nr_2.
|
||||
auto input_vec_2 =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec_2->push_back(nr_0);
|
||||
input_vec_2->push_back(nr_1);
|
||||
input_vec_2->push_back(nr_2);
|
||||
runner.MutableInputs()->Index(2).packets.push_back(
|
||||
Adopt(input_vec_2.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::NormalizedRect>>();
|
||||
|
||||
// nr_3 overlaps with nr_5, so nr_5 is removed. nr_0 overlaps with nr_3, so
|
||||
// nr_3 is removed as nr_0 is in higher priority for keeping. nr_2 overlaps
|
||||
// with nr_4 so nr_4 is removed as nr_2 is higher priority for keeping.
|
||||
// The final output therefore contains 3 elements.
|
||||
EXPECT_EQ(3, assoc_rects.size());
|
||||
// Outputs are in same order as inputs.
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(nr_0));
|
||||
EXPECT_THAT(assoc_rects[1], EqualsProto(nr_1));
|
||||
EXPECT_THAT(assoc_rects[2], EqualsProto(nr_2));
|
||||
}
|
||||
|
||||
TEST_F(AssociationNormRectCalculatorTest, NormRectAssocSingleInputStream) {
|
||||
CalculatorRunner runner(ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "AssociationNormRectCalculator"
|
||||
input_stream: "input_vec"
|
||||
output_stream: "output_vec"
|
||||
options {
|
||||
[mediapipe.AssociationCalculatorOptions.ext] {
|
||||
min_similarity_threshold: 0.1
|
||||
}
|
||||
}
|
||||
)"));
|
||||
|
||||
// Input Stream : nr_3, nr_5.
|
||||
auto input_vec =
|
||||
absl::make_unique<std::vector<::mediapipe::NormalizedRect>>();
|
||||
input_vec->push_back(nr_3);
|
||||
input_vec->push_back(nr_5);
|
||||
runner.MutableInputs()->Index(0).packets.push_back(
|
||||
Adopt(input_vec.release()).At(Timestamp(1)));
|
||||
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Index(0).packets;
|
||||
EXPECT_EQ(1, output.size());
|
||||
const auto& assoc_rects =
|
||||
output[0].Get<std::vector<::mediapipe::NormalizedRect>>();
|
||||
|
||||
// nr_5 overlaps with nr_3. Since nr_5 is after nr_3 in the same input stream
|
||||
// we remove nr_3 and keep nr_5.
|
||||
// The final output therefore contains 1 elements.
|
||||
EXPECT_EQ(1, assoc_rects.size());
|
||||
EXPECT_THAT(assoc_rects[0], EqualsProto(nr_5));
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,77 @@
|
||||
// Copyright 2019 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/association_calculator.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/location.h"
|
||||
#include "mediapipe/framework/port/rectangle.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A subclass of AssociationCalculator<T> for Detection. Example:
|
||||
// node {
|
||||
// calculator: "AssociationDetectionCalculator"
|
||||
// input_stream: "PREV:input_vec_0"
|
||||
// input_stream: "input_vec_1"
|
||||
// input_stream: "input_vec_2"
|
||||
// output_stream: "output_vec"
|
||||
// options {
|
||||
// [mediapipe.AssociationCalculatorOptions.ext] {
|
||||
// min_similarity_threshold: 0.1
|
||||
// }
|
||||
// }
|
||||
class AssociationDetectionCalculator
|
||||
: public AssociationCalculator<::mediapipe::Detection> {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
return AssociationCalculator<::mediapipe::Detection>::GetContract(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::Detection>::Open(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::Detection>::Process(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::Detection>::Close(cc);
|
||||
}
|
||||
|
||||
protected:
|
||||
::mediapipe::StatusOr<Rectangle_f> GetRectangle(
|
||||
const ::mediapipe::Detection& input) override {
|
||||
if (!input.has_location_data()) {
|
||||
return ::mediapipe::InternalError("Missing location_data in Detection");
|
||||
}
|
||||
const Location location(input.location_data());
|
||||
return location.GetRelativeBBox();
|
||||
}
|
||||
|
||||
std::pair<bool, int> GetId(const ::mediapipe::Detection& input) override {
|
||||
return {input.has_detection_id(), input.detection_id()};
|
||||
}
|
||||
|
||||
void SetId(::mediapipe::Detection* input, int id) override {
|
||||
input->set_detection_id(id);
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(AssociationDetectionCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,72 @@
|
||||
// Copyright 2019 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/association_calculator.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/port/rectangle.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A subclass of AssociationCalculator<T> for NormalizedRect. Example use case:
|
||||
// node {
|
||||
// calculator: "AssociationNormRectCalculator"
|
||||
// input_stream: "input_vec_0"
|
||||
// input_stream: "input_vec_1"
|
||||
// input_stream: "input_vec_2"
|
||||
// output_stream: "output_vec"
|
||||
// options {
|
||||
// [mediapipe.AssociationCalculatorOptions.ext] {
|
||||
// min_similarity_threshold: 0.1
|
||||
// }
|
||||
// }
|
||||
class AssociationNormRectCalculator
|
||||
: public AssociationCalculator<::mediapipe::NormalizedRect> {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
return AssociationCalculator<::mediapipe::NormalizedRect>::GetContract(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::NormalizedRect>::Open(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::NormalizedRect>::Process(cc);
|
||||
}
|
||||
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override {
|
||||
return AssociationCalculator<::mediapipe::NormalizedRect>::Close(cc);
|
||||
}
|
||||
|
||||
protected:
|
||||
::mediapipe::StatusOr<Rectangle_f> GetRectangle(
|
||||
const ::mediapipe::NormalizedRect& input) override {
|
||||
if (!input.has_x_center() || !input.has_y_center() || !input.has_width() ||
|
||||
!input.has_height()) {
|
||||
return ::mediapipe::InternalError(
|
||||
"Missing dimensions in NormalizedRect.");
|
||||
}
|
||||
const float xmin = input.x_center() - input.width() / 2.0;
|
||||
const float ymin = input.y_center() - input.height() / 2.0;
|
||||
// TODO: Support rotation for rectangle.
|
||||
return Rectangle_f(xmin, ymin, input.width(), input.height());
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(AssociationNormRectCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,26 @@
|
||||
|
||||
// Copyright 2019 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/collection_has_min_size_calculator.h"
|
||||
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef CollectionHasMinSizeCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||
NormalizedRectVectorHasMinSizeCalculator;
|
||||
REGISTER_CALCULATOR(NormalizedRectVectorHasMinSizeCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,84 @@
|
||||
// Copyright 2019 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_COLLECTION_HAS_MIN_SIZE_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_UTIL_COLLECTION_HAS_MIN_SIZE_CALCULATOR_H_
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/calculators/util/collection_has_min_size_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Deterimines if an input iterable collection has a minimum size, specified
|
||||
// in CollectionHasMinSizeCalculatorOptions. Example usage:
|
||||
// node {
|
||||
// calculator: "IntVectorHasMinSizeCalculator"
|
||||
// input_stream: "ITERABLE:input_int_vector"
|
||||
// output_stream: "has_min_ints"
|
||||
// options {
|
||||
// [mediapipe.CollectionHasMinSizeCalculatorOptions.ext] {
|
||||
// min_size: 2
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
template <typename IterableT>
|
||||
class CollectionHasMinSizeCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag("ITERABLE"));
|
||||
RET_CHECK_EQ(1, cc->Inputs().NumEntries());
|
||||
|
||||
RET_CHECK_EQ(1, cc->Outputs().NumEntries());
|
||||
|
||||
RET_CHECK_GE(
|
||||
cc->Options<::mediapipe::CollectionHasMinSizeCalculatorOptions>()
|
||||
.min_size(),
|
||||
0);
|
||||
|
||||
cc->Inputs().Tag("ITERABLE").Set<IterableT>();
|
||||
cc->Outputs().Index(0).Set<bool>();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
min_size_ =
|
||||
cc->Options<::mediapipe::CollectionHasMinSizeCalculatorOptions>()
|
||||
.min_size();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
const IterableT& input = cc->Inputs().Tag("ITERABLE").Get<IterableT>();
|
||||
bool has_min_size = input.size() >= min_size_;
|
||||
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
MakePacket<bool>(has_min_size).At(cc->InputTimestamp()));
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
int min_size_ = 0;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_UTIL_COLLECTION_HAS_MIN_SIZE_CALCULATOR_H_
|
||||
@@ -0,0 +1,29 @@
|
||||
// Copyright 2019 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;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
|
||||
message CollectionHasMinSizeCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional CollectionHasMinSizeCalculatorOptions ext = 259397840;
|
||||
}
|
||||
|
||||
// The minimum size an input iterable collection should have for the
|
||||
// calculator to output true.
|
||||
optional int32 min_size = 1 [default = 0];
|
||||
}
|
||||
@@ -19,8 +19,7 @@
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
|
||||
#if defined(MEDIAPIPE_LITE) || defined(__ANDROID__) || \
|
||||
(defined(__APPLE__) && !TARGET_OS_OSX)
|
||||
#if defined(MEDIAPIPE_MOBILE)
|
||||
#include "mediapipe/util/android/file/base/file.h"
|
||||
#include "mediapipe/util/android/file/base/helpers.h"
|
||||
#else
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
// Copyright 2019 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/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr char kDetectionsTag[] = "DETECTIONS";
|
||||
constexpr char kDetectionListTag[] = "DETECTION_LIST";
|
||||
|
||||
// Each detection processed by DetectionUniqueIDCalculator will be assigned an
|
||||
// unique id that starts from 1. If a detection already has an ID other than 0,
|
||||
// the ID will be overwritten.
|
||||
static int64 detection_id = 0;
|
||||
|
||||
inline int GetNextDetectionId() { return ++detection_id; }
|
||||
|
||||
} // namespace
|
||||
|
||||
// Assign a unique id to detections.
|
||||
// Note that the calculator will consume the input vector of Detection or
|
||||
// DetectionList. So the input stream can not be connected to other calculators.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "DetectionUniqueIdCalculator"
|
||||
// input_stream: "DETECTIONS:detections"
|
||||
// output_stream: "DETECTIONS:output_detections"
|
||||
// }
|
||||
class DetectionUniqueIdCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag(kDetectionListTag) ||
|
||||
cc->Inputs().HasTag(kDetectionsTag))
|
||||
<< "None of the input streams are provided.";
|
||||
|
||||
if (cc->Inputs().HasTag(kDetectionListTag)) {
|
||||
RET_CHECK(cc->Outputs().HasTag(kDetectionListTag));
|
||||
cc->Inputs().Tag(kDetectionListTag).Set<DetectionList>();
|
||||
cc->Outputs().Tag(kDetectionListTag).Set<DetectionList>();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kDetectionsTag)) {
|
||||
RET_CHECK(cc->Outputs().HasTag(kDetectionsTag));
|
||||
cc->Inputs().Tag(kDetectionsTag).Set<std::vector<Detection>>();
|
||||
cc->Outputs().Tag(kDetectionsTag).Set<std::vector<Detection>>();
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(::mediapipe::TimestampDiff(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
};
|
||||
REGISTER_CALCULATOR(DetectionUniqueIdCalculator);
|
||||
|
||||
::mediapipe::Status DetectionUniqueIdCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
if (cc->Inputs().HasTag(kDetectionListTag) &&
|
||||
!cc->Inputs().Tag(kDetectionListTag).IsEmpty()) {
|
||||
auto result =
|
||||
cc->Inputs().Tag(kDetectionListTag).Value().Consume<DetectionList>();
|
||||
if (result.ok()) {
|
||||
auto detection_list = std::move(result).ValueOrDie();
|
||||
for (Detection& detection : *detection_list->mutable_detection()) {
|
||||
detection.set_detection_id(GetNextDetectionId());
|
||||
}
|
||||
cc->Outputs()
|
||||
.Tag(kDetectionListTag)
|
||||
.Add(detection_list.release(), cc->InputTimestamp());
|
||||
}
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag(kDetectionsTag) &&
|
||||
!cc->Inputs().Tag(kDetectionsTag).IsEmpty()) {
|
||||
auto result = cc->Inputs()
|
||||
.Tag(kDetectionsTag)
|
||||
.Value()
|
||||
.Consume<std::vector<Detection>>();
|
||||
if (result.ok()) {
|
||||
auto detections = std::move(result).ValueOrDie();
|
||||
for (Detection& detection : *detections) {
|
||||
detection.set_detection_id(GetNextDetectionId());
|
||||
}
|
||||
cc->Outputs()
|
||||
.Tag(kDetectionsTag)
|
||||
.Add(detections.release(), cc->InputTimestamp());
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -11,6 +11,8 @@
|
||||
// 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/detections_to_rects_calculator.h"
|
||||
|
||||
#include <cmath>
|
||||
|
||||
#include "mediapipe/calculators/util/detections_to_rects_calculator.pb.h"
|
||||
@@ -24,8 +26,6 @@
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
using mediapipe::DetectionsToRectsCalculatorOptions;
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr char kDetectionTag[] = "DETECTION";
|
||||
@@ -36,7 +36,10 @@ constexpr char kNormRectTag[] = "NORM_RECT";
|
||||
constexpr char kRectsTag[] = "RECTS";
|
||||
constexpr char kNormRectsTag[] = "NORM_RECTS";
|
||||
|
||||
::mediapipe::Status DetectionToRect(const Detection& detection, Rect* rect) {
|
||||
} // namespace
|
||||
|
||||
::mediapipe::Status DetectionsToRectsCalculator::DetectionToRect(
|
||||
const Detection& detection, Rect* rect) {
|
||||
const LocationData location_data = detection.location_data();
|
||||
RET_CHECK(location_data.format() == LocationData::BOUNDING_BOX)
|
||||
<< "Only Detection with formats of BOUNDING_BOX can be converted to Rect";
|
||||
@@ -48,8 +51,8 @@ constexpr char kNormRectsTag[] = "NORM_RECTS";
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status DetectionToNormalizedRect(const Detection& detection,
|
||||
NormalizedRect* rect) {
|
||||
::mediapipe::Status DetectionsToRectsCalculator::DetectionToNormalizedRect(
|
||||
const Detection& detection, NormalizedRect* rect) {
|
||||
const LocationData location_data = detection.location_data();
|
||||
RET_CHECK(location_data.format() == LocationData::RELATIVE_BOUNDING_BOX)
|
||||
<< "Only Detection with formats of RELATIVE_BOUNDING_BOX can be "
|
||||
@@ -63,79 +66,6 @@ constexpr char kNormRectsTag[] = "NORM_RECTS";
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
// Wraps around an angle in radians to within -M_PI and M_PI.
|
||||
inline float NormalizeRadians(float angle) {
|
||||
return angle - 2 * M_PI * std::floor((angle - (-M_PI)) / (2 * M_PI));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// A calculator that converts Detection proto to Rect proto.
|
||||
//
|
||||
// Detection is the format for encoding one or more detections in an image.
|
||||
// The input can be a single Detection or std::vector<Detection>. The output can
|
||||
// be either a single Rect or NormalizedRect, or std::vector<Rect> or
|
||||
// std::vector<NormalizedRect>. If Rect is used, the LocationData format is
|
||||
// expected to be BOUNDING_BOX, and if NormalizedRect is used it is expected to
|
||||
// be RELATIVE_BOUNDING_BOX.
|
||||
//
|
||||
// When the input is std::vector<Detection> and the output is a Rect or
|
||||
// NormalizedRect, only the first detection is converted. When the input is a
|
||||
// single Detection and the output is a std::vector<Rect> or
|
||||
// std::vector<NormalizedRect>, the output is a vector of size 1.
|
||||
//
|
||||
// Inputs:
|
||||
//
|
||||
// One of the following:
|
||||
// DETECTION: A Detection proto.
|
||||
// DETECTIONS: An std::vector<Detection>.
|
||||
//
|
||||
// IMAGE_SIZE (optional): A std::pair<int, int> represention image width and
|
||||
// height. This is required only when rotation needs to be computed (see
|
||||
// calculator options).
|
||||
//
|
||||
// Output:
|
||||
// One of the following:
|
||||
// RECT: A Rect proto.
|
||||
// NORM_RECT: A NormalizedRect proto.
|
||||
// RECTS: An std::vector<Rect>.
|
||||
// NORM_RECTS: An std::vector<NormalizedRect>.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "DetectionsToRectsCalculator"
|
||||
// input_stream: "DETECTIONS:detections"
|
||||
// input_stream: "IMAGE_SIZE:image_size"
|
||||
// output_stream: "NORM_RECT:rect"
|
||||
// options: {
|
||||
// [mediapipe.DetectionsToRectCalculatorOptions.ext] {
|
||||
// rotation_vector_start_keypoint_index: 0
|
||||
// rotation_vector_end_keypoint_index: 2
|
||||
// rotation_vector_target_angle_degrees: 90
|
||||
// output_zero_rect_for_empty_detections: true
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
class DetectionsToRectsCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
float ComputeRotation(const Detection& detection,
|
||||
const std::pair<int, int> image_size);
|
||||
|
||||
DetectionsToRectsCalculatorOptions options_;
|
||||
int start_keypoint_index_;
|
||||
int end_keypoint_index_;
|
||||
float target_angle_; // In radians.
|
||||
bool rotate_;
|
||||
bool output_zero_rect_for_empty_detections_;
|
||||
};
|
||||
REGISTER_CALCULATOR(DetectionsToRectsCalculator);
|
||||
|
||||
::mediapipe::Status DetectionsToRectsCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag(kDetectionTag) ^
|
||||
@@ -232,6 +162,13 @@ REGISTER_CALCULATOR(DetectionsToRectsCalculator);
|
||||
.Tag(kNormRectTag)
|
||||
.AddPacket(MakePacket<NormalizedRect>().At(cc->InputTimestamp()));
|
||||
}
|
||||
if (cc->Outputs().HasTag(kNormRectsTag)) {
|
||||
auto rect_vector = absl::make_unique<std::vector<NormalizedRect>>();
|
||||
rect_vector->emplace_back(NormalizedRect());
|
||||
cc->Outputs()
|
||||
.Tag(kNormRectsTag)
|
||||
.Add(rect_vector.release(), cc->InputTimestamp());
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -312,4 +249,6 @@ float DetectionsToRectsCalculator::ComputeRotation(
|
||||
return NormalizeRadians(rotation);
|
||||
}
|
||||
|
||||
REGISTER_CALCULATOR(DetectionsToRectsCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
// Copyright 2019 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_DETECTIONS_TO_RECTS_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_UTIL_DETECTIONS_TO_RECTS_CALCULATOR_H_
|
||||
|
||||
#include <cmath>
|
||||
|
||||
#include "mediapipe/calculators/util/detections_to_rects_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_options.pb.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/location_data.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator that converts Detection proto to Rect proto.
|
||||
//
|
||||
// Detection is the format for encoding one or more detections in an image.
|
||||
// The input can be a single Detection or std::vector<Detection>. The output can
|
||||
// be either a single Rect or NormalizedRect, or std::vector<Rect> or
|
||||
// std::vector<NormalizedRect>. If Rect is used, the LocationData format is
|
||||
// expected to be BOUNDING_BOX, and if NormalizedRect is used it is expected to
|
||||
// be RELATIVE_BOUNDING_BOX.
|
||||
//
|
||||
// When the input is std::vector<Detection> and the output is a Rect or
|
||||
// NormalizedRect, only the first detection is converted. When the input is a
|
||||
// single Detection and the output is a std::vector<Rect> or
|
||||
// std::vector<NormalizedRect>, the output is a vector of size 1.
|
||||
//
|
||||
// Inputs:
|
||||
//
|
||||
// One of the following:
|
||||
// DETECTION: A Detection proto.
|
||||
// DETECTIONS: An std::vector<Detection>.
|
||||
//
|
||||
// IMAGE_SIZE (optional): A std::pair<int, int> represention image width and
|
||||
// height. This is required only when rotation needs to be computed (see
|
||||
// calculator options).
|
||||
//
|
||||
// Output:
|
||||
// One of the following:
|
||||
// RECT: A Rect proto.
|
||||
// NORM_RECT: A NormalizedRect proto.
|
||||
// RECTS: An std::vector<Rect>.
|
||||
// NORM_RECTS: An std::vector<NormalizedRect>.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "DetectionsToRectsCalculator"
|
||||
// input_stream: "DETECTIONS:detections"
|
||||
// input_stream: "IMAGE_SIZE:image_size"
|
||||
// output_stream: "NORM_RECT:rect"
|
||||
// options: {
|
||||
// [mediapipe.DetectionsToRectCalculatorOptions.ext] {
|
||||
// rotation_vector_start_keypoint_index: 0
|
||||
// rotation_vector_end_keypoint_index: 2
|
||||
// rotation_vector_target_angle_degrees: 90
|
||||
// output_zero_rect_for_empty_detections: true
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
class DetectionsToRectsCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
protected:
|
||||
virtual float ComputeRotation(const ::mediapipe::Detection& detection,
|
||||
const std::pair<int, int> image_size);
|
||||
virtual ::mediapipe::Status DetectionToRect(
|
||||
const ::mediapipe::Detection& detection, ::mediapipe::Rect* rect);
|
||||
virtual ::mediapipe::Status DetectionToNormalizedRect(
|
||||
const ::mediapipe::Detection& detection,
|
||||
::mediapipe::NormalizedRect* rect);
|
||||
|
||||
static inline float NormalizeRadians(float angle) {
|
||||
return angle - 2 * M_PI * std::floor((angle - (-M_PI)) / (2 * M_PI));
|
||||
}
|
||||
|
||||
::mediapipe::DetectionsToRectsCalculatorOptions options_;
|
||||
int start_keypoint_index_;
|
||||
int end_keypoint_index_;
|
||||
float target_angle_ = 0.0f; // In radians.
|
||||
bool rotate_;
|
||||
bool output_zero_rect_for_empty_detections_;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
#endif // MEDIAPIPE_CALCULATORS_UTIL_DETECTIONS_TO_RECTS_CALCULATOR_H_
|
||||
@@ -39,6 +39,8 @@ constexpr char kKeypointLabel[] = "KEYPOINT";
|
||||
// The ratio of detection label font height to the height of detection bounding
|
||||
// box.
|
||||
constexpr double kLabelToBoundingBoxRatio = 0.1;
|
||||
// Perserve 2 decimal digits.
|
||||
constexpr float kNumScoreDecimalDigitsMultipler = 100;
|
||||
|
||||
} // namespace
|
||||
|
||||
@@ -235,18 +237,26 @@ void DetectionsToRenderDataCalculator::AddLabels(
|
||||
std::string label_str = detection.label().empty()
|
||||
? absl::StrCat(detection.label_id(i))
|
||||
: detection.label(i);
|
||||
const float rounded_score =
|
||||
std::round(detection.score(i) * kNumScoreDecimalDigitsMultipler) /
|
||||
kNumScoreDecimalDigitsMultipler;
|
||||
std::string label_and_score =
|
||||
absl::StrCat(label_str, options.text_delimiter(), detection.score(i),
|
||||
absl::StrCat(label_str, options.text_delimiter(), rounded_score,
|
||||
options.text_delimiter());
|
||||
label_and_scores.push_back(label_and_score);
|
||||
}
|
||||
std::vector<std::string> labels;
|
||||
if (options.render_detection_id()) {
|
||||
const std::string detection_id_str =
|
||||
absl::StrCat("Id: ", detection.detection_id());
|
||||
labels.push_back(detection_id_str);
|
||||
}
|
||||
if (options.one_label_per_line()) {
|
||||
labels.swap(label_and_scores);
|
||||
labels.insert(labels.end(), label_and_scores.begin(),
|
||||
label_and_scores.end());
|
||||
} else {
|
||||
labels.push_back(absl::StrJoin(label_and_scores, ""));
|
||||
}
|
||||
|
||||
// Add the render annotations for "label(_id),score".
|
||||
for (int i = 0; i < labels.size(); ++i) {
|
||||
auto label = labels.at(i);
|
||||
|
||||
@@ -53,4 +53,7 @@ message DetectionsToRenderDataCalculatorOptions {
|
||||
// instances of this calculator are present in the graph, this value
|
||||
// should be unique among them.
|
||||
optional string scene_class = 7 [default = "DETECTION"];
|
||||
|
||||
// If true, renders the detection id in the first line before the labels.
|
||||
optional bool render_detection_id = 8 [default = false];
|
||||
}
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
// Copyright 2019 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/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/location_data.pb.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/util/tracking/box_tracker.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr char kDetectionsTag[] = "DETECTIONS";
|
||||
constexpr char kDetectionListTag[] = "DETECTION_LIST";
|
||||
constexpr char kBoxesTag[] = "BOXES";
|
||||
|
||||
} // namespace
|
||||
|
||||
// A calculator that converts Detection proto to TimedBoxList proto for
|
||||
// tracking.
|
||||
//
|
||||
// Please note that only Location Data formats of RELATIVE_BOUNDING_BOX are
|
||||
// supported.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "DetectionsToTimedBoxListCalculator"
|
||||
// input_stream: "DETECTIONS:detections"
|
||||
// output_stream: "BOXES:boxes"
|
||||
// }
|
||||
class DetectionsToTimedBoxListCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag(kDetectionListTag) ||
|
||||
cc->Inputs().HasTag(kDetectionsTag))
|
||||
<< "None of the input streams are provided.";
|
||||
if (cc->Inputs().HasTag(kDetectionListTag)) {
|
||||
cc->Inputs().Tag(kDetectionListTag).Set<DetectionList>();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kDetectionsTag)) {
|
||||
cc->Inputs().Tag(kDetectionsTag).Set<std::vector<Detection>>();
|
||||
}
|
||||
cc->Outputs().Tag(kBoxesTag).Set<TimedBoxProtoList>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
void ConvertDetectionToTimedBox(const Detection& detection,
|
||||
TimedBoxProto* box, CalculatorContext* cc);
|
||||
};
|
||||
REGISTER_CALCULATOR(DetectionsToTimedBoxListCalculator);
|
||||
|
||||
::mediapipe::Status DetectionsToTimedBoxListCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
auto output_timed_box_list = absl::make_unique<TimedBoxProtoList>();
|
||||
|
||||
if (cc->Inputs().HasTag(kDetectionListTag)) {
|
||||
const auto& detection_list =
|
||||
cc->Inputs().Tag(kDetectionListTag).Get<DetectionList>();
|
||||
for (const auto& detection : detection_list.detection()) {
|
||||
TimedBoxProto* box = output_timed_box_list->add_box();
|
||||
ConvertDetectionToTimedBox(detection, box, cc);
|
||||
}
|
||||
}
|
||||
if (cc->Inputs().HasTag(kDetectionsTag)) {
|
||||
const auto& detections =
|
||||
cc->Inputs().Tag(kDetectionsTag).Get<std::vector<Detection>>();
|
||||
for (const auto& detection : detections) {
|
||||
TimedBoxProto* box = output_timed_box_list->add_box();
|
||||
ConvertDetectionToTimedBox(detection, box, cc);
|
||||
}
|
||||
}
|
||||
|
||||
cc->Outputs().Tag(kBoxesTag).Add(output_timed_box_list.release(),
|
||||
cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
void DetectionsToTimedBoxListCalculator::ConvertDetectionToTimedBox(
|
||||
const Detection& detection, TimedBoxProto* box, CalculatorContext* cc) {
|
||||
const auto& relative_bounding_box =
|
||||
detection.location_data().relative_bounding_box();
|
||||
box->set_left(relative_bounding_box.xmin());
|
||||
box->set_right(relative_bounding_box.xmin() + relative_bounding_box.width());
|
||||
box->set_top(relative_bounding_box.ymin());
|
||||
box->set_bottom(relative_bounding_box.ymin() +
|
||||
relative_bounding_box.height());
|
||||
box->set_id(detection.detection_id());
|
||||
box->set_time_msec(cc->InputTimestamp().Microseconds() / 1000);
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,34 @@
|
||||
|
||||
// Copyright 2019 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/filter_collection_calculator.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
typedef FilterCollectionCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||
FilterNormalizedRectCollectionCalculator;
|
||||
REGISTER_CALCULATOR(FilterNormalizedRectCollectionCalculator);
|
||||
|
||||
typedef FilterCollectionCalculator<
|
||||
std::vector<::mediapipe::NormalizedLandmarkList>>
|
||||
FilterLandmarkListCollectionCalculator;
|
||||
REGISTER_CALCULATOR(FilterLandmarkListCollectionCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,109 @@
|
||||
// Copyright 2019 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_FILTER_VECTOR_CALCULATOR_H_
|
||||
#define MEDIAPIPE_CALCULATORS_UTIL_FILTER_VECTOR_CALCULATOR_H_
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator that gates elements of an input collection based on
|
||||
// corresponding boolean values of the "CONDITION" vector. If there is no input
|
||||
// collection or "CONDITION" vector, the calculator forwards timestamp bounds
|
||||
// for downstream calculators. If the "CONDITION" vector has false values for
|
||||
// all elements of the input collection, the calculator outputs a packet
|
||||
// containing an empty collection.
|
||||
// Example usage:
|
||||
// node {
|
||||
// calculator: "FilterCollectionCalculator"
|
||||
// input_stream: "ITERABLE:input_collection"
|
||||
// input_stream: "CONDITION:condition_vector"
|
||||
// output_stream: "ITERABLE:output_collection"
|
||||
// }
|
||||
// This calculator is able to handle collections of copyable types T.
|
||||
template <typename IterableT>
|
||||
class FilterCollectionCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag("ITERABLE"));
|
||||
RET_CHECK(cc->Inputs().HasTag("CONDITION"));
|
||||
RET_CHECK(cc->Outputs().HasTag("ITERABLE"));
|
||||
|
||||
cc->Inputs().Tag("ITERABLE").Set<IterableT>();
|
||||
cc->Inputs().Tag("CONDITION").Set<std::vector<bool>>();
|
||||
|
||||
cc->Outputs().Tag("ITERABLE").Set<IterableT>();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
if (cc->Inputs().Tag("ITERABLE").IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
if (cc->Inputs().Tag("CONDITION").IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
const std::vector<bool>& filter_by =
|
||||
cc->Inputs().Tag("CONDITION").Get<std::vector<bool>>();
|
||||
|
||||
return FilterCollection<IterableT>(
|
||||
std::is_copy_constructible<typename IterableT::value_type>(), cc,
|
||||
filter_by);
|
||||
}
|
||||
|
||||
template <typename IterableU>
|
||||
::mediapipe::Status FilterCollection(std::true_type, CalculatorContext* cc,
|
||||
const std::vector<bool>& filter_by) {
|
||||
const IterableU& input = cc->Inputs().Tag("ITERABLE").Get<IterableU>();
|
||||
if (input.size() != filter_by.size()) {
|
||||
return ::mediapipe::InternalError(absl::StrCat(
|
||||
"Input vector size: ", input.size(),
|
||||
" doesn't mach condition vector size: ", filter_by.size()));
|
||||
}
|
||||
|
||||
auto output = absl::make_unique<IterableU>();
|
||||
for (int i = 0; i < input.size(); ++i) {
|
||||
if (filter_by[i]) {
|
||||
output->push_back(input[i]);
|
||||
}
|
||||
}
|
||||
cc->Outputs().Tag("ITERABLE").Add(output.release(), cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
template <typename IterableU>
|
||||
::mediapipe::Status FilterCollection(std::false_type, CalculatorContext* cc,
|
||||
const std::vector<bool>& filter_by) {
|
||||
return ::mediapipe::InternalError(
|
||||
"Cannot copy input collection to filter it.");
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_UTIL_FILTER_VECTOR_CALCULATOR_H_
|
||||
@@ -0,0 +1,182 @@
|
||||
// Copyright 2019 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 <math.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "mediapipe/calculators/util/labels_to_render_data_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/classification.pb.h"
|
||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/statusor.h"
|
||||
#include "mediapipe/util/color.pb.h"
|
||||
#include "mediapipe/util/render_data.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
constexpr float kFontHeightScale = 1.25f;
|
||||
|
||||
// A calculator takes in pairs of labels and scores or classifications, outputs
|
||||
// generates render data. Either both "LABELS" and "SCORES" or "CLASSIFICATIONS"
|
||||
// must be present.
|
||||
//
|
||||
// Usage example:
|
||||
// node {
|
||||
// calculator: "LabelsToRenderDataCalculator"
|
||||
// input_stream: "LABELS:labels"
|
||||
// input_stream: "SCORES:scores"
|
||||
// output_stream: "VIDEO_PRESTREAM:video_header"
|
||||
// options {
|
||||
// [LabelsToRenderDataCalculatorOptions.ext] {
|
||||
// color { r: 255 g: 0 b: 0 }
|
||||
// color { r: 0 g: 255 b: 0 }
|
||||
// color { r: 0 g: 0 b: 255 }
|
||||
// thickness: 2.0
|
||||
// font_height_px: 20
|
||||
// max_num_labels: 3
|
||||
// font_face: 1
|
||||
// location: TOP_LEFT
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
class LabelsToRenderDataCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
LabelsToRenderDataCalculatorOptions options_;
|
||||
int num_colors_ = 0;
|
||||
int video_width_ = 0;
|
||||
int video_height_ = 0;
|
||||
int label_height_px_ = 0;
|
||||
int label_left_px_ = 0;
|
||||
};
|
||||
REGISTER_CALCULATOR(LabelsToRenderDataCalculator);
|
||||
|
||||
::mediapipe::Status LabelsToRenderDataCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
if (cc->Inputs().HasTag("CLASSIFICATIONS")) {
|
||||
cc->Inputs().Tag("CLASSIFICATIONS").Set<ClassificationList>();
|
||||
} else {
|
||||
RET_CHECK(cc->Inputs().HasTag("LABELS"))
|
||||
<< "Must provide input stream \"LABELS\"";
|
||||
cc->Inputs().Tag("LABELS").Set<std::vector<std::string>>();
|
||||
if (cc->Inputs().HasTag("SCORES")) {
|
||||
cc->Inputs().Tag("SCORES").Set<std::vector<float>>();
|
||||
}
|
||||
}
|
||||
if (cc->Inputs().HasTag("VIDEO_PRESTREAM")) {
|
||||
cc->Inputs().Tag("VIDEO_PRESTREAM").Set<VideoHeader>();
|
||||
}
|
||||
cc->Outputs().Tag("RENDER_DATA").Set<RenderData>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status LabelsToRenderDataCalculator::Open(CalculatorContext* cc) {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
options_ = cc->Options<LabelsToRenderDataCalculatorOptions>();
|
||||
num_colors_ = options_.color_size();
|
||||
label_height_px_ = std::ceil(options_.font_height_px() * kFontHeightScale);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status LabelsToRenderDataCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
if (cc->Inputs().HasTag("VIDEO_PRESTREAM") &&
|
||||
cc->InputTimestamp() == Timestamp::PreStream()) {
|
||||
const VideoHeader& video_header =
|
||||
cc->Inputs().Tag("VIDEO_PRESTREAM").Get<VideoHeader>();
|
||||
video_width_ = video_header.width;
|
||||
video_height_ = video_header.height;
|
||||
return ::mediapipe::OkStatus();
|
||||
} else {
|
||||
CHECK_EQ(options_.location(), LabelsToRenderDataCalculatorOptions::TOP_LEFT)
|
||||
<< "Only TOP_LEFT is supported without VIDEO_PRESTREAM.";
|
||||
}
|
||||
|
||||
std::vector<std::string> labels;
|
||||
std::vector<float> scores;
|
||||
if (cc->Inputs().HasTag("CLASSIFICATIONS")) {
|
||||
const ClassificationList& classifications =
|
||||
cc->Inputs().Tag("CLASSIFICATIONS").Get<ClassificationList>();
|
||||
labels.resize(classifications.classification_size());
|
||||
scores.resize(classifications.classification_size());
|
||||
for (int i = 0; i < classifications.classification_size(); ++i) {
|
||||
labels[i] = classifications.classification(i).label();
|
||||
scores[i] = classifications.classification(i).score();
|
||||
}
|
||||
} else {
|
||||
const std::vector<std::string>& label_vector =
|
||||
cc->Inputs().Tag("LABELS").Get<std::vector<std::string>>();
|
||||
std::vector<float> score_vector;
|
||||
if (cc->Inputs().HasTag("SCORES")) {
|
||||
score_vector = cc->Inputs().Tag("SCORES").Get<std::vector<float>>();
|
||||
}
|
||||
CHECK_EQ(label_vector.size(), score_vector.size());
|
||||
labels.resize(label_vector.size());
|
||||
scores.resize(label_vector.size());
|
||||
for (int i = 0; i < label_vector.size(); ++i) {
|
||||
labels[i] = label_vector[i];
|
||||
scores[i] = score_vector[i];
|
||||
}
|
||||
}
|
||||
|
||||
RenderData render_data;
|
||||
int num_label = std::min((int)labels.size(), options_.max_num_labels());
|
||||
int label_baseline_px = options_.vertical_offset_px();
|
||||
if (options_.location() == LabelsToRenderDataCalculatorOptions::TOP_LEFT) {
|
||||
label_baseline_px += label_height_px_;
|
||||
} else if (options_.location() ==
|
||||
LabelsToRenderDataCalculatorOptions::BOTTOM_LEFT) {
|
||||
label_baseline_px += video_height_ - label_height_px_ * (num_label - 1);
|
||||
}
|
||||
label_left_px_ = options_.horizontal_offset_px();
|
||||
for (int i = 0; i < num_label; ++i) {
|
||||
auto* label_annotation = render_data.add_render_annotations();
|
||||
label_annotation->set_thickness(options_.thickness());
|
||||
if (num_colors_ > 0) {
|
||||
*(label_annotation->mutable_color()) = options_.color(i % num_colors_);
|
||||
} else {
|
||||
label_annotation->mutable_color()->set_r(255);
|
||||
label_annotation->mutable_color()->set_g(0);
|
||||
label_annotation->mutable_color()->set_b(0);
|
||||
}
|
||||
|
||||
auto* text = label_annotation->mutable_text();
|
||||
std::string display_text = labels[i];
|
||||
if (cc->Inputs().HasTag("SCORES")) {
|
||||
absl::StrAppend(&display_text, ":", scores[i]);
|
||||
}
|
||||
text->set_display_text(display_text);
|
||||
text->set_font_height(options_.font_height_px());
|
||||
text->set_left(label_left_px_);
|
||||
text->set_baseline(label_baseline_px + i * label_height_px_);
|
||||
text->set_font_face(options_.font_face());
|
||||
}
|
||||
cc->Outputs()
|
||||
.Tag("RENDER_DATA")
|
||||
.AddPacket(MakePacket<RenderData>(render_data).At(cc->InputTimestamp()));
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,62 @@
|
||||
// Copyright 2019 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;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
import "mediapipe/util/color.proto";
|
||||
|
||||
message LabelsToRenderDataCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional LabelsToRenderDataCalculatorOptions ext = 271660364;
|
||||
}
|
||||
|
||||
// Colors for drawing the label(s).
|
||||
repeated Color color = 1;
|
||||
|
||||
// Thickness for drawing the label(s).
|
||||
optional double thickness = 2 [default = 2];
|
||||
|
||||
// The font height in absolute pixels.
|
||||
optional int32 font_height_px = 3 [default = 50];
|
||||
|
||||
// The offset of the starting text in horizontal direction in absolute pixels.
|
||||
optional int32 horizontal_offset_px = 7 [default = 0];
|
||||
// The offset of the starting text in vertical direction in absolute pixels.
|
||||
optional int32 vertical_offset_px = 8 [default = 0];
|
||||
|
||||
// The maximum number of labels to display.
|
||||
optional int32 max_num_labels = 4 [default = 1];
|
||||
|
||||
// Specifies the font for the text. Font must be one of the following from
|
||||
// OpenCV:
|
||||
// cv::FONT_HERSHEY_SIMPLEX (0)
|
||||
// cv::FONT_HERSHEY_PLAIN (1)
|
||||
// cv::FONT_HERSHEY_DUPLEX (2)
|
||||
// cv::FONT_HERSHEY_COMPLEX (3)
|
||||
// cv::FONT_HERSHEY_TRIPLEX (4)
|
||||
// cv::FONT_HERSHEY_COMPLEX_SMALL (5)
|
||||
// cv::FONT_HERSHEY_SCRIPT_SIMPLEX (6)
|
||||
// cv::FONT_HERSHEY_SCRIPT_COMPLEX (7)
|
||||
optional int32 font_face = 5 [default = 0];
|
||||
|
||||
// Label location.
|
||||
enum Location {
|
||||
TOP_LEFT = 0;
|
||||
BOTTOM_LEFT = 1;
|
||||
}
|
||||
optional Location location = 6 [default = TOP_LEFT];
|
||||
}
|
||||
@@ -49,7 +49,7 @@ constexpr char kLetterboxPaddingTag[] = "LETTERBOX_PADDING";
|
||||
// corresponding input image before letterboxing.
|
||||
//
|
||||
// Input:
|
||||
// LANDMARKS: An std::vector<NormalizedLandmark> representing landmarks on an
|
||||
// LANDMARKS: A NormalizedLandmarkList representing landmarks on an
|
||||
// letterboxed image.
|
||||
//
|
||||
// LETTERBOX_PADDING: An std::array<float, 4> representing the letterbox
|
||||
@@ -57,7 +57,7 @@ constexpr char kLetterboxPaddingTag[] = "LETTERBOX_PADDING";
|
||||
// image, normalized to [0.f, 1.f] by the letterboxed image dimensions.
|
||||
//
|
||||
// Output:
|
||||
// LANDMARKS: An std::vector<NormalizedLandmark> representing landmarks with
|
||||
// LANDMARKS: An NormalizedLandmarkList proto representing landmarks with
|
||||
// their locations adjusted to the letterbox-removed (non-padded) image.
|
||||
//
|
||||
// Usage example:
|
||||
@@ -74,10 +74,10 @@ class LandmarkLetterboxRemovalCalculator : public CalculatorBase {
|
||||
cc->Inputs().HasTag(kLetterboxPaddingTag))
|
||||
<< "Missing one or more input streams.";
|
||||
|
||||
cc->Inputs().Tag(kLandmarksTag).Set<std::vector<NormalizedLandmark>>();
|
||||
cc->Inputs().Tag(kLandmarksTag).Set<NormalizedLandmarkList>();
|
||||
cc->Inputs().Tag(kLetterboxPaddingTag).Set<std::array<float, 4>>();
|
||||
|
||||
cc->Outputs().Tag(kLandmarksTag).Set<std::vector<NormalizedLandmark>>();
|
||||
cc->Outputs().Tag(kLandmarksTag).Set<NormalizedLandmarkList>();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -94,8 +94,8 @@ class LandmarkLetterboxRemovalCalculator : public CalculatorBase {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
const auto& input_landmarks =
|
||||
cc->Inputs().Tag(kLandmarksTag).Get<std::vector<NormalizedLandmark>>();
|
||||
const NormalizedLandmarkList& input_landmarks =
|
||||
cc->Inputs().Tag(kLandmarksTag).Get<NormalizedLandmarkList>();
|
||||
const auto& letterbox_padding =
|
||||
cc->Inputs().Tag(kLetterboxPaddingTag).Get<std::array<float, 4>>();
|
||||
|
||||
@@ -104,24 +104,23 @@ class LandmarkLetterboxRemovalCalculator : public CalculatorBase {
|
||||
const float left_and_right = letterbox_padding[0] + letterbox_padding[2];
|
||||
const float top_and_bottom = letterbox_padding[1] + letterbox_padding[3];
|
||||
|
||||
auto output_landmarks =
|
||||
absl::make_unique<std::vector<NormalizedLandmark>>();
|
||||
for (const auto& landmark : input_landmarks) {
|
||||
NormalizedLandmark new_landmark;
|
||||
NormalizedLandmarkList output_landmarks;
|
||||
for (int i = 0; i < input_landmarks.landmark_size(); ++i) {
|
||||
const NormalizedLandmark& landmark = input_landmarks.landmark(i);
|
||||
NormalizedLandmark* new_landmark = output_landmarks.add_landmark();
|
||||
const float new_x = (landmark.x() - left) / (1.0f - left_and_right);
|
||||
const float new_y = (landmark.y() - top) / (1.0f - top_and_bottom);
|
||||
|
||||
new_landmark.set_x(new_x);
|
||||
new_landmark.set_y(new_y);
|
||||
new_landmark->set_x(new_x);
|
||||
new_landmark->set_y(new_y);
|
||||
// Keep z-coord as is.
|
||||
new_landmark.set_z(landmark.z());
|
||||
|
||||
output_landmarks->emplace_back(new_landmark);
|
||||
new_landmark->set_z(landmark.z());
|
||||
}
|
||||
|
||||
cc->Outputs()
|
||||
.Tag(kLandmarksTag)
|
||||
.Add(output_landmarks.release(), cc->InputTimestamp());
|
||||
.AddPacket(MakePacket<NormalizedLandmarkList>(output_landmarks)
|
||||
.At(cc->InputTimestamp()));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
@@ -43,10 +43,10 @@ CalculatorGraphConfig::Node GetDefaultNode() {
|
||||
TEST(LandmarkLetterboxRemovalCalculatorTest, PaddingLeftRight) {
|
||||
CalculatorRunner runner(GetDefaultNode());
|
||||
|
||||
auto landmarks = absl::make_unique<std::vector<NormalizedLandmark>>();
|
||||
landmarks->push_back(CreateLandmark(0.5f, 0.5f));
|
||||
landmarks->push_back(CreateLandmark(0.2f, 0.2f));
|
||||
landmarks->push_back(CreateLandmark(0.7f, 0.7f));
|
||||
auto landmarks = absl::make_unique<NormalizedLandmarkList>();
|
||||
*landmarks->add_landmark() = CreateLandmark(0.5f, 0.5f);
|
||||
*landmarks->add_landmark() = CreateLandmark(0.2f, 0.2f);
|
||||
*landmarks->add_landmark() = CreateLandmark(0.7f, 0.7f);
|
||||
runner.MutableInputs()
|
||||
->Tag("LANDMARKS")
|
||||
.packets.push_back(
|
||||
@@ -61,26 +61,28 @@ TEST(LandmarkLetterboxRemovalCalculatorTest, PaddingLeftRight) {
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Tag("LANDMARKS").packets;
|
||||
ASSERT_EQ(1, output.size());
|
||||
const auto& output_landmarks =
|
||||
output[0].Get<std::vector<NormalizedLandmark>>();
|
||||
const auto& output_landmarks = output[0].Get<NormalizedLandmarkList>();
|
||||
|
||||
EXPECT_EQ(output_landmarks.size(), 3);
|
||||
EXPECT_EQ(output_landmarks.landmark_size(), 3);
|
||||
|
||||
EXPECT_THAT(output_landmarks[0].x(), testing::FloatNear(0.6f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks[0].y(), testing::FloatNear(0.5f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks[1].x(), testing::FloatNear(0.0f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks[1].y(), testing::FloatNear(0.2f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks[2].x(), testing::FloatNear(1.0f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks[2].y(), testing::FloatNear(0.7f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(0).x(), testing::FloatNear(0.6f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(0).y(), testing::FloatNear(0.5f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(1).x(), testing::FloatNear(0.0f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(1).y(), testing::FloatNear(0.2f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(2).x(), testing::FloatNear(1.0f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(2).y(), testing::FloatNear(0.7f, 1e-5));
|
||||
}
|
||||
|
||||
TEST(LandmarkLetterboxRemovalCalculatorTest, PaddingTopBottom) {
|
||||
CalculatorRunner runner(GetDefaultNode());
|
||||
|
||||
auto landmarks = absl::make_unique<std::vector<NormalizedLandmark>>();
|
||||
landmarks->push_back(CreateLandmark(0.5f, 0.5f));
|
||||
landmarks->push_back(CreateLandmark(0.2f, 0.2f));
|
||||
landmarks->push_back(CreateLandmark(0.7f, 0.7f));
|
||||
auto landmarks = absl::make_unique<NormalizedLandmarkList>();
|
||||
NormalizedLandmark* landmark = landmarks->add_landmark();
|
||||
*landmark = CreateLandmark(0.5f, 0.5f);
|
||||
landmark = landmarks->add_landmark();
|
||||
*landmark = CreateLandmark(0.2f, 0.2f);
|
||||
landmark = landmarks->add_landmark();
|
||||
*landmark = CreateLandmark(0.7f, 0.7f);
|
||||
runner.MutableInputs()
|
||||
->Tag("LANDMARKS")
|
||||
.packets.push_back(
|
||||
@@ -95,17 +97,16 @@ TEST(LandmarkLetterboxRemovalCalculatorTest, PaddingTopBottom) {
|
||||
MP_ASSERT_OK(runner.Run()) << "Calculator execution failed.";
|
||||
const std::vector<Packet>& output = runner.Outputs().Tag("LANDMARKS").packets;
|
||||
ASSERT_EQ(1, output.size());
|
||||
const auto& output_landmarks =
|
||||
output[0].Get<std::vector<NormalizedLandmark>>();
|
||||
const auto& output_landmarks = output[0].Get<NormalizedLandmarkList>();
|
||||
|
||||
EXPECT_EQ(output_landmarks.size(), 3);
|
||||
EXPECT_EQ(output_landmarks.landmark_size(), 3);
|
||||
|
||||
EXPECT_THAT(output_landmarks[0].x(), testing::FloatNear(0.5f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks[0].y(), testing::FloatNear(0.6f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks[1].x(), testing::FloatNear(0.2f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks[1].y(), testing::FloatNear(0.0f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks[2].x(), testing::FloatNear(0.7f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks[2].y(), testing::FloatNear(1.0f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(0).x(), testing::FloatNear(0.5f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(0).y(), testing::FloatNear(0.6f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(1).x(), testing::FloatNear(0.2f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(1).y(), testing::FloatNear(0.0f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(2).x(), testing::FloatNear(0.7f, 1e-5));
|
||||
EXPECT_THAT(output_landmarks.landmark(2).y(), testing::FloatNear(1.0f, 1e-5));
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -47,13 +47,13 @@ constexpr char kRectTag[] = "NORM_RECT";
|
||||
// Projects normalized landmarks in a rectangle to its original coordinates. The
|
||||
// rectangle must also be in normalized coordinates.
|
||||
// Input:
|
||||
// NORM_LANDMARKS: An std::vector<NormalizedLandmark> representing landmarks
|
||||
// NORM_LANDMARKS: A NormalizedLandmarkList representing landmarks
|
||||
// in a normalized rectangle.
|
||||
// NORM_RECT: An NormalizedRect representing a normalized rectangle in image
|
||||
// coordinates.
|
||||
//
|
||||
// Output:
|
||||
// NORM_LANDMARKS: An std::vector<NormalizedLandmark> representing landmarks
|
||||
// NORM_LANDMARKS: A NormalizedLandmarkList representing landmarks
|
||||
// with their locations adjusted to the image.
|
||||
//
|
||||
// Usage example:
|
||||
@@ -70,10 +70,10 @@ class LandmarkProjectionCalculator : public CalculatorBase {
|
||||
cc->Inputs().HasTag(kRectTag))
|
||||
<< "Missing one or more input streams.";
|
||||
|
||||
cc->Inputs().Tag(kLandmarksTag).Set<std::vector<NormalizedLandmark>>();
|
||||
cc->Inputs().Tag(kLandmarksTag).Set<NormalizedLandmarkList>();
|
||||
cc->Inputs().Tag(kRectTag).Set<NormalizedRect>();
|
||||
|
||||
cc->Outputs().Tag(kLandmarksTag).Set<std::vector<NormalizedLandmark>>();
|
||||
cc->Outputs().Tag(kLandmarksTag).Set<NormalizedLandmarkList>();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -92,14 +92,14 @@ class LandmarkProjectionCalculator : public CalculatorBase {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
const auto& input_landmarks =
|
||||
cc->Inputs().Tag(kLandmarksTag).Get<std::vector<NormalizedLandmark>>();
|
||||
const NormalizedLandmarkList& input_landmarks =
|
||||
cc->Inputs().Tag(kLandmarksTag).Get<NormalizedLandmarkList>();
|
||||
const auto& input_rect = cc->Inputs().Tag(kRectTag).Get<NormalizedRect>();
|
||||
|
||||
auto output_landmarks =
|
||||
absl::make_unique<std::vector<NormalizedLandmark>>();
|
||||
for (const auto& landmark : input_landmarks) {
|
||||
NormalizedLandmark new_landmark;
|
||||
NormalizedLandmarkList output_landmarks;
|
||||
for (int i = 0; i < input_landmarks.landmark_size(); ++i) {
|
||||
const NormalizedLandmark& landmark = input_landmarks.landmark(i);
|
||||
NormalizedLandmark* new_landmark = output_landmarks.add_landmark();
|
||||
|
||||
const float x = landmark.x() - 0.5f;
|
||||
const float y = landmark.y() - 0.5f;
|
||||
@@ -110,17 +110,16 @@ class LandmarkProjectionCalculator : public CalculatorBase {
|
||||
new_x = new_x * input_rect.width() + input_rect.x_center();
|
||||
new_y = new_y * input_rect.height() + input_rect.y_center();
|
||||
|
||||
new_landmark.set_x(new_x);
|
||||
new_landmark.set_y(new_y);
|
||||
new_landmark->set_x(new_x);
|
||||
new_landmark->set_y(new_y);
|
||||
// Keep z-coord as is.
|
||||
new_landmark.set_z(landmark.z());
|
||||
|
||||
output_landmarks->emplace_back(new_landmark);
|
||||
new_landmark->set_z(landmark.z());
|
||||
}
|
||||
|
||||
cc->Outputs()
|
||||
.Tag(kLandmarksTag)
|
||||
.Add(output_landmarks.release(), cc->InputTimestamp());
|
||||
.AddPacket(MakePacket<NormalizedLandmarkList>(output_landmarks)
|
||||
.At(cc->InputTimestamp()));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
@@ -28,8 +28,7 @@ namespace {
|
||||
constexpr char kDetectionTag[] = "DETECTION";
|
||||
constexpr char kNormalizedLandmarksTag[] = "NORM_LANDMARKS";
|
||||
|
||||
Detection ConvertLandmarksToDetection(
|
||||
const std::vector<NormalizedLandmark>& landmarks) {
|
||||
Detection ConvertLandmarksToDetection(const NormalizedLandmarkList& landmarks) {
|
||||
Detection detection;
|
||||
LocationData* location_data = detection.mutable_location_data();
|
||||
|
||||
@@ -37,7 +36,8 @@ Detection ConvertLandmarksToDetection(
|
||||
float x_max = std::numeric_limits<float>::min();
|
||||
float y_min = std::numeric_limits<float>::max();
|
||||
float y_max = std::numeric_limits<float>::min();
|
||||
for (const auto& landmark : landmarks) {
|
||||
for (int i = 0; i < landmarks.landmark_size(); ++i) {
|
||||
const NormalizedLandmark& landmark = landmarks.landmark(i);
|
||||
x_min = std::min(x_min, landmark.x());
|
||||
x_max = std::max(x_max, landmark.x());
|
||||
y_min = std::min(y_min, landmark.y());
|
||||
@@ -67,7 +67,7 @@ Detection ConvertLandmarksToDetection(
|
||||
// to specify a subset of landmarks for creating the detection.
|
||||
//
|
||||
// Input:
|
||||
// NOMR_LANDMARKS: A vector of NormalizedLandmark.
|
||||
// NOMR_LANDMARKS: A NormalizedLandmarkList proto.
|
||||
//
|
||||
// Output:
|
||||
// DETECTION: A Detection proto.
|
||||
@@ -95,9 +95,7 @@ REGISTER_CALCULATOR(LandmarksToDetectionCalculator);
|
||||
RET_CHECK(cc->Inputs().HasTag(kNormalizedLandmarksTag));
|
||||
RET_CHECK(cc->Outputs().HasTag(kDetectionTag));
|
||||
// TODO: Also support converting Landmark to Detection.
|
||||
cc->Inputs()
|
||||
.Tag(kNormalizedLandmarksTag)
|
||||
.Set<std::vector<NormalizedLandmark>>();
|
||||
cc->Inputs().Tag(kNormalizedLandmarksTag).Set<NormalizedLandmarkList>();
|
||||
cc->Outputs().Tag(kDetectionTag).Set<Detection>();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
@@ -113,19 +111,20 @@ REGISTER_CALCULATOR(LandmarksToDetectionCalculator);
|
||||
|
||||
::mediapipe::Status LandmarksToDetectionCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
const auto& landmarks = cc->Inputs()
|
||||
.Tag(kNormalizedLandmarksTag)
|
||||
.Get<std::vector<NormalizedLandmark>>();
|
||||
RET_CHECK_GT(landmarks.size(), 0) << "Input landmark vector is empty.";
|
||||
const auto& landmarks =
|
||||
cc->Inputs().Tag(kNormalizedLandmarksTag).Get<NormalizedLandmarkList>();
|
||||
RET_CHECK_GT(landmarks.landmark_size(), 0)
|
||||
<< "Input landmark vector is empty.";
|
||||
|
||||
auto detection = absl::make_unique<Detection>();
|
||||
if (options_.selected_landmark_indices_size()) {
|
||||
std::vector<NormalizedLandmark> subset_landmarks(
|
||||
options_.selected_landmark_indices_size());
|
||||
for (int i = 0; i < subset_landmarks.size(); ++i) {
|
||||
RET_CHECK_LT(options_.selected_landmark_indices(i), landmarks.size())
|
||||
NormalizedLandmarkList subset_landmarks;
|
||||
for (int i = 0; i < options_.selected_landmark_indices_size(); ++i) {
|
||||
RET_CHECK_LT(options_.selected_landmark_indices(i),
|
||||
landmarks.landmark_size())
|
||||
<< "Index of landmark subset is out of range.";
|
||||
subset_landmarks[i] = landmarks[options_.selected_landmark_indices(i)];
|
||||
*subset_landmarks.add_landmark() =
|
||||
landmarks.landmark(options_.selected_landmark_indices(i));
|
||||
}
|
||||
*detection = ConvertLandmarksToDetection(subset_landmarks);
|
||||
} else {
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
// Copyright 2019 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.
|
||||
|
||||
// Copyright 2019 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 <cmath>
|
||||
#include <vector>
|
||||
|
||||
#include "Eigen/Core"
|
||||
#include "mediapipe/calculators/util/landmarks_to_floats_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/matrix.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr char kLandmarksTag[] = "NORM_LANDMARKS";
|
||||
constexpr char kFloatsTag[] = "FLOATS";
|
||||
constexpr char kMatrixTag[] = "MATRIX";
|
||||
|
||||
} // namespace
|
||||
|
||||
// Converts a vector of landmarks to a vector of floats or a matrix.
|
||||
// Input:
|
||||
// NORM_LANDMARKS: A NormalizedLandmarkList proto.
|
||||
//
|
||||
// Output:
|
||||
// FLOATS(optional): A vector of floats from flattened landmarks.
|
||||
// MATRIX(optional): A matrix of floats of the landmarks.
|
||||
//
|
||||
// Usage example:
|
||||
// node {
|
||||
// calculator: "LandmarksToFloatsCalculator"
|
||||
// input_stream: "NORM_LANDMARKS:landmarks"
|
||||
// output_stream: "MATRIX:landmark_matrix"
|
||||
// }
|
||||
class LandmarksToFloatsCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Tag(kLandmarksTag).Set<NormalizedLandmarkList>();
|
||||
RET_CHECK(cc->Outputs().HasTag(kFloatsTag) ||
|
||||
cc->Outputs().HasTag(kMatrixTag));
|
||||
if (cc->Outputs().HasTag(kFloatsTag)) {
|
||||
cc->Outputs().Tag(kFloatsTag).Set<std::vector<float>>();
|
||||
}
|
||||
if (cc->Outputs().HasTag(kMatrixTag)) {
|
||||
cc->Outputs().Tag(kMatrixTag).Set<Matrix>();
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::LandmarksToFloatsCalculatorOptions>();
|
||||
num_dimensions_ = options.num_dimensions();
|
||||
// Currently number of dimensions must be within [1, 3].
|
||||
RET_CHECK_GE(num_dimensions_, 1);
|
||||
RET_CHECK_LE(num_dimensions_, 3);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
// Only process if there's input landmarks.
|
||||
if (cc->Inputs().Tag(kLandmarksTag).IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
const auto& input_landmarks =
|
||||
cc->Inputs().Tag(kLandmarksTag).Get<NormalizedLandmarkList>();
|
||||
|
||||
if (cc->Outputs().HasTag(kFloatsTag)) {
|
||||
auto output_floats = absl::make_unique<std::vector<float>>();
|
||||
for (int i = 0; i < input_landmarks.landmark_size(); ++i) {
|
||||
const NormalizedLandmark& landmark = input_landmarks.landmark(i);
|
||||
output_floats->emplace_back(landmark.x());
|
||||
if (num_dimensions_ > 1) {
|
||||
output_floats->emplace_back(landmark.y());
|
||||
}
|
||||
if (num_dimensions_ > 2) {
|
||||
output_floats->emplace_back(landmark.z());
|
||||
}
|
||||
}
|
||||
|
||||
cc->Outputs()
|
||||
.Tag(kFloatsTag)
|
||||
.Add(output_floats.release(), cc->InputTimestamp());
|
||||
} else {
|
||||
auto output_matrix = absl::make_unique<Matrix>();
|
||||
output_matrix->setZero(num_dimensions_, input_landmarks.landmark_size());
|
||||
for (int i = 0; i < input_landmarks.landmark_size(); ++i) {
|
||||
(*output_matrix)(0, i) = input_landmarks.landmark(i).x();
|
||||
if (num_dimensions_ > 1) {
|
||||
(*output_matrix)(1, i) = input_landmarks.landmark(i).y();
|
||||
}
|
||||
if (num_dimensions_ > 2) {
|
||||
(*output_matrix)(2, i) = input_landmarks.landmark(i).z();
|
||||
}
|
||||
}
|
||||
cc->Outputs()
|
||||
.Tag(kMatrixTag)
|
||||
.Add(output_matrix.release(), cc->InputTimestamp());
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
int num_dimensions_ = 0;
|
||||
};
|
||||
REGISTER_CALCULATOR(LandmarksToFloatsCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright 2019 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;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
|
||||
message LandmarksToFloatsCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional LandmarksToFloatsCalculatorOptions ext = 274035660;
|
||||
}
|
||||
|
||||
// Number of dimensions to convert. Must within [1, 3].
|
||||
optional int32 num_dimensions = 1 [default = 2];
|
||||
}
|
||||
@@ -46,12 +46,13 @@ inline float Remap(float x, float lo, float hi, float scale) {
|
||||
return (x - lo) / (hi - lo + 1e-6) * scale;
|
||||
}
|
||||
|
||||
template <class LandmarkType>
|
||||
inline void GetMinMaxZ(const std::vector<LandmarkType>& landmarks, float* z_min,
|
||||
template <class LandmarkListType, class LandmarkType>
|
||||
inline void GetMinMaxZ(const LandmarkListType& landmarks, float* z_min,
|
||||
float* z_max) {
|
||||
*z_min = std::numeric_limits<float>::max();
|
||||
*z_max = std::numeric_limits<float>::min();
|
||||
for (const auto& landmark : landmarks) {
|
||||
for (int i = 0; i < landmarks.landmark_size(); ++i) {
|
||||
const LandmarkType& landmark = landmarks.landmark(i);
|
||||
*z_min = std::min(landmark.z(), *z_min);
|
||||
*z_max = std::max(landmark.z(), *z_max);
|
||||
}
|
||||
@@ -73,7 +74,7 @@ void SetColorSizeValueFromZ(float z, float z_min, float z_max,
|
||||
} // namespace
|
||||
|
||||
// A calculator that converts Landmark proto to RenderData proto for
|
||||
// visualization. The input should be std::vector<Landmark>. It is also possible
|
||||
// visualization. The input should be LandmarkList proto. It is also possible
|
||||
// to specify the connections between landmarks.
|
||||
//
|
||||
// Example config:
|
||||
@@ -121,11 +122,11 @@ class LandmarksToRenderDataCalculator : public CalculatorBase {
|
||||
const LandmarksToRenderDataCalculatorOptions& options, bool normalized,
|
||||
int gray_val1, int gray_val2, RenderData* render_data);
|
||||
|
||||
template <class LandmarkType>
|
||||
void AddConnections(const std::vector<LandmarkType>& landmarks,
|
||||
bool normalized, RenderData* render_data);
|
||||
template <class LandmarkType>
|
||||
void AddConnectionsWithDepth(const std::vector<LandmarkType>& landmarks,
|
||||
template <class LandmarkListType>
|
||||
void AddConnections(const LandmarkListType& landmarks, bool normalized,
|
||||
RenderData* render_data);
|
||||
template <class LandmarkListType>
|
||||
void AddConnectionsWithDepth(const LandmarkListType& landmarks,
|
||||
bool normalized, float min_z, float max_z,
|
||||
RenderData* render_data);
|
||||
|
||||
@@ -144,10 +145,10 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
"normalized landmarks.";
|
||||
|
||||
if (cc->Inputs().HasTag(kLandmarksTag)) {
|
||||
cc->Inputs().Tag(kLandmarksTag).Set<std::vector<Landmark>>();
|
||||
cc->Inputs().Tag(kLandmarksTag).Set<LandmarkList>();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kNormLandmarksTag)) {
|
||||
cc->Inputs().Tag(kNormLandmarksTag).Set<std::vector<NormalizedLandmark>>();
|
||||
cc->Inputs().Tag(kNormLandmarksTag).Set<NormalizedLandmarkList>();
|
||||
}
|
||||
cc->Outputs().Tag(kRenderDataTag).Set<RenderData>();
|
||||
return ::mediapipe::OkStatus();
|
||||
@@ -169,16 +170,17 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
float z_max = 0.f;
|
||||
|
||||
if (cc->Inputs().HasTag(kLandmarksTag)) {
|
||||
const auto& landmarks =
|
||||
cc->Inputs().Tag(kLandmarksTag).Get<std::vector<Landmark>>();
|
||||
const LandmarkList& landmarks =
|
||||
cc->Inputs().Tag(kLandmarksTag).Get<LandmarkList>();
|
||||
RET_CHECK_EQ(options_.landmark_connections_size() % 2, 0)
|
||||
<< "Number of entries in landmark connections must be a multiple of 2";
|
||||
if (visualize_depth) {
|
||||
GetMinMaxZ(landmarks, &z_min, &z_max);
|
||||
GetMinMaxZ<LandmarkList, Landmark>(landmarks, &z_min, &z_max);
|
||||
}
|
||||
// Only change rendering if there are actually z values other than 0.
|
||||
visualize_depth &= ((z_max - z_min) > 1e-3);
|
||||
for (const auto& landmark : landmarks) {
|
||||
for (int i = 0; i < landmarks.landmark_size(); ++i) {
|
||||
const Landmark& landmark = landmarks.landmark(i);
|
||||
auto* landmark_data_render =
|
||||
AddPointRenderData(options_, render_data.get());
|
||||
if (visualize_depth) {
|
||||
@@ -191,25 +193,27 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
landmark_data->set_y(landmark.y());
|
||||
}
|
||||
if (visualize_depth) {
|
||||
AddConnectionsWithDepth(landmarks, /*normalized=*/false, z_min, z_max,
|
||||
render_data.get());
|
||||
AddConnectionsWithDepth<LandmarkList>(landmarks, /*normalized=*/false,
|
||||
z_min, z_max, render_data.get());
|
||||
} else {
|
||||
AddConnections(landmarks, /*normalized=*/false, render_data.get());
|
||||
AddConnections<LandmarkList>(landmarks, /*normalized=*/false,
|
||||
render_data.get());
|
||||
}
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag(kNormLandmarksTag)) {
|
||||
const auto& landmarks = cc->Inputs()
|
||||
.Tag(kNormLandmarksTag)
|
||||
.Get<std::vector<NormalizedLandmark>>();
|
||||
const NormalizedLandmarkList& landmarks =
|
||||
cc->Inputs().Tag(kNormLandmarksTag).Get<NormalizedLandmarkList>();
|
||||
RET_CHECK_EQ(options_.landmark_connections_size() % 2, 0)
|
||||
<< "Number of entries in landmark connections must be a multiple of 2";
|
||||
if (visualize_depth) {
|
||||
GetMinMaxZ(landmarks, &z_min, &z_max);
|
||||
GetMinMaxZ<NormalizedLandmarkList, NormalizedLandmark>(landmarks, &z_min,
|
||||
&z_max);
|
||||
}
|
||||
// Only change rendering if there are actually z values other than 0.
|
||||
visualize_depth &= ((z_max - z_min) > 1e-3);
|
||||
for (const auto& landmark : landmarks) {
|
||||
for (int i = 0; i < landmarks.landmark_size(); ++i) {
|
||||
const NormalizedLandmark& landmark = landmarks.landmark(i);
|
||||
auto* landmark_data_render =
|
||||
AddPointRenderData(options_, render_data.get());
|
||||
if (visualize_depth) {
|
||||
@@ -222,10 +226,11 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
landmark_data->set_y(landmark.y());
|
||||
}
|
||||
if (visualize_depth) {
|
||||
AddConnectionsWithDepth(landmarks, /*normalized=*/true, z_min, z_max,
|
||||
render_data.get());
|
||||
AddConnectionsWithDepth<NormalizedLandmarkList>(
|
||||
landmarks, /*normalized=*/true, z_min, z_max, render_data.get());
|
||||
} else {
|
||||
AddConnections(landmarks, /*normalized=*/true, render_data.get());
|
||||
AddConnections<NormalizedLandmarkList>(landmarks, /*normalized=*/true,
|
||||
render_data.get());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -235,13 +240,13 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
template <class LandmarkType>
|
||||
template <class LandmarkListType>
|
||||
void LandmarksToRenderDataCalculator::AddConnectionsWithDepth(
|
||||
const std::vector<LandmarkType>& landmarks, bool normalized, float min_z,
|
||||
const LandmarkListType& landmarks, bool normalized, float min_z,
|
||||
float max_z, RenderData* render_data) {
|
||||
for (int i = 0; i < options_.landmark_connections_size(); i += 2) {
|
||||
const auto& ld0 = landmarks[options_.landmark_connections(i)];
|
||||
const auto& ld1 = landmarks[options_.landmark_connections(i + 1)];
|
||||
const auto& ld0 = landmarks.landmark(options_.landmark_connections(i));
|
||||
const auto& ld1 = landmarks.landmark(options_.landmark_connections(i + 1));
|
||||
const int gray_val1 =
|
||||
255 - static_cast<int>(Remap(ld0.z(), min_z, max_z, 255));
|
||||
const int gray_val2 =
|
||||
@@ -272,13 +277,13 @@ void LandmarksToRenderDataCalculator::AddConnectionToRenderData(
|
||||
connection_annotation->set_thickness(options.thickness());
|
||||
}
|
||||
|
||||
template <class LandmarkType>
|
||||
template <class LandmarkListType>
|
||||
void LandmarksToRenderDataCalculator::AddConnections(
|
||||
const std::vector<LandmarkType>& landmarks, bool normalized,
|
||||
const LandmarkListType& landmarks, bool normalized,
|
||||
RenderData* render_data) {
|
||||
for (int i = 0; i < options_.landmark_connections_size(); i += 2) {
|
||||
const auto& ld0 = landmarks[options_.landmark_connections(i)];
|
||||
const auto& ld1 = landmarks[options_.landmark_connections(i + 1)];
|
||||
const auto& ld0 = landmarks.landmark(options_.landmark_connections(i));
|
||||
const auto& ld1 = landmarks.landmark(options_.landmark_connections(i + 1));
|
||||
AddConnectionToRenderData(ld0.x(), ld0.y(), ld1.x(), ld1.y(), options_,
|
||||
normalized, render_data);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
// Copyright 2019 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 "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/file_helpers.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
|
||||
namespace mediapipe {
|
||||
// The calculator takes the path to the local file as an input side packet and
|
||||
// outputs the contents of that file.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "LocalFileContentsCalculator"
|
||||
// input_side_packet: "FILE_PATH:file_path"
|
||||
// output_side_packet: "CONTENTS:contents"
|
||||
// }
|
||||
class LocalFileContentsCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->InputSidePackets().Tag("FILE_PATH").Set<std::string>();
|
||||
cc->OutputSidePackets().Tag("CONTENTS").Set<std::string>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
std::string contents;
|
||||
MP_RETURN_IF_ERROR(mediapipe::file::GetContents(
|
||||
cc->InputSidePackets().Tag("FILE_PATH").Get<std::string>(), &contents));
|
||||
cc->OutputSidePackets()
|
||||
.Tag("CONTENTS")
|
||||
.Set(MakePacket<std::string>(std::move(contents)));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(LocalFileContentsCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -23,7 +23,9 @@ namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
constexpr char kNormRectTag[] = "NORM_RECT";
|
||||
constexpr char kNormRectsTag[] = "NORM_RECTS";
|
||||
constexpr char kRectTag[] = "RECT";
|
||||
constexpr char kRectsTag[] = "RECTS";
|
||||
constexpr char kImageSizeTag[] = "IMAGE_SIZE";
|
||||
|
||||
// Wraps around an angle in radians to within -M_PI and M_PI.
|
||||
@@ -72,17 +74,31 @@ REGISTER_CALCULATOR(RectTransformationCalculator);
|
||||
|
||||
::mediapipe::Status RectTransformationCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag(kNormRectTag) ^ cc->Inputs().HasTag(kRectTag));
|
||||
RET_CHECK_EQ((cc->Inputs().HasTag(kNormRectTag) ? 1 : 0) +
|
||||
(cc->Inputs().HasTag(kNormRectsTag) ? 1 : 0) +
|
||||
(cc->Inputs().HasTag(kRectTag) ? 1 : 0) +
|
||||
(cc->Inputs().HasTag(kRectsTag) ? 1 : 0),
|
||||
1);
|
||||
if (cc->Inputs().HasTag(kRectTag)) {
|
||||
cc->Inputs().Tag(kRectTag).Set<Rect>();
|
||||
cc->Outputs().Index(0).Set<Rect>();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kRectsTag)) {
|
||||
cc->Inputs().Tag(kRectsTag).Set<std::vector<Rect>>();
|
||||
cc->Outputs().Index(0).Set<std::vector<Rect>>();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kNormRectTag)) {
|
||||
RET_CHECK(cc->Inputs().HasTag(kImageSizeTag));
|
||||
cc->Inputs().Tag(kNormRectTag).Set<NormalizedRect>();
|
||||
cc->Inputs().Tag(kImageSizeTag).Set<std::pair<int, int>>();
|
||||
cc->Outputs().Index(0).Set<NormalizedRect>();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kNormRectsTag)) {
|
||||
RET_CHECK(cc->Inputs().HasTag(kImageSizeTag));
|
||||
cc->Inputs().Tag(kNormRectsTag).Set<std::vector<NormalizedRect>>();
|
||||
cc->Inputs().Tag(kImageSizeTag).Set<std::pair<int, int>>();
|
||||
cc->Outputs().Index(0).Set<std::vector<NormalizedRect>>();
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -105,7 +121,17 @@ REGISTER_CALCULATOR(RectTransformationCalculator);
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
MakePacket<Rect>(rect).At(cc->InputTimestamp()));
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag(kRectsTag) &&
|
||||
!cc->Inputs().Tag(kRectsTag).IsEmpty()) {
|
||||
auto rects = cc->Inputs().Tag(kRectsTag).Get<std::vector<Rect>>();
|
||||
auto output_rects = absl::make_unique<std::vector<Rect>>(rects.size());
|
||||
for (int i = 0; i < rects.size(); ++i) {
|
||||
output_rects->at(i) = rects[i];
|
||||
auto it = output_rects->begin() + i;
|
||||
TransformRect(&(*it));
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output_rects.release(), cc->InputTimestamp());
|
||||
}
|
||||
if (cc->Inputs().HasTag(kNormRectTag) &&
|
||||
!cc->Inputs().Tag(kNormRectTag).IsEmpty()) {
|
||||
auto rect = cc->Inputs().Tag(kNormRectTag).Get<NormalizedRect>();
|
||||
@@ -115,6 +141,21 @@ REGISTER_CALCULATOR(RectTransformationCalculator);
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
MakePacket<NormalizedRect>(rect).At(cc->InputTimestamp()));
|
||||
}
|
||||
if (cc->Inputs().HasTag(kNormRectsTag) &&
|
||||
!cc->Inputs().Tag(kNormRectsTag).IsEmpty()) {
|
||||
auto rects =
|
||||
cc->Inputs().Tag(kNormRectsTag).Get<std::vector<NormalizedRect>>();
|
||||
const auto& image_size =
|
||||
cc->Inputs().Tag(kImageSizeTag).Get<std::pair<int, int>>();
|
||||
auto output_rects =
|
||||
absl::make_unique<std::vector<NormalizedRect>>(rects.size());
|
||||
for (int i = 0; i < rects.size(); ++i) {
|
||||
output_rects->at(i) = rects[i];
|
||||
auto it = output_rects->begin() + i;
|
||||
TransformNormalizedRect(&(*it), image_size.first, image_size.second);
|
||||
}
|
||||
cc->Outputs().Index(0).Add(output_rects.release(), cc->InputTimestamp());
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -23,13 +23,13 @@
|
||||
|
||||
#include "mediapipe/calculators/util/top_k_scores_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/classification.pb.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/statusor.h"
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
|
||||
#if defined(MEDIAPIPE_LITE) || defined(__ANDROID__) || \
|
||||
(defined(__APPLE__) && !TARGET_OS_OSX)
|
||||
#if defined(MEDIAPIPE_MOBILE)
|
||||
#include "mediapipe/util/android/file/base/file.h"
|
||||
#include "mediapipe/util/android/file/base/helpers.h"
|
||||
#else
|
||||
@@ -37,8 +37,10 @@
|
||||
#endif
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator that takes a vector of scores and returns the indexes, scores,
|
||||
// labels of the top k elements.
|
||||
// labels of the top k elements, classification protos, and summary std::string
|
||||
// (in csv format).
|
||||
//
|
||||
// Usage example:
|
||||
// node {
|
||||
@@ -47,6 +49,8 @@ namespace mediapipe {
|
||||
// output_stream: "TOP_K_INDEXES:top_k_indexes"
|
||||
// output_stream: "TOP_K_SCORES:top_k_scores"
|
||||
// output_stream: "TOP_K_LABELS:top_k_labels"
|
||||
// output_stream: "TOP_K_CLASSIFICATIONS:top_k_classes"
|
||||
// output_stream: "SUMMARY:summary"
|
||||
// options: {
|
||||
// [mediapipe.TopKScoresCalculatorOptions.ext] {
|
||||
// top_k: 5
|
||||
@@ -69,6 +73,7 @@ class TopKScoresCalculator : public CalculatorBase {
|
||||
int top_k_ = -1;
|
||||
float threshold_ = 0.0;
|
||||
std::unordered_map<int, std::string> label_map_;
|
||||
bool label_map_loaded_ = false;
|
||||
};
|
||||
REGISTER_CALCULATOR(TopKScoresCalculator);
|
||||
|
||||
@@ -84,6 +89,12 @@ REGISTER_CALCULATOR(TopKScoresCalculator);
|
||||
if (cc->Outputs().HasTag("TOP_K_LABELS")) {
|
||||
cc->Outputs().Tag("TOP_K_LABELS").Set<std::vector<std::string>>();
|
||||
}
|
||||
if (cc->Outputs().HasTag("CLASSIFICATIONS")) {
|
||||
cc->Outputs().Tag("CLASSIFICATIONS").Set<ClassificationList>();
|
||||
}
|
||||
if (cc->Outputs().HasTag("SUMMARY")) {
|
||||
cc->Outputs().Tag("SUMMARY").Set<std::string>();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -149,7 +160,7 @@ REGISTER_CALCULATOR(TopKScoresCalculator);
|
||||
reverse(top_k_indexes.begin(), top_k_indexes.end());
|
||||
reverse(top_k_scores.begin(), top_k_scores.end());
|
||||
|
||||
if (cc->Outputs().HasTag("TOP_K_LABELS")) {
|
||||
if (label_map_loaded_) {
|
||||
for (int index : top_k_indexes) {
|
||||
top_k_labels.push_back(label_map_[index]);
|
||||
}
|
||||
@@ -172,6 +183,35 @@ REGISTER_CALCULATOR(TopKScoresCalculator);
|
||||
.AddPacket(MakePacket<std::vector<std::string>>(top_k_labels)
|
||||
.At(cc->InputTimestamp()));
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("SUMMARY")) {
|
||||
std::vector<std::string> results;
|
||||
for (int index = 0; index < top_k_indexes.size(); ++index) {
|
||||
if (label_map_loaded_) {
|
||||
results.push_back(
|
||||
absl::StrCat(top_k_labels[index], ":", top_k_scores[index]));
|
||||
} else {
|
||||
results.push_back(
|
||||
absl::StrCat(top_k_indexes[index], ":", top_k_scores[index]));
|
||||
}
|
||||
}
|
||||
cc->Outputs().Tag("SUMMARY").AddPacket(
|
||||
MakePacket<std::string>(absl::StrJoin(results, ","))
|
||||
.At(cc->InputTimestamp()));
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("TOP_K_CLASSIFICATION")) {
|
||||
auto classification_list = absl::make_unique<ClassificationList>();
|
||||
for (int index = 0; index < top_k_indexes.size(); ++index) {
|
||||
Classification* classification =
|
||||
classification_list->add_classification();
|
||||
classification->set_index(top_k_indexes[index]);
|
||||
classification->set_score(top_k_scores[index]);
|
||||
if (label_map_loaded_) {
|
||||
classification->set_label(top_k_labels[index]);
|
||||
}
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -188,6 +228,7 @@ REGISTER_CALCULATOR(TopKScoresCalculator);
|
||||
while (std::getline(stream, line)) {
|
||||
label_map_[i++] = line;
|
||||
}
|
||||
label_map_loaded_ = true;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
|
||||
@@ -13,12 +13,16 @@
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
load(
|
||||
"//mediapipe/framework/tool:mediapipe_graph.bzl",
|
||||
"mediapipe_binary_graph",
|
||||
)
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(default_visibility = ["//visibility:private"])
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "flow_to_image_calculator_proto",
|
||||
srcs = ["flow_to_image_calculator.proto"],
|
||||
@@ -33,6 +37,86 @@ proto_library(
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "motion_analysis_calculator_proto",
|
||||
srcs = ["motion_analysis_calculator.proto"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/util/tracking:motion_analysis_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "flow_packager_calculator_proto",
|
||||
srcs = ["flow_packager_calculator.proto"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/util/tracking:flow_packager_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "box_tracker_calculator_proto",
|
||||
srcs = ["box_tracker_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/util/tracking:box_tracker_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "video_pre_stream_calculator_proto",
|
||||
srcs = ["video_pre_stream_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "motion_analysis_calculator_cc_proto",
|
||||
srcs = ["motion_analysis_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/util/tracking:motion_analysis_cc_proto",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":motion_analysis_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "flow_packager_calculator_cc_proto",
|
||||
srcs = ["flow_packager_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/util/tracking:flow_packager_cc_proto",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":flow_packager_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "box_tracker_calculator_cc_proto",
|
||||
srcs = ["box_tracker_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/util/tracking:box_tracker_cc_proto",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":box_tracker_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "video_pre_stream_calculator_cc_proto",
|
||||
srcs = ["video_pre_stream_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":video_pre_stream_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "flow_to_image_calculator_cc_proto",
|
||||
srcs = ["flow_to_image_calculator.proto"],
|
||||
@@ -52,9 +136,7 @@ mediapipe_cc_proto_library(
|
||||
cc_library(
|
||||
name = "flow_to_image_calculator",
|
||||
srcs = ["flow_to_image_calculator.cc"],
|
||||
visibility = [
|
||||
"//visibility:public",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/video:flow_to_image_calculator_cc_proto",
|
||||
"//mediapipe/calculators/video/tool:flow_quantizer_model",
|
||||
@@ -129,10 +211,121 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "motion_analysis_calculator",
|
||||
srcs = ["motion_analysis_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":motion_analysis_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/util/tracking:camera_motion",
|
||||
"//mediapipe/util/tracking:camera_motion_cc_proto",
|
||||
"//mediapipe/util/tracking:frame_selection_cc_proto",
|
||||
"//mediapipe/util/tracking:motion_analysis",
|
||||
"//mediapipe/util/tracking:motion_estimation",
|
||||
"//mediapipe/util/tracking:motion_models",
|
||||
"//mediapipe/util/tracking:region_flow_cc_proto",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "flow_packager_calculator",
|
||||
srcs = ["flow_packager_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":flow_packager_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/util/tracking:camera_motion_cc_proto",
|
||||
"//mediapipe/util/tracking:flow_packager",
|
||||
"//mediapipe/util/tracking:region_flow_cc_proto",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/strings:str_format",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "box_tracker_calculator",
|
||||
srcs = ["box_tracker_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":box_tracker_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/formats:video_stream_header", # fixdeps: keep -- required for exobazel build.
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/tool:options_util",
|
||||
"//mediapipe/util/tracking",
|
||||
"//mediapipe/util/tracking:box_tracker",
|
||||
"//mediapipe/util/tracking:tracking_visualization_utilities",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "tracked_detection_manager_calculator",
|
||||
srcs = ["tracked_detection_manager_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/formats:location_data_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/util/tracking",
|
||||
"//mediapipe/util/tracking:box_tracker",
|
||||
"//mediapipe/util/tracking:tracked_detection",
|
||||
"//mediapipe/util/tracking:tracked_detection_manager",
|
||||
"//mediapipe/util/tracking:tracking_visualization_utilities",
|
||||
"@com_google_absl//absl/container:node_hash_map",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "video_pre_stream_calculator",
|
||||
srcs = ["video_pre_stream_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":video_pre_stream_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
filegroup(
|
||||
name = "test_videos",
|
||||
srcs = [
|
||||
"testdata/format_FLV_H264_AAC.video",
|
||||
"testdata/format_MKV_VP8_VORBIS.video",
|
||||
"testdata/format_MP4_AVC720P_AAC.video",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "opencv_video_decoder_calculator_test",
|
||||
srcs = ["opencv_video_decoder_calculator_test.cc"],
|
||||
data = ["//mediapipe/calculators/video/testdata:test_videos"],
|
||||
data = [":test_videos"],
|
||||
deps = [
|
||||
":opencv_video_decoder_calculator",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
@@ -151,7 +344,7 @@ cc_test(
|
||||
cc_test(
|
||||
name = "opencv_video_encoder_calculator_test",
|
||||
srcs = ["opencv_video_encoder_calculator_test.cc"],
|
||||
data = ["//mediapipe/calculators/video/testdata:test_videos"],
|
||||
data = [":test_videos"],
|
||||
deps = [
|
||||
":opencv_video_decoder_calculator",
|
||||
":opencv_video_encoder_calculator",
|
||||
@@ -175,7 +368,6 @@ cc_test(
|
||||
cc_test(
|
||||
name = "tvl1_optical_flow_calculator_test",
|
||||
srcs = ["tvl1_optical_flow_calculator_test.cc"],
|
||||
data = ["//mediapipe/calculators/image/testdata:test_images"],
|
||||
deps = [
|
||||
":tvl1_optical_flow_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
@@ -190,3 +382,78 @@ cc_test(
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
],
|
||||
)
|
||||
|
||||
MEDIAPIPE_DEPS = [
|
||||
"//mediapipe/calculators/video:box_tracker_calculator",
|
||||
"//mediapipe/calculators/video:flow_packager_calculator",
|
||||
"//mediapipe/calculators/video:motion_analysis_calculator",
|
||||
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
|
||||
"//mediapipe/framework/stream_handler:sync_set_input_stream_handler",
|
||||
]
|
||||
|
||||
mediapipe_binary_graph(
|
||||
name = "parallel_tracker_binarypb",
|
||||
graph = "testdata/parallel_tracker_graph.pbtxt",
|
||||
output_name = "testdata/parallel_tracker.binarypb",
|
||||
visibility = ["//visibility:public"],
|
||||
deps = MEDIAPIPE_DEPS,
|
||||
)
|
||||
|
||||
mediapipe_binary_graph(
|
||||
name = "tracker_binarypb",
|
||||
graph = "testdata/tracker_graph.pbtxt",
|
||||
output_name = "testdata/tracker.binarypb",
|
||||
visibility = ["//visibility:public"],
|
||||
deps = MEDIAPIPE_DEPS,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "tracking_graph_test",
|
||||
size = "small",
|
||||
srcs = ["tracking_graph_test.cc"],
|
||||
copts = ["-DPARALLEL_INVOKER_ACTIVE"] + select({
|
||||
"//mediapipe:apple": [],
|
||||
"//mediapipe:android": [],
|
||||
"//conditions:default": [],
|
||||
}),
|
||||
data = [
|
||||
":testdata/lenna.png",
|
||||
":testdata/parallel_tracker.binarypb",
|
||||
":testdata/tracker.binarypb",
|
||||
],
|
||||
deps = [
|
||||
":box_tracker_calculator",
|
||||
":box_tracker_calculator_cc_proto",
|
||||
":flow_packager_calculator",
|
||||
":motion_analysis_calculator",
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/deps:file_path",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/port:advanced_proto",
|
||||
"//mediapipe/framework/port:core_proto",
|
||||
"//mediapipe/framework/port:file_helpers",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:opencv_highgui",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
|
||||
"//mediapipe/framework/stream_handler:sync_set_input_stream_handler",
|
||||
"//mediapipe/util/tracking:box_tracker_cc_proto",
|
||||
"//mediapipe/util/tracking:tracking_cc_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "video_pre_stream_calculator_test",
|
||||
srcs = ["video_pre_stream_calculator_test.cc"],
|
||||
deps = [
|
||||
":video_pre_stream_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/port:status",
|
||||
],
|
||||
)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user