Compare commits

..
16 Commits
Author SHA1 Message Date
MediaPipe TeamandHadon Nash ae6be10afe Project import generated by Copybara.
GitOrigin-RevId: 0517756260533d374df93679965ca662d0ec6943
2020-01-10 13:13:24 -08:00
MediaPipe Teamandjqtang 38ee2603a7 Project import generated by Copybara.
GitOrigin-RevId: 87e46800807001e01d686fd7bcc2533714556920
2019-12-09 13:11:22 -08:00
MediaPipe Teamandjqtang 86b3283b2f Project import generated by Copybara.
GitOrigin-RevId: 831b7eb6038549a3a5047e7a113d6a11956e2de9
2019-12-06 16:17:14 -08:00
MediaPipe Teamandjqtang 7d470a1335 Project import generated by Copybara.
GitOrigin-RevId: 398d8577074c6e93041c01ed34bd6f27b2773c4f
2019-12-06 16:07:44 -08:00
MediaPipe Teamandmgyong d16cc3be5b Project import generated by Copybara.
GitOrigin-RevId: d91373b4d4d10abef49cab410caa6aadf0875049
2019-12-06 15:57:20 -08:00
MediaPipe Teamandjqtang 137867d088 Project import generated by Copybara.
GitOrigin-RevId: e3566e5029af25b0fc4b1071a49e49ae20aa5df6
2019-12-02 17:54:10 -08:00
MediaPipe Teamandmgyong 446d7cf6b6 Project import generated by Copybara.
GitOrigin-RevId: b02a6442fa6234cd2c15fa19f09accd8767adbee
2019-11-21 14:48:32 -08:00
MediaPipe Teamandmgyong 90f72bd851 Project import generated by Copybara.
GitOrigin-RevId: 5aa039c4a51ab7b4a1c58c17ad13af4c833e25e7
2019-11-21 14:35:46 -08:00
MediaPipe Teamandmgyong 4285aeddfc Project import generated by Copybara.
GitOrigin-RevId: 651ba7a75bb696877570a8a1b4244b34d59088f8
2019-11-21 14:24:17 -08:00
MediaPipe Teamandmgyong 37287925b0 Project import generated by Copybara.
GitOrigin-RevId: ba1d851bc868c2f8037a6fa96ee90e4b8ab9bd40
2019-11-21 14:10:52 -08:00
MediaPipe Teamandmgyong 48bcbb115f Project import generated by Copybara.
GitOrigin-RevId: 50714fe28298d7b707eff7304547d89d6ec34a54
2019-11-21 13:20:47 -08:00
MediaPipe Teamandjqtang 9437483827 Project import generated by Copybara.
GitOrigin-RevId: 5aca6b3f07b67e09988a901f50f595ca5f566e67
2019-11-15 13:10:50 -08:00
MediaPipe Teamandjqtang d030c13931 Project import generated by Copybara.
GitOrigin-RevId: dab808e56f90d1ad93e6014869f7fe4646b67fe0
2019-11-11 22:59:01 -08:00
MediaPipe Teamandjqtang fce372d153 Project import generated by Copybara.
GitOrigin-RevId: ac03a471f5b9df34de46dd684202e4365c5ceac3
2019-10-29 15:59:27 -07:00
MediaPipe Teamandjqtang c6fea4c9d9 Project import generated by Copybara.
GitOrigin-RevId: 1a0caa03bbf3673dbe772c8045b687c6b6821bcc
2019-10-25 14:50:09 -07:00
MediaPipe Teamandjqtang 259b48e082 Project import generated by Copybara.
GitOrigin-RevId: b137378673f7d66d41bcd46e4fc3a0d9ef254894
2019-10-25 14:29:15 -07:00
542 changed files with 75090 additions and 1598 deletions
+5 -2
View File
@@ -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
View File
@@ -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" && \
+16 -3
View File
@@ -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)
![hand_tracking](mediapipe/docs/images/mobile/hand_tracking_3d_android_gpu_small.gif)
![multi-hand_tracking](mediapipe/docs/images/mobile/multi_hand_tracking_android_gpu_small.gif)
![face_detection](mediapipe/docs/images/mobile/face_detection_android_gpu_small.gif)
![hair_segmentation](mediapipe/docs/images/mobile/hair_segmentation_android_gpu_small.gif)
![object_detection](mediapipe/docs/images/mobile/object_detection_android_gpu_small.gif)
![object_tracking](mediapipe/docs/images/mobile/object_tracking_android_gpu_small.gif)
## 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.
+51 -33
View File
@@ -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];
}
+273 -7
View File
@@ -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
+6 -4
View File
@@ -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
+124 -56
View File
@@ -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
+8 -2
View File
@@ -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,
+265 -4
View File
@@ -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();
}
+275 -8
View File
@@ -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