Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7bad8fce62 | ||
|
|
024f7bf0f1 | ||
|
|
16e5d7242d | ||
|
|
4c68eb4a70 | ||
|
|
a3d36eee32 | ||
|
|
1722d4b8a2 |
@@ -12,9 +12,9 @@ build --copt='-Wno-comment'
|
||||
build --copt='-Wno-return-type'
|
||||
build --copt='-Wno-unused-local-typedefs'
|
||||
build --copt='-Wno-ignored-attributes'
|
||||
# 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
|
||||
|
||||
# Tensorflow needs remote repo
|
||||
build --experimental_repo_remote_exec
|
||||
|
||||
# Sets the default Apple platform to macOS.
|
||||
build --apple_platform_type=macos
|
||||
@@ -23,6 +23,10 @@ build --apple_platform_type=macos
|
||||
build --apple_generate_dsym
|
||||
|
||||
# Android configs.
|
||||
# Note: the documentation tells us to use @androidndk//:default_crosstool, but
|
||||
# the automatic configuration transition uses //external:android/crosstool.
|
||||
# Using it here works and spares us from having two different config_settings
|
||||
# for Android.
|
||||
build:android --crosstool_top=//external:android/crosstool
|
||||
build:android --host_crosstool_top=@bazel_tools//tools/cpp:toolchain
|
||||
build:android --linkopt=-landroid
|
||||
|
||||
@@ -48,7 +48,7 @@ RUN pip install future
|
||||
RUN pip3 install six
|
||||
|
||||
# Install bazel
|
||||
ARG BAZEL_VERSION=1.1.0
|
||||
ARG BAZEL_VERSION=2.0.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" && \
|
||||
|
||||
@@ -9,26 +9,34 @@
|
||||
|
||||
## ML Solutions in MediaPipe
|
||||
|
||||
* [Face Detection](mediapipe/docs/face_detection_mobile_gpu.md) [[Web Demo]](https://viz.mediapipe.dev/runner/demos/face_detection/face_detection.html)
|
||||
* [Face Detection](mediapipe/docs/face_detection_mobile_gpu.md) [(web demo)](https://viz.mediapipe.dev/runner/demos/face_detection/face_detection.html)
|
||||
* [Face Mesh](mediapipe/docs/face_mesh_mobile_gpu.md)
|
||||
* [Hand Detection](mediapipe/docs/hand_detection_mobile_gpu.md)
|
||||
* [Hand Tracking](mediapipe/docs/hand_tracking_mobile_gpu.md) [(web demo)](https://viz.mediapipe.dev/runner/demos/hand_tracking/hand_tracking.html)
|
||||
* [Multi-hand Tracking](mediapipe/docs/multi_hand_tracking_mobile_gpu.md)
|
||||
* [Hand Tracking](mediapipe/docs/hand_tracking_mobile_gpu.md) [[Web Demo]](https://viz.mediapipe.dev/runner/demos/hand_tracking/hand_tracking.html)
|
||||
* [Hair Segmentation](mediapipe/docs/hair_segmentation_mobile_gpu.md) [[Web Demo]](https://viz.mediapipe.dev/runner/demos/hair_segmentation/hair_segmentation.html)
|
||||
* [Hair Segmentation](mediapipe/docs/hair_segmentation_mobile_gpu.md) [(web demo)](https://viz.mediapipe.dev/runner/demos/hair_segmentation/hair_segmentation.html)
|
||||
* [Object Detection](mediapipe/docs/object_detection_mobile_gpu.md)
|
||||
* [Object Detection and Tracking](mediapipe/docs/object_tracking_mobile_gpu.md)
|
||||
* [Objectron: 3D Object Detection and Tracking](mediapipe/docs/objectron_mobile_gpu.md)
|
||||
* [AutoFlip](mediapipe/docs/autoflip.md)
|
||||
* [AutoFlip: Intelligent Video Reframing](mediapipe/docs/autoflip.md)
|
||||
* [KNIFT: Template Matching with Neural Image Features](mediapipe/docs/template_matching_mobile_cpu.md)
|
||||
|
||||

|
||||

|
||||

|
||||

|
||||

|
||||

|
||||

|
||||

|
||||

|
||||

|
||||

|
||||

|
||||
|
||||
## Installation
|
||||
Follow these [instructions](mediapipe/docs/install.md).
|
||||
|
||||
## Getting started
|
||||
See mobile, desktop and Google Coral [examples](mediapipe/docs/examples.md).
|
||||
See mobile, desktop, web and Google Coral [examples](mediapipe/docs/examples.md).
|
||||
|
||||
Check out some web demos [[Edge detection]](https://viz.mediapipe.dev/runner/demos/edge_detection/edge_detection.html) [[Face detection]](https://viz.mediapipe.dev/runner/demos/face_detection/face_detection.html) [[Hand Tracking]](https://viz.mediapipe.dev/runner/demos/hand_tracking/hand_tracking.html)
|
||||
|
||||
@@ -40,10 +48,15 @@ Check out the [Examples page](https://mediapipe.readthedocs.io/en/latest/example
|
||||
## Visualizing MediaPipe graphs
|
||||
A web-based visualizer is hosted on [viz.mediapipe.dev](https://viz.mediapipe.dev/). Please also see instructions [here](mediapipe/docs/visualizer.md).
|
||||
|
||||
## Google Open Source Code search
|
||||
Search MediaPipe Github repository using [Google Open Source code search](https://t.co/LSZnbMUUnT?amp=1)
|
||||
|
||||
## Videos
|
||||
* [YouTube Channel](https://www.youtube.com/channel/UCObqmpuSMx-usADtL_qdMAw)
|
||||
|
||||
## Publications
|
||||
* [MediaPipe KNIFT: Template-based Feature Matching](https://mediapipe.page.link/knift-blog)
|
||||
* [Alfred Camera: Smart camera features using MediaPipe](https://developers.googleblog.com/2020/03/alfred-camera-smart-camera-features-using-mediapipe.html)
|
||||
* [MediaPipe Objectron: Real-time 3D Object Detection on Mobile Devices](https://mediapipe.page.link/objectron-aiblog)
|
||||
* [AutoFlip: An Open Source Framework for Intelligent Video Reframing](https://mediapipe.page.link/autoflip)
|
||||
* [Google Developer Blog: MediaPipe on the Web](https://mediapipe.page.link/webdevblog)
|
||||
@@ -52,6 +65,7 @@ A web-based visualizer is hosted on [viz.mediapipe.dev](https://viz.mediapipe.de
|
||||
* [MediaPipe: A Framework for Building Perception Pipelines](https://arxiv.org/abs/1906.08172)
|
||||
|
||||
## Events
|
||||
* [MediaPipe Seattle Meetup, Google Building Waterside, 13 Feb 2020](https://mediapipe.page.link/seattle2020)
|
||||
* [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)
|
||||
|
||||
@@ -2,16 +2,15 @@ workspace(name = "mediapipe")
|
||||
|
||||
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
|
||||
|
||||
skylib_version = "0.8.0"
|
||||
skylib_version = "0.9.0"
|
||||
http_archive(
|
||||
name = "bazel_skylib",
|
||||
type = "tar.gz",
|
||||
url = "https://github.com/bazelbuild/bazel-skylib/releases/download/{}/bazel-skylib.{}.tar.gz".format (skylib_version, skylib_version),
|
||||
sha256 = "2ef429f5d7ce7111263289644d233707dba35e39696377ebab8b0bc701f7818e",
|
||||
url = "https://github.com/bazelbuild/bazel-skylib/releases/download/{}/bazel_skylib-{}.tar.gz".format (skylib_version, skylib_version),
|
||||
sha256 = "1dde365491125a3db70731e25658dfdd3bc5dbdfd11b840b3e987ecf043c7ca0",
|
||||
)
|
||||
load("@bazel_skylib//lib:versions.bzl", "versions")
|
||||
versions.check(minimum_bazel_version = "1.0.0",
|
||||
maximum_bazel_version = "1.2.1")
|
||||
versions.check(minimum_bazel_version = "2.0.0")
|
||||
|
||||
|
||||
# ABSL cpp library lts_2020_02_25
|
||||
@@ -93,11 +92,13 @@ http_archive(
|
||||
build_file = "@//third_party:libyuv.BUILD",
|
||||
)
|
||||
|
||||
# Note: protobuf-javalite is no longer released as a separate download, it's included in the main Java download.
|
||||
# ...but the Java download is currently broken, so we use the "source" download.
|
||||
http_archive(
|
||||
name = "com_google_protobuf_javalite",
|
||||
sha256 = "79d102c61e2a479a0b7e5fc167bcfaa4832a0c6aad4a75fa7da0480564931bcc",
|
||||
strip_prefix = "protobuf-384989534b2246d413dbcd750744faab2607b516",
|
||||
urls = ["https://github.com/google/protobuf/archive/384989534b2246d413dbcd750744faab2607b516.zip"],
|
||||
sha256 = "a79d19dcdf9139fa4b81206e318e33d245c4c9da1ffed21c87288ed4380426f9",
|
||||
strip_prefix = "protobuf-3.11.4",
|
||||
urls = ["https://github.com/protocolbuffers/protobuf/archive/v3.11.4.tar.gz"],
|
||||
)
|
||||
|
||||
http_archive(
|
||||
@@ -117,19 +118,18 @@ http_archive(
|
||||
],
|
||||
)
|
||||
|
||||
# 2020-02-12
|
||||
# The last commit before TensorFlow switched to Bazel 2.0
|
||||
_TENSORFLOW_GIT_COMMIT = "77e9ffb9b2bfb1a4f7056e62d84039626923e328"
|
||||
_TENSORFLOW_SHA256= "176ccd82f7dd17c5e117b50d353603b129c7a6ccbfebd522ca47cc2a40f33f13"
|
||||
# 2020-04-01
|
||||
_TENSORFLOW_GIT_COMMIT = "805e47cea96c7e8c6fccf494d40a2392dc99fdd8"
|
||||
_TENSORFLOW_SHA256= "9ee3ae604c2e1345ac60345becee6d659364721513f9cb8652eb2e7138320ca5"
|
||||
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,
|
||||
],
|
||||
# A compatibility patch
|
||||
patches = [
|
||||
"@//third_party:org_tensorflow_528e22eae8bf3206189a066032c66e9e5c9b4a61.diff"
|
||||
"@//third_party:org_tensorflow_compatibility_fixes.diff",
|
||||
"@//third_party:org_tensorflow_protobuf_updates.diff",
|
||||
],
|
||||
patch_args = [
|
||||
"-p1",
|
||||
@@ -154,10 +154,6 @@ http_archive(
|
||||
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",
|
||||
build_file = "@//third_party:opencv_linux.BUILD",
|
||||
@@ -170,7 +166,6 @@ new_local_repository(
|
||||
path = "/usr"
|
||||
)
|
||||
|
||||
# Please run $ brew install opencv@3
|
||||
new_local_repository(
|
||||
name = "macos_opencv",
|
||||
build_file = "@//third_party:opencv_macos.BUILD",
|
||||
@@ -203,79 +198,6 @@ http_archive(
|
||||
url = "https://github.com/opencv/opencv/releases/download/3.2.0/opencv-3.2.0-ios-framework.zip",
|
||||
)
|
||||
|
||||
RULES_JVM_EXTERNAL_TAG = "2.2"
|
||||
RULES_JVM_EXTERNAL_SHA = "f1203ce04e232ab6fdd81897cf0ff76f2c04c0741424d192f28e65ae752ce2d6"
|
||||
|
||||
http_archive(
|
||||
name = "rules_jvm_external",
|
||||
strip_prefix = "rules_jvm_external-%s" % RULES_JVM_EXTERNAL_TAG,
|
||||
sha256 = RULES_JVM_EXTERNAL_SHA,
|
||||
url = "https://github.com/bazelbuild/rules_jvm_external/archive/%s.zip" % RULES_JVM_EXTERNAL_TAG,
|
||||
)
|
||||
|
||||
load("@rules_jvm_external//:defs.bzl", "maven_install")
|
||||
|
||||
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",
|
||||
"https://repo1.maven.org/maven2",
|
||||
],
|
||||
)
|
||||
|
||||
maven_server(
|
||||
name = "google_server",
|
||||
url = "https://dl.google.com/dl/android/maven2",
|
||||
)
|
||||
|
||||
maven_jar(
|
||||
name = "androidx_lifecycle",
|
||||
artifact = "androidx.lifecycle:lifecycle-common:2.0.0",
|
||||
sha1 = "e070ffae07452331bc5684734fce6831d531785c",
|
||||
server = "google_server",
|
||||
)
|
||||
|
||||
maven_jar(
|
||||
name = "androidx_concurrent_futures",
|
||||
artifact = "androidx.concurrent:concurrent-futures:1.0.0-alpha03",
|
||||
sha1 = "b528df95c7e2fefa2210c0c742bf3e491c1818ae",
|
||||
server = "google_server",
|
||||
)
|
||||
|
||||
maven_jar(
|
||||
name = "com_google_guava_android",
|
||||
artifact = "com.google.guava:guava:27.0.1-android",
|
||||
sha1 = "b7e1c37f66ef193796ccd7ea6e80c2b05426182d",
|
||||
)
|
||||
|
||||
maven_jar(
|
||||
name = "com_google_common_flogger",
|
||||
artifact = "com.google.flogger:flogger:0.3.1",
|
||||
sha1 = "585030fe1ec709760cbef997a459729fb965df0e",
|
||||
)
|
||||
|
||||
maven_jar(
|
||||
name = "com_google_common_flogger_system_backend",
|
||||
artifact = "com.google.flogger:flogger-system-backend:0.3.1",
|
||||
sha1 = "287b569d76abcd82f9de87fe41829fbc7ebd8ac9",
|
||||
)
|
||||
|
||||
maven_jar(
|
||||
name = "com_google_code_findbugs",
|
||||
artifact = "com.google.code.findbugs:jsr305:3.0.2",
|
||||
sha1 = "25ea2e8b0c338a877313bd4672d3fe056ea78f0d",
|
||||
)
|
||||
|
||||
# You may run setup_android.sh to install Android SDK and NDK.
|
||||
android_ndk_repository(
|
||||
name = "androidndk",
|
||||
@@ -289,9 +211,15 @@ android_sdk_repository(
|
||||
|
||||
http_archive(
|
||||
name = "build_bazel_rules_apple",
|
||||
sha256 = "bdc8e66e70b8a75da23b79f1f8c6207356df07d041d96d2189add7ee0780cf4e",
|
||||
strip_prefix = "rules_apple-b869b0d3868d78a1d4ffd866ccb304fb68aa12c3",
|
||||
url = "https://github.com/bazelbuild/rules_apple/archive/b869b0d3868d78a1d4ffd866ccb304fb68aa12c3.tar.gz",
|
||||
sha256 = "7a7afdd4869bb201c9352eed2daf37294d42b093579b70423490c1b4d4f6ce42",
|
||||
url = "https://github.com/bazelbuild/rules_apple/releases/download/0.19.0/rules_apple.0.19.0.tar.gz",
|
||||
patches = [
|
||||
# Bypass checking ios unit test runner when building MP ios applications.
|
||||
"@//third_party:build_bazel_rules_apple_bypass_test_runner_check.diff"
|
||||
],
|
||||
patch_args = [
|
||||
"-p1",
|
||||
],
|
||||
)
|
||||
|
||||
load(
|
||||
@@ -325,3 +253,49 @@ http_archive(
|
||||
build_file = "@//third_party:google_toolbox_for_mac.BUILD",
|
||||
)
|
||||
|
||||
# Maven dependencies.
|
||||
|
||||
RULES_JVM_EXTERNAL_TAG = "3.2"
|
||||
RULES_JVM_EXTERNAL_SHA = "82262ff4223c5fda6fb7ff8bd63db8131b51b413d26eb49e3131037e79e324af"
|
||||
|
||||
http_archive(
|
||||
name = "rules_jvm_external",
|
||||
strip_prefix = "rules_jvm_external-%s" % RULES_JVM_EXTERNAL_TAG,
|
||||
sha256 = RULES_JVM_EXTERNAL_SHA,
|
||||
url = "https://github.com/bazelbuild/rules_jvm_external/archive/%s.zip" % RULES_JVM_EXTERNAL_TAG,
|
||||
)
|
||||
|
||||
load("@rules_jvm_external//:defs.bzl", "maven_install")
|
||||
|
||||
# Important: there can only be one maven_install rule. Add new maven deps here.
|
||||
maven_install(
|
||||
name = "maven",
|
||||
artifacts = [
|
||||
"junit:junit:4.12",
|
||||
"androidx.test.espresso:espresso-core:3.1.1",
|
||||
"org.hamcrest:hamcrest-library:1.3",
|
||||
"androidx.concurrent:concurrent-futures:1.0.0-alpha03",
|
||||
"androidx.lifecycle:lifecycle-common:2.2.0",
|
||||
"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",
|
||||
"com.google.code.findbugs:jsr305:3.0.2",
|
||||
"com.google.flogger:flogger-system-backend:0.3.1",
|
||||
"com.google.flogger:flogger:0.3.1",
|
||||
"com.google.guava:guava:27.0.1-android",
|
||||
],
|
||||
repositories = [
|
||||
"https://jcenter.bintray.com",
|
||||
"https://maven.google.com",
|
||||
"https://dl.google.com/dl/android/maven2",
|
||||
"https://repo1.maven.org/maven2",
|
||||
],
|
||||
fetch_sources = True,
|
||||
version_conflict_policy = "pinned",
|
||||
)
|
||||
|
||||
@@ -14,6 +14,9 @@
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
# Note: yes, these need to use "//external:android/crosstool", not
|
||||
# @androidndk//:default_crosstool.
|
||||
|
||||
config_setting(
|
||||
name = "android",
|
||||
values = {"crosstool_top": "//external:android/crosstool"},
|
||||
|
||||
@@ -184,23 +184,14 @@ const float SpectrogramCalculator::kLnPowerToDb = 4.342944819032518;
|
||||
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"
|
||||
"frame_duration_seconds: "
|
||||
<< spectrogram_options.frame_overlap_seconds();
|
||||
// TODO: return an error.
|
||||
}
|
||||
if (spectrogram_options.frame_overlap_seconds() >=
|
||||
spectrogram_options.frame_duration_seconds()) {
|
||||
::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Invalid frame_overlap_seconds.\nframe_overlap_seconds: "
|
||||
<< spectrogram_options.frame_overlap_seconds()
|
||||
<< "\nframe_duration_seconds: "
|
||||
<< spectrogram_options.frame_duration_seconds();
|
||||
// TODO: return an error.
|
||||
}
|
||||
if (spectrogram_options.frame_overlap_seconds() < 0.0) {
|
||||
::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Frame_overlap_seconds is < 0.0.\nframe_overlap_seconds: "
|
||||
<< spectrogram_options.frame_overlap_seconds();
|
||||
// TODO: return an error.
|
||||
}
|
||||
|
||||
TimeSeriesHeader input_header;
|
||||
@@ -212,9 +203,7 @@ const float SpectrogramCalculator::kLnPowerToDb = 4.342944819032518;
|
||||
|
||||
if (!spectrogram_options.allow_multichannel_input() &&
|
||||
num_input_channels_ != 1) {
|
||||
::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "The current setting only supports single-channel input. Please set "
|
||||
"allow_multichannel_input.\n";
|
||||
// TODO: return an error.
|
||||
}
|
||||
|
||||
frame_duration_samples_ =
|
||||
@@ -293,10 +282,7 @@ const float SpectrogramCalculator::kLnPowerToDb = 4.342944819032518;
|
||||
|
||||
const Matrix& input_stream = cc->Inputs().Index(0).Get<Matrix>();
|
||||
if (input_stream.rows() != num_input_channels_) {
|
||||
::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||
<< "Number of input channels do not correspond to the number of rows "
|
||||
<< "in the input matrix: " << num_input_channels_ << "channels vs "
|
||||
<< input_stream.rows() << " rows";
|
||||
// TODO: return an error.
|
||||
}
|
||||
|
||||
cumulative_input_samples_ += input_stream.cols();
|
||||
|
||||
@@ -228,7 +228,9 @@ cc_library(
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
@@ -249,12 +251,14 @@ cc_library(
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/formats:classification_cc_proto",
|
||||
"//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",
|
||||
"@org_tensorflow//tensorflow/lite:framework",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -265,10 +269,11 @@ cc_test(
|
||||
deps = [
|
||||
":begin_loop_calculator",
|
||||
":end_loop_calculator",
|
||||
"//mediapipe/calculators/core:packet_cloner_calculator",
|
||||
":gate_calculator",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_contract",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
@@ -334,6 +339,7 @@ cc_library(
|
||||
deps = [
|
||||
":clip_vector_size_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
@@ -693,15 +699,17 @@ cc_test(
|
||||
name = "previous_loopback_calculator_test",
|
||||
srcs = ["previous_loopback_calculator_test.cc"],
|
||||
deps = [
|
||||
":gate_calculator",
|
||||
":make_pair_calculator",
|
||||
":pass_through_calculator",
|
||||
":previous_loopback_calculator",
|
||||
"//mediapipe/calculators/core:make_pair_calculator",
|
||||
"//mediapipe/calculators/core:pass_through_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/stream_handler:immediate_input_stream_handler",
|
||||
"//mediapipe/framework/tool:sink",
|
||||
"@com_google_absl//absl/time",
|
||||
@@ -769,9 +777,11 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":split_vector_calculator_cc_proto",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/formats:matrix",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/util:resource_util",
|
||||
@@ -805,6 +815,38 @@ cc_test(
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "split_normalized_landmark_list_calculator",
|
||||
srcs = ["split_normalized_landmark_list_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":split_vector_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/util:resource_util",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "split_normalized_landmark_list_calculator_test",
|
||||
srcs = ["split_normalized_landmark_list_calculator_test.cc"],
|
||||
deps = [
|
||||
":split_normalized_landmark_list_calculator",
|
||||
":split_vector_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/deps:file_path",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/tool:validate_type",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "dequantize_byte_array_calculator",
|
||||
srcs = ["dequantize_byte_array_calculator.cc"],
|
||||
|
||||
@@ -20,6 +20,8 @@
|
||||
#include "mediapipe/calculators/core/end_loop_calculator.h"
|
||||
#include "mediapipe/framework/calculator_contract.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
@@ -28,6 +30,13 @@
|
||||
namespace mediapipe {
|
||||
namespace {
|
||||
|
||||
MATCHER_P2(PacketOfIntsEq, timestamp, value, "") {
|
||||
Timestamp actual_timestamp = arg.Timestamp();
|
||||
const auto& actual_value = arg.template Get<std::vector<int>>();
|
||||
return testing::Value(actual_timestamp, testing::Eq(timestamp)) &&
|
||||
testing::Value(actual_value, testing::ElementsAreArray(value));
|
||||
}
|
||||
|
||||
typedef BeginLoopCalculator<std::vector<int>> BeginLoopIntegerCalculator;
|
||||
REGISTER_CALCULATOR(BeginLoopIntegerCalculator);
|
||||
|
||||
@@ -59,8 +68,8 @@ REGISTER_CALCULATOR(EndLoopIntegersCalculator);
|
||||
|
||||
class BeginEndLoopCalculatorGraphTest : public ::testing::Test {
|
||||
protected:
|
||||
BeginEndLoopCalculatorGraphTest() {
|
||||
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
void SetUp() override {
|
||||
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
num_threads: 4
|
||||
input_stream: "ints"
|
||||
@@ -82,94 +91,222 @@ class BeginEndLoopCalculatorGraphTest : public ::testing::Test {
|
||||
output_stream: "ITERABLE:ints_plus_one"
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("ints_plus_one", &graph_config_, &output_packets_);
|
||||
tool::AddVectorSink("ints_plus_one", &graph_config, &output_packets_);
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
}
|
||||
|
||||
CalculatorGraphConfig graph_config_;
|
||||
void SendPacketOfInts(Timestamp timestamp, std::vector<int> ints) {
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"ints", MakePacket<std::vector<int>>(std::move(ints)).At(timestamp)));
|
||||
}
|
||||
|
||||
CalculatorGraph graph_;
|
||||
std::vector<Packet> output_packets_;
|
||||
};
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphTest, InputStreamForIterableIsEmpty) {
|
||||
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||
|
||||
// EndLoopCalc will forward the timestamp bound because there are no packets
|
||||
// to process.
|
||||
ASSERT_EQ(0, output_packets_.size());
|
||||
|
||||
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
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());
|
||||
SendPacketOfInts(Timestamp(0), {});
|
||||
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());
|
||||
EXPECT_TRUE(output_packets_.empty());
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
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());
|
||||
SendPacketOfInts(input_timestamp, {0, 1, 2});
|
||||
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>>());
|
||||
EXPECT_THAT(output_packets_,
|
||||
testing::ElementsAre(
|
||||
PacketOfIntsEq(input_timestamp, std::vector<int>{1, 2, 3})));
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
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)));
|
||||
SendPacketOfInts(input_timestamp0, {0, 1});
|
||||
|
||||
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)));
|
||||
SendPacketOfInts(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)));
|
||||
SendPacketOfInts(input_timestamp2, {2, 3});
|
||||
|
||||
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>>());
|
||||
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||
|
||||
// At input_timestamp1, EndLoopCalc will forward timestamp bound as there are
|
||||
// no elements in vector to process.
|
||||
EXPECT_THAT(output_packets_,
|
||||
testing::ElementsAre(
|
||||
PacketOfIntsEq(input_timestamp0, std::vector<int>{1, 2}),
|
||||
PacketOfIntsEq(input_timestamp2, std::vector<int>{3, 4})));
|
||||
}
|
||||
|
||||
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>>());
|
||||
// Passes non empty vector through or outputs empty vector in case of timestamp
|
||||
// bound update.
|
||||
class PassThroughOrEmptyVectorCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->SetProcessTimestampBounds(true);
|
||||
cc->Inputs().Index(0).Set<std::vector<int>>();
|
||||
cc->Outputs().Index(0).Set<std::vector<int>>();
|
||||
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().Index(0).IsEmpty()) {
|
||||
cc->Outputs().Index(0).AddPacket(cc->Inputs().Index(0).Value());
|
||||
} else {
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
MakePacket<std::vector<int>>(std::vector<int>())
|
||||
.At(cc->InputTimestamp()));
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(PassThroughOrEmptyVectorCalculator);
|
||||
|
||||
class BeginEndLoopCalculatorGraphProcessingEmptyPacketsTest
|
||||
: public ::testing::Test {
|
||||
protected:
|
||||
void SetUp() override {
|
||||
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
num_threads: 4
|
||||
input_stream: "ints"
|
||||
input_stream: "force_ints_to_be_timestamp_bound_update"
|
||||
node {
|
||||
calculator: "GateCalculator"
|
||||
input_stream: "ints"
|
||||
input_stream: "DISALLOW:force_ints_to_be_timestamp_bound_update"
|
||||
output_stream: "ints_passed_through"
|
||||
}
|
||||
node {
|
||||
calculator: "BeginLoopIntegerCalculator"
|
||||
input_stream: "ITERABLE:ints_passed_through"
|
||||
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"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughOrEmptyVectorCalculator"
|
||||
input_stream: "ints_plus_one"
|
||||
output_stream: "ints_plus_one_passed_through"
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("ints_plus_one_passed_through", &graph_config,
|
||||
&output_packets_);
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
}
|
||||
|
||||
void SendPacketOfIntsOrBound(Timestamp timestamp, std::vector<int> ints) {
|
||||
// All "ints" packets which are empty are forced to be just timestamp
|
||||
// bound updates for begin loop calculator.
|
||||
bool force_ints_to_be_timestamp_bound_update = ints.empty();
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"force_ints_to_be_timestamp_bound_update",
|
||||
MakePacket<bool>(force_ints_to_be_timestamp_bound_update)
|
||||
.At(timestamp)));
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"ints", MakePacket<std::vector<int>>(std::move(ints)).At(timestamp)));
|
||||
}
|
||||
|
||||
CalculatorGraph graph_;
|
||||
std::vector<Packet> output_packets_;
|
||||
};
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphProcessingEmptyPacketsTest,
|
||||
SingleEmptyVector) {
|
||||
SendPacketOfIntsOrBound(Timestamp(0), {});
|
||||
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||
|
||||
EXPECT_THAT(output_packets_, testing::ElementsAre(PacketOfIntsEq(
|
||||
Timestamp(0), std::vector<int>{})));
|
||||
|
||||
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphProcessingEmptyPacketsTest,
|
||||
SingleNonEmptyVector) {
|
||||
SendPacketOfIntsOrBound(Timestamp(0), {0, 1, 2});
|
||||
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||
|
||||
EXPECT_THAT(output_packets_, testing::ElementsAre(PacketOfIntsEq(
|
||||
Timestamp(0), std::vector<int>{1, 2, 3})));
|
||||
|
||||
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphProcessingEmptyPacketsTest, MultipleVectors) {
|
||||
SendPacketOfIntsOrBound(Timestamp(0), {});
|
||||
// Waiting until idle to guarantee all timestamp bound updates are processed
|
||||
// individually. (Timestamp bounds updates occur in the provide config only
|
||||
// if input is an empty vector.)
|
||||
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||
|
||||
SendPacketOfIntsOrBound(Timestamp(1), {0, 1});
|
||||
SendPacketOfIntsOrBound(Timestamp(2), {});
|
||||
// Waiting until idle to guarantee all timestamp bound updates are processed
|
||||
// individually. (Timestamp bounds updates occur in the provide config only
|
||||
// if input is an empty vector.)
|
||||
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||
|
||||
SendPacketOfIntsOrBound(Timestamp(3), {2, 3});
|
||||
SendPacketOfIntsOrBound(Timestamp(4), {});
|
||||
// Waiting until idle to guarantee all timestamp bound updates are processed
|
||||
// individually. (Timestamp bounds updates occur in the provide config only
|
||||
// if input is an empty vector.)
|
||||
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||
|
||||
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
testing::ElementsAre(PacketOfIntsEq(Timestamp(0), std::vector<int>{}),
|
||||
PacketOfIntsEq(Timestamp(1), std::vector<int>{1, 2}),
|
||||
PacketOfIntsEq(Timestamp(2), std::vector<int>{}),
|
||||
PacketOfIntsEq(Timestamp(3), std::vector<int>{3, 4}),
|
||||
PacketOfIntsEq(Timestamp(4), std::vector<int>{})));
|
||||
}
|
||||
|
||||
class MultiplierCalculator : public CalculatorBase {
|
||||
@@ -199,8 +336,8 @@ REGISTER_CALCULATOR(MultiplierCalculator);
|
||||
|
||||
class BeginEndLoopCalculatorGraphWithClonedInputsTest : public ::testing::Test {
|
||||
protected:
|
||||
BeginEndLoopCalculatorGraphWithClonedInputsTest() {
|
||||
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
void SetUp() override {
|
||||
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
num_threads: 4
|
||||
input_stream: "ints"
|
||||
@@ -226,109 +363,85 @@ class BeginEndLoopCalculatorGraphWithClonedInputsTest : public ::testing::Test {
|
||||
output_stream: "ITERABLE:multiplied_ints"
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("multiplied_ints", &graph_config_, &output_packets_);
|
||||
tool::AddVectorSink("multiplied_ints", &graph_config, &output_packets_);
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config));
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
}
|
||||
|
||||
CalculatorGraphConfig graph_config_;
|
||||
void SendPackets(Timestamp timestamp, int multiplier, std::vector<int> ints) {
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"ints", MakePacket<std::vector<int>>(std::move(ints)).At(timestamp)));
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"multiplier", MakePacket<int>(multiplier).At(timestamp)));
|
||||
}
|
||||
|
||||
void SendMultiplier(Timestamp timestamp, int multiplier) {
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"multiplier", MakePacket<int>(multiplier).At(timestamp)));
|
||||
}
|
||||
|
||||
CalculatorGraph graph_;
|
||||
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>>();
|
||||
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest,
|
||||
InputStreamForIterableIsEmpty) {
|
||||
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());
|
||||
SendMultiplier(input_timestamp, /*multiplier=*/2);
|
||||
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||
|
||||
// EndLoopCalc will forward the timestamp bound because there are no packets
|
||||
// to process.
|
||||
ASSERT_EQ(0, output_packets_.size());
|
||||
|
||||
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, SingleEmptyVector) {
|
||||
SendPackets(Timestamp(0), /*multiplier=*/2, /*ints=*/{});
|
||||
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());
|
||||
EXPECT_TRUE(output_packets_.empty());
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
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());
|
||||
SendPackets(input_timestamp, /*multiplier=*/2, /*ints=*/{0, 1, 2});
|
||||
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>>());
|
||||
EXPECT_THAT(output_packets_,
|
||||
testing::ElementsAre(
|
||||
PacketOfIntsEq(input_timestamp, std::vector<int>{0, 2, 4})));
|
||||
|
||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
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)));
|
||||
SendPackets(input_timestamp0, /*multiplier=*/2, /*ints=*/{0, 1});
|
||||
|
||||
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)));
|
||||
SendPackets(input_timestamp1, /*multiplier=*/2, /*ints=*/{});
|
||||
|
||||
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)));
|
||||
SendPackets(input_timestamp2, /*multiplier=*/3, /*ints=*/{2, 3});
|
||||
|
||||
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>>());
|
||||
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||
|
||||
// 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>>());
|
||||
EXPECT_THAT(output_packets_,
|
||||
testing::ElementsAre(
|
||||
PacketOfIntsEq(input_timestamp0, std::vector<int>{0, 2}),
|
||||
PacketOfIntsEq(input_timestamp2, std::vector<int>{6, 9})));
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
@@ -16,7 +16,9 @@
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/matrix.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
@@ -31,4 +33,13 @@ typedef BeginLoopCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||
BeginLoopNormalizedRectCalculator;
|
||||
REGISTER_CALCULATOR(BeginLoopNormalizedRectCalculator);
|
||||
|
||||
// A calculator to process std::vector<Detection>.
|
||||
typedef BeginLoopCalculator<std::vector<::mediapipe::Detection>>
|
||||
BeginLoopDetectionCalculator;
|
||||
REGISTER_CALCULATOR(BeginLoopDetectionCalculator);
|
||||
|
||||
// A calculator to process std::vector<Matrix>.
|
||||
typedef BeginLoopCalculator<std::vector<Matrix>> BeginLoopMatrixCalculator;
|
||||
REGISTER_CALCULATOR(BeginLoopMatrixCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -52,20 +52,28 @@ namespace mediapipe {
|
||||
// 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.
|
||||
// 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) {
|
||||
// The below enables processing of timestamp bound updates, and that enables
|
||||
// correct timestamp propagation by the companion EndLoopCalculator.
|
||||
//
|
||||
// For instance, Process() function will be still invoked even if upstream
|
||||
// calculator has updated timestamp bound for ITERABLE input instead of
|
||||
// providing actual value.
|
||||
cc->SetProcessTimestampBounds(true);
|
||||
|
||||
// A non-empty packet in the optional "TICK" input stream wakes up the
|
||||
// calculator.
|
||||
// DEPRECATED as timestamp bound updates are processed by default in this
|
||||
// calculator.
|
||||
if (cc->Inputs().HasTag("TICK")) {
|
||||
cc->Inputs().Tag("TICK").SetAny();
|
||||
}
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
@@ -25,4 +26,8 @@ typedef ClipVectorSizeCalculator<::mediapipe::NormalizedRect>
|
||||
ClipNormalizedRectVectorSizeCalculator;
|
||||
REGISTER_CALCULATOR(ClipNormalizedRectVectorSizeCalculator);
|
||||
|
||||
typedef ClipVectorSizeCalculator<::mediapipe::Detection>
|
||||
ClipDetectionVectorSizeCalculator;
|
||||
REGISTER_CALCULATOR(ClipDetectionVectorSizeCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -51,8 +51,8 @@ namespace mediapipe {
|
||||
class ConstantSidePacketCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
const auto& options = cc->Options().GetExtension(
|
||||
::mediapipe::ConstantSidePacketCalculatorOptions::ext);
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::ConstantSidePacketCalculatorOptions>();
|
||||
RET_CHECK_EQ(cc->OutputSidePackets().NumEntries(kPacketTag),
|
||||
options.packet_size())
|
||||
<< "Number of output side packets has to be same as number of packets "
|
||||
@@ -80,8 +80,8 @@ class ConstantSidePacketCalculator : public CalculatorBase {
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
const auto& options = cc->Options().GetExtension(
|
||||
::mediapipe::ConstantSidePacketCalculatorOptions::ext);
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::ConstantSidePacketCalculatorOptions>();
|
||||
int index = 0;
|
||||
for (CollectionItemId id = cc->OutputSidePackets().BeginId(kPacketTag);
|
||||
id != cc->OutputSidePackets().EndId(kPacketTag); ++id, ++index) {
|
||||
|
||||
@@ -16,9 +16,11 @@
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/classification.pb.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "mediapipe/util/render_data.pb.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
@@ -37,4 +39,11 @@ typedef EndLoopCalculator<std::vector<::mediapipe::RenderData>>
|
||||
EndLoopRenderDataCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopRenderDataCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<::mediapipe::ClassificationList>>
|
||||
EndLoopClassificationListCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopClassificationListCalculator);
|
||||
|
||||
typedef EndLoopCalculator<std::vector<TfLiteTensor>> EndLoopTensorCalculator;
|
||||
REGISTER_CALCULATOR(EndLoopTensorCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -25,13 +25,17 @@ namespace mediapipe {
|
||||
// together with some previous output.
|
||||
//
|
||||
// For the first packet that arrives on the MAIN input, the timestamp bound is
|
||||
// advanced on the output. Downstream calculators will see this as an empty
|
||||
// advanced on the PREV_LOOP. Downstream calculators will see this as an empty
|
||||
// packet. This way they are not kept waiting for the previous output, which
|
||||
// for the first iteration does not exist.
|
||||
//
|
||||
// Thereafter, each packet received on MAIN is matched with a packet received
|
||||
// on LOOP; the LOOP packet's timestamp is changed to that of the MAIN packet,
|
||||
// and it is output on PREV_LOOP.
|
||||
// Thereafter,
|
||||
// - Each non-empty MAIN packet results in:
|
||||
// a) a PREV_LOOP packet with contents of the LOOP packet received at the
|
||||
// timestamp of the previous non-empty MAIN packet
|
||||
// b) or in a PREV_LOOP timestamp bound update if the LOOP packet was empty.
|
||||
// - Each empty MAIN packet indicating timestamp bound update results in a
|
||||
// PREV_LOOP timestamp bound update.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
@@ -56,83 +60,115 @@ class PreviousLoopbackCalculator : public CalculatorBase {
|
||||
// TODO: an optional PREV_TIMESTAMP output could be added to
|
||||
// carry the original timestamp of the packet on PREV_LOOP.
|
||||
cc->SetInputStreamHandler("ImmediateInputStreamHandler");
|
||||
// Process() function is invoked in response to MAIN/LOOP stream timestamp
|
||||
// bound updates.
|
||||
cc->SetProcessTimestampBounds(true);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) final {
|
||||
main_id_ = cc->Inputs().GetId("MAIN", 0);
|
||||
loop_id_ = cc->Inputs().GetId("LOOP", 0);
|
||||
loop_out_id_ = cc->Outputs().GetId("PREV_LOOP", 0);
|
||||
prev_loop_id_ = cc->Outputs().GetId("PREV_LOOP", 0);
|
||||
cc->Outputs()
|
||||
.Get(loop_out_id_)
|
||||
.Get(prev_loop_id_)
|
||||
.SetHeader(cc->Inputs().Get(loop_id_).Header());
|
||||
|
||||
// Use an empty packet for the first round, since there is no previous
|
||||
// output.
|
||||
loopback_packets_.push_back({});
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||
Packet& main_packet = cc->Inputs().Get(main_id_).Value();
|
||||
if (!main_packet.IsEmpty()) {
|
||||
main_ts_.push_back(main_packet.Timestamp());
|
||||
}
|
||||
Packet& loopback_packet = cc->Inputs().Get(loop_id_).Value();
|
||||
if (!loopback_packet.IsEmpty()) {
|
||||
loopback_packets_.push_back(loopback_packet);
|
||||
while (!main_ts_.empty() &&
|
||||
main_ts_.front() <= loopback_packets_.front().Timestamp()) {
|
||||
main_ts_.pop_front();
|
||||
}
|
||||
}
|
||||
auto& loop_out = cc->Outputs().Get(loop_out_id_);
|
||||
// Non-empty packets and empty packets indicating timestamp bound updates
|
||||
// are guaranteed to have timestamps greater than timestamps of previous
|
||||
// packets within the same stream. Calculator tracks and operates on such
|
||||
// packets.
|
||||
|
||||
while (!main_ts_.empty() && !loopback_packets_.empty()) {
|
||||
Timestamp main_timestamp = main_ts_.front();
|
||||
main_ts_.pop_front();
|
||||
Packet previous_loopback = loopback_packets_.front().At(main_timestamp);
|
||||
loopback_packets_.pop_front();
|
||||
|
||||
if (previous_loopback.IsEmpty()) {
|
||||
// TODO: SetCompleteTimestampBound would be more useful.
|
||||
loop_out.SetNextTimestampBound(main_timestamp + 1);
|
||||
const Packet& main_packet = cc->Inputs().Get(main_id_).Value();
|
||||
if (prev_main_ts_ < main_packet.Timestamp()) {
|
||||
Timestamp loop_timestamp;
|
||||
if (!main_packet.IsEmpty()) {
|
||||
loop_timestamp = prev_non_empty_main_ts_;
|
||||
prev_non_empty_main_ts_ = main_packet.Timestamp();
|
||||
} else {
|
||||
loop_out.AddPacket(std::move(previous_loopback));
|
||||
// Calculator advances PREV_LOOP timestamp bound in response to empty
|
||||
// MAIN packet, hence not caring about corresponding loop packet.
|
||||
loop_timestamp = Timestamp::Unset();
|
||||
}
|
||||
main_packet_specs_.push_back({.timestamp = main_packet.Timestamp(),
|
||||
.loop_timestamp = loop_timestamp});
|
||||
prev_main_ts_ = main_packet.Timestamp();
|
||||
}
|
||||
|
||||
const Packet& loop_packet = cc->Inputs().Get(loop_id_).Value();
|
||||
if (prev_loop_ts_ < loop_packet.Timestamp()) {
|
||||
loop_packets_.push_back(loop_packet);
|
||||
prev_loop_ts_ = loop_packet.Timestamp();
|
||||
}
|
||||
|
||||
auto& prev_loop = cc->Outputs().Get(prev_loop_id_);
|
||||
while (!main_packet_specs_.empty() && !loop_packets_.empty()) {
|
||||
// The earliest MAIN packet.
|
||||
const MainPacketSpec& main_spec = main_packet_specs_.front();
|
||||
// The earliest LOOP packet.
|
||||
const Packet& loop_candidate = loop_packets_.front();
|
||||
// Match LOOP and MAIN packets.
|
||||
if (main_spec.loop_timestamp < loop_candidate.Timestamp()) {
|
||||
// No LOOP packet can match the MAIN packet under review.
|
||||
prev_loop.SetNextTimestampBound(main_spec.timestamp + 1);
|
||||
main_packet_specs_.pop_front();
|
||||
} else if (main_spec.loop_timestamp > loop_candidate.Timestamp()) {
|
||||
// No MAIN packet can match the LOOP packet under review.
|
||||
loop_packets_.pop_front();
|
||||
} else {
|
||||
// Exact match found.
|
||||
if (loop_candidate.IsEmpty()) {
|
||||
// However, LOOP packet is empty.
|
||||
prev_loop.SetNextTimestampBound(main_spec.timestamp + 1);
|
||||
} else {
|
||||
prev_loop.AddPacket(loop_candidate.At(main_spec.timestamp));
|
||||
}
|
||||
loop_packets_.pop_front();
|
||||
main_packet_specs_.pop_front();
|
||||
}
|
||||
}
|
||||
|
||||
// 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();
|
||||
if (main_packet_specs_.empty() && cc->Inputs().Get(main_id_).IsDone()) {
|
||||
prev_loop.Close();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
struct MainPacketSpec {
|
||||
Timestamp timestamp;
|
||||
// Expected timestamp of the packet from LOOP stream that corresponds to the
|
||||
// packet from MAIN stream descirbed by this spec.
|
||||
Timestamp loop_timestamp;
|
||||
};
|
||||
|
||||
CollectionItemId main_id_;
|
||||
CollectionItemId loop_id_;
|
||||
CollectionItemId loop_out_id_;
|
||||
CollectionItemId prev_loop_id_;
|
||||
|
||||
std::deque<Timestamp> main_ts_;
|
||||
std::deque<Packet> loopback_packets_;
|
||||
// Contains specs for MAIN packets which only can be:
|
||||
// - non-empty packets
|
||||
// - empty packets indicating timestamp bound updates
|
||||
//
|
||||
// Sorted according to packet timestamps.
|
||||
std::deque<MainPacketSpec> main_packet_specs_;
|
||||
Timestamp prev_main_ts_ = Timestamp::Unstarted();
|
||||
Timestamp prev_non_empty_main_ts_ = Timestamp::Unstarted();
|
||||
|
||||
// Contains LOOP packets which only can be:
|
||||
// - the very first empty packet
|
||||
// - non empty packets
|
||||
// - empty packets indicating timestamp bound updates
|
||||
//
|
||||
// Sorted according to packet timestamps.
|
||||
std::deque<Packet> loop_packets_;
|
||||
// Using "Timestamp::Unset" instead of "Timestamp::Unstarted" in order to
|
||||
// allow addition of the very first empty packet (which doesn't indicate
|
||||
// timestamp bound change necessarily).
|
||||
Timestamp prev_loop_ts_ = Timestamp::Unset();
|
||||
};
|
||||
REGISTER_CALCULATOR(PreviousLoopbackCalculator);
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
// limitations under the License.
|
||||
|
||||
#include <algorithm>
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
@@ -25,12 +26,17 @@
|
||||
#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.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
#include "mediapipe/framework/tool/sink.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
using ::testing::ElementsAre;
|
||||
using ::testing::Eq;
|
||||
using ::testing::Pair;
|
||||
using ::testing::Value;
|
||||
namespace {
|
||||
|
||||
// Returns the timestamp values for a vector of Packets.
|
||||
@@ -43,6 +49,23 @@ std::vector<int64> TimestampValues(const std::vector<Packet>& packets) {
|
||||
return result;
|
||||
}
|
||||
|
||||
MATCHER(EmptyPacket, negation ? "isn't empty" : "is empty") {
|
||||
if (arg.IsEmpty()) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
MATCHER_P(IntPacket, value, "") {
|
||||
return Value(arg.template Get<int>(), Eq(value));
|
||||
}
|
||||
|
||||
MATCHER_P2(PairPacket, timestamp, pair, "") {
|
||||
Timestamp actual_timestamp = arg.Timestamp();
|
||||
const auto& actual_pair = arg.template Get<std::pair<Packet, Packet>>();
|
||||
return Value(actual_timestamp, Eq(timestamp)) && Value(actual_pair, pair);
|
||||
}
|
||||
|
||||
TEST(PreviousLoopbackCalculator, CorrectTimestamps) {
|
||||
std::vector<Packet> in_prev;
|
||||
CalculatorGraphConfig graph_config_ =
|
||||
@@ -81,32 +104,30 @@ TEST(PreviousLoopbackCalculator, CorrectTimestamps) {
|
||||
MP_EXPECT_OK(graph_.AddPacketToInputStream(
|
||||
input_name, MakePacket<int>(n).At(Timestamp(n))));
|
||||
};
|
||||
auto pair_values = [](const Packet& packet) {
|
||||
auto pair = packet.Get<std::pair<Packet, Packet>>();
|
||||
int first = pair.first.IsEmpty() ? -1 : pair.first.Get<int>();
|
||||
int second = pair.second.IsEmpty() ? -1 : pair.second.Get<int>();
|
||||
return std::make_pair(first, second);
|
||||
};
|
||||
|
||||
send_packet("in", 1);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1}));
|
||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(1, -1));
|
||||
EXPECT_THAT(TimestampValues(in_prev), ElementsAre(1));
|
||||
EXPECT_THAT(in_prev.back(),
|
||||
PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())));
|
||||
|
||||
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));
|
||||
EXPECT_THAT(TimestampValues(in_prev), ElementsAre(1, 2));
|
||||
EXPECT_THAT(in_prev.back(),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), IntPacket(1))));
|
||||
|
||||
send_packet("in", 5);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 2, 5}));
|
||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(5, 2));
|
||||
EXPECT_THAT(TimestampValues(in_prev), ElementsAre(1, 2, 5));
|
||||
EXPECT_THAT(in_prev.back(),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), IntPacket(2))));
|
||||
|
||||
send_packet("in", 15);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 2, 5, 15}));
|
||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(15, 5));
|
||||
EXPECT_THAT(TimestampValues(in_prev), ElementsAre(1, 2, 5, 15));
|
||||
EXPECT_THAT(in_prev.back(),
|
||||
PairPacket(Timestamp(15), Pair(IntPacket(15), IntPacket(5))));
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
@@ -185,24 +206,24 @@ TEST(PreviousLoopbackCalculator, ClosesCorrectly) {
|
||||
|
||||
send_packet("in", 1);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1}));
|
||||
EXPECT_THAT(TimestampValues(outputs), ElementsAre(1));
|
||||
|
||||
send_packet("in", 2);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1, 2}));
|
||||
EXPECT_THAT(TimestampValues(outputs), ElementsAre(1, 2));
|
||||
|
||||
send_packet("in", 5);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1, 2, 5}));
|
||||
EXPECT_THAT(TimestampValues(outputs), ElementsAre(1, 2, 5));
|
||||
|
||||
send_packet("in", 15);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1, 2, 5, 15}));
|
||||
EXPECT_THAT(TimestampValues(outputs), ElementsAre(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()}));
|
||||
EXPECT_THAT(TimestampValues(outputs),
|
||||
ElementsAre(1, 2, 5, 15, Timestamp::Max().Value()));
|
||||
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
@@ -247,16 +268,12 @@ TEST(PreviousLoopbackCalculator, EmptyLoopForever) {
|
||||
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) {
|
||||
for (int main_ts = 0; 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) {
|
||||
for (int j = 0; j < main_ts + 1; ++j) {
|
||||
EXPECT_EQ(ts_values[j], j);
|
||||
}
|
||||
}
|
||||
@@ -266,5 +283,487 @@ TEST(PreviousLoopbackCalculator, EmptyLoopForever) {
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
class PreviousLoopbackCalculatorProcessingTimestampsTest
|
||||
: public testing::Test {
|
||||
protected:
|
||||
void SetUp() override {
|
||||
CalculatorGraphConfig graph_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: 'input'
|
||||
input_stream: 'force_main_empty'
|
||||
input_stream: 'force_loop_empty'
|
||||
# Used to indicate "main" timestamp bound updates.
|
||||
node {
|
||||
calculator: 'GateCalculator'
|
||||
input_stream: 'input'
|
||||
input_stream: 'DISALLOW:force_main_empty'
|
||||
output_stream: 'main'
|
||||
}
|
||||
node {
|
||||
calculator: 'PreviousLoopbackCalculator'
|
||||
input_stream: 'MAIN:main'
|
||||
input_stream: 'LOOP:loop'
|
||||
input_stream_info: { tag_index: 'LOOP' back_edge: true }
|
||||
output_stream: 'PREV_LOOP:prev_loop'
|
||||
}
|
||||
node {
|
||||
calculator: 'PassThroughCalculator'
|
||||
input_stream: 'input'
|
||||
input_stream: 'prev_loop'
|
||||
output_stream: 'passed_through_input'
|
||||
output_stream: 'passed_through_prev_loop'
|
||||
}
|
||||
# Used to indicate "loop" timestamp bound updates.
|
||||
node {
|
||||
calculator: 'GateCalculator'
|
||||
input_stream: 'input'
|
||||
input_stream: 'DISALLOW:force_loop_empty'
|
||||
output_stream: 'loop'
|
||||
}
|
||||
node {
|
||||
calculator: 'MakePairCalculator'
|
||||
input_stream: 'passed_through_input'
|
||||
input_stream: 'passed_through_prev_loop'
|
||||
output_stream: 'passed_through_input_and_prev_loop'
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("passed_through_input_and_prev_loop", &graph_config,
|
||||
&output_packets_);
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config, {}));
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
}
|
||||
|
||||
void SendPackets(int timestamp, int input, bool force_main_empty,
|
||||
bool force_loop_empty) {
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"input", MakePacket<int>(input).At(Timestamp(timestamp))));
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"force_main_empty",
|
||||
MakePacket<bool>(force_main_empty).At(Timestamp(timestamp))));
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"force_loop_empty",
|
||||
MakePacket<bool>(force_loop_empty).At(Timestamp(timestamp))));
|
||||
}
|
||||
|
||||
CalculatorGraph graph_;
|
||||
std::vector<Packet> output_packets_;
|
||||
};
|
||||
|
||||
TEST_F(PreviousLoopbackCalculatorProcessingTimestampsTest,
|
||||
MultiplePacketsEmptyMainNonEmptyLoop) {
|
||||
SendPackets(/*timestamp=*/1, /*input=*/1, /*force_main_empty=*/true,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/2, /*input=*/2, /*force_main_empty=*/true,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/3, /*input=*/3, /*force_main_empty=*/true,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket())),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/5, /*input=*/5, /*force_main_empty=*/true,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket())),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket())),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/15, /*input=*/15,
|
||||
/*force_main_empty=*/true,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(
|
||||
PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket())),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket())),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), EmptyPacket())),
|
||||
PairPacket(Timestamp(15), Pair(IntPacket(15), EmptyPacket()))));
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(PreviousLoopbackCalculatorProcessingTimestampsTest,
|
||||
MultiplePacketsNonEmptyMainEmptyLoop) {
|
||||
SendPackets(/*timestamp=*/1, /*input=*/1,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/2, /*input=*/2,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/3, /*input=*/3,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket())),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/5, /*input=*/5,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket())),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket())),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/15, /*input=*/15,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(
|
||||
PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket())),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket())),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), EmptyPacket())),
|
||||
PairPacket(Timestamp(15), Pair(IntPacket(15), EmptyPacket()))));
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(PreviousLoopbackCalculatorProcessingTimestampsTest,
|
||||
MultiplePacketsAlteringMainNonEmptyLoop) {
|
||||
SendPackets(/*timestamp=*/1, /*input=*/1,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/2, /*input=*/2, /*force_main_empty=*/true,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/3, /*input=*/3,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket())),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), IntPacket(1)))));
|
||||
|
||||
SendPackets(/*timestamp=*/5, /*input=*/5,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket())),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), IntPacket(1))),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), IntPacket(3)))));
|
||||
|
||||
SendPackets(/*timestamp=*/15, /*input=*/15,
|
||||
/*force_main_empty=*/true,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(
|
||||
PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket())),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), IntPacket(1))),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), IntPacket(3))),
|
||||
PairPacket(Timestamp(15), Pair(IntPacket(15), EmptyPacket()))));
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(PreviousLoopbackCalculatorProcessingTimestampsTest,
|
||||
MultiplePacketsNonEmptyMainAlteringLoop) {
|
||||
SendPackets(/*timestamp=*/1, /*input=*/1,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/2, /*input=*/2,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), IntPacket(1)))));
|
||||
|
||||
SendPackets(/*timestamp=*/3, /*input=*/3,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), IntPacket(1))),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/5, /*input=*/5,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), IntPacket(1))),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket())),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), IntPacket(3)))));
|
||||
|
||||
SendPackets(/*timestamp=*/15, /*input=*/15,
|
||||
/*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(
|
||||
PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), IntPacket(1))),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket())),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), IntPacket(3))),
|
||||
PairPacket(Timestamp(15), Pair(IntPacket(15), EmptyPacket()))));
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(PreviousLoopbackCalculatorProcessingTimestampsTest,
|
||||
MultiplePacketsCheckIfLastCorrectAlteringMainAlteringLoop) {
|
||||
int num_packets = 1000;
|
||||
for (int i = 0; i < num_packets; ++i) {
|
||||
bool force_main_empty = i % 3 == 0 ? true : false;
|
||||
bool force_loop_empty = i % 2 == 0 ? true : false;
|
||||
SendPackets(/*timestamp=*/i + 1, /*input=*/i + 1, force_main_empty,
|
||||
force_loop_empty);
|
||||
}
|
||||
SendPackets(/*timestamp=*/num_packets + 1,
|
||||
/*input=*/num_packets + 1, /*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/false);
|
||||
SendPackets(/*timestamp=*/num_packets + 2,
|
||||
/*input=*/num_packets + 2, /*force_main_empty=*/false,
|
||||
/*force_loop_empty=*/false);
|
||||
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
ASSERT_FALSE(output_packets_.empty());
|
||||
EXPECT_THAT(
|
||||
output_packets_.back(),
|
||||
PairPacket(Timestamp(num_packets + 2),
|
||||
Pair(IntPacket(num_packets + 2), IntPacket(num_packets + 1))));
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
// Similar to GateCalculator, but it doesn't propagate timestamp bound updates.
|
||||
class DroppingGateCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Index(0).SetAny();
|
||||
cc->Inputs().Tag("DISALLOW").Set<bool>();
|
||||
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||
if (!cc->Inputs().Index(0).IsEmpty() &&
|
||||
!cc->Inputs().Tag("DISALLOW").Get<bool>()) {
|
||||
cc->Outputs().Index(0).AddPacket(cc->Inputs().Index(0).Value());
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
REGISTER_CALCULATOR(DroppingGateCalculator);
|
||||
|
||||
// Tests PreviousLoopbackCalculator in cases when there are no "LOOP" timestamp
|
||||
// bound updates and non-empty packets for a while and the aforementioned start
|
||||
// to arrive at some point. So, "PREV_LOOP" is delayed for a couple of inputs.
|
||||
class PreviousLoopbackCalculatorDelayBehaviorTest : public testing::Test {
|
||||
protected:
|
||||
void SetUp() override {
|
||||
CalculatorGraphConfig graph_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: 'input'
|
||||
# Drops "loop" when set to "true", delaying output of prev_loop, hence
|
||||
# delaying output of the graph.
|
||||
input_stream: 'delay_next_output'
|
||||
node {
|
||||
calculator: 'PreviousLoopbackCalculator'
|
||||
input_stream: 'MAIN:input'
|
||||
input_stream: 'LOOP:loop'
|
||||
input_stream_info: { tag_index: 'LOOP' back_edge: true }
|
||||
output_stream: 'PREV_LOOP:prev_loop'
|
||||
}
|
||||
node {
|
||||
calculator: 'PassThroughCalculator'
|
||||
input_stream: 'input'
|
||||
input_stream: 'prev_loop'
|
||||
output_stream: 'passed_through_input'
|
||||
output_stream: 'passed_through_prev_loop'
|
||||
}
|
||||
node {
|
||||
calculator: 'DroppingGateCalculator'
|
||||
input_stream: 'input'
|
||||
input_stream: 'DISALLOW:delay_next_output'
|
||||
output_stream: 'loop'
|
||||
}
|
||||
node {
|
||||
calculator: 'MakePairCalculator'
|
||||
input_stream: 'passed_through_input'
|
||||
input_stream: 'passed_through_prev_loop'
|
||||
output_stream: 'passed_through_input_and_prev_loop'
|
||||
}
|
||||
)");
|
||||
tool::AddVectorSink("passed_through_input_and_prev_loop", &graph_config,
|
||||
&output_packets_);
|
||||
MP_ASSERT_OK(graph_.Initialize(graph_config, {}));
|
||||
MP_ASSERT_OK(graph_.StartRun({}));
|
||||
}
|
||||
|
||||
void SendPackets(int timestamp, int input, bool delay_next_output) {
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"input", MakePacket<int>(input).At(Timestamp(timestamp))));
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"delay_next_output",
|
||||
MakePacket<bool>(delay_next_output).At(Timestamp(timestamp))));
|
||||
}
|
||||
|
||||
CalculatorGraph graph_;
|
||||
std::vector<Packet> output_packets_;
|
||||
};
|
||||
|
||||
TEST_F(PreviousLoopbackCalculatorDelayBehaviorTest, MultipleDelayedOutputs) {
|
||||
SendPackets(/*timestamp=*/1, /*input=*/1, /*delay_next_output=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/2, /*input=*/2, /*delay_next_output=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/3, /*input=*/3, /*delay_next_output=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/5, /*input=*/5, /*delay_next_output=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket())),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket())),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/15, /*input=*/15, /*delay_next_output=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(
|
||||
PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), EmptyPacket())),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket())),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), EmptyPacket())),
|
||||
PairPacket(Timestamp(15), Pair(IntPacket(15), IntPacket(5)))));
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(PreviousLoopbackCalculatorDelayBehaviorTest,
|
||||
NonDelayedOutputFollowedByMultipleDelayedOutputs) {
|
||||
SendPackets(/*timestamp=*/1, /*input=*/1, /*delay_next_output=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/2, /*input=*/2, /*delay_next_output=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), IntPacket(1)))));
|
||||
|
||||
SendPackets(/*timestamp=*/3, /*input=*/3, /*delay_next_output=*/true);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), IntPacket(1)))));
|
||||
|
||||
SendPackets(/*timestamp=*/5, /*input=*/5, /*delay_next_output=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), IntPacket(1))),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket())),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), EmptyPacket()))));
|
||||
|
||||
SendPackets(/*timestamp=*/15, /*input=*/15, /*delay_next_output=*/false);
|
||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||
EXPECT_THAT(
|
||||
output_packets_,
|
||||
ElementsAre(
|
||||
PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())),
|
||||
PairPacket(Timestamp(2), Pair(IntPacket(2), IntPacket(1))),
|
||||
PairPacket(Timestamp(3), Pair(IntPacket(3), EmptyPacket())),
|
||||
PairPacket(Timestamp(5), Pair(IntPacket(5), EmptyPacket())),
|
||||
PairPacket(Timestamp(15), Pair(IntPacket(15), IntPacket(5)))));
|
||||
|
||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
// 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_SPLIT_NORMALIZED_LANDMARK_LIST_CALCULATOR_H_ // NOLINT
|
||||
#define MEDIAPIPE_CALCULATORS_CORE_SPLIT_NORMALIZED_LANDMARK_LIST_CALCULATOR_H_ // NOLINT
|
||||
|
||||
#include "mediapipe/calculators/core/split_vector_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/port/canonical_errors.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Splits an input packet with NormalizedLandmarkList into
|
||||
// multiple NormalizedLandmarkList 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 NormalizedLandmark. If "element_only" is false, ranges can be
|
||||
// non-zero in size and all outputs will be of type NormalizedLandmarkList.
|
||||
// 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
|
||||
// NormalizedLandmarkList.
|
||||
class SplitNormalizedLandmarkListCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().NumEntries() == 1);
|
||||
RET_CHECK(cc->Outputs().NumEntries() != 0);
|
||||
|
||||
cc->Inputs().Index(0).Set<NormalizedLandmarkList>();
|
||||
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
|
||||
|
||||
if (options.combine_outputs()) {
|
||||
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1);
|
||||
cc->Outputs().Index(0).Set<NormalizedLandmarkList>();
|
||||
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.");
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
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.");
|
||||
}
|
||||
|
||||
// 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(
|
||||
"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<NormalizedLandmark>();
|
||||
} else {
|
||||
cc->Outputs().Index(i).Set<NormalizedLandmarkList>();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
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();
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
const NormalizedLandmarkList& input =
|
||||
cc->Inputs().Index(0).Get<NormalizedLandmarkList>();
|
||||
RET_CHECK_GE(input.landmark_size(), max_range_end_);
|
||||
|
||||
if (combine_outputs_) {
|
||||
NormalizedLandmarkList output;
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
for (int j = ranges_[i].first; j < ranges_[i].second; ++j) {
|
||||
const NormalizedLandmark& input_landmark = input.landmark(j);
|
||||
*output.add_landmark() = input_landmark;
|
||||
}
|
||||
}
|
||||
RET_CHECK_EQ(output.landmark_size(), total_elements_);
|
||||
cc->Outputs().Index(0).AddPacket(
|
||||
MakePacket<NormalizedLandmarkList>(output).At(cc->InputTimestamp()));
|
||||
} else {
|
||||
if (element_only_) {
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
cc->Outputs().Index(i).AddPacket(
|
||||
MakePacket<NormalizedLandmark>(input.landmark(ranges_[i].first))
|
||||
.At(cc->InputTimestamp()));
|
||||
}
|
||||
} else {
|
||||
for (int i = 0; i < ranges_.size(); ++i) {
|
||||
NormalizedLandmarkList output;
|
||||
for (int j = ranges_[i].first; j < ranges_[i].second; ++j) {
|
||||
const NormalizedLandmark& input_landmark = input.landmark(j);
|
||||
*output.add_landmark() = input_landmark;
|
||||
}
|
||||
cc->Outputs().Index(i).AddPacket(
|
||||
MakePacket<NormalizedLandmarkList>(output).At(
|
||||
cc->InputTimestamp()));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
std::vector<std::pair<int32, int32>> ranges_;
|
||||
int32 max_range_end_ = -1;
|
||||
int32 total_elements_ = 0;
|
||||
bool element_only_ = false;
|
||||
bool combine_outputs_ = false;
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(SplitNormalizedLandmarkListCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
// NOLINTNEXTLINE
|
||||
#endif // MEDIAPIPE_CALCULATORS_CORE_SPLIT_NORMALIZED_LANDMARK_LIST_CALCULATOR_H_
|
||||
@@ -0,0 +1,404 @@
|
||||
// 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 "mediapipe/calculators/core/split_vector_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/calculator_runner.h"
|
||||
#include "mediapipe/framework/deps/file_path.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/port/gmock.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
|
||||
#include "mediapipe/framework/tool/validate_type.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
constexpr float kLocationVal = 3;
|
||||
|
||||
class SplitNormalizedLandmarkListCalculatorTest : public ::testing::Test {
|
||||
protected:
|
||||
void TearDown() { expected_landmarks_.reset(); }
|
||||
|
||||
void PrepareNormalizedLandmarkList(int list_size) {
|
||||
// Prepare input landmark list.
|
||||
input_landmarks_ = absl::make_unique<NormalizedLandmarkList>();
|
||||
expected_landmarks_ = absl::make_unique<NormalizedLandmarkList>();
|
||||
for (int i = 0; i < list_size; ++i) {
|
||||
NormalizedLandmark* landmark = input_landmarks_->add_landmark();
|
||||
landmark->set_x(i * kLocationVal);
|
||||
landmark->set_y(i * kLocationVal);
|
||||
landmark->set_z(i * kLocationVal);
|
||||
// Save the landmarks for comparison after the graph runs.
|
||||
*expected_landmarks_->add_landmark() = *landmark;
|
||||
}
|
||||
}
|
||||
|
||||
void ValidateListOutput(std::vector<Packet>& output_packets,
|
||||
int expected_elements, int input_begin_index) {
|
||||
ASSERT_EQ(1, output_packets.size());
|
||||
const NormalizedLandmarkList& output_landmarks =
|
||||
output_packets[0].Get<NormalizedLandmarkList>();
|
||||
ASSERT_EQ(expected_elements, output_landmarks.landmark_size());
|
||||
|
||||
for (int i = 0; i < expected_elements; ++i) {
|
||||
const NormalizedLandmark& expected_landmark =
|
||||
expected_landmarks_->landmark(input_begin_index + i);
|
||||
const NormalizedLandmark& result = output_landmarks.landmark(i);
|
||||
EXPECT_FLOAT_EQ(expected_landmark.x(), result.x());
|
||||
EXPECT_FLOAT_EQ(expected_landmark.y(), result.y());
|
||||
EXPECT_FLOAT_EQ(expected_landmark.z(), result.z());
|
||||
}
|
||||
}
|
||||
|
||||
void ValidateCombinedListOutput(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 NormalizedLandmarkList& output_landmarks =
|
||||
output_packets[0].Get<NormalizedLandmarkList>();
|
||||
ASSERT_EQ(expected_elements, output_landmarks.landmark_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 NormalizedLandmark& expected_landmark =
|
||||
expected_landmarks_->landmark(i);
|
||||
const NormalizedLandmark& result =
|
||||
output_landmarks.landmark(element_id);
|
||||
EXPECT_FLOAT_EQ(expected_landmark.x(), result.x());
|
||||
EXPECT_FLOAT_EQ(expected_landmark.y(), result.y());
|
||||
EXPECT_FLOAT_EQ(expected_landmark.z(), result.z());
|
||||
element_id++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void ValidateElementOutput(std::vector<Packet>& output_packets,
|
||||
int input_begin_index) {
|
||||
ASSERT_EQ(1, output_packets.size());
|
||||
|
||||
const NormalizedLandmark& output_landmark =
|
||||
output_packets[0].Get<NormalizedLandmark>();
|
||||
ASSERT_TRUE(output_landmark.IsInitialized());
|
||||
|
||||
const NormalizedLandmark& expected_landmark =
|
||||
expected_landmarks_->landmark(input_begin_index);
|
||||
|
||||
EXPECT_FLOAT_EQ(expected_landmark.x(), output_landmark.x());
|
||||
EXPECT_FLOAT_EQ(expected_landmark.y(), output_landmark.y());
|
||||
EXPECT_FLOAT_EQ(expected_landmark.z(), output_landmark.z());
|
||||
}
|
||||
|
||||
std::unique_ptr<NormalizedLandmarkList> input_landmarks_ = nullptr;
|
||||
std::unique_ptr<NormalizedLandmarkList> expected_landmarks_ = nullptr;
|
||||
std::unique_ptr<CalculatorRunner> runner_ = nullptr;
|
||||
};
|
||||
|
||||
TEST_F(SplitNormalizedLandmarkListCalculatorTest, SmokeTest) {
|
||||
PrepareNormalizedLandmarkList(/*list_size=*/5);
|
||||
ASSERT_NE(input_landmarks_, nullptr);
|
||||
|
||||
// Prepare a graph to use the SplitNormalizedLandmarkListCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "landmarks_in"
|
||||
node {
|
||||
calculator: "SplitNormalizedLandmarkListCalculator"
|
||||
input_stream: "landmarks_in"
|
||||
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({}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"landmarks_in", Adopt(input_landmarks_.release()).At(Timestamp(0))));
|
||||
// Wait until the calculator finishes processing.
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
ValidateListOutput(range_0_packets, /*expected_elements=*/1,
|
||||
/*input_begin_index=*/0);
|
||||
ValidateListOutput(range_1_packets, /*expected_elements=*/3,
|
||||
/*input_begin_index=*/1);
|
||||
ValidateListOutput(range_2_packets, /*expected_elements=*/1,
|
||||
/*input_begin_index=*/4);
|
||||
|
||||
// Fully close the graph at the end.
|
||||
MP_ASSERT_OK(graph.CloseInputStream("landmarks_in"));
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(SplitNormalizedLandmarkListCalculatorTest, InvalidRangeTest) {
|
||||
// Prepare a graph to use the SplitNormalizedLandmarkListCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "landmarks_in"
|
||||
node {
|
||||
calculator: "SplitNormalizedLandmarkListCalculator"
|
||||
input_stream: "landmarks_in"
|
||||
output_stream: "range_0"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 0 }
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
// The graph should fail running because of an invalid range (begin == end).
|
||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||
}
|
||||
|
||||
TEST_F(SplitNormalizedLandmarkListCalculatorTest,
|
||||
InvalidOutputStreamCountTest) {
|
||||
// Prepare a graph to use the SplitNormalizedLandmarkListCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "landmarks_in"
|
||||
node {
|
||||
calculator: "SplitNormalizedLandmarkListCalculator"
|
||||
input_stream: "landmarks_in"
|
||||
output_stream: "range_0"
|
||||
output_stream: "range_1"
|
||||
options {
|
||||
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||
ranges: { begin: 0 end: 1 }
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
// 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(SplitNormalizedLandmarkListCalculatorTest,
|
||||
InvalidCombineOutputsMultipleOutputsTest) {
|
||||
// Prepare a graph to use the SplitNormalizedLandmarkListCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "landmarks_in"
|
||||
node {
|
||||
calculator: "SplitNormalizedLandmarkListCalculator"
|
||||
input_stream: "landmarks_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(SplitNormalizedLandmarkListCalculatorTest,
|
||||
InvalidOverlappingRangesTest) {
|
||||
// Prepare a graph to use the SplitNormalizedLandmarkListCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "landmarks_in"
|
||||
node {
|
||||
calculator: "SplitNormalizedLandmarkListCalculator"
|
||||
input_stream: "landmarks_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(SplitNormalizedLandmarkListCalculatorTest, SmokeTestElementOnly) {
|
||||
PrepareNormalizedLandmarkList(/*list_size=*/5);
|
||||
ASSERT_NE(input_landmarks_, nullptr);
|
||||
|
||||
// Prepare a graph to use the SplitNormalizedLandmarkListCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "landmarks_in"
|
||||
node {
|
||||
calculator: "SplitNormalizedLandmarkListCalculator"
|
||||
input_stream: "landmarks_in"
|
||||
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({}));
|
||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"landmarks_in", Adopt(input_landmarks_.release()).At(Timestamp(0))));
|
||||
// Wait until the calculator finishes processing.
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
ValidateElementOutput(range_0_packets,
|
||||
/*input_begin_index=*/0);
|
||||
ValidateElementOutput(range_1_packets,
|
||||
/*input_begin_index=*/2);
|
||||
ValidateElementOutput(range_2_packets,
|
||||
/*input_begin_index=*/4);
|
||||
|
||||
// Fully close the graph at the end.
|
||||
MP_ASSERT_OK(graph.CloseInputStream("landmarks_in"));
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(SplitNormalizedLandmarkListCalculatorTest, SmokeTestCombiningOutputs) {
|
||||
PrepareNormalizedLandmarkList(/*list_size=*/5);
|
||||
ASSERT_NE(input_landmarks_, nullptr);
|
||||
|
||||
// Prepare a graph to use the SplitNormalizedLandmarkListCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "landmarks_in"
|
||||
node {
|
||||
calculator: "SplitNormalizedLandmarkListCalculator"
|
||||
input_stream: "landmarks_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(
|
||||
"landmarks_in", Adopt(input_landmarks_.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};
|
||||
ValidateCombinedListOutput(range_0_packets, /*expected_elements=*/3,
|
||||
input_begin_indices, input_end_indices);
|
||||
|
||||
// Fully close the graph at the end.
|
||||
MP_ASSERT_OK(graph.CloseInputStream("landmarks_in"));
|
||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
TEST_F(SplitNormalizedLandmarkListCalculatorTest,
|
||||
ElementOnlyDisablesVectorOutputs) {
|
||||
// Prepare a graph to use the SplitNormalizedLandmarkListCalculator.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "landmarks_in"
|
||||
node {
|
||||
calculator: "SplitNormalizedLandmarkListCalculator"
|
||||
input_stream: "landmarks_in"
|
||||
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 }
|
||||
element_only: true
|
||||
}
|
||||
}
|
||||
}
|
||||
)");
|
||||
|
||||
// Run the graph.
|
||||
CalculatorGraph graph;
|
||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -16,7 +16,9 @@
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/matrix.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
|
||||
@@ -48,13 +50,25 @@ typedef SplitVectorCalculator<::mediapipe::NormalizedLandmark, false>
|
||||
SplitLandmarkVectorCalculator;
|
||||
REGISTER_CALCULATOR(SplitLandmarkVectorCalculator);
|
||||
|
||||
typedef SplitVectorCalculator<::mediapipe::NormalizedLandmarkList, false>
|
||||
SplitNormalizedLandmarkListVectorCalculator;
|
||||
REGISTER_CALCULATOR(SplitNormalizedLandmarkListVectorCalculator);
|
||||
|
||||
typedef SplitVectorCalculator<::mediapipe::NormalizedRect, false>
|
||||
SplitNormalizedRectVectorCalculator;
|
||||
REGISTER_CALCULATOR(SplitNormalizedRectVectorCalculator);
|
||||
|
||||
typedef SplitVectorCalculator<Matrix, false> SplitMatrixVectorCalculator;
|
||||
REGISTER_CALCULATOR(SplitMatrixVectorCalculator);
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
typedef SplitVectorCalculator<::tflite::gpu::gl::GlBuffer, true>
|
||||
MovableSplitGlBufferVectorCalculator;
|
||||
REGISTER_CALCULATOR(MovableSplitGlBufferVectorCalculator);
|
||||
#endif
|
||||
|
||||
typedef SplitVectorCalculator<::mediapipe::Detection, false>
|
||||
SplitDetectionVectorCalculator;
|
||||
REGISTER_CALCULATOR(SplitDetectionVectorCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -86,6 +86,14 @@ mediapipe_cc_proto_library(
|
||||
deps = [":opencv_image_encoder_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "opencv_encoded_image_to_image_frame_calculator_cc_proto",
|
||||
srcs = ["opencv_encoded_image_to_image_frame_calculator.proto"],
|
||||
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":opencv_encoded_image_to_image_frame_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "mask_overlay_calculator_cc_proto",
|
||||
srcs = ["mask_overlay_calculator.proto"],
|
||||
@@ -172,6 +180,7 @@ cc_library(
|
||||
srcs = ["opencv_encoded_image_to_image_frame_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":opencv_encoded_image_to_image_frame_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/port:opencv_imgcodecs",
|
||||
@@ -422,9 +431,12 @@ cc_library(
|
||||
":recolor_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/util:color_cc_proto",
|
||||
"//mediapipe/framework/port:opencv_core",
|
||||
"//mediapipe/framework/port:opencv_imgproc",
|
||||
] + select({
|
||||
"//mediapipe/gpu:disable_gpu": [],
|
||||
"//conditions:default": [
|
||||
@@ -554,6 +566,27 @@ proto_library(
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "opencv_encoded_image_to_image_frame_calculator_proto",
|
||||
srcs = ["opencv_encoded_image_to_image_frame_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "feature_detector_calculator_proto",
|
||||
srcs = ["feature_detector_calculator.proto"],
|
||||
deps = ["//mediapipe/framework:calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "feature_detector_calculator_cc_proto",
|
||||
srcs = ["feature_detector_calculator.proto"],
|
||||
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":feature_detector_calculator_proto"],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "mask_overlay_calculator",
|
||||
srcs = ["mask_overlay_calculator.cc"],
|
||||
@@ -569,3 +602,30 @@ cc_library(
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "feature_detector_calculator",
|
||||
srcs = ["feature_detector_calculator.cc"],
|
||||
visibility = ["//mediapipe:__subpackages__"],
|
||||
deps = [
|
||||
":feature_detector_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_frame",
|
||||
"//mediapipe/framework/formats:image_frame_opencv",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:video_stream_header",
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:opencv_core",
|
||||
"//mediapipe/framework/port:opencv_features2d",
|
||||
"//mediapipe/framework/port:opencv_imgproc",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/port:threadpool",
|
||||
"//mediapipe/framework/tool:options_util",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/synchronization",
|
||||
"@org_tensorflow//tensorflow/lite:framework",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,210 @@
|
||||
// Copyright 2020 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 <vector>
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "absl/synchronization/blocking_counter.h"
|
||||
#include "mediapipe/calculators/image/feature_detector_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.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/opencv_core_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_features2d_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/port/threadpool.h"
|
||||
#include "mediapipe/framework/tool/options_util.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
const char kOptionsTag[] = "OPTIONS";
|
||||
const int kPatchSize = 32;
|
||||
const int kNumThreads = 16;
|
||||
|
||||
// A calculator to apply local feature detection.
|
||||
// Input stream:
|
||||
// IMAGE: Input image frame of type ImageFrame from video stream.
|
||||
// Output streams:
|
||||
// FEATURES: The detected keypoints from input image as vector<cv::KeyPoint>.
|
||||
// PATCHES: Optional output the extracted patches as vector<cv::Mat>
|
||||
class FeatureDetectorCalculator : public CalculatorBase {
|
||||
public:
|
||||
~FeatureDetectorCalculator() override = default;
|
||||
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
FeatureDetectorCalculatorOptions options_;
|
||||
cv::Ptr<cv::Feature2D> feature_detector_;
|
||||
std::unique_ptr<::mediapipe::ThreadPool> pool_;
|
||||
|
||||
// Create image pyramid based on input image.
|
||||
void ComputeImagePyramid(const cv::Mat& input_image,
|
||||
std::vector<cv::Mat>* image_pyramid);
|
||||
|
||||
// Extract the patch for single feature with image pyramid.
|
||||
cv::Mat ExtractPatch(const cv::KeyPoint& feature,
|
||||
const std::vector<cv::Mat>& image_pyramid);
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(FeatureDetectorCalculator);
|
||||
|
||||
::mediapipe::Status FeatureDetectorCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
if (cc->Inputs().HasTag("IMAGE")) {
|
||||
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
||||
}
|
||||
if (cc->Outputs().HasTag("FEATURES")) {
|
||||
cc->Outputs().Tag("FEATURES").Set<std::vector<cv::KeyPoint>>();
|
||||
}
|
||||
if (cc->Outputs().HasTag("LANDMARKS")) {
|
||||
cc->Outputs().Tag("LANDMARKS").Set<NormalizedLandmarkList>();
|
||||
}
|
||||
if (cc->Outputs().HasTag("PATCHES")) {
|
||||
cc->Outputs().Tag("PATCHES").Set<std::vector<TfLiteTensor>>();
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status FeatureDetectorCalculator::Open(CalculatorContext* cc) {
|
||||
options_ =
|
||||
tool::RetrieveOptions(cc->Options(), cc->InputSidePackets(), kOptionsTag)
|
||||
.GetExtension(FeatureDetectorCalculatorOptions::ext);
|
||||
feature_detector_ = cv::ORB::create(
|
||||
options_.max_features(), options_.scale_factor(),
|
||||
options_.pyramid_level(), kPatchSize - 1, 0, 2, cv::ORB::FAST_SCORE);
|
||||
pool_ = absl::make_unique<::mediapipe::ThreadPool>("ThreadPool", kNumThreads);
|
||||
pool_->StartWorkers();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status FeatureDetectorCalculator::Process(CalculatorContext* cc) {
|
||||
const Timestamp& timestamp = cc->InputTimestamp();
|
||||
if (timestamp == Timestamp::PreStream()) {
|
||||
// Indicator packet.
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
InputStream* input_frame = &(cc->Inputs().Tag("IMAGE"));
|
||||
cv::Mat input_view = formats::MatView(&input_frame->Get<ImageFrame>());
|
||||
cv::Mat grayscale_view;
|
||||
cv::cvtColor(input_view, grayscale_view, cv::COLOR_RGB2GRAY);
|
||||
|
||||
std::vector<cv::KeyPoint> keypoints;
|
||||
feature_detector_->detect(grayscale_view, keypoints);
|
||||
if (keypoints.size() > options_.max_features()) {
|
||||
keypoints.resize(options_.max_features());
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("FEATURES")) {
|
||||
auto features_ptr = absl::make_unique<std::vector<cv::KeyPoint>>(keypoints);
|
||||
cc->Outputs().Tag("FEATURES").Add(features_ptr.release(), timestamp);
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("LANDMARKS")) {
|
||||
auto landmarks_ptr = absl::make_unique<NormalizedLandmarkList>();
|
||||
for (int j = 0; j < keypoints.size(); ++j) {
|
||||
auto feature_landmark = landmarks_ptr->add_landmark();
|
||||
feature_landmark->set_x(keypoints[j].pt.x / grayscale_view.cols);
|
||||
feature_landmark->set_y(keypoints[j].pt.y / grayscale_view.rows);
|
||||
}
|
||||
cc->Outputs().Tag("LANDMARKS").Add(landmarks_ptr.release(), timestamp);
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("PATCHES")) {
|
||||
std::vector<cv::Mat> image_pyramid;
|
||||
ComputeImagePyramid(grayscale_view, &image_pyramid);
|
||||
std::vector<cv::Mat> patch_mat;
|
||||
patch_mat.resize(keypoints.size());
|
||||
absl::BlockingCounter counter(keypoints.size());
|
||||
for (int i = 0; i < keypoints.size(); i++) {
|
||||
pool_->Schedule(
|
||||
[this, &image_pyramid, &keypoints, &patch_mat, i, &counter] {
|
||||
patch_mat[i] = ExtractPatch(keypoints[i], image_pyramid);
|
||||
counter.DecrementCount();
|
||||
});
|
||||
}
|
||||
counter.Wait();
|
||||
const int batch_size = options_.max_features();
|
||||
auto patches = absl::make_unique<std::vector<TfLiteTensor>>();
|
||||
TfLiteTensor tensor;
|
||||
tensor.type = kTfLiteFloat32;
|
||||
tensor.dims = TfLiteIntArrayCreate(4);
|
||||
tensor.dims->data[0] = batch_size;
|
||||
tensor.dims->data[1] = kPatchSize;
|
||||
tensor.dims->data[2] = kPatchSize;
|
||||
tensor.dims->data[3] = 1;
|
||||
int num_bytes = batch_size * kPatchSize * kPatchSize * sizeof(float);
|
||||
tensor.data.data = malloc(num_bytes);
|
||||
tensor.bytes = num_bytes;
|
||||
tensor.allocation_type = kTfLiteArenaRw;
|
||||
float* tensor_buffer = tensor.data.f;
|
||||
for (int i = 0; i < keypoints.size(); i++) {
|
||||
for (int j = 0; j < patch_mat[i].rows; ++j) {
|
||||
for (int k = 0; k < patch_mat[i].cols; ++k) {
|
||||
*tensor_buffer++ = patch_mat[i].at<uchar>(j, k) / 128.0f - 1.0f;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (int i = keypoints.size() * kPatchSize * kPatchSize; i < num_bytes / 4;
|
||||
i++) {
|
||||
*tensor_buffer++ = 0;
|
||||
}
|
||||
|
||||
patches->emplace_back(tensor);
|
||||
cc->Outputs().Tag("PATCHES").Add(patches.release(), timestamp);
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
void FeatureDetectorCalculator::ComputeImagePyramid(
|
||||
const cv::Mat& input_image, std::vector<cv::Mat>* image_pyramid) {
|
||||
cv::Mat tmp_image = input_image;
|
||||
cv::Mat src_image = input_image;
|
||||
for (int i = 0; i < options_.pyramid_level(); ++i) {
|
||||
image_pyramid->push_back(src_image);
|
||||
cv::resize(src_image, tmp_image, cv::Size(), 1.0f / options_.scale_factor(),
|
||||
1.0f / options_.scale_factor());
|
||||
src_image = tmp_image;
|
||||
}
|
||||
}
|
||||
|
||||
cv::Mat FeatureDetectorCalculator::ExtractPatch(
|
||||
const cv::KeyPoint& feature, const std::vector<cv::Mat>& image_pyramid) {
|
||||
cv::Mat img = image_pyramid[feature.octave];
|
||||
float scale_factor = 1 / pow(options_.scale_factor(), feature.octave);
|
||||
cv::Point2f center =
|
||||
cv::Point2f(feature.pt.x * scale_factor, feature.pt.y * scale_factor);
|
||||
cv::Mat rot = cv::getRotationMatrix2D(center, feature.angle, 1.0);
|
||||
rot.at<double>(0, 2) += kPatchSize / 2 - center.x;
|
||||
rot.at<double>(1, 2) += kPatchSize / 2 - center.y;
|
||||
cv::Mat cropped_img;
|
||||
// perform the affine transformation
|
||||
cv::warpAffine(img, cropped_img, rot, cv::Size(kPatchSize, kPatchSize),
|
||||
cv::INTER_LINEAR);
|
||||
return cropped_img;
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,24 @@
|
||||
// Options for FeatureDetectorCalculator
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe;
|
||||
|
||||
import "mediapipe/framework/calculator.proto";
|
||||
|
||||
message FeatureDetectorCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional FeatureDetectorCalculatorOptions ext = 278741680;
|
||||
}
|
||||
|
||||
// Set to true if output patches, otherwise only output cv::KeyPoint
|
||||
optional bool output_patch = 1;
|
||||
|
||||
// The max number of detected features.
|
||||
optional int32 max_features = 2 [default = 200];
|
||||
|
||||
// The number of pyramid levels.
|
||||
optional int32 pyramid_level = 3 [default = 4];
|
||||
|
||||
// Pyramid decimation ratio.
|
||||
optional float scale_factor = 4 [default = 1.2];
|
||||
}
|
||||
@@ -219,8 +219,10 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
|
||||
const auto& input_img = cc->Inputs().Tag(kImageTag).Get<ImageFrame>();
|
||||
cv::Mat input_mat = formats::MatView(&input_img);
|
||||
|
||||
auto [target_width, target_height, rect_center_x, rect_center_y, rotation] =
|
||||
GetCropSpecs(cc, input_img.Width(), input_img.Height());
|
||||
RectSpec specs = GetCropSpecs(cc, input_img.Width(), input_img.Height());
|
||||
int target_width = specs.width, target_height = specs.height,
|
||||
rect_center_x = specs.center_x, rect_center_y = specs.center_y;
|
||||
float rotation = specs.rotation;
|
||||
|
||||
// Get border mode and value for OpenCV.
|
||||
int border_mode;
|
||||
@@ -403,8 +405,10 @@ void ImageCroppingCalculator::GetOutputDimensions(CalculatorContext* cc,
|
||||
int src_width, int src_height,
|
||||
int* dst_width,
|
||||
int* dst_height) {
|
||||
auto [crop_width, crop_height, x_center, y_center, rotation] =
|
||||
GetCropSpecs(cc, src_width, src_height);
|
||||
RectSpec specs = GetCropSpecs(cc, src_width, src_height);
|
||||
int crop_width = specs.width, crop_height = specs.height,
|
||||
x_center = specs.center_x, y_center = specs.center_y;
|
||||
float rotation = specs.rotation;
|
||||
|
||||
const float half_width = crop_width / 2.0f;
|
||||
const float half_height = crop_height / 2.0f;
|
||||
|
||||
@@ -19,6 +19,11 @@
|
||||
#include "mediapipe/gpu/gpu_buffer.h"
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
namespace {
|
||||
constexpr char kImageFrameTag[] = "IMAGE";
|
||||
constexpr char kGpuBufferTag[] = "IMAGE_GPU";
|
||||
} // namespace
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Extracts image properties from the input image and outputs the properties.
|
||||
@@ -40,13 +45,14 @@ namespace mediapipe {
|
||||
class ImagePropertiesCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag("IMAGE") ^ cc->Inputs().HasTag("IMAGE_GPU"));
|
||||
if (cc->Inputs().HasTag("IMAGE")) {
|
||||
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
||||
RET_CHECK(cc->Inputs().HasTag(kImageFrameTag) ^
|
||||
cc->Inputs().HasTag(kGpuBufferTag));
|
||||
if (cc->Inputs().HasTag(kImageFrameTag)) {
|
||||
cc->Inputs().Tag(kImageFrameTag).Set<ImageFrame>();
|
||||
}
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||
cc->Inputs().Tag("IMAGE_GPU").Set<::mediapipe::GpuBuffer>();
|
||||
if (cc->Inputs().HasTag(kGpuBufferTag)) {
|
||||
cc->Inputs().Tag(kGpuBufferTag).Set<::mediapipe::GpuBuffer>();
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
@@ -66,16 +72,17 @@ class ImagePropertiesCalculator : public CalculatorBase {
|
||||
int width;
|
||||
int height;
|
||||
|
||||
if (cc->Inputs().HasTag("IMAGE") && !cc->Inputs().Tag("IMAGE").IsEmpty()) {
|
||||
const auto& image = cc->Inputs().Tag("IMAGE").Get<ImageFrame>();
|
||||
if (cc->Inputs().HasTag(kImageFrameTag) &&
|
||||
!cc->Inputs().Tag(kImageFrameTag).IsEmpty()) {
|
||||
const auto& image = cc->Inputs().Tag(kImageFrameTag).Get<ImageFrame>();
|
||||
width = image.Width();
|
||||
height = image.Height();
|
||||
}
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU") &&
|
||||
!cc->Inputs().Tag("IMAGE_GPU").IsEmpty()) {
|
||||
if (cc->Inputs().HasTag(kGpuBufferTag) &&
|
||||
!cc->Inputs().Tag(kGpuBufferTag).IsEmpty()) {
|
||||
const auto& image =
|
||||
cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
||||
cc->Inputs().Tag(kGpuBufferTag).Get<mediapipe::GpuBuffer>();
|
||||
width = image.width();
|
||||
height = image.height();
|
||||
}
|
||||
|
||||
@@ -47,6 +47,9 @@ namespace mediapipe {
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
namespace {
|
||||
constexpr char kImageFrameTag[] = "IMAGE";
|
||||
constexpr char kGpuBufferTag[] = "IMAGE_GPU";
|
||||
|
||||
int RotationModeToDegrees(mediapipe::RotationMode_Mode rotation) {
|
||||
switch (rotation) {
|
||||
case mediapipe::RotationMode_Mode_UNKNOWN:
|
||||
@@ -95,7 +98,7 @@ mediapipe::ScaleMode_Mode ParseScaleMode(
|
||||
// Scales, rotates, and flips images horizontally or vertically.
|
||||
//
|
||||
// Input:
|
||||
// One of the following two tags:
|
||||
// One of the following tags:
|
||||
// IMAGE: ImageFrame representing the input image.
|
||||
// IMAGE_GPU: GpuBuffer representing the input image.
|
||||
//
|
||||
@@ -113,7 +116,7 @@ mediapipe::ScaleMode_Mode ParseScaleMode(
|
||||
// corresponding field in the calculator options.
|
||||
//
|
||||
// Output:
|
||||
// One of the following two tags:
|
||||
// One of the following tags:
|
||||
// IMAGE - ImageFrame representing the output image.
|
||||
// IMAGE_GPU - GpuBuffer representing the output image.
|
||||
//
|
||||
@@ -152,7 +155,8 @@ mediapipe::ScaleMode_Mode ParseScaleMode(
|
||||
// Note: To enable horizontal or vertical flipping, specify them in the
|
||||
// calculator options. Flipping is applied after rotation.
|
||||
//
|
||||
// Note: Only scale mode STRETCH is currently supported on CPU.
|
||||
// Note: Input defines output, so only matchig types supported:
|
||||
// IMAGE -> IMAGE or IMAGE_GPU -> IMAGE_GPU
|
||||
//
|
||||
class ImageTransformationCalculator : public CalculatorBase {
|
||||
public:
|
||||
@@ -186,7 +190,7 @@ class ImageTransformationCalculator : public CalculatorBase {
|
||||
|
||||
bool use_gpu_ = false;
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
GlCalculatorHelper helper_;
|
||||
GlCalculatorHelper gpu_helper_;
|
||||
std::unique_ptr<QuadRenderer> rgb_renderer_;
|
||||
std::unique_ptr<QuadRenderer> yuv_renderer_;
|
||||
std::unique_ptr<QuadRenderer> ext_rgb_renderer_;
|
||||
@@ -197,21 +201,22 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
// static
|
||||
::mediapipe::Status ImageTransformationCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag("IMAGE") ^ cc->Inputs().HasTag("IMAGE_GPU"));
|
||||
RET_CHECK(cc->Outputs().HasTag("IMAGE") ^ cc->Outputs().HasTag("IMAGE_GPU"));
|
||||
// Only one input can be set, and the output type must match.
|
||||
RET_CHECK(cc->Inputs().HasTag(kImageFrameTag) ^
|
||||
cc->Inputs().HasTag(kGpuBufferTag));
|
||||
|
||||
bool use_gpu = false;
|
||||
|
||||
if (cc->Inputs().HasTag("IMAGE")) {
|
||||
RET_CHECK(cc->Outputs().HasTag("IMAGE"));
|
||||
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
||||
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
|
||||
if (cc->Inputs().HasTag(kImageFrameTag)) {
|
||||
RET_CHECK(cc->Outputs().HasTag(kImageFrameTag));
|
||||
cc->Inputs().Tag(kImageFrameTag).Set<ImageFrame>();
|
||||
cc->Outputs().Tag(kImageFrameTag).Set<ImageFrame>();
|
||||
}
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||
RET_CHECK(cc->Outputs().HasTag("IMAGE_GPU"));
|
||||
cc->Inputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
|
||||
cc->Outputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
|
||||
if (cc->Inputs().HasTag(kGpuBufferTag)) {
|
||||
RET_CHECK(cc->Outputs().HasTag(kGpuBufferTag));
|
||||
cc->Inputs().Tag(kGpuBufferTag).Set<GpuBuffer>();
|
||||
cc->Outputs().Tag(kGpuBufferTag).Set<GpuBuffer>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
@@ -259,7 +264,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
|
||||
options_ = cc->Options<ImageTransformationCalculatorOptions>();
|
||||
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||
if (cc->Inputs().HasTag(kGpuBufferTag)) {
|
||||
use_gpu_ = true;
|
||||
}
|
||||
|
||||
@@ -300,7 +305,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
if (use_gpu_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
// Let the helper access the GL context information.
|
||||
MP_RETURN_IF_ERROR(helper_.Open(cc));
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||
#else
|
||||
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
@@ -328,18 +333,14 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
|
||||
if (use_gpu_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
if (cc->Inputs().Tag("IMAGE_GPU").IsEmpty()) {
|
||||
// Image is missing, hence no way to produce output image. (Timestamp
|
||||
// bound will be updated automatically.)
|
||||
if (cc->Inputs().Tag(kGpuBufferTag).IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
return helper_.RunInGlContext(
|
||||
return gpu_helper_.RunInGlContext(
|
||||
[this, cc]() -> ::mediapipe::Status { return RenderGpu(cc); });
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
} else {
|
||||
if (cc->Inputs().Tag("IMAGE").IsEmpty()) {
|
||||
// Image is missing, hence no way to produce output image. (Timestamp
|
||||
// bound will be updated automatically.)
|
||||
if (cc->Inputs().Tag(kImageFrameTag).IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
return RenderCpu(cc);
|
||||
@@ -354,7 +355,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
QuadRenderer* rgb_renderer = rgb_renderer_.release();
|
||||
QuadRenderer* yuv_renderer = yuv_renderer_.release();
|
||||
QuadRenderer* ext_rgb_renderer = ext_rgb_renderer_.release();
|
||||
helper_.RunInGlContext([rgb_renderer, yuv_renderer, ext_rgb_renderer] {
|
||||
gpu_helper_.RunInGlContext([rgb_renderer, yuv_renderer, ext_rgb_renderer] {
|
||||
if (rgb_renderer) {
|
||||
rgb_renderer->GlTeardown();
|
||||
delete rgb_renderer;
|
||||
@@ -376,18 +377,21 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
|
||||
::mediapipe::Status ImageTransformationCalculator::RenderCpu(
|
||||
CalculatorContext* cc) {
|
||||
int input_width = cc->Inputs().Tag("IMAGE").Get<ImageFrame>().Width();
|
||||
int input_height = cc->Inputs().Tag("IMAGE").Get<ImageFrame>().Height();
|
||||
cv::Mat input_mat;
|
||||
mediapipe::ImageFormat::Format format;
|
||||
|
||||
const auto& input_img = cc->Inputs().Tag("IMAGE").Get<ImageFrame>();
|
||||
cv::Mat input_mat = formats::MatView(&input_img);
|
||||
cv::Mat scaled_mat;
|
||||
const auto& input = cc->Inputs().Tag(kImageFrameTag).Get<ImageFrame>();
|
||||
input_mat = formats::MatView(&input);
|
||||
format = input.Format();
|
||||
|
||||
const int input_width = input_mat.cols;
|
||||
const int input_height = input_mat.rows;
|
||||
if (!output_height_ || !output_width_) {
|
||||
output_height_ = input_height;
|
||||
output_width_ = input_width;
|
||||
}
|
||||
|
||||
cv::Mat scaled_mat;
|
||||
if (scale_mode_ == mediapipe::ScaleMode_Mode_STRETCH) {
|
||||
cv::resize(input_mat, scaled_mat, cv::Size(output_width_, output_height_));
|
||||
} else {
|
||||
@@ -444,10 +448,12 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
}
|
||||
|
||||
std::unique_ptr<ImageFrame> output_frame(
|
||||
new ImageFrame(input_img.Format(), output_width, output_height));
|
||||
new ImageFrame(format, output_width, output_height));
|
||||
cv::Mat output_mat = formats::MatView(output_frame.get());
|
||||
flipped_mat.copyTo(output_mat);
|
||||
cc->Outputs().Tag("IMAGE").Add(output_frame.release(), cc->InputTimestamp());
|
||||
cc->Outputs()
|
||||
.Tag(kImageFrameTag)
|
||||
.Add(output_frame.release(), cc->InputTimestamp());
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -455,8 +461,9 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
::mediapipe::Status ImageTransformationCalculator::RenderGpu(
|
||||
CalculatorContext* cc) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
int input_width = cc->Inputs().Tag("IMAGE_GPU").Get<GpuBuffer>().width();
|
||||
int input_height = cc->Inputs().Tag("IMAGE_GPU").Get<GpuBuffer>().height();
|
||||
const auto& input = cc->Inputs().Tag(kGpuBufferTag).Get<GpuBuffer>();
|
||||
const int input_width = input.width();
|
||||
const int input_height = input.height();
|
||||
|
||||
int output_width;
|
||||
int output_height;
|
||||
@@ -472,7 +479,6 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
.Add(padding.release(), cc->InputTimestamp());
|
||||
}
|
||||
|
||||
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<GpuBuffer>();
|
||||
QuadRenderer* renderer = nullptr;
|
||||
GlTexture src1;
|
||||
|
||||
@@ -486,11 +492,11 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
{"video_frame_y", "video_frame_uv"}));
|
||||
}
|
||||
renderer = yuv_renderer_.get();
|
||||
src1 = helper_.CreateSourceTexture(input, 0);
|
||||
src1 = gpu_helper_.CreateSourceTexture(input, 0);
|
||||
} else // NOLINT(readability/braces)
|
||||
#endif // iOS
|
||||
{
|
||||
src1 = helper_.CreateSourceTexture(input);
|
||||
src1 = gpu_helper_.CreateSourceTexture(input);
|
||||
#if defined(TEXTURE_EXTERNAL_OES)
|
||||
if (src1.target() == GL_TEXTURE_EXTERNAL_OES) {
|
||||
if (!ext_rgb_renderer_) {
|
||||
@@ -516,10 +522,10 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
mediapipe::FrameRotation rotation =
|
||||
mediapipe::FrameRotationFromDegrees(RotationModeToDegrees(rotation_));
|
||||
|
||||
auto dst = helper_.CreateDestinationTexture(output_width, output_height,
|
||||
input.format());
|
||||
auto dst = gpu_helper_.CreateDestinationTexture(output_width, output_height,
|
||||
input.format());
|
||||
|
||||
helper_.BindFramebuffer(dst); // GL_TEXTURE0
|
||||
gpu_helper_.BindFramebuffer(dst); // GL_TEXTURE0
|
||||
glActiveTexture(GL_TEXTURE1);
|
||||
glBindTexture(src1.target(), src1.name());
|
||||
|
||||
@@ -534,8 +540,8 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||
// Execute GL commands, before getting result.
|
||||
glFlush();
|
||||
|
||||
auto output = dst.GetFrame<GpuBuffer>();
|
||||
cc->Outputs().Tag("IMAGE_GPU").Add(output.release(), cc->InputTimestamp());
|
||||
auto output = dst.template GetFrame<GpuBuffer>();
|
||||
cc->Outputs().Tag(kGpuBufferTag).Add(output.release(), cc->InputTimestamp());
|
||||
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "mediapipe/calculators/image/opencv_encoded_image_to_image_frame_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||
#include "mediapipe/framework/port/opencv_imgcodecs_inc.h"
|
||||
@@ -34,7 +35,11 @@ namespace mediapipe {
|
||||
class OpenCvEncodedImageToImageFrameCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
mediapipe::OpenCvEncodedImageToImageFrameCalculatorOptions options_;
|
||||
};
|
||||
|
||||
::mediapipe::Status OpenCvEncodedImageToImageFrameCalculator::GetContract(
|
||||
@@ -44,13 +49,29 @@ class OpenCvEncodedImageToImageFrameCalculator : public CalculatorBase {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status OpenCvEncodedImageToImageFrameCalculator::Open(
|
||||
CalculatorContext* cc) {
|
||||
options_ =
|
||||
cc->Options<mediapipe::OpenCvEncodedImageToImageFrameCalculatorOptions>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status OpenCvEncodedImageToImageFrameCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
const std::string& contents = cc->Inputs().Index(0).Get<std::string>();
|
||||
const std::vector<char> contents_vector(contents.begin(), contents.end());
|
||||
cv::Mat decoded_mat =
|
||||
cv::imdecode(contents_vector, -1 /* return the loaded image as-is */);
|
||||
|
||||
cv::Mat decoded_mat;
|
||||
if (options_.apply_orientation_from_exif_data()) {
|
||||
// We want to respect the orientation from the EXIF data, which
|
||||
// IMREAD_UNCHANGED ignores, but otherwise we want to be as permissive as
|
||||
// possible with our reading flags. Therefore, we use IMREAD_ANYCOLOR and
|
||||
// IMREAD_ANYDEPTH.
|
||||
decoded_mat = cv::imdecode(contents_vector,
|
||||
cv::IMREAD_ANYCOLOR | cv::IMREAD_ANYDEPTH);
|
||||
} else {
|
||||
// Return the loaded image as-is
|
||||
decoded_mat = cv::imdecode(contents_vector, cv::IMREAD_UNCHANGED);
|
||||
}
|
||||
ImageFormat::Format image_format = ImageFormat::UNKNOWN;
|
||||
cv::Mat output_mat;
|
||||
switch (decoded_mat.channels()) {
|
||||
@@ -70,7 +91,8 @@ class OpenCvEncodedImageToImageFrameCalculator : public CalculatorBase {
|
||||
<< "Unsupported number of channels: " << decoded_mat.channels();
|
||||
}
|
||||
std::unique_ptr<ImageFrame> output_frame = absl::make_unique<ImageFrame>(
|
||||
image_format, decoded_mat.size().width, decoded_mat.size().height);
|
||||
image_format, decoded_mat.size().width, decoded_mat.size().height,
|
||||
ImageFrame::kGlDefaultAlignmentBoundary);
|
||||
output_mat.copyTo(formats::MatView(output_frame.get()));
|
||||
cc->Outputs().Index(0).Add(output_frame.release(), cc->InputTimestamp());
|
||||
return ::mediapipe::OkStatus();
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright 2020 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 OpenCvEncodedImageToImageFrameCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional OpenCvEncodedImageToImageFrameCalculatorOptions ext = 303447308;
|
||||
}
|
||||
|
||||
// If set, we will attempt to automatically apply the orientation specified by
|
||||
// the image's EXIF data when loading the image. Otherwise, the image data
|
||||
// will be loaded as-is.
|
||||
optional bool apply_orientation_from_exif_data = 1 [default = false];
|
||||
}
|
||||
@@ -17,6 +17,9 @@
|
||||
#include "mediapipe/calculators/image/recolor_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||
#include "mediapipe/framework/port/opencv_core_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/util/color.pb.h"
|
||||
@@ -29,6 +32,11 @@
|
||||
|
||||
namespace {
|
||||
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
|
||||
|
||||
constexpr char kImageFrameTag[] = "IMAGE";
|
||||
constexpr char kMaskCpuTag[] = "MASK";
|
||||
constexpr char kGpuBufferTag[] = "IMAGE_GPU";
|
||||
constexpr char kMaskGpuTag[] = "MASK_GPU";
|
||||
} // namespace
|
||||
|
||||
namespace mediapipe {
|
||||
@@ -39,8 +47,6 @@ namespace mediapipe {
|
||||
// The luminance of the input image is used to adjust the blending weight,
|
||||
// to help preserve image textures.
|
||||
//
|
||||
// TODO implement cpu support.
|
||||
//
|
||||
// Inputs:
|
||||
// One of the following IMAGE tags:
|
||||
// IMAGE: An ImageFrame input image, RGB or RGBA.
|
||||
@@ -71,6 +77,8 @@ namespace mediapipe {
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// Note: Cannot mix-match CPU & GPU inputs/outputs.
|
||||
// CPU-in & CPU-out <or> GPU-in & GPU-out
|
||||
class RecolorCalculator : public CalculatorBase {
|
||||
public:
|
||||
RecolorCalculator() = default;
|
||||
@@ -109,35 +117,42 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
||||
bool use_gpu = false;
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||
cc->Inputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
|
||||
if (cc->Inputs().HasTag(kGpuBufferTag)) {
|
||||
cc->Inputs().Tag(kGpuBufferTag).Set<mediapipe::GpuBuffer>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
if (cc->Inputs().HasTag("IMAGE")) {
|
||||
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
||||
if (cc->Inputs().HasTag(kImageFrameTag)) {
|
||||
cc->Inputs().Tag(kImageFrameTag).Set<ImageFrame>();
|
||||
}
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
if (cc->Inputs().HasTag("MASK_GPU")) {
|
||||
cc->Inputs().Tag("MASK_GPU").Set<mediapipe::GpuBuffer>();
|
||||
if (cc->Inputs().HasTag(kMaskGpuTag)) {
|
||||
cc->Inputs().Tag(kMaskGpuTag).Set<mediapipe::GpuBuffer>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
if (cc->Inputs().HasTag("MASK")) {
|
||||
cc->Inputs().Tag("MASK").Set<ImageFrame>();
|
||||
if (cc->Inputs().HasTag(kMaskCpuTag)) {
|
||||
cc->Inputs().Tag(kMaskCpuTag).Set<ImageFrame>();
|
||||
}
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
if (cc->Outputs().HasTag("IMAGE_GPU")) {
|
||||
cc->Outputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
|
||||
if (cc->Outputs().HasTag(kGpuBufferTag)) {
|
||||
cc->Outputs().Tag(kGpuBufferTag).Set<mediapipe::GpuBuffer>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
if (cc->Outputs().HasTag("IMAGE")) {
|
||||
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
|
||||
if (cc->Outputs().HasTag(kImageFrameTag)) {
|
||||
cc->Outputs().Tag(kImageFrameTag).Set<ImageFrame>();
|
||||
}
|
||||
|
||||
// Confirm only one of the input streams is present.
|
||||
RET_CHECK(cc->Inputs().HasTag(kImageFrameTag) ^
|
||||
cc->Inputs().HasTag(kGpuBufferTag));
|
||||
// Confirm only one of the output streams is present.
|
||||
RET_CHECK(cc->Outputs().HasTag(kImageFrameTag) ^
|
||||
cc->Outputs().HasTag(kGpuBufferTag));
|
||||
|
||||
if (use_gpu) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||
@@ -150,7 +165,7 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
||||
::mediapipe::Status RecolorCalculator::Open(CalculatorContext* cc) {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||
if (cc->Inputs().HasTag(kGpuBufferTag)) {
|
||||
use_gpu_ = true;
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||
@@ -193,17 +208,74 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
||||
}
|
||||
|
||||
::mediapipe::Status RecolorCalculator::RenderCpu(CalculatorContext* cc) {
|
||||
return ::mediapipe::UnimplementedError("CPU support is not implemented yet.");
|
||||
if (cc->Inputs().Tag(kMaskCpuTag).IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
// Get inputs and setup output.
|
||||
const auto& input_img = cc->Inputs().Tag(kImageFrameTag).Get<ImageFrame>();
|
||||
const auto& mask_img = cc->Inputs().Tag(kMaskCpuTag).Get<ImageFrame>();
|
||||
|
||||
cv::Mat input_mat = formats::MatView(&input_img);
|
||||
cv::Mat mask_mat = formats::MatView(&mask_img);
|
||||
|
||||
RET_CHECK(input_mat.channels() == 3); // RGB only.
|
||||
|
||||
if (mask_mat.channels() > 1) {
|
||||
std::vector<cv::Mat> channels;
|
||||
cv::split(mask_mat, channels);
|
||||
if (mask_channel_ == mediapipe::RecolorCalculatorOptions_MaskChannel_ALPHA)
|
||||
mask_mat = channels[3];
|
||||
else
|
||||
mask_mat = channels[0];
|
||||
}
|
||||
cv::Mat mask_full;
|
||||
cv::resize(mask_mat, mask_full, input_mat.size());
|
||||
|
||||
auto output_img = absl::make_unique<ImageFrame>(
|
||||
input_img.Format(), input_mat.cols, input_mat.rows);
|
||||
cv::Mat output_mat = mediapipe::formats::MatView(output_img.get());
|
||||
|
||||
// From GPU shader:
|
||||
/*
|
||||
vec4 weight = texture2D(mask, sample_coordinate);
|
||||
vec4 color1 = texture2D(frame, sample_coordinate);
|
||||
vec4 color2 = vec4(recolor, 1.0);
|
||||
|
||||
float luminance = dot(color1.rgb, vec3(0.299, 0.587, 0.114));
|
||||
float mix_value = weight.MASK_COMPONENT * luminance;
|
||||
|
||||
fragColor = mix(color1, color2, mix_value);
|
||||
*/
|
||||
for (int i = 0; i < output_mat.rows; ++i) {
|
||||
for (int j = 0; j < output_mat.cols; ++j) {
|
||||
float weight = mask_full.at<uchar>(i, j) * (1.0 / 255.0);
|
||||
cv::Vec3f color1 = input_mat.at<cv::Vec3b>(i, j);
|
||||
cv::Vec3f color2 = {color_[0], color_[1], color_[2]};
|
||||
|
||||
float luminance =
|
||||
(color1[0] * 0.299 + color1[1] * 0.587 + color1[2] * 0.114) / 255;
|
||||
float mix_value = weight * luminance;
|
||||
|
||||
cv::Vec3b mix_color = color1 * (1.0 - mix_value) + color2 * mix_value;
|
||||
output_mat.at<cv::Vec3b>(i, j) = mix_color;
|
||||
}
|
||||
}
|
||||
|
||||
cc->Outputs()
|
||||
.Tag(kImageFrameTag)
|
||||
.Add(output_img.release(), cc->InputTimestamp());
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status RecolorCalculator::RenderGpu(CalculatorContext* cc) {
|
||||
if (cc->Inputs().Tag("MASK_GPU").IsEmpty()) {
|
||||
if (cc->Inputs().Tag(kMaskGpuTag).IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||
// Get inputs and setup output.
|
||||
const Packet& input_packet = cc->Inputs().Tag("IMAGE_GPU").Value();
|
||||
const Packet& mask_packet = cc->Inputs().Tag("MASK_GPU").Value();
|
||||
const Packet& input_packet = cc->Inputs().Tag(kGpuBufferTag).Value();
|
||||
const Packet& mask_packet = cc->Inputs().Tag(kMaskGpuTag).Value();
|
||||
|
||||
const auto& input_buffer = input_packet.Get<mediapipe::GpuBuffer>();
|
||||
const auto& mask_buffer = mask_packet.Get<mediapipe::GpuBuffer>();
|
||||
@@ -233,7 +305,7 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
||||
|
||||
// Send result image in GPU packet.
|
||||
auto output = dst_tex.GetFrame<mediapipe::GpuBuffer>();
|
||||
cc->Outputs().Tag("IMAGE_GPU").Add(output.release(), cc->InputTimestamp());
|
||||
cc->Outputs().Tag(kGpuBufferTag).Add(output.release(), cc->InputTimestamp());
|
||||
|
||||
// Cleanup
|
||||
img_tex.Release();
|
||||
@@ -303,9 +375,9 @@ void RecolorCalculator::GlRender() {
|
||||
|
||||
if (!options.has_color()) RET_CHECK_FAIL() << "Missing color option.";
|
||||
|
||||
color_.push_back(options.color().r() / 255.0);
|
||||
color_.push_back(options.color().g() / 255.0);
|
||||
color_.push_back(options.color().b() / 255.0);
|
||||
color_.push_back(options.color().r());
|
||||
color_.push_back(options.color().g());
|
||||
color_.push_back(options.color().b());
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -378,8 +450,8 @@ void RecolorCalculator::GlRender() {
|
||||
glUseProgram(program_);
|
||||
glUniform1i(glGetUniformLocation(program_, "frame"), 1);
|
||||
glUniform1i(glGetUniformLocation(program_, "mask"), 2);
|
||||
glUniform3f(glGetUniformLocation(program_, "recolor"), color_[0], color_[1],
|
||||
color_[2]);
|
||||
glUniform3f(glGetUniformLocation(program_, "recolor"), color_[0] / 255.0,
|
||||
color_[1] / 255.0, color_[2] / 255.0);
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
|
||||
@@ -1110,6 +1110,7 @@ cc_test(
|
||||
],
|
||||
"//mediapipe:android": [
|
||||
"@org_tensorflow//tensorflow/core:android_tensorflow_lib_with_ops_lite_proto_no_rtti_lib",
|
||||
"@org_tensorflow//tensorflow/core:android_tensorflow_test_lib",
|
||||
],
|
||||
"//mediapipe:ios": [
|
||||
"@org_tensorflow//tensorflow/core:ios_tensorflow_test_lib",
|
||||
|
||||
@@ -244,7 +244,7 @@ class PackMediaSequenceCalculator : public CalculatorBase {
|
||||
::mediapipe::Status VerifySequence() {
|
||||
std::string error_msg = "Missing features - ";
|
||||
bool all_present = true;
|
||||
for (auto iter : features_present_) {
|
||||
for (const auto& iter : features_present_) {
|
||||
if (!iter.second) {
|
||||
all_present = false;
|
||||
absl::StrAppend(&error_msg, iter.first, ", ");
|
||||
|
||||
@@ -51,10 +51,11 @@ static constexpr char kStringSavedModelPath[] = "STRING_SAVED_MODEL_PATH";
|
||||
#endif
|
||||
}
|
||||
|
||||
// If options.convert_signature_to_tags() will convert letters to uppercase
|
||||
// and replace /'s with _'s. If set, this enables the standard SavedModel
|
||||
// classification, regression, and prediction signatures to be used as
|
||||
// uppercase INPUTS and OUTPUTS tags for streams.
|
||||
// If options.convert_signature_to_tags() is set, will convert letters to
|
||||
// uppercase and replace /'s and -'s with _'s. This enables the standard
|
||||
// SavedModel classification, regression, and prediction signatures to be used
|
||||
// as uppercase INPUTS and OUTPUTS tags for streams and supports other common
|
||||
// patterns.
|
||||
const std::string MaybeConvertSignatureToTag(
|
||||
const std::string& name,
|
||||
const TensorFlowSessionFromSavedModelCalculatorOptions& options) {
|
||||
@@ -64,6 +65,7 @@ const std::string MaybeConvertSignatureToTag(
|
||||
std::transform(name.begin(), name.end(), output.begin(),
|
||||
[](unsigned char c) { return std::toupper(c); });
|
||||
output = absl::StrReplaceAll(output, {{"/", "_"}});
|
||||
output = absl::StrReplaceAll(output, {{"-", "_"}});
|
||||
return output;
|
||||
} else {
|
||||
return name;
|
||||
@@ -124,7 +126,7 @@ class TensorFlowSessionFromSavedModelCalculator : public CalculatorBase {
|
||||
// Set user specified tags properly.
|
||||
// If no tags specified will use tensorflow::kSavedModelTagServe by default.
|
||||
std::unordered_set<std::string> tags_set;
|
||||
for (std::string tag : options.saved_model_tag()) {
|
||||
for (const std::string& tag : options.saved_model_tag()) {
|
||||
tags_set.insert(tag);
|
||||
}
|
||||
if (tags_set.empty()) {
|
||||
|
||||
@@ -32,8 +32,8 @@ message TensorFlowSessionFromSavedModelCalculatorOptions {
|
||||
// The name of the generic signature to load into the mapping from tags to
|
||||
// tensor names.
|
||||
optional string signature_name = 2 [default = "serving_default"];
|
||||
// Whether to convert the signature keys to uppercase and switch /'s to
|
||||
// _'s, which enables standard signatures to be used as Tags.
|
||||
// Whether to convert the signature keys to uppercase as well as switch /'s
|
||||
// and -'s to _'s, which enables common signatures to be used as Tags.
|
||||
optional bool convert_signature_to_tags = 3 [default = true];
|
||||
// If true, saved_model_path can have multiple exported models in
|
||||
// subdirectories saved_model_path/%08d and the alphabetically last (i.e.,
|
||||
|
||||
@@ -53,10 +53,11 @@ static constexpr char kStringSavedModelPath[] = "STRING_SAVED_MODEL_PATH";
|
||||
#endif
|
||||
}
|
||||
|
||||
// If options.convert_signature_to_tags() will convert letters to uppercase
|
||||
// and replace /'s with _'s. If set, this enables the standard SavedModel
|
||||
// classification, regression, and prediction signatures to be used as
|
||||
// uppercase INPUTS and OUTPUTS tags for streams.
|
||||
// If options.convert_signature_to_tags() is set, will convert letters to
|
||||
// uppercase and replace /'s and -'s with _'s. This enables the standard
|
||||
// SavedModel classification, regression, and prediction signatures to be used
|
||||
// as uppercase INPUTS and OUTPUTS tags for streams and supports other common
|
||||
// patterns.
|
||||
const std::string MaybeConvertSignatureToTag(
|
||||
const std::string& name,
|
||||
const TensorFlowSessionFromSavedModelGeneratorOptions& options) {
|
||||
@@ -66,6 +67,7 @@ const std::string MaybeConvertSignatureToTag(
|
||||
std::transform(name.begin(), name.end(), output.begin(),
|
||||
[](unsigned char c) { return std::toupper(c); });
|
||||
output = absl::StrReplaceAll(output, {{"/", "_"}});
|
||||
output = absl::StrReplaceAll(output, {{"-", "_"}});
|
||||
return output;
|
||||
} else {
|
||||
return name;
|
||||
@@ -119,7 +121,7 @@ class TensorFlowSessionFromSavedModelGenerator : public PacketGenerator {
|
||||
// Set user specified tags properly.
|
||||
// If no tags specified will use tensorflow::kSavedModelTagServe by default.
|
||||
std::unordered_set<std::string> tags_set;
|
||||
for (std::string tag : options.saved_model_tag()) {
|
||||
for (const std::string& tag : options.saved_model_tag()) {
|
||||
tags_set.insert(tag);
|
||||
}
|
||||
if (tags_set.empty()) {
|
||||
|
||||
@@ -32,8 +32,8 @@ message TensorFlowSessionFromSavedModelGeneratorOptions {
|
||||
// The name of the generic signature to load into the mapping from tags to
|
||||
// tensor names.
|
||||
optional string signature_name = 2 [default = "serving_default"];
|
||||
// Whether to convert the signature keys to uppercase and switch /'s to
|
||||
// _'s, which enables standard signatures to be used as Tags.
|
||||
// Whether to convert the signature keys to uppercase as well as switch /'s
|
||||
// and -'s to _'s, which enables common signatures to be used as Tags.
|
||||
optional bool convert_signature_to_tags = 3 [default = true];
|
||||
// If true, saved_model_path can have multiple exported models in
|
||||
// subdirectories saved_model_path/%08d and the alphabetically last (i.e.,
|
||||
|
||||
@@ -222,9 +222,11 @@ cc_library(
|
||||
deps = [
|
||||
":util",
|
||||
":tflite_inference_calculator_cc_proto",
|
||||
"@com_google_absl//absl/memory",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/util:resource_util",
|
||||
"@org_tensorflow//tensorflow/lite:framework",
|
||||
"@org_tensorflow//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
|
||||
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
||||
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
@@ -241,6 +243,7 @@ cc_library(
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate_internal",
|
||||
],
|
||||
"//conditions:default": [
|
||||
"//mediapipe/util/tflite:tflite_gpu_runner",
|
||||
"//mediapipe/gpu:gl_calculator_helper",
|
||||
"//mediapipe/gpu:gpu_buffer",
|
||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
|
||||
@@ -254,6 +257,10 @@ cc_library(
|
||||
"//mediapipe:android": [
|
||||
"@org_tensorflow//tensorflow/lite/delegates/nnapi:nnapi_delegate",
|
||||
],
|
||||
}) + select({
|
||||
"//conditions:default": [
|
||||
"//mediapipe/util:cpu_util",
|
||||
],
|
||||
}),
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -308,6 +315,20 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "tflite_model_calculator",
|
||||
srcs = ["tflite_model_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":util",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"@org_tensorflow//tensorflow/lite:framework",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "tflite_tensors_to_segmentation_calculator",
|
||||
srcs = ["tflite_tensors_to_segmentation_calculator.cc"],
|
||||
@@ -431,7 +452,7 @@ cc_library(
|
||||
"//mediapipe:android": [
|
||||
"//mediapipe/util/android/file/base",
|
||||
],
|
||||
"//mediapipe:apple": [
|
||||
"//mediapipe:ios": [
|
||||
"//mediapipe/util/android/file/base",
|
||||
],
|
||||
"//mediapipe:macos": [
|
||||
@@ -478,6 +499,9 @@ cc_test(
|
||||
deps = [
|
||||
":tflite_inference_calculator",
|
||||
":tflite_inference_calculator_cc_proto",
|
||||
":tflite_model_calculator",
|
||||
"//mediapipe/calculators/core:constant_side_packet_calculator",
|
||||
"//mediapipe/calculators/util:local_file_contents_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/deps:file_path",
|
||||
@@ -485,7 +509,9 @@ cc_test(
|
||||
"//mediapipe/framework/port:integral_types",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/tool:validate_type",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/types:optional",
|
||||
"@org_tensorflow//tensorflow/lite:framework",
|
||||
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
||||
],
|
||||
@@ -511,3 +537,19 @@ cc_test(
|
||||
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "tflite_model_calculator_test",
|
||||
srcs = ["tflite_model_calculator_test.cc"],
|
||||
data = ["testdata/add.bin"],
|
||||
deps = [
|
||||
":tflite_model_calculator",
|
||||
"//mediapipe/calculators/core:constant_side_packet_calculator",
|
||||
"//mediapipe/calculators/util:local_file_contents_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"@org_tensorflow//tensorflow/lite:framework",
|
||||
],
|
||||
)
|
||||
|
||||
@@ -63,6 +63,10 @@ typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>
|
||||
typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>
|
||||
ColMajorMatrixXf;
|
||||
|
||||
constexpr char kImageFrameTag[] = "IMAGE";
|
||||
constexpr char kGpuBufferTag[] = "IMAGE_GPU";
|
||||
constexpr char kTensorsTag[] = "TENSORS";
|
||||
constexpr char kTensorsGpuTag[] = "TENSORS_GPU";
|
||||
} // namespace
|
||||
|
||||
namespace mediapipe {
|
||||
@@ -124,6 +128,9 @@ struct GPUData {
|
||||
// GPU tensors are currently only supported on mobile platforms.
|
||||
// This calculator uses FixedSizeInputStreamHandler by default.
|
||||
//
|
||||
// Note: Input defines output, so only these type sets are supported:
|
||||
// IMAGE -> TENSORS | IMAGE_GPU -> TENSORS_GPU | MATRIX -> TENSORS
|
||||
//
|
||||
class TfLiteConverterCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
@@ -138,9 +145,9 @@ class TfLiteConverterCalculator : public CalculatorBase {
|
||||
template <class T>
|
||||
::mediapipe::Status NormalizeImage(const ImageFrame& image_frame,
|
||||
bool zero_center, bool flip_vertically,
|
||||
float* tensor_buffer);
|
||||
float* tensor_ptr);
|
||||
::mediapipe::Status CopyMatrixToTensor(const Matrix& matrix,
|
||||
float* tensor_buffer);
|
||||
float* tensor_ptr);
|
||||
::mediapipe::Status ProcessCPU(CalculatorContext* cc);
|
||||
::mediapipe::Status ProcessGPU(CalculatorContext* cc);
|
||||
|
||||
@@ -166,33 +173,35 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteConverterCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
const bool has_image_tag = cc->Inputs().HasTag("IMAGE");
|
||||
const bool has_image_gpu_tag = cc->Inputs().HasTag("IMAGE_GPU");
|
||||
const bool has_matrix_tag = cc->Inputs().HasTag("MATRIX");
|
||||
// Confirm only one of the input streams is present.
|
||||
RET_CHECK(has_image_tag ^ has_image_gpu_tag ^ has_matrix_tag &&
|
||||
!(has_image_tag && has_image_gpu_tag && has_matrix_tag));
|
||||
RET_CHECK(cc->Inputs().HasTag(kImageFrameTag) ^
|
||||
cc->Inputs().HasTag(kGpuBufferTag) ^ cc->Inputs().HasTag("MATRIX"));
|
||||
|
||||
// Confirm only one of the output streams is present.
|
||||
RET_CHECK(cc->Outputs().HasTag("TENSORS") ^
|
||||
cc->Outputs().HasTag("TENSORS_GPU"));
|
||||
RET_CHECK(cc->Outputs().HasTag(kTensorsTag) ^
|
||||
cc->Outputs().HasTag(kTensorsGpuTag));
|
||||
|
||||
bool use_gpu = false;
|
||||
|
||||
if (cc->Inputs().HasTag("IMAGE")) cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
||||
if (cc->Inputs().HasTag("MATRIX")) cc->Inputs().Tag("MATRIX").Set<Matrix>();
|
||||
if (cc->Inputs().HasTag(kImageFrameTag)) {
|
||||
cc->Inputs().Tag(kImageFrameTag).Set<ImageFrame>();
|
||||
}
|
||||
if (cc->Inputs().HasTag("MATRIX")) {
|
||||
cc->Inputs().Tag("MATRIX").Set<Matrix>();
|
||||
}
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||
cc->Inputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
|
||||
if (cc->Inputs().HasTag(kGpuBufferTag)) {
|
||||
cc->Inputs().Tag(kGpuBufferTag).Set<mediapipe::GpuBuffer>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
if (cc->Outputs().HasTag("TENSORS"))
|
||||
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||
if (cc->Outputs().HasTag(kTensorsTag)) {
|
||||
cc->Outputs().Tag(kTensorsTag).Set<std::vector<TfLiteTensor>>();
|
||||
}
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
if (cc->Outputs().HasTag("TENSORS_GPU")) {
|
||||
cc->Outputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||
if (cc->Outputs().HasTag(kTensorsGpuTag)) {
|
||||
cc->Outputs().Tag(kTensorsGpuTag).Set<std::vector<GpuTensor>>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
@@ -216,8 +225,8 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
|
||||
MP_RETURN_IF_ERROR(LoadOptions(cc));
|
||||
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU") ||
|
||||
cc->Outputs().HasTag("IMAGE_OUT_GPU")) {
|
||||
if (cc->Inputs().HasTag(kGpuBufferTag) ||
|
||||
cc->Outputs().HasTag(kGpuBufferTag)) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
use_gpu_ = true;
|
||||
#else
|
||||
@@ -227,8 +236,8 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
|
||||
if (use_gpu_) {
|
||||
// Cannot mix CPU/GPU streams.
|
||||
RET_CHECK(cc->Inputs().HasTag("IMAGE_GPU") &&
|
||||
cc->Outputs().HasTag("TENSORS_GPU"));
|
||||
RET_CHECK(cc->Inputs().HasTag(kGpuBufferTag) &&
|
||||
cc->Outputs().HasTag(kTensorsGpuTag));
|
||||
// Cannot use quantization.
|
||||
use_quantized_tensors_ = false;
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
@@ -248,7 +257,6 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteConverterCalculator::Process(CalculatorContext* cc) {
|
||||
if (use_gpu_) {
|
||||
// GpuBuffer to tflite::gpu::GlBuffer conversion.
|
||||
if (!initialized_) {
|
||||
MP_RETURN_IF_ERROR(InitGpu(cc));
|
||||
initialized_ = true;
|
||||
@@ -259,7 +267,6 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
// Convert to CPU tensors or Matrix type.
|
||||
MP_RETURN_IF_ERROR(ProcessCPU(cc));
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -275,24 +282,26 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteConverterCalculator::ProcessCPU(
|
||||
CalculatorContext* cc) {
|
||||
if (cc->Inputs().HasTag("IMAGE")) {
|
||||
if (cc->Inputs().HasTag(kImageFrameTag)) {
|
||||
// CPU ImageFrame to TfLiteTensor conversion.
|
||||
|
||||
const auto& image_frame = cc->Inputs().Tag("IMAGE").Get<ImageFrame>();
|
||||
const auto& image_frame =
|
||||
cc->Inputs().Tag(kImageFrameTag).Get<ImageFrame>();
|
||||
const int height = image_frame.Height();
|
||||
const int width = image_frame.Width();
|
||||
const int channels = image_frame.NumberOfChannels();
|
||||
const int channels_preserved = std::min(channels, max_num_channels_);
|
||||
const mediapipe::ImageFormat::Format format = image_frame.Format();
|
||||
|
||||
if (!initialized_) {
|
||||
if (!(image_frame.Format() == mediapipe::ImageFormat::SRGBA ||
|
||||
image_frame.Format() == mediapipe::ImageFormat::SRGB ||
|
||||
image_frame.Format() == mediapipe::ImageFormat::GRAY8 ||
|
||||
image_frame.Format() == mediapipe::ImageFormat::VEC32F1))
|
||||
if (!(format == mediapipe::ImageFormat::SRGBA ||
|
||||
format == mediapipe::ImageFormat::SRGB ||
|
||||
format == mediapipe::ImageFormat::GRAY8 ||
|
||||
format == mediapipe::ImageFormat::VEC32F1))
|
||||
RET_CHECK_FAIL() << "Unsupported CPU input format.";
|
||||
TfLiteQuantization quant;
|
||||
if (use_quantized_tensors_) {
|
||||
RET_CHECK(image_frame.Format() != mediapipe::ImageFormat::VEC32F1)
|
||||
RET_CHECK(format != mediapipe::ImageFormat::VEC32F1)
|
||||
<< "Only 8-bit input images are supported for quantization.";
|
||||
quant.type = kTfLiteAffineQuantization;
|
||||
quant.params = nullptr;
|
||||
@@ -349,8 +358,9 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
|
||||
auto output_tensors = absl::make_unique<std::vector<TfLiteTensor>>();
|
||||
output_tensors->emplace_back(*tensor);
|
||||
cc->Outputs().Tag("TENSORS").Add(output_tensors.release(),
|
||||
cc->InputTimestamp());
|
||||
cc->Outputs()
|
||||
.Tag(kTensorsTag)
|
||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||
} else if (cc->Inputs().HasTag("MATRIX")) {
|
||||
// CPU Matrix to TfLiteTensor conversion.
|
||||
|
||||
@@ -371,15 +381,16 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
interpreter_->ResizeInputTensor(tensor_idx, {height, width, channels});
|
||||
interpreter_->AllocateTensors();
|
||||
|
||||
float* tensor_buffer = tensor->data.f;
|
||||
RET_CHECK(tensor_buffer);
|
||||
float* tensor_ptr = tensor->data.f;
|
||||
RET_CHECK(tensor_ptr);
|
||||
|
||||
MP_RETURN_IF_ERROR(CopyMatrixToTensor(matrix, tensor_buffer));
|
||||
MP_RETURN_IF_ERROR(CopyMatrixToTensor(matrix, tensor_ptr));
|
||||
|
||||
auto output_tensors = absl::make_unique<std::vector<TfLiteTensor>>();
|
||||
output_tensors->emplace_back(*tensor);
|
||||
cc->Outputs().Tag("TENSORS").Add(output_tensors.release(),
|
||||
cc->InputTimestamp());
|
||||
cc->Outputs()
|
||||
.Tag(kTensorsTag)
|
||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
@@ -389,7 +400,8 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
CalculatorContext* cc) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
// GpuBuffer to tflite::gpu::GlBuffer conversion.
|
||||
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
||||
const auto& input =
|
||||
cc->Inputs().Tag(kGpuBufferTag).Get<mediapipe::GpuBuffer>();
|
||||
MP_RETURN_IF_ERROR(
|
||||
gpu_helper_.RunInGlContext([this, &input]() -> ::mediapipe::Status {
|
||||
// Convert GL texture into TfLite GlBuffer (SSBO).
|
||||
@@ -421,11 +433,12 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
cc->Outputs()
|
||||
.Tag("TENSORS_GPU")
|
||||
.Tag(kTensorsGpuTag)
|
||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
// GpuBuffer to id<MTLBuffer> conversion.
|
||||
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
||||
const auto& input =
|
||||
cc->Inputs().Tag(kGpuBufferTag).Get<mediapipe::GpuBuffer>();
|
||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||
|
||||
id<MTLTexture> src_texture = [gpu_helper_ metalTextureWithGpuBuffer:input];
|
||||
@@ -457,7 +470,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
commandBuffer:command_buffer];
|
||||
|
||||
cc->Outputs()
|
||||
.Tag("TENSORS_GPU")
|
||||
.Tag(kTensorsGpuTag)
|
||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||
#else
|
||||
RET_CHECK_FAIL() << "GPU processing is not enabled.";
|
||||
@@ -469,7 +482,8 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
::mediapipe::Status TfLiteConverterCalculator::InitGpu(CalculatorContext* cc) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
// Get input image sizes.
|
||||
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
||||
const auto& input =
|
||||
cc->Inputs().Tag(kGpuBufferTag).Get<mediapipe::GpuBuffer>();
|
||||
mediapipe::ImageFormat::Format format =
|
||||
mediapipe::ImageFormatForGpuBufferFormat(input.format());
|
||||
gpu_data_out_ = absl::make_unique<GPUData>();
|
||||
@@ -612,7 +626,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
CHECK_LE(max_num_channels_, 4);
|
||||
CHECK_NE(max_num_channels_, 2);
|
||||
#if defined(MEDIAPIPE_IOS)
|
||||
if (cc->Inputs().HasTag("IMAGE_GPU"))
|
||||
if (cc->Inputs().HasTag(kGpuBufferTag))
|
||||
// Currently on iOS, tflite gpu input tensor must be 4 channels,
|
||||
// so input image must be 4 channels also (checked in InitGpu).
|
||||
max_num_channels_ = 4;
|
||||
@@ -627,7 +641,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
||||
template <class T>
|
||||
::mediapipe::Status TfLiteConverterCalculator::NormalizeImage(
|
||||
const ImageFrame& image_frame, bool zero_center, bool flip_vertically,
|
||||
float* tensor_buffer) {
|
||||
float* tensor_ptr) {
|
||||
const int height = image_frame.Height();
|
||||
const int width = image_frame.Width();
|
||||
const int channels = image_frame.NumberOfChannels();
|
||||
@@ -651,7 +665,7 @@ template <class T>
|
||||
(flip_vertically ? height - 1 - i : i) * image_frame.WidthStep());
|
||||
for (int j = 0; j < width; ++j) {
|
||||
for (int c = 0; c < channels_preserved; ++c) {
|
||||
*tensor_buffer++ = *image_ptr++ / div - sub;
|
||||
*tensor_ptr++ = *image_ptr++ / div - sub;
|
||||
}
|
||||
image_ptr += channels_ignored;
|
||||
}
|
||||
@@ -661,14 +675,14 @@ template <class T>
|
||||
}
|
||||
|
||||
::mediapipe::Status TfLiteConverterCalculator::CopyMatrixToTensor(
|
||||
const Matrix& matrix, float* tensor_buffer) {
|
||||
const Matrix& matrix, float* tensor_ptr) {
|
||||
if (row_major_matrix_) {
|
||||
auto matrix_map = Eigen::Map<RowMajorMatrixXf>(tensor_buffer, matrix.rows(),
|
||||
matrix.cols());
|
||||
auto matrix_map =
|
||||
Eigen::Map<RowMajorMatrixXf>(tensor_ptr, matrix.rows(), matrix.cols());
|
||||
matrix_map = matrix;
|
||||
} else {
|
||||
auto matrix_map = Eigen::Map<ColMajorMatrixXf>(tensor_buffer, matrix.rows(),
|
||||
matrix.cols());
|
||||
auto matrix_map =
|
||||
Eigen::Map<ColMajorMatrixXf>(tensor_ptr, matrix.rows(), matrix.cols());
|
||||
matrix_map = matrix;
|
||||
}
|
||||
|
||||
|
||||
@@ -17,10 +17,16 @@
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "mediapipe/calculators/tflite/tflite_inference_calculator.pb.h"
|
||||
#include "mediapipe/calculators/tflite/util.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
|
||||
#if !defined(__EMSCRIPTEN__) || defined(__EMSCRIPTEN_PTHREADS__)
|
||||
#include "mediapipe/util/cpu_util.h"
|
||||
#endif // !__EMSCRIPTEN__ || __EMSCRIPTEN_PTHREADS__
|
||||
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
#include "tensorflow/lite/error_reporter.h"
|
||||
#include "tensorflow/lite/interpreter.h"
|
||||
@@ -30,6 +36,7 @@
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||
#include "mediapipe/gpu/gpu_buffer.h"
|
||||
#include "mediapipe/util/tflite/tflite_gpu_runner.h"
|
||||
#include "tensorflow/lite/delegates/gpu/common/shape.h"
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
||||
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
|
||||
@@ -50,7 +57,7 @@
|
||||
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
|
||||
#include "tensorflow/lite/delegates/gpu/metal_delegate_internal.h"
|
||||
#endif // iOS
|
||||
|
||||
#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
|
||||
#if defined(MEDIAPIPE_ANDROID)
|
||||
#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
|
||||
#endif // ANDROID
|
||||
@@ -69,6 +76,9 @@ typedef id<MTLBuffer> GpuTensor;
|
||||
|
||||
// Round up n to next multiple of m.
|
||||
size_t RoundUp(size_t n, size_t m) { return ((n + m - 1) / m) * m; } // NOLINT
|
||||
|
||||
constexpr char kTensorsTag[] = "TENSORS";
|
||||
constexpr char kTensorsGpuTag[] = "TENSORS_GPU";
|
||||
} // namespace
|
||||
|
||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||
@@ -113,6 +123,23 @@ struct GPUData {
|
||||
};
|
||||
#endif
|
||||
|
||||
// Returns number of threads to configure XNNPACK delegate with.
|
||||
// (Equal to user provided value if specified. Otherwise, it returns number of
|
||||
// high cores (hard-coded to 1 for Emscripten without Threads extension))
|
||||
int GetXnnpackNumThreads(
|
||||
const mediapipe::TfLiteInferenceCalculatorOptions& opts) {
|
||||
static constexpr int kDefaultNumThreads = -1;
|
||||
if (opts.has_delegate() && opts.delegate().has_xnnpack() &&
|
||||
opts.delegate().xnnpack().num_threads() != kDefaultNumThreads) {
|
||||
return opts.delegate().xnnpack().num_threads();
|
||||
}
|
||||
#if !defined(__EMSCRIPTEN__) || defined(__EMSCRIPTEN_PTHREADS__)
|
||||
return InferHigherCoreIds().size();
|
||||
#else
|
||||
return 1;
|
||||
#endif // !__EMSCRIPTEN__ || __EMSCRIPTEN_PTHREADS__
|
||||
}
|
||||
|
||||
// Calculator Header Section
|
||||
|
||||
// Runs inference on the provided input TFLite tensors and TFLite model.
|
||||
@@ -139,6 +166,9 @@ struct GPUData {
|
||||
// Input side packet:
|
||||
// CUSTOM_OP_RESOLVER (optional) - Use a custom op resolver,
|
||||
// instead of the builtin one.
|
||||
// MODEL (optional) - Use to specify TfLite model
|
||||
// (std::unique_ptr<tflite::FlatBufferModel,
|
||||
// std::function<void(tflite::FlatBufferModel*)>>)
|
||||
//
|
||||
// Example use:
|
||||
// node {
|
||||
@@ -153,6 +183,20 @@ struct GPUData {
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// or
|
||||
//
|
||||
// node {
|
||||
// calculator: "TfLiteInferenceCalculator"
|
||||
// input_stream: "TENSORS:tensor_image"
|
||||
// input_side_packet: "MODEL:model"
|
||||
// output_stream: "TENSORS:tensors"
|
||||
// options: {
|
||||
// [mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
||||
// delegate { gpu {} }
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// IMPORTANT Notes:
|
||||
// Tensors are assumed to be ordered correctly (sequentially added to model).
|
||||
// Input tensors are assumed to be of the correct size and already normalized.
|
||||
@@ -165,6 +209,9 @@ class TfLiteInferenceCalculator : public CalculatorBase {
|
||||
public:
|
||||
using TfLiteDelegatePtr =
|
||||
std::unique_ptr<TfLiteDelegate, std::function<void(TfLiteDelegate*)>>;
|
||||
using TfLiteModelPtr =
|
||||
std::unique_ptr<tflite::FlatBufferModel,
|
||||
std::function<void(tflite::FlatBufferModel*)>>;
|
||||
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
@@ -173,18 +220,20 @@ class TfLiteInferenceCalculator : public CalculatorBase {
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
::mediapipe::Status LoadOptions(CalculatorContext* cc);
|
||||
::mediapipe::Status LoadModel(CalculatorContext* cc);
|
||||
::mediapipe::StatusOr<Packet> GetModelAsPacket(const CalculatorContext& cc);
|
||||
::mediapipe::Status LoadDelegate(CalculatorContext* cc);
|
||||
::mediapipe::Status InitTFLiteGPURunner();
|
||||
|
||||
Packet model_packet_;
|
||||
std::unique_ptr<tflite::Interpreter> interpreter_;
|
||||
std::unique_ptr<tflite::FlatBufferModel> model_;
|
||||
TfLiteDelegatePtr delegate_;
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||
std::vector<std::unique_ptr<GPUData>> gpu_data_in_;
|
||||
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
|
||||
std::unique_ptr<tflite::gpu::TFLiteGPURunner> tflite_gpu_runner_;
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
MPPMetalHelper* gpu_helper_ = nullptr;
|
||||
std::vector<std::unique_ptr<GPUData>> gpu_data_in_;
|
||||
@@ -198,11 +247,12 @@ class TfLiteInferenceCalculator : public CalculatorBase {
|
||||
edgetpu::EdgeTpuManager::GetSingleton()->OpenDevice();
|
||||
#endif
|
||||
|
||||
std::string model_path_ = "";
|
||||
bool gpu_inference_ = false;
|
||||
bool gpu_input_ = false;
|
||||
bool gpu_output_ = false;
|
||||
bool use_quantized_tensors_ = false;
|
||||
|
||||
bool use_advanced_gpu_api_ = false;
|
||||
};
|
||||
REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
@@ -210,36 +260,40 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteInferenceCalculator::GetContract(
|
||||
CalculatorContract* cc) {
|
||||
RET_CHECK(cc->Inputs().HasTag("TENSORS") ^
|
||||
cc->Inputs().HasTag("TENSORS_GPU"));
|
||||
RET_CHECK(cc->Outputs().HasTag("TENSORS") ^
|
||||
cc->Outputs().HasTag("TENSORS_GPU"));
|
||||
RET_CHECK(cc->Inputs().HasTag(kTensorsTag) ^
|
||||
cc->Inputs().HasTag(kTensorsGpuTag));
|
||||
RET_CHECK(cc->Outputs().HasTag(kTensorsTag) ^
|
||||
cc->Outputs().HasTag(kTensorsGpuTag));
|
||||
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::TfLiteInferenceCalculatorOptions>();
|
||||
RET_CHECK(!options.model_path().empty() ^
|
||||
cc->InputSidePackets().HasTag("MODEL"))
|
||||
<< "Either model as side packet or model path in options is required.";
|
||||
|
||||
bool use_gpu =
|
||||
options.has_delegate() ? options.delegate().has_gpu() : options.use_gpu();
|
||||
|
||||
if (cc->Inputs().HasTag("TENSORS"))
|
||||
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||
if (cc->Inputs().HasTag(kTensorsTag))
|
||||
cc->Inputs().Tag(kTensorsTag).Set<std::vector<TfLiteTensor>>();
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
if (cc->Inputs().HasTag(kTensorsGpuTag)) {
|
||||
RET_CHECK(!options.has_delegate() || options.delegate().has_gpu())
|
||||
<< "GPU input is compatible with GPU delegate only.";
|
||||
|
||||
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||
cc->Inputs().Tag(kTensorsGpuTag).Set<std::vector<GpuTensor>>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
if (cc->Outputs().HasTag("TENSORS"))
|
||||
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||
if (cc->Outputs().HasTag(kTensorsTag))
|
||||
cc->Outputs().Tag(kTensorsTag).Set<std::vector<TfLiteTensor>>();
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
if (cc->Outputs().HasTag("TENSORS_GPU")) {
|
||||
if (cc->Outputs().HasTag(kTensorsGpuTag)) {
|
||||
RET_CHECK(!options.has_delegate() || options.delegate().has_gpu())
|
||||
<< "GPU output is compatible with GPU delegate only.";
|
||||
|
||||
cc->Outputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||
cc->Outputs().Tag(kTensorsGpuTag).Set<std::vector<GpuTensor>>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
@@ -249,6 +303,9 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
.Tag("CUSTOM_OP_RESOLVER")
|
||||
.Set<tflite::ops::builtin::BuiltinOpResolver>();
|
||||
}
|
||||
if (cc->InputSidePackets().HasTag("MODEL")) {
|
||||
cc->InputSidePackets().Tag("MODEL").Set<TfLiteModelPtr>();
|
||||
}
|
||||
|
||||
if (use_gpu) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
@@ -267,29 +324,35 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
::mediapipe::Status TfLiteInferenceCalculator::Open(CalculatorContext* cc) {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
MP_RETURN_IF_ERROR(LoadOptions(cc));
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::TfLiteInferenceCalculatorOptions>();
|
||||
gpu_inference_ = options.use_gpu();
|
||||
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
if (cc->Inputs().HasTag(kTensorsGpuTag)) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
gpu_input_ = true;
|
||||
gpu_inference_ = true; // Inference must be on GPU also.
|
||||
#else
|
||||
RET_CHECK(!cc->Inputs().HasTag("TENSORS_GPU"))
|
||||
RET_CHECK(!cc->Inputs().HasTag(kTensorsGpuTag))
|
||||
<< "GPU processing not enabled.";
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("TENSORS_GPU")) {
|
||||
if (cc->Outputs().HasTag(kTensorsGpuTag)) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
gpu_output_ = true;
|
||||
RET_CHECK(cc->Inputs().HasTag("TENSORS_GPU"))
|
||||
RET_CHECK(cc->Inputs().HasTag(kTensorsGpuTag))
|
||||
<< "GPU output must also have GPU Input.";
|
||||
#else
|
||||
RET_CHECK(!cc->Inputs().HasTag("TENSORS_GPU"))
|
||||
RET_CHECK(!cc->Inputs().HasTag(kTensorsGpuTag))
|
||||
<< "GPU processing not enabled.";
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
}
|
||||
|
||||
const auto& calculator_opts =
|
||||
cc->Options<mediapipe::TfLiteInferenceCalculatorOptions>();
|
||||
use_advanced_gpu_api_ = false;
|
||||
|
||||
MP_RETURN_IF_ERROR(LoadModel(cc));
|
||||
|
||||
if (gpu_inference_) {
|
||||
@@ -301,8 +364,12 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
#endif
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||
[this, &cc]() -> ::mediapipe::Status { return LoadDelegate(cc); }));
|
||||
MP_RETURN_IF_ERROR(
|
||||
gpu_helper_.RunInGlContext([this, &cc]() -> ::mediapipe::Status {
|
||||
return use_advanced_gpu_api_ ? InitTFLiteGPURunner()
|
||||
: LoadDelegate(cc);
|
||||
}));
|
||||
if (use_advanced_gpu_api_) return ::mediapipe::OkStatus();
|
||||
#else
|
||||
MP_RETURN_IF_ERROR(LoadDelegate(cc));
|
||||
#endif
|
||||
@@ -314,13 +381,51 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status TfLiteInferenceCalculator::InitTFLiteGPURunner() {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
// Create and bind OpenGL buffers for outputs.
|
||||
// These buffers are created onve and later their ids are jut passed to the
|
||||
// calculator outputs.
|
||||
|
||||
gpu_data_out_.resize(tflite_gpu_runner_->outputs_size());
|
||||
for (int i = 0; i < tflite_gpu_runner_->outputs_size(); ++i) {
|
||||
gpu_data_out_[i] = absl::make_unique<GPUData>();
|
||||
ASSIGN_OR_RETURN(gpu_data_out_[i]->elements,
|
||||
tflite_gpu_runner_->GetOutputElements(i));
|
||||
// Create and bind input buffer.
|
||||
RET_CHECK_CALL(::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
|
||||
gpu_data_out_[i]->elements, &gpu_data_out_[i]->buffer));
|
||||
}
|
||||
RET_CHECK_CALL(tflite_gpu_runner_->Build());
|
||||
#endif
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status TfLiteInferenceCalculator::Process(CalculatorContext* cc) {
|
||||
// 1. Receive pre-processed tensor inputs.
|
||||
if (gpu_input_) {
|
||||
// Read GPU input into SSBO.
|
||||
if (use_advanced_gpu_api_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||
RET_CHECK(input_tensors.empty());
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||
[this, &input_tensors]() -> ::mediapipe::Status {
|
||||
for (int i = 0; i < input_tensors.size(); ++i) {
|
||||
MP_RETURN_IF_ERROR(tflite_gpu_runner_->BindSSBOToInputTensor(
|
||||
input_tensors[i].id(), i));
|
||||
}
|
||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||
MP_RETURN_IF_ERROR(tflite_gpu_runner_->BindSSBOToOutputTensor(
|
||||
gpu_data_out_[i]->buffer.id(), i));
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
#endif
|
||||
} else if (gpu_input_) {
|
||||
// Read GPU input into SSBO.
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag(kTensorsGpuTag).Get<std::vector<GpuTensor>>();
|
||||
RET_CHECK_GT(input_tensors.size(), 0);
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||
[this, &input_tensors]() -> ::mediapipe::Status {
|
||||
@@ -335,7 +440,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
}));
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||
cc->Inputs().Tag(kTensorsGpuTag).Get<std::vector<GpuTensor>>();
|
||||
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());
|
||||
@@ -362,7 +467,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
} else {
|
||||
// Read CPU input into tensors.
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS").Get<std::vector<TfLiteTensor>>();
|
||||
cc->Inputs().Tag(kTensorsTag).Get<std::vector<TfLiteTensor>>();
|
||||
RET_CHECK_GT(input_tensors.size(), 0);
|
||||
for (int i = 0; i < input_tensors.size(); ++i) {
|
||||
const TfLiteTensor* input_tensor = &input_tensors[i];
|
||||
@@ -386,7 +491,11 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(
|
||||
gpu_helper_.RunInGlContext([this]() -> ::mediapipe::Status {
|
||||
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
|
||||
if (use_advanced_gpu_api_) {
|
||||
RET_CHECK(tflite_gpu_runner_->Invoke().ok());
|
||||
} else {
|
||||
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
@@ -397,7 +506,18 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
}
|
||||
|
||||
// 3. Output processed tensors.
|
||||
if (gpu_output_) {
|
||||
if (use_advanced_gpu_api_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
||||
output_tensors->resize(gpu_data_out_.size());
|
||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||
output_tensors->at(i) = gpu_data_out_[0]->buffer.MakeRef();
|
||||
}
|
||||
cc->Outputs()
|
||||
.Tag("TENSORS_GPU")
|
||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||
#endif
|
||||
} else if (gpu_output_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
// Output result tensors (GPU).
|
||||
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
||||
@@ -413,7 +533,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
cc->Outputs()
|
||||
.Tag("TENSORS_GPU")
|
||||
.Tag(kTensorsGpuTag)
|
||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
// Output result tensors (GPU).
|
||||
@@ -437,7 +557,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
[convert_command endEncoding];
|
||||
[command_buffer commit];
|
||||
cc->Outputs()
|
||||
.Tag("TENSORS_GPU")
|
||||
.Tag(kTensorsGpuTag)
|
||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||
#else
|
||||
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
||||
@@ -450,8 +570,9 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
TfLiteTensor* tensor = interpreter_->tensor(tensor_indexes[i]);
|
||||
output_tensors->emplace_back(*tensor);
|
||||
}
|
||||
cc->Outputs().Tag("TENSORS").Add(output_tensors.release(),
|
||||
cc->InputTimestamp());
|
||||
cc->Outputs()
|
||||
.Tag(kTensorsTag)
|
||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
@@ -462,6 +583,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
if (gpu_inference_) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
|
||||
interpreter_ = nullptr;
|
||||
delegate_ = nullptr;
|
||||
for (int i = 0; i < gpu_data_in_.size(); ++i) {
|
||||
gpu_data_in_[i].reset();
|
||||
@@ -472,6 +594,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
interpreter_ = nullptr;
|
||||
delegate_ = nullptr;
|
||||
for (int i = 0; i < gpu_data_in_.size(); ++i) {
|
||||
gpu_data_in_[i].reset();
|
||||
@@ -481,6 +604,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
}
|
||||
#endif
|
||||
} else {
|
||||
interpreter_ = nullptr;
|
||||
delegate_ = nullptr;
|
||||
}
|
||||
}
|
||||
@@ -492,34 +616,10 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
|
||||
// Calculator Auxiliary Section
|
||||
|
||||
::mediapipe::Status TfLiteInferenceCalculator::LoadOptions(
|
||||
CalculatorContext* cc) {
|
||||
// Get calculator options specified in the graph.
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::TfLiteInferenceCalculatorOptions>();
|
||||
|
||||
// Get model name.
|
||||
if (!options.model_path().empty()) {
|
||||
std::string model_path = options.model_path();
|
||||
|
||||
ASSIGN_OR_RETURN(model_path_, mediapipe::PathToResourceAsFile(model_path));
|
||||
} else {
|
||||
LOG(ERROR) << "Must specify path to TFLite model.";
|
||||
return ::mediapipe::Status(::mediapipe::StatusCode::kNotFound,
|
||||
"Must specify path to TFLite model.");
|
||||
}
|
||||
|
||||
// Get execution modes.
|
||||
gpu_inference_ =
|
||||
options.has_delegate() ? options.delegate().has_gpu() : options.use_gpu();
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status TfLiteInferenceCalculator::LoadModel(
|
||||
CalculatorContext* cc) {
|
||||
model_ = tflite::FlatBufferModel::BuildFromFile(model_path_.c_str());
|
||||
RET_CHECK(model_);
|
||||
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(*cc));
|
||||
const auto& model = *model_packet_.Get<TfLiteModelPtr>();
|
||||
|
||||
tflite::ops::builtin::BuiltinOpResolver op_resolver;
|
||||
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
|
||||
@@ -527,11 +627,25 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
.Tag("CUSTOM_OP_RESOLVER")
|
||||
.Get<tflite::ops::builtin::BuiltinOpResolver>();
|
||||
}
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
if (use_advanced_gpu_api_) {
|
||||
tflite::gpu::InferenceOptions options;
|
||||
options.priority1 = tflite::gpu::InferencePriority::MIN_LATENCY;
|
||||
options.priority2 = tflite::gpu::InferencePriority::AUTO;
|
||||
options.priority3 = tflite::gpu::InferencePriority::AUTO;
|
||||
options.usage = tflite::gpu::InferenceUsage::SUSTAINED_SPEED;
|
||||
tflite_gpu_runner_ =
|
||||
std::make_unique<tflite::gpu::TFLiteGPURunner>(options);
|
||||
return tflite_gpu_runner_->InitializeWithModel(model);
|
||||
}
|
||||
#endif
|
||||
|
||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||
interpreter_ =
|
||||
BuildEdgeTpuInterpreter(*model_, &op_resolver, edgetpu_context_.get());
|
||||
BuildEdgeTpuInterpreter(model, &op_resolver, edgetpu_context_.get());
|
||||
#else
|
||||
tflite::InterpreterBuilder(*model_, op_resolver)(&interpreter_);
|
||||
tflite::InterpreterBuilder(model, op_resolver)(&interpreter_);
|
||||
#endif // MEDIAPIPE_EDGE_TPU
|
||||
|
||||
RET_CHECK(interpreter_);
|
||||
@@ -557,6 +671,28 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::StatusOr<Packet> TfLiteInferenceCalculator::GetModelAsPacket(
|
||||
const CalculatorContext& cc) {
|
||||
const auto& options =
|
||||
cc.Options<mediapipe::TfLiteInferenceCalculatorOptions>();
|
||||
if (!options.model_path().empty()) {
|
||||
std::string model_path = options.model_path();
|
||||
|
||||
ASSIGN_OR_RETURN(model_path, mediapipe::PathToResourceAsFile(model_path));
|
||||
|
||||
auto model = tflite::FlatBufferModel::BuildFromFile(model_path.c_str());
|
||||
RET_CHECK(model) << "Failed to load model from path.";
|
||||
return MakePacket<TfLiteModelPtr>(TfLiteModelPtr(
|
||||
model.release(), [](tflite::FlatBufferModel* model) { delete model; }));
|
||||
}
|
||||
if (cc.InputSidePackets().HasTag("MODEL")) {
|
||||
return cc.InputSidePackets().Tag("MODEL");
|
||||
}
|
||||
return ::mediapipe::Status(
|
||||
::mediapipe::StatusCode::kNotFound,
|
||||
"Must specify TFLite model as path or loaded model.");
|
||||
}
|
||||
|
||||
::mediapipe::Status TfLiteInferenceCalculator::LoadDelegate(
|
||||
CalculatorContext* cc) {
|
||||
const auto& calculator_opts =
|
||||
@@ -587,6 +723,22 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
}
|
||||
#endif // MEDIAPIPE_ANDROID
|
||||
|
||||
#if defined(__EMSCRIPTEN__)
|
||||
const bool xnnpack_requested = true;
|
||||
#else
|
||||
const bool xnnpack_requested = calculator_opts.has_delegate() &&
|
||||
calculator_opts.delegate().has_xnnpack();
|
||||
#endif // __EMSCRIPTEN__
|
||||
|
||||
if (xnnpack_requested) {
|
||||
TfLiteXNNPackDelegateOptions xnnpack_opts{};
|
||||
xnnpack_opts.num_threads = GetXnnpackNumThreads(calculator_opts);
|
||||
delegate_ = TfLiteDelegatePtr(TfLiteXNNPackDelegateCreate(&xnnpack_opts),
|
||||
&TfLiteXNNPackDelegateDelete);
|
||||
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
|
||||
kTfLiteOk);
|
||||
}
|
||||
|
||||
// Return, no need for GPU delegate below.
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -608,7 +760,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
||||
const auto& input_indices = interpreter_->inputs();
|
||||
gpu_data_in_.resize(input_indices.size());
|
||||
for (int i = 0; i < input_indices.size(); ++i) {
|
||||
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[0]);
|
||||
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[i]);
|
||||
gpu_data_in_[i] = absl::make_unique<GPUData>();
|
||||
gpu_data_in_[i]->elements = 1;
|
||||
for (int d = 0; d < tensor->dims->size; ++d) {
|
||||
|
||||
@@ -42,14 +42,24 @@ message TfLiteInferenceCalculatorOptions {
|
||||
message TfLite {}
|
||||
// Delegate to run GPU inference depending on the device.
|
||||
// (Can use OpenGl, OpenCl, Metal depending on the device.)
|
||||
message Gpu {}
|
||||
message Gpu {
|
||||
// Experimental, Android/Linux only. Use TFLite GPU delegate API2 for
|
||||
// the NN inference.
|
||||
optional bool use_advanced_gpu_api = 1 [default = false];
|
||||
}
|
||||
// Android only.
|
||||
message Nnapi {}
|
||||
message Xnnpack {
|
||||
// Number of threads for XNNPACK delegate. (By default, calculator tries
|
||||
// to choose optimal number of threads depending on the device.)
|
||||
optional int32 num_threads = 1 [default = -1];
|
||||
}
|
||||
|
||||
oneof delegate {
|
||||
TfLite tflite = 1;
|
||||
Gpu gpu = 2;
|
||||
Nnapi nnapi = 3;
|
||||
Xnnpack xnnpack = 4;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -41,7 +41,7 @@ namespace mediapipe {
|
||||
|
||||
using ::tflite::Interpreter;
|
||||
|
||||
void DoSmokeTest(absl::string_view delegate) {
|
||||
void DoSmokeTest(const std::string& graph_proto) {
|
||||
const int width = 8;
|
||||
const int height = 8;
|
||||
const int channels = 3;
|
||||
@@ -69,24 +69,9 @@ void DoSmokeTest(absl::string_view delegate) {
|
||||
auto input_vec = absl::make_unique<std::vector<TfLiteTensor>>();
|
||||
input_vec->emplace_back(*tensor);
|
||||
|
||||
std::string graph_proto = R"(
|
||||
input_stream: "tensor_in"
|
||||
node {
|
||||
calculator: "TfLiteInferenceCalculator"
|
||||
input_stream: "TENSORS:tensor_in"
|
||||
output_stream: "TENSORS:tensor_out"
|
||||
options {
|
||||
[mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
||||
model_path: "mediapipe/calculators/tflite/testdata/add.bin"
|
||||
$delegate
|
||||
}
|
||||
}
|
||||
}
|
||||
)";
|
||||
ASSERT_EQ(absl::StrReplaceAll({{"$delegate", delegate}}, &graph_proto), 1);
|
||||
// Prepare single calculator graph to and wait for packets.
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(graph_proto);
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig>(graph_proto);
|
||||
std::vector<Packet> output_packets;
|
||||
tool::AddVectorSink("tensor_out", &graph_config, &output_packets);
|
||||
CalculatorGraph graph(graph_config);
|
||||
@@ -119,8 +104,70 @@ void DoSmokeTest(absl::string_view delegate) {
|
||||
|
||||
// Tests a simple add model that adds an input tensor to itself.
|
||||
TEST(TfLiteInferenceCalculatorTest, SmokeTest) {
|
||||
DoSmokeTest(/*delegate=*/"");
|
||||
DoSmokeTest(/*delegate=*/"delegate { tflite {} }");
|
||||
std::string graph_proto = R"(
|
||||
input_stream: "tensor_in"
|
||||
node {
|
||||
calculator: "TfLiteInferenceCalculator"
|
||||
input_stream: "TENSORS:tensor_in"
|
||||
output_stream: "TENSORS:tensor_out"
|
||||
options {
|
||||
[mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
||||
model_path: "mediapipe/calculators/tflite/testdata/add.bin"
|
||||
$delegate
|
||||
}
|
||||
}
|
||||
}
|
||||
)";
|
||||
DoSmokeTest(
|
||||
/*graph_proto=*/absl::StrReplaceAll(graph_proto, {{"$delegate", ""}}));
|
||||
DoSmokeTest(/*graph_proto=*/absl::StrReplaceAll(
|
||||
graph_proto, {{"$delegate", "delegate { tflite {} }"}}));
|
||||
DoSmokeTest(/*graph_proto=*/absl::StrReplaceAll(
|
||||
graph_proto, {{"$delegate", "delegate { xnnpack {} }"}}));
|
||||
DoSmokeTest(/*graph_proto=*/absl::StrReplaceAll(
|
||||
graph_proto,
|
||||
{{"$delegate", "delegate { xnnpack { num_threads: 10 } }"}}));
|
||||
}
|
||||
|
||||
TEST(TfLiteInferenceCalculatorTest, SmokeTest_ModelAsInputSidePacket) {
|
||||
std::string graph_proto = R"(
|
||||
input_stream: "tensor_in"
|
||||
|
||||
node {
|
||||
calculator: "ConstantSidePacketCalculator"
|
||||
output_side_packet: "PACKET:model_path"
|
||||
options: {
|
||||
[mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
|
||||
packet { string_value: "mediapipe/calculators/tflite/testdata/add.bin" }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
node {
|
||||
calculator: "LocalFileContentsCalculator"
|
||||
input_side_packet: "FILE_PATH:model_path"
|
||||
output_side_packet: "CONTENTS:model_blob"
|
||||
}
|
||||
|
||||
node {
|
||||
calculator: "TfLiteModelCalculator"
|
||||
input_side_packet: "MODEL_BLOB:model_blob"
|
||||
output_side_packet: "MODEL:model"
|
||||
}
|
||||
|
||||
node {
|
||||
calculator: "TfLiteInferenceCalculator"
|
||||
input_stream: "TENSORS:tensor_in"
|
||||
output_stream: "TENSORS:tensor_out"
|
||||
input_side_packet: "MODEL:model"
|
||||
options {
|
||||
[mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
||||
use_gpu: false
|
||||
}
|
||||
}
|
||||
}
|
||||
)";
|
||||
DoSmokeTest(graph_proto);
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
// Copyright 2020 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 <functional>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "tensorflow/lite/model.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Loads TfLite model from model blob specified as input side packet and outputs
|
||||
// corresponding side packet.
|
||||
//
|
||||
// Input side packets:
|
||||
// MODEL_BLOB - TfLite model blob/file-contents (std::string). You can read
|
||||
// model blob from file (using whatever APIs you have) and pass
|
||||
// it to the graph as input side packet or you can use some of
|
||||
// calculators like LocalFileContentsCalculator to get model
|
||||
// blob and use it as input here.
|
||||
//
|
||||
// Output side packets:
|
||||
// MODEL - TfLite model. (std::unique_ptr<tflite::FlatBufferModel,
|
||||
// std::function<void(tflite::FlatBufferModel*)>>)
|
||||
//
|
||||
// Example use:
|
||||
//
|
||||
// node {
|
||||
// calculator: "TfLiteModelCalculator"
|
||||
// input_side_packet: "MODEL_BLOB:model_blob"
|
||||
// output_side_packet: "MODEL:model"
|
||||
// }
|
||||
//
|
||||
class TfLiteModelCalculator : public CalculatorBase {
|
||||
public:
|
||||
using TfLiteModelPtr =
|
||||
std::unique_ptr<tflite::FlatBufferModel,
|
||||
std::function<void(tflite::FlatBufferModel*)>>;
|
||||
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->InputSidePackets().Tag("MODEL_BLOB").Set<std::string>();
|
||||
cc->OutputSidePackets().Tag("MODEL").Set<TfLiteModelPtr>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
const Packet& model_packet = cc->InputSidePackets().Tag("MODEL_BLOB");
|
||||
const std::string& model_blob = model_packet.Get<std::string>();
|
||||
std::unique_ptr<tflite::FlatBufferModel> model =
|
||||
tflite::FlatBufferModel::BuildFromBuffer(model_blob.data(),
|
||||
model_blob.size());
|
||||
RET_CHECK(model) << "Failed to load TfLite model from blob.";
|
||||
|
||||
cc->OutputSidePackets().Tag("MODEL").Set(
|
||||
MakePacket<TfLiteModelPtr>(TfLiteModelPtr(
|
||||
model.release(), [model_packet](tflite::FlatBufferModel* model) {
|
||||
// Keeping model_packet in order to keep underlying model blob
|
||||
// which can be released only after TfLite model is not needed
|
||||
// anymore (deleted).
|
||||
delete model;
|
||||
})));
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
REGISTER_CALCULATOR(TfLiteModelCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,88 @@
|
||||
// Copyright 2020 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 "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
|
||||
#include "tensorflow/lite/model.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
TEST(TfLiteModelCalculatorTest, SmokeTest) {
|
||||
// Prepare single calculator graph to and wait for packets.
|
||||
CalculatorGraphConfig graph_config = ParseTextProtoOrDie<
|
||||
CalculatorGraphConfig>(
|
||||
R"(
|
||||
node {
|
||||
calculator: "ConstantSidePacketCalculator"
|
||||
output_side_packet: "PACKET:model_path"
|
||||
options: {
|
||||
[mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
|
||||
packet {
|
||||
string_value: "mediapipe/calculators/tflite/testdata/add.bin"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
node {
|
||||
calculator: "LocalFileContentsCalculator"
|
||||
input_side_packet: "FILE_PATH:model_path"
|
||||
output_side_packet: "CONTENTS:model_blob"
|
||||
}
|
||||
|
||||
node {
|
||||
calculator: "TfLiteModelCalculator"
|
||||
input_side_packet: "MODEL_BLOB:model_blob"
|
||||
output_side_packet: "MODEL:model"
|
||||
}
|
||||
)");
|
||||
CalculatorGraph graph(graph_config);
|
||||
MP_ASSERT_OK(graph.StartRun({}));
|
||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||
auto status_or_packet = graph.GetOutputSidePacket("model");
|
||||
MP_ASSERT_OK(status_or_packet);
|
||||
auto model_packet = status_or_packet.ValueOrDie();
|
||||
const auto& model = model_packet.Get<
|
||||
std::unique_ptr<tflite::FlatBufferModel,
|
||||
std::function<void(tflite::FlatBufferModel*)>>>();
|
||||
|
||||
auto expected_model = tflite::FlatBufferModel::BuildFromFile(
|
||||
"mediapipe/calculators/tflite/testdata/add.bin");
|
||||
|
||||
EXPECT_EQ(model->GetModel()->version(),
|
||||
expected_model->GetModel()->version());
|
||||
EXPECT_EQ(model->GetModel()->buffers()->size(),
|
||||
expected_model->GetModel()->buffers()->size());
|
||||
const int num_subgraphs = expected_model->GetModel()->subgraphs()->size();
|
||||
EXPECT_EQ(model->GetModel()->subgraphs()->size(), num_subgraphs);
|
||||
for (int i = 0; i < num_subgraphs; ++i) {
|
||||
const auto* expected_subgraph =
|
||||
expected_model->GetModel()->subgraphs()->Get(i);
|
||||
const auto* subgraph = model->GetModel()->subgraphs()->Get(i);
|
||||
const int num_tensors = expected_subgraph->tensors()->size();
|
||||
EXPECT_EQ(subgraph->tensors()->size(), num_tensors);
|
||||
for (int j = 0; j < num_tensors; ++j) {
|
||||
EXPECT_EQ(subgraph->tensors()->Get(j)->name()->str(),
|
||||
expected_subgraph->tensors()->Get(j)->name()->str());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -129,22 +129,43 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
||||
num_classes *= raw_score_tensor->dims->data[i];
|
||||
}
|
||||
|
||||
if (options_.binary_classification()) {
|
||||
RET_CHECK_EQ(num_classes, 1);
|
||||
// Number of classes for binary classification.
|
||||
num_classes = 2;
|
||||
}
|
||||
if (label_map_loaded_) {
|
||||
RET_CHECK_EQ(num_classes, label_map_.size());
|
||||
}
|
||||
const float* raw_scores = raw_score_tensor->data.f;
|
||||
|
||||
auto classification_list = absl::make_unique<ClassificationList>();
|
||||
for (int i = 0; i < num_classes; ++i) {
|
||||
if (options_.has_min_score_threshold() &&
|
||||
raw_scores[i] < options_.min_score_threshold()) {
|
||||
continue;
|
||||
}
|
||||
Classification* classification = classification_list->add_classification();
|
||||
classification->set_index(i);
|
||||
classification->set_score(raw_scores[i]);
|
||||
if (options_.binary_classification()) {
|
||||
Classification* class_first = classification_list->add_classification();
|
||||
Classification* class_second = classification_list->add_classification();
|
||||
class_first->set_index(0);
|
||||
class_second->set_index(1);
|
||||
class_first->set_score(raw_scores[0]);
|
||||
class_second->set_score(1. - raw_scores[0]);
|
||||
|
||||
if (label_map_loaded_) {
|
||||
classification->set_label(label_map_[i]);
|
||||
class_first->set_label(label_map_[0]);
|
||||
class_second->set_label(label_map_[1]);
|
||||
}
|
||||
} else {
|
||||
for (int i = 0; i < num_classes; ++i) {
|
||||
if (options_.has_min_score_threshold() &&
|
||||
raw_scores[i] < options_.min_score_threshold()) {
|
||||
continue;
|
||||
}
|
||||
Classification* classification =
|
||||
classification_list->add_classification();
|
||||
classification->set_index(i);
|
||||
classification->set_score(raw_scores[i]);
|
||||
|
||||
if (label_map_loaded_) {
|
||||
classification->set_label(label_map_[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -32,4 +32,10 @@ message TfLiteTensorsToClassificationCalculatorOptions {
|
||||
optional int32 top_k = 2;
|
||||
// Path to a label map file for getting the actual name of class ids.
|
||||
optional string label_map_path = 3;
|
||||
// Whether the input is a single float for binary classification.
|
||||
// When true, only a single float is expected in the input tensor and the
|
||||
// label map, if provided, is expected to have exactly two labels.
|
||||
// The single score(float) represent the probability of first label, and
|
||||
// 1 - score is the probabilility of the second label.
|
||||
optional bool binary_classification = 4;
|
||||
}
|
||||
|
||||
@@ -47,10 +47,11 @@
|
||||
#endif // iOS
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kNumInputTensorsWithAnchors = 3;
|
||||
constexpr int kNumCoordsPerBox = 4;
|
||||
|
||||
constexpr char kTensorsTag[] = "TENSORS";
|
||||
constexpr char kTensorsGpuTag[] = "TENSORS_GPU";
|
||||
} // namespace
|
||||
|
||||
namespace mediapipe {
|
||||
@@ -99,7 +100,7 @@ void ConvertAnchorsToRawValues(const std::vector<Anchor>& anchors,
|
||||
int num_boxes, float* raw_anchors) {
|
||||
CHECK_EQ(anchors.size(), num_boxes);
|
||||
int box = 0;
|
||||
for (auto anchor : anchors) {
|
||||
for (const auto& anchor : anchors) {
|
||||
raw_anchors[box * kNumCoordsPerBox + 0] = anchor.y_center();
|
||||
raw_anchors[box * kNumCoordsPerBox + 1] = anchor.x_center();
|
||||
raw_anchors[box * kNumCoordsPerBox + 2] = anchor.h();
|
||||
@@ -200,13 +201,13 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
|
||||
bool use_gpu = false;
|
||||
|
||||
if (cc->Inputs().HasTag("TENSORS")) {
|
||||
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||
if (cc->Inputs().HasTag(kTensorsTag)) {
|
||||
cc->Inputs().Tag(kTensorsTag).Set<std::vector<TfLiteTensor>>();
|
||||
}
|
||||
|
||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||
if (cc->Inputs().HasTag(kTensorsGpuTag)) {
|
||||
cc->Inputs().Tag(kTensorsGpuTag).Set<std::vector<GpuTensor>>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
@@ -236,7 +237,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
CalculatorContext* cc) {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
if (cc->Inputs().HasTag(kTensorsGpuTag)) {
|
||||
gpu_input_ = true;
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||
@@ -258,8 +259,8 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
if ((!gpu_input_ && cc->Inputs().Tag("TENSORS").IsEmpty()) ||
|
||||
(gpu_input_ && cc->Inputs().Tag("TENSORS_GPU").IsEmpty())) {
|
||||
if ((!gpu_input_ && cc->Inputs().Tag(kTensorsTag).IsEmpty()) ||
|
||||
(gpu_input_ && cc->Inputs().Tag(kTensorsGpuTag).IsEmpty())) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -284,7 +285,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::ProcessCPU(
|
||||
CalculatorContext* cc, std::vector<Detection>* output_detections) {
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS").Get<std::vector<TfLiteTensor>>();
|
||||
cc->Inputs().Tag(kTensorsTag).Get<std::vector<TfLiteTensor>>();
|
||||
|
||||
if (input_tensors.size() == 2 ||
|
||||
input_tensors.size() == kNumInputTensorsWithAnchors) {
|
||||
@@ -402,7 +403,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
CalculatorContext* cc, std::vector<Detection>* output_detections) {
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||
cc->Inputs().Tag(kTensorsGpuTag).Get<std::vector<GpuTensor>>();
|
||||
RET_CHECK_GE(input_tensors.size(), 2);
|
||||
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, &input_tensors, &cc,
|
||||
@@ -466,7 +467,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
||||
#elif defined(MEDIAPIPE_IOS)
|
||||
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||
cc->Inputs().Tag(kTensorsGpuTag).Get<std::vector<GpuTensor>>();
|
||||
RET_CHECK_GE(input_tensors.size(), 2);
|
||||
|
||||
// Copy inputs.
|
||||
|
||||
@@ -145,7 +145,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToLandmarksCalculator);
|
||||
? cc->InputSidePackets().Tag("FLIP_HORIZONTALLY").Get<bool>()
|
||||
: options_.flip_horizontally();
|
||||
|
||||
flip_horizontally_ =
|
||||
flip_vertically_ =
|
||||
cc->InputSidePackets().HasTag("FLIP_VERTICALLY")
|
||||
? cc->InputSidePackets().Tag("FLIP_VERTICALLY").Get<bool>()
|
||||
: options_.flip_vertically();
|
||||
|
||||
@@ -49,6 +49,16 @@ int NumGroups(const int size, const int group_size) { // NOLINT
|
||||
float Clamp(float val, float min, float max) {
|
||||
return std::min(std::max(val, min), max);
|
||||
}
|
||||
|
||||
constexpr char kTensorsTag[] = "TENSORS";
|
||||
constexpr char kTensorsGpuTag[] = "TENSORS_GPU";
|
||||
constexpr char kSizeImageTag[] = "REFERENCE_IMAGE";
|
||||
constexpr char kSizeImageGpuTag[] = "REFERENCE_IMAGE_GPU";
|
||||
constexpr char kMaskTag[] = "MASK";
|
||||
constexpr char kMaskGpuTag[] = "MASK_GPU";
|
||||
constexpr char kPrevMaskTag[] = "PREV_MASK";
|
||||
constexpr char kPrevMaskGpuTag[] = "PREV_MASK_GPU";
|
||||
|
||||
} // namespace
|
||||
|
||||
namespace mediapipe {
|
||||
@@ -148,39 +158,39 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
bool use_gpu = false;
|
||||
|
||||
// Inputs CPU.
|
||||
if (cc->Inputs().HasTag("TENSORS")) {
|
||||
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||
if (cc->Inputs().HasTag(kTensorsTag)) {
|
||||
cc->Inputs().Tag(kTensorsTag).Set<std::vector<TfLiteTensor>>();
|
||||
}
|
||||
if (cc->Inputs().HasTag("PREV_MASK")) {
|
||||
cc->Inputs().Tag("PREV_MASK").Set<ImageFrame>();
|
||||
if (cc->Inputs().HasTag(kPrevMaskTag)) {
|
||||
cc->Inputs().Tag(kPrevMaskTag).Set<ImageFrame>();
|
||||
}
|
||||
if (cc->Inputs().HasTag("REFERENCE_IMAGE")) {
|
||||
cc->Inputs().Tag("REFERENCE_IMAGE").Set<ImageFrame>();
|
||||
if (cc->Inputs().HasTag(kSizeImageTag)) {
|
||||
cc->Inputs().Tag(kSizeImageTag).Set<ImageFrame>();
|
||||
}
|
||||
|
||||
// Inputs GPU.
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GlBuffer>>();
|
||||
if (cc->Inputs().HasTag(kTensorsGpuTag)) {
|
||||
cc->Inputs().Tag(kTensorsGpuTag).Set<std::vector<GlBuffer>>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
if (cc->Inputs().HasTag("PREV_MASK_GPU")) {
|
||||
cc->Inputs().Tag("PREV_MASK_GPU").Set<mediapipe::GpuBuffer>();
|
||||
if (cc->Inputs().HasTag(kPrevMaskGpuTag)) {
|
||||
cc->Inputs().Tag(kPrevMaskGpuTag).Set<mediapipe::GpuBuffer>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
if (cc->Inputs().HasTag("REFERENCE_IMAGE_GPU")) {
|
||||
cc->Inputs().Tag("REFERENCE_IMAGE_GPU").Set<mediapipe::GpuBuffer>();
|
||||
if (cc->Inputs().HasTag(kSizeImageGpuTag)) {
|
||||
cc->Inputs().Tag(kSizeImageGpuTag).Set<mediapipe::GpuBuffer>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
|
||||
// Outputs.
|
||||
if (cc->Outputs().HasTag("MASK")) {
|
||||
cc->Outputs().Tag("MASK").Set<ImageFrame>();
|
||||
if (cc->Outputs().HasTag(kMaskTag)) {
|
||||
cc->Outputs().Tag(kMaskTag).Set<ImageFrame>();
|
||||
}
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
if (cc->Outputs().HasTag("MASK_GPU")) {
|
||||
cc->Outputs().Tag("MASK_GPU").Set<mediapipe::GpuBuffer>();
|
||||
if (cc->Outputs().HasTag(kMaskGpuTag)) {
|
||||
cc->Outputs().Tag(kMaskGpuTag).Set<mediapipe::GpuBuffer>();
|
||||
use_gpu |= true;
|
||||
}
|
||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||
@@ -197,7 +207,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
CalculatorContext* cc) {
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
|
||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||
if (cc->Inputs().HasTag(kTensorsGpuTag)) {
|
||||
use_gpu_ = true;
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||
@@ -255,23 +265,22 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
|
||||
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::ProcessCpu(
|
||||
CalculatorContext* cc) {
|
||||
if (cc->Inputs().Tag("TENSORS").IsEmpty()) {
|
||||
if (cc->Inputs().Tag(kTensorsTag).IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
// Get input streams.
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS").Get<std::vector<TfLiteTensor>>();
|
||||
const bool has_prev_mask = cc->Inputs().HasTag("PREV_MASK") &&
|
||||
!cc->Inputs().Tag("PREV_MASK").IsEmpty();
|
||||
cc->Inputs().Tag(kTensorsTag).Get<std::vector<TfLiteTensor>>();
|
||||
const bool has_prev_mask = cc->Inputs().HasTag(kPrevMaskTag) &&
|
||||
!cc->Inputs().Tag(kPrevMaskTag).IsEmpty();
|
||||
const ImageFrame placeholder;
|
||||
const auto& input_mask = has_prev_mask
|
||||
? cc->Inputs().Tag("PREV_MASK").Get<ImageFrame>()
|
||||
: placeholder;
|
||||
const auto& input_mask =
|
||||
has_prev_mask ? cc->Inputs().Tag(kPrevMaskTag).Get<ImageFrame>()
|
||||
: placeholder;
|
||||
int output_width = tensor_width_, output_height = tensor_height_;
|
||||
if (cc->Inputs().HasTag("REFERENCE_IMAGE")) {
|
||||
const auto& input_image =
|
||||
cc->Inputs().Tag("REFERENCE_IMAGE").Get<ImageFrame>();
|
||||
if (cc->Inputs().HasTag(kSizeImageTag)) {
|
||||
const auto& input_image = cc->Inputs().Tag(kSizeImageTag).Get<ImageFrame>();
|
||||
output_width = input_image.Width();
|
||||
output_height = input_image.Height();
|
||||
}
|
||||
@@ -353,7 +362,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
ImageFormat::SRGBA, output_width, output_height);
|
||||
cv::Mat output_mat = formats::MatView(output_mask.get());
|
||||
large_mask_mat.copyTo(output_mat);
|
||||
cc->Outputs().Tag("MASK").Add(output_mask.release(), cc->InputTimestamp());
|
||||
cc->Outputs().Tag(kMaskTag).Add(output_mask.release(), cc->InputTimestamp());
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -364,23 +373,23 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
// 3. upsample small mask into output mask to be same size as input image
|
||||
::mediapipe::Status TfLiteTensorsToSegmentationCalculator::ProcessGpu(
|
||||
CalculatorContext* cc) {
|
||||
if (cc->Inputs().Tag("TENSORS_GPU").IsEmpty()) {
|
||||
if (cc->Inputs().Tag(kTensorsGpuTag).IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||
// Get input streams.
|
||||
const auto& input_tensors =
|
||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GlBuffer>>();
|
||||
const bool has_prev_mask = cc->Inputs().HasTag("PREV_MASK_GPU") &&
|
||||
!cc->Inputs().Tag("PREV_MASK_GPU").IsEmpty();
|
||||
cc->Inputs().Tag(kTensorsGpuTag).Get<std::vector<GlBuffer>>();
|
||||
const bool has_prev_mask = cc->Inputs().HasTag(kPrevMaskGpuTag) &&
|
||||
!cc->Inputs().Tag(kPrevMaskGpuTag).IsEmpty();
|
||||
const auto& input_mask =
|
||||
has_prev_mask
|
||||
? cc->Inputs().Tag("PREV_MASK_GPU").Get<mediapipe::GpuBuffer>()
|
||||
? cc->Inputs().Tag(kPrevMaskGpuTag).Get<mediapipe::GpuBuffer>()
|
||||
: mediapipe::GpuBuffer();
|
||||
int output_width = tensor_width_, output_height = tensor_height_;
|
||||
if (cc->Inputs().HasTag("REFERENCE_IMAGE_GPU")) {
|
||||
if (cc->Inputs().HasTag(kSizeImageGpuTag)) {
|
||||
const auto& input_image =
|
||||
cc->Inputs().Tag("REFERENCE_IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
||||
cc->Inputs().Tag(kSizeImageGpuTag).Get<mediapipe::GpuBuffer>();
|
||||
output_width = input_image.width();
|
||||
output_height = input_image.height();
|
||||
}
|
||||
@@ -441,7 +450,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToSegmentationCalculator);
|
||||
// Send out image as GPU packet.
|
||||
auto output_image = output_texture.GetFrame<mediapipe::GpuBuffer>();
|
||||
cc->Outputs()
|
||||
.Tag("MASK_GPU")
|
||||
.Tag(kMaskGpuTag)
|
||||
.Add(output_image.release(), cc->InputTimestamp());
|
||||
|
||||
// Cleanup
|
||||
|
||||
@@ -15,11 +15,11 @@
|
||||
#ifndef MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
|
||||
#define MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
|
||||
|
||||
#define RET_CHECK_CALL(call) \
|
||||
do { \
|
||||
const auto status = (call); \
|
||||
if (ABSL_PREDICT_FALSE(!status.ok())) \
|
||||
return ::mediapipe::InternalError(status.error_message()); \
|
||||
#define RET_CHECK_CALL(call) \
|
||||
do { \
|
||||
const auto status = (call); \
|
||||
if (ABSL_PREDICT_FALSE(!status.ok())) \
|
||||
return ::mediapipe::InternalError(status.message()); \
|
||||
} while (0);
|
||||
|
||||
#endif // MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
|
||||
|
||||
@@ -321,7 +321,7 @@ cc_library(
|
||||
"//mediapipe:android": [
|
||||
"//mediapipe/util/android/file/base",
|
||||
],
|
||||
"//mediapipe:apple": [
|
||||
"//mediapipe:ios": [
|
||||
"//mediapipe/util/android/file/base",
|
||||
],
|
||||
"//mediapipe:macos": [
|
||||
@@ -349,7 +349,7 @@ cc_library(
|
||||
"//mediapipe:android": [
|
||||
"//mediapipe/util/android/file/base",
|
||||
],
|
||||
"//mediapipe:apple": [
|
||||
"//mediapipe:ios": [
|
||||
"//mediapipe/util/android/file/base",
|
||||
],
|
||||
"//mediapipe:macos": [
|
||||
@@ -926,7 +926,7 @@ cc_library(
|
||||
"//mediapipe:android": [
|
||||
"//mediapipe/util/android/file/base",
|
||||
],
|
||||
"//mediapipe:apple": [
|
||||
"//mediapipe:ios": [
|
||||
"//mediapipe/util/android/file/base",
|
||||
],
|
||||
"//mediapipe:macos": [
|
||||
@@ -971,9 +971,9 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:file_helpers",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/util:resource_util",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
@@ -998,6 +998,7 @@ cc_library(
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:classification_cc_proto",
|
||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||
"//mediapipe/framework/formats:rect_cc_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
@@ -1015,6 +1016,7 @@ cc_library(
|
||||
deps = [
|
||||
":collection_has_min_size_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",
|
||||
@@ -1022,6 +1024,18 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "collection_has_min_size_calculator_test",
|
||||
srcs = ["collection_has_min_size_calculator_test.cc"],
|
||||
deps = [
|
||||
":collection_has_min_size_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:calculator_runner",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "association_calculator",
|
||||
hdrs = ["association_calculator.h"],
|
||||
|
||||
@@ -55,7 +55,7 @@ size_t RoundUp(size_t n, size_t m) { return ((n + m - 1) / m) * m; } // NOLINT
|
||||
// When using GPU, this color will become transparent when the calculator
|
||||
// merges the annotation overlay with the image frame. As a result, drawing in
|
||||
// this color is not supported and it should be set to something unlikely used.
|
||||
constexpr int kAnnotationBackgroundColor[] = {100, 101, 102};
|
||||
constexpr uchar kAnnotationBackgroundColor = 2; // Grayscale value.
|
||||
} // namespace
|
||||
|
||||
// A calculator for rendering data on images.
|
||||
@@ -491,11 +491,9 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
|
||||
if (format != mediapipe::ImageFormat::SRGBA &&
|
||||
format != mediapipe::ImageFormat::SRGB)
|
||||
RET_CHECK_FAIL() << "Unsupported GPU input format: " << format;
|
||||
|
||||
image_mat = absl::make_unique<cv::Mat>(
|
||||
height_, width_, CV_8UC3,
|
||||
cv::Scalar(kAnnotationBackgroundColor[0], kAnnotationBackgroundColor[1],
|
||||
kAnnotationBackgroundColor[2]));
|
||||
image_mat = absl::make_unique<cv::Mat>(height_, width_, CV_8UC3);
|
||||
memset(image_mat->data, kAnnotationBackgroundColor,
|
||||
height_ * width_ * image_mat->elemSize());
|
||||
} else {
|
||||
image_mat = absl::make_unique<cv::Mat>(
|
||||
options_.canvas_height_px(), options_.canvas_width_px(), CV_8UC3,
|
||||
@@ -617,9 +615,9 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
|
||||
glUniform1i(glGetUniformLocation(program_, "input_frame"), 1);
|
||||
glUniform1i(glGetUniformLocation(program_, "overlay"), 2);
|
||||
glUniform3f(glGetUniformLocation(program_, "transparent_color"),
|
||||
kAnnotationBackgroundColor[0] / 255.0,
|
||||
kAnnotationBackgroundColor[1] / 255.0,
|
||||
kAnnotationBackgroundColor[2] / 255.0);
|
||||
kAnnotationBackgroundColor / 255.0,
|
||||
kAnnotationBackgroundColor / 255.0,
|
||||
kAnnotationBackgroundColor / 255.0);
|
||||
|
||||
// Init texture for opencv rendered frame.
|
||||
const auto& input_frame =
|
||||
|
||||
@@ -15,6 +15,9 @@
|
||||
|
||||
#include "mediapipe/calculators/util/collection_has_min_size_calculator.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
@@ -23,4 +26,9 @@ typedef CollectionHasMinSizeCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||
NormalizedRectVectorHasMinSizeCalculator;
|
||||
REGISTER_CALCULATOR(NormalizedRectVectorHasMinSizeCalculator);
|
||||
|
||||
typedef CollectionHasMinSizeCalculator<
|
||||
std::vector<::mediapipe::NormalizedLandmarkList>>
|
||||
NormalizedLandmarkListVectorHasMinSizeCalculator;
|
||||
REGISTER_CALCULATOR(NormalizedLandmarkListVectorHasMinSizeCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
// Copyright 2020 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 <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 CollectionHasMinSizeCalculator<std::vector<int>>
|
||||
TestIntCollectionHasMinSizeCalculator;
|
||||
REGISTER_CALCULATOR(TestIntCollectionHasMinSizeCalculator);
|
||||
|
||||
void AddInputVector(const std::vector<int>& input, int64 timestamp,
|
||||
CalculatorRunner* runner) {
|
||||
runner->MutableInputs()
|
||||
->Tag("ITERABLE")
|
||||
.packets.push_back(
|
||||
MakePacket<std::vector<int>>(input).At(Timestamp(timestamp)));
|
||||
}
|
||||
|
||||
TEST(TestIntCollectionHasMinSizeCalculator, DoesHaveMinSize) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "TestIntCollectionHasMinSizeCalculator"
|
||||
input_stream: "ITERABLE:input_vector"
|
||||
output_stream: "output_vector"
|
||||
options {
|
||||
[mediapipe.CollectionHasMinSizeCalculatorOptions.ext] { min_size: 2 }
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||
|
||||
AddInputVector({1, 2}, /*timestamp=*/1, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||
EXPECT_TRUE(outputs[0].Get<bool>());
|
||||
|
||||
AddInputVector({1, 2, 3}, /*timestamp=*/2, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
EXPECT_EQ(2, outputs.size());
|
||||
EXPECT_EQ(Timestamp(2), outputs[1].Timestamp());
|
||||
EXPECT_TRUE(outputs[1].Get<bool>());
|
||||
}
|
||||
|
||||
TEST(TestIntCollectionHasMinSizeCalculator,
|
||||
DoesHaveMinSize_MinSizeAsSidePacket) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "TestIntCollectionHasMinSizeCalculator"
|
||||
input_stream: "ITERABLE:input_vector"
|
||||
input_side_packet: "min_size"
|
||||
output_stream: "output_vector"
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||
|
||||
runner.MutableSidePackets()->Index(0) = MakePacket<int>(2);
|
||||
|
||||
AddInputVector({1, 2}, /*timestamp=*/1, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||
EXPECT_TRUE(outputs[0].Get<bool>());
|
||||
|
||||
AddInputVector({1, 2, 3}, /*timestamp=*/2, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
EXPECT_EQ(2, outputs.size());
|
||||
EXPECT_EQ(Timestamp(2), outputs[1].Timestamp());
|
||||
EXPECT_TRUE(outputs[1].Get<bool>());
|
||||
}
|
||||
|
||||
TEST(TestIntCollectionHasMinSizeCalculator, DoesNotHaveMinSize) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "TestIntCollectionHasMinSizeCalculator"
|
||||
input_stream: "ITERABLE:input_vector"
|
||||
output_stream: "output_vector"
|
||||
options {
|
||||
[mediapipe.CollectionHasMinSizeCalculatorOptions.ext] { min_size: 3 }
|
||||
}
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||
|
||||
AddInputVector({1}, /*timestamp=*/1, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||
EXPECT_FALSE(outputs[0].Get<bool>());
|
||||
|
||||
AddInputVector({1, 2}, /*timestamp=*/2, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
EXPECT_EQ(2, outputs.size());
|
||||
EXPECT_EQ(Timestamp(2), outputs[1].Timestamp());
|
||||
EXPECT_FALSE(outputs[1].Get<bool>());
|
||||
}
|
||||
|
||||
TEST(TestIntCollectionHasMinSizeCalculator,
|
||||
DoesNotHaveMinSize_MinSizeAsSidePacket) {
|
||||
CalculatorGraphConfig::Node node_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||
calculator: "TestIntCollectionHasMinSizeCalculator"
|
||||
input_stream: "ITERABLE:input_vector"
|
||||
input_side_packet: "min_size"
|
||||
output_stream: "output_vector"
|
||||
)");
|
||||
CalculatorRunner runner(node_config);
|
||||
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||
|
||||
runner.MutableSidePackets()->Index(0) = MakePacket<int>(3);
|
||||
|
||||
AddInputVector({1}, /*timestamp=*/1, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
EXPECT_EQ(1, outputs.size());
|
||||
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||
EXPECT_FALSE(outputs[0].Get<bool>());
|
||||
|
||||
AddInputVector({1, 2}, /*timestamp=*/2, &runner);
|
||||
MP_ASSERT_OK(runner.Run());
|
||||
|
||||
EXPECT_EQ(2, outputs.size());
|
||||
EXPECT_EQ(Timestamp(2), outputs[1].Timestamp());
|
||||
EXPECT_FALSE(outputs[1].Get<bool>());
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -71,16 +71,23 @@ REGISTER_CALCULATOR(DetectionLabelIdToTextCalculator);
|
||||
const auto& options =
|
||||
cc->Options<::mediapipe::DetectionLabelIdToTextCalculatorOptions>();
|
||||
|
||||
std::string string_path;
|
||||
ASSIGN_OR_RETURN(string_path, PathToResourceAsFile(options.label_map_path()));
|
||||
std::string label_map_string;
|
||||
MP_RETURN_IF_ERROR(file::GetContents(string_path, &label_map_string));
|
||||
if (options.has_label_map_path()) {
|
||||
std::string string_path;
|
||||
ASSIGN_OR_RETURN(string_path,
|
||||
PathToResourceAsFile(options.label_map_path()));
|
||||
std::string label_map_string;
|
||||
MP_RETURN_IF_ERROR(file::GetContents(string_path, &label_map_string));
|
||||
|
||||
std::istringstream stream(label_map_string);
|
||||
std::string line;
|
||||
int i = 0;
|
||||
while (std::getline(stream, line)) {
|
||||
label_map_[i++] = line;
|
||||
std::istringstream stream(label_map_string);
|
||||
std::string line;
|
||||
int i = 0;
|
||||
while (std::getline(stream, line)) {
|
||||
label_map_[i++] = line;
|
||||
}
|
||||
} else {
|
||||
for (int i = 0; i < options.label_size(); ++i) {
|
||||
label_map_[i] = options.label(i);
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -25,4 +25,10 @@ message DetectionLabelIdToTextCalculatorOptions {
|
||||
|
||||
// Path to a label map file for getting the actual name of detected classes.
|
||||
optional string label_map_path = 1;
|
||||
|
||||
// Alternative way to specify label map
|
||||
// label: "label for id 0"
|
||||
// label: "label for id 1"
|
||||
// ...
|
||||
repeated string label = 2;
|
||||
}
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/formats/classification.pb.h"
|
||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||
#include "mediapipe/framework/formats/rect.pb.h"
|
||||
|
||||
@@ -31,4 +32,8 @@ typedef FilterCollectionCalculator<
|
||||
FilterLandmarkListCollectionCalculator;
|
||||
REGISTER_CALCULATOR(FilterLandmarkListCollectionCalculator);
|
||||
|
||||
typedef FilterCollectionCalculator<std::vector<::mediapipe::ClassificationList>>
|
||||
FilterClassificationListCollectionCalculator;
|
||||
REGISTER_CALCULATOR(FilterClassificationListCollectionCalculator);
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -128,16 +128,19 @@ REGISTER_CALCULATOR(LabelsToRenderDataCalculator);
|
||||
} 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];
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag("SCORES")) {
|
||||
std::vector<float> score_vector =
|
||||
cc->Inputs().Tag("SCORES").Get<std::vector<float>>();
|
||||
CHECK_EQ(label_vector.size(), score_vector.size());
|
||||
scores.resize(label_vector.size());
|
||||
for (int i = 0; i < label_vector.size(); ++i) {
|
||||
scores[i] = score_vector[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -29,6 +29,7 @@ namespace {
|
||||
|
||||
constexpr char kLandmarksTag[] = "LANDMARKS";
|
||||
constexpr char kNormLandmarksTag[] = "NORM_LANDMARKS";
|
||||
constexpr char kRenderScaleTag[] = "RENDER_SCALE";
|
||||
constexpr char kRenderDataTag[] = "RENDER_DATA";
|
||||
constexpr char kLandmarkLabel[] = "KEYPOINT";
|
||||
constexpr int kMaxLandmarkThickness = 18;
|
||||
@@ -71,6 +72,83 @@ void SetColorSizeValueFromZ(float z, float z_min, float z_max,
|
||||
render_annotation->set_thickness(thickness);
|
||||
}
|
||||
|
||||
template <class LandmarkType>
|
||||
void AddConnectionToRenderData(const LandmarkType& start,
|
||||
const LandmarkType& end, int gray_val1,
|
||||
int gray_val2, float thickness, bool normalized,
|
||||
RenderData* render_data) {
|
||||
auto* connection_annotation = render_data->add_render_annotations();
|
||||
RenderAnnotation::GradientLine* line =
|
||||
connection_annotation->mutable_gradient_line();
|
||||
line->set_x_start(start.x());
|
||||
line->set_y_start(start.y());
|
||||
line->set_x_end(end.x());
|
||||
line->set_y_end(end.y());
|
||||
line->set_normalized(normalized);
|
||||
line->mutable_color1()->set_r(gray_val1);
|
||||
line->mutable_color1()->set_g(gray_val1);
|
||||
line->mutable_color1()->set_b(gray_val1);
|
||||
line->mutable_color2()->set_r(gray_val2);
|
||||
line->mutable_color2()->set_g(gray_val2);
|
||||
line->mutable_color2()->set_b(gray_val2);
|
||||
connection_annotation->set_thickness(thickness);
|
||||
}
|
||||
|
||||
template <class LandmarkListType, class LandmarkType>
|
||||
void AddConnectionsWithDepth(const LandmarkListType& landmarks,
|
||||
const std::vector<int>& landmark_connections,
|
||||
float thickness, bool normalized, float min_z,
|
||||
float max_z, RenderData* render_data) {
|
||||
for (int i = 0; i < landmark_connections.size(); i += 2) {
|
||||
const auto& ld0 = landmarks.landmark(landmark_connections[i]);
|
||||
const auto& ld1 = landmarks.landmark(landmark_connections[i + 1]);
|
||||
const int gray_val1 =
|
||||
255 - static_cast<int>(Remap(ld0.z(), min_z, max_z, 255));
|
||||
const int gray_val2 =
|
||||
255 - static_cast<int>(Remap(ld1.z(), min_z, max_z, 255));
|
||||
AddConnectionToRenderData<LandmarkType>(ld0, ld1, gray_val1, gray_val2,
|
||||
thickness, normalized, render_data);
|
||||
}
|
||||
}
|
||||
|
||||
template <class LandmarkType>
|
||||
void AddConnectionToRenderData(const LandmarkType& start,
|
||||
const LandmarkType& end,
|
||||
const Color& connection_color, float thickness,
|
||||
bool normalized, RenderData* render_data) {
|
||||
auto* connection_annotation = render_data->add_render_annotations();
|
||||
RenderAnnotation::Line* line = connection_annotation->mutable_line();
|
||||
line->set_x_start(start.x());
|
||||
line->set_y_start(start.y());
|
||||
line->set_x_end(end.x());
|
||||
line->set_y_end(end.y());
|
||||
line->set_normalized(normalized);
|
||||
SetColor(connection_annotation, connection_color);
|
||||
connection_annotation->set_thickness(thickness);
|
||||
}
|
||||
|
||||
template <class LandmarkListType, class LandmarkType>
|
||||
void AddConnections(const LandmarkListType& landmarks,
|
||||
const std::vector<int>& landmark_connections,
|
||||
const Color& connection_color, float thickness,
|
||||
bool normalized, RenderData* render_data) {
|
||||
for (int i = 0; i < landmark_connections.size(); i += 2) {
|
||||
const auto& ld0 = landmarks.landmark(landmark_connections[i]);
|
||||
const auto& ld1 = landmarks.landmark(landmark_connections[i + 1]);
|
||||
AddConnectionToRenderData<LandmarkType>(ld0, ld1, connection_color,
|
||||
thickness, normalized, render_data);
|
||||
}
|
||||
}
|
||||
|
||||
RenderAnnotation* AddPointRenderData(const Color& landmark_color,
|
||||
float thickness, RenderData* render_data) {
|
||||
auto* landmark_data_annotation = render_data->add_render_annotations();
|
||||
landmark_data_annotation->set_scene_tag(kLandmarkLabel);
|
||||
SetColor(landmark_data_annotation, landmark_color);
|
||||
landmark_data_annotation->set_thickness(thickness);
|
||||
return landmark_data_annotation;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// A calculator that converts Landmark proto to RenderData proto for
|
||||
@@ -107,30 +185,8 @@ class LandmarksToRenderDataCalculator : public CalculatorBase {
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
static void AddConnectionToRenderData(
|
||||
float start_x, float start_y, float end_x, float end_y,
|
||||
const LandmarksToRenderDataCalculatorOptions& options, bool normalized,
|
||||
RenderData* render_data);
|
||||
static void SetRenderAnnotationColorThickness(
|
||||
const LandmarksToRenderDataCalculatorOptions& options,
|
||||
RenderAnnotation* render_annotation);
|
||||
static RenderAnnotation* AddPointRenderData(
|
||||
const LandmarksToRenderDataCalculatorOptions& options,
|
||||
RenderData* render_data);
|
||||
static void AddConnectionToRenderData(
|
||||
float start_x, float start_y, float end_x, float end_y,
|
||||
const LandmarksToRenderDataCalculatorOptions& options, bool normalized,
|
||||
int gray_val1, int gray_val2, RenderData* render_data);
|
||||
|
||||
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);
|
||||
|
||||
LandmarksToRenderDataCalculatorOptions options_;
|
||||
std::vector<int> landmark_connections_;
|
||||
};
|
||||
REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
|
||||
@@ -150,6 +206,9 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
if (cc->Inputs().HasTag(kNormLandmarksTag)) {
|
||||
cc->Inputs().Tag(kNormLandmarksTag).Set<NormalizedLandmarkList>();
|
||||
}
|
||||
if (cc->Inputs().HasTag(kRenderScaleTag)) {
|
||||
cc->Inputs().Tag(kRenderScaleTag).Set<float>();
|
||||
}
|
||||
cc->Outputs().Tag(kRenderDataTag).Set<RenderData>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
@@ -159,6 +218,14 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
options_ = cc->Options<LandmarksToRenderDataCalculatorOptions>();
|
||||
|
||||
// Parse landmarks connections to a vector.
|
||||
RET_CHECK_EQ(options_.landmark_connections_size() % 2, 0)
|
||||
<< "Number of entries in landmark connections must be a multiple of 2";
|
||||
|
||||
for (int i = 0; i < options_.landmark_connections_size(); ++i) {
|
||||
landmark_connections_.push_back(options_.landmark_connections(i));
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
@@ -169,20 +236,36 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
float z_min = 0.f;
|
||||
float z_max = 0.f;
|
||||
|
||||
// Apply scale to `thickness` of rendered landmarks and connections to make
|
||||
// them bigger when object (e.g. pose, hand or face) is closer/bigger and
|
||||
// snaller when object is further/smaller.
|
||||
float thickness = options_.thickness();
|
||||
if (cc->Inputs().HasTag(kRenderScaleTag)) {
|
||||
const float render_scale = cc->Inputs().Tag(kRenderScaleTag).Get<float>();
|
||||
thickness *= render_scale;
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag(kLandmarksTag)) {
|
||||
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<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);
|
||||
if (visualize_depth) {
|
||||
AddConnectionsWithDepth<LandmarkList, Landmark>(
|
||||
landmarks, landmark_connections_, thickness, /*normalized=*/false,
|
||||
z_min, z_max, render_data.get());
|
||||
} else {
|
||||
AddConnections<LandmarkList, Landmark>(
|
||||
landmarks, landmark_connections_, options_.connection_color(),
|
||||
thickness, /*normalized=*/false, render_data.get());
|
||||
}
|
||||
for (int i = 0; i < landmarks.landmark_size(); ++i) {
|
||||
const Landmark& landmark = landmarks.landmark(i);
|
||||
auto* landmark_data_render =
|
||||
AddPointRenderData(options_, render_data.get());
|
||||
auto* landmark_data_render = AddPointRenderData(
|
||||
options_.landmark_color(), thickness, render_data.get());
|
||||
if (visualize_depth) {
|
||||
SetColorSizeValueFromZ(landmark.z(), z_min, z_max,
|
||||
landmark_data_render);
|
||||
@@ -192,30 +275,30 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
landmark_data->set_x(landmark.x());
|
||||
landmark_data->set_y(landmark.y());
|
||||
}
|
||||
if (visualize_depth) {
|
||||
AddConnectionsWithDepth<LandmarkList>(landmarks, /*normalized=*/false,
|
||||
z_min, z_max, render_data.get());
|
||||
} else {
|
||||
AddConnections<LandmarkList>(landmarks, /*normalized=*/false,
|
||||
render_data.get());
|
||||
}
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag(kNormLandmarksTag)) {
|
||||
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<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);
|
||||
if (visualize_depth) {
|
||||
AddConnectionsWithDepth<NormalizedLandmarkList, NormalizedLandmark>(
|
||||
landmarks, landmark_connections_, thickness, /*normalized=*/true,
|
||||
z_min, z_max, render_data.get());
|
||||
} else {
|
||||
AddConnections<NormalizedLandmarkList, NormalizedLandmark>(
|
||||
landmarks, landmark_connections_, options_.connection_color(),
|
||||
thickness, /*normalized=*/true, render_data.get());
|
||||
}
|
||||
for (int i = 0; i < landmarks.landmark_size(); ++i) {
|
||||
const NormalizedLandmark& landmark = landmarks.landmark(i);
|
||||
auto* landmark_data_render =
|
||||
AddPointRenderData(options_, render_data.get());
|
||||
auto* landmark_data_render = AddPointRenderData(
|
||||
options_.landmark_color(), thickness, render_data.get());
|
||||
if (visualize_depth) {
|
||||
SetColorSizeValueFromZ(landmark.z(), z_min, z_max,
|
||||
landmark_data_render);
|
||||
@@ -225,13 +308,6 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
landmark_data->set_x(landmark.x());
|
||||
landmark_data->set_y(landmark.y());
|
||||
}
|
||||
if (visualize_depth) {
|
||||
AddConnectionsWithDepth<NormalizedLandmarkList>(
|
||||
landmarks, /*normalized=*/true, z_min, z_max, render_data.get());
|
||||
} else {
|
||||
AddConnections<NormalizedLandmarkList>(landmarks, /*normalized=*/true,
|
||||
render_data.get());
|
||||
}
|
||||
}
|
||||
|
||||
cc->Outputs()
|
||||
@@ -240,84 +316,4 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
template <class LandmarkListType>
|
||||
void LandmarksToRenderDataCalculator::AddConnectionsWithDepth(
|
||||
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.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 =
|
||||
255 - static_cast<int>(Remap(ld1.z(), min_z, max_z, 255));
|
||||
AddConnectionToRenderData(ld0.x(), ld0.y(), ld1.x(), ld1.y(), options_,
|
||||
normalized, gray_val1, gray_val2, render_data);
|
||||
}
|
||||
}
|
||||
|
||||
void LandmarksToRenderDataCalculator::AddConnectionToRenderData(
|
||||
float start_x, float start_y, float end_x, float end_y,
|
||||
const LandmarksToRenderDataCalculatorOptions& options, bool normalized,
|
||||
int gray_val1, int gray_val2, RenderData* render_data) {
|
||||
auto* connection_annotation = render_data->add_render_annotations();
|
||||
RenderAnnotation::GradientLine* line =
|
||||
connection_annotation->mutable_gradient_line();
|
||||
line->set_x_start(start_x);
|
||||
line->set_y_start(start_y);
|
||||
line->set_x_end(end_x);
|
||||
line->set_y_end(end_y);
|
||||
line->set_normalized(normalized);
|
||||
line->mutable_color1()->set_r(gray_val1);
|
||||
line->mutable_color1()->set_g(gray_val1);
|
||||
line->mutable_color1()->set_b(gray_val1);
|
||||
line->mutable_color2()->set_r(gray_val2);
|
||||
line->mutable_color2()->set_g(gray_val2);
|
||||
line->mutable_color2()->set_b(gray_val2);
|
||||
connection_annotation->set_thickness(options.thickness());
|
||||
}
|
||||
|
||||
template <class LandmarkListType>
|
||||
void LandmarksToRenderDataCalculator::AddConnections(
|
||||
const LandmarkListType& landmarks, bool normalized,
|
||||
RenderData* render_data) {
|
||||
for (int i = 0; i < options_.landmark_connections_size(); i += 2) {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
void LandmarksToRenderDataCalculator::AddConnectionToRenderData(
|
||||
float start_x, float start_y, float end_x, float end_y,
|
||||
const LandmarksToRenderDataCalculatorOptions& options, bool normalized,
|
||||
RenderData* render_data) {
|
||||
auto* connection_annotation = render_data->add_render_annotations();
|
||||
RenderAnnotation::Line* line = connection_annotation->mutable_line();
|
||||
line->set_x_start(start_x);
|
||||
line->set_y_start(start_y);
|
||||
line->set_x_end(end_x);
|
||||
line->set_y_end(end_y);
|
||||
line->set_normalized(normalized);
|
||||
SetColor(connection_annotation, options.connection_color());
|
||||
connection_annotation->set_thickness(options.thickness());
|
||||
}
|
||||
|
||||
RenderAnnotation* LandmarksToRenderDataCalculator::AddPointRenderData(
|
||||
const LandmarksToRenderDataCalculatorOptions& options,
|
||||
RenderData* render_data) {
|
||||
auto* landmark_data_annotation = render_data->add_render_annotations();
|
||||
landmark_data_annotation->set_scene_tag(kLandmarkLabel);
|
||||
SetRenderAnnotationColorThickness(options, landmark_data_annotation);
|
||||
return landmark_data_annotation;
|
||||
}
|
||||
|
||||
void LandmarksToRenderDataCalculator::SetRenderAnnotationColorThickness(
|
||||
const LandmarksToRenderDataCalculatorOptions& options,
|
||||
RenderAnnotation* render_annotation) {
|
||||
SetColor(render_annotation, options.landmark_color());
|
||||
render_annotation->set_thickness(options.thickness());
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
@@ -16,34 +16,80 @@
|
||||
#include <string>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/file_helpers.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr char kFilePathTag[] = "FILE_PATH";
|
||||
constexpr char kContentsTag[] = "CONTENTS";
|
||||
|
||||
} // namespace
|
||||
|
||||
// The calculator takes the path to the local file as an input side packet and
|
||||
// outputs the contents of that file.
|
||||
//
|
||||
// NOTE: file loading can be batched by providing multiple input/output side
|
||||
// packets.
|
||||
//
|
||||
// Example config:
|
||||
// node {
|
||||
// calculator: "LocalFileContentsCalculator"
|
||||
// input_side_packet: "FILE_PATH:file_path"
|
||||
// output_side_packet: "CONTENTS:contents"
|
||||
// }
|
||||
//
|
||||
// node {
|
||||
// calculator: "LocalFileContentsCalculator"
|
||||
// input_side_packet: "FILE_PATH:0:file_path1"
|
||||
// input_side_packet: "FILE_PATH:1:file_path2"
|
||||
// ...
|
||||
// output_side_packet: "CONTENTS:0:contents1"
|
||||
// output_side_packet: "CONTENTS:1:contents2"
|
||||
// ...
|
||||
// }
|
||||
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>();
|
||||
RET_CHECK(cc->InputSidePackets().HasTag(kFilePathTag))
|
||||
<< "Missing PATH input side packet(s)";
|
||||
RET_CHECK(cc->OutputSidePackets().HasTag(kContentsTag))
|
||||
<< "Missing CONTENTS output side packet(s)";
|
||||
|
||||
RET_CHECK_EQ(cc->InputSidePackets().NumEntries(kFilePathTag),
|
||||
cc->OutputSidePackets().NumEntries(kContentsTag))
|
||||
<< "Same number of input streams and output streams is required.";
|
||||
|
||||
for (CollectionItemId id = cc->InputSidePackets().BeginId(kFilePathTag);
|
||||
id != cc->InputSidePackets().EndId(kFilePathTag); ++id) {
|
||||
cc->InputSidePackets().Get(id).Set<std::string>();
|
||||
}
|
||||
|
||||
for (CollectionItemId id = cc->OutputSidePackets().BeginId(kContentsTag);
|
||||
id != cc->OutputSidePackets().EndId(kContentsTag); ++id) {
|
||||
cc->OutputSidePackets().Get(id).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)));
|
||||
CollectionItemId input_id = cc->InputSidePackets().BeginId(kFilePathTag);
|
||||
CollectionItemId output_id = cc->OutputSidePackets().BeginId(kContentsTag);
|
||||
// Number of inputs and outpus is the same according to the contract.
|
||||
for (; input_id != cc->InputSidePackets().EndId(kFilePathTag);
|
||||
++input_id, ++output_id) {
|
||||
std::string file_path =
|
||||
cc->InputSidePackets().Get(input_id).Get<std::string>();
|
||||
ASSIGN_OR_RETURN(file_path, PathToResourceAsFile(file_path));
|
||||
|
||||
std::string contents;
|
||||
MP_RETURN_IF_ERROR(GetResourceContents(file_path, &contents));
|
||||
cc->OutputSidePackets().Get(output_id).Set(
|
||||
MakePacket<std::string>(std::move(contents)));
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
|
||||
@@ -276,6 +276,7 @@ TEST_F(PacketLatencyCalculatorTest, DoesNotOutputUntilReferencePacketReceived) {
|
||||
"delayed_packet_0", Adopt(new double()).At(Timestamp(2))));
|
||||
|
||||
// Send a reference packet with timestamp 10 usec.
|
||||
simulation_clock_->Sleep(absl::Microseconds(1));
|
||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||
"camera_frames", Adopt(new double()).At(Timestamp(10))));
|
||||
simulation_clock_->Sleep(absl::Microseconds(1));
|
||||
|
||||
@@ -12,6 +12,8 @@
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "absl/strings/str_join.h"
|
||||
@@ -76,14 +78,16 @@ void AddTimedBoxProtoToRenderData(
|
||||
RenderAnnotation::Text* text = label_annotation->mutable_text();
|
||||
text->set_display_text(box_proto.label());
|
||||
text->set_normalized(true);
|
||||
constexpr float text_left_start = 0.3f;
|
||||
constexpr float text_left_start = 0.2f;
|
||||
text->set_left((1.0f - text_left_start) * box_proto.left() +
|
||||
text_left_start * box_proto.right());
|
||||
constexpr float text_baseline = 0.6f;
|
||||
text->set_baseline(text_baseline * box_proto.bottom() +
|
||||
(1.0f - text_baseline) * box_proto.top());
|
||||
constexpr float text_height = 0.2f;
|
||||
text->set_font_height((box_proto.bottom() - box_proto.top()) * text_height);
|
||||
constexpr float text_height = 0.1f;
|
||||
text->set_font_height(std::min(box_proto.bottom() - box_proto.top(),
|
||||
box_proto.right() - box_proto.left()) *
|
||||
text_height);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -65,6 +65,26 @@ proto_library(
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "tracked_detection_manager_calculator_proto",
|
||||
srcs = ["tracked_detection_manager_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/util/tracking:tracked_detection_manager_config_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "box_detector_calculator_proto",
|
||||
srcs = ["box_detector_calculator.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_proto",
|
||||
"//mediapipe/util/tracking:box_detector_proto",
|
||||
],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "video_pre_stream_calculator_proto",
|
||||
srcs = ["video_pre_stream_calculator.proto"],
|
||||
@@ -107,6 +127,28 @@ mediapipe_cc_proto_library(
|
||||
deps = [":box_tracker_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "tracked_detection_manager_calculator_cc_proto",
|
||||
srcs = ["tracked_detection_manager_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/util/tracking:tracked_detection_manager_config_cc_proto",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":tracked_detection_manager_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "box_detector_calculator_cc_proto",
|
||||
srcs = ["box_detector_calculator.proto"],
|
||||
cc_deps = [
|
||||
"//mediapipe/framework:calculator_cc_proto",
|
||||
"//mediapipe/util/tracking:box_detector_cc_proto",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [":box_detector_calculator_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "video_pre_stream_calculator_cc_proto",
|
||||
srcs = ["video_pre_stream_calculator.proto"],
|
||||
@@ -138,7 +180,7 @@ cc_library(
|
||||
srcs = ["flow_to_image_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/calculators/video:flow_to_image_calculator_cc_proto",
|
||||
":flow_to_image_calculator_cc_proto",
|
||||
"//mediapipe/calculators/video/tool:flow_quantizer_model",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:image_format_cc_proto",
|
||||
@@ -279,11 +321,54 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "box_detector_calculator",
|
||||
srcs = ["box_detector_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":box_detector_calculator_cc_proto",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/strings",
|
||||
"//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:opencv_core",
|
||||
"//mediapipe/framework/port:opencv_features2d",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/util:resource_util",
|
||||
"//mediapipe/util/tracking",
|
||||
"//mediapipe/util/tracking:box_detector",
|
||||
"//mediapipe/util/tracking:box_tracker",
|
||||
"//mediapipe/util/tracking:box_tracker_cc_proto",
|
||||
"//mediapipe/util/tracking:flow_packager_cc_proto",
|
||||
"//mediapipe/util/tracking:tracking_visualization_utilities",
|
||||
] + select({
|
||||
"//mediapipe:android": [
|
||||
"//mediapipe/util/android/file/base",
|
||||
],
|
||||
"//mediapipe:ios": [
|
||||
"//mediapipe/util/android/file/base",
|
||||
],
|
||||
"//mediapipe:macos": [
|
||||
"//mediapipe/framework/port:file_helpers",
|
||||
],
|
||||
"//conditions:default": [
|
||||
"//mediapipe/framework/port:file_helpers",
|
||||
],
|
||||
}),
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "tracked_detection_manager_calculator",
|
||||
srcs = ["tracked_detection_manager_calculator.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":tracked_detection_manager_calculator_cc_proto",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/formats:detection_cc_proto",
|
||||
"//mediapipe/framework/formats:location_data_cc_proto",
|
||||
@@ -384,20 +469,18 @@ cc_test(
|
||||
],
|
||||
)
|
||||
|
||||
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,
|
||||
deps = [
|
||||
":box_tracker_calculator",
|
||||
":flow_packager_calculator",
|
||||
":motion_analysis_calculator",
|
||||
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
|
||||
"//mediapipe/framework/stream_handler:sync_set_input_stream_handler",
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_binary_graph(
|
||||
@@ -405,7 +488,13 @@ mediapipe_binary_graph(
|
||||
graph = "testdata/tracker_graph.pbtxt",
|
||||
output_name = "testdata/tracker.binarypb",
|
||||
visibility = ["//visibility:public"],
|
||||
deps = MEDIAPIPE_DEPS,
|
||||
deps = [
|
||||
":box_tracker_calculator",
|
||||
":flow_packager_calculator",
|
||||
":motion_analysis_calculator",
|
||||
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
|
||||
"//mediapipe/framework/stream_handler:sync_set_input_stream_handler",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
|
||||
@@ -0,0 +1,393 @@
|
||||
// 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 <stdio.h>
|
||||
|
||||
#include <memory>
|
||||
#include <unordered_set>
|
||||
|
||||
#include "absl/memory/memory.h"
|
||||
#include "absl/strings/numbers.h"
|
||||
#include "mediapipe/calculators/video/box_detector_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/image_frame.h"
|
||||
#include "mediapipe/framework/formats/image_frame_opencv.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/opencv_core_inc.h"
|
||||
#include "mediapipe/framework/port/opencv_features2d_inc.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/util/resource_util.h"
|
||||
#include "mediapipe/util/tracking/box_detector.h"
|
||||
#include "mediapipe/util/tracking/box_tracker.h"
|
||||
#include "mediapipe/util/tracking/box_tracker.pb.h"
|
||||
#include "mediapipe/util/tracking/flow_packager.pb.h"
|
||||
#include "mediapipe/util/tracking/tracking.h"
|
||||
#include "mediapipe/util/tracking/tracking_visualization_utilities.h"
|
||||
|
||||
#if defined(MEDIAPIPE_MOBILE)
|
||||
#include "mediapipe/util/android/file/base/file.h"
|
||||
#include "mediapipe/util/android/file/base/helpers.h"
|
||||
#else
|
||||
#include "mediapipe/framework/port/file_helpers.h"
|
||||
#endif
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// A calculator to detect reappeared box positions from single frame.
|
||||
//
|
||||
// Input stream:
|
||||
// TRACKING: Input tracking data (proto TrackingData) containing features and
|
||||
// descriptors.
|
||||
// VIDEO: Optional input video stream tracked boxes are rendered over
|
||||
// (Required if VIZ is specified).
|
||||
// FEATURES: Input feature points (std::vector<cv::KeyPoint>) in the original
|
||||
// pixel space.
|
||||
// DESCRIPTORS: Input feature descriptors (std::vector<float>). Actual feature
|
||||
// dimension needs to be specified in detector_options.
|
||||
// IMAGE_SIZE: Input image dimension.
|
||||
// TRACKED_BOXES : input box tracking result (proto TimedBoxProtoList) from
|
||||
// BoxTrackerCalculator.
|
||||
// ADD_INDEX: Optional std::string containing binary format proto of type
|
||||
// BoxDetectorIndex. Used for adding target index to the detector
|
||||
// search index during runtime.
|
||||
// CANCEL_OBJECT_ID: Optional id of box to be removed. This is recommended
|
||||
// to be used with SyncSetInputStreamHandler.
|
||||
// REACQ_SWITCH: Optional bool for swithcing on and off reacquisition
|
||||
// functionality. User should initialize a graph with box detector
|
||||
// calculator and be able to switch it on and off in runtime.
|
||||
//
|
||||
// Output streams:
|
||||
// VIZ: Optional output video stream with rendered box positions
|
||||
// (requires VIDEO to be present)
|
||||
// BOXES: Optional output stream of type TimedBoxProtoList for each lost box.
|
||||
//
|
||||
// Imput side packets:
|
||||
// INDEX_PROTO_STRING: Optional std::string containing binary format proto of
|
||||
// type
|
||||
// BoxDetectorIndex. Used for initializing box_detector
|
||||
// with predefined template images.
|
||||
// FRAME_ALIGNMENT: Optional integer to indicate alignment_boundary for
|
||||
// outputing ImageFrame in "VIZ" stream.
|
||||
// Set to ImageFrame::kDefaultAlignmentBoundary for
|
||||
// offline pipeline to be compatible with FFmpeg.
|
||||
// Set to ImageFrame::kGlDefaultAlignmentBoundary for Apps
|
||||
// to be compatible with GL renderer.
|
||||
// OUTPUT_INDEX_FILENAME: File path to the output index file.
|
||||
|
||||
class BoxDetectorCalculator : public CalculatorBase {
|
||||
public:
|
||||
~BoxDetectorCalculator() override = default;
|
||||
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||
|
||||
private:
|
||||
BoxDetectorCalculatorOptions options_;
|
||||
std::unique_ptr<BoxDetectorInterface> box_detector_;
|
||||
bool detector_switch_ = true;
|
||||
uint32 frame_alignment_ = ImageFrame::kDefaultAlignmentBoundary;
|
||||
bool write_index_ = false;
|
||||
int box_id_ = 0;
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(BoxDetectorCalculator);
|
||||
|
||||
::mediapipe::Status BoxDetectorCalculator::GetContract(CalculatorContract* cc) {
|
||||
if (cc->Inputs().HasTag("TRACKING")) {
|
||||
cc->Inputs().Tag("TRACKING").Set<TrackingData>();
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag("TRACKED_BOXES")) {
|
||||
cc->Inputs().Tag("TRACKED_BOXES").Set<TimedBoxProtoList>();
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag("VIDEO")) {
|
||||
cc->Inputs().Tag("VIDEO").Set<ImageFrame>();
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag("FEATURES")) {
|
||||
RET_CHECK(cc->Inputs().HasTag("DESCRIPTORS"))
|
||||
<< "FEATURES and DESCRIPTORS need to be specified together.";
|
||||
cc->Inputs().Tag("FEATURES").Set<std::vector<cv::KeyPoint>>();
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag("DESCRIPTORS")) {
|
||||
RET_CHECK(cc->Inputs().HasTag("FEATURES"))
|
||||
<< "FEATURES and DESCRIPTORS need to be specified together.";
|
||||
cc->Inputs().Tag("DESCRIPTORS").Set<std::vector<float>>();
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag("IMAGE_SIZE")) {
|
||||
cc->Inputs().Tag("IMAGE_SIZE").Set<std::pair<int, int>>();
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag("ADD_INDEX")) {
|
||||
cc->Inputs().Tag("ADD_INDEX").Set<std::string>();
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag("CANCEL_OBJECT_ID")) {
|
||||
cc->Inputs().Tag("CANCEL_OBJECT_ID").Set<int>();
|
||||
}
|
||||
|
||||
if (cc->Inputs().HasTag("REACQ_SWITCH")) {
|
||||
cc->Inputs().Tag("REACQ_SWITCH").Set<bool>();
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("BOXES")) {
|
||||
cc->Outputs().Tag("BOXES").Set<TimedBoxProtoList>();
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("VIZ")) {
|
||||
RET_CHECK(cc->Inputs().HasTag("VIDEO"))
|
||||
<< "Output stream VIZ requires VIDEO to be present.";
|
||||
cc->Outputs().Tag("VIZ").Set<ImageFrame>();
|
||||
}
|
||||
|
||||
if (cc->InputSidePackets().HasTag("INDEX_PROTO_STRING")) {
|
||||
cc->InputSidePackets().Tag("INDEX_PROTO_STRING").Set<std::string>();
|
||||
}
|
||||
|
||||
if (cc->InputSidePackets().HasTag("OUTPUT_INDEX_FILENAME")) {
|
||||
cc->InputSidePackets().Tag("OUTPUT_INDEX_FILENAME").Set<std::string>();
|
||||
}
|
||||
|
||||
if (cc->InputSidePackets().HasTag("FRAME_ALIGNMENT")) {
|
||||
cc->InputSidePackets().Tag("FRAME_ALIGNMENT").Set<int>();
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status BoxDetectorCalculator::Open(CalculatorContext* cc) {
|
||||
options_ = cc->Options<BoxDetectorCalculatorOptions>();
|
||||
box_detector_ = BoxDetectorInterface::Create(options_.detector_options());
|
||||
|
||||
if (cc->InputSidePackets().HasTag("INDEX_PROTO_STRING")) {
|
||||
BoxDetectorIndex predefined_index;
|
||||
if (!predefined_index.ParseFromString(cc->InputSidePackets()
|
||||
.Tag("INDEX_PROTO_STRING")
|
||||
.Get<std::string>())) {
|
||||
LOG(FATAL) << "failed to parse BoxDetectorIndex from INDEX_PROTO_STRING";
|
||||
}
|
||||
box_detector_->AddBoxDetectorIndex(predefined_index);
|
||||
}
|
||||
|
||||
for (const auto& filename : options_.index_proto_filename()) {
|
||||
std::string string_path;
|
||||
ASSIGN_OR_RETURN(string_path, PathToResourceAsFile(filename));
|
||||
std::string index_string;
|
||||
MP_RETURN_IF_ERROR(file::GetContents(string_path, &index_string));
|
||||
BoxDetectorIndex predefined_index;
|
||||
if (!predefined_index.ParseFromString(index_string)) {
|
||||
LOG(FATAL)
|
||||
<< "failed to parse BoxDetectorIndex from index_proto_filename";
|
||||
}
|
||||
box_detector_->AddBoxDetectorIndex(predefined_index);
|
||||
}
|
||||
|
||||
if (cc->InputSidePackets().HasTag("OUTPUT_INDEX_FILENAME")) {
|
||||
write_index_ = true;
|
||||
}
|
||||
|
||||
if (cc->InputSidePackets().HasTag("FRAME_ALIGNMENT")) {
|
||||
frame_alignment_ = cc->InputSidePackets().Tag("FRAME_ALIGNMENT").Get<int>();
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status BoxDetectorCalculator::Process(CalculatorContext* cc) {
|
||||
const Timestamp timestamp = cc->InputTimestamp();
|
||||
const int64 timestamp_msec = timestamp.Value() / 1000;
|
||||
|
||||
InputStream* cancel_object_id_stream =
|
||||
cc->Inputs().HasTag("CANCEL_OBJECT_ID")
|
||||
? &(cc->Inputs().Tag("CANCEL_OBJECT_ID"))
|
||||
: nullptr;
|
||||
if (cancel_object_id_stream && !cancel_object_id_stream->IsEmpty()) {
|
||||
const int cancel_object_id = cancel_object_id_stream->Get<int>();
|
||||
box_detector_->CancelBoxDetection(cancel_object_id);
|
||||
}
|
||||
|
||||
InputStream* add_index_stream = cc->Inputs().HasTag("ADD_INDEX")
|
||||
? &(cc->Inputs().Tag("ADD_INDEX"))
|
||||
: nullptr;
|
||||
if (add_index_stream && !add_index_stream->IsEmpty()) {
|
||||
BoxDetectorIndex predefined_index;
|
||||
if (!predefined_index.ParseFromString(
|
||||
add_index_stream->Get<std::string>())) {
|
||||
LOG(FATAL) << "failed to parse BoxDetectorIndex from ADD_INDEX";
|
||||
}
|
||||
box_detector_->AddBoxDetectorIndex(predefined_index);
|
||||
}
|
||||
|
||||
InputStream* reacq_switch_stream = cc->Inputs().HasTag("REACQ_SWITCH")
|
||||
? &(cc->Inputs().Tag("REACQ_SWITCH"))
|
||||
: nullptr;
|
||||
if (reacq_switch_stream && !reacq_switch_stream->IsEmpty()) {
|
||||
detector_switch_ = reacq_switch_stream->Get<bool>();
|
||||
}
|
||||
|
||||
if (!detector_switch_) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
InputStream* track_stream = cc->Inputs().HasTag("TRACKING")
|
||||
? &(cc->Inputs().Tag("TRACKING"))
|
||||
: nullptr;
|
||||
InputStream* video_stream =
|
||||
cc->Inputs().HasTag("VIDEO") ? &(cc->Inputs().Tag("VIDEO")) : nullptr;
|
||||
InputStream* feature_stream = cc->Inputs().HasTag("FEATURES")
|
||||
? &(cc->Inputs().Tag("FEATURES"))
|
||||
: nullptr;
|
||||
InputStream* descriptor_stream = cc->Inputs().HasTag("DESCRIPTORS")
|
||||
? &(cc->Inputs().Tag("DESCRIPTORS"))
|
||||
: nullptr;
|
||||
|
||||
CHECK(track_stream != nullptr || video_stream != nullptr ||
|
||||
(feature_stream != nullptr && descriptor_stream != nullptr))
|
||||
<< "One and only one of {tracking_data, input image frame, "
|
||||
"feature/descriptor} need to be valid.";
|
||||
|
||||
InputStream* tracked_boxes_stream = cc->Inputs().HasTag("TRACKED_BOXES")
|
||||
? &(cc->Inputs().Tag("TRACKED_BOXES"))
|
||||
: nullptr;
|
||||
std::unique_ptr<TimedBoxProtoList> detected_boxes(new TimedBoxProtoList());
|
||||
|
||||
if (track_stream != nullptr) {
|
||||
// Detect from tracking data
|
||||
if (track_stream->IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
const TrackingData& tracking_data = track_stream->Get<TrackingData>();
|
||||
|
||||
CHECK(tracked_boxes_stream != nullptr) << "tracked_boxes needed.";
|
||||
|
||||
const TimedBoxProtoList tracked_boxes =
|
||||
tracked_boxes_stream->Get<TimedBoxProtoList>();
|
||||
|
||||
box_detector_->DetectAndAddBox(tracking_data, tracked_boxes, timestamp_msec,
|
||||
detected_boxes.get());
|
||||
} else if (video_stream != nullptr) {
|
||||
// Detect from input frame
|
||||
if (video_stream->IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
TimedBoxProtoList tracked_boxes;
|
||||
if (tracked_boxes_stream != nullptr && !tracked_boxes_stream->IsEmpty()) {
|
||||
tracked_boxes = tracked_boxes_stream->Get<TimedBoxProtoList>();
|
||||
}
|
||||
|
||||
// Just directly pass along the image frame data as-is for detection; we
|
||||
// don't need to worry about conforming to a specific alignment here.
|
||||
const cv::Mat input_view =
|
||||
formats::MatView(&video_stream->Get<ImageFrame>());
|
||||
box_detector_->DetectAndAddBox(input_view, tracked_boxes, timestamp_msec,
|
||||
detected_boxes.get());
|
||||
} else {
|
||||
if (feature_stream->IsEmpty() || descriptor_stream->IsEmpty()) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
const auto& image_size =
|
||||
cc->Inputs().Tag("IMAGE_SIZE").Get<std::pair<int, int>>();
|
||||
float inv_scale = 1.0f / std::max(image_size.first, image_size.second);
|
||||
|
||||
TimedBoxProtoList tracked_boxes;
|
||||
if (tracked_boxes_stream != nullptr && !tracked_boxes_stream->IsEmpty()) {
|
||||
tracked_boxes = tracked_boxes_stream->Get<TimedBoxProtoList>();
|
||||
} else if (write_index_) {
|
||||
auto* box_ptr = tracked_boxes.add_box();
|
||||
box_ptr->set_id(box_id_);
|
||||
box_ptr->set_reacquisition(true);
|
||||
box_ptr->set_aspect_ratio((float)image_size.first /
|
||||
(float)image_size.second);
|
||||
|
||||
box_ptr->mutable_quad()->add_vertices(0);
|
||||
box_ptr->mutable_quad()->add_vertices(0);
|
||||
|
||||
box_ptr->mutable_quad()->add_vertices(0);
|
||||
box_ptr->mutable_quad()->add_vertices(1);
|
||||
|
||||
box_ptr->mutable_quad()->add_vertices(1);
|
||||
box_ptr->mutable_quad()->add_vertices(1);
|
||||
|
||||
box_ptr->mutable_quad()->add_vertices(1);
|
||||
box_ptr->mutable_quad()->add_vertices(0);
|
||||
|
||||
++box_id_;
|
||||
}
|
||||
|
||||
const auto& features = feature_stream->Get<std::vector<cv::KeyPoint>>();
|
||||
const int feature_size = features.size();
|
||||
std::vector<Vector2_f> features_vec(feature_size);
|
||||
|
||||
const auto& descriptors = descriptor_stream->Get<std::vector<float>>();
|
||||
const int dims = options_.detector_options().descriptor_dims();
|
||||
CHECK_GE(descriptors.size(), feature_size * dims);
|
||||
cv::Mat descriptors_mat(feature_size, dims, CV_32F);
|
||||
for (int j = 0; j < feature_size; ++j) {
|
||||
features_vec[j].Set(features[j].pt.x * inv_scale,
|
||||
features[j].pt.y * inv_scale);
|
||||
for (int i = 0; i < dims; ++i) {
|
||||
descriptors_mat.at<float>(j, i) = descriptors[j * dims + i];
|
||||
}
|
||||
}
|
||||
|
||||
box_detector_->DetectAndAddBoxFromFeatures(
|
||||
features_vec, descriptors_mat, tracked_boxes, timestamp_msec,
|
||||
image_size.first * inv_scale, image_size.second * inv_scale,
|
||||
detected_boxes.get());
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("VIZ")) {
|
||||
cv::Mat viz_view;
|
||||
std::unique_ptr<ImageFrame> viz_frame;
|
||||
if (video_stream != nullptr && !video_stream->IsEmpty()) {
|
||||
viz_frame = absl::make_unique<ImageFrame>();
|
||||
viz_frame->CopyFrom(video_stream->Get<ImageFrame>(), frame_alignment_);
|
||||
viz_view = formats::MatView(viz_frame.get());
|
||||
}
|
||||
for (const auto& box : detected_boxes->box()) {
|
||||
RenderBox(box, &viz_view);
|
||||
}
|
||||
cc->Outputs().Tag("VIZ").Add(viz_frame.release(), timestamp);
|
||||
}
|
||||
|
||||
if (cc->Outputs().HasTag("BOXES")) {
|
||||
cc->Outputs().Tag("BOXES").Add(detected_boxes.release(), timestamp);
|
||||
}
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status BoxDetectorCalculator::Close(CalculatorContext* cc) {
|
||||
if (write_index_) {
|
||||
BoxDetectorIndex index = box_detector_->ObtainBoxDetectorIndex();
|
||||
MEDIAPIPE_CHECK_OK(mediapipe::file::SetContents(
|
||||
cc->InputSidePackets().Tag("OUTPUT_INDEX_FILENAME").Get<std::string>(),
|
||||
index.SerializeAsString()));
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,31 @@
|
||||
// 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/tracking/box_detector.proto";
|
||||
|
||||
message BoxDetectorCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional BoxDetectorCalculatorOptions ext = 289746530;
|
||||
}
|
||||
|
||||
optional BoxDetectorOptions detector_options = 1;
|
||||
|
||||
// File path to the template index files.
|
||||
repeated string index_proto_filename = 2;
|
||||
}
|
||||
@@ -18,6 +18,7 @@
|
||||
#include <vector>
|
||||
|
||||
#include "absl/container/node_hash_map.h"
|
||||
#include "mediapipe/calculators/video/tracked_detection_manager_calculator.pb.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/formats/detection.pb.h"
|
||||
#include "mediapipe/framework/formats/location_data.pb.h"
|
||||
@@ -139,6 +140,7 @@ Detection GetAxisAlignedDetectionFromTrackedDetection(
|
||||
class TrackedDetectionManagerCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||
|
||||
@@ -184,6 +186,15 @@ REGISTER_CALCULATOR(TrackedDetectionManagerCalculator);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status TrackedDetectionManagerCalculator::Open(
|
||||
CalculatorContext* cc) {
|
||||
mediapipe::TrackedDetectionManagerCalculatorOptions options =
|
||||
cc->Options<mediapipe::TrackedDetectionManagerCalculatorOptions>();
|
||||
tracked_detection_manager_.SetConfig(
|
||||
options.tracked_detection_manager_options());
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status TrackedDetectionManagerCalculator::Process(
|
||||
CalculatorContext* cc) {
|
||||
if (cc->Inputs().HasTag("TRACKING_BOXES")) {
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright 2020 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/tracking/tracked_detection_manager_config.proto";
|
||||
|
||||
message TrackedDetectionManagerCalculatorOptions {
|
||||
extend CalculatorOptions {
|
||||
optional TrackedDetectionManagerCalculatorOptions ext = 301970230;
|
||||
}
|
||||
|
||||
optional TrackedDetectionManagerConfig tracked_detection_manager_options = 1;
|
||||
}
|
||||
@@ -32,13 +32,14 @@ project.
|
||||
2. Run the Bazel build command to generate the AAR.
|
||||
|
||||
```bash
|
||||
bazel build -c opt --fat_apk_cpu=arm64-v8a,armeabi-v7a //path/to/the/aar/build/file:aar_name
|
||||
bazel build -c opt --host_crosstool_top=@bazel_tools//tools/cpp:toolchain --fat_apk_cpu=arm64-v8a,armeabi-v7a \
|
||||
//path/to/the/aar/build/file:aar_name
|
||||
```
|
||||
|
||||
For the face detection AAR target we made in the step 1, run:
|
||||
|
||||
```bash
|
||||
bazel build -c opt --fat_apk_cpu=arm64-v8a,armeabi-v7a \
|
||||
bazel build -c opt --host_crosstool_top=@bazel_tools//tools/cpp:toolchain --fat_apk_cpu=arm64-v8a,armeabi-v7a \
|
||||
//mediapipe/examples/android/src/java/com/google/mediapipe/apps/aar_example:mp_face_detection_aar
|
||||
|
||||
# It should print:
|
||||
|
||||
@@ -73,6 +73,18 @@ can be easily adapted to run on CPU v.s. GPU.
|
||||
* [Android](./face_detection_mobile_cpu.md)
|
||||
* [iOS](./face_detection_mobile_cpu.md)
|
||||
|
||||
### Face Mesh with GPU
|
||||
|
||||
[Face Mesh with GPU](./face_mesh_mobile_gpu.md) illustrates how to run the
|
||||
MediaPipe Face Mesh pipeline to perform 3D face landmark estimation in real-time
|
||||
on mobile devices, utilizing GPU acceleration. The pipeline is based on
|
||||
["Real-time Facial Surface Geometry from Monocular Video on Mobile GPUs"](https://arxiv.org/abs/1907.06724),
|
||||
and details of the underlying ML models are described in the
|
||||
[model card](https://drive.google.com/file/d/1VFC_wIpw4O7xBOiTgUldl79d9LA-LsnA/view).
|
||||
|
||||
* [Android](./face_mesh_mobile_gpu.md)
|
||||
* [iOS](./face_mesh_mobile_gpu.md)
|
||||
|
||||
### Hand Detection with GPU
|
||||
|
||||
[Hand Detection with GPU](./hand_detection_mobile_gpu.md) illustrates how to use
|
||||
@@ -84,7 +96,7 @@ MediaPipe with a TFLite model for hand detection in a GPU-accelerated pipeline.
|
||||
### Hand Tracking with GPU
|
||||
|
||||
[Hand Tracking with GPU](./hand_tracking_mobile_gpu.md) illustrates how to use
|
||||
MediaPipe with a TFLite model for hand tracking in a GPU-accelerated pipeline.
|
||||
MediaPipe with TFLite models for hand tracking in a GPU-accelerated pipeline.
|
||||
|
||||
* [Android](./hand_tracking_mobile_gpu.md)
|
||||
* [iOS](./hand_tracking_mobile_gpu.md)
|
||||
@@ -92,7 +104,7 @@ MediaPipe with a TFLite model for hand tracking in a GPU-accelerated pipeline.
|
||||
### Multi-Hand Tracking with GPU
|
||||
|
||||
[Multi-Hand Tracking with GPU](./multi_hand_tracking_mobile_gpu.md) illustrates
|
||||
how to use MediaPipe with a TFLite model for multi-hand tracking in a
|
||||
how to use MediaPipe with TFLite models for multi-hand tracking in a
|
||||
GPU-accelerated pipeline.
|
||||
|
||||
* [Android](./multi_hand_tracking_mobile_gpu.md)
|
||||
@@ -109,6 +121,14 @@ and model details are described in the
|
||||
|
||||
* [Android](./hair_segmentation_mobile_gpu.md)
|
||||
|
||||
### Template Matching using KNIFT with CPU
|
||||
|
||||
[Template Matching using KNIFT on Mobile](./template_matching_mobile_cpu.md)
|
||||
shows how to use MediaPipe with TFLite model for template matching using Knift
|
||||
on mobile using CPU.
|
||||
|
||||
* [Android](./template_matching_mobile_cpu.md)
|
||||
|
||||
## Desktop
|
||||
|
||||
### Hello World for C++
|
||||
@@ -150,11 +170,19 @@ GPU with live video from a webcam.
|
||||
* [Desktop GPU](./face_detection_desktop.md)
|
||||
* [Desktop CPU](./face_detection_desktop.md)
|
||||
|
||||
### Face Mesh on Desktop with Webcam
|
||||
|
||||
[Face Mesh on Desktop with Webcam](./face_mesh_desktop.md) shows how to run the
|
||||
MediaPipe Face Mesh pipeline to perform 3D face landmark estimation in real-time
|
||||
on desktop with webcam input.
|
||||
|
||||
* [Desktop GPU](./face_mesh_desktop.md)
|
||||
* [Desktop CPU](./face_mesh_desktop.md)
|
||||
|
||||
### Hand Tracking on Desktop with Webcam
|
||||
|
||||
[Hand Tracking on Desktop with Webcam](./hand_tracking_desktop.md) shows how to
|
||||
use MediaPipe with a TFLite model for hand tracking on desktop using CPU or GPU
|
||||
use MediaPipe with TFLite models for hand tracking on desktop using CPU or GPU
|
||||
with live video from a webcam.
|
||||
|
||||
* [Desktop GPU](./hand_tracking_desktop.md)
|
||||
@@ -163,8 +191,8 @@ with live video from a webcam.
|
||||
### Multi-Hand Tracking on Desktop with Webcam
|
||||
|
||||
[Multi-Hand Tracking on Desktop with Webcam](./multi_hand_tracking_desktop.md)
|
||||
shows how to use MediaPipe with a TFLite model for multi-hand tracking on
|
||||
desktop using CPU or GPU with live video from a webcam.
|
||||
shows how to use MediaPipe with TFLite models for multi-hand tracking on desktop
|
||||
using CPU or GPU with live video from a webcam.
|
||||
|
||||
* [Desktop GPU](./multi_hand_tracking_desktop.md)
|
||||
* [Desktop CPU](./multi_hand_tracking_desktop.md)
|
||||
@@ -177,7 +205,7 @@ GPU with live video from a webcam.
|
||||
|
||||
* [Desktop GPU](./hair_segmentation_desktop.md)
|
||||
|
||||
## Google Coral (machine learning acceleration with Google EdgeTPU)
|
||||
## Google Coral (ML acceleration with Google EdgeTPU)
|
||||
|
||||
Below are code samples on how to run MediaPipe on Google Coral Dev Board.
|
||||
|
||||
|
||||
@@ -51,7 +51,7 @@ To build and run the TensorFlow Lite example on desktop (GPU) with Webcam, run:
|
||||
```bash
|
||||
# Video from webcam running on desktop GPU
|
||||
# This works only for linux currently
|
||||
$ bazel build -c opt --copt -DMESA_EGL_NO_X11_HEADERS \
|
||||
$ bazel build -c opt --copt -DMESA_EGL_NO_X11_HEADERS --copt -DEGL_NO_X11 \
|
||||
mediapipe/examples/desktop/face_detection:face_detection_gpu
|
||||
|
||||
# It should print:
|
||||
|
||||
@@ -4,6 +4,8 @@ This doc focuses on the
|
||||
[example graph](https://github.com/google/mediapipe/tree/master/mediapipe/graphs/face_detection/face_detection_mobile_cpu.pbtxt)
|
||||
that performs face detection with TensorFlow Lite on CPU.
|
||||
|
||||

|
||||
|
||||
## Android
|
||||
|
||||
[Source](https://github.com/google/mediapipe/tree/master/mediapipe/examples/android/src/java/com/google/mediapipe/apps/facedetectioncpu)
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
## Face Mesh on Desktop with Webcam
|
||||
|
||||
This doc focuses on running the **MediaPipe Face Mesh** pipeline to perform 3D
|
||||
face landmark estimation in real-time on desktop with webcam input. The pipeline
|
||||
internally incorporates TensorFlow Lite models. To know more about the models,
|
||||
please refer to the model
|
||||
[README file](https://github.com/google/mediapipe/tree/master/mediapipe/models/README.md#face-mesh).
|
||||
Moreover, if you are interested in running the same pipeline on Android/iOS,
|
||||
please see [Face Mesh on Android/iOS](face_mesh_mobile_gpu.md).
|
||||
|
||||
- [Face Mesh on Desktop with Webcam (CPU)](#face-mesh-on-desktop-with-webcam-cpu)
|
||||
|
||||
- [Face Mesh on Desktop with Webcam (GPU)](#face-mesh-on-desktop-with-webcam-gpu)
|
||||
|
||||
Note: Desktop GPU works only on Linux. Mesa drivers need to be installed. Please
|
||||
see
|
||||
[step 4 of "Installing on Debian and Ubuntu" in the installation guide](./install.md).
|
||||
|
||||
Note: If MediaPipe depends on OpenCV 2, please see the [known issues with OpenCV 2](#known-issues-with-opencv-2) section.
|
||||
|
||||
### Face Mesh on Desktop with Webcam (CPU)
|
||||
|
||||
To build and run Face Mesh on desktop with webcam (CPU), run:
|
||||
|
||||
```bash
|
||||
$ bazel build -c opt --define MEDIAPIPE_DISABLE_GPU=1 \
|
||||
mediapipe/examples/desktop/face_mesh:face_mesh_cpu
|
||||
|
||||
# It should print:
|
||||
# Target //mediapipe/examples/desktop/face_mesh:face_mesh_cpu up-to-date:
|
||||
# bazel-bin/mediapipe/examples/desktop/face_mesh/face_mesh_cpu
|
||||
|
||||
# This will open up your webcam as long as it is connected. Errors are likely
|
||||
# due to your webcam being not accessible.
|
||||
$ GLOG_logtostderr=1 bazel-bin/mediapipe/examples/desktop/face_mesh/face_mesh_cpu \
|
||||
--calculator_graph_config_file=mediapipe/graphs/face_mesh/face_mesh_desktop_live.pbtxt
|
||||
```
|
||||
|
||||
### Face Mesh on Desktop with Webcam (GPU)
|
||||
|
||||
Note: please first [check that your GPU is supported](gpu.md#desktop-gpu-linux).
|
||||
|
||||
To build and run Face Mesh on desktop with webcam (GPU), run:
|
||||
|
||||
```bash
|
||||
# This works only for linux currently
|
||||
$ bazel build -c opt --copt -DMESA_EGL_NO_X11_HEADERS --copt -DEGL_NO_X11 \
|
||||
mediapipe/examples/desktop/face_mesh:face_mesh_gpu
|
||||
|
||||
# It should print:
|
||||
# Target //mediapipe/examples/desktop/face_mesh:face_mesh_gpu up-to-date:
|
||||
# bazel-bin/mediapipe/examples/desktop/face_mesh/face_mesh_gpu
|
||||
|
||||
# This will open up your webcam as long as it is connected. Errors are likely
|
||||
# due to your webcam being not accessible, or GPU drivers not setup properly.
|
||||
$ GLOG_logtostderr=1 bazel-bin/mediapipe/examples/desktop/face_mesh/face_mesh_gpu \
|
||||
--calculator_graph_config_file=mediapipe/graphs/face_mesh/face_mesh_desktop_live_gpu.pbtxt
|
||||
```
|
||||
@@ -0,0 +1,90 @@
|
||||
# Face Mesh (GPU)
|
||||
|
||||
This example focuses on running the **MediaPipe Face Mesh** pipeline on mobile
|
||||
devices to perform 3D face landmark estimation in real-time, utilizing GPU
|
||||
acceleration. The pipeline internally incorporates TensorFlow Lite models. To
|
||||
know more about the models, please refer to the model
|
||||
[README file](https://github.com/google/mediapipe/tree/master/mediapipe/models/README.md#face-mesh).
|
||||
The pipeline is related to the
|
||||
[face detection example](./face_detection_mobile_gpu.md) as it internally
|
||||
utilizes face detection and performs landmark estimation only within the
|
||||
detected region.
|
||||
|
||||

|
||||
|
||||
**MediaPipe Face Mesh** generates 468 3D face landmarks in real-time on mobile
|
||||
devices. In the visualization above, the red dots represent the landmarks, and
|
||||
the green lines connecting landmarks illustrate the contours around the eyes,
|
||||
eyebrows, lips and the entire face.
|
||||
|
||||
## Android
|
||||
|
||||
[Source](https://github.com/google/mediapipe/tree/master/mediapipe/examples/android/src/java/com/google/mediapipe/apps/facemeshgpu)
|
||||
|
||||
A prebuilt arm64 APK can be
|
||||
[downloaded here](https://drive.google.com/open?id=1pUmd7CXCL_onYMbsZo5p91cH0oNnR4gi).
|
||||
|
||||
To build the app yourself, run:
|
||||
|
||||
```bash
|
||||
bazel build -c opt --config=android_arm64 mediapipe/examples/android/src/java/com/google/mediapipe/apps/facemeshgpu
|
||||
```
|
||||
|
||||
Once the app is built, install it on Android device with:
|
||||
|
||||
```bash
|
||||
adb install bazel-bin/mediapipe/examples/android/src/java/com/google/mediapipe/apps/facemeshgpu/facemeshgpu.apk
|
||||
```
|
||||
|
||||
## iOS
|
||||
|
||||
[Source](https://github.com/google/mediapipe/tree/master/mediapipe/examples/ios/facemeshgpu).
|
||||
|
||||
See the general [instructions](./mediapipe_ios_setup.md) for building iOS
|
||||
examples and generating an Xcode project. This will be the FaceMeshGpuApp
|
||||
target.
|
||||
|
||||
To build on the command line:
|
||||
|
||||
```bash
|
||||
bazel build -c opt --config=ios_arm64 mediapipe/examples/ios/facemeshgpu:FaceMeshGpuApp
|
||||
```
|
||||
|
||||
## Graph
|
||||
|
||||
The face mesh [main graph](#main-graph) utilizes a
|
||||
[face landmark subgraph](#face-landmark-subgraph) from the
|
||||
[face landmark module](https://github.com/google/mediapipe/tree/master/mediapipe/modules/face_landmark),
|
||||
and renders using a dedicated [face renderer subgraph](#face-renderer-subgraph).
|
||||
|
||||
The subgraphs show up in the main graph visualization as nodes colored in
|
||||
purple, and the subgraph itself can also be visualized just like a regular
|
||||
graph. For more information on how to visualize a graph that includes subgraphs,
|
||||
see the Visualizing Subgraphs section in the
|
||||
[visualizer documentation](./visualizer.md).
|
||||
|
||||
### Main Graph
|
||||
|
||||

|
||||
|
||||
[Source pbtxt file](https://github.com/google/mediapipe/tree/master/mediapipe/graphs/face_mesh/face_mesh_mobile.pbtxt)
|
||||
|
||||
### Face Landmark Subgraph
|
||||
|
||||
The
|
||||
[face landmark module](https://github.com/google/mediapipe/tree/master/mediapipe/modules/face_landmark)
|
||||
contains several subgraphs that can be used to detect and track face landmarks.
|
||||
In particular, in this example the
|
||||
[FaceLandmarkFrontGPU](https://github.com/google/mediapipe/tree/master/mediapipe/modules/face_landmark/face_landmark_front_gpu.pbtxt)
|
||||
subgraph, suitable for images from front-facing cameras (i.e., selfie images)
|
||||
and utilizing GPU acceleration, is selected.
|
||||
|
||||

|
||||
|
||||
[Source pbtxt file](https://github.com/google/mediapipe/tree/master/mediapipe/modules/face_landmark/face_landmark_front_gpu.pbtxt)
|
||||
|
||||
### Face Renderer Subgraph
|
||||
|
||||

|
||||
|
||||
[Source pbtxt file](https://github.com/google/mediapipe/tree/master/mediapipe/graphs/face_mesh/subgraphs/face_renderer_gpu.pbtxt)
|
||||
@@ -190,10 +190,10 @@ within the MediaPipe framework:
|
||||
bazel build --define MEDIAPIPE_DISABLE_GPU=1 <my-target>
|
||||
|
||||
# to enable full GPU support (OpenGL ES 3.1+ & Metal)
|
||||
bazel build --copt -DMESA_EGL_NO_X11_HEADERS <my-target>
|
||||
bazel build --copt -DMESA_EGL_NO_X11_HEADERS --copt -DEGL_NO_X11 <my-target>
|
||||
|
||||
# to enable only OpenGL ES 3.0 and below (no GLES 3.1+ features)
|
||||
bazel build --copt -DMESA_EGL_NO_X11_HEADERS --copt -DMEDIAPIPE_DISABLE_GL_COMPUTE <my-target>
|
||||
bazel build --copt -DMESA_EGL_NO_X11_HEADERS --copt -DEGL_NO_X11 --copt -DMEDIAPIPE_DISABLE_GL_COMPUTE <my-target>
|
||||
```
|
||||
|
||||
Note *MEDIAPIPE_DISABLE_GL_COMPUTE* is automatically defined on all Apple
|
||||
|
||||
@@ -26,7 +26,7 @@ To build and run the TensorFlow Lite example on desktop (GPU) with Webcam, run:
|
||||
```bash
|
||||
# Video from webcam running on desktop GPU
|
||||
# This works only for linux currently
|
||||
$ bazel build -c opt --copt -DMESA_EGL_NO_X11_HEADERS \
|
||||
$ bazel build -c opt --copt -DMESA_EGL_NO_X11_HEADERS --copt -DEGL_NO_X11 \
|
||||
mediapipe/examples/desktop/hair_segmentation:hair_segmentation_gpu
|
||||
|
||||
# It should print:
|
||||
|
||||
@@ -48,7 +48,7 @@ To build and run the TensorFlow Lite example on desktop (GPU) with Webcam, run:
|
||||
```bash
|
||||
# Video from webcam running on desktop GPU
|
||||
# This works only for linux currently
|
||||
$ bazel build -c opt --copt -DMESA_EGL_NO_X11_HEADERS \
|
||||
$ bazel build -c opt --copt -DMESA_EGL_NO_X11_HEADERS --copt -DEGL_NO_X11 \
|
||||
mediapipe/examples/desktop/hand_tracking:hand_tracking_gpu
|
||||
|
||||
# It should print:
|
||||
|
||||
|
After Width: | Height: | Size: 26 KiB |
|
After Width: | Height: | Size: 20 KiB |
|
Before Width: | Height: | Size: 6.8 KiB After Width: | Height: | Size: 9.4 KiB |
|
Before Width: | Height: | Size: 1.8 MiB After Width: | Height: | Size: 2.3 MiB |
|
Before Width: | Height: | Size: 350 KiB After Width: | Height: | Size: 808 KiB |
|
After Width: | Height: | Size: 121 KiB |
|
After Width: | Height: | Size: 3.3 MiB |
|
After Width: | Height: | Size: 1.0 MiB |
|
After Width: | Height: | Size: 32 KiB |
|
After Width: | Height: | Size: 59 KiB |
|
Before Width: | Height: | Size: 2.1 MiB After Width: | Height: | Size: 1.3 MiB |
|
Before Width: | Height: | Size: 430 KiB After Width: | Height: | Size: 460 KiB |
|
Before Width: | Height: | Size: 2.8 MiB After Width: | Height: | Size: 3.1 MiB |
|
After Width: | Height: | Size: 383 KiB |
|
Before Width: | Height: | Size: 3.5 MiB After Width: | Height: | Size: 5.6 MiB |
|
Before Width: | Height: | Size: 1.6 MiB |
|
Before Width: | Height: | Size: 3.6 MiB After Width: | Height: | Size: 4.7 MiB |
|
Before Width: | Height: | Size: 1.4 MiB After Width: | Height: | Size: 448 KiB |
|
Before Width: | Height: | Size: 2.2 MiB After Width: | Height: | Size: 1.5 MiB |
|
After Width: | Height: | Size: 475 KiB |
|
Before Width: | Height: | Size: 2.6 MiB After Width: | Height: | Size: 1.3 MiB |
|
Before Width: | Height: | Size: 1.7 MiB |
|
Before Width: | Height: | Size: 1.6 MiB After Width: | Height: | Size: 1004 KiB |
|
Before Width: | Height: | Size: 1.5 MiB After Width: | Height: | Size: 945 KiB |
|
Before Width: | Height: | Size: 282 KiB After Width: | Height: | Size: 336 KiB |
|
Before Width: | Height: | Size: 8.6 MiB After Width: | Height: | Size: 2.2 MiB |
|
After Width: | Height: | Size: 2.2 MiB |
|
Before Width: | Height: | Size: 1.3 MiB After Width: | Height: | Size: 730 KiB |