Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4c68eb4a70 | ||
|
|
a3d36eee32 | ||
|
|
1722d4b8a2 | ||
|
|
3b6d3c4058 | ||
|
|
252a5713c7 | ||
|
|
de4fbc10e6 | ||
|
|
d144e564d8 | ||
|
|
dd02df1dbe | ||
|
|
66b377c825 | ||
|
|
bf5185f122 | ||
|
|
a2823541e6 | ||
|
|
ae6be10afe | ||
|
|
38ee2603a7 | ||
|
|
86b3283b2f | ||
|
|
7d470a1335 | ||
|
|
d16cc3be5b |
@@ -12,14 +12,21 @@ build --copt='-Wno-comment'
|
|||||||
build --copt='-Wno-return-type'
|
build --copt='-Wno-return-type'
|
||||||
build --copt='-Wno-unused-local-typedefs'
|
build --copt='-Wno-unused-local-typedefs'
|
||||||
build --copt='-Wno-ignored-attributes'
|
build --copt='-Wno-ignored-attributes'
|
||||||
# Temporarily set the incompatiblity flag for Bazel 0.27.0 and above
|
|
||||||
build --incompatible_disable_deprecated_attr_params=false
|
# Tensorflow needs remote repo
|
||||||
build --incompatible_depset_is_not_iterable=false
|
build --experimental_repo_remote_exec
|
||||||
|
|
||||||
# Sets the default Apple platform to macOS.
|
# Sets the default Apple platform to macOS.
|
||||||
build --apple_platform_type=macos
|
build --apple_platform_type=macos
|
||||||
|
|
||||||
|
# Allow debugging with XCODE
|
||||||
|
build --apple_generate_dsym
|
||||||
|
|
||||||
# Android configs.
|
# 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 --crosstool_top=//external:android/crosstool
|
||||||
build:android --host_crosstool_top=@bazel_tools//tools/cpp:toolchain
|
build:android --host_crosstool_top=@bazel_tools//tools/cpp:toolchain
|
||||||
build:android --linkopt=-landroid
|
build:android --linkopt=-landroid
|
||||||
|
|||||||
+3
-1
@@ -35,6 +35,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
|
|||||||
libopencv-highgui-dev \
|
libopencv-highgui-dev \
|
||||||
libopencv-imgproc-dev \
|
libopencv-imgproc-dev \
|
||||||
libopencv-video-dev \
|
libopencv-video-dev \
|
||||||
|
libopencv-calib3d-dev \
|
||||||
|
libopencv-features2d-dev \
|
||||||
software-properties-common && \
|
software-properties-common && \
|
||||||
add-apt-repository -y ppa:openjdk-r/ppa && \
|
add-apt-repository -y ppa:openjdk-r/ppa && \
|
||||||
apt-get update && apt-get install -y openjdk-8-jdk && \
|
apt-get update && apt-get install -y openjdk-8-jdk && \
|
||||||
@@ -46,7 +48,7 @@ RUN pip install future
|
|||||||
RUN pip3 install six
|
RUN pip3 install six
|
||||||
|
|
||||||
# Install bazel
|
# Install bazel
|
||||||
ARG BAZEL_VERSION=1.1.0
|
ARG BAZEL_VERSION=2.0.0
|
||||||
RUN mkdir /bazel && \
|
RUN mkdir /bazel && \
|
||||||
wget --no-check-certificate -O /bazel/installer.sh "https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION}/b\
|
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" && \
|
azel-${BAZEL_VERSION}-installer-linux-x86_64.sh" && \
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||

|

|
||||||
=======================================================================
|
=======================================================================
|
||||||
|
|
||||||
[MediaPipe](http://mediapipe.dev) is a framework for building multimodal (eg. video, audio, any time series data) applied ML pipelines. With MediaPipe, a perception pipeline can be built as a graph of modular components, including, for instance, inference models (e.g., TensorFlow, TFLite) and media processing functions.
|
[MediaPipe](http://mediapipe.dev) is a framework for building multimodal (eg. video, audio, any time series data), cross platform (i.e Android, iOS, web, edge devices) applied ML pipelines. With MediaPipe, a perception pipeline can be built as a graph of modular components, including, for instance, inference models (e.g., TensorFlow, TFLite) and media processing functions.
|
||||||
|
|
||||||

|

|
||||||
|
|
||||||
@@ -9,17 +9,20 @@
|
|||||||
|
|
||||||
## ML Solutions in MediaPipe
|
## ML Solutions in MediaPipe
|
||||||
|
|
||||||
* [Hand Tracking](mediapipe/docs/hand_tracking_mobile_gpu.md)
|
* [Face Detection](mediapipe/docs/face_detection_mobile_gpu.md) [[Web Demo]](https://viz.mediapipe.dev/runner/demos/face_detection/face_detection.html)
|
||||||
* [Multi-hand Tracking](mediapipe/docs/multi_hand_tracking_mobile_gpu.md)
|
* [Multi-hand Tracking](mediapipe/docs/multi_hand_tracking_mobile_gpu.md)
|
||||||
* [Face Detection](mediapipe/docs/face_detection_mobile_gpu.md)
|
* [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)
|
* [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](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)
|
||||||
|
|
||||||

|
|
||||||

|
|
||||||

|

|
||||||
|

|
||||||
|

|
||||||

|

|
||||||

|

|
||||||
|
|
||||||
## Installation
|
## Installation
|
||||||
Follow these [instructions](mediapipe/docs/install.md).
|
Follow these [instructions](mediapipe/docs/install.md).
|
||||||
@@ -27,6 +30,8 @@ Follow these [instructions](mediapipe/docs/install.md).
|
|||||||
## Getting started
|
## Getting started
|
||||||
See mobile, desktop and Google Coral [examples](mediapipe/docs/examples.md).
|
See mobile, desktop 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)
|
||||||
|
|
||||||
## Documentation
|
## Documentation
|
||||||
[MediaPipe Read-the-Docs](https://mediapipe.readthedocs.io/) or [docs.mediapipe.dev](https://docs.mediapipe.dev)
|
[MediaPipe Read-the-Docs](https://mediapipe.readthedocs.io/) or [docs.mediapipe.dev](https://docs.mediapipe.dev)
|
||||||
|
|
||||||
@@ -35,10 +40,14 @@ Check out the [Examples page](https://mediapipe.readthedocs.io/en/latest/example
|
|||||||
## Visualizing MediaPipe graphs
|
## 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).
|
A web-based visualizer is hosted on [viz.mediapipe.dev](https://viz.mediapipe.dev/). Please also see instructions [here](mediapipe/docs/visualizer.md).
|
||||||
|
|
||||||
## Community forum
|
## Videos
|
||||||
* [Discuss](https://groups.google.com/forum/#!forum/mediapipe) - General community discussion around MediaPipe
|
* [YouTube Channel](https://www.youtube.com/channel/UCObqmpuSMx-usADtL_qdMAw)
|
||||||
|
|
||||||
## Publications
|
## Publications
|
||||||
|
* [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)
|
||||||
|
* [Google Developer Blog: Object Detection and Tracking using MediaPipe](https://mediapipe.page.link/objecttrackingblog)
|
||||||
* [On-Device, Real-Time Hand Tracking with MediaPipe](https://ai.googleblog.com/2019/08/on-device-real-time-hand-tracking-with.html)
|
* [On-Device, Real-Time Hand Tracking with MediaPipe](https://ai.googleblog.com/2019/08/on-device-real-time-hand-tracking-with.html)
|
||||||
* [MediaPipe: A Framework for Building Perception Pipelines](https://arxiv.org/abs/1906.08172)
|
* [MediaPipe: A Framework for Building Perception Pipelines](https://arxiv.org/abs/1906.08172)
|
||||||
|
|
||||||
@@ -53,8 +62,11 @@ A web-based visualizer is hosted on [viz.mediapipe.dev](https://viz.mediapipe.de
|
|||||||
* [Google Industry Workshop at ICIP 2019](http://2019.ieeeicip.org/?action=page4&id=14#Google) [Presentation](https://docs.google.com/presentation/d/e/2PACX-1vRIBBbO_LO9v2YmvbHHEt1cwyqH6EjDxiILjuT0foXy1E7g6uyh4CesB2DkkEwlRDO9_lWfuKMZx98T/pub?start=false&loop=false&delayms=3000&slide=id.g556cc1a659_0_5) on Sept 24 in Taipei, Taiwan
|
* [Google Industry Workshop at ICIP 2019](http://2019.ieeeicip.org/?action=page4&id=14#Google) [Presentation](https://docs.google.com/presentation/d/e/2PACX-1vRIBBbO_LO9v2YmvbHHEt1cwyqH6EjDxiILjuT0foXy1E7g6uyh4CesB2DkkEwlRDO9_lWfuKMZx98T/pub?start=false&loop=false&delayms=3000&slide=id.g556cc1a659_0_5) on Sept 24 in Taipei, Taiwan
|
||||||
* [Open sourced at CVPR 2019](https://sites.google.com/corp/view/perception-cv4arvr/mediapipe) on June 17~20 in Long Beach, CA
|
* [Open sourced at CVPR 2019](https://sites.google.com/corp/view/perception-cv4arvr/mediapipe) on June 17~20 in Long Beach, CA
|
||||||
|
|
||||||
|
## Community forum
|
||||||
|
* [Discuss](https://groups.google.com/forum/#!forum/mediapipe) - General community discussion around MediaPipe
|
||||||
|
|
||||||
## Alpha Disclaimer
|
## Alpha Disclaimer
|
||||||
MediaPipe is currently in alpha for v0.6. We are still making breaking API changes and expect to get to stable API by v1.0.
|
MediaPipe is currently in alpha for v0.7. We are still making breaking API changes and expect to get to stable API by v1.0.
|
||||||
|
|
||||||
## Contributing
|
## Contributing
|
||||||
We welcome contributions. Please follow these [guidelines](./CONTRIBUTING.md).
|
We welcome contributions. Please follow these [guidelines](./CONTRIBUTING.md).
|
||||||
|
|||||||
@@ -2,21 +2,22 @@ workspace(name = "mediapipe")
|
|||||||
|
|
||||||
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
|
load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
|
||||||
|
|
||||||
skylib_version = "0.8.0"
|
skylib_version = "0.9.0"
|
||||||
http_archive(
|
http_archive(
|
||||||
name = "bazel_skylib",
|
name = "bazel_skylib",
|
||||||
type = "tar.gz",
|
type = "tar.gz",
|
||||||
url = "https://github.com/bazelbuild/bazel-skylib/releases/download/{}/bazel-skylib.{}.tar.gz".format (skylib_version, skylib_version),
|
url = "https://github.com/bazelbuild/bazel-skylib/releases/download/{}/bazel_skylib-{}.tar.gz".format (skylib_version, skylib_version),
|
||||||
sha256 = "2ef429f5d7ce7111263289644d233707dba35e39696377ebab8b0bc701f7818e",
|
sha256 = "1dde365491125a3db70731e25658dfdd3bc5dbdfd11b840b3e987ecf043c7ca0",
|
||||||
)
|
)
|
||||||
load("@bazel_skylib//lib:versions.bzl", "versions")
|
load("@bazel_skylib//lib:versions.bzl", "versions")
|
||||||
versions.check(minimum_bazel_version = "0.24.1")
|
versions.check(minimum_bazel_version = "2.0.0")
|
||||||
|
|
||||||
# ABSL cpp library lts_2019_08_08.
|
|
||||||
|
# ABSL cpp library lts_2020_02_25
|
||||||
http_archive(
|
http_archive(
|
||||||
name = "com_google_absl",
|
name = "com_google_absl",
|
||||||
urls = [
|
urls = [
|
||||||
"https://github.com/abseil/abseil-cpp/archive/20190808.tar.gz",
|
"https://github.com/abseil/abseil-cpp/archive/20200225.tar.gz",
|
||||||
],
|
],
|
||||||
# Remove after https://github.com/abseil/abseil-cpp/issues/326 is solved.
|
# Remove after https://github.com/abseil/abseil-cpp/issues/326 is solved.
|
||||||
patches = [
|
patches = [
|
||||||
@@ -25,8 +26,8 @@ http_archive(
|
|||||||
patch_args = [
|
patch_args = [
|
||||||
"-p1",
|
"-p1",
|
||||||
],
|
],
|
||||||
strip_prefix = "abseil-cpp-20190808",
|
strip_prefix = "abseil-cpp-20200225",
|
||||||
sha256 = "8100085dada279bf3ee00cd064d43b5f55e5d913be0dfe2906f06f8f28d5b37e"
|
sha256 = "728a813291bdec2aa46eab8356ace9f75ac2ed9dfe2df5ab603c4e6c09f1c353"
|
||||||
)
|
)
|
||||||
|
|
||||||
http_archive(
|
http_archive(
|
||||||
@@ -76,6 +77,14 @@ http_archive(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# easyexif
|
||||||
|
http_archive(
|
||||||
|
name = "easyexif",
|
||||||
|
url = "https://github.com/mayanklahiri/easyexif/archive/master.zip",
|
||||||
|
strip_prefix = "easyexif-master",
|
||||||
|
build_file = "@//third_party:easyexif.BUILD",
|
||||||
|
)
|
||||||
|
|
||||||
# libyuv
|
# libyuv
|
||||||
http_archive(
|
http_archive(
|
||||||
name = "libyuv",
|
name = "libyuv",
|
||||||
@@ -83,11 +92,13 @@ http_archive(
|
|||||||
build_file = "@//third_party:libyuv.BUILD",
|
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(
|
http_archive(
|
||||||
name = "com_google_protobuf_javalite",
|
name = "com_google_protobuf_javalite",
|
||||||
sha256 = "79d102c61e2a479a0b7e5fc167bcfaa4832a0c6aad4a75fa7da0480564931bcc",
|
sha256 = "a79d19dcdf9139fa4b81206e318e33d245c4c9da1ffed21c87288ed4380426f9",
|
||||||
strip_prefix = "protobuf-384989534b2246d413dbcd750744faab2607b516",
|
strip_prefix = "protobuf-3.11.4",
|
||||||
urls = ["https://github.com/google/protobuf/archive/384989534b2246d413dbcd750744faab2607b516.zip"],
|
urls = ["https://github.com/protocolbuffers/protobuf/archive/v3.11.4.tar.gz"],
|
||||||
)
|
)
|
||||||
|
|
||||||
http_archive(
|
http_archive(
|
||||||
@@ -107,15 +118,22 @@ http_archive(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
# 2019-11-21
|
# 2020-04-01
|
||||||
_TENSORFLOW_GIT_COMMIT = "f482488b481a799ca07e7e2d153cf47b8e91a60c"
|
_TENSORFLOW_GIT_COMMIT = "805e47cea96c7e8c6fccf494d40a2392dc99fdd8"
|
||||||
_TENSORFLOW_SHA256= "8d9118c2ce186c7e1403f04b96982fe72c184060c7f7a93e30a28dca358694f0"
|
_TENSORFLOW_SHA256= "9ee3ae604c2e1345ac60345becee6d659364721513f9cb8652eb2e7138320ca5"
|
||||||
http_archive(
|
http_archive(
|
||||||
name = "org_tensorflow",
|
name = "org_tensorflow",
|
||||||
urls = [
|
urls = [
|
||||||
"https://mirror.bazel.build/github.com/tensorflow/tensorflow/archive/%s.tar.gz" % _TENSORFLOW_GIT_COMMIT,
|
"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,
|
"https://github.com/tensorflow/tensorflow/archive/%s.tar.gz" % _TENSORFLOW_GIT_COMMIT,
|
||||||
],
|
],
|
||||||
|
patches = [
|
||||||
|
"@//third_party:org_tensorflow_compatibility_fixes.diff",
|
||||||
|
"@//third_party:org_tensorflow_protobuf_updates.diff",
|
||||||
|
],
|
||||||
|
patch_args = [
|
||||||
|
"-p1",
|
||||||
|
],
|
||||||
strip_prefix = "tensorflow-%s" % _TENSORFLOW_GIT_COMMIT,
|
strip_prefix = "tensorflow-%s" % _TENSORFLOW_GIT_COMMIT,
|
||||||
sha256 = _TENSORFLOW_SHA256,
|
sha256 = _TENSORFLOW_SHA256,
|
||||||
)
|
)
|
||||||
@@ -123,9 +141,19 @@ http_archive(
|
|||||||
load("@org_tensorflow//tensorflow:workspace.bzl", "tf_workspace")
|
load("@org_tensorflow//tensorflow:workspace.bzl", "tf_workspace")
|
||||||
tf_workspace(tf_repo_name = "org_tensorflow")
|
tf_workspace(tf_repo_name = "org_tensorflow")
|
||||||
|
|
||||||
# Please run
|
http_archive(
|
||||||
# $ sudo apt-get install libopencv-core-dev libopencv-highgui-dev \
|
name = "ceres_solver",
|
||||||
# libopencv-imgproc-dev libopencv-video-dev
|
url = "https://github.com/ceres-solver/ceres-solver/archive/1.14.0.zip",
|
||||||
|
patches = [
|
||||||
|
"@//third_party:ceres_solver_9bf9588988236279e1262f75d7f4d85711dfa172.diff"
|
||||||
|
],
|
||||||
|
patch_args = [
|
||||||
|
"-p1",
|
||||||
|
],
|
||||||
|
strip_prefix = "ceres-solver-1.14.0",
|
||||||
|
sha256 = "5ba6d0db4e784621fda44a50c58bb23b0892684692f0c623e2063f9c19f192f1"
|
||||||
|
)
|
||||||
|
|
||||||
new_local_repository(
|
new_local_repository(
|
||||||
name = "linux_opencv",
|
name = "linux_opencv",
|
||||||
build_file = "@//third_party:opencv_linux.BUILD",
|
build_file = "@//third_party:opencv_linux.BUILD",
|
||||||
@@ -138,7 +166,6 @@ new_local_repository(
|
|||||||
path = "/usr"
|
path = "/usr"
|
||||||
)
|
)
|
||||||
|
|
||||||
# Please run $ brew install opencv@3
|
|
||||||
new_local_repository(
|
new_local_repository(
|
||||||
name = "macos_opencv",
|
name = "macos_opencv",
|
||||||
build_file = "@//third_party:opencv_macos.BUILD",
|
build_file = "@//third_party:opencv_macos.BUILD",
|
||||||
@@ -171,79 +198,6 @@ http_archive(
|
|||||||
url = "https://github.com/opencv/opencv/releases/download/3.2.0/opencv-3.2.0-ios-framework.zip",
|
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.
|
# You may run setup_android.sh to install Android SDK and NDK.
|
||||||
android_ndk_repository(
|
android_ndk_repository(
|
||||||
name = "androidndk",
|
name = "androidndk",
|
||||||
@@ -257,9 +211,15 @@ android_sdk_repository(
|
|||||||
|
|
||||||
http_archive(
|
http_archive(
|
||||||
name = "build_bazel_rules_apple",
|
name = "build_bazel_rules_apple",
|
||||||
sha256 = "bdc8e66e70b8a75da23b79f1f8c6207356df07d041d96d2189add7ee0780cf4e",
|
sha256 = "7a7afdd4869bb201c9352eed2daf37294d42b093579b70423490c1b4d4f6ce42",
|
||||||
strip_prefix = "rules_apple-b869b0d3868d78a1d4ffd866ccb304fb68aa12c3",
|
url = "https://github.com/bazelbuild/rules_apple/releases/download/0.19.0/rules_apple.0.19.0.tar.gz",
|
||||||
url = "https://github.com/bazelbuild/rules_apple/archive/b869b0d3868d78a1d4ffd866ccb304fb68aa12c3.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(
|
load(
|
||||||
@@ -292,3 +252,50 @@ http_archive(
|
|||||||
strip_prefix = "google-toolbox-for-mac-2.2.1",
|
strip_prefix = "google-toolbox-for-mac-2.2.1",
|
||||||
build_file = "@//third_party:google_toolbox_for_mac.BUILD",
|
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
|
licenses(["notice"]) # Apache 2.0
|
||||||
|
|
||||||
|
# Note: yes, these need to use "//external:android/crosstool", not
|
||||||
|
# @androidndk//:default_crosstool.
|
||||||
|
|
||||||
config_setting(
|
config_setting(
|
||||||
name = "android",
|
name = "android",
|
||||||
values = {"crosstool_top": "//external:android/crosstool"},
|
values = {"crosstool_top": "//external:android/crosstool"},
|
||||||
|
|||||||
@@ -9,6 +9,7 @@
|
|||||||
"mediapipe/examples/ios/facedetectiongpu/BUILD",
|
"mediapipe/examples/ios/facedetectiongpu/BUILD",
|
||||||
"mediapipe/examples/ios/handdetectiongpu/BUILD",
|
"mediapipe/examples/ios/handdetectiongpu/BUILD",
|
||||||
"mediapipe/examples/ios/handtrackinggpu/BUILD",
|
"mediapipe/examples/ios/handtrackinggpu/BUILD",
|
||||||
|
"mediapipe/examples/ios/multihandtrackinggpu/BUILD",
|
||||||
"mediapipe/examples/ios/objectdetectioncpu/BUILD",
|
"mediapipe/examples/ios/objectdetectioncpu/BUILD",
|
||||||
"mediapipe/examples/ios/objectdetectiongpu/BUILD"
|
"mediapipe/examples/ios/objectdetectiongpu/BUILD"
|
||||||
],
|
],
|
||||||
@@ -18,6 +19,7 @@
|
|||||||
"//mediapipe/examples/ios/facedetectiongpu:FaceDetectionGpuApp",
|
"//mediapipe/examples/ios/facedetectiongpu:FaceDetectionGpuApp",
|
||||||
"//mediapipe/examples/ios/handdetectiongpu:HandDetectionGpuApp",
|
"//mediapipe/examples/ios/handdetectiongpu:HandDetectionGpuApp",
|
||||||
"//mediapipe/examples/ios/handtrackinggpu:HandTrackingGpuApp",
|
"//mediapipe/examples/ios/handtrackinggpu:HandTrackingGpuApp",
|
||||||
|
"//mediapipe/examples/ios/multihandtrackinggpu:MultiHandTrackingGpuApp",
|
||||||
"//mediapipe/examples/ios/objectdetectioncpu:ObjectDetectionCpuApp",
|
"//mediapipe/examples/ios/objectdetectioncpu:ObjectDetectionCpuApp",
|
||||||
"//mediapipe/examples/ios/objectdetectiongpu:ObjectDetectionGpuApp",
|
"//mediapipe/examples/ios/objectdetectiongpu:ObjectDetectionGpuApp",
|
||||||
"//mediapipe/objc:mediapipe_framework_ios"
|
"//mediapipe/objc:mediapipe_framework_ios"
|
||||||
@@ -84,6 +86,8 @@
|
|||||||
"mediapipe/examples/ios/handdetectiongpu/Base.lproj",
|
"mediapipe/examples/ios/handdetectiongpu/Base.lproj",
|
||||||
"mediapipe/examples/ios/handtrackinggpu",
|
"mediapipe/examples/ios/handtrackinggpu",
|
||||||
"mediapipe/examples/ios/handtrackinggpu/Base.lproj",
|
"mediapipe/examples/ios/handtrackinggpu/Base.lproj",
|
||||||
|
"mediapipe/examples/ios/multihandtrackinggpu",
|
||||||
|
"mediapipe/examples/ios/multihandtrackinggpu/Base.lproj",
|
||||||
"mediapipe/examples/ios/objectdetectioncpu",
|
"mediapipe/examples/ios/objectdetectioncpu",
|
||||||
"mediapipe/examples/ios/objectdetectioncpu/Base.lproj",
|
"mediapipe/examples/ios/objectdetectioncpu/Base.lproj",
|
||||||
"mediapipe/examples/ios/objectdetectiongpu",
|
"mediapipe/examples/ios/objectdetectiongpu",
|
||||||
|
|||||||
@@ -16,6 +16,7 @@
|
|||||||
"mediapipe/examples/ios/facedetectiongpu",
|
"mediapipe/examples/ios/facedetectiongpu",
|
||||||
"mediapipe/examples/ios/handdetectiongpu",
|
"mediapipe/examples/ios/handdetectiongpu",
|
||||||
"mediapipe/examples/ios/handtrackinggpu",
|
"mediapipe/examples/ios/handtrackinggpu",
|
||||||
|
"mediapipe/examples/ios/multihandtrackinggpu",
|
||||||
"mediapipe/examples/ios/objectdetectioncpu",
|
"mediapipe/examples/ios/objectdetectioncpu",
|
||||||
"mediapipe/examples/ios/objectdetectiongpu"
|
"mediapipe/examples/ios/objectdetectiongpu"
|
||||||
],
|
],
|
||||||
|
|||||||
@@ -47,6 +47,13 @@ proto_library(
|
|||||||
deps = ["//mediapipe/framework:calculator_proto"],
|
deps = ["//mediapipe/framework:calculator_proto"],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
proto_library(
|
||||||
|
name = "packet_thinner_calculator_proto",
|
||||||
|
srcs = ["packet_thinner_calculator.proto"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = ["//mediapipe/framework:calculator_proto"],
|
||||||
|
)
|
||||||
|
|
||||||
proto_library(
|
proto_library(
|
||||||
name = "split_vector_calculator_proto",
|
name = "split_vector_calculator_proto",
|
||||||
srcs = ["split_vector_calculator.proto"],
|
srcs = ["split_vector_calculator.proto"],
|
||||||
@@ -79,6 +86,15 @@ proto_library(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
proto_library(
|
||||||
|
name = "constant_side_packet_calculator_proto",
|
||||||
|
srcs = ["constant_side_packet_calculator.proto"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_proto",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
proto_library(
|
proto_library(
|
||||||
name = "clip_vector_size_calculator_proto",
|
name = "clip_vector_size_calculator_proto",
|
||||||
srcs = ["clip_vector_size_calculator.proto"],
|
srcs = ["clip_vector_size_calculator.proto"],
|
||||||
@@ -102,6 +118,14 @@ mediapipe_cc_proto_library(
|
|||||||
deps = [":packet_resampler_calculator_proto"],
|
deps = [":packet_resampler_calculator_proto"],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
mediapipe_cc_proto_library(
|
||||||
|
name = "packet_thinner_calculator_cc_proto",
|
||||||
|
srcs = ["packet_thinner_calculator.proto"],
|
||||||
|
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [":packet_thinner_calculator_proto"],
|
||||||
|
)
|
||||||
|
|
||||||
mediapipe_cc_proto_library(
|
mediapipe_cc_proto_library(
|
||||||
name = "split_vector_calculator_cc_proto",
|
name = "split_vector_calculator_cc_proto",
|
||||||
srcs = ["split_vector_calculator.proto"],
|
srcs = ["split_vector_calculator.proto"],
|
||||||
@@ -158,6 +182,14 @@ mediapipe_cc_proto_library(
|
|||||||
deps = [":gate_calculator_proto"],
|
deps = [":gate_calculator_proto"],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
mediapipe_cc_proto_library(
|
||||||
|
name = "constant_side_packet_calculator_cc_proto",
|
||||||
|
srcs = ["constant_side_packet_calculator.proto"],
|
||||||
|
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [":constant_side_packet_calculator_proto"],
|
||||||
|
)
|
||||||
|
|
||||||
cc_library(
|
cc_library(
|
||||||
name = "add_header_calculator",
|
name = "add_header_calculator",
|
||||||
srcs = ["add_header_calculator.cc"],
|
srcs = ["add_header_calculator.cc"],
|
||||||
@@ -165,6 +197,7 @@ cc_library(
|
|||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/port:logging",
|
"//mediapipe/framework/port:logging",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
],
|
],
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -195,7 +228,9 @@ cc_library(
|
|||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:collection_item_id",
|
"//mediapipe/framework:collection_item_id",
|
||||||
"//mediapipe/framework:packet",
|
"//mediapipe/framework:packet",
|
||||||
|
"//mediapipe/framework/formats:detection_cc_proto",
|
||||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||||
|
"//mediapipe/framework/formats:matrix",
|
||||||
"//mediapipe/framework/formats:rect_cc_proto",
|
"//mediapipe/framework/formats:rect_cc_proto",
|
||||||
"//mediapipe/framework/port:integral_types",
|
"//mediapipe/framework/port:integral_types",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
@@ -216,12 +251,14 @@ cc_library(
|
|||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:collection_item_id",
|
"//mediapipe/framework:collection_item_id",
|
||||||
"//mediapipe/framework:packet",
|
"//mediapipe/framework:packet",
|
||||||
|
"//mediapipe/framework/formats:classification_cc_proto",
|
||||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||||
"//mediapipe/framework/formats:rect_cc_proto",
|
"//mediapipe/framework/formats:rect_cc_proto",
|
||||||
"//mediapipe/framework/port:integral_types",
|
"//mediapipe/framework/port:integral_types",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
"//mediapipe/util:render_data_cc_proto",
|
"//mediapipe/util:render_data_cc_proto",
|
||||||
|
"@org_tensorflow//tensorflow/lite:framework",
|
||||||
],
|
],
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -232,10 +269,11 @@ cc_test(
|
|||||||
deps = [
|
deps = [
|
||||||
":begin_loop_calculator",
|
":begin_loop_calculator",
|
||||||
":end_loop_calculator",
|
":end_loop_calculator",
|
||||||
"//mediapipe/calculators/core:packet_cloner_calculator",
|
":gate_calculator",
|
||||||
"//mediapipe/framework:calculator_context",
|
"//mediapipe/framework:calculator_context",
|
||||||
"//mediapipe/framework:calculator_contract",
|
"//mediapipe/framework:calculator_contract",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:packet",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"//mediapipe/framework/port:integral_types",
|
"//mediapipe/framework/port:integral_types",
|
||||||
"//mediapipe/framework/port:parse_text_proto",
|
"//mediapipe/framework/port:parse_text_proto",
|
||||||
@@ -283,7 +321,6 @@ cc_test(
|
|||||||
srcs = ["concatenate_vector_calculator_test.cc"],
|
srcs = ["concatenate_vector_calculator_test.cc"],
|
||||||
deps = [
|
deps = [
|
||||||
":concatenate_vector_calculator",
|
":concatenate_vector_calculator",
|
||||||
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
|
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework:timestamp",
|
"//mediapipe/framework:timestamp",
|
||||||
@@ -302,6 +339,7 @@ cc_library(
|
|||||||
deps = [
|
deps = [
|
||||||
":clip_vector_size_calculator_cc_proto",
|
":clip_vector_size_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/formats:detection_cc_proto",
|
||||||
"//mediapipe/framework/formats:rect_cc_proto",
|
"//mediapipe/framework/formats:rect_cc_proto",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
@@ -450,6 +488,37 @@ cc_test(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "packet_thinner_calculator",
|
||||||
|
srcs = ["packet_thinner_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/calculators/core:packet_thinner_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework:calculator_context",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/formats:video_stream_header",
|
||||||
|
"//mediapipe/framework/port:integral_types",
|
||||||
|
"//mediapipe/framework/port:logging",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_test(
|
||||||
|
name = "packet_thinner_calculator_test",
|
||||||
|
srcs = ["packet_thinner_calculator_test.cc"],
|
||||||
|
deps = [
|
||||||
|
":packet_thinner_calculator",
|
||||||
|
"//mediapipe/calculators/core:packet_thinner_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:calculator_runner",
|
||||||
|
"//mediapipe/framework/formats:video_stream_header",
|
||||||
|
"//mediapipe/framework/port:gtest_main",
|
||||||
|
"//mediapipe/framework/port:integral_types",
|
||||||
|
"@com_google_absl//absl/strings",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
cc_library(
|
cc_library(
|
||||||
name = "pass_through_calculator",
|
name = "pass_through_calculator",
|
||||||
srcs = ["pass_through_calculator.cc"],
|
srcs = ["pass_through_calculator.cc"],
|
||||||
@@ -547,6 +616,22 @@ cc_library(
|
|||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_test(
|
||||||
|
name = "side_packet_to_stream_calculator_test",
|
||||||
|
srcs = ["side_packet_to_stream_calculator_test.cc"],
|
||||||
|
deps = [
|
||||||
|
":side_packet_to_stream_calculator",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/port:gtest_main",
|
||||||
|
"//mediapipe/framework/port:integral_types",
|
||||||
|
"//mediapipe/framework/port:parse_text_proto",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"//mediapipe/framework/tool:options_util",
|
||||||
|
"@com_google_absl//absl/memory",
|
||||||
|
"@com_google_absl//absl/strings",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
cc_test(
|
cc_test(
|
||||||
name = "immediate_mux_calculator_test",
|
name = "immediate_mux_calculator_test",
|
||||||
srcs = ["immediate_mux_calculator_test.cc"],
|
srcs = ["immediate_mux_calculator_test.cc"],
|
||||||
@@ -571,6 +656,7 @@ cc_test(
|
|||||||
cc_library(
|
cc_library(
|
||||||
name = "packet_resampler_calculator",
|
name = "packet_resampler_calculator",
|
||||||
srcs = ["packet_resampler_calculator.cc"],
|
srcs = ["packet_resampler_calculator.cc"],
|
||||||
|
hdrs = ["packet_resampler_calculator.h"],
|
||||||
visibility = [
|
visibility = [
|
||||||
"//visibility:public",
|
"//visibility:public",
|
||||||
],
|
],
|
||||||
@@ -594,17 +680,17 @@ cc_library(
|
|||||||
cc_test(
|
cc_test(
|
||||||
name = "packet_resampler_calculator_test",
|
name = "packet_resampler_calculator_test",
|
||||||
timeout = "short",
|
timeout = "short",
|
||||||
srcs = ["packet_resampler_calculator_test.cc"],
|
srcs = [
|
||||||
|
"packet_resampler_calculator_test.cc",
|
||||||
|
],
|
||||||
deps = [
|
deps = [
|
||||||
":packet_resampler_calculator",
|
":packet_resampler_calculator",
|
||||||
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
|
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework:timestamp",
|
|
||||||
"//mediapipe/framework/formats:video_stream_header",
|
"//mediapipe/framework/formats:video_stream_header",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"//mediapipe/framework/port:parse_text_proto",
|
"//mediapipe/framework/port:parse_text_proto",
|
||||||
"//mediapipe/framework/port:status",
|
|
||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
@@ -613,15 +699,17 @@ cc_test(
|
|||||||
name = "previous_loopback_calculator_test",
|
name = "previous_loopback_calculator_test",
|
||||||
srcs = ["previous_loopback_calculator_test.cc"],
|
srcs = ["previous_loopback_calculator_test.cc"],
|
||||||
deps = [
|
deps = [
|
||||||
|
":gate_calculator",
|
||||||
|
":make_pair_calculator",
|
||||||
|
":pass_through_calculator",
|
||||||
":previous_loopback_calculator",
|
":previous_loopback_calculator",
|
||||||
"//mediapipe/calculators/core:make_pair_calculator",
|
|
||||||
"//mediapipe/calculators/core:pass_through_calculator",
|
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework:timestamp",
|
"//mediapipe/framework:timestamp",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"//mediapipe/framework/port:integral_types",
|
"//mediapipe/framework/port:integral_types",
|
||||||
"//mediapipe/framework/port:parse_text_proto",
|
"//mediapipe/framework/port:parse_text_proto",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
"//mediapipe/framework/stream_handler:immediate_input_stream_handler",
|
"//mediapipe/framework/stream_handler:immediate_input_stream_handler",
|
||||||
"//mediapipe/framework/tool:sink",
|
"//mediapipe/framework/tool:sink",
|
||||||
"@com_google_absl//absl/time",
|
"@com_google_absl//absl/time",
|
||||||
@@ -689,15 +777,23 @@ cc_library(
|
|||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
":split_vector_calculator_cc_proto",
|
":split_vector_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework/formats:detection_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||||
"//mediapipe/framework/formats:rect_cc_proto",
|
"//mediapipe/framework/formats:rect_cc_proto",
|
||||||
|
"//mediapipe/framework/formats:matrix",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
"//mediapipe/util:resource_util",
|
"//mediapipe/util:resource_util",
|
||||||
"@org_tensorflow//tensorflow/lite:framework",
|
"@org_tensorflow//tensorflow/lite:framework",
|
||||||
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
||||||
],
|
] + select({
|
||||||
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
|
"//mediapipe:ios": [],
|
||||||
|
"//conditions:default": [
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
|
||||||
|
],
|
||||||
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -907,3 +1003,30 @@ cc_test(
|
|||||||
"@com_google_absl//absl/memory",
|
"@com_google_absl//absl/memory",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "constant_side_packet_calculator",
|
||||||
|
srcs = ["constant_side_packet_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
":constant_side_packet_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:collection_item_id",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_test(
|
||||||
|
name = "constant_side_packet_calculator_test",
|
||||||
|
srcs = ["constant_side_packet_calculator_test.cc"],
|
||||||
|
deps = [
|
||||||
|
":constant_side_packet_calculator",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/port:gtest_main",
|
||||||
|
"//mediapipe/framework/port:parse_text_proto",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"@com_google_absl//absl/strings",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|||||||
@@ -13,11 +13,12 @@
|
|||||||
// limitations under the License.
|
// limitations under the License.
|
||||||
|
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/canonical_errors.h"
|
||||||
#include "mediapipe/framework/port/logging.h"
|
#include "mediapipe/framework/port/logging.h"
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
// Attach the header from one stream to another stream.
|
// Attach the header from a stream or side input to another stream.
|
||||||
//
|
//
|
||||||
// The header stream (tag HEADER) must not have any packets in it.
|
// The header stream (tag HEADER) must not have any packets in it.
|
||||||
//
|
//
|
||||||
@@ -25,17 +26,53 @@ namespace mediapipe {
|
|||||||
// calculator to not need a header or to accept a separate stream with
|
// calculator to not need a header or to accept a separate stream with
|
||||||
// a header, that would be more future proof.
|
// a header, that would be more future proof.
|
||||||
//
|
//
|
||||||
|
// Example usage 1:
|
||||||
|
// node {
|
||||||
|
// calculator: "AddHeaderCalculator"
|
||||||
|
// input_stream: "DATA:audio"
|
||||||
|
// input_stream: "HEADER:audio_header"
|
||||||
|
// output_stream: "audio_with_header"
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// Example usage 2:
|
||||||
|
// node {
|
||||||
|
// calculator: "AddHeaderCalculator"
|
||||||
|
// input_stream: "DATA:audio"
|
||||||
|
// input_side_packet: "HEADER:audio_header"
|
||||||
|
// output_stream: "audio_with_header"
|
||||||
|
// }
|
||||||
|
//
|
||||||
class AddHeaderCalculator : public CalculatorBase {
|
class AddHeaderCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
cc->Inputs().Tag("HEADER").SetNone();
|
bool has_side_input = false;
|
||||||
|
bool has_header_stream = false;
|
||||||
|
if (cc->InputSidePackets().HasTag("HEADER")) {
|
||||||
|
cc->InputSidePackets().Tag("HEADER").SetAny();
|
||||||
|
has_side_input = true;
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag("HEADER")) {
|
||||||
|
cc->Inputs().Tag("HEADER").SetNone();
|
||||||
|
has_header_stream = true;
|
||||||
|
}
|
||||||
|
if (has_side_input == has_header_stream) {
|
||||||
|
return mediapipe::InvalidArgumentError(
|
||||||
|
"Header must be provided via exactly one of side input and input "
|
||||||
|
"stream");
|
||||||
|
}
|
||||||
cc->Inputs().Tag("DATA").SetAny();
|
cc->Inputs().Tag("DATA").SetAny();
|
||||||
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Tag("DATA"));
|
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Tag("DATA"));
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
const Packet& header = cc->Inputs().Tag("HEADER").Header();
|
Packet header;
|
||||||
|
if (cc->InputSidePackets().HasTag("HEADER")) {
|
||||||
|
header = cc->InputSidePackets().Tag("HEADER");
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag("HEADER")) {
|
||||||
|
header = cc->Inputs().Tag("HEADER").Header();
|
||||||
|
}
|
||||||
if (!header.IsEmpty()) {
|
if (!header.IsEmpty()) {
|
||||||
cc->Outputs().Index(0).SetHeader(header);
|
cc->Outputs().Index(0).SetHeader(header);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -14,8 +14,10 @@
|
|||||||
|
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/calculator_runner.h"
|
#include "mediapipe/framework/calculator_runner.h"
|
||||||
|
#include "mediapipe/framework/port/canonical_errors.h"
|
||||||
#include "mediapipe/framework/port/gmock.h"
|
#include "mediapipe/framework/port/gmock.h"
|
||||||
#include "mediapipe/framework/port/gtest.h"
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
#include "mediapipe/framework/port/status_matchers.h"
|
#include "mediapipe/framework/port/status_matchers.h"
|
||||||
#include "mediapipe/framework/timestamp.h"
|
#include "mediapipe/framework/timestamp.h"
|
||||||
#include "mediapipe/framework/tool/validate_type.h"
|
#include "mediapipe/framework/tool/validate_type.h"
|
||||||
@@ -24,7 +26,7 @@ namespace mediapipe {
|
|||||||
|
|
||||||
class AddHeaderCalculatorTest : public ::testing::Test {};
|
class AddHeaderCalculatorTest : public ::testing::Test {};
|
||||||
|
|
||||||
TEST_F(AddHeaderCalculatorTest, Works) {
|
TEST_F(AddHeaderCalculatorTest, HeaderStream) {
|
||||||
CalculatorGraphConfig::Node node;
|
CalculatorGraphConfig::Node node;
|
||||||
node.set_calculator("AddHeaderCalculator");
|
node.set_calculator("AddHeaderCalculator");
|
||||||
node.add_input_stream("HEADER:header_stream");
|
node.add_input_stream("HEADER:header_stream");
|
||||||
@@ -96,4 +98,62 @@ TEST_F(AddHeaderCalculatorTest, NoPacketsOnHeaderStream) {
|
|||||||
ASSERT_FALSE(runner.Run().ok());
|
ASSERT_FALSE(runner.Run().ok());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
TEST_F(AddHeaderCalculatorTest, InputSidePacket) {
|
||||||
|
CalculatorGraphConfig::Node node;
|
||||||
|
node.set_calculator("AddHeaderCalculator");
|
||||||
|
node.add_input_stream("DATA:data_stream");
|
||||||
|
node.add_output_stream("merged_stream");
|
||||||
|
node.add_input_side_packet("HEADER:header");
|
||||||
|
|
||||||
|
CalculatorRunner runner(node);
|
||||||
|
|
||||||
|
// Set header and add 5 packets.
|
||||||
|
runner.MutableSidePackets()->Tag("HEADER") =
|
||||||
|
Adopt(new std::string("my_header"));
|
||||||
|
for (int i = 0; i < 5; ++i) {
|
||||||
|
Packet packet = Adopt(new int(i)).At(Timestamp(i * 1000));
|
||||||
|
runner.MutableInputs()->Tag("DATA").packets.push_back(packet);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Run calculator.
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
ASSERT_EQ(1, runner.Outputs().NumEntries());
|
||||||
|
|
||||||
|
// Test output.
|
||||||
|
EXPECT_EQ(std::string("my_header"),
|
||||||
|
runner.Outputs().Index(0).header.Get<std::string>());
|
||||||
|
const std::vector<Packet>& output_packets = runner.Outputs().Index(0).packets;
|
||||||
|
ASSERT_EQ(5, output_packets.size());
|
||||||
|
for (int i = 0; i < 5; ++i) {
|
||||||
|
const int val = output_packets[i].Get<int>();
|
||||||
|
EXPECT_EQ(i, val);
|
||||||
|
EXPECT_EQ(Timestamp(i * 1000), output_packets[i].Timestamp());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(AddHeaderCalculatorTest, UsingBothSideInputAndStream) {
|
||||||
|
CalculatorGraphConfig::Node node;
|
||||||
|
node.set_calculator("AddHeaderCalculator");
|
||||||
|
node.add_input_stream("HEADER:header_stream");
|
||||||
|
node.add_input_stream("DATA:data_stream");
|
||||||
|
node.add_output_stream("merged_stream");
|
||||||
|
node.add_input_side_packet("HEADER:header");
|
||||||
|
|
||||||
|
CalculatorRunner runner(node);
|
||||||
|
|
||||||
|
// Set both headers and add 5 packets.
|
||||||
|
runner.MutableSidePackets()->Tag("HEADER") =
|
||||||
|
Adopt(new std::string("my_header"));
|
||||||
|
runner.MutableSidePackets()->Tag("HEADER") =
|
||||||
|
Adopt(new std::string("my_header"));
|
||||||
|
for (int i = 0; i < 5; ++i) {
|
||||||
|
Packet packet = Adopt(new int(i)).At(Timestamp(i * 1000));
|
||||||
|
runner.MutableInputs()->Tag("DATA").packets.push_back(packet);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Run should fail because header can only be provided one way.
|
||||||
|
EXPECT_EQ(runner.Run().code(), ::mediapipe::InvalidArgumentError("").code());
|
||||||
|
}
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -20,6 +20,8 @@
|
|||||||
#include "mediapipe/calculators/core/end_loop_calculator.h"
|
#include "mediapipe/calculators/core/end_loop_calculator.h"
|
||||||
#include "mediapipe/framework/calculator_contract.h"
|
#include "mediapipe/framework/calculator_contract.h"
|
||||||
#include "mediapipe/framework/calculator_framework.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/gtest.h"
|
||||||
#include "mediapipe/framework/port/integral_types.h"
|
#include "mediapipe/framework/port/integral_types.h"
|
||||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||||
@@ -28,6 +30,13 @@
|
|||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
namespace {
|
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;
|
typedef BeginLoopCalculator<std::vector<int>> BeginLoopIntegerCalculator;
|
||||||
REGISTER_CALCULATOR(BeginLoopIntegerCalculator);
|
REGISTER_CALCULATOR(BeginLoopIntegerCalculator);
|
||||||
|
|
||||||
@@ -59,8 +68,8 @@ REGISTER_CALCULATOR(EndLoopIntegersCalculator);
|
|||||||
|
|
||||||
class BeginEndLoopCalculatorGraphTest : public ::testing::Test {
|
class BeginEndLoopCalculatorGraphTest : public ::testing::Test {
|
||||||
protected:
|
protected:
|
||||||
BeginEndLoopCalculatorGraphTest() {
|
void SetUp() override {
|
||||||
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
R"(
|
R"(
|
||||||
num_threads: 4
|
num_threads: 4
|
||||||
input_stream: "ints"
|
input_stream: "ints"
|
||||||
@@ -82,94 +91,222 @@ class BeginEndLoopCalculatorGraphTest : public ::testing::Test {
|
|||||||
output_stream: "ITERABLE:ints_plus_one"
|
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_;
|
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) {
|
TEST_F(BeginEndLoopCalculatorGraphTest, SingleEmptyVector) {
|
||||||
CalculatorGraph graph;
|
SendPacketOfInts(Timestamp(0), {});
|
||||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||||
MP_EXPECT_OK(graph.StartRun({}));
|
|
||||||
auto input_vector = absl::make_unique<std::vector<int>>();
|
|
||||||
Timestamp input_timestamp = Timestamp(0);
|
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
|
||||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
|
||||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
|
||||||
|
|
||||||
// EndLoopCalc will forward the timestamp bound because there are no elements
|
// EndLoopCalc will forward the timestamp bound because there are no elements
|
||||||
// in collection to output.
|
// in collection to output.
|
||||||
ASSERT_EQ(0, output_packets_.size());
|
EXPECT_TRUE(output_packets_.empty());
|
||||||
|
|
||||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||||
}
|
}
|
||||||
|
|
||||||
TEST_F(BeginEndLoopCalculatorGraphTest, SingleNonEmptyVector) {
|
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);
|
Timestamp input_timestamp = Timestamp(0);
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
SendPacketOfInts(input_timestamp, {0, 1, 2});
|
||||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
|
||||||
|
|
||||||
ASSERT_EQ(1, output_packets_.size());
|
EXPECT_THAT(output_packets_,
|
||||||
EXPECT_EQ(input_timestamp, output_packets_[0].Timestamp());
|
testing::ElementsAre(
|
||||||
std::vector<int> expected_output_vector = {1, 2, 3};
|
PacketOfIntsEq(input_timestamp, std::vector<int>{1, 2, 3})));
|
||||||
EXPECT_EQ(expected_output_vector, output_packets_[0].Get<std::vector<int>>());
|
|
||||||
|
|
||||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||||
}
|
}
|
||||||
|
|
||||||
TEST_F(BeginEndLoopCalculatorGraphTest, MultipleVectors) {
|
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);
|
Timestamp input_timestamp0 = Timestamp(0);
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
SendPacketOfInts(input_timestamp0, {0, 1});
|
||||||
"ints", Adopt(input_vector0.release()).At(input_timestamp0)));
|
|
||||||
|
|
||||||
auto input_vector1 = absl::make_unique<std::vector<int>>();
|
|
||||||
Timestamp input_timestamp1 = Timestamp(1);
|
Timestamp input_timestamp1 = Timestamp(1);
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
SendPacketOfInts(input_timestamp1, {});
|
||||||
"ints", Adopt(input_vector1.release()).At(input_timestamp1)));
|
|
||||||
|
|
||||||
auto input_vector2 = absl::make_unique<std::vector<int>>();
|
|
||||||
input_vector2->emplace_back(2);
|
|
||||||
input_vector2->emplace_back(3);
|
|
||||||
Timestamp input_timestamp2 = Timestamp(2);
|
Timestamp input_timestamp2 = Timestamp(2);
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
SendPacketOfInts(input_timestamp2, {2, 3});
|
||||||
"ints", Adopt(input_vector2.release()).At(input_timestamp2)));
|
|
||||||
|
|
||||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||||
|
|
||||||
ASSERT_EQ(2, output_packets_.size());
|
|
||||||
|
|
||||||
EXPECT_EQ(input_timestamp0, output_packets_[0].Timestamp());
|
|
||||||
std::vector<int> expected_output_vector0 = {1, 2};
|
|
||||||
EXPECT_EQ(expected_output_vector0,
|
|
||||||
output_packets_[0].Get<std::vector<int>>());
|
|
||||||
|
|
||||||
// At input_timestamp1, EndLoopCalc will forward timestamp bound as there are
|
// At input_timestamp1, EndLoopCalc will forward timestamp bound as there are
|
||||||
// no elements in vector to process.
|
// 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());
|
// Passes non empty vector through or outputs empty vector in case of timestamp
|
||||||
std::vector<int> expected_output_vector2 = {3, 4};
|
// bound update.
|
||||||
EXPECT_EQ(expected_output_vector2,
|
class PassThroughOrEmptyVectorCalculator : public CalculatorBase {
|
||||||
output_packets_[1].Get<std::vector<int>>());
|
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 {
|
class MultiplierCalculator : public CalculatorBase {
|
||||||
@@ -199,8 +336,8 @@ REGISTER_CALCULATOR(MultiplierCalculator);
|
|||||||
|
|
||||||
class BeginEndLoopCalculatorGraphWithClonedInputsTest : public ::testing::Test {
|
class BeginEndLoopCalculatorGraphWithClonedInputsTest : public ::testing::Test {
|
||||||
protected:
|
protected:
|
||||||
BeginEndLoopCalculatorGraphWithClonedInputsTest() {
|
void SetUp() override {
|
||||||
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
R"(
|
R"(
|
||||||
num_threads: 4
|
num_threads: 4
|
||||||
input_stream: "ints"
|
input_stream: "ints"
|
||||||
@@ -226,109 +363,85 @@ class BeginEndLoopCalculatorGraphWithClonedInputsTest : public ::testing::Test {
|
|||||||
output_stream: "ITERABLE:multiplied_ints"
|
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_;
|
std::vector<Packet> output_packets_;
|
||||||
};
|
};
|
||||||
|
|
||||||
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, SingleEmptyVector) {
|
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest,
|
||||||
CalculatorGraph graph;
|
InputStreamForIterableIsEmpty) {
|
||||||
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
|
||||||
MP_EXPECT_OK(graph.StartRun({}));
|
|
||||||
auto input_vector = absl::make_unique<std::vector<int>>();
|
|
||||||
Timestamp input_timestamp = Timestamp(42);
|
Timestamp input_timestamp = Timestamp(42);
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
SendMultiplier(input_timestamp, /*multiplier=*/2);
|
||||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||||
auto multiplier = absl::make_unique<int>(2);
|
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
// EndLoopCalc will forward the timestamp bound because there are no packets
|
||||||
"multiplier", Adopt(multiplier.release()).At(input_timestamp)));
|
// to process.
|
||||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
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
|
// EndLoopCalc will forward the timestamp bound because there are no elements
|
||||||
// in collection to output.
|
// in collection to output.
|
||||||
ASSERT_EQ(0, output_packets_.size());
|
EXPECT_TRUE(output_packets_.empty());
|
||||||
|
|
||||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||||
}
|
}
|
||||||
|
|
||||||
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, SingleNonEmptyVector) {
|
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);
|
Timestamp input_timestamp = Timestamp(42);
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
SendPackets(input_timestamp, /*multiplier=*/2, /*ints=*/{0, 1, 2});
|
||||||
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
MP_ASSERT_OK(graph_.WaitUntilIdle());
|
||||||
auto multiplier = absl::make_unique<int>(2);
|
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
|
||||||
"multiplier", Adopt(multiplier.release()).At(input_timestamp)));
|
|
||||||
MP_ASSERT_OK(graph.WaitUntilIdle());
|
|
||||||
|
|
||||||
ASSERT_EQ(1, output_packets_.size());
|
EXPECT_THAT(output_packets_,
|
||||||
EXPECT_EQ(input_timestamp, output_packets_[0].Timestamp());
|
testing::ElementsAre(
|
||||||
std::vector<int> expected_output_vector = {0, 2, 4};
|
PacketOfIntsEq(input_timestamp, std::vector<int>{0, 2, 4})));
|
||||||
EXPECT_EQ(expected_output_vector, output_packets_[0].Get<std::vector<int>>());
|
|
||||||
|
|
||||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||||
}
|
}
|
||||||
|
|
||||||
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, MultipleVectors) {
|
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);
|
Timestamp input_timestamp0 = Timestamp(42);
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
SendPackets(input_timestamp0, /*multiplier=*/2, /*ints=*/{0, 1});
|
||||||
"ints", Adopt(input_vector0.release()).At(input_timestamp0)));
|
|
||||||
auto multiplier0 = absl::make_unique<int>(2);
|
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
|
||||||
"multiplier", Adopt(multiplier0.release()).At(input_timestamp0)));
|
|
||||||
|
|
||||||
auto input_vector1 = absl::make_unique<std::vector<int>>();
|
|
||||||
Timestamp input_timestamp1 = Timestamp(43);
|
Timestamp input_timestamp1 = Timestamp(43);
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
SendPackets(input_timestamp1, /*multiplier=*/2, /*ints=*/{});
|
||||||
"ints", Adopt(input_vector1.release()).At(input_timestamp1)));
|
|
||||||
auto multiplier1 = absl::make_unique<int>(2);
|
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
|
||||||
"multiplier", Adopt(multiplier1.release()).At(input_timestamp1)));
|
|
||||||
|
|
||||||
auto input_vector2 = absl::make_unique<std::vector<int>>();
|
|
||||||
input_vector2->emplace_back(2);
|
|
||||||
input_vector2->emplace_back(3);
|
|
||||||
Timestamp input_timestamp2 = Timestamp(44);
|
Timestamp input_timestamp2 = Timestamp(44);
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
SendPackets(input_timestamp2, /*multiplier=*/3, /*ints=*/{2, 3});
|
||||||
"ints", Adopt(input_vector2.release()).At(input_timestamp2)));
|
|
||||||
auto multiplier2 = absl::make_unique<int>(3);
|
|
||||||
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
|
||||||
"multiplier", Adopt(multiplier2.release()).At(input_timestamp2)));
|
|
||||||
|
|
||||||
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
MP_ASSERT_OK(graph_.CloseAllPacketSources());
|
||||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
MP_ASSERT_OK(graph_.WaitUntilDone());
|
||||||
|
|
||||||
ASSERT_EQ(2, output_packets_.size());
|
|
||||||
|
|
||||||
EXPECT_EQ(input_timestamp0, output_packets_[0].Timestamp());
|
|
||||||
std::vector<int> expected_output_vector0 = {0, 2};
|
|
||||||
EXPECT_EQ(expected_output_vector0,
|
|
||||||
output_packets_[0].Get<std::vector<int>>());
|
|
||||||
|
|
||||||
// At input_timestamp1, EndLoopCalc will forward timestamp bound as there are
|
// At input_timestamp1, EndLoopCalc will forward timestamp bound as there are
|
||||||
// no elements in vector to process.
|
// no elements in vector to process.
|
||||||
|
EXPECT_THAT(output_packets_,
|
||||||
EXPECT_EQ(input_timestamp2, output_packets_[1].Timestamp());
|
testing::ElementsAre(
|
||||||
std::vector<int> expected_output_vector2 = {6, 9};
|
PacketOfIntsEq(input_timestamp0, std::vector<int>{0, 2}),
|
||||||
EXPECT_EQ(expected_output_vector2,
|
PacketOfIntsEq(input_timestamp2, std::vector<int>{6, 9})));
|
||||||
output_packets_[1].Get<std::vector<int>>());
|
|
||||||
}
|
}
|
||||||
|
|
||||||
} // namespace
|
} // namespace
|
||||||
|
|||||||
@@ -16,7 +16,9 @@
|
|||||||
|
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/formats/detection.pb.h"
|
||||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||||
|
#include "mediapipe/framework/formats/matrix.h"
|
||||||
#include "mediapipe/framework/formats/rect.pb.h"
|
#include "mediapipe/framework/formats/rect.pb.h"
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
@@ -31,4 +33,13 @@ typedef BeginLoopCalculator<std::vector<::mediapipe::NormalizedRect>>
|
|||||||
BeginLoopNormalizedRectCalculator;
|
BeginLoopNormalizedRectCalculator;
|
||||||
REGISTER_CALCULATOR(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
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -52,20 +52,28 @@ namespace mediapipe {
|
|||||||
// output_stream: "OUTPUT:aggregated_result" # IterableU @ext_ts
|
// output_stream: "OUTPUT:aggregated_result" # IterableU @ext_ts
|
||||||
// }
|
// }
|
||||||
//
|
//
|
||||||
// BeginLoopCalculator accepts an optional input stream tagged with "TICK"
|
// Input streams tagged with "CLONE" are cloned to the corresponding output
|
||||||
// which if non-empty, wakes up the calculator and calls
|
// streams at loop timestamps. This ensures that a MediaPipe graph or sub-graph
|
||||||
// BeginLoopCalculator::Process(). Input streams tagged with "CLONE" are cloned
|
// can run multiple times, once per element in the "ITERABLE" for each pakcet
|
||||||
// to the corresponding output streams at loop timestamps. This ensures that a
|
// clone of the packets in the "CLONE" input streams.
|
||||||
// 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>
|
template <typename IterableT>
|
||||||
class BeginLoopCalculator : public CalculatorBase {
|
class BeginLoopCalculator : public CalculatorBase {
|
||||||
using ItemT = typename IterableT::value_type;
|
using ItemT = typename IterableT::value_type;
|
||||||
|
|
||||||
public:
|
public:
|
||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
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
|
// A non-empty packet in the optional "TICK" input stream wakes up the
|
||||||
// calculator.
|
// calculator.
|
||||||
|
// DEPRECATED as timestamp bound updates are processed by default in this
|
||||||
|
// calculator.
|
||||||
if (cc->Inputs().HasTag("TICK")) {
|
if (cc->Inputs().HasTag("TICK")) {
|
||||||
cc->Inputs().Tag("TICK").SetAny();
|
cc->Inputs().Tag("TICK").SetAny();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -17,6 +17,7 @@
|
|||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/formats/detection.pb.h"
|
||||||
#include "mediapipe/framework/formats/rect.pb.h"
|
#include "mediapipe/framework/formats/rect.pb.h"
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
@@ -25,4 +26,8 @@ typedef ClipVectorSizeCalculator<::mediapipe::NormalizedRect>
|
|||||||
ClipNormalizedRectVectorSizeCalculator;
|
ClipNormalizedRectVectorSizeCalculator;
|
||||||
REGISTER_CALCULATOR(ClipNormalizedRectVectorSizeCalculator);
|
REGISTER_CALCULATOR(ClipNormalizedRectVectorSizeCalculator);
|
||||||
|
|
||||||
|
typedef ClipVectorSizeCalculator<::mediapipe::Detection>
|
||||||
|
ClipDetectionVectorSizeCalculator;
|
||||||
|
REGISTER_CALCULATOR(ClipDetectionVectorSizeCalculator);
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -38,6 +38,8 @@ namespace mediapipe {
|
|||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
|
// Optionally, you can pass in a side packet that will override `max_vec_size`
|
||||||
|
// that is specified in the options.
|
||||||
template <typename T>
|
template <typename T>
|
||||||
class ClipVectorSizeCalculator : public CalculatorBase {
|
class ClipVectorSizeCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
@@ -53,6 +55,10 @@ class ClipVectorSizeCalculator : public CalculatorBase {
|
|||||||
|
|
||||||
cc->Inputs().Index(0).Set<std::vector<T>>();
|
cc->Inputs().Index(0).Set<std::vector<T>>();
|
||||||
cc->Outputs().Index(0).Set<std::vector<T>>();
|
cc->Outputs().Index(0).Set<std::vector<T>>();
|
||||||
|
// Optional input side packet that determines `max_vec_size`.
|
||||||
|
if (cc->InputSidePackets().NumEntries() > 0) {
|
||||||
|
cc->InputSidePackets().Index(0).Set<int>();
|
||||||
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -61,6 +67,11 @@ class ClipVectorSizeCalculator : public CalculatorBase {
|
|||||||
cc->SetOffset(TimestampDiff(0));
|
cc->SetOffset(TimestampDiff(0));
|
||||||
max_vec_size_ = cc->Options<::mediapipe::ClipVectorSizeCalculatorOptions>()
|
max_vec_size_ = cc->Options<::mediapipe::ClipVectorSizeCalculatorOptions>()
|
||||||
.max_vec_size();
|
.max_vec_size();
|
||||||
|
// Override `max_vec_size` if passed as side packet.
|
||||||
|
if (cc->InputSidePackets().NumEntries() > 0 &&
|
||||||
|
!cc->InputSidePackets().Index(0).IsEmpty()) {
|
||||||
|
max_vec_size_ = cc->InputSidePackets().Index(0).Get<int>();
|
||||||
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -176,4 +176,31 @@ TEST(TestClipUniqueIntPtrVectorSizeCalculatorTest, ConsumeOneTimestamp) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
TEST(TestClipIntVectorSizeCalculatorTest, SidePacket) {
|
||||||
|
CalculatorGraphConfig::Node node_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||||
|
calculator: "TestClipIntVectorSizeCalculator"
|
||||||
|
input_stream: "input_vector"
|
||||||
|
input_side_packet: "max_vec_size"
|
||||||
|
output_stream: "output_vector"
|
||||||
|
options {
|
||||||
|
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 1 }
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorRunner runner(node_config);
|
||||||
|
// This should override the default of 1 set in the options.
|
||||||
|
runner.MutableSidePackets()->Index(0) = Adopt(new int(2));
|
||||||
|
std::vector<int> input = {0, 1, 2, 3};
|
||||||
|
AddInputVector(input, /*timestamp=*/1, &runner);
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||||
|
EXPECT_EQ(1, outputs.size());
|
||||||
|
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||||
|
const std::vector<int>& output = outputs[0].Get<std::vector<int>>();
|
||||||
|
EXPECT_EQ(2, output.size());
|
||||||
|
std::vector<int> expected_vector = {0, 1};
|
||||||
|
EXPECT_EQ(expected_vector, output);
|
||||||
|
}
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -35,6 +35,16 @@ namespace mediapipe {
|
|||||||
typedef ConcatenateVectorCalculator<float> ConcatenateFloatVectorCalculator;
|
typedef ConcatenateVectorCalculator<float> ConcatenateFloatVectorCalculator;
|
||||||
REGISTER_CALCULATOR(ConcatenateFloatVectorCalculator);
|
REGISTER_CALCULATOR(ConcatenateFloatVectorCalculator);
|
||||||
|
|
||||||
|
// Example config:
|
||||||
|
// node {
|
||||||
|
// calculator: "ConcatenateInt32VectorCalculator"
|
||||||
|
// input_stream: "int32_vector_1"
|
||||||
|
// input_stream: "int32_vector_2"
|
||||||
|
// output_stream: "concatenated_int32_vector"
|
||||||
|
// }
|
||||||
|
typedef ConcatenateVectorCalculator<int32> ConcatenateInt32VectorCalculator;
|
||||||
|
REGISTER_CALCULATOR(ConcatenateInt32VectorCalculator);
|
||||||
|
|
||||||
// Example config:
|
// Example config:
|
||||||
// node {
|
// node {
|
||||||
// calculator: "ConcatenateTfLiteTensorVectorCalculator"
|
// calculator: "ConcatenateTfLiteTensorVectorCalculator"
|
||||||
|
|||||||
@@ -0,0 +1,116 @@
|
|||||||
|
// 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 <string>
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/core/constant_side_packet_calculator.pb.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/collection_item_id.h"
|
||||||
|
#include "mediapipe/framework/port/canonical_errors.h"
|
||||||
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
// Generates an output side packet or multiple output side packets according to
|
||||||
|
// the specified options.
|
||||||
|
//
|
||||||
|
// Example configs:
|
||||||
|
// node {
|
||||||
|
// calculator: "ConstantSidePacketCalculator"
|
||||||
|
// output_side_packet: "PACKET:packet"
|
||||||
|
// options: {
|
||||||
|
// [mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
|
||||||
|
// packet { int_value: 2 }
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "ConstantSidePacketCalculator"
|
||||||
|
// output_side_packet: "PACKET:0:int_packet"
|
||||||
|
// output_side_packet: "PACKET:1:bool_packet"
|
||||||
|
// options: {
|
||||||
|
// [mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
|
||||||
|
// packet { int_value: 2 }
|
||||||
|
// packet { bool_value: true }
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
class ConstantSidePacketCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
const auto& options = cc->Options().GetExtension(
|
||||||
|
::mediapipe::ConstantSidePacketCalculatorOptions::ext);
|
||||||
|
RET_CHECK_EQ(cc->OutputSidePackets().NumEntries(kPacketTag),
|
||||||
|
options.packet_size())
|
||||||
|
<< "Number of output side packets has to be same as number of packets "
|
||||||
|
"configured in options.";
|
||||||
|
|
||||||
|
int index = 0;
|
||||||
|
for (CollectionItemId id = cc->OutputSidePackets().BeginId(kPacketTag);
|
||||||
|
id != cc->OutputSidePackets().EndId(kPacketTag); ++id, ++index) {
|
||||||
|
const auto& packet_options = options.packet(index);
|
||||||
|
auto& packet = cc->OutputSidePackets().Get(id);
|
||||||
|
if (packet_options.has_int_value()) {
|
||||||
|
packet.Set<int>();
|
||||||
|
} else if (packet_options.has_float_value()) {
|
||||||
|
packet.Set<float>();
|
||||||
|
} else if (packet_options.has_bool_value()) {
|
||||||
|
packet.Set<bool>();
|
||||||
|
} else if (packet_options.has_string_value()) {
|
||||||
|
packet.Set<std::string>();
|
||||||
|
} else {
|
||||||
|
return ::mediapipe::InvalidArgumentError(
|
||||||
|
"None of supported values were specified in options.");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
|
const auto& options = cc->Options().GetExtension(
|
||||||
|
::mediapipe::ConstantSidePacketCalculatorOptions::ext);
|
||||||
|
int index = 0;
|
||||||
|
for (CollectionItemId id = cc->OutputSidePackets().BeginId(kPacketTag);
|
||||||
|
id != cc->OutputSidePackets().EndId(kPacketTag); ++id, ++index) {
|
||||||
|
auto& packet = cc->OutputSidePackets().Get(id);
|
||||||
|
const auto& packet_options = options.packet(index);
|
||||||
|
if (packet_options.has_int_value()) {
|
||||||
|
packet.Set(MakePacket<int>(packet_options.int_value()));
|
||||||
|
} else if (packet_options.has_float_value()) {
|
||||||
|
packet.Set(MakePacket<float>(packet_options.float_value()));
|
||||||
|
} else if (packet_options.has_bool_value()) {
|
||||||
|
packet.Set(MakePacket<bool>(packet_options.bool_value()));
|
||||||
|
} else if (packet_options.has_string_value()) {
|
||||||
|
packet.Set(MakePacket<std::string>(packet_options.string_value()));
|
||||||
|
} else {
|
||||||
|
return ::mediapipe::InvalidArgumentError(
|
||||||
|
"None of supported values were specified in options.");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
static constexpr const char* kPacketTag = "PACKET";
|
||||||
|
};
|
||||||
|
|
||||||
|
REGISTER_CALCULATOR(ConstantSidePacketCalculator);
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,36 @@
|
|||||||
|
// 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 ConstantSidePacketCalculatorOptions {
|
||||||
|
extend CalculatorOptions {
|
||||||
|
optional ConstantSidePacketCalculatorOptions ext = 291214597;
|
||||||
|
}
|
||||||
|
|
||||||
|
message ConstantSidePacket {
|
||||||
|
oneof value {
|
||||||
|
int32 int_value = 1;
|
||||||
|
float float_value = 2;
|
||||||
|
bool bool_value = 3;
|
||||||
|
string string_value = 4;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
repeated ConstantSidePacket packet = 1;
|
||||||
|
}
|
||||||
@@ -0,0 +1,196 @@
|
|||||||
|
// 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 <string>
|
||||||
|
|
||||||
|
#include "absl/strings/string_view.h"
|
||||||
|
#include "absl/strings/substitute.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/gmock.h"
|
||||||
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
#include "mediapipe/framework/port/status_matchers.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
void DoTestSingleSidePacket(absl::string_view packet_spec,
|
||||||
|
const T& expected_value) {
|
||||||
|
static constexpr absl::string_view graph_config_template = R"(
|
||||||
|
node {
|
||||||
|
calculator: "ConstantSidePacketCalculator"
|
||||||
|
output_side_packet: "PACKET:packet"
|
||||||
|
options: {
|
||||||
|
[mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
|
||||||
|
packet $0
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)";
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
absl::Substitute(graph_config_template, packet_spec));
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||||
|
MP_ASSERT_OK(graph.StartRun({}));
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.GetOutputSidePacket("packet"));
|
||||||
|
auto actual_value =
|
||||||
|
graph.GetOutputSidePacket("packet").ValueOrDie().template Get<T>();
|
||||||
|
EXPECT_EQ(actual_value, expected_value);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(ConstantSidePacketCalculatorTest, EveryPossibleType) {
|
||||||
|
DoTestSingleSidePacket("{ int_value: 2 }", 2);
|
||||||
|
DoTestSingleSidePacket("{ float_value: 6.5f }", 6.5f);
|
||||||
|
DoTestSingleSidePacket("{ bool_value: true }", true);
|
||||||
|
DoTestSingleSidePacket<std::string>(R"({ string_value: "str" })", "str");
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(ConstantSidePacketCalculatorTest, MultiplePackets) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||||
|
node {
|
||||||
|
calculator: "ConstantSidePacketCalculator"
|
||||||
|
output_side_packet: "PACKET:0:int_packet"
|
||||||
|
output_side_packet: "PACKET:1:float_packet"
|
||||||
|
output_side_packet: "PACKET:2:bool_packet"
|
||||||
|
output_side_packet: "PACKET:3:string_packet"
|
||||||
|
output_side_packet: "PACKET:4:another_string_packet"
|
||||||
|
output_side_packet: "PACKET:5:another_int_packet"
|
||||||
|
options: {
|
||||||
|
[mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
|
||||||
|
packet { int_value: 256 }
|
||||||
|
packet { float_value: 0.5f }
|
||||||
|
packet { bool_value: false }
|
||||||
|
packet { string_value: "string" }
|
||||||
|
packet { string_value: "another string" }
|
||||||
|
packet { int_value: 128 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||||
|
MP_ASSERT_OK(graph.StartRun({}));
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.GetOutputSidePacket("int_packet"));
|
||||||
|
EXPECT_EQ(graph.GetOutputSidePacket("int_packet").ValueOrDie().Get<int>(),
|
||||||
|
256);
|
||||||
|
MP_ASSERT_OK(graph.GetOutputSidePacket("float_packet"));
|
||||||
|
EXPECT_EQ(graph.GetOutputSidePacket("float_packet").ValueOrDie().Get<float>(),
|
||||||
|
0.5f);
|
||||||
|
MP_ASSERT_OK(graph.GetOutputSidePacket("bool_packet"));
|
||||||
|
EXPECT_FALSE(
|
||||||
|
graph.GetOutputSidePacket("bool_packet").ValueOrDie().Get<bool>());
|
||||||
|
MP_ASSERT_OK(graph.GetOutputSidePacket("string_packet"));
|
||||||
|
EXPECT_EQ(graph.GetOutputSidePacket("string_packet")
|
||||||
|
.ValueOrDie()
|
||||||
|
.Get<std::string>(),
|
||||||
|
"string");
|
||||||
|
MP_ASSERT_OK(graph.GetOutputSidePacket("another_string_packet"));
|
||||||
|
EXPECT_EQ(graph.GetOutputSidePacket("another_string_packet")
|
||||||
|
.ValueOrDie()
|
||||||
|
.Get<std::string>(),
|
||||||
|
"another string");
|
||||||
|
MP_ASSERT_OK(graph.GetOutputSidePacket("another_int_packet"));
|
||||||
|
EXPECT_EQ(
|
||||||
|
graph.GetOutputSidePacket("another_int_packet").ValueOrDie().Get<int>(),
|
||||||
|
128);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(ConstantSidePacketCalculatorTest, ProcessingPacketsWithCorrectTagOnly) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||||
|
node {
|
||||||
|
calculator: "ConstantSidePacketCalculator"
|
||||||
|
output_side_packet: "PACKET:0:int_packet"
|
||||||
|
output_side_packet: "no_tag0"
|
||||||
|
output_side_packet: "PACKET:1:float_packet"
|
||||||
|
output_side_packet: "INCORRECT_TAG:0:name1"
|
||||||
|
output_side_packet: "PACKET:2:bool_packet"
|
||||||
|
output_side_packet: "PACKET:3:string_packet"
|
||||||
|
output_side_packet: "no_tag2"
|
||||||
|
output_side_packet: "INCORRECT_TAG:1:name2"
|
||||||
|
options: {
|
||||||
|
[mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
|
||||||
|
packet { int_value: 256 }
|
||||||
|
packet { float_value: 0.5f }
|
||||||
|
packet { bool_value: false }
|
||||||
|
packet { string_value: "string" }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||||
|
MP_ASSERT_OK(graph.StartRun({}));
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.GetOutputSidePacket("int_packet"));
|
||||||
|
EXPECT_EQ(graph.GetOutputSidePacket("int_packet").ValueOrDie().Get<int>(),
|
||||||
|
256);
|
||||||
|
MP_ASSERT_OK(graph.GetOutputSidePacket("float_packet"));
|
||||||
|
EXPECT_EQ(graph.GetOutputSidePacket("float_packet").ValueOrDie().Get<float>(),
|
||||||
|
0.5f);
|
||||||
|
MP_ASSERT_OK(graph.GetOutputSidePacket("bool_packet"));
|
||||||
|
EXPECT_FALSE(
|
||||||
|
graph.GetOutputSidePacket("bool_packet").ValueOrDie().Get<bool>());
|
||||||
|
MP_ASSERT_OK(graph.GetOutputSidePacket("string_packet"));
|
||||||
|
EXPECT_EQ(graph.GetOutputSidePacket("string_packet")
|
||||||
|
.ValueOrDie()
|
||||||
|
.Get<std::string>(),
|
||||||
|
"string");
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(ConstantSidePacketCalculatorTest, IncorrectConfig_MoreOptionsThanPackets) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||||
|
node {
|
||||||
|
calculator: "ConstantSidePacketCalculator"
|
||||||
|
output_side_packet: "PACKET:int_packet"
|
||||||
|
options: {
|
||||||
|
[mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
|
||||||
|
packet { int_value: 256 }
|
||||||
|
packet { float_value: 0.5f }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorGraph graph;
|
||||||
|
EXPECT_FALSE(graph.Initialize(graph_config).ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(ConstantSidePacketCalculatorTest, IncorrectConfig_MorePacketsThanOptions) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||||
|
node {
|
||||||
|
calculator: "ConstantSidePacketCalculator"
|
||||||
|
output_side_packet: "PACKET:0:int_packet"
|
||||||
|
output_side_packet: "PACKET:1:float_packet"
|
||||||
|
options: {
|
||||||
|
[mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
|
||||||
|
packet { int_value: 256 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorGraph graph;
|
||||||
|
EXPECT_FALSE(graph.Initialize(graph_config).ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -16,9 +16,11 @@
|
|||||||
|
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/formats/classification.pb.h"
|
||||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||||
#include "mediapipe/framework/formats/rect.pb.h"
|
#include "mediapipe/framework/formats/rect.pb.h"
|
||||||
#include "mediapipe/util/render_data.pb.h"
|
#include "mediapipe/util/render_data.pb.h"
|
||||||
|
#include "tensorflow/lite/interpreter.h"
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
@@ -37,4 +39,11 @@ typedef EndLoopCalculator<std::vector<::mediapipe::RenderData>>
|
|||||||
EndLoopRenderDataCalculator;
|
EndLoopRenderDataCalculator;
|
||||||
REGISTER_CALCULATOR(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
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -12,25 +12,17 @@
|
|||||||
// See the License for the specific language governing permissions and
|
// See the License for the specific language governing permissions and
|
||||||
// limitations under the License.
|
// limitations under the License.
|
||||||
|
|
||||||
#include <cstdlib>
|
#include "mediapipe/calculators/core/packet_resampler_calculator.h"
|
||||||
#include <memory>
|
|
||||||
#include <string>
|
|
||||||
|
|
||||||
#include "absl/strings/str_cat.h"
|
#include <memory>
|
||||||
#include "mediapipe/calculators/core/packet_resampler_calculator.pb.h"
|
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
|
||||||
#include "mediapipe/framework/collection_item_id.h"
|
|
||||||
#include "mediapipe/framework/deps/mathutil.h"
|
|
||||||
#include "mediapipe/framework/deps/random_base.h"
|
|
||||||
#include "mediapipe/framework/formats/video_stream_header.h"
|
|
||||||
#include "mediapipe/framework/port/integral_types.h"
|
|
||||||
#include "mediapipe/framework/port/logging.h"
|
|
||||||
#include "mediapipe/framework/port/ret_check.h"
|
|
||||||
#include "mediapipe/framework/port/status.h"
|
|
||||||
#include "mediapipe/framework/port/status_macros.h"
|
|
||||||
#include "mediapipe/framework/tool/options_util.h"
|
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
|
// Reflect an integer against the lower and upper bound of an interval.
|
||||||
|
int64 ReflectBetween(int64 ts, int64 ts_min, int64 ts_max) {
|
||||||
|
if (ts < ts_min) return 2 * ts_min - ts - 1;
|
||||||
|
if (ts >= ts_max) return 2 * ts_max - ts - 1;
|
||||||
|
return ts;
|
||||||
|
}
|
||||||
|
|
||||||
// Creates a secure random number generator for use in ProcessWithJitter.
|
// Creates a secure random number generator for use in ProcessWithJitter.
|
||||||
// If no secure random number generator can be constructed, the jitter
|
// If no secure random number generator can be constructed, the jitter
|
||||||
@@ -45,120 +37,7 @@ std::unique_ptr<RandomBase> CreateSecureRandom(const std::string& seed) {
|
|||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
// This calculator is used to normalize the frequency of the packets
|
|
||||||
// out of a stream. Given a desired frame rate, packets are going to be
|
|
||||||
// removed or added to achieve it.
|
|
||||||
//
|
|
||||||
// The jitter feature is disabled by default. To enable it, you need to
|
|
||||||
// implement CreateSecureRandom(const std::string&).
|
|
||||||
//
|
|
||||||
// The data stream may be either specified as the only stream (by index)
|
|
||||||
// or as the stream with tag "DATA".
|
|
||||||
//
|
|
||||||
// The input and output streams may be accompanied by a VIDEO_HEADER
|
|
||||||
// stream. This stream includes a VideoHeader at Timestamp::PreStream().
|
|
||||||
// The input VideoHeader on the VIDEO_HEADER stream will always be updated
|
|
||||||
// with the resampler frame rate no matter what the options value for
|
|
||||||
// output_header is before being output on the output VIDEO_HEADER stream.
|
|
||||||
// If the input VideoHeader is not available, then only the frame rate
|
|
||||||
// value will be set in the output.
|
|
||||||
//
|
|
||||||
// Related:
|
|
||||||
// packet_downsampler_calculator.cc: skips packets regardless of timestamps.
|
|
||||||
class PacketResamplerCalculator : public CalculatorBase {
|
|
||||||
public:
|
|
||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
|
||||||
|
|
||||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
|
||||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
|
||||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
|
||||||
|
|
||||||
private:
|
|
||||||
// Calculates the first sampled timestamp that incorporates a jittering
|
|
||||||
// offset.
|
|
||||||
void InitializeNextOutputTimestampWithJitter();
|
|
||||||
// Calculates the next sampled timestamp that incorporates a jittering offset.
|
|
||||||
void UpdateNextOutputTimestampWithJitter();
|
|
||||||
|
|
||||||
// Logic for Process() when jitter_ != 0.0.
|
|
||||||
::mediapipe::Status ProcessWithJitter(CalculatorContext* cc);
|
|
||||||
|
|
||||||
// Logic for Process() when jitter_ == 0.0.
|
|
||||||
::mediapipe::Status ProcessWithoutJitter(CalculatorContext* cc);
|
|
||||||
|
|
||||||
// Given the current count of periods that have passed, this returns
|
|
||||||
// the next valid timestamp of the middle point of the next period:
|
|
||||||
// if count is 0, it returns the first_timestamp_.
|
|
||||||
// if count is 1, it returns the first_timestamp_ + period (corresponding
|
|
||||||
// to the first tick using exact fps)
|
|
||||||
// e.g. for frame_rate=30 and first_timestamp_=0:
|
|
||||||
// 0: 0
|
|
||||||
// 1: 33333
|
|
||||||
// 2: 66667
|
|
||||||
// 3: 100000
|
|
||||||
//
|
|
||||||
// Can only be used if jitter_ equals zero.
|
|
||||||
Timestamp PeriodIndexToTimestamp(int64 index) const;
|
|
||||||
|
|
||||||
// Given a Timestamp, finds the closest sync Timestamp based on
|
|
||||||
// first_timestamp_ and the desired fps.
|
|
||||||
//
|
|
||||||
// Can only be used if jitter_ equals zero.
|
|
||||||
int64 TimestampToPeriodIndex(Timestamp timestamp) const;
|
|
||||||
|
|
||||||
// Outputs a packet if it is in range (start_time_, end_time_).
|
|
||||||
void OutputWithinLimits(CalculatorContext* cc, const Packet& packet) const;
|
|
||||||
|
|
||||||
// The timestamp of the first packet received.
|
|
||||||
Timestamp first_timestamp_;
|
|
||||||
|
|
||||||
// Number of frames per second (desired output frequency).
|
|
||||||
double frame_rate_;
|
|
||||||
|
|
||||||
// Inverse of frame_rate_.
|
|
||||||
int64 frame_time_usec_;
|
|
||||||
|
|
||||||
// Number of periods that have passed (= #packets sent to the output).
|
|
||||||
//
|
|
||||||
// Can only be used if jitter_ equals zero.
|
|
||||||
int64 period_count_;
|
|
||||||
|
|
||||||
// The last packet that was received.
|
|
||||||
Packet last_packet_;
|
|
||||||
|
|
||||||
VideoHeader video_header_;
|
|
||||||
// The "DATA" input stream.
|
|
||||||
CollectionItemId input_data_id_;
|
|
||||||
// The "DATA" output stream.
|
|
||||||
CollectionItemId output_data_id_;
|
|
||||||
|
|
||||||
// Indicator whether to flush last packet even if its timestamp is greater
|
|
||||||
// than the final stream timestamp. Set to false when jitter_ is non-zero.
|
|
||||||
bool flush_last_packet_;
|
|
||||||
|
|
||||||
// Jitter-related variables.
|
|
||||||
std::unique_ptr<RandomBase> random_;
|
|
||||||
double jitter_ = 0.0;
|
|
||||||
Timestamp next_output_timestamp_;
|
|
||||||
|
|
||||||
// If specified, output timestamps are aligned with base_timestamp.
|
|
||||||
// Otherwise, they are aligned with the first input timestamp.
|
|
||||||
Timestamp base_timestamp_;
|
|
||||||
|
|
||||||
// If specified, only outputs at/after start_time are included.
|
|
||||||
Timestamp start_time_;
|
|
||||||
|
|
||||||
// If specified, only outputs before end_time are included.
|
|
||||||
Timestamp end_time_;
|
|
||||||
|
|
||||||
// If set, the output timestamps nearest to start_time and end_time
|
|
||||||
// are included in the output, even if the nearest timestamp is not
|
|
||||||
// between start_time and end_time.
|
|
||||||
bool round_limits_;
|
|
||||||
};
|
|
||||||
|
|
||||||
REGISTER_CALCULATOR(PacketResamplerCalculator);
|
REGISTER_CALCULATOR(PacketResamplerCalculator);
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
// Returns a TimestampDiff (assuming microseconds) corresponding to the
|
// Returns a TimestampDiff (assuming microseconds) corresponding to the
|
||||||
// given time in seconds.
|
// given time in seconds.
|
||||||
@@ -209,6 +88,7 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
|||||||
|
|
||||||
flush_last_packet_ = resampler_options.flush_last_packet();
|
flush_last_packet_ = resampler_options.flush_last_packet();
|
||||||
jitter_ = resampler_options.jitter();
|
jitter_ = resampler_options.jitter();
|
||||||
|
jitter_with_reflection_ = resampler_options.jitter_with_reflection();
|
||||||
|
|
||||||
input_data_id_ = cc->Inputs().GetId("DATA", 0);
|
input_data_id_ = cc->Inputs().GetId("DATA", 0);
|
||||||
if (!input_data_id_.IsValid()) {
|
if (!input_data_id_.IsValid()) {
|
||||||
@@ -239,6 +119,8 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
|||||||
<< Timestamp::kTimestampUnitsPerSecond;
|
<< Timestamp::kTimestampUnitsPerSecond;
|
||||||
|
|
||||||
frame_time_usec_ = static_cast<int64>(1000000.0 / frame_rate_);
|
frame_time_usec_ = static_cast<int64>(1000000.0 / frame_rate_);
|
||||||
|
jitter_usec_ = static_cast<int64>(1000000.0 * jitter_ / frame_rate_);
|
||||||
|
RET_CHECK_LE(jitter_usec_, frame_time_usec_);
|
||||||
|
|
||||||
video_header_.frame_rate = frame_rate_;
|
video_header_.frame_rate = frame_rate_;
|
||||||
|
|
||||||
@@ -279,7 +161,10 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
|||||||
"SecureRandom is not available. With \"jitter\" specified, "
|
"SecureRandom is not available. With \"jitter\" specified, "
|
||||||
"PacketResamplerCalculator processing cannot proceed.");
|
"PacketResamplerCalculator processing cannot proceed.");
|
||||||
}
|
}
|
||||||
|
packet_reservoir_random_ = CreateSecureRandom(seed);
|
||||||
}
|
}
|
||||||
|
packet_reservoir_ =
|
||||||
|
std::make_unique<PacketReservoir>(packet_reservoir_random_.get());
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -294,6 +179,14 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
if (jitter_ != 0.0 && random_ != nullptr) {
|
if (jitter_ != 0.0 && random_ != nullptr) {
|
||||||
|
// Packet reservior is used to make sure there's an output for every period,
|
||||||
|
// e.g. partial period at the end of the stream.
|
||||||
|
if (packet_reservoir_->IsEnabled() &&
|
||||||
|
(first_timestamp_ == Timestamp::Unset() ||
|
||||||
|
(cc->InputTimestamp() - next_output_timestamp_min_).Value() >= 0)) {
|
||||||
|
auto curr_packet = cc->Inputs().Get(input_data_id_).Value();
|
||||||
|
packet_reservoir_->AddSample(curr_packet);
|
||||||
|
}
|
||||||
MP_RETURN_IF_ERROR(ProcessWithJitter(cc));
|
MP_RETURN_IF_ERROR(ProcessWithJitter(cc));
|
||||||
} else {
|
} else {
|
||||||
MP_RETURN_IF_ERROR(ProcessWithoutJitter(cc));
|
MP_RETURN_IF_ERROR(ProcessWithoutJitter(cc));
|
||||||
@@ -303,11 +196,34 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void PacketResamplerCalculator::InitializeNextOutputTimestampWithJitter() {
|
void PacketResamplerCalculator::InitializeNextOutputTimestampWithJitter() {
|
||||||
|
next_output_timestamp_min_ = first_timestamp_;
|
||||||
|
if (jitter_with_reflection_) {
|
||||||
|
next_output_timestamp_ =
|
||||||
|
first_timestamp_ + random_->UnbiasedUniform64(frame_time_usec_);
|
||||||
|
return;
|
||||||
|
}
|
||||||
next_output_timestamp_ =
|
next_output_timestamp_ =
|
||||||
first_timestamp_ + frame_time_usec_ * random_->RandFloat();
|
first_timestamp_ + frame_time_usec_ * random_->RandFloat();
|
||||||
}
|
}
|
||||||
|
|
||||||
void PacketResamplerCalculator::UpdateNextOutputTimestampWithJitter() {
|
void PacketResamplerCalculator::UpdateNextOutputTimestampWithJitter() {
|
||||||
|
packet_reservoir_->Clear();
|
||||||
|
if (jitter_with_reflection_) {
|
||||||
|
next_output_timestamp_min_ += frame_time_usec_;
|
||||||
|
Timestamp next_output_timestamp_max_ =
|
||||||
|
next_output_timestamp_min_ + frame_time_usec_;
|
||||||
|
|
||||||
|
next_output_timestamp_ += frame_time_usec_ +
|
||||||
|
random_->UnbiasedUniform64(2 * jitter_usec_ + 1) -
|
||||||
|
jitter_usec_;
|
||||||
|
next_output_timestamp_ = Timestamp(ReflectBetween(
|
||||||
|
next_output_timestamp_.Value(), next_output_timestamp_min_.Value(),
|
||||||
|
next_output_timestamp_max_.Value()));
|
||||||
|
CHECK_GE(next_output_timestamp_, next_output_timestamp_min_);
|
||||||
|
CHECK_LT(next_output_timestamp_, next_output_timestamp_max_);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
packet_reservoir_->Disable();
|
||||||
next_output_timestamp_ +=
|
next_output_timestamp_ +=
|
||||||
frame_time_usec_ *
|
frame_time_usec_ *
|
||||||
((1.0 - jitter_) + 2.0 * jitter_ * random_->RandFloat());
|
((1.0 - jitter_) + 2.0 * jitter_ * random_->RandFloat());
|
||||||
@@ -330,22 +246,27 @@ void PacketResamplerCalculator::UpdateNextOutputTimestampWithJitter() {
|
|||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
LOG_IF(WARNING, frame_time_usec_ <
|
if (frame_time_usec_ <
|
||||||
(cc->InputTimestamp() - last_packet_.Timestamp()).Value())
|
(cc->InputTimestamp() - last_packet_.Timestamp()).Value()) {
|
||||||
<< "Adding jitter is meaningless when upsampling.";
|
LOG_FIRST_N(WARNING, 2)
|
||||||
|
<< "Adding jitter is not very useful when upsampling.";
|
||||||
|
}
|
||||||
|
|
||||||
const int64 curr_diff =
|
while (true) {
|
||||||
(next_output_timestamp_ - cc->InputTimestamp()).Value();
|
const int64 last_diff =
|
||||||
const int64 last_diff =
|
(next_output_timestamp_ - last_packet_.Timestamp()).Value();
|
||||||
(next_output_timestamp_ - last_packet_.Timestamp()).Value();
|
RET_CHECK_GT(last_diff, 0);
|
||||||
if (curr_diff * last_diff > 0) {
|
const int64 curr_diff =
|
||||||
return ::mediapipe::OkStatus();
|
(next_output_timestamp_ - cc->InputTimestamp()).Value();
|
||||||
|
if (curr_diff > 0) {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
OutputWithinLimits(cc, (std::abs(curr_diff) > last_diff
|
||||||
|
? last_packet_
|
||||||
|
: cc->Inputs().Get(input_data_id_).Value())
|
||||||
|
.At(next_output_timestamp_));
|
||||||
|
UpdateNextOutputTimestampWithJitter();
|
||||||
}
|
}
|
||||||
OutputWithinLimits(cc, (std::abs(curr_diff) > std::abs(last_diff)
|
|
||||||
? last_packet_
|
|
||||||
: cc->Inputs().Get(input_data_id_).Value())
|
|
||||||
.At(next_output_timestamp_));
|
|
||||||
UpdateNextOutputTimestampWithJitter();
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -426,6 +347,9 @@ void PacketResamplerCalculator::UpdateNextOutputTimestampWithJitter() {
|
|||||||
OutputWithinLimits(cc,
|
OutputWithinLimits(cc,
|
||||||
last_packet_.At(PeriodIndexToTimestamp(period_count_)));
|
last_packet_.At(PeriodIndexToTimestamp(period_count_)));
|
||||||
}
|
}
|
||||||
|
if (!packet_reservoir_->IsEmpty()) {
|
||||||
|
OutputWithinLimits(cc, packet_reservoir_->GetSample());
|
||||||
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,205 @@
|
|||||||
|
#ifndef MEDIAPIPE_CALCULATORS_CORE_PACKET_RESAMPLER_CALCULATOR_H_
|
||||||
|
#define MEDIAPIPE_CALCULATORS_CORE_PACKET_RESAMPLER_CALCULATOR_H_
|
||||||
|
|
||||||
|
#include <cstdlib>
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
#include "absl/strings/str_cat.h"
|
||||||
|
#include "mediapipe/calculators/core/packet_resampler_calculator.pb.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/collection_item_id.h"
|
||||||
|
#include "mediapipe/framework/deps/mathutil.h"
|
||||||
|
#include "mediapipe/framework/deps/random_base.h"
|
||||||
|
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||||
|
#include "mediapipe/framework/port/integral_types.h"
|
||||||
|
#include "mediapipe/framework/port/logging.h"
|
||||||
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
#include "mediapipe/framework/port/status_macros.h"
|
||||||
|
#include "mediapipe/framework/tool/options_util.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
class PacketReservoir {
|
||||||
|
public:
|
||||||
|
PacketReservoir(RandomBase* rng) : rng_(rng) {}
|
||||||
|
// Replace candidate with current packet with 1/count_ probability.
|
||||||
|
void AddSample(Packet sample) {
|
||||||
|
if (rng_->UnbiasedUniform(++count_) == 0) {
|
||||||
|
reservoir_ = sample;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
bool IsEnabled() { return rng_ && enabled_; }
|
||||||
|
void Disable() {
|
||||||
|
if (enabled_) enabled_ = false;
|
||||||
|
}
|
||||||
|
void Clear() { count_ = 0; }
|
||||||
|
bool IsEmpty() { return count_ == 0; }
|
||||||
|
Packet GetSample() { return reservoir_; }
|
||||||
|
|
||||||
|
private:
|
||||||
|
RandomBase* rng_;
|
||||||
|
bool enabled_ = true;
|
||||||
|
int32 count_ = 0;
|
||||||
|
Packet reservoir_;
|
||||||
|
};
|
||||||
|
|
||||||
|
// This calculator is used to normalize the frequency of the packets
|
||||||
|
// out of a stream. Given a desired frame rate, packets are going to be
|
||||||
|
// removed or added to achieve it.
|
||||||
|
//
|
||||||
|
// If jitter_ is specified:
|
||||||
|
// - The first packet is chosen randomly (uniform distribution) among frames
|
||||||
|
// that correspond to timestamps [0, 1/frame_rate). Let the chosen packet
|
||||||
|
// correspond to timestamp t.
|
||||||
|
// - The next packet is chosen randomly (uniform distribution) among frames
|
||||||
|
// that correspond to [t+(1-jitter)/frame_rate, t+(1+jitter)/frame_rate].
|
||||||
|
// - if jitter_with_reflection_ is true, the timestamp will be reflected
|
||||||
|
// against the boundaries of [t_0 + (k-1)/frame_rate, t_0 + k/frame_rate)
|
||||||
|
// so that its marginal distribution is uniform within this interval.
|
||||||
|
// In the formula, t_0 is the timestamp of the first sampled
|
||||||
|
// packet, and the k is the packet index.
|
||||||
|
// See paper (https://arxiv.org/abs/2002.01147) for details.
|
||||||
|
// - t is updated and the process is repeated.
|
||||||
|
// - Note that seed is specified as input side packet for reproducibility of
|
||||||
|
// the resampling. For Cloud ML Video Intelligence API, the hash of the
|
||||||
|
// input video should serve this purpose. For YouTube, either video ID or
|
||||||
|
// content hex ID of the input video should do.
|
||||||
|
//
|
||||||
|
// If jitter_ is not specified:
|
||||||
|
// - The first packet defines the first_timestamp of the output stream,
|
||||||
|
// so it is always emitted.
|
||||||
|
// - If more packets are emitted, they will have timestamp equal to
|
||||||
|
// round(first_timestamp + k * period) , where k is a positive
|
||||||
|
// integer and the period is defined by the frame rate.
|
||||||
|
// Example: first_timestamp=0, fps=30, then the output stream
|
||||||
|
// will have timestamps: 0, 33333, 66667, 100000, etc...
|
||||||
|
// - The packets selected for the output stream are the ones closer
|
||||||
|
// to the exact middle point (33333.33, 66666.67 in our previous
|
||||||
|
// example). In case of ties, later packets are chosen.
|
||||||
|
// - 'Empty' periods happen when there are no packets for a long time
|
||||||
|
// (greater than a period). In this case, we send a copy of the last
|
||||||
|
// packet received before the empty period.
|
||||||
|
// The jitter feature is disabled by default. To enable it, you need to
|
||||||
|
// implement CreateSecureRandom(const std::string&).
|
||||||
|
//
|
||||||
|
// The data stream may be either specified as the only stream (by index)
|
||||||
|
// or as the stream with tag "DATA".
|
||||||
|
//
|
||||||
|
// The input and output streams may be accompanied by a VIDEO_HEADER
|
||||||
|
// stream. This stream includes a VideoHeader at Timestamp::PreStream().
|
||||||
|
// The input VideoHeader on the VIDEO_HEADER stream will always be updated
|
||||||
|
// with the resampler frame rate no matter what the options value for
|
||||||
|
// output_header is before being output on the output VIDEO_HEADER stream.
|
||||||
|
// If the input VideoHeader is not available, then only the frame rate
|
||||||
|
// value will be set in the output.
|
||||||
|
//
|
||||||
|
// Related:
|
||||||
|
// packet_downsampler_calculator.cc: skips packets regardless of timestamps.
|
||||||
|
class PacketResamplerCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||||
|
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||||
|
|
||||||
|
private:
|
||||||
|
// Calculates the first sampled timestamp that incorporates a jittering
|
||||||
|
// offset.
|
||||||
|
void InitializeNextOutputTimestampWithJitter();
|
||||||
|
// Calculates the next sampled timestamp that incorporates a jittering offset.
|
||||||
|
void UpdateNextOutputTimestampWithJitter();
|
||||||
|
|
||||||
|
// Logic for Process() when jitter_ != 0.0.
|
||||||
|
::mediapipe::Status ProcessWithJitter(CalculatorContext* cc);
|
||||||
|
|
||||||
|
// Logic for Process() when jitter_ == 0.0.
|
||||||
|
::mediapipe::Status ProcessWithoutJitter(CalculatorContext* cc);
|
||||||
|
|
||||||
|
// Given the current count of periods that have passed, this returns
|
||||||
|
// the next valid timestamp of the middle point of the next period:
|
||||||
|
// if count is 0, it returns the first_timestamp_.
|
||||||
|
// if count is 1, it returns the first_timestamp_ + period (corresponding
|
||||||
|
// to the first tick using exact fps)
|
||||||
|
// e.g. for frame_rate=30 and first_timestamp_=0:
|
||||||
|
// 0: 0
|
||||||
|
// 1: 33333
|
||||||
|
// 2: 66667
|
||||||
|
// 3: 100000
|
||||||
|
//
|
||||||
|
// Can only be used if jitter_ equals zero.
|
||||||
|
Timestamp PeriodIndexToTimestamp(int64 index) const;
|
||||||
|
|
||||||
|
// Given a Timestamp, finds the closest sync Timestamp based on
|
||||||
|
// first_timestamp_ and the desired fps.
|
||||||
|
//
|
||||||
|
// Can only be used if jitter_ equals zero.
|
||||||
|
int64 TimestampToPeriodIndex(Timestamp timestamp) const;
|
||||||
|
|
||||||
|
// Outputs a packet if it is in range (start_time_, end_time_).
|
||||||
|
void OutputWithinLimits(CalculatorContext* cc, const Packet& packet) const;
|
||||||
|
|
||||||
|
// The timestamp of the first packet received.
|
||||||
|
Timestamp first_timestamp_;
|
||||||
|
|
||||||
|
// Number of frames per second (desired output frequency).
|
||||||
|
double frame_rate_;
|
||||||
|
|
||||||
|
// Inverse of frame_rate_.
|
||||||
|
int64 frame_time_usec_;
|
||||||
|
|
||||||
|
// Number of periods that have passed (= #packets sent to the output).
|
||||||
|
//
|
||||||
|
// Can only be used if jitter_ equals zero.
|
||||||
|
int64 period_count_;
|
||||||
|
|
||||||
|
// The last packet that was received.
|
||||||
|
Packet last_packet_;
|
||||||
|
|
||||||
|
VideoHeader video_header_;
|
||||||
|
// The "DATA" input stream.
|
||||||
|
CollectionItemId input_data_id_;
|
||||||
|
// The "DATA" output stream.
|
||||||
|
CollectionItemId output_data_id_;
|
||||||
|
|
||||||
|
// Indicator whether to flush last packet even if its timestamp is greater
|
||||||
|
// than the final stream timestamp. Set to false when jitter_ is non-zero.
|
||||||
|
bool flush_last_packet_;
|
||||||
|
|
||||||
|
// Jitter-related variables.
|
||||||
|
std::unique_ptr<RandomBase> random_;
|
||||||
|
double jitter_ = 0.0;
|
||||||
|
bool jitter_with_reflection_;
|
||||||
|
int64 jitter_usec_;
|
||||||
|
Timestamp next_output_timestamp_;
|
||||||
|
// If jittering_with_reflection_ is true, next_output_timestamp_ will be
|
||||||
|
// kept within the interval
|
||||||
|
// [next_output_timestamp_min_, next_output_timestamp_min_ + frame_time_usec_)
|
||||||
|
Timestamp next_output_timestamp_min_;
|
||||||
|
|
||||||
|
// If specified, output timestamps are aligned with base_timestamp.
|
||||||
|
// Otherwise, they are aligned with the first input timestamp.
|
||||||
|
Timestamp base_timestamp_;
|
||||||
|
|
||||||
|
// If specified, only outputs at/after start_time are included.
|
||||||
|
Timestamp start_time_;
|
||||||
|
|
||||||
|
// If specified, only outputs before end_time are included.
|
||||||
|
Timestamp end_time_;
|
||||||
|
|
||||||
|
// If set, the output timestamps nearest to start_time and end_time
|
||||||
|
// are included in the output, even if the nearest timestamp is not
|
||||||
|
// between start_time and end_time.W
|
||||||
|
bool round_limits_;
|
||||||
|
|
||||||
|
// packet reservior used for sampling random packet out of partial
|
||||||
|
// period when jitter is enabled
|
||||||
|
std::unique_ptr<PacketReservoir> packet_reservoir_;
|
||||||
|
// random number generator used in packet_reservior_.
|
||||||
|
std::unique_ptr<RandomBase> packet_reservoir_random_;
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
|
#endif // MEDIAPIPE_CALCULATORS_CORE_PACKET_RESAMPLER_CALCULATOR_H_
|
||||||
@@ -66,6 +66,7 @@ message PacketResamplerCalculatorOptions {
|
|||||||
// pseudo-random number generator does its job and the number of frames is
|
// pseudo-random number generator does its job and the number of frames is
|
||||||
// sufficiently large, the average frame rate will be close to this value.
|
// sufficiently large, the average frame rate will be close to this value.
|
||||||
optional double jitter = 4;
|
optional double jitter = 4;
|
||||||
|
optional bool jitter_with_reflection = 9 [default = false];
|
||||||
|
|
||||||
// If specified, output timestamps are aligned with base_timestamp.
|
// If specified, output timestamps are aligned with base_timestamp.
|
||||||
// Otherwise, they are aligned with the first input timestamp.
|
// Otherwise, they are aligned with the first input timestamp.
|
||||||
|
|||||||
@@ -12,6 +12,8 @@
|
|||||||
// See the License for the specific language governing permissions and
|
// See the License for the specific language governing permissions and
|
||||||
// limitations under the License.
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/core/packet_resampler_calculator.h"
|
||||||
|
|
||||||
#include <memory>
|
#include <memory>
|
||||||
#include <string>
|
#include <string>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
@@ -29,7 +31,6 @@
|
|||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
|
|
||||||
// A simple version of CalculatorRunner with built-in convenience
|
// A simple version of CalculatorRunner with built-in convenience
|
||||||
// methods for setting inputs from a vector and checking outputs
|
// methods for setting inputs from a vector and checking outputs
|
||||||
// against expected outputs (both timestamps and contents).
|
// against expected outputs (both timestamps and contents).
|
||||||
|
|||||||
@@ -0,0 +1,304 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
//
|
||||||
|
// Declaration of PacketThinnerCalculator.
|
||||||
|
|
||||||
|
#include <cmath> // for ceil
|
||||||
|
#include <memory>
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/core/packet_thinner_calculator.pb.h"
|
||||||
|
#include "mediapipe/framework/calculator_context.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||||
|
#include "mediapipe/framework/port/integral_types.h"
|
||||||
|
#include "mediapipe/framework/port/logging.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
const double kTimebaseUs = 1000000; // Microseconds.
|
||||||
|
const char* const kPeriodTag = "PERIOD";
|
||||||
|
} // namespace
|
||||||
|
|
||||||
|
// This calculator is used to thin an input stream of Packets.
|
||||||
|
// An example application would be to sample decoded frames of video
|
||||||
|
// at a coarser temporal resolution. Unless otherwise stated, all
|
||||||
|
// timestamps are in units of microseconds.
|
||||||
|
//
|
||||||
|
// Thinning can be accomplished in one of two ways:
|
||||||
|
// 1) asynchronous thinning (known below as async):
|
||||||
|
// Algorithm does not rely on a master clock and is parameterized only
|
||||||
|
// by a single option -- the period. Once a packet is emitted, the
|
||||||
|
// thinner will discard subsequent packets for the duration of the period
|
||||||
|
// [Analogous to a refractory period during which packet emission is
|
||||||
|
// suppressed.]
|
||||||
|
// Packets arriving before start_time are discarded, as are packets
|
||||||
|
// arriving at or after end_time.
|
||||||
|
// 2) synchronous thinning (known below as sync):
|
||||||
|
// There are two variants of this algorithm, both parameterized by a
|
||||||
|
// start_time and a period. As in (1), packets arriving before start_time
|
||||||
|
// or at/after end_time are discarded. Otherwise, at most one packet is
|
||||||
|
// emitted during a period, centered at timestamps generated by the
|
||||||
|
// expression:
|
||||||
|
// start_time + i * period [where i is a non-negative integer]
|
||||||
|
// During each period, the packet closest to the generated timestamp is
|
||||||
|
// emitted (latest in the case of ties). In the first variant
|
||||||
|
// (sync_output_timestamps = true), the emitted packet is output at the
|
||||||
|
// generated timestamp. In the second variant, the packet is output at
|
||||||
|
// its original timestamp. Both variants emit exactly the same packets,
|
||||||
|
// but at different timestamps.
|
||||||
|
//
|
||||||
|
// Thinning period can be provided in the calculator options or via a
|
||||||
|
// side packet with the tag "PERIOD".
|
||||||
|
//
|
||||||
|
// Example config:
|
||||||
|
// node {
|
||||||
|
// calculator: "PacketThinnerCalculator"
|
||||||
|
// input_stream: "signal"
|
||||||
|
// output_stream: "output"
|
||||||
|
// options {
|
||||||
|
// [mediapipe.PacketThinnerCalculatorOptions.ext] {
|
||||||
|
// thinner_type: SYNC
|
||||||
|
// period: 10
|
||||||
|
// sync_output_timestamps: true
|
||||||
|
// update_frame_rate: false
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
class PacketThinnerCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
PacketThinnerCalculator() {}
|
||||||
|
~PacketThinnerCalculator() override {}
|
||||||
|
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
cc->Inputs().Index(0).SetAny();
|
||||||
|
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
|
||||||
|
if (cc->InputSidePackets().HasTag(kPeriodTag)) {
|
||||||
|
cc->InputSidePackets().Tag(kPeriodTag).Set<int64>();
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||||
|
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
|
if (cc->InputTimestamp() < start_time_) {
|
||||||
|
return ::mediapipe::OkStatus(); // Drop packets before start_time_.
|
||||||
|
} else if (cc->InputTimestamp() >= end_time_) {
|
||||||
|
if (!cc->Outputs().Index(0).IsClosed()) {
|
||||||
|
cc->Outputs()
|
||||||
|
.Index(0)
|
||||||
|
.Close(); // No more Packets will be output after end_time_.
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
} else {
|
||||||
|
return thinner_type_ == PacketThinnerCalculatorOptions::ASYNC
|
||||||
|
? AsyncThinnerProcess(cc)
|
||||||
|
: SyncThinnerProcess(cc);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
// Implementation of ASYNC and SYNC versions of thinner algorithm.
|
||||||
|
::mediapipe::Status AsyncThinnerProcess(CalculatorContext* cc);
|
||||||
|
::mediapipe::Status SyncThinnerProcess(CalculatorContext* cc);
|
||||||
|
|
||||||
|
// Cached option.
|
||||||
|
PacketThinnerCalculatorOptions::ThinnerType thinner_type_;
|
||||||
|
|
||||||
|
// Given a Timestamp, finds the closest sync Timestamp
|
||||||
|
// based on start_time_ and period_. This can be earlier or
|
||||||
|
// later than given Timestamp, but is guaranteed to be within
|
||||||
|
// half a period_.
|
||||||
|
Timestamp NearestSyncTimestamp(Timestamp now) const;
|
||||||
|
|
||||||
|
// Cached option used by both async and sync thinners.
|
||||||
|
TimestampDiff period_; // Interval during which only one packet is emitted.
|
||||||
|
Timestamp start_time_; // Cached option - default Timestamp::Min()
|
||||||
|
Timestamp end_time_; // Cached option - default Timestamp::Max()
|
||||||
|
|
||||||
|
// Only used by async thinner:
|
||||||
|
Timestamp next_valid_timestamp_; // Suppress packets until this timestamp.
|
||||||
|
|
||||||
|
// Only used by sync thinner:
|
||||||
|
Packet saved_packet_; // Best packet not yet emitted.
|
||||||
|
bool sync_output_timestamps_; // Cached option.
|
||||||
|
};
|
||||||
|
REGISTER_CALCULATOR(PacketThinnerCalculator);
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
TimestampDiff abs(TimestampDiff t) { return t < 0 ? -t : t; }
|
||||||
|
} // namespace
|
||||||
|
|
||||||
|
::mediapipe::Status PacketThinnerCalculator::Open(CalculatorContext* cc) {
|
||||||
|
auto& options = cc->Options<PacketThinnerCalculatorOptions>();
|
||||||
|
|
||||||
|
thinner_type_ = options.thinner_type();
|
||||||
|
// This check enables us to assume only two thinner types exist in Process()
|
||||||
|
CHECK(thinner_type_ == PacketThinnerCalculatorOptions::ASYNC ||
|
||||||
|
thinner_type_ == PacketThinnerCalculatorOptions::SYNC)
|
||||||
|
<< "Unsupported thinner type.";
|
||||||
|
|
||||||
|
if (thinner_type_ == PacketThinnerCalculatorOptions::ASYNC) {
|
||||||
|
// ASYNC thinner outputs packets with the same timestamp as their input so
|
||||||
|
// its safe to SetOffset(0). SYNC thinner manipulates timestamps of its
|
||||||
|
// output so we don't do this for that case.
|
||||||
|
cc->SetOffset(0);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (cc->InputSidePackets().HasTag(kPeriodTag)) {
|
||||||
|
period_ =
|
||||||
|
TimestampDiff(cc->InputSidePackets().Tag(kPeriodTag).Get<int64>());
|
||||||
|
} else {
|
||||||
|
period_ = TimestampDiff(options.period());
|
||||||
|
}
|
||||||
|
CHECK_LT(TimestampDiff(0), period_) << "Specified period must be positive.";
|
||||||
|
|
||||||
|
if (options.has_start_time()) {
|
||||||
|
start_time_ = Timestamp(options.start_time());
|
||||||
|
} else if (thinner_type_ == PacketThinnerCalculatorOptions::ASYNC) {
|
||||||
|
start_time_ = Timestamp::Min();
|
||||||
|
} else {
|
||||||
|
start_time_ = Timestamp(0);
|
||||||
|
}
|
||||||
|
|
||||||
|
end_time_ =
|
||||||
|
options.has_end_time() ? Timestamp(options.end_time()) : Timestamp::Max();
|
||||||
|
CHECK_LT(start_time_, end_time_)
|
||||||
|
<< "Invalid PacketThinner: start_time must be earlier than end_time";
|
||||||
|
|
||||||
|
sync_output_timestamps_ = options.sync_output_timestamps();
|
||||||
|
|
||||||
|
next_valid_timestamp_ = start_time_;
|
||||||
|
// Drop packets until this time.
|
||||||
|
cc->Outputs().Index(0).SetNextTimestampBound(start_time_);
|
||||||
|
|
||||||
|
if (!cc->Inputs().Index(0).Header().IsEmpty()) {
|
||||||
|
if (options.update_frame_rate()) {
|
||||||
|
const VideoHeader& video_header =
|
||||||
|
cc->Inputs().Index(0).Header().Get<VideoHeader>();
|
||||||
|
double new_frame_rate;
|
||||||
|
if (thinner_type_ == PacketThinnerCalculatorOptions::ASYNC) {
|
||||||
|
new_frame_rate =
|
||||||
|
video_header.frame_rate /
|
||||||
|
ceil(video_header.frame_rate * options.period() / kTimebaseUs);
|
||||||
|
} else {
|
||||||
|
const double sampling_rate = kTimebaseUs / options.period();
|
||||||
|
new_frame_rate = video_header.frame_rate < sampling_rate
|
||||||
|
? video_header.frame_rate
|
||||||
|
: sampling_rate;
|
||||||
|
}
|
||||||
|
std::unique_ptr<VideoHeader> header(new VideoHeader);
|
||||||
|
header->format = video_header.format;
|
||||||
|
header->width = video_header.width;
|
||||||
|
header->height = video_header.height;
|
||||||
|
header->frame_rate = new_frame_rate;
|
||||||
|
cc->Outputs().Index(0).SetHeader(Adopt(header.release()));
|
||||||
|
} else {
|
||||||
|
cc->Outputs().Index(0).SetHeader(cc->Inputs().Index(0).Header());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status PacketThinnerCalculator::Close(CalculatorContext* cc) {
|
||||||
|
// Emit any saved packets before quitting.
|
||||||
|
if (!saved_packet_.IsEmpty()) {
|
||||||
|
// Only sync thinner should have saved packets.
|
||||||
|
CHECK_EQ(PacketThinnerCalculatorOptions::SYNC, thinner_type_);
|
||||||
|
if (sync_output_timestamps_) {
|
||||||
|
cc->Outputs().Index(0).AddPacket(
|
||||||
|
saved_packet_.At(NearestSyncTimestamp(saved_packet_.Timestamp())));
|
||||||
|
} else {
|
||||||
|
cc->Outputs().Index(0).AddPacket(saved_packet_);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status PacketThinnerCalculator::AsyncThinnerProcess(
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
if (cc->InputTimestamp() >= next_valid_timestamp_) {
|
||||||
|
cc->Outputs().Index(0).AddPacket(
|
||||||
|
cc->Inputs().Index(0).Value()); // Emit current packet.
|
||||||
|
next_valid_timestamp_ = cc->InputTimestamp() + period_;
|
||||||
|
// Guaranteed not to emit packets seen during refractory period.
|
||||||
|
cc->Outputs().Index(0).SetNextTimestampBound(next_valid_timestamp_);
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status PacketThinnerCalculator::SyncThinnerProcess(
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
if (saved_packet_.IsEmpty()) {
|
||||||
|
// If no packet has been saved, store the current packet.
|
||||||
|
saved_packet_ = cc->Inputs().Index(0).Value();
|
||||||
|
cc->Outputs().Index(0).SetNextTimestampBound(
|
||||||
|
sync_output_timestamps_ ? NearestSyncTimestamp(cc->InputTimestamp())
|
||||||
|
: cc->InputTimestamp());
|
||||||
|
} else {
|
||||||
|
// Saved packet exists -- update or emit.
|
||||||
|
const Timestamp saved = saved_packet_.Timestamp();
|
||||||
|
const Timestamp saved_sync = NearestSyncTimestamp(saved);
|
||||||
|
const Timestamp now = cc->InputTimestamp();
|
||||||
|
const Timestamp now_sync = NearestSyncTimestamp(now);
|
||||||
|
CHECK_LE(saved_sync, now_sync);
|
||||||
|
if (saved_sync == now_sync) {
|
||||||
|
// Saved Packet is in same interval as current packet.
|
||||||
|
// Replace saved packet with current if it is at least as
|
||||||
|
// central as the saved packet wrt temporal interval.
|
||||||
|
// [We break ties in favor of fresher packets]
|
||||||
|
if (abs(now - now_sync) <= abs(saved - saved_sync)) {
|
||||||
|
saved_packet_ = cc->Inputs().Index(0).Value();
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// Saved packet is the best packet from earlier interval: emit!
|
||||||
|
if (sync_output_timestamps_) {
|
||||||
|
cc->Outputs().Index(0).AddPacket(saved_packet_.At(saved_sync));
|
||||||
|
cc->Outputs().Index(0).SetNextTimestampBound(now_sync);
|
||||||
|
} else {
|
||||||
|
cc->Outputs().Index(0).AddPacket(saved_packet_);
|
||||||
|
cc->Outputs().Index(0).SetNextTimestampBound(now);
|
||||||
|
}
|
||||||
|
// Current packet is the first one we've seen from new interval -- save!
|
||||||
|
saved_packet_ = cc->Inputs().Index(0).Value();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
Timestamp PacketThinnerCalculator::NearestSyncTimestamp(Timestamp now) const {
|
||||||
|
CHECK_NE(start_time_, Timestamp::Unset())
|
||||||
|
<< "Method only valid for sync thinner calculator.";
|
||||||
|
|
||||||
|
// Computation is done using int64 arithmetic. No easy way to avoid
|
||||||
|
// since Timestamps don't support div and multiply.
|
||||||
|
const int64 now64 = now.Value();
|
||||||
|
const int64 start64 = start_time_.Value();
|
||||||
|
const int64 period64 = period_.Value();
|
||||||
|
CHECK_LE(0, period64);
|
||||||
|
|
||||||
|
// Round now64 to its closest interval (units of period64).
|
||||||
|
int64 sync64 =
|
||||||
|
(now64 - start64 + period64 / 2) / period64 * period64 + start64;
|
||||||
|
CHECK_LE(abs(now64 - sync64), period64 / 2)
|
||||||
|
<< "start64: " << start64 << "; now64: " << now64
|
||||||
|
<< "; sync64: " << sync64;
|
||||||
|
|
||||||
|
return Timestamp(sync64);
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,66 @@
|
|||||||
|
// Copyright 2018 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
syntax = "proto2";
|
||||||
|
|
||||||
|
package mediapipe;
|
||||||
|
|
||||||
|
import "mediapipe/framework/calculator.proto";
|
||||||
|
|
||||||
|
message PacketThinnerCalculatorOptions {
|
||||||
|
extend CalculatorOptions {
|
||||||
|
optional PacketThinnerCalculatorOptions ext = 288533508;
|
||||||
|
}
|
||||||
|
|
||||||
|
enum ThinnerType {
|
||||||
|
ASYNC = 1; // Asynchronous thinner, described below [default].
|
||||||
|
SYNC = 2; // Synchronous thinner, also described below.
|
||||||
|
}
|
||||||
|
optional ThinnerType thinner_type = 1 [default = ASYNC];
|
||||||
|
|
||||||
|
// The period (in microsecond) specifies the temporal interval during which
|
||||||
|
// only a single packet is emitted in the output stream. Has subtly different
|
||||||
|
// semantics depending on the thinner type, as follows.
|
||||||
|
//
|
||||||
|
// Async thinner: this option is a refractory period -- once a packet is
|
||||||
|
// emitted, we guarantee that no packets will be emitted for period ticks.
|
||||||
|
//
|
||||||
|
// Sync thinner: the period specifies a temporal interval during which
|
||||||
|
// only one packet is emitted. The emitted packet is guaranteed to be
|
||||||
|
// the one closest to the center of the temporal interval (no guarantee on
|
||||||
|
// how ties are broken). More specifically,
|
||||||
|
// intervals are centered at start_time + i * period
|
||||||
|
// (for non-negative integers i).
|
||||||
|
// Thus, each interval extends period/2 ticks before and after its center.
|
||||||
|
// Additionally, in the sync thinner any packets earlier than start_time
|
||||||
|
// are discarded and the thinner calls Close() once timestamp equals or
|
||||||
|
// exceeds end_time.
|
||||||
|
optional int64 period = 2 [default = 1];
|
||||||
|
|
||||||
|
// Packets before start_time and at/after end_time are discarded.
|
||||||
|
// Additionally, for a sync thinner, start time specifies the center of
|
||||||
|
// time invervals as described above and therefore should be set explicitly.
|
||||||
|
optional int64 start_time = 3; // If not specified, set to 0 for SYNC type,
|
||||||
|
// and set to Timestamp::Min() for ASYNC type.
|
||||||
|
optional int64 end_time = 4; // Set to Timestamp::Max() if not specified.
|
||||||
|
|
||||||
|
// Whether the timestamps of packets emitted by sync thinner should
|
||||||
|
// correspond to the center of their corresponding temporal interval.
|
||||||
|
// If false, packets emitted using original timestamp (as in async thinner).
|
||||||
|
optional bool sync_output_timestamps = 5 [default = true];
|
||||||
|
|
||||||
|
// If true, update the frame rate in the header, if it's available, to an
|
||||||
|
// estimated frame rate due to the sampling.
|
||||||
|
optional bool update_frame_rate = 6 [default = false];
|
||||||
|
}
|
||||||
@@ -0,0 +1,357 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "absl/strings/str_cat.h"
|
||||||
|
#include "mediapipe/calculators/core/packet_thinner_calculator.pb.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/calculator_runner.h"
|
||||||
|
#include "mediapipe/framework/formats/video_stream_header.h"
|
||||||
|
#include "mediapipe/framework/port/gmock.h"
|
||||||
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "mediapipe/framework/port/integral_types.h"
|
||||||
|
#include "mediapipe/framework/port/status_matchers.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
// A simple version of CalculatorRunner with built-in convenience methods for
|
||||||
|
// setting inputs from a vector and checking outputs against a vector of
|
||||||
|
// expected outputs.
|
||||||
|
class SimpleRunner : public CalculatorRunner {
|
||||||
|
public:
|
||||||
|
explicit SimpleRunner(const CalculatorOptions& options)
|
||||||
|
: CalculatorRunner("PacketThinnerCalculator", options) {
|
||||||
|
SetNumInputs(1);
|
||||||
|
SetNumOutputs(1);
|
||||||
|
SetNumInputSidePackets(0);
|
||||||
|
}
|
||||||
|
|
||||||
|
explicit SimpleRunner(const CalculatorGraphConfig::Node& node)
|
||||||
|
: CalculatorRunner(node) {}
|
||||||
|
|
||||||
|
void SetInput(const std::vector<int>& timestamp_list) {
|
||||||
|
MutableInputs()->Index(0).packets.clear();
|
||||||
|
for (const int ts : timestamp_list) {
|
||||||
|
MutableInputs()->Index(0).packets.push_back(
|
||||||
|
MakePacket<std::string>(absl::StrCat("Frame #", ts))
|
||||||
|
.At(Timestamp(ts)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
void SetFrameRate(const double frame_rate) {
|
||||||
|
auto video_header = absl::make_unique<VideoHeader>();
|
||||||
|
video_header->frame_rate = frame_rate;
|
||||||
|
MutableInputs()->Index(0).header = Adopt(video_header.release());
|
||||||
|
}
|
||||||
|
|
||||||
|
std::vector<int64> GetOutputTimestamps() const {
|
||||||
|
std::vector<int64> timestamps;
|
||||||
|
for (const Packet& packet : Outputs().Index(0).packets) {
|
||||||
|
timestamps.emplace_back(packet.Timestamp().Value());
|
||||||
|
}
|
||||||
|
return timestamps;
|
||||||
|
}
|
||||||
|
|
||||||
|
double GetFrameRate() const {
|
||||||
|
CHECK(!Outputs().Index(0).header.IsEmpty());
|
||||||
|
return Outputs().Index(0).header.Get<VideoHeader>().frame_rate;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
// Check that thinner respects start_time and end_time options.
|
||||||
|
// We only test with one thinner because the logic for start & end time
|
||||||
|
// handling is shared across both types of thinner in Process().
|
||||||
|
TEST(PacketThinnerCalculatorTest, StartAndEndTimeTest) {
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||||
|
extension->set_period(5);
|
||||||
|
extension->set_start_time(4);
|
||||||
|
extension->set_end_time(12);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
runner.SetInput({2, 3, 5, 7, 11, 13, 17, 19, 23, 29});
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {5, 11};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, AsyncUniformStreamThinningTest) {
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||||
|
extension->set_period(5);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {2, 8, 14};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, ASyncUniformStreamThinningTestBySidePacket) {
|
||||||
|
// Note: sync runner but outputting *original* timestamps.
|
||||||
|
CalculatorGraphConfig::Node node;
|
||||||
|
node.set_calculator("PacketThinnerCalculator");
|
||||||
|
node.add_input_side_packet("PERIOD:period");
|
||||||
|
node.add_input_stream("input_stream");
|
||||||
|
node.add_output_stream("output_stream");
|
||||||
|
auto* extension = node.mutable_options()->MutableExtension(
|
||||||
|
PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||||
|
extension->set_start_time(0);
|
||||||
|
extension->set_sync_output_timestamps(false);
|
||||||
|
|
||||||
|
SimpleRunner runner(node);
|
||||||
|
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||||
|
runner.MutableSidePackets()->Tag("PERIOD") = MakePacket<int64>(5);
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {2, 8, 14};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, SyncUniformStreamThinningTest1) {
|
||||||
|
// Note: sync runner but outputting *original* timestamps.
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||||
|
extension->set_start_time(0);
|
||||||
|
extension->set_period(5);
|
||||||
|
extension->set_sync_output_timestamps(false);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {2, 6, 10, 14};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, SyncUniformStreamThinningTestBySidePacket1) {
|
||||||
|
// Note: sync runner but outputting *original* timestamps.
|
||||||
|
CalculatorGraphConfig::Node node;
|
||||||
|
node.set_calculator("PacketThinnerCalculator");
|
||||||
|
node.add_input_side_packet("PERIOD:period");
|
||||||
|
node.add_input_stream("input_stream");
|
||||||
|
node.add_output_stream("output_stream");
|
||||||
|
auto* extension = node.mutable_options()->MutableExtension(
|
||||||
|
PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||||
|
extension->set_start_time(0);
|
||||||
|
extension->set_sync_output_timestamps(false);
|
||||||
|
|
||||||
|
SimpleRunner runner(node);
|
||||||
|
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||||
|
runner.MutableSidePackets()->Tag("PERIOD") = MakePacket<int64>(5);
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {2, 6, 10, 14};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, SyncUniformStreamThinningTest2) {
|
||||||
|
// Same test but now with synced timestamps.
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||||
|
extension->set_start_time(0);
|
||||||
|
extension->set_period(5);
|
||||||
|
extension->set_sync_output_timestamps(true);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {0, 5, 10, 15};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test: Given a stream with timestamps corresponding to first ten prime numbers
|
||||||
|
// and period of 5, confirm whether timestamps of thinner stream matches
|
||||||
|
// expectations.
|
||||||
|
TEST(PacketThinnerCalculatorTest, PrimeStreamThinningTest1) {
|
||||||
|
// ASYNC thinner.
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||||
|
extension->set_period(5);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
runner.SetInput({2, 3, 5, 7, 11, 13, 17, 19, 23, 29});
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {2, 7, 13, 19, 29};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, PrimeStreamThinningTest2) {
|
||||||
|
// SYNC with original timestamps.
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||||
|
extension->set_start_time(0);
|
||||||
|
extension->set_period(5);
|
||||||
|
extension->set_sync_output_timestamps(false);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
runner.SetInput({2, 3, 5, 7, 11, 13, 17, 19, 23, 29});
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {2, 5, 11, 17, 19, 23, 29};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Confirm that Calculator correctly handles boundary cases.
|
||||||
|
TEST(PacketThinnerCalculatorTest, BoundaryTimestampTest1) {
|
||||||
|
// Odd period, negative start_time
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||||
|
extension->set_start_time(-10);
|
||||||
|
extension->set_period(5);
|
||||||
|
extension->set_sync_output_timestamps(true);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
// Two timestamps falling on either side of a period boundary.
|
||||||
|
runner.SetInput({2, 3});
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {0, 5};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, BoundaryTimestampTest2) {
|
||||||
|
// Even period, negative start_time, negative packet timestamps.
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||||
|
extension->set_start_time(-144);
|
||||||
|
extension->set_period(6);
|
||||||
|
extension->set_sync_output_timestamps(true);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
// Two timestamps falling on either side of a period boundary.
|
||||||
|
runner.SetInput({-4, -3, 8, 9});
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {-6, 0, 6, 12};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, FrameRateTest1) {
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||||
|
extension->set_period(5);
|
||||||
|
extension->set_update_frame_rate(true);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||||
|
runner.SetFrameRate(1000000.0 / 2);
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {2, 8, 14};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
// The true sampling period is 6.
|
||||||
|
EXPECT_DOUBLE_EQ(1000000.0 / 6, runner.GetFrameRate());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, FrameRateTest2) {
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::ASYNC);
|
||||||
|
extension->set_period(5);
|
||||||
|
extension->set_update_frame_rate(true);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
runner.SetInput({8, 16, 24, 32, 40, 48, 56});
|
||||||
|
runner.SetFrameRate(1000000.0 / 8);
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
const std::vector<int64> expected_timestamps = {8, 16, 24, 32, 40, 48, 56};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
// The true sampling period is still 8.
|
||||||
|
EXPECT_DOUBLE_EQ(1000000.0 / 8, runner.GetFrameRate());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, FrameRateTest3) {
|
||||||
|
// Note: sync runner but outputting *original* timestamps.
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||||
|
extension->set_start_time(0);
|
||||||
|
extension->set_period(5);
|
||||||
|
extension->set_sync_output_timestamps(false);
|
||||||
|
extension->set_update_frame_rate(true);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||||
|
runner.SetFrameRate(1000000.0 / 2);
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {2, 6, 10, 14};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
// The true (long-run) sampling period is 5.
|
||||||
|
EXPECT_DOUBLE_EQ(1000000.0 / 5, runner.GetFrameRate());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, FrameRateTest4) {
|
||||||
|
// Same test but now with synced timestamps.
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||||
|
extension->set_start_time(0);
|
||||||
|
extension->set_period(5);
|
||||||
|
extension->set_sync_output_timestamps(true);
|
||||||
|
extension->set_update_frame_rate(true);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
runner.SetInput({2, 4, 6, 8, 10, 12, 14});
|
||||||
|
runner.SetFrameRate(1000000.0 / 2);
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {0, 5, 10, 15};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
// The true (long-run) sampling period is 5.
|
||||||
|
EXPECT_DOUBLE_EQ(1000000.0 / 5, runner.GetFrameRate());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(PacketThinnerCalculatorTest, FrameRateTest5) {
|
||||||
|
CalculatorOptions options;
|
||||||
|
auto* extension =
|
||||||
|
options.MutableExtension(PacketThinnerCalculatorOptions::ext);
|
||||||
|
extension->set_thinner_type(PacketThinnerCalculatorOptions::SYNC);
|
||||||
|
extension->set_start_time(0);
|
||||||
|
extension->set_period(5);
|
||||||
|
extension->set_sync_output_timestamps(true);
|
||||||
|
extension->set_update_frame_rate(true);
|
||||||
|
SimpleRunner runner(options);
|
||||||
|
runner.SetInput({8, 16, 24, 32, 40, 48, 56});
|
||||||
|
runner.SetFrameRate(1000000.0 / 8);
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<int64> expected_timestamps = {10, 15, 25, 30, 40, 50, 55};
|
||||||
|
EXPECT_EQ(expected_timestamps, runner.GetOutputTimestamps());
|
||||||
|
// The true (long-run) sampling period is 8.
|
||||||
|
EXPECT_DOUBLE_EQ(1000000.0 / 8, runner.GetFrameRate());
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -17,6 +17,7 @@
|
|||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/port/ret_check.h"
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
#include "mediapipe/framework/port/status.h"
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
#include "mediapipe/framework/timestamp.h"
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
@@ -24,13 +25,17 @@ namespace mediapipe {
|
|||||||
// together with some previous output.
|
// together with some previous output.
|
||||||
//
|
//
|
||||||
// For the first packet that arrives on the MAIN input, the timestamp bound is
|
// 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
|
// packet. This way they are not kept waiting for the previous output, which
|
||||||
// for the first iteration does not exist.
|
// for the first iteration does not exist.
|
||||||
//
|
//
|
||||||
// Thereafter, each packet received on MAIN is matched with a packet received
|
// Thereafter,
|
||||||
// on LOOP; the LOOP packet's timestamp is changed to that of the MAIN packet,
|
// - Each non-empty MAIN packet results in:
|
||||||
// and it is output on PREV_LOOP.
|
// 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:
|
// Example config:
|
||||||
// node {
|
// node {
|
||||||
@@ -55,69 +60,115 @@ class PreviousLoopbackCalculator : public CalculatorBase {
|
|||||||
// TODO: an optional PREV_TIMESTAMP output could be added to
|
// TODO: an optional PREV_TIMESTAMP output could be added to
|
||||||
// carry the original timestamp of the packet on PREV_LOOP.
|
// carry the original timestamp of the packet on PREV_LOOP.
|
||||||
cc->SetInputStreamHandler("ImmediateInputStreamHandler");
|
cc->SetInputStreamHandler("ImmediateInputStreamHandler");
|
||||||
|
// Process() function is invoked in response to MAIN/LOOP stream timestamp
|
||||||
|
// bound updates.
|
||||||
|
cc->SetProcessTimestampBounds(true);
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status Open(CalculatorContext* cc) final {
|
::mediapipe::Status Open(CalculatorContext* cc) final {
|
||||||
main_id_ = cc->Inputs().GetId("MAIN", 0);
|
main_id_ = cc->Inputs().GetId("MAIN", 0);
|
||||||
loop_id_ = cc->Inputs().GetId("LOOP", 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()
|
cc->Outputs()
|
||||||
.Get(loop_out_id_)
|
.Get(prev_loop_id_)
|
||||||
.SetHeader(cc->Inputs().Get(loop_id_).Header());
|
.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();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status Process(CalculatorContext* cc) final {
|
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||||
Packet& main_packet = cc->Inputs().Get(main_id_).Value();
|
// Non-empty packets and empty packets indicating timestamp bound updates
|
||||||
if (!main_packet.IsEmpty()) {
|
// are guaranteed to have timestamps greater than timestamps of previous
|
||||||
main_ts_.push_back(main_packet.Timestamp());
|
// packets within the same stream. Calculator tracks and operates on such
|
||||||
}
|
// packets.
|
||||||
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();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
while (!main_ts_.empty() && !loopback_packets_.empty()) {
|
const Packet& main_packet = cc->Inputs().Get(main_id_).Value();
|
||||||
Timestamp main_timestamp = main_ts_.front();
|
if (prev_main_ts_ < main_packet.Timestamp()) {
|
||||||
main_ts_.pop_front();
|
Timestamp loop_timestamp;
|
||||||
Packet previous_loopback = loopback_packets_.front().At(main_timestamp);
|
if (!main_packet.IsEmpty()) {
|
||||||
loopback_packets_.pop_front();
|
loop_timestamp = prev_non_empty_main_ts_;
|
||||||
|
prev_non_empty_main_ts_ = main_packet.Timestamp();
|
||||||
if (previous_loopback.IsEmpty()) {
|
|
||||||
// TODO: SetCompleteTimestampBound would be more useful.
|
|
||||||
cc->Outputs()
|
|
||||||
.Get(loop_out_id_)
|
|
||||||
.SetNextTimestampBound(main_timestamp + 1);
|
|
||||||
} else {
|
} else {
|
||||||
cc->Outputs().Get(loop_out_id_).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();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if (!main_ts_.empty()) {
|
|
||||||
cc->Outputs().Get(loop_out_id_).SetNextTimestampBound(main_ts_.front());
|
if (main_packet_specs_.empty() && cc->Inputs().Get(main_id_).IsDone()) {
|
||||||
}
|
prev_loop.Close();
|
||||||
if (cc->Inputs().Get(main_id_).IsDone() && main_ts_.empty()) {
|
|
||||||
cc->Outputs().Get(loop_out_id_).Close();
|
|
||||||
}
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
private:
|
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 main_id_;
|
||||||
CollectionItemId loop_id_;
|
CollectionItemId loop_id_;
|
||||||
CollectionItemId loop_out_id_;
|
CollectionItemId prev_loop_id_;
|
||||||
|
|
||||||
std::deque<Timestamp> main_ts_;
|
// Contains specs for MAIN packets which only can be:
|
||||||
std::deque<Packet> loopback_packets_;
|
// - 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);
|
REGISTER_CALCULATOR(PreviousLoopbackCalculator);
|
||||||
|
|
||||||
|
|||||||
@@ -13,6 +13,7 @@
|
|||||||
// limitations under the License.
|
// limitations under the License.
|
||||||
|
|
||||||
#include <algorithm>
|
#include <algorithm>
|
||||||
|
#include <functional>
|
||||||
#include <memory>
|
#include <memory>
|
||||||
#include <string>
|
#include <string>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
@@ -25,12 +26,17 @@
|
|||||||
#include "mediapipe/framework/port/gtest.h"
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
#include "mediapipe/framework/port/integral_types.h"
|
#include "mediapipe/framework/port/integral_types.h"
|
||||||
#include "mediapipe/framework/port/parse_text_proto.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/port/status_matchers.h"
|
||||||
#include "mediapipe/framework/timestamp.h"
|
#include "mediapipe/framework/timestamp.h"
|
||||||
#include "mediapipe/framework/tool/sink.h"
|
#include "mediapipe/framework/tool/sink.h"
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
|
using ::testing::ElementsAre;
|
||||||
|
using ::testing::Eq;
|
||||||
|
using ::testing::Pair;
|
||||||
|
using ::testing::Value;
|
||||||
namespace {
|
namespace {
|
||||||
|
|
||||||
// Returns the timestamp values for a vector of Packets.
|
// Returns the timestamp values for a vector of Packets.
|
||||||
@@ -43,6 +49,23 @@ std::vector<int64> TimestampValues(const std::vector<Packet>& packets) {
|
|||||||
return result;
|
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) {
|
TEST(PreviousLoopbackCalculator, CorrectTimestamps) {
|
||||||
std::vector<Packet> in_prev;
|
std::vector<Packet> in_prev;
|
||||||
CalculatorGraphConfig graph_config_ =
|
CalculatorGraphConfig graph_config_ =
|
||||||
@@ -81,32 +104,30 @@ TEST(PreviousLoopbackCalculator, CorrectTimestamps) {
|
|||||||
MP_EXPECT_OK(graph_.AddPacketToInputStream(
|
MP_EXPECT_OK(graph_.AddPacketToInputStream(
|
||||||
input_name, MakePacket<int>(n).At(Timestamp(n))));
|
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);
|
send_packet("in", 1);
|
||||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1}));
|
EXPECT_THAT(TimestampValues(in_prev), ElementsAre(1));
|
||||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(1, -1));
|
EXPECT_THAT(in_prev.back(),
|
||||||
|
PairPacket(Timestamp(1), Pair(IntPacket(1), EmptyPacket())));
|
||||||
|
|
||||||
send_packet("in", 2);
|
send_packet("in", 2);
|
||||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 2}));
|
EXPECT_THAT(TimestampValues(in_prev), ElementsAre(1, 2));
|
||||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(2, 1));
|
EXPECT_THAT(in_prev.back(),
|
||||||
|
PairPacket(Timestamp(2), Pair(IntPacket(2), IntPacket(1))));
|
||||||
|
|
||||||
send_packet("in", 5);
|
send_packet("in", 5);
|
||||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 2, 5}));
|
EXPECT_THAT(TimestampValues(in_prev), ElementsAre(1, 2, 5));
|
||||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(5, 2));
|
EXPECT_THAT(in_prev.back(),
|
||||||
|
PairPacket(Timestamp(5), Pair(IntPacket(5), IntPacket(2))));
|
||||||
|
|
||||||
send_packet("in", 15);
|
send_packet("in", 15);
|
||||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 2, 5, 15}));
|
EXPECT_THAT(TimestampValues(in_prev), ElementsAre(1, 2, 5, 15));
|
||||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(15, 5));
|
EXPECT_THAT(in_prev.back(),
|
||||||
|
PairPacket(Timestamp(15), Pair(IntPacket(15), IntPacket(5))));
|
||||||
|
|
||||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||||
@@ -185,27 +206,564 @@ TEST(PreviousLoopbackCalculator, ClosesCorrectly) {
|
|||||||
|
|
||||||
send_packet("in", 1);
|
send_packet("in", 1);
|
||||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1}));
|
EXPECT_THAT(TimestampValues(outputs), ElementsAre(1));
|
||||||
|
|
||||||
send_packet("in", 2);
|
send_packet("in", 2);
|
||||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
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);
|
send_packet("in", 5);
|
||||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
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);
|
send_packet("in", 15);
|
||||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
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_.CloseAllInputStreams());
|
||||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
EXPECT_EQ(TimestampValues(outputs),
|
EXPECT_THAT(TimestampValues(outputs),
|
||||||
(std::vector<int64>{1, 2, 5, 15, Timestamp::Max().Value()}));
|
ElementsAre(1, 2, 5, 15, Timestamp::Max().Value()));
|
||||||
|
|
||||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Demonstrates that downstream calculators won't be blocked by
|
||||||
|
// always-empty-LOOP-stream.
|
||||||
|
TEST(PreviousLoopbackCalculator, EmptyLoopForever) {
|
||||||
|
std::vector<Packet> outputs;
|
||||||
|
CalculatorGraphConfig graph_config_ =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||||
|
input_stream: 'in'
|
||||||
|
node {
|
||||||
|
calculator: 'PreviousLoopbackCalculator'
|
||||||
|
input_stream: 'MAIN:in'
|
||||||
|
input_stream: 'LOOP:previous'
|
||||||
|
input_stream_info: { tag_index: 'LOOP' back_edge: true }
|
||||||
|
output_stream: 'PREV_LOOP:previous'
|
||||||
|
}
|
||||||
|
# This calculator synchronizes its inputs as normal, so it is used
|
||||||
|
# to check that both "in" and "previous" are ready.
|
||||||
|
node {
|
||||||
|
calculator: 'PassThroughCalculator'
|
||||||
|
input_stream: 'in'
|
||||||
|
input_stream: 'previous'
|
||||||
|
output_stream: 'out'
|
||||||
|
output_stream: 'previous2'
|
||||||
|
}
|
||||||
|
node {
|
||||||
|
calculator: 'PacketOnCloseCalculator'
|
||||||
|
input_stream: 'out'
|
||||||
|
output_stream: 'close_out'
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
tool::AddVectorSink("close_out", &graph_config_, &outputs);
|
||||||
|
|
||||||
|
CalculatorGraph graph_;
|
||||||
|
MP_ASSERT_OK(graph_.Initialize(graph_config_, {}));
|
||||||
|
MP_ASSERT_OK(graph_.StartRun({}));
|
||||||
|
|
||||||
|
auto send_packet = [&graph_](const std::string& input_name, int n) {
|
||||||
|
MP_EXPECT_OK(graph_.AddPacketToInputStream(
|
||||||
|
input_name, MakePacket<int>(n).At(Timestamp(n))));
|
||||||
|
};
|
||||||
|
|
||||||
|
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 + 1; ++j) {
|
||||||
|
EXPECT_EQ(ts_values[j], j);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
|
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
|
} // anonymous namespace
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -28,55 +28,133 @@ using mediapipe::PacketTypeSet;
|
|||||||
using mediapipe::Timestamp;
|
using mediapipe::Timestamp;
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
|
|
||||||
|
constexpr char kTagAtPreStream[] = "AT_PRESTREAM";
|
||||||
|
constexpr char kTagAtPostStream[] = "AT_POSTSTREAM";
|
||||||
|
constexpr char kTagAtZero[] = "AT_ZERO";
|
||||||
|
constexpr char kTagAtTick[] = "AT_TICK";
|
||||||
|
constexpr char kTagTick[] = "TICK";
|
||||||
|
|
||||||
static std::map<std::string, Timestamp>* kTimestampMap = []() {
|
static std::map<std::string, Timestamp>* kTimestampMap = []() {
|
||||||
auto* res = new std::map<std::string, Timestamp>();
|
auto* res = new std::map<std::string, Timestamp>();
|
||||||
res->emplace("AT_PRESTREAM", Timestamp::PreStream());
|
res->emplace(kTagAtPreStream, Timestamp::PreStream());
|
||||||
res->emplace("AT_POSTSTREAM", Timestamp::PostStream());
|
res->emplace(kTagAtPostStream, Timestamp::PostStream());
|
||||||
res->emplace("AT_ZERO", Timestamp(0));
|
res->emplace(kTagAtZero, Timestamp(0));
|
||||||
|
res->emplace(kTagAtTick, Timestamp::Unset());
|
||||||
return res;
|
return res;
|
||||||
}();
|
}();
|
||||||
|
|
||||||
|
template <typename CC>
|
||||||
|
std::string GetOutputTag(const CC& cc) {
|
||||||
|
// Single output tag only is required by contract.
|
||||||
|
return *cc.Outputs().GetTags().begin();
|
||||||
|
}
|
||||||
|
|
||||||
} // namespace
|
} // namespace
|
||||||
|
|
||||||
// Outputs the single input_side_packet at the timestamp specified in the
|
// Outputs side packet(s) in corresponding output stream(s) with a particular
|
||||||
// output_stream tag. Valid tags are AT_PRESTREAM, AT_POSTSTREAM and AT_ZERO.
|
// timestamp, depending on the tag used to define output stream(s). (One tag can
|
||||||
|
// be used only.)
|
||||||
|
//
|
||||||
|
// Valid tags are AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO and AT_TICK and
|
||||||
|
// corresponding timestamps are Timestamp::PreStream(), Timestamp::PostStream(),
|
||||||
|
// Timestamp(0) and timestamp of a packet received in TICK input.
|
||||||
|
//
|
||||||
|
// Examples:
|
||||||
|
// node {
|
||||||
|
// calculator: "SidePacketToStreamCalculator"
|
||||||
|
// input_side_packet: "side_packet"
|
||||||
|
// output_stream: "AT_PRESTREAM:packet"
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "SidePacketToStreamCalculator"
|
||||||
|
// input_stream: "TICK:tick"
|
||||||
|
// input_side_packet: "side_packet"
|
||||||
|
// output_stream: "AT_TICK:packet"
|
||||||
|
// }
|
||||||
class SidePacketToStreamCalculator : public CalculatorBase {
|
class SidePacketToStreamCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
SidePacketToStreamCalculator() = default;
|
SidePacketToStreamCalculator() = default;
|
||||||
~SidePacketToStreamCalculator() override = default;
|
~SidePacketToStreamCalculator() override = default;
|
||||||
|
|
||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||||
|
|
||||||
|
private:
|
||||||
|
bool is_tick_processing_ = false;
|
||||||
|
std::string output_tag_;
|
||||||
};
|
};
|
||||||
REGISTER_CALCULATOR(SidePacketToStreamCalculator);
|
REGISTER_CALCULATOR(SidePacketToStreamCalculator);
|
||||||
|
|
||||||
::mediapipe::Status SidePacketToStreamCalculator::GetContract(
|
::mediapipe::Status SidePacketToStreamCalculator::GetContract(
|
||||||
CalculatorContract* cc) {
|
CalculatorContract* cc) {
|
||||||
cc->InputSidePackets().Index(0).SetAny();
|
const auto& tags = cc->Outputs().GetTags();
|
||||||
|
RET_CHECK(tags.size() == 1 && kTimestampMap->count(*tags.begin()) == 1)
|
||||||
|
<< "Only one of AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO and AT_TICK tags is "
|
||||||
|
"allowed and required to specify output stream(s).";
|
||||||
|
RET_CHECK(
|
||||||
|
(cc->Outputs().HasTag(kTagAtTick) && cc->Inputs().HasTag(kTagTick)) ||
|
||||||
|
(!cc->Outputs().HasTag(kTagAtTick) && !cc->Inputs().HasTag(kTagTick)))
|
||||||
|
<< "Either both of TICK and AT_TICK should be used or none of them.";
|
||||||
|
const std::string output_tag = GetOutputTag(*cc);
|
||||||
|
const int num_entries = cc->Outputs().NumEntries(output_tag);
|
||||||
|
RET_CHECK_EQ(num_entries, cc->InputSidePackets().NumEntries())
|
||||||
|
<< "Same number of input side packets and output streams is required.";
|
||||||
|
for (int i = 0; i < num_entries; ++i) {
|
||||||
|
cc->InputSidePackets().Index(i).SetAny();
|
||||||
|
cc->Outputs()
|
||||||
|
.Get(output_tag, i)
|
||||||
|
.SetSameAs(cc->InputSidePackets().Index(i).GetSameAs());
|
||||||
|
}
|
||||||
|
|
||||||
std::set<std::string> tags = cc->Outputs().GetTags();
|
if (cc->Inputs().HasTag(kTagTick)) {
|
||||||
RET_CHECK_EQ(tags.size(), 1);
|
cc->Inputs().Tag(kTagTick).SetAny();
|
||||||
|
}
|
||||||
|
|
||||||
RET_CHECK_EQ(kTimestampMap->count(*tags.begin()), 1);
|
return ::mediapipe::OkStatus();
|
||||||
cc->Outputs().Tag(*tags.begin()).SetAny();
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status SidePacketToStreamCalculator::Open(CalculatorContext* cc) {
|
||||||
|
output_tag_ = GetOutputTag(*cc);
|
||||||
|
if (cc->Inputs().HasTag(kTagTick)) {
|
||||||
|
is_tick_processing_ = true;
|
||||||
|
// Set offset, so output timestamp bounds are updated in response to TICK
|
||||||
|
// timestamp bound update.
|
||||||
|
cc->SetOffset(TimestampDiff(0));
|
||||||
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status SidePacketToStreamCalculator::Process(
|
::mediapipe::Status SidePacketToStreamCalculator::Process(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
return mediapipe::tool::StatusStop();
|
if (is_tick_processing_) {
|
||||||
|
// TICK input is guaranteed to be non-empty, as it's the only input stream
|
||||||
|
// for this calculator.
|
||||||
|
const auto& timestamp = cc->Inputs().Tag(kTagTick).Value().Timestamp();
|
||||||
|
for (int i = 0; i < cc->Outputs().NumEntries(output_tag_); ++i) {
|
||||||
|
cc->Outputs()
|
||||||
|
.Get(output_tag_, i)
|
||||||
|
.AddPacket(cc->InputSidePackets().Index(i).At(timestamp));
|
||||||
|
}
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
return ::mediapipe::tool::StatusStop();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status SidePacketToStreamCalculator::Close(CalculatorContext* cc) {
|
::mediapipe::Status SidePacketToStreamCalculator::Close(CalculatorContext* cc) {
|
||||||
std::set<std::string> tags = cc->Outputs().GetTags();
|
if (!cc->Outputs().HasTag(kTagAtTick)) {
|
||||||
RET_CHECK_EQ(tags.size(), 1);
|
const auto& timestamp = kTimestampMap->at(output_tag_);
|
||||||
const std::string& tag = *tags.begin();
|
for (int i = 0; i < cc->Outputs().NumEntries(output_tag_); ++i) {
|
||||||
RET_CHECK_EQ(kTimestampMap->count(tag), 1);
|
cc->Outputs()
|
||||||
cc->Outputs().Tag(tag).AddPacket(
|
.Get(output_tag_, i)
|
||||||
cc->InputSidePackets().Index(0).At(kTimestampMap->at(tag)));
|
.AddPacket(cc->InputSidePackets().Index(i).At(timestamp));
|
||||||
|
}
|
||||||
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,275 @@
|
|||||||
|
// 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 <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "absl/memory/memory.h"
|
||||||
|
#include "absl/strings/match.h"
|
||||||
|
#include "absl/strings/str_replace.h"
|
||||||
|
#include "absl/strings/string_view.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "mediapipe/framework/port/integral_types.h"
|
||||||
|
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
#include "mediapipe/framework/port/status_matchers.h"
|
||||||
|
#include "mediapipe/framework/tool/options_util.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
TEST(SidePacketToStreamCalculator, WrongConfig_MissingTick) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "tick"
|
||||||
|
input_side_packet: "side_packet"
|
||||||
|
output_stream: "packet"
|
||||||
|
node {
|
||||||
|
calculator: "SidePacketToStreamCalculator"
|
||||||
|
input_side_packet: "side_packet"
|
||||||
|
output_stream: "AT_TICK:packet"
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorGraph graph;
|
||||||
|
auto status = graph.Initialize(graph_config);
|
||||||
|
EXPECT_FALSE(status.ok());
|
||||||
|
EXPECT_PRED2(
|
||||||
|
absl::StrContains, status.message(),
|
||||||
|
"Either both of TICK and AT_TICK should be used or none of them.");
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(SidePacketToStreamCalculator, WrongConfig_NonExistentTag) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "tick"
|
||||||
|
input_side_packet: "side_packet"
|
||||||
|
output_stream: "packet"
|
||||||
|
node {
|
||||||
|
calculator: "SidePacketToStreamCalculator"
|
||||||
|
input_side_packet: "side_packet"
|
||||||
|
output_stream: "DOES_NOT_EXIST:packet"
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorGraph graph;
|
||||||
|
auto status = graph.Initialize(graph_config);
|
||||||
|
EXPECT_FALSE(status.ok());
|
||||||
|
EXPECT_PRED2(absl::StrContains, status.message(),
|
||||||
|
"Only one of AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO and AT_TICK "
|
||||||
|
"tags is allowed and required to specify output stream(s).");
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(SidePacketToStreamCalculator, WrongConfig_MixedTags) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "tick"
|
||||||
|
input_side_packet: "side_packet0"
|
||||||
|
input_side_packet: "side_packet1"
|
||||||
|
node {
|
||||||
|
calculator: "SidePacketToStreamCalculator"
|
||||||
|
input_side_packet: "side_packet0"
|
||||||
|
input_side_packet: "side_packet1"
|
||||||
|
output_stream: "AT_TICK:packet0"
|
||||||
|
output_stream: "AT_PRE_STREAM:packet1"
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorGraph graph;
|
||||||
|
auto status = graph.Initialize(graph_config);
|
||||||
|
EXPECT_FALSE(status.ok());
|
||||||
|
EXPECT_PRED2(absl::StrContains, status.message(),
|
||||||
|
"Only one of AT_PRESTREAM, AT_POSTSTREAM, AT_ZERO and AT_TICK "
|
||||||
|
"tags is allowed and required to specify output stream(s).");
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(SidePacketToStreamCalculator, WrongConfig_NotEnoughSidePackets) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_side_packet: "side_packet0"
|
||||||
|
input_side_packet: "side_packet1"
|
||||||
|
node {
|
||||||
|
calculator: "SidePacketToStreamCalculator"
|
||||||
|
input_side_packet: "side_packet0"
|
||||||
|
output_stream: "AT_PRESTREAM:0:packet0"
|
||||||
|
output_stream: "AT_PRESTREAM:1:packet1"
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorGraph graph;
|
||||||
|
auto status = graph.Initialize(graph_config);
|
||||||
|
EXPECT_FALSE(status.ok());
|
||||||
|
EXPECT_PRED2(
|
||||||
|
absl::StrContains, status.message(),
|
||||||
|
"Same number of input side packets and output streams is required.");
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(SidePacketToStreamCalculator, WrongConfig_NotEnoughOutputStreams) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_side_packet: "side_packet0"
|
||||||
|
input_side_packet: "side_packet1"
|
||||||
|
node {
|
||||||
|
calculator: "SidePacketToStreamCalculator"
|
||||||
|
input_side_packet: "side_packet0"
|
||||||
|
input_side_packet: "side_packet1"
|
||||||
|
output_stream: "AT_PRESTREAM:packet0"
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorGraph graph;
|
||||||
|
auto status = graph.Initialize(graph_config);
|
||||||
|
EXPECT_FALSE(status.ok());
|
||||||
|
EXPECT_PRED2(
|
||||||
|
absl::StrContains, status.message(),
|
||||||
|
"Same number of input side packets and output streams is required.");
|
||||||
|
}
|
||||||
|
|
||||||
|
void DoTestNonAtTickOutputTag(absl::string_view tag,
|
||||||
|
Timestamp expected_timestamp) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(absl::StrReplaceAll(
|
||||||
|
R"(
|
||||||
|
input_side_packet: "side_packet"
|
||||||
|
output_stream: "packet"
|
||||||
|
node {
|
||||||
|
calculator: "SidePacketToStreamCalculator"
|
||||||
|
input_side_packet: "side_packet"
|
||||||
|
output_stream: "$tag:packet"
|
||||||
|
}
|
||||||
|
)",
|
||||||
|
{{"$tag", tag}}));
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||||
|
const int expected_value = 10;
|
||||||
|
std::vector<Packet> output_packets;
|
||||||
|
MP_ASSERT_OK(graph.ObserveOutputStream(
|
||||||
|
"packet", [&output_packets](const Packet& packet) {
|
||||||
|
output_packets.push_back(packet);
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}));
|
||||||
|
MP_ASSERT_OK(
|
||||||
|
graph.StartRun({{"side_packet", MakePacket<int>(expected_value)}}));
|
||||||
|
MP_ASSERT_OK(graph.WaitForObservedOutput());
|
||||||
|
|
||||||
|
ASSERT_FALSE(output_packets.empty());
|
||||||
|
EXPECT_EQ(expected_timestamp, output_packets.back().Timestamp());
|
||||||
|
EXPECT_EQ(expected_value, output_packets.back().Get<int>());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(SidePacketToStreamCalculator, NoAtTickOutputTags) {
|
||||||
|
DoTestNonAtTickOutputTag("AT_PRESTREAM", Timestamp::PreStream());
|
||||||
|
DoTestNonAtTickOutputTag("AT_POSTSTREAM", Timestamp::PostStream());
|
||||||
|
DoTestNonAtTickOutputTag("AT_ZERO", Timestamp(0));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(SidePacketToStreamCalculator, AtTick) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "tick"
|
||||||
|
input_side_packet: "side_packet"
|
||||||
|
output_stream: "packet"
|
||||||
|
node {
|
||||||
|
calculator: "SidePacketToStreamCalculator"
|
||||||
|
input_stream: "TICK:tick"
|
||||||
|
input_side_packet: "side_packet"
|
||||||
|
output_stream: "AT_TICK:packet"
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
std::vector<Packet> output_packets;
|
||||||
|
tool::AddVectorSink("packet", &graph_config, &output_packets);
|
||||||
|
CalculatorGraph graph;
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||||
|
const int expected_value = 20;
|
||||||
|
MP_ASSERT_OK(
|
||||||
|
graph.StartRun({{"side_packet", MakePacket<int>(expected_value)}}));
|
||||||
|
|
||||||
|
auto tick_and_verify = [&graph, &output_packets,
|
||||||
|
expected_value](int at_timestamp) {
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"tick",
|
||||||
|
MakePacket<int>(/*doesn't matter*/ 1).At(Timestamp(at_timestamp))));
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
|
||||||
|
ASSERT_FALSE(output_packets.empty());
|
||||||
|
EXPECT_EQ(Timestamp(at_timestamp), output_packets.back().Timestamp());
|
||||||
|
EXPECT_EQ(expected_value, output_packets.back().Get<int>());
|
||||||
|
};
|
||||||
|
|
||||||
|
tick_and_verify(/*at_timestamp=*/0);
|
||||||
|
tick_and_verify(/*at_timestamp=*/1);
|
||||||
|
tick_and_verify(/*at_timestamp=*/128);
|
||||||
|
tick_and_verify(/*at_timestamp=*/1024);
|
||||||
|
tick_and_verify(/*at_timestamp=*/1025);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(SidePacketToStreamCalculator, AtTick_MultipleSidePackets) {
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "tick"
|
||||||
|
input_side_packet: "side_packet0"
|
||||||
|
input_side_packet: "side_packet1"
|
||||||
|
output_stream: "packet0"
|
||||||
|
output_stream: "packet1"
|
||||||
|
node {
|
||||||
|
calculator: "SidePacketToStreamCalculator"
|
||||||
|
input_stream: "TICK:tick"
|
||||||
|
input_side_packet: "side_packet0"
|
||||||
|
input_side_packet: "side_packet1"
|
||||||
|
output_stream: "AT_TICK:0:packet0"
|
||||||
|
output_stream: "AT_TICK:1:packet1"
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
std::vector<Packet> output_packets0;
|
||||||
|
tool::AddVectorSink("packet0", &graph_config, &output_packets0);
|
||||||
|
std::vector<Packet> output_packets1;
|
||||||
|
tool::AddVectorSink("packet1", &graph_config, &output_packets1);
|
||||||
|
CalculatorGraph graph;
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||||
|
const int expected_value0 = 20;
|
||||||
|
const int expected_value1 = 128;
|
||||||
|
MP_ASSERT_OK(
|
||||||
|
graph.StartRun({{"side_packet0", MakePacket<int>(expected_value0)},
|
||||||
|
{"side_packet1", MakePacket<int>(expected_value1)}}));
|
||||||
|
|
||||||
|
auto tick_and_verify = [&graph, &output_packets0, &output_packets1,
|
||||||
|
expected_value0, expected_value1](int at_timestamp) {
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"tick",
|
||||||
|
MakePacket<int>(/*doesn't matter*/ 1).At(Timestamp(at_timestamp))));
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
|
||||||
|
ASSERT_FALSE(output_packets0.empty());
|
||||||
|
ASSERT_FALSE(output_packets1.empty());
|
||||||
|
|
||||||
|
EXPECT_EQ(Timestamp(at_timestamp), output_packets0.back().Timestamp());
|
||||||
|
EXPECT_EQ(expected_value0, output_packets0.back().Get<int>());
|
||||||
|
EXPECT_EQ(Timestamp(at_timestamp), output_packets1.back().Timestamp());
|
||||||
|
EXPECT_EQ(expected_value1, output_packets1.back().Get<int>());
|
||||||
|
};
|
||||||
|
|
||||||
|
tick_and_verify(/*at_timestamp=*/0);
|
||||||
|
tick_and_verify(/*at_timestamp=*/1);
|
||||||
|
tick_and_verify(/*at_timestamp=*/128);
|
||||||
|
tick_and_verify(/*at_timestamp=*/1024);
|
||||||
|
tick_and_verify(/*at_timestamp=*/1025);
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -16,10 +16,16 @@
|
|||||||
|
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/formats/detection.pb.h"
|
||||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||||
|
#include "mediapipe/framework/formats/matrix.h"
|
||||||
#include "mediapipe/framework/formats/rect.pb.h"
|
#include "mediapipe/framework/formats/rect.pb.h"
|
||||||
#include "tensorflow/lite/interpreter.h"
|
#include "tensorflow/lite/interpreter.h"
|
||||||
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
|
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
// Example config:
|
// Example config:
|
||||||
@@ -36,14 +42,33 @@ namespace mediapipe {
|
|||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
typedef SplitVectorCalculator<TfLiteTensor> SplitTfLiteTensorVectorCalculator;
|
typedef SplitVectorCalculator<TfLiteTensor, false>
|
||||||
|
SplitTfLiteTensorVectorCalculator;
|
||||||
REGISTER_CALCULATOR(SplitTfLiteTensorVectorCalculator);
|
REGISTER_CALCULATOR(SplitTfLiteTensorVectorCalculator);
|
||||||
|
|
||||||
typedef SplitVectorCalculator<::mediapipe::NormalizedLandmark>
|
typedef SplitVectorCalculator<::mediapipe::NormalizedLandmark, false>
|
||||||
SplitLandmarkVectorCalculator;
|
SplitLandmarkVectorCalculator;
|
||||||
REGISTER_CALCULATOR(SplitLandmarkVectorCalculator);
|
REGISTER_CALCULATOR(SplitLandmarkVectorCalculator);
|
||||||
|
|
||||||
typedef SplitVectorCalculator<::mediapipe::NormalizedRect>
|
typedef SplitVectorCalculator<::mediapipe::NormalizedLandmarkList, false>
|
||||||
|
SplitNormalizedLandmarkListVectorCalculator;
|
||||||
|
REGISTER_CALCULATOR(SplitNormalizedLandmarkListVectorCalculator);
|
||||||
|
|
||||||
|
typedef SplitVectorCalculator<::mediapipe::NormalizedRect, false>
|
||||||
SplitNormalizedRectVectorCalculator;
|
SplitNormalizedRectVectorCalculator;
|
||||||
REGISTER_CALCULATOR(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
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -15,12 +15,14 @@
|
|||||||
#ifndef MEDIAPIPE_CALCULATORS_CORE_SPLIT_VECTOR_CALCULATOR_H_
|
#ifndef MEDIAPIPE_CALCULATORS_CORE_SPLIT_VECTOR_CALCULATOR_H_
|
||||||
#define MEDIAPIPE_CALCULATORS_CORE_SPLIT_VECTOR_CALCULATOR_H_
|
#define MEDIAPIPE_CALCULATORS_CORE_SPLIT_VECTOR_CALCULATOR_H_
|
||||||
|
|
||||||
|
#include <type_traits>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
#include "mediapipe/calculators/core/split_vector_calculator.pb.h"
|
#include "mediapipe/calculators/core/split_vector_calculator.pb.h"
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/port/canonical_errors.h"
|
#include "mediapipe/framework/port/canonical_errors.h"
|
||||||
#include "mediapipe/framework/port/ret_check.h"
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
#include "mediapipe/util/resource_util.h"
|
#include "mediapipe/util/resource_util.h"
|
||||||
#include "tensorflow/lite/error_reporter.h"
|
#include "tensorflow/lite/error_reporter.h"
|
||||||
#include "tensorflow/lite/interpreter.h"
|
#include "tensorflow/lite/interpreter.h"
|
||||||
@@ -29,6 +31,20 @@
|
|||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
using IsCopyable = std::enable_if_t<std::is_copy_constructible<T>::value, bool>;
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
using IsNotCopyable =
|
||||||
|
std::enable_if_t<!std::is_copy_constructible<T>::value, bool>;
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
using IsMovable = std::enable_if_t<std::is_move_constructible<T>::value, bool>;
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
using IsNotMovable =
|
||||||
|
std::enable_if_t<!std::is_move_constructible<T>::value, bool>;
|
||||||
|
|
||||||
// Splits an input packet with std::vector<T> into multiple std::vector<T>
|
// Splits an input packet with std::vector<T> into multiple std::vector<T>
|
||||||
// output packets using the [begin, end) ranges specified in
|
// output packets using the [begin, end) ranges specified in
|
||||||
// SplitVectorCalculatorOptions. If the option "element_only" is set to true,
|
// SplitVectorCalculatorOptions. If the option "element_only" is set to true,
|
||||||
@@ -39,7 +55,7 @@ namespace mediapipe {
|
|||||||
// combined into one vector.
|
// combined into one vector.
|
||||||
// To use this class for a particular type T, register a calculator using
|
// To use this class for a particular type T, register a calculator using
|
||||||
// SplitVectorCalculator<T>.
|
// SplitVectorCalculator<T>.
|
||||||
template <typename T>
|
template <typename T, bool move_elements>
|
||||||
class SplitVectorCalculator : public CalculatorBase {
|
class SplitVectorCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
@@ -51,23 +67,16 @@ class SplitVectorCalculator : public CalculatorBase {
|
|||||||
const auto& options =
|
const auto& options =
|
||||||
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
|
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
|
||||||
|
|
||||||
|
if (!std::is_copy_constructible<T>::value || move_elements) {
|
||||||
|
// Ranges of elements shouldn't overlap when the vector contains
|
||||||
|
// non-copyable elements.
|
||||||
|
RET_CHECK_OK(checkRangesDontOverlap(options));
|
||||||
|
}
|
||||||
|
|
||||||
if (options.combine_outputs()) {
|
if (options.combine_outputs()) {
|
||||||
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1);
|
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1);
|
||||||
cc->Outputs().Index(0).Set<std::vector<T>>();
|
cc->Outputs().Index(0).Set<std::vector<T>>();
|
||||||
for (int i = 0; i < options.ranges_size() - 1; ++i) {
|
RET_CHECK_OK(checkRangesDontOverlap(options));
|
||||||
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 {
|
} else {
|
||||||
if (cc->Outputs().NumEntries() != options.ranges_size()) {
|
if (cc->Outputs().NumEntries() != options.ranges_size()) {
|
||||||
return ::mediapipe::InvalidArgumentError(
|
return ::mediapipe::InvalidArgumentError(
|
||||||
@@ -117,14 +126,26 @@ class SplitVectorCalculator : public CalculatorBase {
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
const auto& input = cc->Inputs().Index(0).Get<std::vector<T>>();
|
if (cc->Inputs().Index(0).IsEmpty()) return ::mediapipe::OkStatus();
|
||||||
RET_CHECK_GE(input.size(), max_range_end_);
|
|
||||||
|
|
||||||
|
if (move_elements) {
|
||||||
|
return ProcessMovableElements<T>(cc);
|
||||||
|
} else {
|
||||||
|
return ProcessCopyableElements<T>(cc);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename U, IsCopyable<U> = true>
|
||||||
|
::mediapipe::Status ProcessCopyableElements(CalculatorContext* cc) {
|
||||||
|
// static_assert(std::is_copy_constructible<U>::value,
|
||||||
|
// "Cannot copy non-copyable elements");
|
||||||
|
const auto& input = cc->Inputs().Index(0).Get<std::vector<U>>();
|
||||||
|
RET_CHECK_GE(input.size(), max_range_end_);
|
||||||
if (combine_outputs_) {
|
if (combine_outputs_) {
|
||||||
auto output = absl::make_unique<std::vector<T>>();
|
auto output = absl::make_unique<std::vector<U>>();
|
||||||
output->reserve(total_elements_);
|
output->reserve(total_elements_);
|
||||||
for (int i = 0; i < ranges_.size(); ++i) {
|
for (int i = 0; i < ranges_.size(); ++i) {
|
||||||
auto elements = absl::make_unique<std::vector<T>>(
|
auto elements = absl::make_unique<std::vector<U>>(
|
||||||
input.begin() + ranges_[i].first,
|
input.begin() + ranges_[i].first,
|
||||||
input.begin() + ranges_[i].second);
|
input.begin() + ranges_[i].second);
|
||||||
output->insert(output->end(), elements->begin(), elements->end());
|
output->insert(output->end(), elements->begin(), elements->end());
|
||||||
@@ -134,7 +155,7 @@ class SplitVectorCalculator : public CalculatorBase {
|
|||||||
if (element_only_) {
|
if (element_only_) {
|
||||||
for (int i = 0; i < ranges_.size(); ++i) {
|
for (int i = 0; i < ranges_.size(); ++i) {
|
||||||
cc->Outputs().Index(i).AddPacket(
|
cc->Outputs().Index(i).AddPacket(
|
||||||
MakePacket<T>(input[ranges_[i].first]).At(cc->InputTimestamp()));
|
MakePacket<U>(input[ranges_[i].first]).At(cc->InputTimestamp()));
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
for (int i = 0; i < ranges_.size(); ++i) {
|
for (int i = 0; i < ranges_.size(); ++i) {
|
||||||
@@ -149,7 +170,78 @@ class SplitVectorCalculator : public CalculatorBase {
|
|||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
template <typename U, IsNotCopyable<U> = true>
|
||||||
|
::mediapipe::Status ProcessCopyableElements(CalculatorContext* cc) {
|
||||||
|
return ::mediapipe::InternalError("Cannot copy non-copyable elements.");
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename U, IsMovable<U> = true>
|
||||||
|
::mediapipe::Status ProcessMovableElements(CalculatorContext* cc) {
|
||||||
|
::mediapipe::StatusOr<std::unique_ptr<std::vector<U>>> input_status =
|
||||||
|
cc->Inputs().Index(0).Value().Consume<std::vector<U>>();
|
||||||
|
if (!input_status.ok()) return input_status.status();
|
||||||
|
std::unique_ptr<std::vector<U>> input_vector =
|
||||||
|
std::move(input_status).ValueOrDie();
|
||||||
|
RET_CHECK_GE(input_vector->size(), max_range_end_);
|
||||||
|
|
||||||
|
if (combine_outputs_) {
|
||||||
|
auto output = absl::make_unique<std::vector<U>>();
|
||||||
|
output->reserve(total_elements_);
|
||||||
|
for (int i = 0; i < ranges_.size(); ++i) {
|
||||||
|
output->insert(
|
||||||
|
output->end(),
|
||||||
|
std::make_move_iterator(input_vector->begin() + ranges_[i].first),
|
||||||
|
std::make_move_iterator(input_vector->begin() + ranges_[i].second));
|
||||||
|
}
|
||||||
|
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||||
|
} else {
|
||||||
|
if (element_only_) {
|
||||||
|
for (int i = 0; i < ranges_.size(); ++i) {
|
||||||
|
cc->Outputs().Index(i).AddPacket(
|
||||||
|
MakePacket<U>(std::move(input_vector->at(ranges_[i].first)))
|
||||||
|
.At(cc->InputTimestamp()));
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
for (int i = 0; i < ranges_.size(); ++i) {
|
||||||
|
auto output = absl::make_unique<std::vector<T>>();
|
||||||
|
output->insert(
|
||||||
|
output->end(),
|
||||||
|
std::make_move_iterator(input_vector->begin() + ranges_[i].first),
|
||||||
|
std::make_move_iterator(input_vector->begin() +
|
||||||
|
ranges_[i].second));
|
||||||
|
cc->Outputs().Index(i).Add(output.release(), cc->InputTimestamp());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename U, IsNotMovable<U> = true>
|
||||||
|
::mediapipe::Status ProcessMovableElements(CalculatorContext* cc) {
|
||||||
|
return ::mediapipe::InternalError("Cannot move non-movable elements.");
|
||||||
|
}
|
||||||
|
|
||||||
private:
|
private:
|
||||||
|
static ::mediapipe::Status checkRangesDontOverlap(
|
||||||
|
const ::mediapipe::SplitVectorCalculatorOptions& options) {
|
||||||
|
for (int i = 0; i < options.ranges_size() - 1; ++i) {
|
||||||
|
for (int j = i + 1; j < options.ranges_size(); ++j) {
|
||||||
|
const auto& range_0 = options.ranges(i);
|
||||||
|
const auto& range_1 = options.ranges(j);
|
||||||
|
if ((range_0.begin() >= range_1.begin() &&
|
||||||
|
range_0.begin() < range_1.end()) ||
|
||||||
|
(range_1.begin() >= range_0.begin() &&
|
||||||
|
range_1.begin() < range_0.end())) {
|
||||||
|
return ::mediapipe::InvalidArgumentError(
|
||||||
|
"Ranges must be non-overlapping when using combine_outputs "
|
||||||
|
"option.");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
std::vector<std::pair<int32, int32>> ranges_;
|
std::vector<std::pair<int32, int32>> ranges_;
|
||||||
int32 max_range_end_ = -1;
|
int32 max_range_end_ = -1;
|
||||||
int32 total_elements_ = 0;
|
int32 total_elements_ = 0;
|
||||||
|
|||||||
@@ -452,4 +452,243 @@ TEST_F(SplitTfLiteTensorVectorCalculatorTest,
|
|||||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
typedef SplitVectorCalculator<std::unique_ptr<int>, true>
|
||||||
|
MovableSplitUniqueIntPtrCalculator;
|
||||||
|
REGISTER_CALCULATOR(MovableSplitUniqueIntPtrCalculator);
|
||||||
|
|
||||||
|
class MovableSplitUniqueIntPtrCalculatorTest : public ::testing::Test {
|
||||||
|
protected:
|
||||||
|
void ValidateVectorOutput(std::vector<Packet>& output_packets,
|
||||||
|
int expected_elements, int input_begin_index) {
|
||||||
|
ASSERT_EQ(1, output_packets.size());
|
||||||
|
const std::vector<std::unique_ptr<int>>& output_vec =
|
||||||
|
output_packets[0].Get<std::vector<std::unique_ptr<int>>>();
|
||||||
|
ASSERT_EQ(expected_elements, output_vec.size());
|
||||||
|
|
||||||
|
for (int i = 0; i < expected_elements; ++i) {
|
||||||
|
const int expected_value = input_begin_index + i;
|
||||||
|
const std::unique_ptr<int>& result = output_vec[i];
|
||||||
|
ASSERT_NE(result, nullptr);
|
||||||
|
ASSERT_EQ(expected_value, *result);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
void ValidateElementOutput(std::vector<Packet>& output_packets,
|
||||||
|
int expected_value) {
|
||||||
|
ASSERT_EQ(1, output_packets.size());
|
||||||
|
const std::unique_ptr<int>& result =
|
||||||
|
output_packets[0].Get<std::unique_ptr<int>>();
|
||||||
|
ASSERT_NE(result, nullptr);
|
||||||
|
ASSERT_EQ(expected_value, *result);
|
||||||
|
}
|
||||||
|
|
||||||
|
void ValidateCombinedVectorOutput(std::vector<Packet>& output_packets,
|
||||||
|
int expected_elements,
|
||||||
|
std::vector<int>& input_begin_indices,
|
||||||
|
std::vector<int>& input_end_indices) {
|
||||||
|
ASSERT_EQ(1, output_packets.size());
|
||||||
|
ASSERT_EQ(input_begin_indices.size(), input_end_indices.size());
|
||||||
|
const std::vector<std::unique_ptr<int>>& output_vector =
|
||||||
|
output_packets[0].Get<std::vector<std::unique_ptr<int>>>();
|
||||||
|
ASSERT_EQ(expected_elements, output_vector.size());
|
||||||
|
const int num_ranges = input_begin_indices.size();
|
||||||
|
|
||||||
|
int element_id = 0;
|
||||||
|
for (int range_id = 0; range_id < num_ranges; ++range_id) {
|
||||||
|
for (int i = input_begin_indices[range_id];
|
||||||
|
i < input_end_indices[range_id]; ++i) {
|
||||||
|
const int expected_value = i;
|
||||||
|
const std::unique_ptr<int>& result = output_vector[element_id];
|
||||||
|
ASSERT_NE(result, nullptr);
|
||||||
|
ASSERT_EQ(expected_value, *result);
|
||||||
|
++element_id;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, InvalidOverlappingRangesTest) {
|
||||||
|
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "input_vector"
|
||||||
|
node {
|
||||||
|
calculator: "MovableSplitUniqueIntPtrCalculator"
|
||||||
|
input_stream: "input_vector"
|
||||||
|
output_stream: "range_0"
|
||||||
|
options {
|
||||||
|
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||||
|
ranges: { begin: 0 end: 3 }
|
||||||
|
ranges: { begin: 1 end: 4 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
// Run the graph.
|
||||||
|
CalculatorGraph graph;
|
||||||
|
// The graph should fail running because there are overlapping ranges.
|
||||||
|
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, SmokeTest) {
|
||||||
|
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "input_vector"
|
||||||
|
node {
|
||||||
|
calculator: "MovableSplitUniqueIntPtrCalculator"
|
||||||
|
input_stream: "input_vector"
|
||||||
|
output_stream: "range_0"
|
||||||
|
output_stream: "range_1"
|
||||||
|
output_stream: "range_2"
|
||||||
|
options {
|
||||||
|
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||||
|
ranges: { begin: 0 end: 1 }
|
||||||
|
ranges: { begin: 1 end: 4 }
|
||||||
|
ranges: { begin: 4 end: 5 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
std::vector<Packet> range_0_packets;
|
||||||
|
tool::AddVectorSink("range_0", &graph_config, &range_0_packets);
|
||||||
|
std::vector<Packet> range_1_packets;
|
||||||
|
tool::AddVectorSink("range_1", &graph_config, &range_1_packets);
|
||||||
|
std::vector<Packet> range_2_packets;
|
||||||
|
tool::AddVectorSink("range_2", &graph_config, &range_2_packets);
|
||||||
|
|
||||||
|
// Run the graph.
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||||
|
MP_ASSERT_OK(graph.StartRun({}));
|
||||||
|
|
||||||
|
// input_vector : {0, 1, 2, 3, 4, 5}
|
||||||
|
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_vector =
|
||||||
|
absl::make_unique<std::vector<std::unique_ptr<int>>>(6);
|
||||||
|
for (int i = 0; i < 6; ++i) {
|
||||||
|
input_vector->at(i) = absl::make_unique<int>(i);
|
||||||
|
}
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"input_vector", Adopt(input_vector.release()).At(Timestamp(1))));
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
|
|
||||||
|
ValidateVectorOutput(range_0_packets, /*expected_elements=*/1,
|
||||||
|
/*input_begin_index=*/0);
|
||||||
|
ValidateVectorOutput(range_1_packets, /*expected_elements=*/3,
|
||||||
|
/*input_begin_index=*/1);
|
||||||
|
ValidateVectorOutput(range_2_packets, /*expected_elements=*/1,
|
||||||
|
/*input_begin_index=*/4);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, SmokeTestElementOnly) {
|
||||||
|
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "input_vector"
|
||||||
|
node {
|
||||||
|
calculator: "MovableSplitUniqueIntPtrCalculator"
|
||||||
|
input_stream: "input_vector"
|
||||||
|
output_stream: "range_0"
|
||||||
|
output_stream: "range_1"
|
||||||
|
output_stream: "range_2"
|
||||||
|
options {
|
||||||
|
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||||
|
ranges: { begin: 0 end: 1 }
|
||||||
|
ranges: { begin: 2 end: 3 }
|
||||||
|
ranges: { begin: 4 end: 5 }
|
||||||
|
element_only: true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
std::vector<Packet> range_0_packets;
|
||||||
|
tool::AddVectorSink("range_0", &graph_config, &range_0_packets);
|
||||||
|
std::vector<Packet> range_1_packets;
|
||||||
|
tool::AddVectorSink("range_1", &graph_config, &range_1_packets);
|
||||||
|
std::vector<Packet> range_2_packets;
|
||||||
|
tool::AddVectorSink("range_2", &graph_config, &range_2_packets);
|
||||||
|
|
||||||
|
// Run the graph.
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||||
|
MP_ASSERT_OK(graph.StartRun({}));
|
||||||
|
|
||||||
|
// input_vector : {0, 1, 2, 3, 4, 5}
|
||||||
|
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_vector =
|
||||||
|
absl::make_unique<std::vector<std::unique_ptr<int>>>(6);
|
||||||
|
for (int i = 0; i < 6; ++i) {
|
||||||
|
input_vector->at(i) = absl::make_unique<int>(i);
|
||||||
|
}
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"input_vector", Adopt(input_vector.release()).At(Timestamp(1))));
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
|
|
||||||
|
ValidateElementOutput(range_0_packets, /*expected_value=*/0);
|
||||||
|
ValidateElementOutput(range_1_packets, /*expected_value=*/2);
|
||||||
|
ValidateElementOutput(range_2_packets, /*expected_value=*/4);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(MovableSplitUniqueIntPtrCalculatorTest, SmokeTestCombiningOutputs) {
|
||||||
|
// Prepare a graph to use the TestMovableSplitUniqueIntPtrVectorCalculator.
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "input_vector"
|
||||||
|
node {
|
||||||
|
calculator: "MovableSplitUniqueIntPtrCalculator"
|
||||||
|
input_stream: "input_vector"
|
||||||
|
output_stream: "range_0"
|
||||||
|
options {
|
||||||
|
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||||
|
ranges: { begin: 0 end: 1 }
|
||||||
|
ranges: { begin: 2 end: 3 }
|
||||||
|
ranges: { begin: 4 end: 5 }
|
||||||
|
combine_outputs: true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
std::vector<Packet> range_0_packets;
|
||||||
|
tool::AddVectorSink("range_0", &graph_config, &range_0_packets);
|
||||||
|
|
||||||
|
// Run the graph.
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||||
|
MP_ASSERT_OK(graph.StartRun({}));
|
||||||
|
|
||||||
|
// input_vector : {0, 1, 2, 3, 4, 5}
|
||||||
|
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_vector =
|
||||||
|
absl::make_unique<std::vector<std::unique_ptr<int>>>(6);
|
||||||
|
for (int i = 0; i < 6; ++i) {
|
||||||
|
input_vector->at(i) = absl::make_unique<int>(i);
|
||||||
|
}
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"input_vector", Adopt(input_vector.release()).At(Timestamp(1))));
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
|
|
||||||
|
std::vector<int> input_begin_indices = {0, 2, 4};
|
||||||
|
std::vector<int> input_end_indices = {1, 3, 5};
|
||||||
|
ValidateCombinedVectorOutput(range_0_packets, /*expected_elements=*/3,
|
||||||
|
input_begin_indices, input_end_indices);
|
||||||
|
}
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -80,10 +80,20 @@ mediapipe_cc_proto_library(
|
|||||||
name = "opencv_image_encoder_calculator_cc_proto",
|
name = "opencv_image_encoder_calculator_cc_proto",
|
||||||
srcs = ["opencv_image_encoder_calculator.proto"],
|
srcs = ["opencv_image_encoder_calculator.proto"],
|
||||||
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||||
visibility = ["//visibility:public"],
|
visibility = [
|
||||||
|
"//visibility:public",
|
||||||
|
],
|
||||||
deps = [":opencv_image_encoder_calculator_proto"],
|
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(
|
mediapipe_cc_proto_library(
|
||||||
name = "mask_overlay_calculator_cc_proto",
|
name = "mask_overlay_calculator_cc_proto",
|
||||||
srcs = ["mask_overlay_calculator.proto"],
|
srcs = ["mask_overlay_calculator.proto"],
|
||||||
@@ -170,6 +180,7 @@ cc_library(
|
|||||||
srcs = ["opencv_encoded_image_to_image_frame_calculator.cc"],
|
srcs = ["opencv_encoded_image_to_image_frame_calculator.cc"],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
|
":opencv_encoded_image_to_image_frame_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/formats:image_frame_opencv",
|
"//mediapipe/framework/formats:image_frame_opencv",
|
||||||
"//mediapipe/framework/port:opencv_imgcodecs",
|
"//mediapipe/framework/port:opencv_imgcodecs",
|
||||||
@@ -330,6 +341,7 @@ cc_library(
|
|||||||
cc_library(
|
cc_library(
|
||||||
name = "image_cropping_calculator",
|
name = "image_cropping_calculator",
|
||||||
srcs = ["image_cropping_calculator.cc"],
|
srcs = ["image_cropping_calculator.cc"],
|
||||||
|
hdrs = ["image_cropping_calculator.h"],
|
||||||
copts = select({
|
copts = select({
|
||||||
"//mediapipe:apple": [
|
"//mediapipe:apple": [
|
||||||
"-x objective-c++",
|
"-x objective-c++",
|
||||||
@@ -343,9 +355,7 @@ cc_library(
|
|||||||
],
|
],
|
||||||
"//conditions:default": [],
|
"//conditions:default": [],
|
||||||
}),
|
}),
|
||||||
visibility = [
|
visibility = ["//visibility:public"],
|
||||||
"//visibility:public",
|
|
||||||
],
|
|
||||||
deps = [
|
deps = [
|
||||||
":image_cropping_calculator_cc_proto",
|
":image_cropping_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
@@ -369,6 +379,22 @@ cc_library(
|
|||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_test(
|
||||||
|
name = "image_cropping_calculator_test",
|
||||||
|
srcs = ["image_cropping_calculator_test.cc"],
|
||||||
|
deps = [
|
||||||
|
":image_cropping_calculator",
|
||||||
|
":image_cropping_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/formats:rect_cc_proto",
|
||||||
|
"//mediapipe/framework/port:gtest_main",
|
||||||
|
"//mediapipe/framework/port:parse_text_proto",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"//mediapipe/framework/tool:tag_map",
|
||||||
|
"//mediapipe/framework/tool:tag_map_helper",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
cc_library(
|
cc_library(
|
||||||
name = "luminance_calculator",
|
name = "luminance_calculator",
|
||||||
srcs = ["luminance_calculator.cc"],
|
srcs = ["luminance_calculator.cc"],
|
||||||
@@ -405,9 +431,12 @@ cc_library(
|
|||||||
":recolor_calculator_cc_proto",
|
":recolor_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/formats:image_frame",
|
"//mediapipe/framework/formats:image_frame",
|
||||||
|
"//mediapipe/framework/formats:image_frame_opencv",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/util:color_cc_proto",
|
"//mediapipe/util:color_cc_proto",
|
||||||
|
"//mediapipe/framework/port:opencv_core",
|
||||||
|
"//mediapipe/framework/port:opencv_imgproc",
|
||||||
] + select({
|
] + select({
|
||||||
"//mediapipe/gpu:disable_gpu": [],
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
"//conditions:default": [
|
"//conditions:default": [
|
||||||
@@ -537,6 +566,27 @@ proto_library(
|
|||||||
deps = ["//mediapipe/framework:calculator_proto"],
|
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(
|
cc_library(
|
||||||
name = "mask_overlay_calculator",
|
name = "mask_overlay_calculator",
|
||||||
srcs = ["mask_overlay_calculator.cc"],
|
srcs = ["mask_overlay_calculator.cc"],
|
||||||
@@ -552,3 +602,30 @@ cc_library(
|
|||||||
],
|
],
|
||||||
alwayslink = 1,
|
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,
|
||||||
|
)
|
||||||
|
|||||||
@@ -75,6 +75,11 @@ class ColorConvertCalculator : public CalculatorBase {
|
|||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
|
cc->SetOffset(TimestampDiff(0));
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
private:
|
private:
|
||||||
// Wrangles the appropriate inputs and outputs to perform the color
|
// Wrangles the appropriate inputs and outputs to perform the color
|
||||||
// conversion. The ImageFrame on input_tag is converted using the
|
// conversion. The ImageFrame on input_tag is converted using the
|
||||||
|
|||||||
@@ -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];
|
||||||
|
}
|
||||||
@@ -12,10 +12,10 @@
|
|||||||
// See the License for the specific language governing permissions and
|
// See the License for the specific language governing permissions and
|
||||||
// limitations under the License.
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/image/image_cropping_calculator.h"
|
||||||
|
|
||||||
#include <cmath>
|
#include <cmath>
|
||||||
|
|
||||||
#include "mediapipe/calculators/image/image_cropping_calculator.pb.h"
|
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
|
||||||
#include "mediapipe/framework/formats/image_frame.h"
|
#include "mediapipe/framework/formats/image_frame.h"
|
||||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||||
#include "mediapipe/framework/formats/rect.pb.h"
|
#include "mediapipe/framework/formats/rect.pb.h"
|
||||||
@@ -25,7 +25,6 @@
|
|||||||
#include "mediapipe/framework/port/status.h"
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
|
||||||
#include "mediapipe/gpu/gl_simple_shaders.h"
|
#include "mediapipe/gpu/gl_simple_shaders.h"
|
||||||
#include "mediapipe/gpu/gpu_buffer.h"
|
#include "mediapipe/gpu/gpu_buffer.h"
|
||||||
#include "mediapipe/gpu/shader_util.h"
|
#include "mediapipe/gpu/shader_util.h"
|
||||||
@@ -52,62 +51,6 @@ constexpr char kWidthTag[] = "WIDTH";
|
|||||||
|
|
||||||
} // namespace
|
} // namespace
|
||||||
|
|
||||||
// Crops the input texture to the given rectangle region. The rectangle can
|
|
||||||
// be at arbitrary location on the image with rotation. If there's rotation, the
|
|
||||||
// output texture will have the size of the input rectangle. The rotation should
|
|
||||||
// be in radian, see rect.proto for detail.
|
|
||||||
//
|
|
||||||
// Input:
|
|
||||||
// One of the following two tags:
|
|
||||||
// IMAGE - ImageFrame representing the input image.
|
|
||||||
// IMAGE_GPU - GpuBuffer representing the input image.
|
|
||||||
// One of the following two tags (optional if WIDTH/HEIGHT is specified):
|
|
||||||
// RECT - A Rect proto specifying the width/height and location of the
|
|
||||||
// cropping rectangle.
|
|
||||||
// NORM_RECT - A NormalizedRect proto specifying the width/height and location
|
|
||||||
// of the cropping rectangle in normalized coordinates.
|
|
||||||
// Alternative tags to RECT (optional if RECT/NORM_RECT is specified):
|
|
||||||
// WIDTH - The desired width of the output cropped image,
|
|
||||||
// based on image center
|
|
||||||
// HEIGHT - The desired height of the output cropped image,
|
|
||||||
// based on image center
|
|
||||||
//
|
|
||||||
// Output:
|
|
||||||
// One of the following two tags:
|
|
||||||
// IMAGE - Cropped ImageFrame
|
|
||||||
// IMAGE_GPU - Cropped GpuBuffer.
|
|
||||||
//
|
|
||||||
// Note: input_stream values take precedence over options defined in the graph.
|
|
||||||
//
|
|
||||||
class ImageCroppingCalculator : public CalculatorBase {
|
|
||||||
public:
|
|
||||||
ImageCroppingCalculator() = default;
|
|
||||||
~ImageCroppingCalculator() 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:
|
|
||||||
::mediapipe::Status RenderCpu(CalculatorContext* cc);
|
|
||||||
::mediapipe::Status RenderGpu(CalculatorContext* cc);
|
|
||||||
::mediapipe::Status InitGpu(CalculatorContext* cc);
|
|
||||||
void GlRender();
|
|
||||||
void GetOutputDimensions(CalculatorContext* cc, int src_width, int src_height,
|
|
||||||
int* dst_width, int* dst_height);
|
|
||||||
|
|
||||||
mediapipe::ImageCroppingCalculatorOptions options_;
|
|
||||||
|
|
||||||
bool use_gpu_ = false;
|
|
||||||
// Output texture corners (4) after transoformation in normalized coordinates.
|
|
||||||
float transformed_points_[8];
|
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
|
||||||
bool gpu_initialized_ = false;
|
|
||||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
|
||||||
GLuint program_ = 0;
|
|
||||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
|
||||||
};
|
|
||||||
REGISTER_CALCULATOR(ImageCroppingCalculator);
|
REGISTER_CALCULATOR(ImageCroppingCalculator);
|
||||||
|
|
||||||
::mediapipe::Status ImageCroppingCalculator::GetContract(
|
::mediapipe::Status ImageCroppingCalculator::GetContract(
|
||||||
@@ -132,7 +75,28 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
|
|||||||
}
|
}
|
||||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
RET_CHECK(cc->Inputs().HasTag(kRectTag) ^ cc->Inputs().HasTag(kNormRectTag));
|
int flags = 0;
|
||||||
|
if (cc->Inputs().HasTag(kRectTag)) {
|
||||||
|
++flags;
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag(kWidthTag) && cc->Inputs().HasTag(kHeightTag)) {
|
||||||
|
++flags;
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag(kNormRectTag)) {
|
||||||
|
++flags;
|
||||||
|
}
|
||||||
|
if (cc->Options<mediapipe::ImageCroppingCalculatorOptions>()
|
||||||
|
.has_norm_width() &&
|
||||||
|
cc->Options<mediapipe::ImageCroppingCalculatorOptions>()
|
||||||
|
.has_norm_height()) {
|
||||||
|
++flags;
|
||||||
|
}
|
||||||
|
if (cc->Options<mediapipe::ImageCroppingCalculatorOptions>().has_width() &&
|
||||||
|
cc->Options<mediapipe::ImageCroppingCalculatorOptions>().has_height()) {
|
||||||
|
++flags;
|
||||||
|
}
|
||||||
|
RET_CHECK(flags == 1) << "Illegal combination of input streams/options.";
|
||||||
|
|
||||||
if (cc->Inputs().HasTag(kRectTag)) {
|
if (cc->Inputs().HasTag(kRectTag)) {
|
||||||
cc->Inputs().Tag(kRectTag).Set<Rect>();
|
cc->Inputs().Tag(kRectTag).Set<Rect>();
|
||||||
}
|
}
|
||||||
@@ -172,6 +136,13 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
|
|||||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Validate border mode.
|
||||||
|
if (use_gpu_) {
|
||||||
|
MP_RETURN_IF_ERROR(ValidateBorderModeForGPU(cc));
|
||||||
|
} else {
|
||||||
|
MP_RETURN_IF_ERROR(ValidateBorderModeForCPU(cc));
|
||||||
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -215,6 +186,32 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
|
|||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status ImageCroppingCalculator::ValidateBorderModeForCPU(
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
int border_mode;
|
||||||
|
return GetBorderModeForOpenCV(cc, &border_mode);
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status ImageCroppingCalculator::ValidateBorderModeForGPU(
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
mediapipe::ImageCroppingCalculatorOptions options =
|
||||||
|
cc->Options<mediapipe::ImageCroppingCalculatorOptions>();
|
||||||
|
|
||||||
|
switch (options.border_mode()) {
|
||||||
|
case mediapipe::ImageCroppingCalculatorOptions::BORDER_ZERO:
|
||||||
|
LOG(WARNING) << "BORDER_ZERO mode is not supported by GPU "
|
||||||
|
<< "implementation and will fall back into BORDER_REPLICATE";
|
||||||
|
break;
|
||||||
|
case mediapipe::ImageCroppingCalculatorOptions::BORDER_REPLICATE:
|
||||||
|
break;
|
||||||
|
default:
|
||||||
|
RET_CHECK_FAIL() << "Unsupported border mode for GPU: "
|
||||||
|
<< options.border_mode();
|
||||||
|
}
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
::mediapipe::Status ImageCroppingCalculator::RenderCpu(CalculatorContext* cc) {
|
::mediapipe::Status ImageCroppingCalculator::RenderCpu(CalculatorContext* cc) {
|
||||||
if (cc->Inputs().Tag(kImageTag).IsEmpty()) {
|
if (cc->Inputs().Tag(kImageTag).IsEmpty()) {
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
@@ -222,41 +219,14 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
|
|||||||
const auto& input_img = cc->Inputs().Tag(kImageTag).Get<ImageFrame>();
|
const auto& input_img = cc->Inputs().Tag(kImageTag).Get<ImageFrame>();
|
||||||
cv::Mat input_mat = formats::MatView(&input_img);
|
cv::Mat input_mat = formats::MatView(&input_img);
|
||||||
|
|
||||||
float rect_center_x = input_img.Width() / 2.0f;
|
RectSpec specs = GetCropSpecs(cc, input_img.Width(), input_img.Height());
|
||||||
float rect_center_y = input_img.Height() / 2.0f;
|
int target_width = specs.width, target_height = specs.height,
|
||||||
float rotation = 0.0f;
|
rect_center_x = specs.center_x, rect_center_y = specs.center_y;
|
||||||
int target_width = input_img.Width();
|
float rotation = specs.rotation;
|
||||||
int target_height = input_img.Height();
|
|
||||||
if (cc->Inputs().HasTag(kRectTag)) {
|
// Get border mode and value for OpenCV.
|
||||||
const auto& rect = cc->Inputs().Tag(kRectTag).Get<Rect>();
|
int border_mode;
|
||||||
if (rect.width() > 0 && rect.height() > 0 && rect.x_center() >= 0 &&
|
MP_RETURN_IF_ERROR(GetBorderModeForOpenCV(cc, &border_mode));
|
||||||
rect.y_center() >= 0) {
|
|
||||||
rect_center_x = rect.x_center();
|
|
||||||
rect_center_y = rect.y_center();
|
|
||||||
target_width = rect.width();
|
|
||||||
target_height = rect.height();
|
|
||||||
rotation = rect.rotation();
|
|
||||||
}
|
|
||||||
} else if (cc->Inputs().HasTag(kNormRectTag)) {
|
|
||||||
const auto& rect = cc->Inputs().Tag(kNormRectTag).Get<NormalizedRect>();
|
|
||||||
if (rect.width() > 0.0 && rect.height() > 0.0 && rect.x_center() >= 0.0 &&
|
|
||||||
rect.y_center() >= 0.0) {
|
|
||||||
rect_center_x = std::round(rect.x_center() * input_img.Width());
|
|
||||||
rect_center_y = std::round(rect.y_center() * input_img.Height());
|
|
||||||
target_width = std::round(rect.width() * input_img.Width());
|
|
||||||
target_height = std::round(rect.height() * input_img.Height());
|
|
||||||
rotation = rect.rotation();
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
if (cc->Inputs().HasTag(kWidthTag) && cc->Inputs().HasTag(kHeightTag)) {
|
|
||||||
target_width = cc->Inputs().Tag(kWidthTag).Get<int>();
|
|
||||||
target_height = cc->Inputs().Tag(kHeightTag).Get<int>();
|
|
||||||
} else if (options_.has_width() && options_.has_height()) {
|
|
||||||
target_width = options_.width();
|
|
||||||
target_height = options_.height();
|
|
||||||
}
|
|
||||||
rotation = options_.rotation();
|
|
||||||
}
|
|
||||||
|
|
||||||
const cv::RotatedRect min_rect(cv::Point2f(rect_center_x, rect_center_y),
|
const cv::RotatedRect min_rect(cv::Point2f(rect_center_x, rect_center_y),
|
||||||
cv::Size2f(target_width, target_height),
|
cv::Size2f(target_width, target_height),
|
||||||
@@ -277,7 +247,9 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
|
|||||||
cv::getPerspectiveTransform(src_points, dst_points);
|
cv::getPerspectiveTransform(src_points, dst_points);
|
||||||
cv::Mat cropped_image;
|
cv::Mat cropped_image;
|
||||||
cv::warpPerspective(input_mat, cropped_image, projection_matrix,
|
cv::warpPerspective(input_mat, cropped_image, projection_matrix,
|
||||||
cv::Size(min_rect.size.width, min_rect.size.height));
|
cv::Size(min_rect.size.width, min_rect.size.height),
|
||||||
|
/* flags = */ 0,
|
||||||
|
/* borderMode = */ border_mode);
|
||||||
|
|
||||||
std::unique_ptr<ImageFrame> output_frame(new ImageFrame(
|
std::unique_ptr<ImageFrame> output_frame(new ImageFrame(
|
||||||
input_img.Format(), cropped_image.cols, cropped_image.rows));
|
input_img.Format(), cropped_image.cols, cropped_image.rows));
|
||||||
@@ -433,46 +405,10 @@ void ImageCroppingCalculator::GetOutputDimensions(CalculatorContext* cc,
|
|||||||
int src_width, int src_height,
|
int src_width, int src_height,
|
||||||
int* dst_width,
|
int* dst_width,
|
||||||
int* dst_height) {
|
int* dst_height) {
|
||||||
// Get the size of the cropping box.
|
RectSpec specs = GetCropSpecs(cc, src_width, src_height);
|
||||||
int crop_width = src_width;
|
int crop_width = specs.width, crop_height = specs.height,
|
||||||
int crop_height = src_height;
|
x_center = specs.center_x, y_center = specs.center_y;
|
||||||
// Get the center of cropping box. Default is the at the center.
|
float rotation = specs.rotation;
|
||||||
int x_center = src_width / 2;
|
|
||||||
int y_center = src_height / 2;
|
|
||||||
// Get the rotation of the cropping box.
|
|
||||||
float rotation = 0.0f;
|
|
||||||
if (cc->Inputs().HasTag(kRectTag)) {
|
|
||||||
const auto& rect = cc->Inputs().Tag(kRectTag).Get<Rect>();
|
|
||||||
// Only use the rect if it is valid.
|
|
||||||
if (rect.width() > 0 && rect.height() > 0 && rect.x_center() >= 0 &&
|
|
||||||
rect.y_center() >= 0) {
|
|
||||||
x_center = rect.x_center();
|
|
||||||
y_center = rect.y_center();
|
|
||||||
crop_width = rect.width();
|
|
||||||
crop_height = rect.height();
|
|
||||||
rotation = rect.rotation();
|
|
||||||
}
|
|
||||||
} else if (cc->Inputs().HasTag(kNormRectTag)) {
|
|
||||||
const auto& rect = cc->Inputs().Tag(kNormRectTag).Get<NormalizedRect>();
|
|
||||||
// Only use the rect if it is valid.
|
|
||||||
if (rect.width() > 0.0 && rect.height() > 0.0 && rect.x_center() >= 0.0 &&
|
|
||||||
rect.y_center() >= 0.0) {
|
|
||||||
x_center = std::round(rect.x_center() * src_width);
|
|
||||||
y_center = std::round(rect.y_center() * src_height);
|
|
||||||
crop_width = std::round(rect.width() * src_width);
|
|
||||||
crop_height = std::round(rect.height() * src_height);
|
|
||||||
rotation = rect.rotation();
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
if (cc->Inputs().HasTag(kWidthTag) && cc->Inputs().HasTag(kHeightTag)) {
|
|
||||||
crop_width = cc->Inputs().Tag(kWidthTag).Get<int>();
|
|
||||||
crop_height = cc->Inputs().Tag(kHeightTag).Get<int>();
|
|
||||||
} else if (options_.has_width() && options_.has_height()) {
|
|
||||||
crop_width = options_.width();
|
|
||||||
crop_height = options_.height();
|
|
||||||
}
|
|
||||||
rotation = options_.rotation();
|
|
||||||
}
|
|
||||||
|
|
||||||
const float half_width = crop_width / 2.0f;
|
const float half_width = crop_width / 2.0f;
|
||||||
const float half_height = crop_height / 2.0f;
|
const float half_height = crop_height / 2.0f;
|
||||||
@@ -501,8 +437,110 @@ void ImageCroppingCalculator::GetOutputDimensions(CalculatorContext* cc,
|
|||||||
row_max = std::max(row_max, transformed_points_[i * 2 + 1]);
|
row_max = std::max(row_max, transformed_points_[i * 2 + 1]);
|
||||||
}
|
}
|
||||||
|
|
||||||
*dst_width = std::round((col_max - col_min) * src_width);
|
int width = static_cast<int>(std::round((col_max - col_min) * src_width));
|
||||||
*dst_height = std::round((row_max - row_min) * src_height);
|
int height = static_cast<int>(std::round((row_max - row_min) * src_height));
|
||||||
|
// Minimum output dimension 1x1 prevents creation of textures with 0x0.
|
||||||
|
*dst_width = std::max(1, width);
|
||||||
|
*dst_height = std::max(1, height);
|
||||||
|
}
|
||||||
|
|
||||||
|
RectSpec ImageCroppingCalculator::GetCropSpecs(const CalculatorContext* cc,
|
||||||
|
int src_width, int src_height) {
|
||||||
|
// Get the size of the cropping box.
|
||||||
|
int crop_width = src_width;
|
||||||
|
int crop_height = src_height;
|
||||||
|
// Get the center of cropping box. Default is the at the center.
|
||||||
|
int x_center = src_width / 2;
|
||||||
|
int y_center = src_height / 2;
|
||||||
|
// Get the rotation of the cropping box.
|
||||||
|
float rotation = 0.0f;
|
||||||
|
// Get the normalized width and height if specified by the inputs or options.
|
||||||
|
float normalized_width = 0.0f;
|
||||||
|
float normalized_height = 0.0f;
|
||||||
|
|
||||||
|
mediapipe::ImageCroppingCalculatorOptions options =
|
||||||
|
cc->Options<mediapipe::ImageCroppingCalculatorOptions>();
|
||||||
|
|
||||||
|
// width/height, norm_width/norm_height from input streams take precednece.
|
||||||
|
if (cc->Inputs().HasTag(kRectTag)) {
|
||||||
|
const auto& rect = cc->Inputs().Tag(kRectTag).Get<Rect>();
|
||||||
|
// Only use the rect if it is valid.
|
||||||
|
if (rect.width() > 0 && rect.height() > 0 && rect.x_center() >= 0 &&
|
||||||
|
rect.y_center() >= 0) {
|
||||||
|
x_center = rect.x_center();
|
||||||
|
y_center = rect.y_center();
|
||||||
|
crop_width = rect.width();
|
||||||
|
crop_height = rect.height();
|
||||||
|
rotation = rect.rotation();
|
||||||
|
}
|
||||||
|
} else if (cc->Inputs().HasTag(kNormRectTag)) {
|
||||||
|
const auto& norm_rect =
|
||||||
|
cc->Inputs().Tag(kNormRectTag).Get<NormalizedRect>();
|
||||||
|
if (norm_rect.width() > 0.0 && norm_rect.height() > 0.0) {
|
||||||
|
normalized_width = norm_rect.width();
|
||||||
|
normalized_height = norm_rect.height();
|
||||||
|
x_center = std::round(norm_rect.x_center() * src_width);
|
||||||
|
y_center = std::round(norm_rect.y_center() * src_height);
|
||||||
|
rotation = norm_rect.rotation();
|
||||||
|
}
|
||||||
|
} else if (cc->Inputs().HasTag(kWidthTag) &&
|
||||||
|
cc->Inputs().HasTag(kHeightTag)) {
|
||||||
|
crop_width = cc->Inputs().Tag(kWidthTag).Get<int>();
|
||||||
|
crop_height = cc->Inputs().Tag(kHeightTag).Get<int>();
|
||||||
|
} else if (options.has_width() && options.has_height()) {
|
||||||
|
crop_width = options.width();
|
||||||
|
crop_height = options.height();
|
||||||
|
} else if (options.has_norm_width() && options.has_norm_height()) {
|
||||||
|
normalized_width = options.norm_width();
|
||||||
|
normalized_height = options.norm_height();
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get the crop width and height from the normalized width and height.
|
||||||
|
if (normalized_width > 0 && normalized_height > 0) {
|
||||||
|
crop_width = std::round(normalized_width * src_width);
|
||||||
|
crop_height = std::round(normalized_height * src_height);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Rotation and center values from input streams take precedence, so only
|
||||||
|
// look at those values in the options if kRectTag and kNormRectTag are not
|
||||||
|
// present from the inputs.
|
||||||
|
if (!cc->Inputs().HasTag(kRectTag) && !cc->Inputs().HasTag(kNormRectTag)) {
|
||||||
|
if (options.has_norm_center_x() && options.has_norm_center_y()) {
|
||||||
|
x_center = std::round(options.norm_center_x() * src_width);
|
||||||
|
y_center = std::round(options.norm_center_y() * src_height);
|
||||||
|
}
|
||||||
|
if (options.has_rotation()) {
|
||||||
|
rotation = options.rotation();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return {
|
||||||
|
.width = crop_width,
|
||||||
|
.height = crop_height,
|
||||||
|
.center_x = x_center,
|
||||||
|
.center_y = y_center,
|
||||||
|
.rotation = rotation,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status ImageCroppingCalculator::GetBorderModeForOpenCV(
|
||||||
|
CalculatorContext* cc, int* border_mode) {
|
||||||
|
mediapipe::ImageCroppingCalculatorOptions options =
|
||||||
|
cc->Options<mediapipe::ImageCroppingCalculatorOptions>();
|
||||||
|
|
||||||
|
switch (options.border_mode()) {
|
||||||
|
case mediapipe::ImageCroppingCalculatorOptions::BORDER_ZERO:
|
||||||
|
*border_mode = cv::BORDER_CONSTANT;
|
||||||
|
break;
|
||||||
|
case mediapipe::ImageCroppingCalculatorOptions::BORDER_REPLICATE:
|
||||||
|
*border_mode = cv::BORDER_REPLICATE;
|
||||||
|
break;
|
||||||
|
default:
|
||||||
|
RET_CHECK_FAIL() << "Unsupported border mode for CPU: "
|
||||||
|
<< options.border_mode();
|
||||||
|
}
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -0,0 +1,91 @@
|
|||||||
|
#ifndef MEDIAPIPE_CALCULATORS_IMAGE_IMAGE_CROPPING_CALCULATOR_H_
|
||||||
|
#define MEDIAPIPE_CALCULATORS_IMAGE_IMAGE_CROPPING_CALCULATOR_H_
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/image/image_cropping_calculator.pb.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
|
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
|
// Crops the input texture to the given rectangle region. The rectangle can
|
||||||
|
// be at arbitrary location on the image with rotation. If there's rotation, the
|
||||||
|
// output texture will have the size of the input rectangle. The rotation should
|
||||||
|
// be in radian, see rect.proto for detail.
|
||||||
|
//
|
||||||
|
// Input:
|
||||||
|
// One of the following two tags:
|
||||||
|
// IMAGE - ImageFrame representing the input image.
|
||||||
|
// IMAGE_GPU - GpuBuffer representing the input image.
|
||||||
|
// One of the following two tags (optional if WIDTH/HEIGHT is specified):
|
||||||
|
// RECT - A Rect proto specifying the width/height and location of the
|
||||||
|
// cropping rectangle.
|
||||||
|
// NORM_RECT - A NormalizedRect proto specifying the width/height and location
|
||||||
|
// of the cropping rectangle in normalized coordinates.
|
||||||
|
// Alternative tags to RECT (optional if RECT/NORM_RECT is specified):
|
||||||
|
// WIDTH - The desired width of the output cropped image,
|
||||||
|
// based on image center
|
||||||
|
// HEIGHT - The desired height of the output cropped image,
|
||||||
|
// based on image center
|
||||||
|
//
|
||||||
|
// Output:
|
||||||
|
// One of the following two tags:
|
||||||
|
// IMAGE - Cropped ImageFrame
|
||||||
|
// IMAGE_GPU - Cropped GpuBuffer.
|
||||||
|
//
|
||||||
|
// Note: input_stream values take precedence over options defined in the graph.
|
||||||
|
//
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
struct RectSpec {
|
||||||
|
int width;
|
||||||
|
int height;
|
||||||
|
int center_x;
|
||||||
|
int center_y;
|
||||||
|
float rotation;
|
||||||
|
|
||||||
|
bool operator==(const RectSpec& rect) const {
|
||||||
|
return (width == rect.width && height == rect.height &&
|
||||||
|
center_x == rect.center_x && center_y == rect.center_y &&
|
||||||
|
rotation == rect.rotation);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
class ImageCroppingCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
ImageCroppingCalculator() = default;
|
||||||
|
~ImageCroppingCalculator() 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;
|
||||||
|
static RectSpec GetCropSpecs(const CalculatorContext* cc, int src_width,
|
||||||
|
int src_height);
|
||||||
|
|
||||||
|
private:
|
||||||
|
::mediapipe::Status ValidateBorderModeForCPU(CalculatorContext* cc);
|
||||||
|
::mediapipe::Status ValidateBorderModeForGPU(CalculatorContext* cc);
|
||||||
|
::mediapipe::Status RenderCpu(CalculatorContext* cc);
|
||||||
|
::mediapipe::Status RenderGpu(CalculatorContext* cc);
|
||||||
|
::mediapipe::Status InitGpu(CalculatorContext* cc);
|
||||||
|
void GlRender();
|
||||||
|
void GetOutputDimensions(CalculatorContext* cc, int src_width, int src_height,
|
||||||
|
int* dst_width, int* dst_height);
|
||||||
|
::mediapipe::Status GetBorderModeForOpenCV(CalculatorContext* cc,
|
||||||
|
int* border_mode);
|
||||||
|
|
||||||
|
mediapipe::ImageCroppingCalculatorOptions options_;
|
||||||
|
|
||||||
|
bool use_gpu_ = false;
|
||||||
|
// Output texture corners (4) after transoformation in normalized coordinates.
|
||||||
|
float transformed_points_[8];
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
|
bool gpu_initialized_ = false;
|
||||||
|
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||||
|
GLuint program_ = 0;
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
|
#endif // MEDIAPIPE_CALCULATORS_IMAGE_IMAGE_CROPPING_CALCULATOR_H_
|
||||||
@@ -30,4 +30,25 @@ message ImageCroppingCalculatorOptions {
|
|||||||
|
|
||||||
// Rotation angle is counter-clockwise in radian.
|
// Rotation angle is counter-clockwise in radian.
|
||||||
optional float rotation = 3 [default = 0.0];
|
optional float rotation = 3 [default = 0.0];
|
||||||
|
|
||||||
|
// Normalized width and height of the output rect. Value is within [0, 1].
|
||||||
|
optional float norm_width = 4;
|
||||||
|
optional float norm_height = 5;
|
||||||
|
|
||||||
|
// Normalized location of the center of the output
|
||||||
|
// rectangle in image coordinates. Value is within [0, 1].
|
||||||
|
// The (0, 0) point is at the (top, left) corner.
|
||||||
|
optional float norm_center_x = 6 [default = 0];
|
||||||
|
optional float norm_center_y = 7 [default = 0];
|
||||||
|
|
||||||
|
enum BorderMode {
|
||||||
|
// First unspecified value is required by the guideline. See details here:
|
||||||
|
// https://developers.google.com/protocol-buffers/docs/style#enums
|
||||||
|
BORDER_UNSPECIFIED = 0;
|
||||||
|
BORDER_ZERO = 1;
|
||||||
|
BORDER_REPLICATE = 2;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Specifies behaviour for crops that go beyond image borders.
|
||||||
|
optional BorderMode border_mode = 8 [default = BORDER_ZERO];
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,216 @@
|
|||||||
|
// 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/image/image_cropping_calculator.h"
|
||||||
|
|
||||||
|
#include <cmath>
|
||||||
|
#include <memory>
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/image/image_cropping_calculator.pb.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/formats/rect.pb.h"
|
||||||
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||||
|
#include "mediapipe/framework/port/status_matchers.h"
|
||||||
|
#include "mediapipe/framework/tool/tag_map.h"
|
||||||
|
#include "mediapipe/framework/tool/tag_map_helper.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
constexpr int input_width = 100;
|
||||||
|
constexpr int input_height = 100;
|
||||||
|
|
||||||
|
constexpr char kRectTag[] = "RECT";
|
||||||
|
constexpr char kHeightTag[] = "HEIGHT";
|
||||||
|
constexpr char kWidthTag[] = "WIDTH";
|
||||||
|
|
||||||
|
// Test normal case, where norm_width and norm_height in options are set.
|
||||||
|
TEST(ImageCroppingCalculatorTest, GetCroppingDimensionsNormal) {
|
||||||
|
auto calculator_node =
|
||||||
|
ParseTextProtoOrDie<mediapipe::CalculatorGraphConfig::Node>(
|
||||||
|
R"(
|
||||||
|
calculator: "ImageCroppingCalculator"
|
||||||
|
input_stream: "IMAGE_GPU:input_frames"
|
||||||
|
output_stream: "IMAGE_GPU:cropped_output_frames"
|
||||||
|
options: {
|
||||||
|
[mediapipe.ImageCroppingCalculatorOptions.ext] {
|
||||||
|
norm_width: 0.6
|
||||||
|
norm_height: 0.6
|
||||||
|
norm_center_x: 0.5
|
||||||
|
norm_center_y: 0.5
|
||||||
|
rotation: 0.3
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
auto calculator_state = absl::make_unique<CalculatorState>(
|
||||||
|
"Node", 0, "Calculator", calculator_node, nullptr);
|
||||||
|
auto cc = absl::make_unique<CalculatorContext>(
|
||||||
|
calculator_state.get(), tool::CreateTagMap({}).ValueOrDie(),
|
||||||
|
tool::CreateTagMap({}).ValueOrDie());
|
||||||
|
|
||||||
|
RectSpec expectRect = {
|
||||||
|
.width = 60,
|
||||||
|
.height = 60,
|
||||||
|
.center_x = 50,
|
||||||
|
.center_y = 50,
|
||||||
|
.rotation = 0.3,
|
||||||
|
};
|
||||||
|
EXPECT_EQ(ImageCroppingCalculator::GetCropSpecs(cc.get(), input_width,
|
||||||
|
input_height),
|
||||||
|
expectRect);
|
||||||
|
} // TEST
|
||||||
|
|
||||||
|
// Test when (width height) + (norm_width norm_height) are set in options.
|
||||||
|
// width and height should take precedence.
|
||||||
|
TEST(ImageCroppingCalculatorTest, RedundantSpecInOptions) {
|
||||||
|
auto calculator_node =
|
||||||
|
ParseTextProtoOrDie<mediapipe::CalculatorGraphConfig::Node>(
|
||||||
|
R"(
|
||||||
|
calculator: "ImageCroppingCalculator"
|
||||||
|
input_stream: "IMAGE_GPU:input_frames"
|
||||||
|
output_stream: "IMAGE_GPU:cropped_output_frames"
|
||||||
|
options: {
|
||||||
|
[mediapipe.ImageCroppingCalculatorOptions.ext] {
|
||||||
|
width: 50
|
||||||
|
height: 50
|
||||||
|
norm_width: 0.6
|
||||||
|
norm_height: 0.6
|
||||||
|
norm_center_x: 0.5
|
||||||
|
norm_center_y: 0.5
|
||||||
|
rotation: 0.3
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
auto calculator_state = absl::make_unique<CalculatorState>(
|
||||||
|
"Node", 0, "Calculator", calculator_node, nullptr);
|
||||||
|
auto cc = absl::make_unique<CalculatorContext>(
|
||||||
|
calculator_state.get(), tool::CreateTagMap({}).ValueOrDie(),
|
||||||
|
tool::CreateTagMap({}).ValueOrDie());
|
||||||
|
RectSpec expectRect = {
|
||||||
|
.width = 50,
|
||||||
|
.height = 50,
|
||||||
|
.center_x = 50,
|
||||||
|
.center_y = 50,
|
||||||
|
.rotation = 0.3,
|
||||||
|
};
|
||||||
|
EXPECT_EQ(ImageCroppingCalculator::GetCropSpecs(cc.get(), input_width,
|
||||||
|
input_height),
|
||||||
|
expectRect);
|
||||||
|
} // TEST
|
||||||
|
|
||||||
|
// Test when WIDTH HEIGHT are set from input stream,
|
||||||
|
// and options has norm_width/height set.
|
||||||
|
// WIDTH HEIGHT from input stream should take precedence.
|
||||||
|
TEST(ImageCroppingCalculatorTest, RedundantSpectWithInputStream) {
|
||||||
|
auto calculator_node =
|
||||||
|
ParseTextProtoOrDie<mediapipe::CalculatorGraphConfig::Node>(
|
||||||
|
R"(
|
||||||
|
calculator: "ImageCroppingCalculator"
|
||||||
|
input_stream: "IMAGE_GPU:input_frames"
|
||||||
|
input_stream: "WIDTH:crop_width"
|
||||||
|
input_stream: "HEIGHT:crop_height"
|
||||||
|
output_stream: "IMAGE_GPU:cropped_output_frames"
|
||||||
|
options: {
|
||||||
|
[mediapipe.ImageCroppingCalculatorOptions.ext] {
|
||||||
|
width: 50
|
||||||
|
height: 50
|
||||||
|
norm_width: 0.6
|
||||||
|
norm_height: 0.6
|
||||||
|
norm_center_x: 0.5
|
||||||
|
norm_center_y: 0.5
|
||||||
|
rotation: 0.3
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
auto calculator_state = absl::make_unique<CalculatorState>(
|
||||||
|
"Node", 0, "Calculator", calculator_node, nullptr);
|
||||||
|
auto inputTags = tool::CreateTagMap({
|
||||||
|
"HEIGHT:0:crop_height",
|
||||||
|
"WIDTH:0:crop_width",
|
||||||
|
})
|
||||||
|
.ValueOrDie();
|
||||||
|
auto cc = absl::make_unique<CalculatorContext>(
|
||||||
|
calculator_state.get(), inputTags, tool::CreateTagMap({}).ValueOrDie());
|
||||||
|
auto& inputs = cc->Inputs();
|
||||||
|
inputs.Tag(kHeightTag).Value() = MakePacket<int>(1);
|
||||||
|
inputs.Tag(kWidthTag).Value() = MakePacket<int>(1);
|
||||||
|
RectSpec expectRect = {
|
||||||
|
.width = 1,
|
||||||
|
.height = 1,
|
||||||
|
.center_x = 50,
|
||||||
|
.center_y = 50,
|
||||||
|
.rotation = 0.3,
|
||||||
|
};
|
||||||
|
EXPECT_EQ(ImageCroppingCalculator::GetCropSpecs(cc.get(), input_width,
|
||||||
|
input_height),
|
||||||
|
expectRect);
|
||||||
|
} // TEST
|
||||||
|
|
||||||
|
// Test when RECT is set from input stream,
|
||||||
|
// and options has norm_width/height set.
|
||||||
|
// RECT from input stream should take precedence.
|
||||||
|
TEST(ImageCroppingCalculatorTest, RedundantSpecWithInputStream) {
|
||||||
|
auto calculator_node =
|
||||||
|
ParseTextProtoOrDie<mediapipe::CalculatorGraphConfig::Node>(
|
||||||
|
R"(
|
||||||
|
calculator: "ImageCroppingCalculator"
|
||||||
|
input_stream: "IMAGE_GPU:input_frames"
|
||||||
|
input_stream: "RECT:rect"
|
||||||
|
output_stream: "IMAGE_GPU:cropped_output_frames"
|
||||||
|
options: {
|
||||||
|
[mediapipe.ImageCroppingCalculatorOptions.ext] {
|
||||||
|
width: 50
|
||||||
|
height: 50
|
||||||
|
norm_width: 0.6
|
||||||
|
norm_height: 0.6
|
||||||
|
norm_center_x: 0.5
|
||||||
|
norm_center_y: 0.5
|
||||||
|
rotation: 0.3
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
auto calculator_state = absl::make_unique<CalculatorState>(
|
||||||
|
"Node", 0, "Calculator", calculator_node, nullptr);
|
||||||
|
auto inputTags = tool::CreateTagMap({
|
||||||
|
"RECT:0:rect",
|
||||||
|
})
|
||||||
|
.ValueOrDie();
|
||||||
|
auto cc = absl::make_unique<CalculatorContext>(
|
||||||
|
calculator_state.get(), inputTags, tool::CreateTagMap({}).ValueOrDie());
|
||||||
|
auto& inputs = cc->Inputs();
|
||||||
|
mediapipe::Rect rect = ParseTextProtoOrDie<mediapipe::Rect>(
|
||||||
|
R"(
|
||||||
|
width: 1 height: 1 x_center: 40 y_center: 40 rotation: 0.5
|
||||||
|
)");
|
||||||
|
inputs.Tag(kRectTag).Value() = MakePacket<mediapipe::Rect>(rect);
|
||||||
|
RectSpec expectRect = {
|
||||||
|
.width = 1,
|
||||||
|
.height = 1,
|
||||||
|
.center_x = 40,
|
||||||
|
.center_y = 40,
|
||||||
|
.rotation = 0.5,
|
||||||
|
};
|
||||||
|
EXPECT_EQ(ImageCroppingCalculator::GetCropSpecs(cc.get(), input_width,
|
||||||
|
input_height),
|
||||||
|
expectRect);
|
||||||
|
} // TEST
|
||||||
|
|
||||||
|
} // namespace
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -104,6 +104,14 @@ mediapipe::ScaleMode_Mode ParseScaleMode(
|
|||||||
// to be a multiple of 90 degrees. If provided, it overrides the
|
// to be a multiple of 90 degrees. If provided, it overrides the
|
||||||
// ROTATION_DEGREES input side packet.
|
// ROTATION_DEGREES input side packet.
|
||||||
//
|
//
|
||||||
|
// FLIP_HORIZONTALLY (optional): Whether to flip image horizontally or not. If
|
||||||
|
// provided, it overrides the FLIP_HORIZONTALLY input side packet and/or
|
||||||
|
// corresponding field in the calculator options.
|
||||||
|
//
|
||||||
|
// FLIP_VERTICALLY (optional): Whether to flip image vertically or not. If
|
||||||
|
// provided, it overrides the FLIP_VERTICALLY input side packet and/or
|
||||||
|
// corresponding field in the calculator options.
|
||||||
|
//
|
||||||
// Output:
|
// Output:
|
||||||
// One of the following two tags:
|
// One of the following two tags:
|
||||||
// IMAGE - ImageFrame representing the output image.
|
// IMAGE - ImageFrame representing the output image.
|
||||||
@@ -129,6 +137,12 @@ mediapipe::ScaleMode_Mode ParseScaleMode(
|
|||||||
// degrees. It has to be a multiple of 90 degrees. It overrides the
|
// degrees. It has to be a multiple of 90 degrees. It overrides the
|
||||||
// corresponding field in the calculator options.
|
// corresponding field in the calculator options.
|
||||||
//
|
//
|
||||||
|
// FLIP_HORIZONTALLY (optional): Whether to flip image horizontally or not.
|
||||||
|
// It overrides the corresponding field in the calculator options.
|
||||||
|
//
|
||||||
|
// FLIP_VERTICALLY (optional): Whether to flip image vertically or not.
|
||||||
|
// It overrides the corresponding field in the calculator options.
|
||||||
|
//
|
||||||
// Calculator options (see image_transformation_calculator.proto):
|
// Calculator options (see image_transformation_calculator.proto):
|
||||||
// output_width, output_height - (optional) Desired scaled image size.
|
// output_width, output_height - (optional) Desired scaled image size.
|
||||||
// rotation_mode - (optional) Rotation in multiples of 90 degrees.
|
// rotation_mode - (optional) Rotation in multiples of 90 degrees.
|
||||||
@@ -138,8 +152,7 @@ mediapipe::ScaleMode_Mode ParseScaleMode(
|
|||||||
// Note: To enable horizontal or vertical flipping, specify them in the
|
// Note: To enable horizontal or vertical flipping, specify them in the
|
||||||
// calculator options. Flipping is applied after rotation.
|
// calculator options. Flipping is applied after rotation.
|
||||||
//
|
//
|
||||||
// Note: Only scale mode STRETCH is currently supported on CPU,
|
// Note: Only scale mode STRETCH is currently supported on CPU.
|
||||||
// and flipping is not yet supported either.
|
|
||||||
//
|
//
|
||||||
class ImageTransformationCalculator : public CalculatorBase {
|
class ImageTransformationCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
@@ -168,6 +181,8 @@ class ImageTransformationCalculator : public CalculatorBase {
|
|||||||
int output_height_ = 0;
|
int output_height_ = 0;
|
||||||
mediapipe::RotationMode_Mode rotation_;
|
mediapipe::RotationMode_Mode rotation_;
|
||||||
mediapipe::ScaleMode_Mode scale_mode_;
|
mediapipe::ScaleMode_Mode scale_mode_;
|
||||||
|
bool flip_horizontally_ = false;
|
||||||
|
bool flip_vertically_ = false;
|
||||||
|
|
||||||
bool use_gpu_ = false;
|
bool use_gpu_ = false;
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
@@ -204,6 +219,12 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
if (cc->Inputs().HasTag("ROTATION_DEGREES")) {
|
if (cc->Inputs().HasTag("ROTATION_DEGREES")) {
|
||||||
cc->Inputs().Tag("ROTATION_DEGREES").Set<int>();
|
cc->Inputs().Tag("ROTATION_DEGREES").Set<int>();
|
||||||
}
|
}
|
||||||
|
if (cc->Inputs().HasTag("FLIP_HORIZONTALLY")) {
|
||||||
|
cc->Inputs().Tag("FLIP_HORIZONTALLY").Set<bool>();
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag("FLIP_VERTICALLY")) {
|
||||||
|
cc->Inputs().Tag("FLIP_VERTICALLY").Set<bool>();
|
||||||
|
}
|
||||||
|
|
||||||
if (cc->InputSidePackets().HasTag("OUTPUT_DIMENSIONS")) {
|
if (cc->InputSidePackets().HasTag("OUTPUT_DIMENSIONS")) {
|
||||||
cc->InputSidePackets().Tag("OUTPUT_DIMENSIONS").Set<DimensionsPacketType>();
|
cc->InputSidePackets().Tag("OUTPUT_DIMENSIONS").Set<DimensionsPacketType>();
|
||||||
@@ -211,6 +232,12 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
if (cc->InputSidePackets().HasTag("ROTATION_DEGREES")) {
|
if (cc->InputSidePackets().HasTag("ROTATION_DEGREES")) {
|
||||||
cc->InputSidePackets().Tag("ROTATION_DEGREES").Set<int>();
|
cc->InputSidePackets().Tag("ROTATION_DEGREES").Set<int>();
|
||||||
}
|
}
|
||||||
|
if (cc->InputSidePackets().HasTag("FLIP_HORIZONTALLY")) {
|
||||||
|
cc->InputSidePackets().Tag("FLIP_HORIZONTALLY").Set<bool>();
|
||||||
|
}
|
||||||
|
if (cc->InputSidePackets().HasTag("FLIP_VERTICALLY")) {
|
||||||
|
cc->InputSidePackets().Tag("FLIP_VERTICALLY").Set<bool>();
|
||||||
|
}
|
||||||
|
|
||||||
if (cc->Outputs().HasTag("LETTERBOX_PADDING")) {
|
if (cc->Outputs().HasTag("LETTERBOX_PADDING")) {
|
||||||
cc->Outputs().Tag("LETTERBOX_PADDING").Set<std::array<float, 4>>();
|
cc->Outputs().Tag("LETTERBOX_PADDING").Set<std::array<float, 4>>();
|
||||||
@@ -246,6 +273,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
output_width_ = options_.output_width();
|
output_width_ = options_.output_width();
|
||||||
output_height_ = options_.output_height();
|
output_height_ = options_.output_height();
|
||||||
}
|
}
|
||||||
|
|
||||||
if (cc->InputSidePackets().HasTag("ROTATION_DEGREES")) {
|
if (cc->InputSidePackets().HasTag("ROTATION_DEGREES")) {
|
||||||
rotation_ = DegreesToRotationMode(
|
rotation_ = DegreesToRotationMode(
|
||||||
cc->InputSidePackets().Tag("ROTATION_DEGREES").Get<int>());
|
cc->InputSidePackets().Tag("ROTATION_DEGREES").Get<int>());
|
||||||
@@ -253,6 +281,20 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
rotation_ = options_.rotation_mode();
|
rotation_ = options_.rotation_mode();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (cc->InputSidePackets().HasTag("FLIP_HORIZONTALLY")) {
|
||||||
|
flip_horizontally_ =
|
||||||
|
cc->InputSidePackets().Tag("FLIP_HORIZONTALLY").Get<bool>();
|
||||||
|
} else {
|
||||||
|
flip_horizontally_ = options_.flip_horizontally();
|
||||||
|
}
|
||||||
|
|
||||||
|
if (cc->InputSidePackets().HasTag("FLIP_VERTICALLY")) {
|
||||||
|
flip_vertically_ =
|
||||||
|
cc->InputSidePackets().Tag("FLIP_VERTICALLY").Get<bool>();
|
||||||
|
} else {
|
||||||
|
flip_vertically_ = options_.flip_vertically();
|
||||||
|
}
|
||||||
|
|
||||||
scale_mode_ = ParseScaleMode(options_.scale_mode(), DEFAULT_SCALE_MODE);
|
scale_mode_ = ParseScaleMode(options_.scale_mode(), DEFAULT_SCALE_MODE);
|
||||||
|
|
||||||
if (use_gpu_) {
|
if (use_gpu_) {
|
||||||
@@ -269,12 +311,37 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
|
|
||||||
::mediapipe::Status ImageTransformationCalculator::Process(
|
::mediapipe::Status ImageTransformationCalculator::Process(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
|
// Override values if specified so.
|
||||||
|
if (cc->Inputs().HasTag("ROTATION_DEGREES") &&
|
||||||
|
!cc->Inputs().Tag("ROTATION_DEGREES").IsEmpty()) {
|
||||||
|
rotation_ =
|
||||||
|
DegreesToRotationMode(cc->Inputs().Tag("ROTATION_DEGREES").Get<int>());
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag("FLIP_HORIZONTALLY") &&
|
||||||
|
!cc->Inputs().Tag("FLIP_HORIZONTALLY").IsEmpty()) {
|
||||||
|
flip_horizontally_ = cc->Inputs().Tag("FLIP_HORIZONTALLY").Get<bool>();
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag("FLIP_VERTICALLY") &&
|
||||||
|
!cc->Inputs().Tag("FLIP_VERTICALLY").IsEmpty()) {
|
||||||
|
flip_vertically_ = cc->Inputs().Tag("FLIP_VERTICALLY").Get<bool>();
|
||||||
|
}
|
||||||
|
|
||||||
if (use_gpu_) {
|
if (use_gpu_) {
|
||||||
#if !defined(MEDIAPIPE_DISABLE_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.)
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
return helper_.RunInGlContext(
|
return helper_.RunInGlContext(
|
||||||
[this, cc]() -> ::mediapipe::Status { return RenderGpu(cc); });
|
[this, cc]() -> ::mediapipe::Status { return RenderGpu(cc); });
|
||||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
} else {
|
} else {
|
||||||
|
if (cc->Inputs().Tag("IMAGE").IsEmpty()) {
|
||||||
|
// Image is missing, hence no way to produce output image. (Timestamp
|
||||||
|
// bound will be updated automatically.)
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
return RenderCpu(cc);
|
return RenderCpu(cc);
|
||||||
}
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
@@ -316,6 +383,11 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
cv::Mat input_mat = formats::MatView(&input_img);
|
cv::Mat input_mat = formats::MatView(&input_img);
|
||||||
cv::Mat scaled_mat;
|
cv::Mat scaled_mat;
|
||||||
|
|
||||||
|
if (!output_height_ || !output_width_) {
|
||||||
|
output_height_ = input_height;
|
||||||
|
output_width_ = input_width;
|
||||||
|
}
|
||||||
|
|
||||||
if (scale_mode_ == mediapipe::ScaleMode_Mode_STRETCH) {
|
if (scale_mode_ == mediapipe::ScaleMode_Mode_STRETCH) {
|
||||||
cv::resize(input_mat, scaled_mat, cv::Size(output_width_, output_height_));
|
cv::resize(input_mat, scaled_mat, cv::Size(output_width_, output_height_));
|
||||||
} else {
|
} else {
|
||||||
@@ -356,21 +428,25 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
.Add(padding.release(), cc->InputTimestamp());
|
.Add(padding.release(), cc->InputTimestamp());
|
||||||
}
|
}
|
||||||
|
|
||||||
if (cc->InputSidePackets().HasTag("ROTATION_DEGREES")) {
|
|
||||||
rotation_ = DegreesToRotationMode(
|
|
||||||
cc->InputSidePackets().Tag("ROTATION_DEGREES").Get<int>());
|
|
||||||
}
|
|
||||||
|
|
||||||
cv::Mat rotated_mat;
|
cv::Mat rotated_mat;
|
||||||
const int angle = RotationModeToDegrees(rotation_);
|
const int angle = RotationModeToDegrees(rotation_);
|
||||||
cv::Point2f src_center(scaled_mat.cols / 2.0, scaled_mat.rows / 2.0);
|
cv::Point2f src_center(scaled_mat.cols / 2.0, scaled_mat.rows / 2.0);
|
||||||
cv::Mat rotation_mat = cv::getRotationMatrix2D(src_center, angle, 1.0);
|
cv::Mat rotation_mat = cv::getRotationMatrix2D(src_center, angle, 1.0);
|
||||||
cv::warpAffine(scaled_mat, rotated_mat, rotation_mat, scaled_mat.size());
|
cv::warpAffine(scaled_mat, rotated_mat, rotation_mat, scaled_mat.size());
|
||||||
|
|
||||||
|
cv::Mat flipped_mat;
|
||||||
|
if (flip_horizontally_ || flip_vertically_) {
|
||||||
|
const int flip_code =
|
||||||
|
flip_horizontally_ && flip_vertically_ ? -1 : flip_horizontally_;
|
||||||
|
cv::flip(rotated_mat, flipped_mat, flip_code);
|
||||||
|
} else {
|
||||||
|
flipped_mat = rotated_mat;
|
||||||
|
}
|
||||||
|
|
||||||
std::unique_ptr<ImageFrame> output_frame(
|
std::unique_ptr<ImageFrame> output_frame(
|
||||||
new ImageFrame(input_img.Format(), output_width, output_height));
|
new ImageFrame(input_img.Format(), output_width, output_height));
|
||||||
cv::Mat output_mat = formats::MatView(output_frame.get());
|
cv::Mat output_mat = formats::MatView(output_frame.get());
|
||||||
rotated_mat.copyTo(output_mat);
|
flipped_mat.copyTo(output_mat);
|
||||||
cc->Outputs().Tag("IMAGE").Add(output_frame.release(), cc->InputTimestamp());
|
cc->Outputs().Tag("IMAGE").Add(output_frame.release(), cc->InputTimestamp());
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
@@ -435,14 +511,8 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
}
|
}
|
||||||
RET_CHECK(renderer) << "Unsupported input texture type";
|
RET_CHECK(renderer) << "Unsupported input texture type";
|
||||||
|
|
||||||
if (cc->InputSidePackets().HasTag("ROTATION_DEGREES")) {
|
mediapipe::FrameScaleMode scale_mode = mediapipe::FrameScaleModeFromProto(
|
||||||
rotation_ = DegreesToRotationMode(
|
scale_mode_, mediapipe::FrameScaleMode::kStretch);
|
||||||
cc->InputSidePackets().Tag("ROTATION_DEGREES").Get<int>());
|
|
||||||
}
|
|
||||||
|
|
||||||
static mediapipe::FrameScaleMode scale_mode =
|
|
||||||
mediapipe::FrameScaleModeFromProto(scale_mode_,
|
|
||||||
mediapipe::FrameScaleMode::kStretch);
|
|
||||||
mediapipe::FrameRotation rotation =
|
mediapipe::FrameRotation rotation =
|
||||||
mediapipe::FrameRotationFromDegrees(RotationModeToDegrees(rotation_));
|
mediapipe::FrameRotationFromDegrees(RotationModeToDegrees(rotation_));
|
||||||
|
|
||||||
@@ -455,7 +525,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
|
|
||||||
MP_RETURN_IF_ERROR(renderer->GlRender(
|
MP_RETURN_IF_ERROR(renderer->GlRender(
|
||||||
src1.width(), src1.height(), dst.width(), dst.height(), scale_mode,
|
src1.width(), src1.height(), dst.width(), dst.height(), scale_mode,
|
||||||
rotation, options_.flip_horizontally(), options_.flip_vertically(),
|
rotation, flip_horizontally_, flip_vertically_,
|
||||||
/*flip_texture=*/false));
|
/*flip_texture=*/false));
|
||||||
|
|
||||||
glActiveTexture(GL_TEXTURE1);
|
glActiveTexture(GL_TEXTURE1);
|
||||||
|
|||||||
@@ -12,6 +12,7 @@
|
|||||||
// See the License for the specific language governing permissions and
|
// See the License for the specific language governing permissions and
|
||||||
// limitations under the License.
|
// 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/calculator_framework.h"
|
||||||
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
#include "mediapipe/framework/formats/image_frame_opencv.h"
|
||||||
#include "mediapipe/framework/port/opencv_imgcodecs_inc.h"
|
#include "mediapipe/framework/port/opencv_imgcodecs_inc.h"
|
||||||
@@ -34,7 +35,11 @@ namespace mediapipe {
|
|||||||
class OpenCvEncodedImageToImageFrameCalculator : public CalculatorBase {
|
class OpenCvEncodedImageToImageFrameCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||||
|
|
||||||
|
private:
|
||||||
|
mediapipe::OpenCvEncodedImageToImageFrameCalculatorOptions options_;
|
||||||
};
|
};
|
||||||
|
|
||||||
::mediapipe::Status OpenCvEncodedImageToImageFrameCalculator::GetContract(
|
::mediapipe::Status OpenCvEncodedImageToImageFrameCalculator::GetContract(
|
||||||
@@ -44,13 +49,29 @@ class OpenCvEncodedImageToImageFrameCalculator : public CalculatorBase {
|
|||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status OpenCvEncodedImageToImageFrameCalculator::Open(
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
options_ =
|
||||||
|
cc->Options<mediapipe::OpenCvEncodedImageToImageFrameCalculatorOptions>();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
::mediapipe::Status OpenCvEncodedImageToImageFrameCalculator::Process(
|
::mediapipe::Status OpenCvEncodedImageToImageFrameCalculator::Process(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
const std::string& contents = cc->Inputs().Index(0).Get<std::string>();
|
const std::string& contents = cc->Inputs().Index(0).Get<std::string>();
|
||||||
const std::vector<char> contents_vector(contents.begin(), contents.end());
|
const std::vector<char> contents_vector(contents.begin(), contents.end());
|
||||||
cv::Mat decoded_mat =
|
cv::Mat decoded_mat;
|
||||||
cv::imdecode(contents_vector, -1 /* return the loaded image as-is */);
|
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;
|
ImageFormat::Format image_format = ImageFormat::UNKNOWN;
|
||||||
cv::Mat output_mat;
|
cv::Mat output_mat;
|
||||||
switch (decoded_mat.channels()) {
|
switch (decoded_mat.channels()) {
|
||||||
@@ -70,7 +91,8 @@ class OpenCvEncodedImageToImageFrameCalculator : public CalculatorBase {
|
|||||||
<< "Unsupported number of channels: " << decoded_mat.channels();
|
<< "Unsupported number of channels: " << decoded_mat.channels();
|
||||||
}
|
}
|
||||||
std::unique_ptr<ImageFrame> output_frame = absl::make_unique<ImageFrame>(
|
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()));
|
output_mat.copyTo(formats::MatView(output_frame.get()));
|
||||||
cc->Outputs().Index(0).Add(output_frame.release(), cc->InputTimestamp());
|
cc->Outputs().Index(0).Add(output_frame.release(), cc->InputTimestamp());
|
||||||
return ::mediapipe::OkStatus();
|
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/calculators/image/recolor_calculator.pb.h"
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/formats/image_frame.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/ret_check.h"
|
||||||
#include "mediapipe/framework/port/status.h"
|
#include "mediapipe/framework/port/status.h"
|
||||||
#include "mediapipe/util/color.pb.h"
|
#include "mediapipe/util/color.pb.h"
|
||||||
@@ -39,8 +42,6 @@ namespace mediapipe {
|
|||||||
// The luminance of the input image is used to adjust the blending weight,
|
// The luminance of the input image is used to adjust the blending weight,
|
||||||
// to help preserve image textures.
|
// to help preserve image textures.
|
||||||
//
|
//
|
||||||
// TODO implement cpu support.
|
|
||||||
//
|
|
||||||
// Inputs:
|
// Inputs:
|
||||||
// One of the following IMAGE tags:
|
// One of the following IMAGE tags:
|
||||||
// IMAGE: An ImageFrame input image, RGB or RGBA.
|
// IMAGE: An ImageFrame input image, RGB or RGBA.
|
||||||
@@ -71,6 +72,8 @@ namespace mediapipe {
|
|||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
//
|
//
|
||||||
|
// Note: Cannot mix-match CPU & GPU inputs/outputs.
|
||||||
|
// CPU-in & CPU-out <or> GPU-in & GPU-out
|
||||||
class RecolorCalculator : public CalculatorBase {
|
class RecolorCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
RecolorCalculator() = default;
|
RecolorCalculator() = default;
|
||||||
@@ -138,6 +141,11 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
|||||||
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
|
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Confirm only one of the input streams is present.
|
||||||
|
RET_CHECK(cc->Inputs().HasTag("IMAGE") ^ cc->Inputs().HasTag("IMAGE_GPU"));
|
||||||
|
// Confirm only one of the output streams is present.
|
||||||
|
RET_CHECK(cc->Outputs().HasTag("IMAGE") ^ cc->Outputs().HasTag("IMAGE_GPU"));
|
||||||
|
|
||||||
if (use_gpu) {
|
if (use_gpu) {
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||||
@@ -193,7 +201,62 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status RecolorCalculator::RenderCpu(CalculatorContext* cc) {
|
::mediapipe::Status RecolorCalculator::RenderCpu(CalculatorContext* cc) {
|
||||||
return ::mediapipe::UnimplementedError("CPU support is not implemented yet.");
|
if (cc->Inputs().Tag("MASK").IsEmpty()) {
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
// Get inputs and setup output.
|
||||||
|
const auto& input_img = cc->Inputs().Tag("IMAGE").Get<ImageFrame>();
|
||||||
|
const auto& mask_img = cc->Inputs().Tag("MASK").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("IMAGE").Add(output_img.release(), cc->InputTimestamp());
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status RecolorCalculator::RenderGpu(CalculatorContext* cc) {
|
::mediapipe::Status RecolorCalculator::RenderGpu(CalculatorContext* cc) {
|
||||||
@@ -303,9 +366,9 @@ void RecolorCalculator::GlRender() {
|
|||||||
|
|
||||||
if (!options.has_color()) RET_CHECK_FAIL() << "Missing color option.";
|
if (!options.has_color()) RET_CHECK_FAIL() << "Missing color option.";
|
||||||
|
|
||||||
color_.push_back(options.color().r() / 255.0);
|
color_.push_back(options.color().r());
|
||||||
color_.push_back(options.color().g() / 255.0);
|
color_.push_back(options.color().g());
|
||||||
color_.push_back(options.color().b() / 255.0);
|
color_.push_back(options.color().b());
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -378,8 +441,8 @@ void RecolorCalculator::GlRender() {
|
|||||||
glUseProgram(program_);
|
glUseProgram(program_);
|
||||||
glUniform1i(glGetUniformLocation(program_, "frame"), 1);
|
glUniform1i(glGetUniformLocation(program_, "frame"), 1);
|
||||||
glUniform1i(glGetUniformLocation(program_, "mask"), 2);
|
glUniform1i(glGetUniformLocation(program_, "mask"), 2);
|
||||||
glUniform3f(glGetUniformLocation(program_, "recolor"), color_[0], color_[1],
|
glUniform3f(glGetUniformLocation(program_, "recolor"), color_[0] / 255.0,
|
||||||
color_[2]);
|
color_[1] / 255.0, color_[2] / 255.0);
|
||||||
#endif // !MEDIAPIPE_DISABLE_GPU
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
|
|||||||
@@ -260,11 +260,11 @@ ScaleImageCalculator::~ScaleImageCalculator() {}
|
|||||||
&crop_width_, &crop_height_, //
|
&crop_width_, &crop_height_, //
|
||||||
&col_start_, &row_start_));
|
&col_start_, &row_start_));
|
||||||
MP_RETURN_IF_ERROR(
|
MP_RETURN_IF_ERROR(
|
||||||
scale_image::FindOutputDimensions(crop_width_, crop_height_, //
|
scale_image::FindOutputDimensions(crop_width_, crop_height_, //
|
||||||
options_.target_width(), //
|
options_.target_width(), //
|
||||||
options_.target_height(), //
|
options_.target_height(), //
|
||||||
options_.preserve_aspect_ratio(), //
|
options_.preserve_aspect_ratio(), //
|
||||||
options_.scale_to_multiple_of_two(), //
|
options_.scale_to_multiple_of(), //
|
||||||
&output_width_, &output_height_));
|
&output_width_, &output_height_));
|
||||||
MP_RETURN_IF_ERROR(FindInterpolationAlgorithm(options_.algorithm(),
|
MP_RETURN_IF_ERROR(FindInterpolationAlgorithm(options_.algorithm(),
|
||||||
&interpolation_algorithm_));
|
&interpolation_algorithm_));
|
||||||
@@ -361,17 +361,21 @@ ScaleImageCalculator::~ScaleImageCalculator() {}
|
|||||||
output_format_ = input_format_;
|
output_format_ = input_format_;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const bool is_positive_and_even =
|
||||||
|
(options_.scale_to_multiple_of() >= 1) &&
|
||||||
|
(options_.scale_to_multiple_of() % 2 == 0);
|
||||||
|
|
||||||
if (output_format_ == ImageFormat::YCBCR420P) {
|
if (output_format_ == ImageFormat::YCBCR420P) {
|
||||||
RET_CHECK(options_.scale_to_multiple_of_two())
|
RET_CHECK(is_positive_and_even)
|
||||||
<< "ScaleImageCalculator always outputs width and height that are "
|
<< "ScaleImageCalculator always outputs width and height that are "
|
||||||
"divisible by 2 when output format is YCbCr420P. To scale to "
|
"divisible by 2 when output format is YCbCr420P. To scale to "
|
||||||
"width and height of odd numbers, the output format must be SRGB.";
|
"width and height of odd numbers, the output format must be SRGB.";
|
||||||
} else if (options_.preserve_aspect_ratio()) {
|
} else if (options_.preserve_aspect_ratio()) {
|
||||||
RET_CHECK(options_.scale_to_multiple_of_two())
|
RET_CHECK(options_.scale_to_multiple_of() == 2)
|
||||||
<< "ScaleImageCalculator always outputs width and height that are "
|
<< "ScaleImageCalculator always outputs width and height that are "
|
||||||
"divisible by 2 when perserving aspect ratio. To scale to width "
|
"divisible by 2 when preserving aspect ratio. If you'd like to "
|
||||||
"and height of odd numbers, please set "
|
"set scale_to_multiple_of to something other than 2, please "
|
||||||
"preserve_aspect_ratio to false.";
|
"set preserve_aspect_ratio to false.";
|
||||||
}
|
}
|
||||||
|
|
||||||
if (input_width_ > 0 && input_height_ > 0 &&
|
if (input_width_ > 0 && input_height_ > 0 &&
|
||||||
@@ -474,13 +478,20 @@ ScaleImageCalculator::~ScaleImageCalculator() {}
|
|||||||
input_width_, "x", input_height_));
|
input_width_, "x", input_height_));
|
||||||
}
|
}
|
||||||
if (input_format_ != image_frame.Format()) {
|
if (input_format_ != image_frame.Format()) {
|
||||||
|
std::string image_frame_format_desc, input_format_desc;
|
||||||
|
#ifdef MEDIAPIPE_MOBILE
|
||||||
|
image_frame_format_desc = std::to_string(image_frame.Format());
|
||||||
|
input_format_desc = std::to_string(input_format_);
|
||||||
|
#else
|
||||||
const proto_ns::EnumDescriptor* desc = ImageFormat::Format_descriptor();
|
const proto_ns::EnumDescriptor* desc = ImageFormat::Format_descriptor();
|
||||||
|
image_frame_format_desc =
|
||||||
|
desc->FindValueByNumber(image_frame.Format())->DebugString();
|
||||||
|
input_format_desc = desc->FindValueByNumber(input_format_)->DebugString();
|
||||||
|
#endif // MEDIAPIPE_MOBILE
|
||||||
return tool::StatusFail(absl::StrCat(
|
return tool::StatusFail(absl::StrCat(
|
||||||
"If a header specifies a format, then image frames on "
|
"If a header specifies a format, then image frames on "
|
||||||
"the stream must have that format. Actual format ",
|
"the stream must have that format. Actual format ",
|
||||||
desc->FindValueByNumber(image_frame.Format())->DebugString(),
|
image_frame_format_desc, " but expected ", input_format_desc));
|
||||||
" but expected ",
|
|
||||||
desc->FindValueByNumber(input_format_)->DebugString()));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
|
|||||||
@@ -11,9 +11,10 @@ import "mediapipe/framework/formats/image_format.proto";
|
|||||||
// 2) Scale and convert the image to fit inside target_width x target_height
|
// 2) Scale and convert the image to fit inside target_width x target_height
|
||||||
// using the specified scaling algorithm. (maintaining the aspect
|
// using the specified scaling algorithm. (maintaining the aspect
|
||||||
// ratio if preserve_aspect_ratio is true).
|
// ratio if preserve_aspect_ratio is true).
|
||||||
// The output width and height will be divisible by 2. It is possible to output
|
// The output width and height will be divisible by 2, by default. It is
|
||||||
// width and height that are odd number when the output format is SRGB and not
|
// possible to output width and height that are odd numbers when the output
|
||||||
// perserving the aspect ratio. See scale_to_multiple_of_two option for details.
|
// format is SRGB and the aspect ratio is left unpreserved. See
|
||||||
|
// scale_to_multiple_of for details.
|
||||||
message ScaleImageCalculatorOptions {
|
message ScaleImageCalculatorOptions {
|
||||||
extend CalculatorOptions {
|
extend CalculatorOptions {
|
||||||
optional ScaleImageCalculatorOptions ext = 66237115;
|
optional ScaleImageCalculatorOptions ext = 66237115;
|
||||||
@@ -23,7 +24,7 @@ message ScaleImageCalculatorOptions {
|
|||||||
// depending on the other options below. If unset, use the same width
|
// depending on the other options below. If unset, use the same width
|
||||||
// or height as the input. If only one is set then determine the other
|
// or height as the input. If only one is set then determine the other
|
||||||
// from the aspect ratio (after cropping). The output width and height
|
// from the aspect ratio (after cropping). The output width and height
|
||||||
// will be divisible by 2.
|
// will be divisible by 2, by default.
|
||||||
optional int32 target_width = 1;
|
optional int32 target_width = 1;
|
||||||
optional int32 target_height = 2;
|
optional int32 target_height = 2;
|
||||||
|
|
||||||
@@ -31,7 +32,8 @@ message ScaleImageCalculatorOptions {
|
|||||||
// fits inside the box represented by target_width and target_height.
|
// fits inside the box represented by target_width and target_height.
|
||||||
// Otherwise it is scaled to fit target_width and target_height
|
// Otherwise it is scaled to fit target_width and target_height
|
||||||
// completely. In any case, the aspect ratio that is preserved is
|
// completely. In any case, the aspect ratio that is preserved is
|
||||||
// that after cropping to the minimum/maximum aspect ratio.
|
// that after cropping to the minimum/maximum aspect ratio. Additionally, if
|
||||||
|
// true, the output width and height will be divisible by 2.
|
||||||
optional bool preserve_aspect_ratio = 3 [default = true];
|
optional bool preserve_aspect_ratio = 3 [default = true];
|
||||||
|
|
||||||
// If ratio is positive, crop the image to this minimum and maximum
|
// If ratio is positive, crop the image to this minimum and maximum
|
||||||
@@ -95,11 +97,13 @@ message ScaleImageCalculatorOptions {
|
|||||||
// SRGB or YCBCR420P.
|
// SRGB or YCBCR420P.
|
||||||
optional ImageFormat.Format input_format = 12;
|
optional ImageFormat.Format input_format = 12;
|
||||||
|
|
||||||
// If true, the output width and height will be divisible by 2. Otherwise it
|
// If set to 2, the target width and height will be rounded-down
|
||||||
// will use the exact specified output width and height, which is only
|
// to the nearest even number. If set to any positive value other than 2,
|
||||||
// supported when the output format is SRGB and preserve_aspect_ratio option
|
// preserve_aspect_ratio must be false and the target width and height will be
|
||||||
// is set to false.
|
// rounded-down to multiples of the given value. If set to any value less than
|
||||||
optional bool scale_to_multiple_of_two = 13 [default = true];
|
// 1, it will be treated like 1.
|
||||||
|
// NOTE: If set to an odd number, the output format must be SRGB.
|
||||||
|
optional int32 scale_to_multiple_of = 13 [default = 2];
|
||||||
|
|
||||||
// If true, assume the input YUV is BT.709 (this is the HDTV standard, so most
|
// If true, assume the input YUV is BT.709 (this is the HDTV standard, so most
|
||||||
// content is likely using it). If false use the previous assumption of BT.601
|
// content is likely using it). If false use the previous assumption of BT.601
|
||||||
|
|||||||
@@ -88,17 +88,27 @@ double ParseRational(const std::string& rational) {
|
|||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status FindOutputDimensions(int input_width, //
|
::mediapipe::Status FindOutputDimensions(int input_width, //
|
||||||
int input_height, //
|
int input_height, //
|
||||||
int target_width, //
|
int target_width, //
|
||||||
int target_height, //
|
int target_height, //
|
||||||
bool preserve_aspect_ratio, //
|
bool preserve_aspect_ratio, //
|
||||||
bool scale_to_multiple_of_two, //
|
int scale_to_multiple_of, //
|
||||||
int* output_width,
|
int* output_width,
|
||||||
int* output_height) {
|
int* output_height) {
|
||||||
CHECK(output_width);
|
CHECK(output_width);
|
||||||
CHECK(output_height);
|
CHECK(output_height);
|
||||||
|
|
||||||
|
if (preserve_aspect_ratio) {
|
||||||
|
RET_CHECK(scale_to_multiple_of == 2)
|
||||||
|
<< "FindOutputDimensions always outputs width and height that are "
|
||||||
|
"divisible by 2 when preserving aspect ratio. If you'd like to "
|
||||||
|
"set scale_to_multiple_of to something other than 2, please "
|
||||||
|
"set preserve_aspect_ratio to false.";
|
||||||
|
}
|
||||||
|
|
||||||
|
if (scale_to_multiple_of < 1) scale_to_multiple_of = 1;
|
||||||
|
|
||||||
if (!preserve_aspect_ratio || (target_width <= 0 && target_height <= 0)) {
|
if (!preserve_aspect_ratio || (target_width <= 0 && target_height <= 0)) {
|
||||||
if (target_width <= 0) {
|
if (target_width <= 0) {
|
||||||
target_width = input_width;
|
target_width = input_width;
|
||||||
@@ -106,13 +116,13 @@ double ParseRational(const std::string& rational) {
|
|||||||
if (target_height <= 0) {
|
if (target_height <= 0) {
|
||||||
target_height = input_height;
|
target_height = input_height;
|
||||||
}
|
}
|
||||||
if (scale_to_multiple_of_two) {
|
|
||||||
*output_width = (target_width / 2) * 2;
|
target_width -= target_width % scale_to_multiple_of;
|
||||||
*output_height = (target_height / 2) * 2;
|
target_height -= target_height % scale_to_multiple_of;
|
||||||
} else {
|
|
||||||
*output_width = target_width;
|
*output_width = target_width;
|
||||||
*output_height = target_height;
|
*output_height = target_height;
|
||||||
}
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -35,17 +35,19 @@ namespace scale_image {
|
|||||||
int* col_start, int* row_start);
|
int* col_start, int* row_start);
|
||||||
|
|
||||||
// Given an input width and height, a target width and height, whether to
|
// Given an input width and height, a target width and height, whether to
|
||||||
// preserve the aspect ratio, and whether to round down to a multiple of 2,
|
// preserve the aspect ratio, and whether to round-down to the multiple of a
|
||||||
// determine the output width and height. If target_width or target_height is
|
// given number nearest to the targets, determine the output width and height.
|
||||||
// non-positive, then they will be set to the input_width and input_height
|
// If target_width or target_height is non-positive, then they will be set to
|
||||||
// respectively. The output_width and output_height will be reduced as necessary
|
// the input_width and input_height respectively. If scale_to_multiple_of is
|
||||||
// to preserve_aspect_ratio and to scale_to_multipe_of_two if these options are
|
// less than 1, it will be treated like 1. The output_width and
|
||||||
// specified.
|
// output_height will be reduced as necessary to preserve_aspect_ratio if the
|
||||||
|
// option is specified. If preserving the aspect ratio is desired, you must set
|
||||||
|
// scale_to_multiple_of to 2.
|
||||||
::mediapipe::Status FindOutputDimensions(int input_width, int input_height, //
|
::mediapipe::Status FindOutputDimensions(int input_width, int input_height, //
|
||||||
int target_width,
|
int target_width,
|
||||||
int target_height, //
|
int target_height, //
|
||||||
bool preserve_aspect_ratio, //
|
bool preserve_aspect_ratio, //
|
||||||
bool scale_to_multiple_of_two, //
|
int scale_to_multiple_of, //
|
||||||
int* output_width, int* output_height);
|
int* output_width, int* output_height);
|
||||||
|
|
||||||
} // namespace scale_image
|
} // namespace scale_image
|
||||||
|
|||||||
@@ -79,49 +79,49 @@ TEST(ScaleImageUtilsTest, FindOutputDimensionsPreserveRatio) {
|
|||||||
int output_width;
|
int output_width;
|
||||||
int output_height;
|
int output_height;
|
||||||
// Not scale.
|
// Not scale.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, -1, true, true, &output_width,
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, -1, true, 2, &output_width,
|
||||||
&output_height));
|
&output_height));
|
||||||
EXPECT_EQ(200, output_width);
|
EXPECT_EQ(200, output_width);
|
||||||
EXPECT_EQ(100, output_height);
|
EXPECT_EQ(100, output_height);
|
||||||
// Not scale with odd input size.
|
// Not scale with odd input size.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(201, 101, -1, -1, false, false,
|
MP_ASSERT_OK(FindOutputDimensions(201, 101, -1, -1, false, 1, &output_width,
|
||||||
&output_width, &output_height));
|
&output_height));
|
||||||
EXPECT_EQ(201, output_width);
|
EXPECT_EQ(201, output_width);
|
||||||
EXPECT_EQ(101, output_height);
|
EXPECT_EQ(101, output_height);
|
||||||
// Scale down by 1/2.
|
// Scale down by 1/2.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 100, -1, true, true,
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 100, -1, true, 2, &output_width,
|
||||||
&output_width, &output_height));
|
&output_height));
|
||||||
EXPECT_EQ(100, output_width);
|
EXPECT_EQ(100, output_width);
|
||||||
EXPECT_EQ(50, output_height);
|
EXPECT_EQ(50, output_height);
|
||||||
// Scale up, doubling dimensions.
|
// Scale up, doubling dimensions.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, 200, true, true,
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, 200, true, 2, &output_width,
|
||||||
&output_width, &output_height));
|
&output_height));
|
||||||
EXPECT_EQ(400, output_width);
|
EXPECT_EQ(400, output_width);
|
||||||
EXPECT_EQ(200, output_height);
|
EXPECT_EQ(200, output_height);
|
||||||
// Fits a 2:1 image into a 150 x 150 box. Output dimensions are always
|
// Fits a 2:1 image into a 150 x 150 box. Output dimensions are always
|
||||||
// visible by 2.
|
// visible by 2.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 150, 150, true, true,
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 150, 150, true, 2, &output_width,
|
||||||
&output_width, &output_height));
|
&output_height));
|
||||||
EXPECT_EQ(150, output_width);
|
EXPECT_EQ(150, output_width);
|
||||||
EXPECT_EQ(74, output_height);
|
EXPECT_EQ(74, output_height);
|
||||||
// Fits a 2:1 image into a 400 x 50 box.
|
// Fits a 2:1 image into a 400 x 50 box.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 400, 50, true, true,
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 400, 50, true, 2, &output_width,
|
||||||
&output_width, &output_height));
|
&output_height));
|
||||||
EXPECT_EQ(100, output_width);
|
EXPECT_EQ(100, output_width);
|
||||||
EXPECT_EQ(50, output_height);
|
EXPECT_EQ(50, output_height);
|
||||||
// Scale to multiple number with odd targe size.
|
// Scale to multiple number with odd targe size.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 101, -1, true, true,
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 101, -1, true, 2, &output_width,
|
||||||
&output_width, &output_height));
|
&output_height));
|
||||||
EXPECT_EQ(100, output_width);
|
EXPECT_EQ(100, output_width);
|
||||||
EXPECT_EQ(50, output_height);
|
EXPECT_EQ(50, output_height);
|
||||||
// Scale to multiple number with odd targe size.
|
// Scale to multiple number with odd targe size.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 101, -1, true, false,
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 101, -1, true, 2, &output_width,
|
||||||
&output_width, &output_height));
|
&output_height));
|
||||||
EXPECT_EQ(100, output_width);
|
EXPECT_EQ(100, output_width);
|
||||||
EXPECT_EQ(50, output_height);
|
EXPECT_EQ(50, output_height);
|
||||||
// Scale to odd size.
|
// Scale to odd size.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 151, 101, false, false,
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 151, 101, false, 1, &output_width,
|
||||||
&output_width, &output_height));
|
&output_height));
|
||||||
EXPECT_EQ(151, output_width);
|
EXPECT_EQ(151, output_width);
|
||||||
EXPECT_EQ(101, output_height);
|
EXPECT_EQ(101, output_height);
|
||||||
}
|
}
|
||||||
@@ -131,22 +131,62 @@ TEST(ScaleImageUtilsTest, FindOutputDimensionsNoAspectRatio) {
|
|||||||
int output_width;
|
int output_width;
|
||||||
int output_height;
|
int output_height;
|
||||||
// Scale width only.
|
// Scale width only.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 100, -1, false, true,
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 100, -1, false, 2, &output_width,
|
||||||
&output_width, &output_height));
|
&output_height));
|
||||||
EXPECT_EQ(100, output_width);
|
EXPECT_EQ(100, output_width);
|
||||||
EXPECT_EQ(100, output_height);
|
EXPECT_EQ(100, output_height);
|
||||||
// Scale height only.
|
// Scale height only.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, 200, false, true,
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, -1, 200, false, 2, &output_width,
|
||||||
&output_width, &output_height));
|
&output_height));
|
||||||
EXPECT_EQ(200, output_width);
|
EXPECT_EQ(200, output_width);
|
||||||
EXPECT_EQ(200, output_height);
|
EXPECT_EQ(200, output_height);
|
||||||
// Scale both dimensions.
|
// Scale both dimensions.
|
||||||
MP_ASSERT_OK(FindOutputDimensions(200, 100, 150, 200, false, true,
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 150, 200, false, 2, &output_width,
|
||||||
&output_width, &output_height));
|
&output_height));
|
||||||
EXPECT_EQ(150, output_width);
|
EXPECT_EQ(150, output_width);
|
||||||
EXPECT_EQ(200, output_height);
|
EXPECT_EQ(200, output_height);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Tests scale_to_multiple_of.
|
||||||
|
TEST(ScaleImageUtilsTest, FindOutputDimensionsDownScaleToMultipleOf) {
|
||||||
|
int output_width;
|
||||||
|
int output_height;
|
||||||
|
// Set no targets, downscale to a multiple of 8.
|
||||||
|
MP_ASSERT_OK(FindOutputDimensions(100, 100, -1, -1, false, 8, &output_width,
|
||||||
|
&output_height));
|
||||||
|
EXPECT_EQ(96, output_width);
|
||||||
|
EXPECT_EQ(96, output_height);
|
||||||
|
// Set width target, downscale to a multiple of 8.
|
||||||
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 100, -1, false, 8, &output_width,
|
||||||
|
&output_height));
|
||||||
|
EXPECT_EQ(96, output_width);
|
||||||
|
EXPECT_EQ(96, output_height);
|
||||||
|
// Set height target, downscale to a multiple of 8.
|
||||||
|
MP_ASSERT_OK(FindOutputDimensions(201, 101, -1, 201, false, 8, &output_width,
|
||||||
|
&output_height));
|
||||||
|
EXPECT_EQ(200, output_width);
|
||||||
|
EXPECT_EQ(200, output_height);
|
||||||
|
// Set both targets, downscale to a multiple of 8.
|
||||||
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 150, 200, false, 8, &output_width,
|
||||||
|
&output_height));
|
||||||
|
EXPECT_EQ(144, output_width);
|
||||||
|
EXPECT_EQ(200, output_height);
|
||||||
|
// Doesn't throw error if keep aspect is true and downscale multiple is 2.
|
||||||
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 400, 200, true, 2, &output_width,
|
||||||
|
&output_height));
|
||||||
|
EXPECT_EQ(400, output_width);
|
||||||
|
EXPECT_EQ(200, output_height);
|
||||||
|
// Throws error if keep aspect is true, but downscale multiple is not 2.
|
||||||
|
ASSERT_THAT(FindOutputDimensions(200, 100, 400, 200, true, 4, &output_width,
|
||||||
|
&output_height),
|
||||||
|
testing::Not(testing::status::IsOk()));
|
||||||
|
// Downscaling to multiple ignored if multiple is less than 2.
|
||||||
|
MP_ASSERT_OK(FindOutputDimensions(200, 100, 401, 201, false, 1, &output_width,
|
||||||
|
&output_height));
|
||||||
|
EXPECT_EQ(401, output_width);
|
||||||
|
EXPECT_EQ(201, output_height);
|
||||||
|
}
|
||||||
|
|
||||||
} // namespace
|
} // namespace
|
||||||
} // namespace scale_image
|
} // namespace scale_image
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
+1
@@ -21,6 +21,7 @@ filegroup(
|
|||||||
"dino.jpg",
|
"dino.jpg",
|
||||||
"dino_quality_50.jpg",
|
"dino_quality_50.jpg",
|
||||||
"dino_quality_80.jpg",
|
"dino_quality_80.jpg",
|
||||||
|
"front_camera_pixel2.jpg",
|
||||||
],
|
],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
)
|
)
|
||||||
|
|||||||
Binary file not shown.
|
After Width: | Height: | Size: 6.3 MiB |
@@ -138,7 +138,7 @@ mediapipe_cc_proto_library(
|
|||||||
srcs = ["image_frame_to_tensor_calculator.proto"],
|
srcs = ["image_frame_to_tensor_calculator.proto"],
|
||||||
cc_deps = [
|
cc_deps = [
|
||||||
"//mediapipe/framework:calculator_cc_proto",
|
"//mediapipe/framework:calculator_cc_proto",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [":image_frame_to_tensor_calculator_proto"],
|
deps = [":image_frame_to_tensor_calculator_proto"],
|
||||||
@@ -173,7 +173,7 @@ mediapipe_cc_proto_library(
|
|||||||
srcs = ["pack_media_sequence_calculator.proto"],
|
srcs = ["pack_media_sequence_calculator.proto"],
|
||||||
cc_deps = [
|
cc_deps = [
|
||||||
"//mediapipe/framework:calculator_cc_proto",
|
"//mediapipe/framework:calculator_cc_proto",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [":pack_media_sequence_calculator_proto"],
|
deps = [":pack_media_sequence_calculator_proto"],
|
||||||
@@ -192,7 +192,7 @@ mediapipe_cc_proto_library(
|
|||||||
srcs = ["tensorflow_session_from_frozen_graph_generator.proto"],
|
srcs = ["tensorflow_session_from_frozen_graph_generator.proto"],
|
||||||
cc_deps = [
|
cc_deps = [
|
||||||
"//mediapipe/framework:packet_generator_cc_proto",
|
"//mediapipe/framework:packet_generator_cc_proto",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [":tensorflow_session_from_frozen_graph_generator_proto"],
|
deps = [":tensorflow_session_from_frozen_graph_generator_proto"],
|
||||||
@@ -203,7 +203,7 @@ mediapipe_cc_proto_library(
|
|||||||
srcs = ["tensorflow_session_from_frozen_graph_calculator.proto"],
|
srcs = ["tensorflow_session_from_frozen_graph_calculator.proto"],
|
||||||
cc_deps = [
|
cc_deps = [
|
||||||
"//mediapipe/framework:calculator_cc_proto",
|
"//mediapipe/framework:calculator_cc_proto",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [":tensorflow_session_from_frozen_graph_calculator_proto"],
|
deps = [":tensorflow_session_from_frozen_graph_calculator_proto"],
|
||||||
@@ -277,7 +277,7 @@ mediapipe_cc_proto_library(
|
|||||||
srcs = ["vector_int_to_tensor_calculator_options.proto"],
|
srcs = ["vector_int_to_tensor_calculator_options.proto"],
|
||||||
cc_deps = [
|
cc_deps = [
|
||||||
"//mediapipe/framework:calculator_cc_proto",
|
"//mediapipe/framework:calculator_cc_proto",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [":vector_int_to_tensor_calculator_options_proto"],
|
deps = [":vector_int_to_tensor_calculator_options_proto"],
|
||||||
@@ -408,7 +408,7 @@ cc_library(
|
|||||||
"//mediapipe/util/sequence:media_sequence",
|
"//mediapipe/util/sequence:media_sequence",
|
||||||
"//mediapipe/util/sequence:media_sequence_util",
|
"//mediapipe/util/sequence:media_sequence_util",
|
||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -423,7 +423,7 @@ cc_library(
|
|||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -654,7 +654,7 @@ cc_library(
|
|||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
"@org_tensorflow//tensorflow/core:lib",
|
"@org_tensorflow//tensorflow/core:lib",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -695,7 +695,7 @@ cc_library(
|
|||||||
"//mediapipe/util:audio_decoder_cc_proto",
|
"//mediapipe/util:audio_decoder_cc_proto",
|
||||||
"//mediapipe/util/sequence:media_sequence",
|
"//mediapipe/util/sequence:media_sequence",
|
||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -737,7 +737,7 @@ cc_library(
|
|||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:packet",
|
"//mediapipe/framework:packet",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -745,6 +745,7 @@ cc_library(
|
|||||||
cc_test(
|
cc_test(
|
||||||
name = "graph_tensors_packet_generator_test",
|
name = "graph_tensors_packet_generator_test",
|
||||||
srcs = ["graph_tensors_packet_generator_test.cc"],
|
srcs = ["graph_tensors_packet_generator_test.cc"],
|
||||||
|
linkstatic = 1,
|
||||||
deps = [
|
deps = [
|
||||||
":graph_tensors_packet_generator",
|
":graph_tensors_packet_generator",
|
||||||
":graph_tensors_packet_generator_cc_proto",
|
":graph_tensors_packet_generator_cc_proto",
|
||||||
@@ -761,6 +762,7 @@ cc_test(
|
|||||||
name = "image_frame_to_tensor_calculator_test",
|
name = "image_frame_to_tensor_calculator_test",
|
||||||
size = "small",
|
size = "small",
|
||||||
srcs = ["image_frame_to_tensor_calculator_test.cc"],
|
srcs = ["image_frame_to_tensor_calculator_test.cc"],
|
||||||
|
linkstatic = 1,
|
||||||
deps = [
|
deps = [
|
||||||
":image_frame_to_tensor_calculator",
|
":image_frame_to_tensor_calculator",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
@@ -777,6 +779,7 @@ cc_test(
|
|||||||
name = "matrix_to_tensor_calculator_test",
|
name = "matrix_to_tensor_calculator_test",
|
||||||
size = "small",
|
size = "small",
|
||||||
srcs = ["matrix_to_tensor_calculator_test.cc"],
|
srcs = ["matrix_to_tensor_calculator_test.cc"],
|
||||||
|
linkstatic = 1,
|
||||||
deps = [
|
deps = [
|
||||||
":matrix_to_tensor_calculator",
|
":matrix_to_tensor_calculator",
|
||||||
":matrix_to_tensor_calculator_options_cc_proto",
|
":matrix_to_tensor_calculator_options_cc_proto",
|
||||||
@@ -793,6 +796,7 @@ cc_test(
|
|||||||
name = "lapped_tensor_buffer_calculator_test",
|
name = "lapped_tensor_buffer_calculator_test",
|
||||||
size = "small",
|
size = "small",
|
||||||
srcs = ["lapped_tensor_buffer_calculator_test.cc"],
|
srcs = ["lapped_tensor_buffer_calculator_test.cc"],
|
||||||
|
linkstatic = 1,
|
||||||
deps = [
|
deps = [
|
||||||
":lapped_tensor_buffer_calculator",
|
":lapped_tensor_buffer_calculator",
|
||||||
":lapped_tensor_buffer_calculator_cc_proto",
|
":lapped_tensor_buffer_calculator_cc_proto",
|
||||||
@@ -801,7 +805,7 @@ cc_test(
|
|||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"@com_google_absl//absl/memory",
|
"@com_google_absl//absl/memory",
|
||||||
"@org_tensorflow//tensorflow/core:framework",
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -840,7 +844,7 @@ cc_test(
|
|||||||
"//mediapipe/util/sequence:media_sequence",
|
"//mediapipe/util/sequence:media_sequence",
|
||||||
"@com_google_absl//absl/memory",
|
"@com_google_absl//absl/memory",
|
||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -867,7 +871,7 @@ cc_test(
|
|||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
"@org_tensorflow//tensorflow/core:direct_session",
|
"@org_tensorflow//tensorflow/core:direct_session",
|
||||||
"@org_tensorflow//tensorflow/core:framework",
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
"@org_tensorflow//tensorflow/core:testlib",
|
"@org_tensorflow//tensorflow/core:testlib",
|
||||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
||||||
"@org_tensorflow//tensorflow/core/kernels:math",
|
"@org_tensorflow//tensorflow/core/kernels:math",
|
||||||
@@ -897,7 +901,7 @@ cc_test(
|
|||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
"@org_tensorflow//tensorflow/core:direct_session",
|
"@org_tensorflow//tensorflow/core:direct_session",
|
||||||
"@org_tensorflow//tensorflow/core:framework",
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
"@org_tensorflow//tensorflow/core:testlib",
|
"@org_tensorflow//tensorflow/core:testlib",
|
||||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
||||||
"@org_tensorflow//tensorflow/core/kernels:math",
|
"@org_tensorflow//tensorflow/core/kernels:math",
|
||||||
@@ -956,6 +960,7 @@ cc_test(
|
|||||||
cc_test(
|
cc_test(
|
||||||
name = "tensor_squeeze_dimensions_calculator_test",
|
name = "tensor_squeeze_dimensions_calculator_test",
|
||||||
srcs = ["tensor_squeeze_dimensions_calculator_test.cc"],
|
srcs = ["tensor_squeeze_dimensions_calculator_test.cc"],
|
||||||
|
linkstatic = 1,
|
||||||
deps = [
|
deps = [
|
||||||
":tensor_squeeze_dimensions_calculator",
|
":tensor_squeeze_dimensions_calculator",
|
||||||
":tensor_squeeze_dimensions_calculator_cc_proto",
|
":tensor_squeeze_dimensions_calculator_cc_proto",
|
||||||
@@ -963,7 +968,7 @@ cc_test(
|
|||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"@org_tensorflow//tensorflow/core:framework",
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -971,6 +976,7 @@ cc_test(
|
|||||||
name = "tensor_to_image_frame_calculator_test",
|
name = "tensor_to_image_frame_calculator_test",
|
||||||
size = "small",
|
size = "small",
|
||||||
srcs = ["tensor_to_image_frame_calculator_test.cc"],
|
srcs = ["tensor_to_image_frame_calculator_test.cc"],
|
||||||
|
linkstatic = 1,
|
||||||
deps = [
|
deps = [
|
||||||
":tensor_to_image_frame_calculator",
|
":tensor_to_image_frame_calculator",
|
||||||
":tensor_to_image_frame_calculator_cc_proto",
|
":tensor_to_image_frame_calculator_cc_proto",
|
||||||
@@ -979,7 +985,7 @@ cc_test(
|
|||||||
"//mediapipe/framework/formats:image_frame",
|
"//mediapipe/framework/formats:image_frame",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"@org_tensorflow//tensorflow/core:framework",
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -987,6 +993,7 @@ cc_test(
|
|||||||
name = "tensor_to_matrix_calculator_test",
|
name = "tensor_to_matrix_calculator_test",
|
||||||
size = "small",
|
size = "small",
|
||||||
srcs = ["tensor_to_matrix_calculator_test.cc"],
|
srcs = ["tensor_to_matrix_calculator_test.cc"],
|
||||||
|
linkstatic = 1,
|
||||||
deps = [
|
deps = [
|
||||||
":tensor_to_matrix_calculator",
|
":tensor_to_matrix_calculator",
|
||||||
":tensor_to_matrix_calculator_cc_proto",
|
":tensor_to_matrix_calculator_cc_proto",
|
||||||
@@ -996,13 +1003,14 @@ cc_test(
|
|||||||
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"@org_tensorflow//tensorflow/core:framework",
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
cc_test(
|
cc_test(
|
||||||
name = "tensor_to_vector_float_calculator_test",
|
name = "tensor_to_vector_float_calculator_test",
|
||||||
srcs = ["tensor_to_vector_float_calculator_test.cc"],
|
srcs = ["tensor_to_vector_float_calculator_test.cc"],
|
||||||
|
linkstatic = 1,
|
||||||
deps = [
|
deps = [
|
||||||
":tensor_to_vector_float_calculator",
|
":tensor_to_vector_float_calculator",
|
||||||
":tensor_to_vector_float_calculator_options_cc_proto",
|
":tensor_to_vector_float_calculator_options_cc_proto",
|
||||||
@@ -1010,7 +1018,7 @@ cc_test(
|
|||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"@org_tensorflow//tensorflow/core:framework",
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1030,13 +1038,14 @@ cc_test(
|
|||||||
"//mediapipe/util/sequence:media_sequence",
|
"//mediapipe/util/sequence:media_sequence",
|
||||||
"@com_google_absl//absl/memory",
|
"@com_google_absl//absl/memory",
|
||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
cc_test(
|
cc_test(
|
||||||
name = "vector_int_to_tensor_calculator_test",
|
name = "vector_int_to_tensor_calculator_test",
|
||||||
srcs = ["vector_int_to_tensor_calculator_test.cc"],
|
srcs = ["vector_int_to_tensor_calculator_test.cc"],
|
||||||
|
linkstatic = 1,
|
||||||
deps = [
|
deps = [
|
||||||
":vector_int_to_tensor_calculator",
|
":vector_int_to_tensor_calculator",
|
||||||
":vector_int_to_tensor_calculator_options_cc_proto",
|
":vector_int_to_tensor_calculator_options_cc_proto",
|
||||||
@@ -1044,13 +1053,14 @@ cc_test(
|
|||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"@org_tensorflow//tensorflow/core:framework",
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
cc_test(
|
cc_test(
|
||||||
name = "vector_float_to_tensor_calculator_test",
|
name = "vector_float_to_tensor_calculator_test",
|
||||||
srcs = ["vector_float_to_tensor_calculator_test.cc"],
|
srcs = ["vector_float_to_tensor_calculator_test.cc"],
|
||||||
|
linkstatic = 1,
|
||||||
deps = [
|
deps = [
|
||||||
":vector_float_to_tensor_calculator",
|
":vector_float_to_tensor_calculator",
|
||||||
":vector_float_to_tensor_calculator_options_cc_proto",
|
":vector_float_to_tensor_calculator_options_cc_proto",
|
||||||
@@ -1058,7 +1068,7 @@ cc_test(
|
|||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"@org_tensorflow//tensorflow/core:framework",
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
"@org_tensorflow//tensorflow/core:protos_all",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1100,6 +1110,7 @@ cc_test(
|
|||||||
],
|
],
|
||||||
"//mediapipe:android": [
|
"//mediapipe:android": [
|
||||||
"@org_tensorflow//tensorflow/core:android_tensorflow_lib_with_ops_lite_proto_no_rtti_lib",
|
"@org_tensorflow//tensorflow/core:android_tensorflow_lib_with_ops_lite_proto_no_rtti_lib",
|
||||||
|
"@org_tensorflow//tensorflow/core:android_tensorflow_test_lib",
|
||||||
],
|
],
|
||||||
"//mediapipe:ios": [
|
"//mediapipe:ios": [
|
||||||
"@org_tensorflow//tensorflow/core:ios_tensorflow_test_lib",
|
"@org_tensorflow//tensorflow/core:ios_tensorflow_test_lib",
|
||||||
|
|||||||
@@ -34,6 +34,7 @@ namespace mediapipe {
|
|||||||
|
|
||||||
const char kSequenceExampleTag[] = "SEQUENCE_EXAMPLE";
|
const char kSequenceExampleTag[] = "SEQUENCE_EXAMPLE";
|
||||||
const char kImageTag[] = "IMAGE";
|
const char kImageTag[] = "IMAGE";
|
||||||
|
const char kFloatContextFeaturePrefixTag[] = "FLOAT_CONTEXT_FEATURE_";
|
||||||
const char kFloatFeaturePrefixTag[] = "FLOAT_FEATURE_";
|
const char kFloatFeaturePrefixTag[] = "FLOAT_FEATURE_";
|
||||||
const char kForwardFlowEncodedTag[] = "FORWARD_FLOW_ENCODED";
|
const char kForwardFlowEncodedTag[] = "FORWARD_FLOW_ENCODED";
|
||||||
const char kBBoxTag[] = "BBOX";
|
const char kBBoxTag[] = "BBOX";
|
||||||
@@ -145,6 +146,9 @@ class PackMediaSequenceCalculator : public CalculatorBase {
|
|||||||
}
|
}
|
||||||
cc->Inputs().Tag(tag).Set<std::vector<Detection>>();
|
cc->Inputs().Tag(tag).Set<std::vector<Detection>>();
|
||||||
}
|
}
|
||||||
|
if (absl::StartsWith(tag, kFloatContextFeaturePrefixTag)) {
|
||||||
|
cc->Inputs().Tag(tag).Set<std::vector<float>>();
|
||||||
|
}
|
||||||
if (absl::StartsWith(tag, kFloatFeaturePrefixTag)) {
|
if (absl::StartsWith(tag, kFloatFeaturePrefixTag)) {
|
||||||
cc->Inputs().Tag(tag).Set<std::vector<float>>();
|
cc->Inputs().Tag(tag).Set<std::vector<float>>();
|
||||||
}
|
}
|
||||||
@@ -264,7 +268,7 @@ class PackMediaSequenceCalculator : public CalculatorBase {
|
|||||||
if (options.output_only_if_all_present()) {
|
if (options.output_only_if_all_present()) {
|
||||||
::mediapipe::Status status = VerifySequence();
|
::mediapipe::Status status = VerifySequence();
|
||||||
if (!status.ok()) {
|
if (!status.ok()) {
|
||||||
cc->GetCounter(status.error_message())->Increment();
|
cc->GetCounter(status.ToString())->Increment();
|
||||||
return status;
|
return status;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -344,6 +348,17 @@ class PackMediaSequenceCalculator : public CalculatorBase {
|
|||||||
sequence_.get());
|
sequence_.get());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if (absl::StartsWith(tag, kFloatContextFeaturePrefixTag) &&
|
||||||
|
!cc->Inputs().Tag(tag).IsEmpty()) {
|
||||||
|
std::string key =
|
||||||
|
tag.substr(sizeof(kFloatContextFeaturePrefixTag) /
|
||||||
|
sizeof(*kFloatContextFeaturePrefixTag) -
|
||||||
|
1);
|
||||||
|
RET_CHECK_EQ(cc->InputTimestamp(), Timestamp::PostStream());
|
||||||
|
mpms::SetContextFeatureFloats(
|
||||||
|
key, cc->Inputs().Tag(tag).Get<std::vector<float>>(),
|
||||||
|
sequence_.get());
|
||||||
|
}
|
||||||
if (absl::StartsWith(tag, kFloatFeaturePrefixTag) &&
|
if (absl::StartsWith(tag, kFloatFeaturePrefixTag) &&
|
||||||
!cc->Inputs().Tag(tag).IsEmpty()) {
|
!cc->Inputs().Tag(tag).IsEmpty()) {
|
||||||
std::string key = tag.substr(sizeof(kFloatFeaturePrefixTag) /
|
std::string key = tag.substr(sizeof(kFloatFeaturePrefixTag) /
|
||||||
|
|||||||
@@ -194,6 +194,38 @@ TEST_F(PackMediaSequenceCalculatorTest, PacksTwoFloatLists) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
TEST_F(PackMediaSequenceCalculatorTest, PacksTwoContextFloatLists) {
|
||||||
|
SetUpCalculator(
|
||||||
|
{"FLOAT_CONTEXT_FEATURE_TEST:test", "FLOAT_CONTEXT_FEATURE_OTHER:test2"},
|
||||||
|
{}, false, true);
|
||||||
|
auto input_sequence = absl::make_unique<tf::SequenceExample>();
|
||||||
|
|
||||||
|
auto vf_ptr = absl::make_unique<std::vector<float>>(2, 3);
|
||||||
|
runner_->MutableInputs()
|
||||||
|
->Tag("FLOAT_CONTEXT_FEATURE_TEST")
|
||||||
|
.packets.push_back(Adopt(vf_ptr.release()).At(Timestamp::PostStream()));
|
||||||
|
vf_ptr = absl::make_unique<std::vector<float>>(2, 4);
|
||||||
|
runner_->MutableInputs()
|
||||||
|
->Tag("FLOAT_CONTEXT_FEATURE_OTHER")
|
||||||
|
.packets.push_back(Adopt(vf_ptr.release()).At(Timestamp::PostStream()));
|
||||||
|
|
||||||
|
runner_->MutableSidePackets()->Tag("SEQUENCE_EXAMPLE") =
|
||||||
|
Adopt(input_sequence.release());
|
||||||
|
|
||||||
|
MP_ASSERT_OK(runner_->Run());
|
||||||
|
|
||||||
|
const std::vector<Packet>& output_packets =
|
||||||
|
runner_->Outputs().Tag("SEQUENCE_EXAMPLE").packets;
|
||||||
|
ASSERT_EQ(1, output_packets.size());
|
||||||
|
const tf::SequenceExample& output_sequence =
|
||||||
|
output_packets[0].Get<tf::SequenceExample>();
|
||||||
|
|
||||||
|
ASSERT_THAT(mpms::GetContextFeatureFloats("TEST", output_sequence),
|
||||||
|
testing::ElementsAre(3, 3));
|
||||||
|
ASSERT_THAT(mpms::GetContextFeatureFloats("OTHER", output_sequence),
|
||||||
|
testing::ElementsAre(4, 4));
|
||||||
|
}
|
||||||
|
|
||||||
TEST_F(PackMediaSequenceCalculatorTest, PacksAdditionalContext) {
|
TEST_F(PackMediaSequenceCalculatorTest, PacksAdditionalContext) {
|
||||||
tf::Features context;
|
tf::Features context;
|
||||||
(*context.mutable_feature())["TEST"].mutable_bytes_list()->add_value("YES");
|
(*context.mutable_feature())["TEST"].mutable_bytes_list()->add_value("YES");
|
||||||
|
|||||||
@@ -454,7 +454,7 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
|
|||||||
|
|
||||||
// RET_CHECK on the tf::Status object itself in order to print an
|
// RET_CHECK on the tf::Status object itself in order to print an
|
||||||
// informative error message.
|
// informative error message.
|
||||||
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.error_message();
|
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.ToString();
|
||||||
|
|
||||||
const int64 run_end_time = absl::ToUnixMicros(clock_->TimeNow());
|
const int64 run_end_time = absl::ToUnixMicros(clock_->TimeNow());
|
||||||
cc->GetCounter(kTotalSessionRunsTimeUsecsCounterSuffix)
|
cc->GetCounter(kTotalSessionRunsTimeUsecsCounterSuffix)
|
||||||
|
|||||||
@@ -109,7 +109,7 @@ class TensorFlowSessionFromFrozenGraphCalculator : public CalculatorBase {
|
|||||||
|
|
||||||
RET_CHECK(graph_def.ParseFromString(graph_def_serialized));
|
RET_CHECK(graph_def.ParseFromString(graph_def_serialized));
|
||||||
const tf::Status tf_status = session->session->Create(graph_def);
|
const tf::Status tf_status = session->session->Create(graph_def);
|
||||||
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.error_message();
|
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.ToString();
|
||||||
|
|
||||||
for (const auto& key_value : options.tag_to_tensor_names()) {
|
for (const auto& key_value : options.tag_to_tensor_names()) {
|
||||||
session->tag_to_tensor_map[key_value.first] = key_value.second;
|
session->tag_to_tensor_map[key_value.first] = key_value.second;
|
||||||
@@ -119,7 +119,7 @@ class TensorFlowSessionFromFrozenGraphCalculator : public CalculatorBase {
|
|||||||
session->session->Run({}, {}, initialization_op_names, {});
|
session->session->Run({}, {}, initialization_op_names, {});
|
||||||
// RET_CHECK on the tf::Status object itself in order to print an
|
// RET_CHECK on the tf::Status object itself in order to print an
|
||||||
// informative error message.
|
// informative error message.
|
||||||
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.error_message();
|
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.ToString();
|
||||||
}
|
}
|
||||||
|
|
||||||
cc->OutputSidePackets().Tag("SESSION").Set(Adopt(session.release()));
|
cc->OutputSidePackets().Tag("SESSION").Set(Adopt(session.release()));
|
||||||
|
|||||||
@@ -109,7 +109,7 @@ class TensorFlowSessionFromFrozenGraphGenerator : public PacketGenerator {
|
|||||||
|
|
||||||
RET_CHECK(graph_def.ParseFromString(graph_def_serialized));
|
RET_CHECK(graph_def.ParseFromString(graph_def_serialized));
|
||||||
const tf::Status tf_status = session->session->Create(graph_def);
|
const tf::Status tf_status = session->session->Create(graph_def);
|
||||||
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.error_message();
|
RET_CHECK(tf_status.ok()) << "Create failed: " << tf_status.ToString();
|
||||||
|
|
||||||
for (const auto& key_value : options.tag_to_tensor_names()) {
|
for (const auto& key_value : options.tag_to_tensor_names()) {
|
||||||
session->tag_to_tensor_map[key_value.first] = key_value.second;
|
session->tag_to_tensor_map[key_value.first] = key_value.second;
|
||||||
@@ -119,7 +119,7 @@ class TensorFlowSessionFromFrozenGraphGenerator : public PacketGenerator {
|
|||||||
session->session->Run({}, {}, initialization_op_names, {});
|
session->session->Run({}, {}, initialization_op_names, {});
|
||||||
// RET_CHECK on the tf::Status object itself in order to print an
|
// RET_CHECK on the tf::Status object itself in order to print an
|
||||||
// informative error message.
|
// informative error message.
|
||||||
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.error_message();
|
RET_CHECK(tf_status.ok()) << "Run failed: " << tf_status.ToString();
|
||||||
}
|
}
|
||||||
|
|
||||||
output_side_packets->Tag("SESSION") = Adopt(session.release());
|
output_side_packets->Tag("SESSION") = Adopt(session.release());
|
||||||
|
|||||||
@@ -17,7 +17,7 @@
|
|||||||
#if !defined(__ANDROID__)
|
#if !defined(__ANDROID__)
|
||||||
#include "mediapipe/framework/port/file_helpers.h"
|
#include "mediapipe/framework/port/file_helpers.h"
|
||||||
#endif
|
#endif
|
||||||
#include "absl/strings/substitute.h"
|
#include "absl/strings/str_replace.h"
|
||||||
#include "mediapipe/calculators/tensorflow/tensorflow_session.h"
|
#include "mediapipe/calculators/tensorflow/tensorflow_session.h"
|
||||||
#include "mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_calculator.pb.h"
|
#include "mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_calculator.pb.h"
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
@@ -51,10 +51,11 @@ static constexpr char kStringSavedModelPath[] = "STRING_SAVED_MODEL_PATH";
|
|||||||
#endif
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
// If options.convert_signature_to_tags() will convert letters to uppercase
|
// If options.convert_signature_to_tags() is set, will convert letters to
|
||||||
// and replace /'s with _'s. If set, this enables the standard SavedModel
|
// uppercase and replace /'s and -'s with _'s. This enables the standard
|
||||||
// classification, regression, and prediction signatures to be used as
|
// SavedModel classification, regression, and prediction signatures to be used
|
||||||
// uppercase INPUTS and OUTPUTS tags for streams.
|
// as uppercase INPUTS and OUTPUTS tags for streams and supports other common
|
||||||
|
// patterns.
|
||||||
const std::string MaybeConvertSignatureToTag(
|
const std::string MaybeConvertSignatureToTag(
|
||||||
const std::string& name,
|
const std::string& name,
|
||||||
const TensorFlowSessionFromSavedModelCalculatorOptions& options) {
|
const TensorFlowSessionFromSavedModelCalculatorOptions& options) {
|
||||||
@@ -63,7 +64,8 @@ const std::string MaybeConvertSignatureToTag(
|
|||||||
output.resize(name.length());
|
output.resize(name.length());
|
||||||
std::transform(name.begin(), name.end(), output.begin(),
|
std::transform(name.begin(), name.end(), output.begin(),
|
||||||
[](unsigned char c) { return std::toupper(c); });
|
[](unsigned char c) { return std::toupper(c); });
|
||||||
output = absl::Substitute(output, "/", "_");
|
output = absl::StrReplaceAll(output, {{"/", "_"}});
|
||||||
|
output = absl::StrReplaceAll(output, {{"-", "_"}});
|
||||||
return output;
|
return output;
|
||||||
} else {
|
} else {
|
||||||
return name;
|
return name;
|
||||||
@@ -140,7 +142,7 @@ class TensorFlowSessionFromSavedModelCalculator : public CalculatorBase {
|
|||||||
if (!status.ok()) {
|
if (!status.ok()) {
|
||||||
return ::mediapipe::Status(
|
return ::mediapipe::Status(
|
||||||
static_cast<::mediapipe::StatusCode>(status.code()),
|
static_cast<::mediapipe::StatusCode>(status.code()),
|
||||||
status.error_message());
|
status.ToString());
|
||||||
}
|
}
|
||||||
|
|
||||||
auto session = absl::make_unique<TensorFlowSession>();
|
auto session = absl::make_unique<TensorFlowSession>();
|
||||||
|
|||||||
+2
-2
@@ -32,8 +32,8 @@ message TensorFlowSessionFromSavedModelCalculatorOptions {
|
|||||||
// The name of the generic signature to load into the mapping from tags to
|
// The name of the generic signature to load into the mapping from tags to
|
||||||
// tensor names.
|
// tensor names.
|
||||||
optional string signature_name = 2 [default = "serving_default"];
|
optional string signature_name = 2 [default = "serving_default"];
|
||||||
// Whether to convert the signature keys to uppercase and switch /'s to
|
// Whether to convert the signature keys to uppercase as well as switch /'s
|
||||||
// _'s, which enables standard signatures to be used as Tags.
|
// and -'s to _'s, which enables common signatures to be used as Tags.
|
||||||
optional bool convert_signature_to_tags = 3 [default = true];
|
optional bool convert_signature_to_tags = 3 [default = true];
|
||||||
// If true, saved_model_path can have multiple exported models in
|
// If true, saved_model_path can have multiple exported models in
|
||||||
// subdirectories saved_model_path/%08d and the alphabetically last (i.e.,
|
// subdirectories saved_model_path/%08d and the alphabetically last (i.e.,
|
||||||
|
|||||||
+1
-1
@@ -12,7 +12,7 @@
|
|||||||
// See the License for the specific language governing permissions and
|
// See the License for the specific language governing permissions and
|
||||||
// limitations under the License.
|
// limitations under the License.
|
||||||
|
|
||||||
#include "absl/strings/substitute.h"
|
#include "absl/strings/str_replace.h"
|
||||||
#include "mediapipe/calculators/tensorflow/tensorflow_session.h"
|
#include "mediapipe/calculators/tensorflow/tensorflow_session.h"
|
||||||
#include "mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_calculator.pb.h"
|
#include "mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_calculator.pb.h"
|
||||||
#include "mediapipe/framework/calculator.pb.h"
|
#include "mediapipe/framework/calculator.pb.h"
|
||||||
|
|||||||
@@ -17,7 +17,7 @@
|
|||||||
#if !defined(__ANDROID__)
|
#if !defined(__ANDROID__)
|
||||||
#include "mediapipe/framework/port/file_helpers.h"
|
#include "mediapipe/framework/port/file_helpers.h"
|
||||||
#endif
|
#endif
|
||||||
#include "absl/strings/substitute.h"
|
#include "absl/strings/str_replace.h"
|
||||||
#include "mediapipe/calculators/tensorflow/tensorflow_session.h"
|
#include "mediapipe/calculators/tensorflow/tensorflow_session.h"
|
||||||
#include "mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_generator.pb.h"
|
#include "mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_generator.pb.h"
|
||||||
#include "mediapipe/framework/deps/file_path.h"
|
#include "mediapipe/framework/deps/file_path.h"
|
||||||
@@ -53,10 +53,11 @@ static constexpr char kStringSavedModelPath[] = "STRING_SAVED_MODEL_PATH";
|
|||||||
#endif
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
// If options.convert_signature_to_tags() will convert letters to uppercase
|
// If options.convert_signature_to_tags() is set, will convert letters to
|
||||||
// and replace /'s with _'s. If set, this enables the standard SavedModel
|
// uppercase and replace /'s and -'s with _'s. This enables the standard
|
||||||
// classification, regression, and prediction signatures to be used as
|
// SavedModel classification, regression, and prediction signatures to be used
|
||||||
// uppercase INPUTS and OUTPUTS tags for streams.
|
// as uppercase INPUTS and OUTPUTS tags for streams and supports other common
|
||||||
|
// patterns.
|
||||||
const std::string MaybeConvertSignatureToTag(
|
const std::string MaybeConvertSignatureToTag(
|
||||||
const std::string& name,
|
const std::string& name,
|
||||||
const TensorFlowSessionFromSavedModelGeneratorOptions& options) {
|
const TensorFlowSessionFromSavedModelGeneratorOptions& options) {
|
||||||
@@ -65,7 +66,8 @@ const std::string MaybeConvertSignatureToTag(
|
|||||||
output.resize(name.length());
|
output.resize(name.length());
|
||||||
std::transform(name.begin(), name.end(), output.begin(),
|
std::transform(name.begin(), name.end(), output.begin(),
|
||||||
[](unsigned char c) { return std::toupper(c); });
|
[](unsigned char c) { return std::toupper(c); });
|
||||||
output = absl::Substitute(output, "/", "_");
|
output = absl::StrReplaceAll(output, {{"/", "_"}});
|
||||||
|
output = absl::StrReplaceAll(output, {{"-", "_"}});
|
||||||
return output;
|
return output;
|
||||||
} else {
|
} else {
|
||||||
return name;
|
return name;
|
||||||
@@ -135,7 +137,7 @@ class TensorFlowSessionFromSavedModelGenerator : public PacketGenerator {
|
|||||||
if (!status.ok()) {
|
if (!status.ok()) {
|
||||||
return ::mediapipe::Status(
|
return ::mediapipe::Status(
|
||||||
static_cast<::mediapipe::StatusCode>(status.code()),
|
static_cast<::mediapipe::StatusCode>(status.code()),
|
||||||
status.error_message());
|
status.ToString());
|
||||||
}
|
}
|
||||||
|
|
||||||
auto session = absl::make_unique<TensorFlowSession>();
|
auto session = absl::make_unique<TensorFlowSession>();
|
||||||
|
|||||||
+2
-2
@@ -32,8 +32,8 @@ message TensorFlowSessionFromSavedModelGeneratorOptions {
|
|||||||
// The name of the generic signature to load into the mapping from tags to
|
// The name of the generic signature to load into the mapping from tags to
|
||||||
// tensor names.
|
// tensor names.
|
||||||
optional string signature_name = 2 [default = "serving_default"];
|
optional string signature_name = 2 [default = "serving_default"];
|
||||||
// Whether to convert the signature keys to uppercase and switch /'s to
|
// Whether to convert the signature keys to uppercase as well as switch /'s
|
||||||
// _'s, which enables standard signatures to be used as Tags.
|
// and -'s to _'s, which enables common signatures to be used as Tags.
|
||||||
optional bool convert_signature_to_tags = 3 [default = true];
|
optional bool convert_signature_to_tags = 3 [default = true];
|
||||||
// If true, saved_model_path can have multiple exported models in
|
// If true, saved_model_path can have multiple exported models in
|
||||||
// subdirectories saved_model_path/%08d and the alphabetically last (i.e.,
|
// subdirectories saved_model_path/%08d and the alphabetically last (i.e.,
|
||||||
|
|||||||
+1
-1
@@ -12,7 +12,7 @@
|
|||||||
// See the License for the specific language governing permissions and
|
// See the License for the specific language governing permissions and
|
||||||
// limitations under the License.
|
// limitations under the License.
|
||||||
|
|
||||||
#include "absl/strings/substitute.h"
|
#include "absl/strings/str_replace.h"
|
||||||
#include "mediapipe/calculators/tensorflow/tensorflow_session.h"
|
#include "mediapipe/calculators/tensorflow/tensorflow_session.h"
|
||||||
#include "mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_generator.pb.h"
|
#include "mediapipe/calculators/tensorflow/tensorflow_session_from_saved_model_generator.pb.h"
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
|||||||
@@ -81,7 +81,7 @@ class TFRecordReaderCalculator : public CalculatorBase {
|
|||||||
auto tf_status = tensorflow::Env::Default()->NewRandomAccessFile(
|
auto tf_status = tensorflow::Env::Default()->NewRandomAccessFile(
|
||||||
cc->InputSidePackets().Tag(kTFRecordPath).Get<std::string>(), &file);
|
cc->InputSidePackets().Tag(kTFRecordPath).Get<std::string>(), &file);
|
||||||
RET_CHECK(tf_status.ok())
|
RET_CHECK(tf_status.ok())
|
||||||
<< "Failed to open tfrecord file: " << tf_status.error_message();
|
<< "Failed to open tfrecord file: " << tf_status.ToString();
|
||||||
tensorflow::io::RecordReader reader(file.get(),
|
tensorflow::io::RecordReader reader(file.get(),
|
||||||
tensorflow::io::RecordReaderOptions());
|
tensorflow::io::RecordReaderOptions());
|
||||||
tensorflow::uint64 offset = 0;
|
tensorflow::uint64 offset = 0;
|
||||||
@@ -94,7 +94,7 @@ class TFRecordReaderCalculator : public CalculatorBase {
|
|||||||
while (current_idx <= target_idx) {
|
while (current_idx <= target_idx) {
|
||||||
tf_status = reader.ReadRecord(&offset, &example_str);
|
tf_status = reader.ReadRecord(&offset, &example_str);
|
||||||
RET_CHECK(tf_status.ok())
|
RET_CHECK(tf_status.ok())
|
||||||
<< "Failed to read tfrecord: " << tf_status.error_message();
|
<< "Failed to read tfrecord: " << tf_status.ToString();
|
||||||
if (current_idx == target_idx) {
|
if (current_idx == target_idx) {
|
||||||
if (cc->OutputSidePackets().HasTag(kExampleTag)) {
|
if (cc->OutputSidePackets().HasTag(kExampleTag)) {
|
||||||
tensorflow::Example tf_example;
|
tensorflow::Example tf_example;
|
||||||
|
|||||||
@@ -222,9 +222,11 @@ cc_library(
|
|||||||
deps = [
|
deps = [
|
||||||
":util",
|
":util",
|
||||||
":tflite_inference_calculator_cc_proto",
|
":tflite_inference_calculator_cc_proto",
|
||||||
|
"@com_google_absl//absl/memory",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/util:resource_util",
|
"//mediapipe/util:resource_util",
|
||||||
"@org_tensorflow//tensorflow/lite:framework",
|
"@org_tensorflow//tensorflow/lite:framework",
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
|
||||||
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
||||||
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
|
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
@@ -254,6 +256,10 @@ cc_library(
|
|||||||
"//mediapipe:android": [
|
"//mediapipe:android": [
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/nnapi:nnapi_delegate",
|
"@org_tensorflow//tensorflow/lite/delegates/nnapi:nnapi_delegate",
|
||||||
],
|
],
|
||||||
|
}) + select({
|
||||||
|
"//conditions:default": [
|
||||||
|
"//mediapipe/util:cpu_util",
|
||||||
|
],
|
||||||
}),
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -308,6 +314,20 @@ cc_library(
|
|||||||
alwayslink = 1,
|
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(
|
cc_library(
|
||||||
name = "tflite_tensors_to_segmentation_calculator",
|
name = "tflite_tensors_to_segmentation_calculator",
|
||||||
srcs = ["tflite_tensors_to_segmentation_calculator.cc"],
|
srcs = ["tflite_tensors_to_segmentation_calculator.cc"],
|
||||||
@@ -431,7 +451,7 @@ cc_library(
|
|||||||
"//mediapipe:android": [
|
"//mediapipe:android": [
|
||||||
"//mediapipe/util/android/file/base",
|
"//mediapipe/util/android/file/base",
|
||||||
],
|
],
|
||||||
"//mediapipe:apple": [
|
"//mediapipe:ios": [
|
||||||
"//mediapipe/util/android/file/base",
|
"//mediapipe/util/android/file/base",
|
||||||
],
|
],
|
||||||
"//mediapipe:macos": [
|
"//mediapipe:macos": [
|
||||||
@@ -478,6 +498,9 @@ cc_test(
|
|||||||
deps = [
|
deps = [
|
||||||
":tflite_inference_calculator",
|
":tflite_inference_calculator",
|
||||||
":tflite_inference_calculator_cc_proto",
|
":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_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/deps:file_path",
|
"//mediapipe/framework/deps:file_path",
|
||||||
@@ -485,6 +508,9 @@ cc_test(
|
|||||||
"//mediapipe/framework/port:integral_types",
|
"//mediapipe/framework/port:integral_types",
|
||||||
"//mediapipe/framework/port:parse_text_proto",
|
"//mediapipe/framework/port:parse_text_proto",
|
||||||
"//mediapipe/framework/tool:validate_type",
|
"//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:framework",
|
||||||
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
||||||
],
|
],
|
||||||
@@ -510,3 +536,19 @@ cc_test(
|
|||||||
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
"@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",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|||||||
@@ -294,11 +294,15 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
if (use_quantized_tensors_) {
|
if (use_quantized_tensors_) {
|
||||||
RET_CHECK(image_frame.Format() != mediapipe::ImageFormat::VEC32F1)
|
RET_CHECK(image_frame.Format() != mediapipe::ImageFormat::VEC32F1)
|
||||||
<< "Only 8-bit input images are supported for quantization.";
|
<< "Only 8-bit input images are supported for quantization.";
|
||||||
|
quant.type = kTfLiteAffineQuantization;
|
||||||
|
quant.params = nullptr;
|
||||||
// Optional: Set 'quant' quantization params here if needed.
|
// Optional: Set 'quant' quantization params here if needed.
|
||||||
interpreter_->SetTensorParametersReadWrite(0, kTfLiteUInt8, "",
|
interpreter_->SetTensorParametersReadWrite(0, kTfLiteUInt8, "",
|
||||||
{channels_preserved}, quant);
|
{channels_preserved}, quant);
|
||||||
} else {
|
} else {
|
||||||
// Default TfLiteQuantization used for no quantization.
|
// Initialize structure for no quantization.
|
||||||
|
quant.type = kTfLiteNoQuantization;
|
||||||
|
quant.params = nullptr;
|
||||||
interpreter_->SetTensorParametersReadWrite(0, kTfLiteFloat32, "",
|
interpreter_->SetTensorParametersReadWrite(0, kTfLiteFloat32, "",
|
||||||
{channels_preserved}, quant);
|
{channels_preserved}, quant);
|
||||||
}
|
}
|
||||||
@@ -422,40 +426,35 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
#elif defined(MEDIAPIPE_IOS)
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
// GpuBuffer to id<MTLBuffer> conversion.
|
// GpuBuffer to id<MTLBuffer> conversion.
|
||||||
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
||||||
{
|
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||||
id<MTLTexture> src_texture = [gpu_helper_ metalTextureWithGpuBuffer:input];
|
|
||||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
id<MTLTexture> src_texture = [gpu_helper_ metalTextureWithGpuBuffer:input];
|
||||||
command_buffer.label = @"TfLiteConverterCalculatorConvert";
|
command_buffer.label = @"TfLiteConverterCalculatorConvertAndBlit";
|
||||||
id<MTLComputeCommandEncoder> compute_encoder =
|
id<MTLComputeCommandEncoder> compute_encoder =
|
||||||
[command_buffer computeCommandEncoder];
|
[command_buffer computeCommandEncoder];
|
||||||
[compute_encoder setComputePipelineState:gpu_data_out_->pipeline_state];
|
[compute_encoder setComputePipelineState:gpu_data_out_->pipeline_state];
|
||||||
[compute_encoder setTexture:src_texture atIndex:0];
|
[compute_encoder setTexture:src_texture atIndex:0];
|
||||||
[compute_encoder setBuffer:gpu_data_out_->buffer offset:0 atIndex:1];
|
[compute_encoder setBuffer:gpu_data_out_->buffer offset:0 atIndex:1];
|
||||||
MTLSize threads_per_group = MTLSizeMake(kWorkgroupSize, kWorkgroupSize, 1);
|
MTLSize threads_per_group = MTLSizeMake(kWorkgroupSize, kWorkgroupSize, 1);
|
||||||
MTLSize threadgroups =
|
MTLSize threadgroups =
|
||||||
MTLSizeMake(NumGroups(input.width(), kWorkgroupSize),
|
MTLSizeMake(NumGroups(input.width(), kWorkgroupSize),
|
||||||
NumGroups(input.height(), kWorkgroupSize), 1);
|
NumGroups(input.height(), kWorkgroupSize), 1);
|
||||||
[compute_encoder dispatchThreadgroups:threadgroups
|
[compute_encoder dispatchThreadgroups:threadgroups
|
||||||
threadsPerThreadgroup:threads_per_group];
|
threadsPerThreadgroup:threads_per_group];
|
||||||
[compute_encoder endEncoding];
|
[compute_encoder endEncoding];
|
||||||
[command_buffer commit];
|
|
||||||
[command_buffer waitUntilCompleted];
|
|
||||||
}
|
|
||||||
|
|
||||||
// Copy into outputs.
|
// Copy into outputs.
|
||||||
// TODO Avoid this copy.
|
// TODO Avoid this copy.
|
||||||
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
||||||
output_tensors->resize(1);
|
output_tensors->resize(1);
|
||||||
{
|
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
||||||
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
output_tensors->at(0) =
|
||||||
output_tensors->at(0) =
|
[device newBufferWithLength:gpu_data_out_->elements * sizeof(float)
|
||||||
[device newBufferWithLength:gpu_data_out_->elements * sizeof(float)
|
options:MTLResourceStorageModeShared];
|
||||||
options:MTLResourceStorageModeShared];
|
[MPPMetalUtil blitMetalBufferTo:output_tensors->at(0)
|
||||||
[MPPMetalUtil blitMetalBufferTo:output_tensors->at(0)
|
from:gpu_data_out_->buffer
|
||||||
from:gpu_data_out_->buffer
|
blocking:false
|
||||||
blocking:true
|
commandBuffer:command_buffer];
|
||||||
commandBuffer:[gpu_helper_ commandBuffer]];
|
|
||||||
}
|
|
||||||
|
|
||||||
cc->Outputs()
|
cc->Outputs()
|
||||||
.Tag("TENSORS_GPU")
|
.Tag("TENSORS_GPU")
|
||||||
|
|||||||
@@ -17,10 +17,16 @@
|
|||||||
#include <string>
|
#include <string>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
|
#include "absl/memory/memory.h"
|
||||||
#include "mediapipe/calculators/tflite/tflite_inference_calculator.pb.h"
|
#include "mediapipe/calculators/tflite/tflite_inference_calculator.pb.h"
|
||||||
#include "mediapipe/calculators/tflite/util.h"
|
#include "mediapipe/calculators/tflite/util.h"
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/port/ret_check.h"
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
|
|
||||||
|
#if !defined(__EMSCRIPTEN__)
|
||||||
|
#include "mediapipe/util/cpu_util.h"
|
||||||
|
#endif // !__EMSCRIPTEN__
|
||||||
|
|
||||||
#include "mediapipe/util/resource_util.h"
|
#include "mediapipe/util/resource_util.h"
|
||||||
#include "tensorflow/lite/error_reporter.h"
|
#include "tensorflow/lite/error_reporter.h"
|
||||||
#include "tensorflow/lite/interpreter.h"
|
#include "tensorflow/lite/interpreter.h"
|
||||||
@@ -50,12 +56,16 @@
|
|||||||
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
|
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/metal_delegate_internal.h"
|
#include "tensorflow/lite/delegates/gpu/metal_delegate_internal.h"
|
||||||
#endif // iOS
|
#endif // iOS
|
||||||
|
#include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
|
||||||
#if defined(MEDIAPIPE_ANDROID)
|
#if defined(MEDIAPIPE_ANDROID)
|
||||||
#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
|
#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
|
||||||
#endif // ANDROID
|
#endif // ANDROID
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
|
// Commonly used to compute the number of blocks to launch in a kernel.
|
||||||
|
int NumGroups(const int size, const int group_size) { // NOLINT
|
||||||
|
return (size + group_size - 1) / group_size;
|
||||||
|
}
|
||||||
|
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
|
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
|
||||||
@@ -109,6 +119,23 @@ struct GPUData {
|
|||||||
};
|
};
|
||||||
#endif
|
#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__))
|
||||||
|
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__)
|
||||||
|
return InferHigherCoreIds().size();
|
||||||
|
#else
|
||||||
|
return 1;
|
||||||
|
#endif // !__EMSCRIPTEN__
|
||||||
|
}
|
||||||
|
|
||||||
// Calculator Header Section
|
// Calculator Header Section
|
||||||
|
|
||||||
// Runs inference on the provided input TFLite tensors and TFLite model.
|
// Runs inference on the provided input TFLite tensors and TFLite model.
|
||||||
@@ -135,6 +162,9 @@ struct GPUData {
|
|||||||
// Input side packet:
|
// Input side packet:
|
||||||
// CUSTOM_OP_RESOLVER (optional) - Use a custom op resolver,
|
// CUSTOM_OP_RESOLVER (optional) - Use a custom op resolver,
|
||||||
// instead of the builtin one.
|
// instead of the builtin one.
|
||||||
|
// MODEL (optional) - Use to specify TfLite model
|
||||||
|
// (std::unique_ptr<tflite::FlatBufferModel,
|
||||||
|
// std::function<void(tflite::FlatBufferModel*)>>)
|
||||||
//
|
//
|
||||||
// Example use:
|
// Example use:
|
||||||
// node {
|
// node {
|
||||||
@@ -144,7 +174,21 @@ struct GPUData {
|
|||||||
// options: {
|
// options: {
|
||||||
// [mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
// [mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
||||||
// model_path: "modelname.tflite"
|
// model_path: "modelname.tflite"
|
||||||
// use_gpu: true
|
// delegate { gpu {} }
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// or
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "TfLiteInferenceCalculator"
|
||||||
|
// input_stream: "TENSORS:tensor_image"
|
||||||
|
// input_side_packet: "MODEL:model"
|
||||||
|
// output_stream: "TENSORS:tensors"
|
||||||
|
// options: {
|
||||||
|
// [mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
||||||
|
// delegate { gpu {} }
|
||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
@@ -159,6 +203,12 @@ struct GPUData {
|
|||||||
//
|
//
|
||||||
class TfLiteInferenceCalculator : public CalculatorBase {
|
class TfLiteInferenceCalculator : public CalculatorBase {
|
||||||
public:
|
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);
|
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||||
|
|
||||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||||
@@ -166,22 +216,23 @@ class TfLiteInferenceCalculator : public CalculatorBase {
|
|||||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||||
|
|
||||||
private:
|
private:
|
||||||
::mediapipe::Status LoadOptions(CalculatorContext* cc);
|
|
||||||
::mediapipe::Status LoadModel(CalculatorContext* cc);
|
::mediapipe::Status LoadModel(CalculatorContext* cc);
|
||||||
|
::mediapipe::StatusOr<Packet> GetModelAsPacket(const CalculatorContext& cc);
|
||||||
::mediapipe::Status LoadDelegate(CalculatorContext* cc);
|
::mediapipe::Status LoadDelegate(CalculatorContext* cc);
|
||||||
|
|
||||||
|
Packet model_packet_;
|
||||||
std::unique_ptr<tflite::Interpreter> interpreter_;
|
std::unique_ptr<tflite::Interpreter> interpreter_;
|
||||||
std::unique_ptr<tflite::FlatBufferModel> model_;
|
TfLiteDelegatePtr delegate_;
|
||||||
TfLiteDelegate* delegate_ = nullptr;
|
|
||||||
|
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||||
std::unique_ptr<GPUData> gpu_data_in_;
|
std::vector<std::unique_ptr<GPUData>> gpu_data_in_;
|
||||||
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
|
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
|
||||||
#elif defined(MEDIAPIPE_IOS)
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
MPPMetalHelper* gpu_helper_ = nullptr;
|
MPPMetalHelper* gpu_helper_ = nullptr;
|
||||||
std::unique_ptr<GPUData> gpu_data_in_;
|
std::vector<std::unique_ptr<GPUData>> gpu_data_in_;
|
||||||
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
|
std::vector<std::unique_ptr<GPUData>> gpu_data_out_;
|
||||||
|
id<MTLComputePipelineState> fp32_to_fp16_program_;
|
||||||
TFLBufferConvert* converter_from_BPHWC4_ = nil;
|
TFLBufferConvert* converter_from_BPHWC4_ = nil;
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
@@ -190,7 +241,6 @@ class TfLiteInferenceCalculator : public CalculatorBase {
|
|||||||
edgetpu::EdgeTpuManager::GetSingleton()->OpenDevice();
|
edgetpu::EdgeTpuManager::GetSingleton()->OpenDevice();
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
std::string model_path_ = "";
|
|
||||||
bool gpu_inference_ = false;
|
bool gpu_inference_ = false;
|
||||||
bool gpu_input_ = false;
|
bool gpu_input_ = false;
|
||||||
bool gpu_output_ = false;
|
bool gpu_output_ = false;
|
||||||
@@ -207,12 +257,22 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
RET_CHECK(cc->Outputs().HasTag("TENSORS") ^
|
RET_CHECK(cc->Outputs().HasTag("TENSORS") ^
|
||||||
cc->Outputs().HasTag("TENSORS_GPU"));
|
cc->Outputs().HasTag("TENSORS_GPU"));
|
||||||
|
|
||||||
bool use_gpu = false;
|
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"))
|
if (cc->Inputs().HasTag("TENSORS"))
|
||||||
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||||
|
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("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||||
use_gpu |= true;
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
@@ -222,6 +282,9 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
if (cc->Outputs().HasTag("TENSORS_GPU")) {
|
if (cc->Outputs().HasTag("TENSORS_GPU")) {
|
||||||
|
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("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||||
use_gpu |= true;
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
@@ -232,10 +295,9 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
.Tag("CUSTOM_OP_RESOLVER")
|
.Tag("CUSTOM_OP_RESOLVER")
|
||||||
.Set<tflite::ops::builtin::BuiltinOpResolver>();
|
.Set<tflite::ops::builtin::BuiltinOpResolver>();
|
||||||
}
|
}
|
||||||
|
if (cc->InputSidePackets().HasTag("MODEL")) {
|
||||||
const auto& options =
|
cc->InputSidePackets().Tag("MODEL").Set<TfLiteModelPtr>();
|
||||||
cc->Options<::mediapipe::TfLiteInferenceCalculatorOptions>();
|
}
|
||||||
use_gpu |= options.use_gpu();
|
|
||||||
|
|
||||||
if (use_gpu) {
|
if (use_gpu) {
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
@@ -254,7 +316,9 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
::mediapipe::Status TfLiteInferenceCalculator::Open(CalculatorContext* cc) {
|
::mediapipe::Status TfLiteInferenceCalculator::Open(CalculatorContext* cc) {
|
||||||
cc->SetOffset(TimestampDiff(0));
|
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("TENSORS_GPU")) {
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
@@ -308,22 +372,41 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
const auto& input_tensors =
|
const auto& input_tensors =
|
||||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||||
RET_CHECK_EQ(input_tensors.size(), 1);
|
RET_CHECK_GT(input_tensors.size(), 0);
|
||||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||||
[this, &input_tensors]() -> ::mediapipe::Status {
|
[this, &input_tensors]() -> ::mediapipe::Status {
|
||||||
// Explicit copy input.
|
// Explicit copy input.
|
||||||
RET_CHECK_CALL(CopyBuffer(input_tensors[0], gpu_data_in_->buffer));
|
gpu_data_in_.resize(input_tensors.size());
|
||||||
|
for (int i = 0; i < input_tensors.size(); ++i) {
|
||||||
|
RET_CHECK_CALL(
|
||||||
|
CopyBuffer(input_tensors[i], gpu_data_in_[i]->buffer));
|
||||||
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}));
|
}));
|
||||||
#elif defined(MEDIAPIPE_IOS)
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
const auto& input_tensors =
|
const auto& input_tensors =
|
||||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||||
RET_CHECK_EQ(input_tensors.size(), 1);
|
RET_CHECK_GT(input_tensors.size(), 0);
|
||||||
// Explicit copy input.
|
// Explicit copy input with conversion float 32 bits to 16 bits.
|
||||||
[MPPMetalUtil blitMetalBufferTo:gpu_data_in_->buffer
|
gpu_data_in_.resize(input_tensors.size());
|
||||||
from:input_tensors[0]
|
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||||
blocking:true
|
command_buffer.label = @"TfLiteInferenceCalculatorConvert";
|
||||||
commandBuffer:[gpu_helper_ commandBuffer]];
|
id<MTLComputeCommandEncoder> compute_encoder =
|
||||||
|
[command_buffer computeCommandEncoder];
|
||||||
|
[compute_encoder setComputePipelineState:fp32_to_fp16_program_];
|
||||||
|
for (int i = 0; i < input_tensors.size(); ++i) {
|
||||||
|
[compute_encoder setBuffer:input_tensors[i] offset:0 atIndex:0];
|
||||||
|
[compute_encoder setBuffer:gpu_data_in_[i]->buffer offset:0 atIndex:1];
|
||||||
|
constexpr int kWorkgroupSize = 64; // Block size for GPU shader.
|
||||||
|
MTLSize threads_per_group = MTLSizeMake(kWorkgroupSize, 1, 1);
|
||||||
|
const int threadgroups =
|
||||||
|
NumGroups(gpu_data_in_[i]->elements, kWorkgroupSize);
|
||||||
|
[compute_encoder dispatchThreadgroups:MTLSizeMake(threadgroups, 1, 1)
|
||||||
|
threadsPerThreadgroup:threads_per_group];
|
||||||
|
}
|
||||||
|
[compute_encoder endEncoding];
|
||||||
|
[command_buffer commit];
|
||||||
#else
|
#else
|
||||||
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
||||||
#endif
|
#endif
|
||||||
@@ -404,7 +487,6 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
}
|
}
|
||||||
[convert_command endEncoding];
|
[convert_command endEncoding];
|
||||||
[command_buffer commit];
|
[command_buffer commit];
|
||||||
[command_buffer waitUntilCompleted];
|
|
||||||
cc->Outputs()
|
cc->Outputs()
|
||||||
.Tag("TENSORS_GPU")
|
.Tag("TENSORS_GPU")
|
||||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||||
@@ -431,22 +513,27 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
if (gpu_inference_) {
|
if (gpu_inference_) {
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
|
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
|
||||||
TfLiteGpuDelegateDelete(delegate_);
|
delegate_ = nullptr;
|
||||||
gpu_data_in_.reset();
|
for (int i = 0; i < gpu_data_in_.size(); ++i) {
|
||||||
|
gpu_data_in_[i].reset();
|
||||||
|
}
|
||||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||||
gpu_data_out_[i].reset();
|
gpu_data_out_[i].reset();
|
||||||
}
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}));
|
}));
|
||||||
#elif defined(MEDIAPIPE_IOS)
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
TFLGpuDelegateDelete(delegate_);
|
delegate_ = nullptr;
|
||||||
gpu_data_in_.reset();
|
for (int i = 0; i < gpu_data_in_.size(); ++i) {
|
||||||
|
gpu_data_in_[i].reset();
|
||||||
|
}
|
||||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||||
gpu_data_out_[i].reset();
|
gpu_data_out_[i].reset();
|
||||||
}
|
}
|
||||||
#endif
|
#endif
|
||||||
|
} else {
|
||||||
|
delegate_ = nullptr;
|
||||||
}
|
}
|
||||||
delegate_ = nullptr;
|
|
||||||
}
|
}
|
||||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||||
edgetpu_context_.reset();
|
edgetpu_context_.reset();
|
||||||
@@ -456,33 +543,10 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
|
|
||||||
// Calculator Auxiliary Section
|
// 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()) {
|
|
||||||
auto 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.use_gpu();
|
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
|
||||||
}
|
|
||||||
|
|
||||||
::mediapipe::Status TfLiteInferenceCalculator::LoadModel(
|
::mediapipe::Status TfLiteInferenceCalculator::LoadModel(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
model_ = tflite::FlatBufferModel::BuildFromFile(model_path_.c_str());
|
ASSIGN_OR_RETURN(model_packet_, GetModelAsPacket(*cc));
|
||||||
RET_CHECK(model_);
|
const auto& model = *model_packet_.Get<TfLiteModelPtr>();
|
||||||
|
|
||||||
tflite::ops::builtin::BuiltinOpResolver op_resolver;
|
tflite::ops::builtin::BuiltinOpResolver op_resolver;
|
||||||
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
|
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
|
||||||
@@ -492,15 +556,19 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
}
|
}
|
||||||
#if defined(MEDIAPIPE_EDGE_TPU)
|
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||||
interpreter_ =
|
interpreter_ =
|
||||||
BuildEdgeTpuInterpreter(*model_, &op_resolver, edgetpu_context_.get());
|
BuildEdgeTpuInterpreter(model, &op_resolver, edgetpu_context_.get());
|
||||||
#else
|
#else
|
||||||
tflite::InterpreterBuilder(*model_, op_resolver)(&interpreter_);
|
tflite::InterpreterBuilder(model, op_resolver)(&interpreter_);
|
||||||
#endif // MEDIAPIPE_EDGE_TPU
|
#endif // MEDIAPIPE_EDGE_TPU
|
||||||
|
|
||||||
RET_CHECK(interpreter_);
|
RET_CHECK(interpreter_);
|
||||||
|
|
||||||
#if defined(__EMSCRIPTEN__)
|
#if defined(__EMSCRIPTEN__) || defined(MEDIAPIPE_EDGE_TPU)
|
||||||
interpreter_->SetNumThreads(1);
|
interpreter_->SetNumThreads(1);
|
||||||
|
#else
|
||||||
|
interpreter_->SetNumThreads(
|
||||||
|
cc->Options<mediapipe::TfLiteInferenceCalculatorOptions>()
|
||||||
|
.cpu_num_thread());
|
||||||
#endif // __EMSCRIPTEN__
|
#endif // __EMSCRIPTEN__
|
||||||
|
|
||||||
if (gpu_output_) {
|
if (gpu_output_) {
|
||||||
@@ -516,22 +584,77 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
return ::mediapipe::OkStatus();
|
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(
|
::mediapipe::Status TfLiteInferenceCalculator::LoadDelegate(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
#if defined(MEDIAPIPE_ANDROID)
|
const auto& calculator_opts =
|
||||||
|
cc->Options<mediapipe::TfLiteInferenceCalculatorOptions>();
|
||||||
|
if (calculator_opts.has_delegate() &&
|
||||||
|
calculator_opts.delegate().has_tflite()) {
|
||||||
|
// Default tflite inference requeqsted - no need to modify graph.
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
if (!gpu_inference_) {
|
if (!gpu_inference_) {
|
||||||
if (cc->Options<mediapipe::TfLiteInferenceCalculatorOptions>()
|
#if defined(MEDIAPIPE_ANDROID)
|
||||||
.use_nnapi()) {
|
const bool nnapi_requested = calculator_opts.has_delegate()
|
||||||
|
? calculator_opts.delegate().has_nnapi()
|
||||||
|
: calculator_opts.use_nnapi();
|
||||||
|
if (nnapi_requested) {
|
||||||
// Attempt to use NNAPI.
|
// Attempt to use NNAPI.
|
||||||
// If not supported, the default CPU delegate will be created and used.
|
// If not supported, the default CPU delegate will be created and used.
|
||||||
interpreter_->SetAllowFp16PrecisionForFp32(1);
|
interpreter_->SetAllowFp16PrecisionForFp32(1);
|
||||||
delegate_ = tflite::NnApiDelegate();
|
delegate_ =
|
||||||
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
|
TfLiteDelegatePtr(tflite::NnApiDelegate(), [](TfLiteDelegate*) {
|
||||||
|
// No need to free according to tflite::NnApiDelegate()
|
||||||
|
// documentation.
|
||||||
|
});
|
||||||
|
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
|
||||||
|
kTfLiteOk);
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
#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, no need for GPU delegate below.
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
#endif // ANDROID
|
|
||||||
|
|
||||||
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
// Configure and create the delegate.
|
// Configure and create the delegate.
|
||||||
@@ -541,28 +664,30 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
TFLITE_GL_OBJECT_TYPE_FASTEST;
|
TFLITE_GL_OBJECT_TYPE_FASTEST;
|
||||||
options.compile_options.dynamic_batch_enabled = 0;
|
options.compile_options.dynamic_batch_enabled = 0;
|
||||||
options.compile_options.inline_parameters = 1;
|
options.compile_options.inline_parameters = 1;
|
||||||
if (!delegate_) delegate_ = TfLiteGpuDelegateCreate(&options);
|
if (!delegate_)
|
||||||
|
delegate_ = TfLiteDelegatePtr(TfLiteGpuDelegateCreate(&options),
|
||||||
|
&TfLiteGpuDelegateDelete);
|
||||||
|
|
||||||
if (gpu_input_) {
|
if (gpu_input_) {
|
||||||
// Get input image sizes.
|
// Get input image sizes.
|
||||||
gpu_data_in_ = absl::make_unique<GPUData>();
|
|
||||||
const auto& input_indices = interpreter_->inputs();
|
const auto& input_indices = interpreter_->inputs();
|
||||||
RET_CHECK_EQ(input_indices.size(), 1); // TODO accept > 1.
|
gpu_data_in_.resize(input_indices.size());
|
||||||
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[0]);
|
for (int i = 0; i < input_indices.size(); ++i) {
|
||||||
gpu_data_in_->elements = 1;
|
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[i]);
|
||||||
for (int d = 0; d < tensor->dims->size; ++d) {
|
gpu_data_in_[i] = absl::make_unique<GPUData>();
|
||||||
gpu_data_in_->elements *= tensor->dims->data[d];
|
gpu_data_in_[i]->elements = 1;
|
||||||
|
for (int d = 0; d < tensor->dims->size; ++d) {
|
||||||
|
gpu_data_in_[i]->elements *= tensor->dims->data[d];
|
||||||
|
}
|
||||||
|
// Create and bind input buffer.
|
||||||
|
RET_CHECK_CALL(
|
||||||
|
::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
|
||||||
|
gpu_data_in_[i]->elements, &gpu_data_in_[i]->buffer));
|
||||||
|
RET_CHECK_EQ(TfLiteGpuDelegateBindBufferToTensor(
|
||||||
|
delegate_.get(), gpu_data_in_[i]->buffer.id(),
|
||||||
|
interpreter_->inputs()[i]),
|
||||||
|
kTfLiteOk);
|
||||||
}
|
}
|
||||||
CHECK_GE(tensor->dims->data[3], 1);
|
|
||||||
CHECK_LE(tensor->dims->data[3], 4);
|
|
||||||
CHECK_NE(tensor->dims->data[3], 2);
|
|
||||||
// Create and bind input buffer.
|
|
||||||
RET_CHECK_CALL(::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
|
|
||||||
gpu_data_in_->elements, &gpu_data_in_->buffer));
|
|
||||||
RET_CHECK_EQ(TfLiteGpuDelegateBindBufferToTensor(
|
|
||||||
delegate_, gpu_data_in_->buffer.id(),
|
|
||||||
interpreter_->inputs()[0]), // First tensor only
|
|
||||||
kTfLiteOk);
|
|
||||||
}
|
}
|
||||||
if (gpu_output_) {
|
if (gpu_output_) {
|
||||||
// Get output image sizes.
|
// Get output image sizes.
|
||||||
@@ -582,53 +707,85 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||||
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
|
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
|
||||||
gpu_data_out_[i]->elements, &gpu_data_out_[i]->buffer));
|
gpu_data_out_[i]->elements, &gpu_data_out_[i]->buffer));
|
||||||
RET_CHECK_EQ(
|
RET_CHECK_EQ(TfLiteGpuDelegateBindBufferToTensor(
|
||||||
TfLiteGpuDelegateBindBufferToTensor(
|
delegate_.get(), gpu_data_out_[i]->buffer.id(),
|
||||||
delegate_, gpu_data_out_[i]->buffer.id(), output_indices[i]),
|
output_indices[i]),
|
||||||
kTfLiteOk);
|
kTfLiteOk);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Must call this last.
|
// Must call this last.
|
||||||
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
|
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
|
||||||
|
kTfLiteOk);
|
||||||
#endif // OpenGL
|
#endif // OpenGL
|
||||||
|
|
||||||
#if defined(MEDIAPIPE_IOS)
|
#if defined(MEDIAPIPE_IOS)
|
||||||
|
const int kHalfSize = 2; // sizeof(half)
|
||||||
// Configure and create the delegate.
|
// Configure and create the delegate.
|
||||||
TFLGpuDelegateOptions options;
|
TFLGpuDelegateOptions options;
|
||||||
options.allow_precision_loss = false; // Must match converter, F=float/T=half
|
options.allow_precision_loss = true;
|
||||||
options.wait_type = TFLGpuDelegateWaitType::TFLGpuDelegateWaitTypePassive;
|
options.wait_type = TFLGpuDelegateWaitType::TFLGpuDelegateWaitTypePassive;
|
||||||
if (!delegate_) delegate_ = TFLGpuDelegateCreate(&options);
|
if (!delegate_)
|
||||||
|
delegate_ = TfLiteDelegatePtr(TFLGpuDelegateCreate(&options),
|
||||||
|
&TFLGpuDelegateDelete);
|
||||||
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
||||||
|
|
||||||
if (gpu_input_) {
|
if (gpu_input_) {
|
||||||
// Get input image sizes.
|
// Get input image sizes.
|
||||||
gpu_data_in_ = absl::make_unique<GPUData>();
|
|
||||||
const auto& input_indices = interpreter_->inputs();
|
const auto& input_indices = interpreter_->inputs();
|
||||||
RET_CHECK_EQ(input_indices.size(), 1);
|
gpu_data_in_.resize(input_indices.size());
|
||||||
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[0]);
|
for (int i = 0; i < input_indices.size(); ++i) {
|
||||||
gpu_data_in_->elements = 1;
|
const TfLiteTensor* tensor = interpreter_->tensor(input_indices[i]);
|
||||||
// On iOS GPU, input must be 4 channels, regardless of what model expects.
|
gpu_data_in_[i] = absl::make_unique<GPUData>();
|
||||||
{
|
gpu_data_in_[i]->shape.b = tensor->dims->data[0];
|
||||||
gpu_data_in_->elements *= tensor->dims->data[0]; // batch
|
gpu_data_in_[i]->shape.h = tensor->dims->data[1];
|
||||||
gpu_data_in_->elements *= tensor->dims->data[1]; // height
|
gpu_data_in_[i]->shape.w = tensor->dims->data[2];
|
||||||
gpu_data_in_->elements *= tensor->dims->data[2]; // width
|
// On iOS GPU, input must be 4 channels, regardless of what model expects.
|
||||||
gpu_data_in_->elements *= 4; // channels
|
gpu_data_in_[i]->shape.c = 4;
|
||||||
|
gpu_data_in_[i]->elements =
|
||||||
|
gpu_data_in_[i]->shape.b * gpu_data_in_[i]->shape.h *
|
||||||
|
gpu_data_in_[i]->shape.w * gpu_data_in_[i]->shape.c;
|
||||||
|
// Input to model can be RGBA only.
|
||||||
|
if (tensor->dims->data[3] != 4) {
|
||||||
|
LOG(WARNING) << "Please ensure input GPU tensor is 4 channels.";
|
||||||
|
}
|
||||||
|
const std::string shader_source =
|
||||||
|
absl::Substitute(R"(#include <metal_stdlib>
|
||||||
|
using namespace metal;
|
||||||
|
kernel void convertKernel(device float4* const input_buffer [[buffer(0)]],
|
||||||
|
device half4* output_buffer [[buffer(1)]],
|
||||||
|
uint gid [[thread_position_in_grid]]) {
|
||||||
|
if (gid >= $0) return;
|
||||||
|
output_buffer[gid] = half4(input_buffer[gid]);
|
||||||
|
})",
|
||||||
|
gpu_data_in_[i]->elements / 4);
|
||||||
|
NSString* library_source =
|
||||||
|
[NSString stringWithUTF8String:shader_source.c_str()];
|
||||||
|
NSError* error = nil;
|
||||||
|
id<MTLLibrary> library =
|
||||||
|
[device newLibraryWithSource:library_source options:nil error:&error];
|
||||||
|
RET_CHECK(library != nil) << "Couldn't create shader library "
|
||||||
|
<< [[error localizedDescription] UTF8String];
|
||||||
|
id<MTLFunction> kernel_func = nil;
|
||||||
|
kernel_func = [library newFunctionWithName:@"convertKernel"];
|
||||||
|
RET_CHECK(kernel_func != nil) << "Couldn't create kernel function.";
|
||||||
|
fp32_to_fp16_program_ =
|
||||||
|
[device newComputePipelineStateWithFunction:kernel_func error:&error];
|
||||||
|
RET_CHECK(fp32_to_fp16_program_ != nil)
|
||||||
|
<< "Couldn't create pipeline state "
|
||||||
|
<< [[error localizedDescription] UTF8String];
|
||||||
|
|
||||||
|
// Create and bind input buffer.
|
||||||
|
gpu_data_in_[i]->buffer =
|
||||||
|
[device newBufferWithLength:gpu_data_in_[i]->elements * kHalfSize
|
||||||
|
options:MTLResourceStorageModeShared];
|
||||||
|
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
|
||||||
|
kTfLiteOk);
|
||||||
|
RET_CHECK_EQ(
|
||||||
|
TFLGpuDelegateBindMetalBufferToTensor(
|
||||||
|
delegate_.get(), input_indices[i], gpu_data_in_[i]->buffer),
|
||||||
|
true);
|
||||||
}
|
}
|
||||||
// Input to model can be RGBA only.
|
|
||||||
if (tensor->dims->data[3] != 4) {
|
|
||||||
LOG(WARNING) << "Please ensure input GPU tensor is 4 channels.";
|
|
||||||
}
|
|
||||||
// Create and bind input buffer.
|
|
||||||
gpu_data_in_->buffer =
|
|
||||||
[device newBufferWithLength:gpu_data_in_->elements * sizeof(float)
|
|
||||||
options:MTLResourceStorageModeShared];
|
|
||||||
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
|
|
||||||
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
|
|
||||||
delegate_,
|
|
||||||
input_indices[0], // First tensor only
|
|
||||||
gpu_data_in_->buffer),
|
|
||||||
true);
|
|
||||||
}
|
}
|
||||||
if (gpu_output_) {
|
if (gpu_output_) {
|
||||||
// Get output image sizes.
|
// Get output image sizes.
|
||||||
@@ -669,15 +826,17 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
interpreter_->SetAllowBufferHandleOutput(true);
|
interpreter_->SetAllowBufferHandleOutput(true);
|
||||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||||
gpu_data_out_[i]->buffer =
|
gpu_data_out_[i]->buffer =
|
||||||
[device newBufferWithLength:gpu_data_out_[i]->elements * sizeof(float)
|
[device newBufferWithLength:gpu_data_out_[i]->elements * kHalfSize
|
||||||
options:MTLResourceStorageModeShared];
|
options:MTLResourceStorageModeShared];
|
||||||
RET_CHECK_EQ(TFLGpuDelegateBindMetalBufferToTensor(
|
RET_CHECK_EQ(
|
||||||
delegate_, output_indices[i], gpu_data_out_[i]->buffer),
|
TFLGpuDelegateBindMetalBufferToTensor(
|
||||||
true);
|
delegate_.get(), output_indices[i], gpu_data_out_[i]->buffer),
|
||||||
|
true);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Create converter for GPU output.
|
// Create converter for GPU output.
|
||||||
converter_from_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
|
converter_from_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
|
||||||
isFloat16:false
|
isFloat16:true
|
||||||
convertToPBHWC4:false];
|
convertToPBHWC4:false];
|
||||||
if (converter_from_BPHWC4_ == nil) {
|
if (converter_from_BPHWC4_ == nil) {
|
||||||
return mediapipe::InternalError(
|
return mediapipe::InternalError(
|
||||||
|
|||||||
@@ -27,7 +27,7 @@ import "mediapipe/framework/calculator.proto";
|
|||||||
// options {
|
// options {
|
||||||
// [mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
// [mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
||||||
// model_path: "model.tflite"
|
// model_path: "model.tflite"
|
||||||
// use_gpu: true
|
// delegate { gpu {} }
|
||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
@@ -37,6 +37,28 @@ message TfLiteInferenceCalculatorOptions {
|
|||||||
optional TfLiteInferenceCalculatorOptions ext = 233867213;
|
optional TfLiteInferenceCalculatorOptions ext = 233867213;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
message Delegate {
|
||||||
|
// Default inference provided by tflite.
|
||||||
|
message TfLite {}
|
||||||
|
// Delegate to run GPU inference depending on the device.
|
||||||
|
// (Can use OpenGl, OpenCl, Metal depending on the device.)
|
||||||
|
message Gpu {}
|
||||||
|
// 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;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// Path to the TF Lite model (ex: /path/to/modelname.tflite).
|
// Path to the TF Lite model (ex: /path/to/modelname.tflite).
|
||||||
// On mobile, this is generally just modelname.tflite.
|
// On mobile, this is generally just modelname.tflite.
|
||||||
optional string model_path = 1;
|
optional string model_path = 1;
|
||||||
@@ -44,10 +66,22 @@ message TfLiteInferenceCalculatorOptions {
|
|||||||
// Whether the TF Lite GPU or CPU backend should be used. Effective only when
|
// Whether the TF Lite GPU or CPU backend should be used. Effective only when
|
||||||
// input tensors are on CPU. For input tensors on GPU, GPU backend is always
|
// input tensors are on CPU. For input tensors on GPU, GPU backend is always
|
||||||
// used.
|
// used.
|
||||||
optional bool use_gpu = 2 [default = false];
|
// DEPRECATED: configure "delegate" instead.
|
||||||
|
optional bool use_gpu = 2 [deprecated = true, default = false];
|
||||||
|
|
||||||
// Android only. When true, an NNAPI delegate will be used for inference.
|
// Android only. When true, an NNAPI delegate will be used for inference.
|
||||||
// If NNAPI is not available, then the default CPU delegate will be used
|
// If NNAPI is not available, then the default CPU delegate will be used
|
||||||
// automatically.
|
// automatically.
|
||||||
optional bool use_nnapi = 3 [default = false];
|
// DEPRECATED: configure "delegate" instead.
|
||||||
|
optional bool use_nnapi = 3 [deprecated = true, default = false];
|
||||||
|
|
||||||
|
// The number of threads available to the interpreter. Effective only when
|
||||||
|
// input tensors are on CPU and 'use_gpu' is false.
|
||||||
|
optional int32 cpu_num_thread = 4 [default = -1];
|
||||||
|
|
||||||
|
// TfLite delegate to run inference.
|
||||||
|
// NOTE: calculator is free to choose delegate if not specified explicitly.
|
||||||
|
// NOTE: use_gpu/use_nnapi are ignored if specified. (Delegate takes
|
||||||
|
// precedence over use_* deprecated options.)
|
||||||
|
optional Delegate delegate = 5;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -16,6 +16,8 @@
|
|||||||
#include <string>
|
#include <string>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
|
#include "absl/strings/str_replace.h"
|
||||||
|
#include "absl/strings/string_view.h"
|
||||||
#include "mediapipe/calculators/tflite/tflite_inference_calculator.pb.h"
|
#include "mediapipe/calculators/tflite/tflite_inference_calculator.pb.h"
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/calculator_runner.h"
|
#include "mediapipe/framework/calculator_runner.h"
|
||||||
@@ -39,13 +41,7 @@ namespace mediapipe {
|
|||||||
|
|
||||||
using ::tflite::Interpreter;
|
using ::tflite::Interpreter;
|
||||||
|
|
||||||
class TfLiteInferenceCalculatorTest : public ::testing::Test {
|
void DoSmokeTest(const std::string& graph_proto) {
|
||||||
protected:
|
|
||||||
std::unique_ptr<CalculatorRunner> runner_ = nullptr;
|
|
||||||
};
|
|
||||||
|
|
||||||
// Tests a simple add model that adds an input tensor to itself.
|
|
||||||
TEST_F(TfLiteInferenceCalculatorTest, SmokeTest) {
|
|
||||||
const int width = 8;
|
const int width = 8;
|
||||||
const int height = 8;
|
const int height = 8;
|
||||||
const int channels = 3;
|
const int channels = 3;
|
||||||
@@ -75,21 +71,7 @@ TEST_F(TfLiteInferenceCalculatorTest, SmokeTest) {
|
|||||||
|
|
||||||
// Prepare single calculator graph to and wait for packets.
|
// Prepare single calculator graph to and wait for packets.
|
||||||
CalculatorGraphConfig graph_config =
|
CalculatorGraphConfig graph_config =
|
||||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
ParseTextProtoOrDie<CalculatorGraphConfig>(graph_proto);
|
||||||
R"(
|
|
||||||
input_stream: "tensor_in"
|
|
||||||
node {
|
|
||||||
calculator: "TfLiteInferenceCalculator"
|
|
||||||
input_stream: "TENSORS:tensor_in"
|
|
||||||
output_stream: "TENSORS:tensor_out"
|
|
||||||
options {
|
|
||||||
[mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
|
||||||
use_gpu: false
|
|
||||||
model_path: "mediapipe/calculators/tflite/testdata/add.bin"
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
)");
|
|
||||||
std::vector<Packet> output_packets;
|
std::vector<Packet> output_packets;
|
||||||
tool::AddVectorSink("tensor_out", &graph_config, &output_packets);
|
tool::AddVectorSink("tensor_out", &graph_config, &output_packets);
|
||||||
CalculatorGraph graph(graph_config);
|
CalculatorGraph graph(graph_config);
|
||||||
@@ -120,4 +102,72 @@ TEST_F(TfLiteInferenceCalculatorTest, SmokeTest) {
|
|||||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Tests a simple add model that adds an input tensor to itself.
|
||||||
|
TEST(TfLiteInferenceCalculatorTest, SmokeTest) {
|
||||||
|
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
|
} // 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];
|
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_) {
|
if (label_map_loaded_) {
|
||||||
RET_CHECK_EQ(num_classes, label_map_.size());
|
RET_CHECK_EQ(num_classes, label_map_.size());
|
||||||
}
|
}
|
||||||
const float* raw_scores = raw_score_tensor->data.f;
|
const float* raw_scores = raw_score_tensor->data.f;
|
||||||
|
|
||||||
auto classification_list = absl::make_unique<ClassificationList>();
|
auto classification_list = absl::make_unique<ClassificationList>();
|
||||||
for (int i = 0; i < num_classes; ++i) {
|
if (options_.binary_classification()) {
|
||||||
if (options_.has_min_score_threshold() &&
|
Classification* class_first = classification_list->add_classification();
|
||||||
raw_scores[i] < options_.min_score_threshold()) {
|
Classification* class_second = classification_list->add_classification();
|
||||||
continue;
|
class_first->set_index(0);
|
||||||
}
|
class_second->set_index(1);
|
||||||
Classification* classification = classification_list->add_classification();
|
class_first->set_score(raw_scores[0]);
|
||||||
classification->set_index(i);
|
class_second->set_score(1. - raw_scores[0]);
|
||||||
classification->set_score(raw_scores[i]);
|
|
||||||
if (label_map_loaded_) {
|
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;
|
optional int32 top_k = 2;
|
||||||
// Path to a label map file for getting the actual name of class ids.
|
// Path to a label map file for getting the actual name of class ids.
|
||||||
optional string label_map_path = 3;
|
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;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -472,11 +472,11 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
// Copy inputs.
|
// Copy inputs.
|
||||||
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_boxes_buffer
|
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_boxes_buffer
|
||||||
from:input_tensors[0]
|
from:input_tensors[0]
|
||||||
blocking:true
|
blocking:false
|
||||||
commandBuffer:[gpu_helper_ commandBuffer]];
|
commandBuffer:[gpu_helper_ commandBuffer]];
|
||||||
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_scores_buffer
|
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_scores_buffer
|
||||||
from:input_tensors[1]
|
from:input_tensors[1]
|
||||||
blocking:true
|
blocking:false
|
||||||
commandBuffer:[gpu_helper_ commandBuffer]];
|
commandBuffer:[gpu_helper_ commandBuffer]];
|
||||||
if (!anchors_init_) {
|
if (!anchors_init_) {
|
||||||
if (side_packet_anchors_) {
|
if (side_packet_anchors_) {
|
||||||
@@ -491,48 +491,37 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
RET_CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
|
RET_CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
|
||||||
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_anchors_buffer
|
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_anchors_buffer
|
||||||
from:input_tensors[2]
|
from:input_tensors[2]
|
||||||
blocking:true
|
blocking:false
|
||||||
commandBuffer:[gpu_helper_ commandBuffer]];
|
commandBuffer:[gpu_helper_ commandBuffer]];
|
||||||
}
|
}
|
||||||
anchors_init_ = true;
|
anchors_init_ = true;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Run shaders.
|
// Run shaders.
|
||||||
{
|
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
command_buffer.label = @"TfLiteDecodeAndScoreBoxes";
|
||||||
command_buffer.label = @"TfLiteDecodeBoxes";
|
id<MTLComputeCommandEncoder> command_encoder =
|
||||||
id<MTLComputeCommandEncoder> decode_command =
|
[command_buffer computeCommandEncoder];
|
||||||
[command_buffer computeCommandEncoder];
|
[command_encoder setComputePipelineState:gpu_data_->decode_program];
|
||||||
[decode_command setComputePipelineState:gpu_data_->decode_program];
|
[command_encoder setBuffer:gpu_data_->decoded_boxes_buffer
|
||||||
[decode_command setBuffer:gpu_data_->decoded_boxes_buffer
|
offset:0
|
||||||
offset:0
|
atIndex:0];
|
||||||
atIndex:0];
|
[command_encoder setBuffer:gpu_data_->raw_boxes_buffer offset:0 atIndex:1];
|
||||||
[decode_command setBuffer:gpu_data_->raw_boxes_buffer offset:0 atIndex:1];
|
[command_encoder setBuffer:gpu_data_->raw_anchors_buffer offset:0 atIndex:2];
|
||||||
[decode_command setBuffer:gpu_data_->raw_anchors_buffer offset:0 atIndex:2];
|
MTLSize decode_threads_per_group = MTLSizeMake(1, 1, 1);
|
||||||
MTLSize decode_threads_per_group = MTLSizeMake(1, 1, 1);
|
MTLSize decode_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
|
||||||
MTLSize decode_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
|
[command_encoder dispatchThreadgroups:decode_threadgroups
|
||||||
[decode_command dispatchThreadgroups:decode_threadgroups
|
threadsPerThreadgroup:decode_threads_per_group];
|
||||||
threadsPerThreadgroup:decode_threads_per_group];
|
|
||||||
[decode_command endEncoding];
|
[command_encoder setComputePipelineState:gpu_data_->score_program];
|
||||||
[command_buffer commit];
|
[command_encoder setBuffer:gpu_data_->scored_boxes_buffer offset:0 atIndex:0];
|
||||||
[command_buffer waitUntilCompleted];
|
[command_encoder setBuffer:gpu_data_->raw_scores_buffer offset:0 atIndex:1];
|
||||||
}
|
MTLSize score_threads_per_group = MTLSizeMake(1, num_classes_, 1);
|
||||||
{
|
MTLSize score_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
|
||||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
[command_encoder dispatchThreadgroups:score_threadgroups
|
||||||
command_buffer.label = @"TfLiteScoreBoxes";
|
|
||||||
id<MTLComputeCommandEncoder> score_command =
|
|
||||||
[command_buffer computeCommandEncoder];
|
|
||||||
[score_command setComputePipelineState:gpu_data_->score_program];
|
|
||||||
[score_command setBuffer:gpu_data_->scored_boxes_buffer offset:0 atIndex:0];
|
|
||||||
[score_command setBuffer:gpu_data_->raw_scores_buffer offset:0 atIndex:1];
|
|
||||||
MTLSize score_threads_per_group = MTLSizeMake(1, num_classes_, 1);
|
|
||||||
MTLSize score_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
|
|
||||||
[score_command dispatchThreadgroups:score_threadgroups
|
|
||||||
threadsPerThreadgroup:score_threads_per_group];
|
threadsPerThreadgroup:score_threads_per_group];
|
||||||
[score_command endEncoding];
|
[command_encoder endEncoding];
|
||||||
[command_buffer commit];
|
[MPPMetalUtil commitCommandBufferAndWait:command_buffer];
|
||||||
[command_buffer waitUntilCompleted];
|
|
||||||
}
|
|
||||||
|
|
||||||
// Copy decoded boxes from GPU to CPU.
|
// Copy decoded boxes from GPU to CPU.
|
||||||
std::vector<float> boxes(num_boxes_ * num_coords_);
|
std::vector<float> boxes(num_boxes_ * num_coords_);
|
||||||
|
|||||||
@@ -28,6 +28,21 @@ namespace mediapipe {
|
|||||||
// TENSORS - Vector of TfLiteTensor of type kTfLiteFloat32. Only the first
|
// TENSORS - Vector of TfLiteTensor of type kTfLiteFloat32. Only the first
|
||||||
// tensor will be used. The size of the values must be
|
// tensor will be used. The size of the values must be
|
||||||
// (num_dimension x num_landmarks).
|
// (num_dimension x num_landmarks).
|
||||||
|
//
|
||||||
|
// FLIP_HORIZONTALLY (optional): Whether to flip landmarks horizontally or
|
||||||
|
// not. Overrides corresponding side packet and/or field in the calculator
|
||||||
|
// options.
|
||||||
|
//
|
||||||
|
// FLIP_VERTICALLY (optional): Whether to flip landmarks vertically or not.
|
||||||
|
// Overrides corresponding side packet and/or field in the calculator options.
|
||||||
|
//
|
||||||
|
// Input side packet:
|
||||||
|
// FLIP_HORIZONTALLY (optional): Whether to flip landmarks horizontally or
|
||||||
|
// not. Overrides the corresponding field in the calculator options.
|
||||||
|
//
|
||||||
|
// FLIP_VERTICALLY (optional): Whether to flip landmarks vertically or not.
|
||||||
|
// Overrides the corresponding field in the calculator options.
|
||||||
|
//
|
||||||
// Output:
|
// Output:
|
||||||
// LANDMARKS(optional) - Result MediaPipe landmarks.
|
// LANDMARKS(optional) - Result MediaPipe landmarks.
|
||||||
// NORM_LANDMARKS(optional) - Result MediaPipe normalized landmarks.
|
// NORM_LANDMARKS(optional) - Result MediaPipe normalized landmarks.
|
||||||
@@ -61,6 +76,8 @@ class TfLiteTensorsToLandmarksCalculator : public CalculatorBase {
|
|||||||
private:
|
private:
|
||||||
::mediapipe::Status LoadOptions(CalculatorContext* cc);
|
::mediapipe::Status LoadOptions(CalculatorContext* cc);
|
||||||
int num_landmarks_ = 0;
|
int num_landmarks_ = 0;
|
||||||
|
bool flip_vertically_ = false;
|
||||||
|
bool flip_horizontally_ = false;
|
||||||
|
|
||||||
::mediapipe::TfLiteTensorsToLandmarksCalculatorOptions options_;
|
::mediapipe::TfLiteTensorsToLandmarksCalculatorOptions options_;
|
||||||
};
|
};
|
||||||
@@ -75,6 +92,22 @@ REGISTER_CALCULATOR(TfLiteTensorsToLandmarksCalculator);
|
|||||||
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
cc->Inputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (cc->Inputs().HasTag("FLIP_HORIZONTALLY")) {
|
||||||
|
cc->Inputs().Tag("FLIP_HORIZONTALLY").Set<bool>();
|
||||||
|
}
|
||||||
|
|
||||||
|
if (cc->Inputs().HasTag("FLIP_VERTICALLY")) {
|
||||||
|
cc->Inputs().Tag("FLIP_VERTICALLY").Set<bool>();
|
||||||
|
}
|
||||||
|
|
||||||
|
if (cc->InputSidePackets().HasTag("FLIP_HORIZONTALLY")) {
|
||||||
|
cc->InputSidePackets().Tag("FLIP_HORIZONTALLY").Set<bool>();
|
||||||
|
}
|
||||||
|
|
||||||
|
if (cc->InputSidePackets().HasTag("FLIP_VERTICALLY")) {
|
||||||
|
cc->InputSidePackets().Tag("FLIP_VERTICALLY").Set<bool>();
|
||||||
|
}
|
||||||
|
|
||||||
if (cc->Outputs().HasTag("LANDMARKS")) {
|
if (cc->Outputs().HasTag("LANDMARKS")) {
|
||||||
cc->Outputs().Tag("LANDMARKS").Set<LandmarkList>();
|
cc->Outputs().Tag("LANDMARKS").Set<LandmarkList>();
|
||||||
}
|
}
|
||||||
@@ -98,17 +131,40 @@ REGISTER_CALCULATOR(TfLiteTensorsToLandmarksCalculator);
|
|||||||
<< "Must provide input with/height for getting normalized landmarks.";
|
<< "Must provide input with/height for getting normalized landmarks.";
|
||||||
}
|
}
|
||||||
if (cc->Outputs().HasTag("LANDMARKS") &&
|
if (cc->Outputs().HasTag("LANDMARKS") &&
|
||||||
(options_.flip_vertically() || options_.flip_horizontally())) {
|
(options_.flip_vertically() || options_.flip_horizontally() ||
|
||||||
|
cc->InputSidePackets().HasTag("FLIP_HORIZONTALLY") ||
|
||||||
|
cc->InputSidePackets().HasTag("FLIP_VERTICALLY"))) {
|
||||||
RET_CHECK(options_.has_input_image_height() &&
|
RET_CHECK(options_.has_input_image_height() &&
|
||||||
options_.has_input_image_width())
|
options_.has_input_image_width())
|
||||||
<< "Must provide input with/height for using flip_vertically option "
|
<< "Must provide input with/height for using flip_vertically option "
|
||||||
"when outputing landmarks in absolute coordinates.";
|
"when outputing landmarks in absolute coordinates.";
|
||||||
}
|
}
|
||||||
|
|
||||||
|
flip_horizontally_ =
|
||||||
|
cc->InputSidePackets().HasTag("FLIP_HORIZONTALLY")
|
||||||
|
? cc->InputSidePackets().Tag("FLIP_HORIZONTALLY").Get<bool>()
|
||||||
|
: options_.flip_horizontally();
|
||||||
|
|
||||||
|
flip_vertically_ =
|
||||||
|
cc->InputSidePackets().HasTag("FLIP_VERTICALLY")
|
||||||
|
? cc->InputSidePackets().Tag("FLIP_VERTICALLY").Get<bool>()
|
||||||
|
: options_.flip_vertically();
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status TfLiteTensorsToLandmarksCalculator::Process(
|
::mediapipe::Status TfLiteTensorsToLandmarksCalculator::Process(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
|
// Override values if specified so.
|
||||||
|
if (cc->Inputs().HasTag("FLIP_HORIZONTALLY") &&
|
||||||
|
!cc->Inputs().Tag("FLIP_HORIZONTALLY").IsEmpty()) {
|
||||||
|
flip_horizontally_ = cc->Inputs().Tag("FLIP_HORIZONTALLY").Get<bool>();
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag("FLIP_VERTICALLY") &&
|
||||||
|
!cc->Inputs().Tag("FLIP_VERTICALLY").IsEmpty()) {
|
||||||
|
flip_vertically_ = cc->Inputs().Tag("FLIP_VERTICALLY").Get<bool>();
|
||||||
|
}
|
||||||
|
|
||||||
if (cc->Inputs().Tag("TENSORS").IsEmpty()) {
|
if (cc->Inputs().Tag("TENSORS").IsEmpty()) {
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -133,13 +189,13 @@ REGISTER_CALCULATOR(TfLiteTensorsToLandmarksCalculator);
|
|||||||
const int offset = ld * num_dimensions;
|
const int offset = ld * num_dimensions;
|
||||||
Landmark* landmark = output_landmarks.add_landmark();
|
Landmark* landmark = output_landmarks.add_landmark();
|
||||||
|
|
||||||
if (options_.flip_horizontally()) {
|
if (flip_horizontally_) {
|
||||||
landmark->set_x(options_.input_image_width() - raw_landmarks[offset]);
|
landmark->set_x(options_.input_image_width() - raw_landmarks[offset]);
|
||||||
} else {
|
} else {
|
||||||
landmark->set_x(raw_landmarks[offset]);
|
landmark->set_x(raw_landmarks[offset]);
|
||||||
}
|
}
|
||||||
if (num_dimensions > 1) {
|
if (num_dimensions > 1) {
|
||||||
if (options_.flip_vertically()) {
|
if (flip_vertically_) {
|
||||||
landmark->set_y(options_.input_image_height() -
|
landmark->set_y(options_.input_image_height() -
|
||||||
raw_landmarks[offset + 1]);
|
raw_landmarks[offset + 1]);
|
||||||
} else {
|
} else {
|
||||||
|
|||||||
@@ -15,11 +15,11 @@
|
|||||||
#ifndef MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
|
#ifndef MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
|
||||||
#define MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
|
#define MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
|
||||||
|
|
||||||
#define RET_CHECK_CALL(call) \
|
#define RET_CHECK_CALL(call) \
|
||||||
do { \
|
do { \
|
||||||
const auto status = (call); \
|
const auto status = (call); \
|
||||||
if (ABSL_PREDICT_FALSE(!status.ok())) \
|
if (ABSL_PREDICT_FALSE(!status.ok())) \
|
||||||
return ::mediapipe::InternalError(status.error_message()); \
|
return ::mediapipe::InternalError(status.message()); \
|
||||||
} while (0);
|
} while (0);
|
||||||
|
|
||||||
#endif // MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
|
#endif // MEDIAPIPE_CALCULATORS_TFLITE_UTIL_H_
|
||||||
|
|||||||
@@ -39,6 +39,15 @@ proto_library(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
proto_library(
|
||||||
|
name = "timed_box_list_id_to_label_calculator_proto",
|
||||||
|
srcs = ["timed_box_list_id_to_label_calculator.proto"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_proto",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
proto_library(
|
proto_library(
|
||||||
name = "latency_proto",
|
name = "latency_proto",
|
||||||
srcs = ["latency.proto"],
|
srcs = ["latency.proto"],
|
||||||
@@ -113,6 +122,18 @@ mediapipe_cc_proto_library(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
mediapipe_cc_proto_library(
|
||||||
|
name = "timed_box_list_id_to_label_calculator_cc_proto",
|
||||||
|
srcs = ["timed_box_list_id_to_label_calculator.proto"],
|
||||||
|
cc_deps = [
|
||||||
|
"//mediapipe/framework:calculator_cc_proto",
|
||||||
|
],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
":timed_box_list_id_to_label_calculator_proto",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
mediapipe_cc_proto_library(
|
mediapipe_cc_proto_library(
|
||||||
name = "latency_cc_proto",
|
name = "latency_cc_proto",
|
||||||
srcs = ["latency.proto"],
|
srcs = ["latency.proto"],
|
||||||
@@ -300,7 +321,35 @@ cc_library(
|
|||||||
"//mediapipe:android": [
|
"//mediapipe:android": [
|
||||||
"//mediapipe/util/android/file/base",
|
"//mediapipe/util/android/file/base",
|
||||||
],
|
],
|
||||||
"//mediapipe:apple": [
|
"//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 = "timed_box_list_id_to_label_calculator",
|
||||||
|
srcs = ["timed_box_list_id_to_label_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
":timed_box_list_id_to_label_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:packet",
|
||||||
|
"//mediapipe/util/tracking:box_tracker_cc_proto",
|
||||||
|
"//mediapipe/util:resource_util",
|
||||||
|
] + select({
|
||||||
|
"//mediapipe:android": [
|
||||||
|
"//mediapipe/util/android/file/base",
|
||||||
|
],
|
||||||
|
"//mediapipe:ios": [
|
||||||
"//mediapipe/util/android/file/base",
|
"//mediapipe/util/android/file/base",
|
||||||
],
|
],
|
||||||
"//mediapipe:macos": [
|
"//mediapipe:macos": [
|
||||||
@@ -437,6 +486,7 @@ cc_library(
|
|||||||
"//mediapipe/framework/formats:rect_cc_proto",
|
"//mediapipe/framework/formats:rect_cc_proto",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
|
"@com_google_absl//absl/types:optional",
|
||||||
],
|
],
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -508,6 +558,17 @@ proto_library(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
proto_library(
|
||||||
|
name = "timed_box_list_to_render_data_calculator_proto",
|
||||||
|
srcs = ["timed_box_list_to_render_data_calculator.proto"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_proto",
|
||||||
|
"//mediapipe/util:color_proto",
|
||||||
|
"//mediapipe/util:render_data_proto",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
proto_library(
|
proto_library(
|
||||||
name = "labels_to_render_data_calculator_proto",
|
name = "labels_to_render_data_calculator_proto",
|
||||||
srcs = ["labels_to_render_data_calculator.proto"],
|
srcs = ["labels_to_render_data_calculator.proto"],
|
||||||
@@ -651,6 +712,37 @@ cc_library(
|
|||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
mediapipe_cc_proto_library(
|
||||||
|
name = "timed_box_list_to_render_data_calculator_cc_proto",
|
||||||
|
srcs = ["timed_box_list_to_render_data_calculator.proto"],
|
||||||
|
cc_deps = [
|
||||||
|
"//mediapipe/framework:calculator_cc_proto",
|
||||||
|
"//mediapipe/util:color_cc_proto",
|
||||||
|
"//mediapipe/util:render_data_cc_proto",
|
||||||
|
],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [":timed_box_list_to_render_data_calculator_proto"],
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "timed_box_list_to_render_data_calculator",
|
||||||
|
srcs = ["timed_box_list_to_render_data_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
":timed_box_list_to_render_data_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:calculator_options_cc_proto",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
|
"//mediapipe/util:color_cc_proto",
|
||||||
|
"//mediapipe/util:render_data_cc_proto",
|
||||||
|
"//mediapipe/util/tracking:box_tracker_cc_proto",
|
||||||
|
"//mediapipe/util/tracking:tracking_cc_proto",
|
||||||
|
"@com_google_absl//absl/memory",
|
||||||
|
"@com_google_absl//absl/strings",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
cc_library(
|
cc_library(
|
||||||
name = "labels_to_render_data_calculator",
|
name = "labels_to_render_data_calculator",
|
||||||
srcs = ["labels_to_render_data_calculator.cc"],
|
srcs = ["labels_to_render_data_calculator.cc"],
|
||||||
@@ -834,7 +926,7 @@ cc_library(
|
|||||||
"//mediapipe:android": [
|
"//mediapipe:android": [
|
||||||
"//mediapipe/util/android/file/base",
|
"//mediapipe/util/android/file/base",
|
||||||
],
|
],
|
||||||
"//mediapipe:apple": [
|
"//mediapipe:ios": [
|
||||||
"//mediapipe/util/android/file/base",
|
"//mediapipe/util/android/file/base",
|
||||||
],
|
],
|
||||||
"//mediapipe:macos": [
|
"//mediapipe:macos": [
|
||||||
@@ -877,6 +969,19 @@ cc_library(
|
|||||||
name = "local_file_contents_calculator",
|
name = "local_file_contents_calculator",
|
||||||
srcs = ["local_file_contents_calculator.cc"],
|
srcs = ["local_file_contents_calculator.cc"],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"//mediapipe/util:resource_util",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "local_file_pattern_contents_calculator",
|
||||||
|
srcs = ["local_file_pattern_contents_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/port:file_helpers",
|
"//mediapipe/framework/port:file_helpers",
|
||||||
@@ -893,6 +998,7 @@ cc_library(
|
|||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/formats:classification_cc_proto",
|
||||||
"//mediapipe/framework/formats:landmark_cc_proto",
|
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||||
"//mediapipe/framework/formats:rect_cc_proto",
|
"//mediapipe/framework/formats:rect_cc_proto",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
@@ -910,6 +1016,7 @@ cc_library(
|
|||||||
deps = [
|
deps = [
|
||||||
":collection_has_min_size_calculator_cc_proto",
|
":collection_has_min_size_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||||
"//mediapipe/framework/formats:rect_cc_proto",
|
"//mediapipe/framework/formats:rect_cc_proto",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
@@ -917,6 +1024,18 @@ cc_library(
|
|||||||
alwayslink = 1,
|
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(
|
cc_library(
|
||||||
name = "association_calculator",
|
name = "association_calculator",
|
||||||
hdrs = ["association_calculator.h"],
|
hdrs = ["association_calculator.h"],
|
||||||
@@ -983,3 +1102,31 @@ cc_test(
|
|||||||
"//mediapipe/framework/port:parse_text_proto",
|
"//mediapipe/framework/port:parse_text_proto",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "detections_to_timed_box_list_calculator",
|
||||||
|
srcs = ["detections_to_timed_box_list_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/formats:detection_cc_proto",
|
||||||
|
"//mediapipe/framework/formats:location_data_cc_proto",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"//mediapipe/util/tracking:box_tracker",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "detection_unique_id_calculator",
|
||||||
|
srcs = ["detection_unique_id_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/formats:detection_cc_proto",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|||||||
@@ -39,13 +39,13 @@ namespace mediapipe {
|
|||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
|
|
||||||
constexpr char kInputFrameTag[] = "INPUT_FRAME";
|
constexpr char kInputFrameTag[] = "IMAGE";
|
||||||
constexpr char kOutputFrameTag[] = "OUTPUT_FRAME";
|
constexpr char kOutputFrameTag[] = "IMAGE";
|
||||||
|
|
||||||
constexpr char kInputVectorTag[] = "VECTOR";
|
constexpr char kInputVectorTag[] = "VECTOR";
|
||||||
|
|
||||||
constexpr char kInputFrameTagGpu[] = "INPUT_FRAME_GPU";
|
constexpr char kInputFrameTagGpu[] = "IMAGE_GPU";
|
||||||
constexpr char kOutputFrameTagGpu[] = "OUTPUT_FRAME_GPU";
|
constexpr char kOutputFrameTagGpu[] = "IMAGE_GPU";
|
||||||
|
|
||||||
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
|
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
|
||||||
|
|
||||||
@@ -55,13 +55,13 @@ 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
|
// When using GPU, this color will become transparent when the calculator
|
||||||
// merges the annotation overlay with the image frame. As a result, drawing in
|
// 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.
|
// 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
|
} // namespace
|
||||||
|
|
||||||
// A calculator for rendering data on images.
|
// A calculator for rendering data on images.
|
||||||
//
|
//
|
||||||
// Inputs:
|
// Inputs:
|
||||||
// 1. INPUT_FRAME or INPUT_FRAME_GPU (optional): An ImageFrame (or GpuBuffer)
|
// 1. IMAGE or IMAGE_GPU (optional): An ImageFrame (or GpuBuffer)
|
||||||
// containing the input image.
|
// containing the input image.
|
||||||
// If output is CPU, and input isn't provided, the renderer creates a
|
// If output is CPU, and input isn't provided, the renderer creates a
|
||||||
// blank canvas with the width, height and color provided in the options.
|
// blank canvas with the width, height and color provided in the options.
|
||||||
@@ -73,7 +73,7 @@ constexpr int kAnnotationBackgroundColor[] = {100, 101, 102};
|
|||||||
// input vector items. These input streams are tagged with "VECTOR".
|
// input vector items. These input streams are tagged with "VECTOR".
|
||||||
//
|
//
|
||||||
// Output:
|
// Output:
|
||||||
// 1. OUTPUT_FRAME or OUTPUT_FRAME_GPU: A rendered ImageFrame (or GpuBuffer).
|
// 1. IMAGE or IMAGE_GPU: A rendered ImageFrame (or GpuBuffer).
|
||||||
//
|
//
|
||||||
// For CPU input frames, only SRGBA, SRGB and GRAY8 format are supported. The
|
// For CPU input frames, only SRGBA, SRGB and GRAY8 format are supported. The
|
||||||
// output format is the same as input except for GRAY8 where the output is in
|
// output format is the same as input except for GRAY8 where the output is in
|
||||||
@@ -87,13 +87,13 @@ constexpr int kAnnotationBackgroundColor[] = {100, 101, 102};
|
|||||||
// Example config (CPU):
|
// Example config (CPU):
|
||||||
// node {
|
// node {
|
||||||
// calculator: "AnnotationOverlayCalculator"
|
// calculator: "AnnotationOverlayCalculator"
|
||||||
// input_stream: "INPUT_FRAME:image_frames"
|
// input_stream: "IMAGE:image_frames"
|
||||||
// input_stream: "render_data_1"
|
// input_stream: "render_data_1"
|
||||||
// input_stream: "render_data_2"
|
// input_stream: "render_data_2"
|
||||||
// input_stream: "render_data_3"
|
// input_stream: "render_data_3"
|
||||||
// input_stream: "VECTOR:0:render_data_vec_0"
|
// input_stream: "VECTOR:0:render_data_vec_0"
|
||||||
// input_stream: "VECTOR:1:render_data_vec_1"
|
// input_stream: "VECTOR:1:render_data_vec_1"
|
||||||
// output_stream: "OUTPUT_FRAME:decorated_frames"
|
// output_stream: "IMAGE:decorated_frames"
|
||||||
// options {
|
// options {
|
||||||
// [mediapipe.AnnotationOverlayCalculatorOptions.ext] {
|
// [mediapipe.AnnotationOverlayCalculatorOptions.ext] {
|
||||||
// }
|
// }
|
||||||
@@ -103,13 +103,13 @@ constexpr int kAnnotationBackgroundColor[] = {100, 101, 102};
|
|||||||
// Example config (GPU):
|
// Example config (GPU):
|
||||||
// node {
|
// node {
|
||||||
// calculator: "AnnotationOverlayCalculator"
|
// calculator: "AnnotationOverlayCalculator"
|
||||||
// input_stream: "INPUT_FRAME_GPU:image_frames"
|
// input_stream: "IMAGE_GPU:image_frames"
|
||||||
// input_stream: "render_data_1"
|
// input_stream: "render_data_1"
|
||||||
// input_stream: "render_data_2"
|
// input_stream: "render_data_2"
|
||||||
// input_stream: "render_data_3"
|
// input_stream: "render_data_3"
|
||||||
// input_stream: "VECTOR:0:render_data_vec_0"
|
// input_stream: "VECTOR:0:render_data_vec_0"
|
||||||
// input_stream: "VECTOR:1:render_data_vec_1"
|
// input_stream: "VECTOR:1:render_data_vec_1"
|
||||||
// output_stream: "OUTPUT_FRAME_GPU:decorated_frames"
|
// output_stream: "IMAGE_GPU:decorated_frames"
|
||||||
// options {
|
// options {
|
||||||
// [mediapipe.AnnotationOverlayCalculatorOptions.ext] {
|
// [mediapipe.AnnotationOverlayCalculatorOptions.ext] {
|
||||||
// }
|
// }
|
||||||
@@ -491,11 +491,9 @@ REGISTER_CALCULATOR(AnnotationOverlayCalculator);
|
|||||||
if (format != mediapipe::ImageFormat::SRGBA &&
|
if (format != mediapipe::ImageFormat::SRGBA &&
|
||||||
format != mediapipe::ImageFormat::SRGB)
|
format != mediapipe::ImageFormat::SRGB)
|
||||||
RET_CHECK_FAIL() << "Unsupported GPU input format: " << format;
|
RET_CHECK_FAIL() << "Unsupported GPU input format: " << format;
|
||||||
|
image_mat = absl::make_unique<cv::Mat>(height_, width_, CV_8UC3);
|
||||||
image_mat = absl::make_unique<cv::Mat>(
|
memset(image_mat->data, kAnnotationBackgroundColor,
|
||||||
height_, width_, CV_8UC3,
|
height_ * width_ * image_mat->elemSize());
|
||||||
cv::Scalar(kAnnotationBackgroundColor[0], kAnnotationBackgroundColor[1],
|
|
||||||
kAnnotationBackgroundColor[2]));
|
|
||||||
} else {
|
} else {
|
||||||
image_mat = absl::make_unique<cv::Mat>(
|
image_mat = absl::make_unique<cv::Mat>(
|
||||||
options_.canvas_height_px(), options_.canvas_width_px(), CV_8UC3,
|
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_, "input_frame"), 1);
|
||||||
glUniform1i(glGetUniformLocation(program_, "overlay"), 2);
|
glUniform1i(glGetUniformLocation(program_, "overlay"), 2);
|
||||||
glUniform3f(glGetUniformLocation(program_, "transparent_color"),
|
glUniform3f(glGetUniformLocation(program_, "transparent_color"),
|
||||||
kAnnotationBackgroundColor[0] / 255.0,
|
kAnnotationBackgroundColor / 255.0,
|
||||||
kAnnotationBackgroundColor[1] / 255.0,
|
kAnnotationBackgroundColor / 255.0,
|
||||||
kAnnotationBackgroundColor[2] / 255.0);
|
kAnnotationBackgroundColor / 255.0);
|
||||||
|
|
||||||
// Init texture for opencv rendered frame.
|
// Init texture for opencv rendered frame.
|
||||||
const auto& input_frame =
|
const auto& input_frame =
|
||||||
|
|||||||
@@ -15,6 +15,9 @@
|
|||||||
|
|
||||||
#include "mediapipe/calculators/util/collection_has_min_size_calculator.h"
|
#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"
|
#include "mediapipe/framework/formats/rect.pb.h"
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
@@ -23,4 +26,9 @@ typedef CollectionHasMinSizeCalculator<std::vector<::mediapipe::NormalizedRect>>
|
|||||||
NormalizedRectVectorHasMinSizeCalculator;
|
NormalizedRectVectorHasMinSizeCalculator;
|
||||||
REGISTER_CALCULATOR(NormalizedRectVectorHasMinSizeCalculator);
|
REGISTER_CALCULATOR(NormalizedRectVectorHasMinSizeCalculator);
|
||||||
|
|
||||||
|
typedef CollectionHasMinSizeCalculator<
|
||||||
|
std::vector<::mediapipe::NormalizedLandmarkList>>
|
||||||
|
NormalizedLandmarkListVectorHasMinSizeCalculator;
|
||||||
|
REGISTER_CALCULATOR(NormalizedLandmarkListVectorHasMinSizeCalculator);
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -37,6 +37,8 @@ namespace mediapipe {
|
|||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
|
// Optionally, uses a side packet to override `min_size` specified in the
|
||||||
|
// calculator options.
|
||||||
template <typename IterableT>
|
template <typename IterableT>
|
||||||
class CollectionHasMinSizeCalculator : public CalculatorBase {
|
class CollectionHasMinSizeCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
@@ -54,6 +56,10 @@ class CollectionHasMinSizeCalculator : public CalculatorBase {
|
|||||||
cc->Inputs().Tag("ITERABLE").Set<IterableT>();
|
cc->Inputs().Tag("ITERABLE").Set<IterableT>();
|
||||||
cc->Outputs().Index(0).Set<bool>();
|
cc->Outputs().Index(0).Set<bool>();
|
||||||
|
|
||||||
|
// Optional input side packet that determines `min_size_`.
|
||||||
|
if (cc->InputSidePackets().NumEntries() > 0) {
|
||||||
|
cc->InputSidePackets().Index(0).Set<int>();
|
||||||
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -62,6 +68,11 @@ class CollectionHasMinSizeCalculator : public CalculatorBase {
|
|||||||
min_size_ =
|
min_size_ =
|
||||||
cc->Options<::mediapipe::CollectionHasMinSizeCalculatorOptions>()
|
cc->Options<::mediapipe::CollectionHasMinSizeCalculatorOptions>()
|
||||||
.min_size();
|
.min_size();
|
||||||
|
// Override `min_size` if passed as side packet.
|
||||||
|
if (cc->InputSidePackets().NumEntries() > 0 &&
|
||||||
|
!cc->InputSidePackets().Index(0).IsEmpty()) {
|
||||||
|
min_size_ = cc->InputSidePackets().Index(0).Get<int>();
|
||||||
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -12,10 +12,10 @@
|
|||||||
// See the License for the specific language governing permissions and
|
// See the License for the specific language governing permissions and
|
||||||
// limitations under the License.
|
// limitations under the License.
|
||||||
|
|
||||||
#include "mediapipe//framework/packet.h"
|
|
||||||
#include "mediapipe/calculators/util/detection_label_id_to_text_calculator.pb.h"
|
#include "mediapipe/calculators/util/detection_label_id_to_text_calculator.pb.h"
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/formats/detection.pb.h"
|
#include "mediapipe/framework/formats/detection.pb.h"
|
||||||
|
#include "mediapipe/framework/packet.h"
|
||||||
#include "mediapipe/framework/port/status.h"
|
#include "mediapipe/framework/port/status.h"
|
||||||
#include "mediapipe/util/resource_util.h"
|
#include "mediapipe/util/resource_util.h"
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,110 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/formats/detection.pb.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
constexpr char kDetectionsTag[] = "DETECTIONS";
|
||||||
|
constexpr char kDetectionListTag[] = "DETECTION_LIST";
|
||||||
|
|
||||||
|
// Each detection processed by DetectionUniqueIDCalculator will be assigned an
|
||||||
|
// unique id that starts from 1. If a detection already has an ID other than 0,
|
||||||
|
// the ID will be overwritten.
|
||||||
|
static int64 detection_id = 0;
|
||||||
|
|
||||||
|
inline int GetNextDetectionId() { return ++detection_id; }
|
||||||
|
|
||||||
|
} // namespace
|
||||||
|
|
||||||
|
// Assign a unique id to detections.
|
||||||
|
// Note that the calculator will consume the input vector of Detection or
|
||||||
|
// DetectionList. So the input stream can not be connected to other calculators.
|
||||||
|
//
|
||||||
|
// Example config:
|
||||||
|
// node {
|
||||||
|
// calculator: "DetectionUniqueIdCalculator"
|
||||||
|
// input_stream: "DETECTIONS:detections"
|
||||||
|
// output_stream: "DETECTIONS:output_detections"
|
||||||
|
// }
|
||||||
|
class DetectionUniqueIdCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
RET_CHECK(cc->Inputs().HasTag(kDetectionListTag) ||
|
||||||
|
cc->Inputs().HasTag(kDetectionsTag))
|
||||||
|
<< "None of the input streams are provided.";
|
||||||
|
|
||||||
|
if (cc->Inputs().HasTag(kDetectionListTag)) {
|
||||||
|
RET_CHECK(cc->Outputs().HasTag(kDetectionListTag));
|
||||||
|
cc->Inputs().Tag(kDetectionListTag).Set<DetectionList>();
|
||||||
|
cc->Outputs().Tag(kDetectionListTag).Set<DetectionList>();
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag(kDetectionsTag)) {
|
||||||
|
RET_CHECK(cc->Outputs().HasTag(kDetectionsTag));
|
||||||
|
cc->Inputs().Tag(kDetectionsTag).Set<std::vector<Detection>>();
|
||||||
|
cc->Outputs().Tag(kDetectionsTag).Set<std::vector<Detection>>();
|
||||||
|
}
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
|
cc->SetOffset(::mediapipe::TimestampDiff(0));
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||||
|
};
|
||||||
|
REGISTER_CALCULATOR(DetectionUniqueIdCalculator);
|
||||||
|
|
||||||
|
::mediapipe::Status DetectionUniqueIdCalculator::Process(
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
if (cc->Inputs().HasTag(kDetectionListTag) &&
|
||||||
|
!cc->Inputs().Tag(kDetectionListTag).IsEmpty()) {
|
||||||
|
auto result =
|
||||||
|
cc->Inputs().Tag(kDetectionListTag).Value().Consume<DetectionList>();
|
||||||
|
if (result.ok()) {
|
||||||
|
auto detection_list = std::move(result).ValueOrDie();
|
||||||
|
for (Detection& detection : *detection_list->mutable_detection()) {
|
||||||
|
detection.set_detection_id(GetNextDetectionId());
|
||||||
|
}
|
||||||
|
cc->Outputs()
|
||||||
|
.Tag(kDetectionListTag)
|
||||||
|
.Add(detection_list.release(), cc->InputTimestamp());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if (cc->Inputs().HasTag(kDetectionsTag) &&
|
||||||
|
!cc->Inputs().Tag(kDetectionsTag).IsEmpty()) {
|
||||||
|
auto result = cc->Inputs()
|
||||||
|
.Tag(kDetectionsTag)
|
||||||
|
.Value()
|
||||||
|
.Consume<std::vector<Detection>>();
|
||||||
|
if (result.ok()) {
|
||||||
|
auto detections = std::move(result).ValueOrDie();
|
||||||
|
for (Detection& detection : *detections) {
|
||||||
|
detection.set_detection_id(GetNextDetectionId());
|
||||||
|
}
|
||||||
|
cc->Outputs()
|
||||||
|
.Tag(kDetectionsTag)
|
||||||
|
.Add(detections.release(), cc->InputTimestamp());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -39,7 +39,8 @@ constexpr char kNormRectsTag[] = "NORM_RECTS";
|
|||||||
} // namespace
|
} // namespace
|
||||||
|
|
||||||
::mediapipe::Status DetectionsToRectsCalculator::DetectionToRect(
|
::mediapipe::Status DetectionsToRectsCalculator::DetectionToRect(
|
||||||
const Detection& detection, Rect* rect) {
|
const Detection& detection, const DetectionSpec& detection_spec,
|
||||||
|
Rect* rect) {
|
||||||
const LocationData location_data = detection.location_data();
|
const LocationData location_data = detection.location_data();
|
||||||
RET_CHECK(location_data.format() == LocationData::BOUNDING_BOX)
|
RET_CHECK(location_data.format() == LocationData::BOUNDING_BOX)
|
||||||
<< "Only Detection with formats of BOUNDING_BOX can be converted to Rect";
|
<< "Only Detection with formats of BOUNDING_BOX can be converted to Rect";
|
||||||
@@ -52,7 +53,8 @@ constexpr char kNormRectsTag[] = "NORM_RECTS";
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status DetectionsToRectsCalculator::DetectionToNormalizedRect(
|
::mediapipe::Status DetectionsToRectsCalculator::DetectionToNormalizedRect(
|
||||||
const Detection& detection, NormalizedRect* rect) {
|
const Detection& detection, const DetectionSpec& detection_spec,
|
||||||
|
NormalizedRect* rect) {
|
||||||
const LocationData location_data = detection.location_data();
|
const LocationData location_data = detection.location_data();
|
||||||
RET_CHECK(location_data.format() == LocationData::RELATIVE_BOUNDING_BOX)
|
RET_CHECK(location_data.format() == LocationData::RELATIVE_BOUNDING_BOX)
|
||||||
<< "Only Detection with formats of RELATIVE_BOUNDING_BOX can be "
|
<< "Only Detection with formats of RELATIVE_BOUNDING_BOX can be "
|
||||||
@@ -174,27 +176,31 @@ constexpr char kNormRectsTag[] = "NORM_RECTS";
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
std::pair<int, int> image_size;
|
// Get dynamic calculator options (e.g. `image_size`).
|
||||||
if (rotate_) {
|
const DetectionSpec detection_spec = GetDetectionSpec(cc);
|
||||||
RET_CHECK(!cc->Inputs().Tag(kImageSizeTag).IsEmpty());
|
|
||||||
image_size = cc->Inputs().Tag(kImageSizeTag).Get<std::pair<int, int>>();
|
|
||||||
}
|
|
||||||
|
|
||||||
if (cc->Outputs().HasTag(kRectTag)) {
|
if (cc->Outputs().HasTag(kRectTag)) {
|
||||||
auto output_rect = absl::make_unique<Rect>();
|
auto output_rect = absl::make_unique<Rect>();
|
||||||
MP_RETURN_IF_ERROR(DetectionToRect(detections[0], output_rect.get()));
|
MP_RETURN_IF_ERROR(
|
||||||
|
DetectionToRect(detections[0], detection_spec, output_rect.get()));
|
||||||
if (rotate_) {
|
if (rotate_) {
|
||||||
output_rect->set_rotation(ComputeRotation(detections[0], image_size));
|
float rotation;
|
||||||
|
MP_RETURN_IF_ERROR(
|
||||||
|
ComputeRotation(detections[0], detection_spec, &rotation));
|
||||||
|
output_rect->set_rotation(rotation);
|
||||||
}
|
}
|
||||||
cc->Outputs().Tag(kRectTag).Add(output_rect.release(),
|
cc->Outputs().Tag(kRectTag).Add(output_rect.release(),
|
||||||
cc->InputTimestamp());
|
cc->InputTimestamp());
|
||||||
}
|
}
|
||||||
if (cc->Outputs().HasTag(kNormRectTag)) {
|
if (cc->Outputs().HasTag(kNormRectTag)) {
|
||||||
auto output_rect = absl::make_unique<NormalizedRect>();
|
auto output_rect = absl::make_unique<NormalizedRect>();
|
||||||
MP_RETURN_IF_ERROR(
|
MP_RETURN_IF_ERROR(DetectionToNormalizedRect(detections[0], detection_spec,
|
||||||
DetectionToNormalizedRect(detections[0], output_rect.get()));
|
output_rect.get()));
|
||||||
if (rotate_) {
|
if (rotate_) {
|
||||||
output_rect->set_rotation(ComputeRotation(detections[0], image_size));
|
float rotation;
|
||||||
|
MP_RETURN_IF_ERROR(
|
||||||
|
ComputeRotation(detections[0], detection_spec, &rotation));
|
||||||
|
output_rect->set_rotation(rotation);
|
||||||
}
|
}
|
||||||
cc->Outputs()
|
cc->Outputs()
|
||||||
.Tag(kNormRectTag)
|
.Tag(kNormRectTag)
|
||||||
@@ -203,11 +209,13 @@ constexpr char kNormRectsTag[] = "NORM_RECTS";
|
|||||||
if (cc->Outputs().HasTag(kRectsTag)) {
|
if (cc->Outputs().HasTag(kRectsTag)) {
|
||||||
auto output_rects = absl::make_unique<std::vector<Rect>>(detections.size());
|
auto output_rects = absl::make_unique<std::vector<Rect>>(detections.size());
|
||||||
for (int i = 0; i < detections.size(); ++i) {
|
for (int i = 0; i < detections.size(); ++i) {
|
||||||
MP_RETURN_IF_ERROR(
|
MP_RETURN_IF_ERROR(DetectionToRect(detections[i], detection_spec,
|
||||||
DetectionToRect(detections[i], &(output_rects->at(i))));
|
&(output_rects->at(i))));
|
||||||
if (rotate_) {
|
if (rotate_) {
|
||||||
output_rects->at(i).set_rotation(
|
float rotation;
|
||||||
ComputeRotation(detections[i], image_size));
|
MP_RETURN_IF_ERROR(
|
||||||
|
ComputeRotation(detections[i], detection_spec, &rotation));
|
||||||
|
output_rects->at(i).set_rotation(rotation);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
cc->Outputs().Tag(kRectsTag).Add(output_rects.release(),
|
cc->Outputs().Tag(kRectsTag).Add(output_rects.release(),
|
||||||
@@ -217,11 +225,13 @@ constexpr char kNormRectsTag[] = "NORM_RECTS";
|
|||||||
auto output_rects =
|
auto output_rects =
|
||||||
absl::make_unique<std::vector<NormalizedRect>>(detections.size());
|
absl::make_unique<std::vector<NormalizedRect>>(detections.size());
|
||||||
for (int i = 0; i < detections.size(); ++i) {
|
for (int i = 0; i < detections.size(); ++i) {
|
||||||
MP_RETURN_IF_ERROR(
|
MP_RETURN_IF_ERROR(DetectionToNormalizedRect(
|
||||||
DetectionToNormalizedRect(detections[i], &(output_rects->at(i))));
|
detections[i], detection_spec, &(output_rects->at(i))));
|
||||||
if (rotate_) {
|
if (rotate_) {
|
||||||
output_rects->at(i).set_rotation(
|
float rotation;
|
||||||
ComputeRotation(detections[i], image_size));
|
MP_RETURN_IF_ERROR(
|
||||||
|
ComputeRotation(detections[i], detection_spec, &rotation));
|
||||||
|
output_rects->at(i).set_rotation(rotation);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
cc->Outputs()
|
cc->Outputs()
|
||||||
@@ -232,21 +242,35 @@ constexpr char kNormRectsTag[] = "NORM_RECTS";
|
|||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
float DetectionsToRectsCalculator::ComputeRotation(
|
::mediapipe::Status DetectionsToRectsCalculator::ComputeRotation(
|
||||||
const Detection& detection, const std::pair<int, int> image_size) {
|
const Detection& detection, const DetectionSpec& detection_spec,
|
||||||
|
float* rotation) {
|
||||||
const auto& location_data = detection.location_data();
|
const auto& location_data = detection.location_data();
|
||||||
|
const auto& image_size = detection_spec.image_size;
|
||||||
|
RET_CHECK(image_size) << "Image size is required to calculate rotation";
|
||||||
|
|
||||||
const float x0 = location_data.relative_keypoints(start_keypoint_index_).x() *
|
const float x0 = location_data.relative_keypoints(start_keypoint_index_).x() *
|
||||||
image_size.first;
|
image_size->first;
|
||||||
const float y0 = location_data.relative_keypoints(start_keypoint_index_).y() *
|
const float y0 = location_data.relative_keypoints(start_keypoint_index_).y() *
|
||||||
image_size.second;
|
image_size->second;
|
||||||
const float x1 = location_data.relative_keypoints(end_keypoint_index_).x() *
|
const float x1 = location_data.relative_keypoints(end_keypoint_index_).x() *
|
||||||
image_size.first;
|
image_size->first;
|
||||||
const float y1 = location_data.relative_keypoints(end_keypoint_index_).y() *
|
const float y1 = location_data.relative_keypoints(end_keypoint_index_).y() *
|
||||||
image_size.second;
|
image_size->second;
|
||||||
|
|
||||||
float rotation = target_angle_ - std::atan2(-(y1 - y0), x1 - x0);
|
*rotation = NormalizeRadians(target_angle_ - std::atan2(-(y1 - y0), x1 - x0));
|
||||||
|
|
||||||
return NormalizeRadians(rotation);
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
DetectionSpec DetectionsToRectsCalculator::GetDetectionSpec(
|
||||||
|
const CalculatorContext* cc) {
|
||||||
|
absl::optional<std::pair<int, int>> image_size;
|
||||||
|
if (cc->Inputs().HasTag(kImageSizeTag)) {
|
||||||
|
image_size = cc->Inputs().Tag(kImageSizeTag).Get<std::pair<int, int>>();
|
||||||
|
}
|
||||||
|
|
||||||
|
return {image_size};
|
||||||
}
|
}
|
||||||
|
|
||||||
REGISTER_CALCULATOR(DetectionsToRectsCalculator);
|
REGISTER_CALCULATOR(DetectionsToRectsCalculator);
|
||||||
|
|||||||
@@ -16,6 +16,7 @@
|
|||||||
|
|
||||||
#include <cmath>
|
#include <cmath>
|
||||||
|
|
||||||
|
#include "absl/types/optional.h"
|
||||||
#include "mediapipe/calculators/util/detections_to_rects_calculator.pb.h"
|
#include "mediapipe/calculators/util/detections_to_rects_calculator.pb.h"
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/calculator_options.pb.h"
|
#include "mediapipe/framework/calculator_options.pb.h"
|
||||||
@@ -27,6 +28,13 @@
|
|||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
|
// Dynamic options passed as calculator `input_stream` that can be used for
|
||||||
|
// calculation of rectangle or rotation for given detection. Does not include
|
||||||
|
// static calculator options which are available via private fields.
|
||||||
|
struct DetectionSpec {
|
||||||
|
absl::optional<std::pair<int, int>> image_size;
|
||||||
|
};
|
||||||
|
|
||||||
// A calculator that converts Detection proto to Rect proto.
|
// A calculator that converts Detection proto to Rect proto.
|
||||||
//
|
//
|
||||||
// Detection is the format for encoding one or more detections in an image.
|
// Detection is the format for encoding one or more detections in an image.
|
||||||
@@ -81,13 +89,16 @@ class DetectionsToRectsCalculator : public CalculatorBase {
|
|||||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||||
|
|
||||||
protected:
|
protected:
|
||||||
virtual float ComputeRotation(const ::mediapipe::Detection& detection,
|
|
||||||
const std::pair<int, int> image_size);
|
|
||||||
virtual ::mediapipe::Status DetectionToRect(
|
virtual ::mediapipe::Status DetectionToRect(
|
||||||
const ::mediapipe::Detection& detection, ::mediapipe::Rect* rect);
|
const ::mediapipe::Detection& detection,
|
||||||
|
const DetectionSpec& detection_spec, ::mediapipe::Rect* rect);
|
||||||
virtual ::mediapipe::Status DetectionToNormalizedRect(
|
virtual ::mediapipe::Status DetectionToNormalizedRect(
|
||||||
const ::mediapipe::Detection& detection,
|
const ::mediapipe::Detection& detection,
|
||||||
::mediapipe::NormalizedRect* rect);
|
const DetectionSpec& detection_spec, ::mediapipe::NormalizedRect* rect);
|
||||||
|
virtual ::mediapipe::Status ComputeRotation(
|
||||||
|
const ::mediapipe::Detection& detection,
|
||||||
|
const DetectionSpec& detection_spec, float* rotation);
|
||||||
|
virtual DetectionSpec GetDetectionSpec(const CalculatorContext* cc);
|
||||||
|
|
||||||
static inline float NormalizeRadians(float angle) {
|
static inline float NormalizeRadians(float angle) {
|
||||||
return angle - 2 * M_PI * std::floor((angle - (-M_PI)) / (2 * M_PI));
|
return angle - 2 * M_PI * std::floor((angle - (-M_PI)) / (2 * M_PI));
|
||||||
|
|||||||
@@ -39,6 +39,8 @@ constexpr char kKeypointLabel[] = "KEYPOINT";
|
|||||||
// The ratio of detection label font height to the height of detection bounding
|
// The ratio of detection label font height to the height of detection bounding
|
||||||
// box.
|
// box.
|
||||||
constexpr double kLabelToBoundingBoxRatio = 0.1;
|
constexpr double kLabelToBoundingBoxRatio = 0.1;
|
||||||
|
// Perserve 2 decimal digits.
|
||||||
|
constexpr float kNumScoreDecimalDigitsMultipler = 100;
|
||||||
|
|
||||||
} // namespace
|
} // namespace
|
||||||
|
|
||||||
@@ -235,18 +237,26 @@ void DetectionsToRenderDataCalculator::AddLabels(
|
|||||||
std::string label_str = detection.label().empty()
|
std::string label_str = detection.label().empty()
|
||||||
? absl::StrCat(detection.label_id(i))
|
? absl::StrCat(detection.label_id(i))
|
||||||
: detection.label(i);
|
: detection.label(i);
|
||||||
|
const float rounded_score =
|
||||||
|
std::round(detection.score(i) * kNumScoreDecimalDigitsMultipler) /
|
||||||
|
kNumScoreDecimalDigitsMultipler;
|
||||||
std::string label_and_score =
|
std::string label_and_score =
|
||||||
absl::StrCat(label_str, options.text_delimiter(), detection.score(i),
|
absl::StrCat(label_str, options.text_delimiter(), rounded_score,
|
||||||
options.text_delimiter());
|
options.text_delimiter());
|
||||||
label_and_scores.push_back(label_and_score);
|
label_and_scores.push_back(label_and_score);
|
||||||
}
|
}
|
||||||
std::vector<std::string> labels;
|
std::vector<std::string> labels;
|
||||||
|
if (options.render_detection_id()) {
|
||||||
|
const std::string detection_id_str =
|
||||||
|
absl::StrCat("Id: ", detection.detection_id());
|
||||||
|
labels.push_back(detection_id_str);
|
||||||
|
}
|
||||||
if (options.one_label_per_line()) {
|
if (options.one_label_per_line()) {
|
||||||
labels.swap(label_and_scores);
|
labels.insert(labels.end(), label_and_scores.begin(),
|
||||||
|
label_and_scores.end());
|
||||||
} else {
|
} else {
|
||||||
labels.push_back(absl::StrJoin(label_and_scores, ""));
|
labels.push_back(absl::StrJoin(label_and_scores, ""));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Add the render annotations for "label(_id),score".
|
// Add the render annotations for "label(_id),score".
|
||||||
for (int i = 0; i < labels.size(); ++i) {
|
for (int i = 0; i < labels.size(); ++i) {
|
||||||
auto label = labels.at(i);
|
auto label = labels.at(i);
|
||||||
|
|||||||
@@ -53,4 +53,7 @@ message DetectionsToRenderDataCalculatorOptions {
|
|||||||
// instances of this calculator are present in the graph, this value
|
// instances of this calculator are present in the graph, this value
|
||||||
// should be unique among them.
|
// should be unique among them.
|
||||||
optional string scene_class = 7 [default = "DETECTION"];
|
optional string scene_class = 7 [default = "DETECTION"];
|
||||||
|
|
||||||
|
// If true, renders the detection id in the first line before the labels.
|
||||||
|
optional bool render_detection_id = 8 [default = false];
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,110 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/formats/detection.pb.h"
|
||||||
|
#include "mediapipe/framework/formats/location_data.pb.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
#include "mediapipe/util/tracking/box_tracker.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
constexpr char kDetectionsTag[] = "DETECTIONS";
|
||||||
|
constexpr char kDetectionListTag[] = "DETECTION_LIST";
|
||||||
|
constexpr char kBoxesTag[] = "BOXES";
|
||||||
|
|
||||||
|
} // namespace
|
||||||
|
|
||||||
|
// A calculator that converts Detection proto to TimedBoxList proto for
|
||||||
|
// tracking.
|
||||||
|
//
|
||||||
|
// Please note that only Location Data formats of RELATIVE_BOUNDING_BOX are
|
||||||
|
// supported.
|
||||||
|
//
|
||||||
|
// Example config:
|
||||||
|
// node {
|
||||||
|
// calculator: "DetectionsToTimedBoxListCalculator"
|
||||||
|
// input_stream: "DETECTIONS:detections"
|
||||||
|
// output_stream: "BOXES:boxes"
|
||||||
|
// }
|
||||||
|
class DetectionsToTimedBoxListCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
RET_CHECK(cc->Inputs().HasTag(kDetectionListTag) ||
|
||||||
|
cc->Inputs().HasTag(kDetectionsTag))
|
||||||
|
<< "None of the input streams are provided.";
|
||||||
|
if (cc->Inputs().HasTag(kDetectionListTag)) {
|
||||||
|
cc->Inputs().Tag(kDetectionListTag).Set<DetectionList>();
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag(kDetectionsTag)) {
|
||||||
|
cc->Inputs().Tag(kDetectionsTag).Set<std::vector<Detection>>();
|
||||||
|
}
|
||||||
|
cc->Outputs().Tag(kBoxesTag).Set<TimedBoxProtoList>();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
|
cc->SetOffset(TimestampDiff(0));
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||||
|
|
||||||
|
private:
|
||||||
|
void ConvertDetectionToTimedBox(const Detection& detection,
|
||||||
|
TimedBoxProto* box, CalculatorContext* cc);
|
||||||
|
};
|
||||||
|
REGISTER_CALCULATOR(DetectionsToTimedBoxListCalculator);
|
||||||
|
|
||||||
|
::mediapipe::Status DetectionsToTimedBoxListCalculator::Process(
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
auto output_timed_box_list = absl::make_unique<TimedBoxProtoList>();
|
||||||
|
|
||||||
|
if (cc->Inputs().HasTag(kDetectionListTag)) {
|
||||||
|
const auto& detection_list =
|
||||||
|
cc->Inputs().Tag(kDetectionListTag).Get<DetectionList>();
|
||||||
|
for (const auto& detection : detection_list.detection()) {
|
||||||
|
TimedBoxProto* box = output_timed_box_list->add_box();
|
||||||
|
ConvertDetectionToTimedBox(detection, box, cc);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag(kDetectionsTag)) {
|
||||||
|
const auto& detections =
|
||||||
|
cc->Inputs().Tag(kDetectionsTag).Get<std::vector<Detection>>();
|
||||||
|
for (const auto& detection : detections) {
|
||||||
|
TimedBoxProto* box = output_timed_box_list->add_box();
|
||||||
|
ConvertDetectionToTimedBox(detection, box, cc);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
cc->Outputs().Tag(kBoxesTag).Add(output_timed_box_list.release(),
|
||||||
|
cc->InputTimestamp());
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
void DetectionsToTimedBoxListCalculator::ConvertDetectionToTimedBox(
|
||||||
|
const Detection& detection, TimedBoxProto* box, CalculatorContext* cc) {
|
||||||
|
const auto& relative_bounding_box =
|
||||||
|
detection.location_data().relative_bounding_box();
|
||||||
|
box->set_left(relative_bounding_box.xmin());
|
||||||
|
box->set_right(relative_bounding_box.xmin() + relative_bounding_box.width());
|
||||||
|
box->set_top(relative_bounding_box.ymin());
|
||||||
|
box->set_bottom(relative_bounding_box.ymin() +
|
||||||
|
relative_bounding_box.height());
|
||||||
|
box->set_id(detection.detection_id());
|
||||||
|
box->set_time_msec(cc->InputTimestamp().Microseconds() / 1000);
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -17,6 +17,7 @@
|
|||||||
|
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/formats/classification.pb.h"
|
||||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||||
#include "mediapipe/framework/formats/rect.pb.h"
|
#include "mediapipe/framework/formats/rect.pb.h"
|
||||||
|
|
||||||
@@ -31,4 +32,8 @@ typedef FilterCollectionCalculator<
|
|||||||
FilterLandmarkListCollectionCalculator;
|
FilterLandmarkListCollectionCalculator;
|
||||||
REGISTER_CALCULATOR(FilterLandmarkListCollectionCalculator);
|
REGISTER_CALCULATOR(FilterLandmarkListCollectionCalculator);
|
||||||
|
|
||||||
|
typedef FilterCollectionCalculator<std::vector<::mediapipe::ClassificationList>>
|
||||||
|
FilterClassificationListCollectionCalculator;
|
||||||
|
REGISTER_CALCULATOR(FilterClassificationListCollectionCalculator);
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -128,16 +128,19 @@ REGISTER_CALCULATOR(LabelsToRenderDataCalculator);
|
|||||||
} else {
|
} else {
|
||||||
const std::vector<std::string>& label_vector =
|
const std::vector<std::string>& label_vector =
|
||||||
cc->Inputs().Tag("LABELS").Get<std::vector<std::string>>();
|
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());
|
labels.resize(label_vector.size());
|
||||||
scores.resize(label_vector.size());
|
|
||||||
for (int i = 0; i < label_vector.size(); ++i) {
|
for (int i = 0; i < label_vector.size(); ++i) {
|
||||||
labels[i] = label_vector[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];
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -12,20 +12,6 @@
|
|||||||
// See the License for the specific language governing permissions and
|
// See the License for the specific language governing permissions and
|
||||||
// limitations under the License.
|
// limitations under the License.
|
||||||
|
|
||||||
// Copyright 2019 The MediaPipe Authors.
|
|
||||||
//
|
|
||||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
||||||
// you may not use this file except in compliance with the License.
|
|
||||||
// You may obtain a copy of the License at
|
|
||||||
//
|
|
||||||
// http://www.apache.org/licenses/LICENSE-2.0
|
|
||||||
//
|
|
||||||
// Unless required by applicable law or agreed to in writing, software
|
|
||||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
||||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
||||||
// See the License for the specific language governing permissions and
|
|
||||||
// limitations under the License.
|
|
||||||
|
|
||||||
#include <cmath>
|
#include <cmath>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
@@ -67,6 +53,15 @@ constexpr char kLetterboxPaddingTag[] = "LETTERBOX_PADDING";
|
|||||||
// input_stream: "LETTERBOX_PADDING:letterbox_padding"
|
// input_stream: "LETTERBOX_PADDING:letterbox_padding"
|
||||||
// output_stream: "LANDMARKS:adjusted_landmarks"
|
// output_stream: "LANDMARKS:adjusted_landmarks"
|
||||||
// }
|
// }
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "LandmarkLetterboxRemovalCalculator"
|
||||||
|
// input_stream: "LANDMARKS:0:landmarks_0"
|
||||||
|
// input_stream: "LANDMARKS:1:landmarks_1"
|
||||||
|
// input_stream: "LETTERBOX_PADDING:letterbox_padding"
|
||||||
|
// output_stream: "LANDMARKS:0:adjusted_landmarks_0"
|
||||||
|
// output_stream: "LANDMARKS:1:adjusted_landmarks_1"
|
||||||
|
// }
|
||||||
class LandmarkLetterboxRemovalCalculator : public CalculatorBase {
|
class LandmarkLetterboxRemovalCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
@@ -74,10 +69,20 @@ class LandmarkLetterboxRemovalCalculator : public CalculatorBase {
|
|||||||
cc->Inputs().HasTag(kLetterboxPaddingTag))
|
cc->Inputs().HasTag(kLetterboxPaddingTag))
|
||||||
<< "Missing one or more input streams.";
|
<< "Missing one or more input streams.";
|
||||||
|
|
||||||
cc->Inputs().Tag(kLandmarksTag).Set<NormalizedLandmarkList>();
|
RET_CHECK_EQ(cc->Inputs().NumEntries(kLandmarksTag),
|
||||||
|
cc->Outputs().NumEntries(kLandmarksTag))
|
||||||
|
<< "Same number of input and output landmarks is required.";
|
||||||
|
|
||||||
|
for (CollectionItemId id = cc->Inputs().BeginId(kLandmarksTag);
|
||||||
|
id != cc->Inputs().EndId(kLandmarksTag); ++id) {
|
||||||
|
cc->Inputs().Get(id).Set<NormalizedLandmarkList>();
|
||||||
|
}
|
||||||
cc->Inputs().Tag(kLetterboxPaddingTag).Set<std::array<float, 4>>();
|
cc->Inputs().Tag(kLetterboxPaddingTag).Set<std::array<float, 4>>();
|
||||||
|
|
||||||
cc->Outputs().Tag(kLandmarksTag).Set<NormalizedLandmarkList>();
|
for (CollectionItemId id = cc->Outputs().BeginId(kLandmarksTag);
|
||||||
|
id != cc->Outputs().EndId(kLandmarksTag); ++id) {
|
||||||
|
cc->Outputs().Get(id).Set<NormalizedLandmarkList>();
|
||||||
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -89,38 +94,45 @@ class LandmarkLetterboxRemovalCalculator : public CalculatorBase {
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
// Only process if there's input landmarks.
|
if (cc->Inputs().Tag(kLetterboxPaddingTag).IsEmpty()) {
|
||||||
if (cc->Inputs().Tag(kLandmarksTag).IsEmpty()) {
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
const NormalizedLandmarkList& input_landmarks =
|
|
||||||
cc->Inputs().Tag(kLandmarksTag).Get<NormalizedLandmarkList>();
|
|
||||||
const auto& letterbox_padding =
|
const auto& letterbox_padding =
|
||||||
cc->Inputs().Tag(kLetterboxPaddingTag).Get<std::array<float, 4>>();
|
cc->Inputs().Tag(kLetterboxPaddingTag).Get<std::array<float, 4>>();
|
||||||
|
|
||||||
const float left = letterbox_padding[0];
|
const float left = letterbox_padding[0];
|
||||||
const float top = letterbox_padding[1];
|
const float top = letterbox_padding[1];
|
||||||
const float left_and_right = letterbox_padding[0] + letterbox_padding[2];
|
const float left_and_right = letterbox_padding[0] + letterbox_padding[2];
|
||||||
const float top_and_bottom = letterbox_padding[1] + letterbox_padding[3];
|
const float top_and_bottom = letterbox_padding[1] + letterbox_padding[3];
|
||||||
|
|
||||||
NormalizedLandmarkList output_landmarks;
|
CollectionItemId input_id = cc->Inputs().BeginId(kLandmarksTag);
|
||||||
for (int i = 0; i < input_landmarks.landmark_size(); ++i) {
|
CollectionItemId output_id = cc->Outputs().BeginId(kLandmarksTag);
|
||||||
const NormalizedLandmark& landmark = input_landmarks.landmark(i);
|
// Number of inputs and outpus is the same according to the contract.
|
||||||
NormalizedLandmark* new_landmark = output_landmarks.add_landmark();
|
for (; input_id != cc->Inputs().EndId(kLandmarksTag);
|
||||||
const float new_x = (landmark.x() - left) / (1.0f - left_and_right);
|
++input_id, ++output_id) {
|
||||||
const float new_y = (landmark.y() - top) / (1.0f - top_and_bottom);
|
const auto& input_packet = cc->Inputs().Get(input_id);
|
||||||
|
if (input_packet.IsEmpty()) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
new_landmark->set_x(new_x);
|
const NormalizedLandmarkList& input_landmarks =
|
||||||
new_landmark->set_y(new_y);
|
input_packet.Get<NormalizedLandmarkList>();
|
||||||
// Keep z-coord as is.
|
NormalizedLandmarkList output_landmarks;
|
||||||
new_landmark->set_z(landmark.z());
|
for (int i = 0; i < input_landmarks.landmark_size(); ++i) {
|
||||||
|
const NormalizedLandmark& landmark = input_landmarks.landmark(i);
|
||||||
|
NormalizedLandmark* new_landmark = output_landmarks.add_landmark();
|
||||||
|
const float new_x = (landmark.x() - left) / (1.0f - left_and_right);
|
||||||
|
const float new_y = (landmark.y() - top) / (1.0f - top_and_bottom);
|
||||||
|
|
||||||
|
new_landmark->set_x(new_x);
|
||||||
|
new_landmark->set_y(new_y);
|
||||||
|
// Keep z-coord as is.
|
||||||
|
new_landmark->set_z(landmark.z());
|
||||||
|
}
|
||||||
|
|
||||||
|
cc->Outputs().Get(output_id).AddPacket(
|
||||||
|
MakePacket<NormalizedLandmarkList>(output_landmarks)
|
||||||
|
.At(cc->InputTimestamp()));
|
||||||
}
|
}
|
||||||
|
|
||||||
cc->Outputs()
|
|
||||||
.Tag(kLandmarksTag)
|
|
||||||
.AddPacket(MakePacket<NormalizedLandmarkList>(output_landmarks)
|
|
||||||
.At(cc->InputTimestamp()));
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -12,20 +12,6 @@
|
|||||||
// See the License for the specific language governing permissions and
|
// See the License for the specific language governing permissions and
|
||||||
// limitations under the License.
|
// limitations under the License.
|
||||||
|
|
||||||
// Copyright 2019 The MediaPipe Authors.
|
|
||||||
//
|
|
||||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
||||||
// you may not use this file except in compliance with the License.
|
|
||||||
// You may obtain a copy of the License at
|
|
||||||
//
|
|
||||||
// http://www.apache.org/licenses/LICENSE-2.0
|
|
||||||
//
|
|
||||||
// Unless required by applicable law or agreed to in writing, software
|
|
||||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
||||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
||||||
// See the License for the specific language governing permissions and
|
|
||||||
// limitations under the License.
|
|
||||||
|
|
||||||
#include <cmath>
|
#include <cmath>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
@@ -63,6 +49,15 @@ constexpr char kRectTag[] = "NORM_RECT";
|
|||||||
// input_stream: "NORM_RECT:rect"
|
// input_stream: "NORM_RECT:rect"
|
||||||
// output_stream: "NORM_LANDMARKS:projected_landmarks"
|
// output_stream: "NORM_LANDMARKS:projected_landmarks"
|
||||||
// }
|
// }
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "LandmarkProjectionCalculator"
|
||||||
|
// input_stream: "NORM_LANDMARKS:0:landmarks_0"
|
||||||
|
// input_stream: "NORM_LANDMARKS:1:landmarks_1"
|
||||||
|
// input_stream: "NORM_RECT:rect"
|
||||||
|
// output_stream: "NORM_LANDMARKS:0:projected_landmarks_0"
|
||||||
|
// output_stream: "NORM_LANDMARKS:1:projected_landmarks_1"
|
||||||
|
// }
|
||||||
class LandmarkProjectionCalculator : public CalculatorBase {
|
class LandmarkProjectionCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
@@ -70,10 +65,20 @@ class LandmarkProjectionCalculator : public CalculatorBase {
|
|||||||
cc->Inputs().HasTag(kRectTag))
|
cc->Inputs().HasTag(kRectTag))
|
||||||
<< "Missing one or more input streams.";
|
<< "Missing one or more input streams.";
|
||||||
|
|
||||||
cc->Inputs().Tag(kLandmarksTag).Set<NormalizedLandmarkList>();
|
RET_CHECK_EQ(cc->Inputs().NumEntries(kLandmarksTag),
|
||||||
|
cc->Outputs().NumEntries(kLandmarksTag))
|
||||||
|
<< "Same number of input and output landmarks is required.";
|
||||||
|
|
||||||
|
for (CollectionItemId id = cc->Inputs().BeginId(kLandmarksTag);
|
||||||
|
id != cc->Inputs().EndId(kLandmarksTag); ++id) {
|
||||||
|
cc->Inputs().Get(id).Set<NormalizedLandmarkList>();
|
||||||
|
}
|
||||||
cc->Inputs().Tag(kRectTag).Set<NormalizedRect>();
|
cc->Inputs().Tag(kRectTag).Set<NormalizedRect>();
|
||||||
|
|
||||||
cc->Outputs().Tag(kLandmarksTag).Set<NormalizedLandmarkList>();
|
for (CollectionItemId id = cc->Outputs().BeginId(kLandmarksTag);
|
||||||
|
id != cc->Outputs().EndId(kLandmarksTag); ++id) {
|
||||||
|
cc->Outputs().Get(id).Set<NormalizedLandmarkList>();
|
||||||
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -85,41 +90,50 @@ class LandmarkProjectionCalculator : public CalculatorBase {
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
const auto& options =
|
if (cc->Inputs().Tag(kRectTag).IsEmpty()) {
|
||||||
cc->Options<::mediapipe::LandmarkProjectionCalculatorOptions>();
|
|
||||||
// Only process if there's input landmarks.
|
|
||||||
if (cc->Inputs().Tag(kLandmarksTag).IsEmpty()) {
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
const NormalizedLandmarkList& input_landmarks =
|
|
||||||
cc->Inputs().Tag(kLandmarksTag).Get<NormalizedLandmarkList>();
|
|
||||||
const auto& input_rect = cc->Inputs().Tag(kRectTag).Get<NormalizedRect>();
|
const auto& input_rect = cc->Inputs().Tag(kRectTag).Get<NormalizedRect>();
|
||||||
|
|
||||||
NormalizedLandmarkList output_landmarks;
|
const auto& options =
|
||||||
for (int i = 0; i < input_landmarks.landmark_size(); ++i) {
|
cc->Options<::mediapipe::LandmarkProjectionCalculatorOptions>();
|
||||||
const NormalizedLandmark& landmark = input_landmarks.landmark(i);
|
|
||||||
NormalizedLandmark* new_landmark = output_landmarks.add_landmark();
|
|
||||||
|
|
||||||
const float x = landmark.x() - 0.5f;
|
CollectionItemId input_id = cc->Inputs().BeginId(kLandmarksTag);
|
||||||
const float y = landmark.y() - 0.5f;
|
CollectionItemId output_id = cc->Outputs().BeginId(kLandmarksTag);
|
||||||
const float angle = options.ignore_rotation() ? 0 : input_rect.rotation();
|
// Number of inputs and outpus is the same according to the contract.
|
||||||
float new_x = std::cos(angle) * x - std::sin(angle) * y;
|
for (; input_id != cc->Inputs().EndId(kLandmarksTag);
|
||||||
float new_y = std::sin(angle) * x + std::cos(angle) * y;
|
++input_id, ++output_id) {
|
||||||
|
const auto& input_packet = cc->Inputs().Get(input_id);
|
||||||
|
if (input_packet.IsEmpty()) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
new_x = new_x * input_rect.width() + input_rect.x_center();
|
const auto& input_landmarks = input_packet.Get<NormalizedLandmarkList>();
|
||||||
new_y = new_y * input_rect.height() + input_rect.y_center();
|
NormalizedLandmarkList output_landmarks;
|
||||||
|
for (int i = 0; i < input_landmarks.landmark_size(); ++i) {
|
||||||
|
const NormalizedLandmark& landmark = input_landmarks.landmark(i);
|
||||||
|
NormalizedLandmark* new_landmark = output_landmarks.add_landmark();
|
||||||
|
|
||||||
new_landmark->set_x(new_x);
|
const float x = landmark.x() - 0.5f;
|
||||||
new_landmark->set_y(new_y);
|
const float y = landmark.y() - 0.5f;
|
||||||
// Keep z-coord as is.
|
const float angle =
|
||||||
new_landmark->set_z(landmark.z());
|
options.ignore_rotation() ? 0 : input_rect.rotation();
|
||||||
|
float new_x = std::cos(angle) * x - std::sin(angle) * y;
|
||||||
|
float new_y = std::sin(angle) * x + std::cos(angle) * y;
|
||||||
|
|
||||||
|
new_x = new_x * input_rect.width() + input_rect.x_center();
|
||||||
|
new_y = new_y * input_rect.height() + input_rect.y_center();
|
||||||
|
|
||||||
|
new_landmark->set_x(new_x);
|
||||||
|
new_landmark->set_y(new_y);
|
||||||
|
// Keep z-coord as is.
|
||||||
|
new_landmark->set_z(landmark.z());
|
||||||
|
}
|
||||||
|
|
||||||
|
cc->Outputs().Get(output_id).AddPacket(
|
||||||
|
MakePacket<NormalizedLandmarkList>(output_landmarks)
|
||||||
|
.At(cc->InputTimestamp()));
|
||||||
}
|
}
|
||||||
|
|
||||||
cc->Outputs()
|
|
||||||
.Tag(kLandmarksTag)
|
|
||||||
.AddPacket(MakePacket<NormalizedLandmarkList>(output_landmarks)
|
|
||||||
.At(cc->InputTimestamp()));
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ namespace {
|
|||||||
|
|
||||||
constexpr char kLandmarksTag[] = "LANDMARKS";
|
constexpr char kLandmarksTag[] = "LANDMARKS";
|
||||||
constexpr char kNormLandmarksTag[] = "NORM_LANDMARKS";
|
constexpr char kNormLandmarksTag[] = "NORM_LANDMARKS";
|
||||||
|
constexpr char kRenderScaleTag[] = "RENDER_SCALE";
|
||||||
constexpr char kRenderDataTag[] = "RENDER_DATA";
|
constexpr char kRenderDataTag[] = "RENDER_DATA";
|
||||||
constexpr char kLandmarkLabel[] = "KEYPOINT";
|
constexpr char kLandmarkLabel[] = "KEYPOINT";
|
||||||
constexpr int kMaxLandmarkThickness = 18;
|
constexpr int kMaxLandmarkThickness = 18;
|
||||||
@@ -71,6 +72,83 @@ void SetColorSizeValueFromZ(float z, float z_min, float z_max,
|
|||||||
render_annotation->set_thickness(thickness);
|
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
|
} // namespace
|
||||||
|
|
||||||
// A calculator that converts Landmark proto to RenderData proto for
|
// A calculator that converts Landmark proto to RenderData proto for
|
||||||
@@ -107,29 +185,6 @@ class LandmarksToRenderDataCalculator : public CalculatorBase {
|
|||||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||||
|
|
||||||
private:
|
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_;
|
LandmarksToRenderDataCalculatorOptions options_;
|
||||||
};
|
};
|
||||||
REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
||||||
@@ -150,6 +205,9 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
|||||||
if (cc->Inputs().HasTag(kNormLandmarksTag)) {
|
if (cc->Inputs().HasTag(kNormLandmarksTag)) {
|
||||||
cc->Inputs().Tag(kNormLandmarksTag).Set<NormalizedLandmarkList>();
|
cc->Inputs().Tag(kNormLandmarksTag).Set<NormalizedLandmarkList>();
|
||||||
}
|
}
|
||||||
|
if (cc->Inputs().HasTag(kRenderScaleTag)) {
|
||||||
|
cc->Inputs().Tag(kRenderScaleTag).Set<float>();
|
||||||
|
}
|
||||||
cc->Outputs().Tag(kRenderDataTag).Set<RenderData>();
|
cc->Outputs().Tag(kRenderDataTag).Set<RenderData>();
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -169,11 +227,26 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
|||||||
float z_min = 0.f;
|
float z_min = 0.f;
|
||||||
float z_max = 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;
|
||||||
|
}
|
||||||
|
|
||||||
|
// 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";
|
||||||
|
std::vector<int> landmark_connections;
|
||||||
|
for (int i = 0; i < options_.landmark_connections_size(); i += 1) {
|
||||||
|
landmark_connections.push_back(options_.landmark_connections(i));
|
||||||
|
}
|
||||||
|
|
||||||
if (cc->Inputs().HasTag(kLandmarksTag)) {
|
if (cc->Inputs().HasTag(kLandmarksTag)) {
|
||||||
const LandmarkList& landmarks =
|
const LandmarkList& landmarks =
|
||||||
cc->Inputs().Tag(kLandmarksTag).Get<LandmarkList>();
|
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) {
|
if (visualize_depth) {
|
||||||
GetMinMaxZ<LandmarkList, Landmark>(landmarks, &z_min, &z_max);
|
GetMinMaxZ<LandmarkList, Landmark>(landmarks, &z_min, &z_max);
|
||||||
}
|
}
|
||||||
@@ -181,8 +254,8 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
|||||||
visualize_depth &= ((z_max - z_min) > 1e-3);
|
visualize_depth &= ((z_max - z_min) > 1e-3);
|
||||||
for (int i = 0; i < landmarks.landmark_size(); ++i) {
|
for (int i = 0; i < landmarks.landmark_size(); ++i) {
|
||||||
const Landmark& landmark = landmarks.landmark(i);
|
const Landmark& landmark = landmarks.landmark(i);
|
||||||
auto* landmark_data_render =
|
auto* landmark_data_render = AddPointRenderData(
|
||||||
AddPointRenderData(options_, render_data.get());
|
options_.landmark_color(), thickness, render_data.get());
|
||||||
if (visualize_depth) {
|
if (visualize_depth) {
|
||||||
SetColorSizeValueFromZ(landmark.z(), z_min, z_max,
|
SetColorSizeValueFromZ(landmark.z(), z_min, z_max,
|
||||||
landmark_data_render);
|
landmark_data_render);
|
||||||
@@ -193,19 +266,19 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
|||||||
landmark_data->set_y(landmark.y());
|
landmark_data->set_y(landmark.y());
|
||||||
}
|
}
|
||||||
if (visualize_depth) {
|
if (visualize_depth) {
|
||||||
AddConnectionsWithDepth<LandmarkList>(landmarks, /*normalized=*/false,
|
AddConnectionsWithDepth<LandmarkList, Landmark>(
|
||||||
z_min, z_max, render_data.get());
|
landmarks, landmark_connections, thickness, /*normalized=*/false,
|
||||||
|
z_min, z_max, render_data.get());
|
||||||
} else {
|
} else {
|
||||||
AddConnections<LandmarkList>(landmarks, /*normalized=*/false,
|
AddConnections<LandmarkList, Landmark>(
|
||||||
render_data.get());
|
landmarks, landmark_connections, options_.connection_color(),
|
||||||
|
thickness, /*normalized=*/false, render_data.get());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if (cc->Inputs().HasTag(kNormLandmarksTag)) {
|
if (cc->Inputs().HasTag(kNormLandmarksTag)) {
|
||||||
const NormalizedLandmarkList& landmarks =
|
const NormalizedLandmarkList& landmarks =
|
||||||
cc->Inputs().Tag(kNormLandmarksTag).Get<NormalizedLandmarkList>();
|
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) {
|
if (visualize_depth) {
|
||||||
GetMinMaxZ<NormalizedLandmarkList, NormalizedLandmark>(landmarks, &z_min,
|
GetMinMaxZ<NormalizedLandmarkList, NormalizedLandmark>(landmarks, &z_min,
|
||||||
&z_max);
|
&z_max);
|
||||||
@@ -214,8 +287,8 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
|||||||
visualize_depth &= ((z_max - z_min) > 1e-3);
|
visualize_depth &= ((z_max - z_min) > 1e-3);
|
||||||
for (int i = 0; i < landmarks.landmark_size(); ++i) {
|
for (int i = 0; i < landmarks.landmark_size(); ++i) {
|
||||||
const NormalizedLandmark& landmark = landmarks.landmark(i);
|
const NormalizedLandmark& landmark = landmarks.landmark(i);
|
||||||
auto* landmark_data_render =
|
auto* landmark_data_render = AddPointRenderData(
|
||||||
AddPointRenderData(options_, render_data.get());
|
options_.landmark_color(), thickness, render_data.get());
|
||||||
if (visualize_depth) {
|
if (visualize_depth) {
|
||||||
SetColorSizeValueFromZ(landmark.z(), z_min, z_max,
|
SetColorSizeValueFromZ(landmark.z(), z_min, z_max,
|
||||||
landmark_data_render);
|
landmark_data_render);
|
||||||
@@ -226,11 +299,13 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
|||||||
landmark_data->set_y(landmark.y());
|
landmark_data->set_y(landmark.y());
|
||||||
}
|
}
|
||||||
if (visualize_depth) {
|
if (visualize_depth) {
|
||||||
AddConnectionsWithDepth<NormalizedLandmarkList>(
|
AddConnectionsWithDepth<NormalizedLandmarkList, NormalizedLandmark>(
|
||||||
landmarks, /*normalized=*/true, z_min, z_max, render_data.get());
|
landmarks, landmark_connections, thickness, /*normalized=*/true,
|
||||||
|
z_min, z_max, render_data.get());
|
||||||
} else {
|
} else {
|
||||||
AddConnections<NormalizedLandmarkList>(landmarks, /*normalized=*/true,
|
AddConnections<NormalizedLandmarkList, NormalizedLandmark>(
|
||||||
render_data.get());
|
landmarks, landmark_connections, options_.connection_color(),
|
||||||
|
thickness, /*normalized=*/true, render_data.get());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -240,84 +315,4 @@ REGISTER_CALCULATOR(LandmarksToRenderDataCalculator);
|
|||||||
return ::mediapipe::OkStatus();
|
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
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -16,34 +16,80 @@
|
|||||||
#include <string>
|
#include <string>
|
||||||
|
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/port/file_helpers.h"
|
|
||||||
#include "mediapipe/framework/port/status.h"
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
#include "mediapipe/util/resource_util.h"
|
||||||
|
|
||||||
namespace mediapipe {
|
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
|
// The calculator takes the path to the local file as an input side packet and
|
||||||
// outputs the contents of that file.
|
// outputs the contents of that file.
|
||||||
//
|
//
|
||||||
|
// NOTE: file loading can be batched by providing multiple input/output side
|
||||||
|
// packets.
|
||||||
|
//
|
||||||
// Example config:
|
// Example config:
|
||||||
// node {
|
// node {
|
||||||
// calculator: "LocalFileContentsCalculator"
|
// calculator: "LocalFileContentsCalculator"
|
||||||
// input_side_packet: "FILE_PATH:file_path"
|
// input_side_packet: "FILE_PATH:file_path"
|
||||||
// output_side_packet: "CONTENTS:contents"
|
// 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 {
|
class LocalFileContentsCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
cc->InputSidePackets().Tag("FILE_PATH").Set<std::string>();
|
RET_CHECK(cc->InputSidePackets().HasTag(kFilePathTag))
|
||||||
cc->OutputSidePackets().Tag("CONTENTS").Set<std::string>();
|
<< "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();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
std::string contents;
|
CollectionItemId input_id = cc->InputSidePackets().BeginId(kFilePathTag);
|
||||||
MP_RETURN_IF_ERROR(mediapipe::file::GetContents(
|
CollectionItemId output_id = cc->OutputSidePackets().BeginId(kContentsTag);
|
||||||
cc->InputSidePackets().Tag("FILE_PATH").Get<std::string>(), &contents));
|
// Number of inputs and outpus is the same according to the contract.
|
||||||
cc->OutputSidePackets()
|
for (; input_id != cc->InputSidePackets().EndId(kFilePathTag);
|
||||||
.Tag("CONTENTS")
|
++input_id, ++output_id) {
|
||||||
.Set(MakePacket<std::string>(std::move(contents)));
|
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();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,75 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/file_helpers.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
// The calculator takes the path to local directory and desired file suffix to
|
||||||
|
// mach as input side packets, and outputs the contents of those files that
|
||||||
|
// match the pattern. Those matched files will be sent sequentially through the
|
||||||
|
// output stream with incremental timestamp difference by 1.
|
||||||
|
//
|
||||||
|
// Example config:
|
||||||
|
// node {
|
||||||
|
// calculator: "LocalFilePatternContentsCalculator"
|
||||||
|
// input_side_packet: "FILE_DIRECTORY:file_directory"
|
||||||
|
// input_side_packet: "FILE_SUFFIX:file_suffix"
|
||||||
|
// output_stream: "CONTENTS:contents"
|
||||||
|
// }
|
||||||
|
class LocalFilePatternContentsCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
cc->InputSidePackets().Tag("FILE_DIRECTORY").Set<std::string>();
|
||||||
|
cc->InputSidePackets().Tag("FILE_SUFFIX").Set<std::string>();
|
||||||
|
cc->Outputs().Tag("CONTENTS").Set<std::string>();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
|
MP_RETURN_IF_ERROR(::mediapipe::file::MatchFileTypeInDirectory(
|
||||||
|
cc->InputSidePackets().Tag("FILE_DIRECTORY").Get<std::string>(),
|
||||||
|
cc->InputSidePackets().Tag("FILE_SUFFIX").Get<std::string>(),
|
||||||
|
&filenames_));
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
|
if (current_output_ < filenames_.size()) {
|
||||||
|
auto contents = absl::make_unique<std::string>();
|
||||||
|
LOG(INFO) << filenames_[current_output_];
|
||||||
|
MP_RETURN_IF_ERROR(mediapipe::file::GetContents(
|
||||||
|
filenames_[current_output_], contents.get()));
|
||||||
|
++current_output_;
|
||||||
|
cc->Outputs()
|
||||||
|
.Tag("CONTENTS")
|
||||||
|
.Add(contents.release(), Timestamp(current_output_));
|
||||||
|
} else {
|
||||||
|
return tool::StatusStop();
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
std::vector<std::string> filenames_;
|
||||||
|
int current_output_ = 0;
|
||||||
|
};
|
||||||
|
|
||||||
|
REGISTER_CALCULATOR(LocalFilePatternContentsCalculator);
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -276,6 +276,7 @@ TEST_F(PacketLatencyCalculatorTest, DoesNotOutputUntilReferencePacketReceived) {
|
|||||||
"delayed_packet_0", Adopt(new double()).At(Timestamp(2))));
|
"delayed_packet_0", Adopt(new double()).At(Timestamp(2))));
|
||||||
|
|
||||||
// Send a reference packet with timestamp 10 usec.
|
// Send a reference packet with timestamp 10 usec.
|
||||||
|
simulation_clock_->Sleep(absl::Microseconds(1));
|
||||||
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
MP_ASSERT_OK(graph_.AddPacketToInputStream(
|
||||||
"camera_frames", Adopt(new double()).At(Timestamp(10))));
|
"camera_frames", Adopt(new double()).At(Timestamp(10))));
|
||||||
simulation_clock_->Sleep(absl::Microseconds(1));
|
simulation_clock_->Sleep(absl::Microseconds(1));
|
||||||
|
|||||||
@@ -45,15 +45,17 @@ RenderAnnotation::Rectangle* NewRect(
|
|||||||
void SetRect(bool normalized, double xmin, double ymin, double width,
|
void SetRect(bool normalized, double xmin, double ymin, double width,
|
||||||
double height, double rotation,
|
double height, double rotation,
|
||||||
RenderAnnotation::Rectangle* rect) {
|
RenderAnnotation::Rectangle* rect) {
|
||||||
if (xmin + width < 0.0 || ymin + height < 0.0) return;
|
if (rotation == 0.0) {
|
||||||
if (normalized) {
|
if (xmin + width < 0.0 || ymin + height < 0.0) return;
|
||||||
if (xmin > 1.0 || ymin > 1.0) return;
|
if (normalized) {
|
||||||
|
if (xmin > 1.0 || ymin > 1.0) return;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
rect->set_normalized(normalized);
|
rect->set_normalized(normalized);
|
||||||
rect->set_left(normalized ? std::max(xmin, 0.0) : xmin);
|
rect->set_left(xmin);
|
||||||
rect->set_top(normalized ? std::max(ymin, 0.0) : ymin);
|
rect->set_top(ymin);
|
||||||
rect->set_right(normalized ? std::min(xmin + width, 1.0) : xmin + width);
|
rect->set_right(xmin + width);
|
||||||
rect->set_bottom(normalized ? std::min(ymin + height, 1.0) : ymin + height);
|
rect->set_bottom(ymin + height);
|
||||||
rect->set_rotation(rotation);
|
rect->set_rotation(rotation);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user