Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3b6d3c4058 | ||
|
|
252a5713c7 | ||
|
|
de4fbc10e6 | ||
|
|
d144e564d8 | ||
|
|
dd02df1dbe | ||
|
|
66b377c825 | ||
|
|
bf5185f122 | ||
|
|
a2823541e6 | ||
|
|
ae6be10afe | ||
|
|
38ee2603a7 | ||
|
|
86b3283b2f | ||
|
|
7d470a1335 | ||
|
|
d16cc3be5b | ||
|
|
137867d088 | ||
|
|
446d7cf6b6 | ||
|
|
90f72bd851 | ||
|
|
4285aeddfc | ||
|
|
37287925b0 | ||
|
|
48bcbb115f | ||
|
|
9437483827 | ||
|
|
d030c13931 | ||
|
|
fce372d153 | ||
|
|
c6fea4c9d9 | ||
|
|
259b48e082 | ||
|
|
a2a63e3876 | ||
|
|
412ab42d1f |
@@ -3,7 +3,7 @@
|
|||||||
# Basic build settings
|
# Basic build settings
|
||||||
build --jobs 128
|
build --jobs 128
|
||||||
build --define='absl=1'
|
build --define='absl=1'
|
||||||
build --cxxopt='-std=c++11'
|
build --cxxopt='-std=c++14'
|
||||||
build --copt='-Wno-sign-compare'
|
build --copt='-Wno-sign-compare'
|
||||||
build --copt='-Wno-unused-function'
|
build --copt='-Wno-unused-function'
|
||||||
build --copt='-Wno-uninitialized'
|
build --copt='-Wno-uninitialized'
|
||||||
@@ -12,13 +12,16 @@ 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
|
# Temporarily set the incompatibility flag for Bazel 0.27.0 and above
|
||||||
build --incompatible_disable_deprecated_attr_params=false
|
build --incompatible_disable_deprecated_attr_params=false
|
||||||
build --incompatible_depset_is_not_iterable=false
|
build --incompatible_depset_is_not_iterable=false
|
||||||
|
|
||||||
# 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.
|
||||||
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
|
||||||
|
|||||||
+5
-1
@@ -30,10 +30,13 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
|
|||||||
unzip \
|
unzip \
|
||||||
python \
|
python \
|
||||||
python-pip \
|
python-pip \
|
||||||
|
python3-pip \
|
||||||
libopencv-core-dev \
|
libopencv-core-dev \
|
||||||
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 && \
|
||||||
@@ -42,9 +45,10 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
|
|||||||
|
|
||||||
RUN pip install --upgrade setuptools
|
RUN pip install --upgrade setuptools
|
||||||
RUN pip install future
|
RUN pip install future
|
||||||
|
RUN pip3 install six
|
||||||
|
|
||||||
# Install bazel
|
# Install bazel
|
||||||
ARG BAZEL_VERSION=0.26.1
|
ARG BAZEL_VERSION=1.1.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,27 +1,36 @@
|
|||||||

|

|
||||||
=======================================================================
|
=======================================================================
|
||||||
|
|
||||||
[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.
|
||||||
|
|
||||||

|

|
||||||
|
|
||||||
|
> "<em>MediaPipe has made it extremely easy to build our 3D person pose reconstruction demo app, facilitating accelerated neural network inference on device and synchronization of our result visualization with the video capture stream. Highly recommended!</em>" - George Papandreou, CTO, [Ariel AI](https://arielai.com)
|
||||||
|
|
||||||
## 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)
|
||||||
* [Face Detection](mediapipe/docs/face_detection_mobile_gpu.md)
|
* [Multi-hand Tracking](mediapipe/docs/multi_hand_tracking_mobile_gpu.md)
|
||||||
* [Hair Segmentation](mediapipe/docs/hair_segmentation_mobile_gpu.md)
|
* [Hand Tracking](mediapipe/docs/hand_tracking_mobile_gpu.md) [[Web Demo]](https://viz.mediapipe.dev/runner/demos/hand_tracking/hand_tracking.html)
|
||||||
|
* [Hair Segmentation](mediapipe/docs/hair_segmentation_mobile_gpu.md) [[Web Demo]](https://viz.mediapipe.dev/runner/demos/hair_segmentation/hair_segmentation.html)
|
||||||
* [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).
|
||||||
|
|
||||||
## Getting started
|
## Getting started
|
||||||
See mobile and desktop [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)
|
||||||
@@ -31,17 +40,33 @@ 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)
|
||||||
* [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)
|
||||||
|
|
||||||
## Events
|
## Events
|
||||||
[Open sourced at CVPR 2019](https://sites.google.com/corp/view/perception-cv4arvr/mediapipe) on June 17~20 in Long Beach, CA
|
* [AI Nextcon 2020, 12-16 Feb 2020, Seattle](http://aisea20.xnextcon.com/)
|
||||||
|
* [MediaPipe Madrid Meetup, 16 Dec 2019](https://www.meetup.com/Madrid-AI-Developers-Group/events/266329088/)
|
||||||
|
* [MediaPipe London Meetup, Google 123 Building, 12 Dec 2019](https://www.meetup.com/London-AI-Tech-Talk/events/266329038)
|
||||||
|
* [ML Conference, Berlin, 11 Dec 2019](https://mlconference.ai/machine-learning-advanced-development/mediapipe-building-real-time-cross-platform-mobile-web-edge-desktop-video-audio-ml-pipelines/)
|
||||||
|
* [MediaPipe Berlin Meetup, Google Berlin, 11 Dec 2019](https://www.meetup.com/Berlin-AI-Tech-Talk/events/266328794/)
|
||||||
|
* [The 3rd Workshop on YouTube-8M Large Scale Video Understanding Workshop](https://research.google.com/youtube8m/workshop2019/index.html) Seoul, Korea ICCV 2019
|
||||||
|
* [AI DevWorld 2019](https://aidevworld.com) on Oct 10 in San Jose, California
|
||||||
|
* [Google Industry Workshop at ICIP 2019](http://2019.ieeeicip.org/?action=page4&id=14#Google) [Presentation](https://docs.google.com/presentation/d/e/2PACX-1vRIBBbO_LO9v2YmvbHHEt1cwyqH6EjDxiILjuT0foXy1E7g6uyh4CesB2DkkEwlRDO9_lWfuKMZx98T/pub?start=false&loop=false&delayms=3000&slide=id.g556cc1a659_0_5) on Sept 24 in Taipei, Taiwan
|
||||||
|
* [Open sourced at CVPR 2019](https://sites.google.com/corp/view/perception-cv4arvr/mediapipe) on June 17~20 in Long Beach, CA
|
||||||
|
|
||||||
|
## 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).
|
||||||
|
|||||||
@@ -10,19 +10,25 @@ http_archive(
|
|||||||
sha256 = "2ef429f5d7ce7111263289644d233707dba35e39696377ebab8b0bc701f7818e",
|
sha256 = "2ef429f5d7ce7111263289644d233707dba35e39696377ebab8b0bc701f7818e",
|
||||||
)
|
)
|
||||||
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 = "1.0.0",
|
||||||
|
maximum_bazel_version = "1.2.1")
|
||||||
|
|
||||||
# ABSL cpp library.
|
|
||||||
|
# ABSL cpp library lts_2020_02_25
|
||||||
http_archive(
|
http_archive(
|
||||||
name = "com_google_absl",
|
name = "com_google_absl",
|
||||||
# Head commit on 2019-04-12.
|
|
||||||
# TODO: Switch to the latest absl version when the problem gets
|
|
||||||
# fixed.
|
|
||||||
urls = [
|
urls = [
|
||||||
"https://github.com/abseil/abseil-cpp/archive/a02f62f456f2c4a7ecf2be3104fe0c6e16fbad9a.tar.gz",
|
"https://github.com/abseil/abseil-cpp/archive/20200225.tar.gz",
|
||||||
],
|
],
|
||||||
sha256 = "d437920d1434c766d22e85773b899c77c672b8b4865d5dc2cd61a29fdff3cf03",
|
# Remove after https://github.com/abseil/abseil-cpp/issues/326 is solved.
|
||||||
strip_prefix = "abseil-cpp-a02f62f456f2c4a7ecf2be3104fe0c6e16fbad9a",
|
patches = [
|
||||||
|
"@//third_party:com_google_absl_f863b622fe13612433fdf43f76547d5edda0c93001.diff"
|
||||||
|
],
|
||||||
|
patch_args = [
|
||||||
|
"-p1",
|
||||||
|
],
|
||||||
|
strip_prefix = "abseil-cpp-20200225",
|
||||||
|
sha256 = "728a813291bdec2aa46eab8356ace9f75ac2ed9dfe2df5ab603c4e6c09f1c353"
|
||||||
)
|
)
|
||||||
|
|
||||||
http_archive(
|
http_archive(
|
||||||
@@ -64,6 +70,20 @@ http_archive(
|
|||||||
sha256 = "267103f8a1e9578978aa1dc256001e6529ef593e5aea38193d31c2872ee025e8",
|
sha256 = "267103f8a1e9578978aa1dc256001e6529ef593e5aea38193d31c2872ee025e8",
|
||||||
strip_prefix = "glog-0.3.5",
|
strip_prefix = "glog-0.3.5",
|
||||||
build_file = "@//third_party:glog.BUILD",
|
build_file = "@//third_party:glog.BUILD",
|
||||||
|
patches = [
|
||||||
|
"@//third_party:com_github_glog_glog_9779e5ea6ef59562b030248947f787d1256132ae.diff"
|
||||||
|
],
|
||||||
|
patch_args = [
|
||||||
|
"-p1",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
# 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
|
||||||
@@ -97,31 +117,46 @@ http_archive(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
# 2019-08-15
|
# 2020-02-12
|
||||||
_TENSORFLOW_GIT_COMMIT = "67def62936e28f97c16182dfcc467d8d1cae02b4"
|
# The last commit before TensorFlow switched to Bazel 2.0
|
||||||
_TENSORFLOW_SHA256= "ddd4e3c056e7c0ff2ef29133b30fa62781dfbf8a903e99efb91a02d292fa9562"
|
_TENSORFLOW_GIT_COMMIT = "77e9ffb9b2bfb1a4f7056e62d84039626923e328"
|
||||||
|
_TENSORFLOW_SHA256= "176ccd82f7dd17c5e117b50d353603b129c7a6ccbfebd522ca47cc2a40f33f13"
|
||||||
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,
|
||||||
],
|
],
|
||||||
strip_prefix = "tensorflow-%s" % _TENSORFLOW_GIT_COMMIT,
|
# A compatibility patch
|
||||||
sha256 = _TENSORFLOW_SHA256,
|
|
||||||
patches = [
|
patches = [
|
||||||
"@//third_party:tensorflow_065c20bf79253257c87bd4614bb9a7fdef015cbb.diff",
|
"@//third_party:org_tensorflow_528e22eae8bf3206189a066032c66e9e5c9b4a61.diff"
|
||||||
"@//third_party:tensorflow_f67fcbefce906cd419e4657f0d41e21019b71abd.diff",
|
|
||||||
],
|
],
|
||||||
patch_args = [
|
patch_args = [
|
||||||
"-p1",
|
"-p1",
|
||||||
],
|
],
|
||||||
|
strip_prefix = "tensorflow-%s" % _TENSORFLOW_GIT_COMMIT,
|
||||||
|
sha256 = _TENSORFLOW_SHA256,
|
||||||
)
|
)
|
||||||
|
|
||||||
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")
|
||||||
|
|
||||||
|
http_archive(
|
||||||
|
name = "ceres_solver",
|
||||||
|
url = "https://github.com/ceres-solver/ceres-solver/archive/1.14.0.zip",
|
||||||
|
patches = [
|
||||||
|
"@//third_party:ceres_solver_9bf9588988236279e1262f75d7f4d85711dfa172.diff"
|
||||||
|
],
|
||||||
|
patch_args = [
|
||||||
|
"-p1",
|
||||||
|
],
|
||||||
|
strip_prefix = "ceres-solver-1.14.0",
|
||||||
|
sha256 = "5ba6d0db4e784621fda44a50c58bb23b0892684692f0c623e2063f9c19f192f1"
|
||||||
|
)
|
||||||
|
|
||||||
# Please run
|
# Please run
|
||||||
# $ sudo apt-get install libopencv-core-dev libopencv-highgui-dev \
|
# $ sudo apt-get install libopencv-core-dev libopencv-highgui-dev \
|
||||||
|
# libopencv-calib3d-dev libopencv-features2d-dev \
|
||||||
# libopencv-imgproc-dev libopencv-video-dev
|
# libopencv-imgproc-dev libopencv-video-dev
|
||||||
new_local_repository(
|
new_local_repository(
|
||||||
name = "linux_opencv",
|
name = "linux_opencv",
|
||||||
@@ -150,11 +185,10 @@ new_local_repository(
|
|||||||
|
|
||||||
http_archive(
|
http_archive(
|
||||||
name = "android_opencv",
|
name = "android_opencv",
|
||||||
sha256 = "056b849842e4fa8751d09edbb64530cfa7a63c84ccd232d0ace330e27ba55d0b",
|
|
||||||
build_file = "@//third_party:opencv_android.BUILD",
|
build_file = "@//third_party:opencv_android.BUILD",
|
||||||
strip_prefix = "OpenCV-android-sdk",
|
strip_prefix = "OpenCV-android-sdk",
|
||||||
type = "zip",
|
type = "zip",
|
||||||
url = "https://github.com/opencv/opencv/releases/download/4.1.0/opencv-4.1.0-android-sdk.zip",
|
url = "https://github.com/opencv/opencv/releases/download/3.4.3/opencv-3.4.3-android-sdk.zip",
|
||||||
)
|
)
|
||||||
|
|
||||||
# After OpenCV 3.2.0, the pre-compiled opencv2.framework has google protobuf symbols, which will
|
# After OpenCV 3.2.0, the pre-compiled opencv2.framework has google protobuf symbols, which will
|
||||||
@@ -185,13 +219,18 @@ maven_install(
|
|||||||
artifacts = [
|
artifacts = [
|
||||||
"androidx.annotation:annotation:aar:1.1.0",
|
"androidx.annotation:annotation:aar:1.1.0",
|
||||||
"androidx.appcompat:appcompat:aar:1.1.0-rc01",
|
"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.constraintlayout:constraintlayout:aar:1.1.3",
|
||||||
"androidx.core:core:aar:1.1.0-rc03",
|
"androidx.core:core:aar:1.1.0-rc03",
|
||||||
"androidx.legacy:legacy-support-v4:aar:1.0.0",
|
"androidx.legacy:legacy-support-v4:aar:1.0.0",
|
||||||
"androidx.recyclerview:recyclerview:aar:1.1.0-beta02",
|
"androidx.recyclerview:recyclerview:aar:1.1.0-beta02",
|
||||||
"com.google.android.material:material:aar:1.0.0-rc01",
|
"com.google.android.material:material:aar:1.0.0-rc01",
|
||||||
],
|
],
|
||||||
repositories = ["https://dl.google.com/dl/android/maven2"],
|
repositories = [
|
||||||
|
"https://dl.google.com/dl/android/maven2",
|
||||||
|
"https://repo1.maven.org/maven2",
|
||||||
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
maven_server(
|
maven_server(
|
||||||
@@ -207,10 +246,10 @@ maven_jar(
|
|||||||
)
|
)
|
||||||
|
|
||||||
maven_jar(
|
maven_jar(
|
||||||
name = "androidx_concurrent_futures",
|
name = "androidx_concurrent_futures",
|
||||||
artifact = "androidx.concurrent:concurrent-futures:1.0.0-alpha03",
|
artifact = "androidx.concurrent:concurrent-futures:1.0.0-alpha03",
|
||||||
sha1 = "b528df95c7e2fefa2210c0c742bf3e491c1818ae",
|
sha1 = "b528df95c7e2fefa2210c0c742bf3e491c1818ae",
|
||||||
server = "google_server",
|
server = "google_server",
|
||||||
)
|
)
|
||||||
|
|
||||||
maven_jar(
|
maven_jar(
|
||||||
@@ -248,18 +287,11 @@ android_sdk_repository(
|
|||||||
|
|
||||||
# iOS basic build deps.
|
# iOS basic build deps.
|
||||||
|
|
||||||
load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
|
http_archive(
|
||||||
|
|
||||||
git_repository(
|
|
||||||
name = "build_bazel_rules_apple",
|
name = "build_bazel_rules_apple",
|
||||||
remote = "https://github.com/bazelbuild/rules_apple.git",
|
sha256 = "bdc8e66e70b8a75da23b79f1f8c6207356df07d041d96d2189add7ee0780cf4e",
|
||||||
tag = "0.18.0",
|
strip_prefix = "rules_apple-b869b0d3868d78a1d4ffd866ccb304fb68aa12c3",
|
||||||
patches = [
|
url = "https://github.com/bazelbuild/rules_apple/archive/b869b0d3868d78a1d4ffd866ccb304fb68aa12c3.tar.gz",
|
||||||
"@//third_party:rules_apple_c0863d0596ae6b769a29fa3fb72ff036444fd249.diff",
|
|
||||||
],
|
|
||||||
patch_args = [
|
|
||||||
"-p1",
|
|
||||||
],
|
|
||||||
)
|
)
|
||||||
|
|
||||||
load(
|
load(
|
||||||
@@ -292,3 +324,4 @@ 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",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -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"
|
||||||
],
|
],
|
||||||
|
|||||||
@@ -61,7 +61,9 @@ class AudioDecoderCalculator : public CalculatorBase {
|
|||||||
::mediapipe::Status AudioDecoderCalculator::GetContract(
|
::mediapipe::Status AudioDecoderCalculator::GetContract(
|
||||||
CalculatorContract* cc) {
|
CalculatorContract* cc) {
|
||||||
cc->InputSidePackets().Tag("INPUT_FILE_PATH").Set<std::string>();
|
cc->InputSidePackets().Tag("INPUT_FILE_PATH").Set<std::string>();
|
||||||
|
if (cc->InputSidePackets().HasTag("OPTIONS")) {
|
||||||
|
cc->InputSidePackets().Tag("OPTIONS").Set<mediapipe::AudioDecoderOptions>();
|
||||||
|
}
|
||||||
cc->Outputs().Tag("AUDIO").Set<Matrix>();
|
cc->Outputs().Tag("AUDIO").Set<Matrix>();
|
||||||
if (cc->Outputs().HasTag("AUDIO_HEADER")) {
|
if (cc->Outputs().HasTag("AUDIO_HEADER")) {
|
||||||
cc->Outputs().Tag("AUDIO_HEADER").SetNone();
|
cc->Outputs().Tag("AUDIO_HEADER").SetNone();
|
||||||
@@ -72,7 +74,9 @@ class AudioDecoderCalculator : public CalculatorBase {
|
|||||||
::mediapipe::Status AudioDecoderCalculator::Open(CalculatorContext* cc) {
|
::mediapipe::Status AudioDecoderCalculator::Open(CalculatorContext* cc) {
|
||||||
const std::string& input_file_path =
|
const std::string& input_file_path =
|
||||||
cc->InputSidePackets().Tag("INPUT_FILE_PATH").Get<std::string>();
|
cc->InputSidePackets().Tag("INPUT_FILE_PATH").Get<std::string>();
|
||||||
const auto& decoder_options = cc->Options<mediapipe::AudioDecoderOptions>();
|
const auto& decoder_options =
|
||||||
|
tool::RetrieveOptions(cc->Options<mediapipe::AudioDecoderOptions>(),
|
||||||
|
cc->InputSidePackets(), "OPTIONS");
|
||||||
decoder_ = absl::make_unique<AudioDecoder>();
|
decoder_ = absl::make_unique<AudioDecoder>();
|
||||||
MP_RETURN_IF_ERROR(decoder_->Initialize(input_file_path, decoder_options));
|
MP_RETURN_IF_ERROR(decoder_->Initialize(input_file_path, decoder_options));
|
||||||
std::unique_ptr<mediapipe::TimeSeriesHeader> header =
|
std::unique_ptr<mediapipe::TimeSeriesHeader> header =
|
||||||
|
|||||||
@@ -113,8 +113,15 @@ class SpectrogramCalculator : public CalculatorBase {
|
|||||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||||
|
|
||||||
private:
|
private:
|
||||||
Timestamp CurrentOutputTimestamp() {
|
Timestamp CurrentOutputTimestamp(CalculatorContext* cc) {
|
||||||
// Current output timestamp is the *center* of the next frame to be
|
if (use_local_timestamp_) {
|
||||||
|
return cc->InputTimestamp();
|
||||||
|
}
|
||||||
|
return CumulativeOutputTimestamp();
|
||||||
|
}
|
||||||
|
|
||||||
|
Timestamp CumulativeOutputTimestamp() {
|
||||||
|
// Cumulative output timestamp is the *center* of the next frame to be
|
||||||
// emitted, hence delayed by half a window duration compared to relevant
|
// emitted, hence delayed by half a window duration compared to relevant
|
||||||
// input timestamp.
|
// input timestamp.
|
||||||
return initial_input_timestamp_ +
|
return initial_input_timestamp_ +
|
||||||
@@ -141,6 +148,7 @@ class SpectrogramCalculator : public CalculatorBase {
|
|||||||
const OutputMatrixType postprocess_output_fn(const OutputMatrixType&),
|
const OutputMatrixType postprocess_output_fn(const OutputMatrixType&),
|
||||||
CalculatorContext* cc);
|
CalculatorContext* cc);
|
||||||
|
|
||||||
|
bool use_local_timestamp_;
|
||||||
double input_sample_rate_;
|
double input_sample_rate_;
|
||||||
bool pad_final_packet_;
|
bool pad_final_packet_;
|
||||||
int frame_duration_samples_;
|
int frame_duration_samples_;
|
||||||
@@ -173,6 +181,8 @@ const float SpectrogramCalculator::kLnPowerToDb = 4.342944819032518;
|
|||||||
SpectrogramCalculatorOptions spectrogram_options =
|
SpectrogramCalculatorOptions spectrogram_options =
|
||||||
cc->Options<SpectrogramCalculatorOptions>();
|
cc->Options<SpectrogramCalculatorOptions>();
|
||||||
|
|
||||||
|
use_local_timestamp_ = spectrogram_options.use_local_timestamp();
|
||||||
|
|
||||||
if (spectrogram_options.frame_duration_seconds() <= 0.0) {
|
if (spectrogram_options.frame_duration_seconds() <= 0.0) {
|
||||||
::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||||
<< "Invalid or missing frame_duration_seconds.\n"
|
<< "Invalid or missing frame_duration_seconds.\n"
|
||||||
@@ -351,11 +361,11 @@ template <class OutputMatrixType>
|
|||||||
<< "Inconsistent number of spectrogram channels.";
|
<< "Inconsistent number of spectrogram channels.";
|
||||||
if (allow_multichannel_input_) {
|
if (allow_multichannel_input_) {
|
||||||
cc->Outputs().Index(0).Add(spectrogram_matrices.release(),
|
cc->Outputs().Index(0).Add(spectrogram_matrices.release(),
|
||||||
CurrentOutputTimestamp());
|
CurrentOutputTimestamp(cc));
|
||||||
} else {
|
} else {
|
||||||
cc->Outputs().Index(0).Add(
|
cc->Outputs().Index(0).Add(
|
||||||
new OutputMatrixType(spectrogram_matrices->at(0)),
|
new OutputMatrixType(spectrogram_matrices->at(0)),
|
||||||
CurrentOutputTimestamp());
|
CurrentOutputTimestamp(cc));
|
||||||
}
|
}
|
||||||
cumulative_completed_frames_ += output_vectors.size();
|
cumulative_completed_frames_ += output_vectors.size();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -66,4 +66,11 @@ message SpectrogramCalculatorOptions {
|
|||||||
// uniformly regardless of output type (i.e., even dBs are multiplied, not
|
// uniformly regardless of output type (i.e., even dBs are multiplied, not
|
||||||
// offset).
|
// offset).
|
||||||
optional double output_scale = 7 [default = 1.0];
|
optional double output_scale = 7 [default = 1.0];
|
||||||
|
|
||||||
|
// If use_local_timestamp is true, the output packet's timestamp is based on
|
||||||
|
// the last sample of the packet and it's inferred from the latest input
|
||||||
|
// packet's timestamp. If false, the output packet's timestamp is based on
|
||||||
|
// the cumulative timestamping, which is inferred from the intial input
|
||||||
|
// timestamp and the cumulative number of samples.
|
||||||
|
optional bool use_local_timestamp = 8 [default = false];
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -75,8 +75,13 @@ class StabilizedLogCalculator : public CalculatorBase {
|
|||||||
|
|
||||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
auto input_matrix = cc->Inputs().Index(0).Get<Matrix>();
|
auto input_matrix = cc->Inputs().Index(0).Get<Matrix>();
|
||||||
|
if (input_matrix.array().isNaN().any()) {
|
||||||
|
return ::mediapipe::InvalidArgumentError("NaN input to log operation.");
|
||||||
|
}
|
||||||
if (check_nonnegativity_) {
|
if (check_nonnegativity_) {
|
||||||
CHECK_GE(input_matrix.minCoeff(), 0);
|
if (input_matrix.minCoeff() < 0.0) {
|
||||||
|
return ::mediapipe::OutOfRangeError("Negative input to log operation.");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
std::unique_ptr<Matrix> output_frame(new Matrix(
|
std::unique_ptr<Matrix> output_frame(new Matrix(
|
||||||
output_scale_ * (input_matrix.array() + stabilizer_).log().matrix()));
|
output_scale_ * (input_matrix.array() + stabilizer_).log().matrix()));
|
||||||
|
|||||||
@@ -11,6 +11,7 @@
|
|||||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
// 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 <cmath>
|
||||||
|
|
||||||
#include "Eigen/Core"
|
#include "Eigen/Core"
|
||||||
#include "mediapipe/calculators/audio/stabilized_log_calculator.pb.h"
|
#include "mediapipe/calculators/audio/stabilized_log_calculator.pb.h"
|
||||||
@@ -108,13 +109,22 @@ TEST_F(StabilizedLogCalculatorTest, ZerosAreStabilized) {
|
|||||||
runner_->Outputs().Index(0).packets[0].Get<Matrix>());
|
runner_->Outputs().Index(0).packets[0].Get<Matrix>());
|
||||||
}
|
}
|
||||||
|
|
||||||
TEST_F(StabilizedLogCalculatorTest, NegativeValuesCheckFail) {
|
TEST_F(StabilizedLogCalculatorTest, NanValuesReturnError) {
|
||||||
|
InitializeGraph();
|
||||||
|
FillInputHeader();
|
||||||
|
AppendInputPacket(
|
||||||
|
new Matrix(Matrix::Constant(kNumChannels, kNumSamples, std::nanf(""))),
|
||||||
|
0 /* timestamp */);
|
||||||
|
ASSERT_FALSE(RunGraph().ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(StabilizedLogCalculatorTest, NegativeValuesReturnError) {
|
||||||
InitializeGraph();
|
InitializeGraph();
|
||||||
FillInputHeader();
|
FillInputHeader();
|
||||||
AppendInputPacket(
|
AppendInputPacket(
|
||||||
new Matrix(Matrix::Constant(kNumChannels, kNumSamples, -1.0)),
|
new Matrix(Matrix::Constant(kNumChannels, kNumSamples, -1.0)),
|
||||||
0 /* timestamp */);
|
0 /* timestamp */);
|
||||||
ASSERT_DEATH(RunGraphNoReturn(), "");
|
ASSERT_FALSE(RunGraph().ok());
|
||||||
}
|
}
|
||||||
|
|
||||||
TEST_F(StabilizedLogCalculatorTest, NegativeValuesDoNotCheckFailIfCheckIsOff) {
|
TEST_F(StabilizedLogCalculatorTest, NegativeValuesDoNotCheckFailIfCheckIsOff) {
|
||||||
|
|||||||
@@ -56,6 +56,14 @@ namespace mediapipe {
|
|||||||
// If pad_final_packet is true, all input samples will be emitted and the final
|
// If pad_final_packet is true, all input samples will be emitted and the final
|
||||||
// packet will be zero padded as necessary. If pad_final_packet is false, some
|
// packet will be zero padded as necessary. If pad_final_packet is false, some
|
||||||
// samples may be dropped at the end of the stream.
|
// samples may be dropped at the end of the stream.
|
||||||
|
//
|
||||||
|
// If use_local_timestamp is true, the output packet's timestamp is based on the
|
||||||
|
// last sample of the packet. The timestamp of this sample is inferred by
|
||||||
|
// input_packet_timesamp + local_sample_index / sampling_rate_. If false, the
|
||||||
|
// output packet's timestamp is based on the cumulative timestamping, which is
|
||||||
|
// done by adopting the timestamp of the first sample of the packet and this
|
||||||
|
// sample's timestamp is inferred by initial_input_timestamp_ +
|
||||||
|
// cumulative_completed_samples / sample_rate_.
|
||||||
class TimeSeriesFramerCalculator : public CalculatorBase {
|
class TimeSeriesFramerCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
@@ -86,11 +94,26 @@ class TimeSeriesFramerCalculator : public CalculatorBase {
|
|||||||
void FrameOutput(CalculatorContext* cc);
|
void FrameOutput(CalculatorContext* cc);
|
||||||
|
|
||||||
Timestamp CurrentOutputTimestamp() {
|
Timestamp CurrentOutputTimestamp() {
|
||||||
|
if (use_local_timestamp_) {
|
||||||
|
return current_timestamp_;
|
||||||
|
}
|
||||||
|
return CumulativeOutputTimestamp();
|
||||||
|
}
|
||||||
|
|
||||||
|
Timestamp CumulativeOutputTimestamp() {
|
||||||
return initial_input_timestamp_ +
|
return initial_input_timestamp_ +
|
||||||
round(cumulative_completed_samples_ / sample_rate_ *
|
round(cumulative_completed_samples_ / sample_rate_ *
|
||||||
Timestamp::kTimestampUnitsPerSecond);
|
Timestamp::kTimestampUnitsPerSecond);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Returns the timestamp of a sample on a base, which is usually the time
|
||||||
|
// stamp of a packet.
|
||||||
|
Timestamp CurrentSampleTimestamp(const Timestamp& timestamp_base,
|
||||||
|
int64 number_of_samples) {
|
||||||
|
return timestamp_base + round(number_of_samples / sample_rate_ *
|
||||||
|
Timestamp::kTimestampUnitsPerSecond);
|
||||||
|
}
|
||||||
|
|
||||||
// The number of input samples to advance after the current output frame is
|
// The number of input samples to advance after the current output frame is
|
||||||
// emitted.
|
// emitted.
|
||||||
int next_frame_step_samples() const {
|
int next_frame_step_samples() const {
|
||||||
@@ -118,14 +141,18 @@ class TimeSeriesFramerCalculator : public CalculatorBase {
|
|||||||
// any overlap).
|
// any overlap).
|
||||||
int64 cumulative_completed_samples_;
|
int64 cumulative_completed_samples_;
|
||||||
Timestamp initial_input_timestamp_;
|
Timestamp initial_input_timestamp_;
|
||||||
|
// The current timestamp is updated along with the incoming packets.
|
||||||
|
Timestamp current_timestamp_;
|
||||||
int num_channels_;
|
int num_channels_;
|
||||||
|
|
||||||
// Each entry in this deque consists of a single sample, i.e. a
|
// Each entry in this deque consists of a single sample, i.e. a
|
||||||
// single column vector.
|
// single column vector, and its timestamp.
|
||||||
std::deque<Matrix> sample_buffer_;
|
std::deque<std::pair<Matrix, Timestamp>> sample_buffer_;
|
||||||
|
|
||||||
bool use_window_;
|
bool use_window_;
|
||||||
Matrix window_;
|
Matrix window_;
|
||||||
|
|
||||||
|
bool use_local_timestamp_;
|
||||||
};
|
};
|
||||||
REGISTER_CALCULATOR(TimeSeriesFramerCalculator);
|
REGISTER_CALCULATOR(TimeSeriesFramerCalculator);
|
||||||
|
|
||||||
@@ -133,7 +160,8 @@ void TimeSeriesFramerCalculator::EnqueueInput(CalculatorContext* cc) {
|
|||||||
const Matrix& input_frame = cc->Inputs().Index(0).Get<Matrix>();
|
const Matrix& input_frame = cc->Inputs().Index(0).Get<Matrix>();
|
||||||
|
|
||||||
for (int i = 0; i < input_frame.cols(); ++i) {
|
for (int i = 0; i < input_frame.cols(); ++i) {
|
||||||
sample_buffer_.emplace_back(input_frame.col(i));
|
sample_buffer_.emplace_back(std::make_pair(
|
||||||
|
input_frame.col(i), CurrentSampleTimestamp(cc->InputTimestamp(), i)));
|
||||||
}
|
}
|
||||||
|
|
||||||
cumulative_input_samples_ += input_frame.cols();
|
cumulative_input_samples_ += input_frame.cols();
|
||||||
@@ -151,14 +179,16 @@ void TimeSeriesFramerCalculator::FrameOutput(CalculatorContext* cc) {
|
|||||||
new Matrix(num_channels_, frame_duration_samples_));
|
new Matrix(num_channels_, frame_duration_samples_));
|
||||||
for (int i = 0; i < std::min(frame_step_samples, frame_duration_samples_);
|
for (int i = 0; i < std::min(frame_step_samples, frame_duration_samples_);
|
||||||
++i) {
|
++i) {
|
||||||
output_frame->col(i) = sample_buffer_.front();
|
output_frame->col(i) = sample_buffer_.front().first;
|
||||||
|
current_timestamp_ = sample_buffer_.front().second;
|
||||||
sample_buffer_.pop_front();
|
sample_buffer_.pop_front();
|
||||||
}
|
}
|
||||||
const int frame_overlap_samples =
|
const int frame_overlap_samples =
|
||||||
frame_duration_samples_ - frame_step_samples;
|
frame_duration_samples_ - frame_step_samples;
|
||||||
if (frame_overlap_samples > 0) {
|
if (frame_overlap_samples > 0) {
|
||||||
for (int i = 0; i < frame_overlap_samples; ++i) {
|
for (int i = 0; i < frame_overlap_samples; ++i) {
|
||||||
output_frame->col(i + frame_step_samples) = sample_buffer_[i];
|
output_frame->col(i + frame_step_samples) = sample_buffer_[i].first;
|
||||||
|
current_timestamp_ = sample_buffer_[i].second;
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
samples_still_to_drop_ = -frame_overlap_samples;
|
samples_still_to_drop_ = -frame_overlap_samples;
|
||||||
@@ -178,6 +208,7 @@ void TimeSeriesFramerCalculator::FrameOutput(CalculatorContext* cc) {
|
|||||||
::mediapipe::Status TimeSeriesFramerCalculator::Process(CalculatorContext* cc) {
|
::mediapipe::Status TimeSeriesFramerCalculator::Process(CalculatorContext* cc) {
|
||||||
if (initial_input_timestamp_ == Timestamp::Unstarted()) {
|
if (initial_input_timestamp_ == Timestamp::Unstarted()) {
|
||||||
initial_input_timestamp_ = cc->InputTimestamp();
|
initial_input_timestamp_ = cc->InputTimestamp();
|
||||||
|
current_timestamp_ = initial_input_timestamp_;
|
||||||
}
|
}
|
||||||
|
|
||||||
EnqueueInput(cc);
|
EnqueueInput(cc);
|
||||||
@@ -195,7 +226,8 @@ void TimeSeriesFramerCalculator::FrameOutput(CalculatorContext* cc) {
|
|||||||
std::unique_ptr<Matrix> output_frame(new Matrix);
|
std::unique_ptr<Matrix> output_frame(new Matrix);
|
||||||
output_frame->setZero(num_channels_, frame_duration_samples_);
|
output_frame->setZero(num_channels_, frame_duration_samples_);
|
||||||
for (int i = 0; i < sample_buffer_.size(); ++i) {
|
for (int i = 0; i < sample_buffer_.size(); ++i) {
|
||||||
output_frame->col(i) = sample_buffer_[i];
|
output_frame->col(i) = sample_buffer_[i].first;
|
||||||
|
current_timestamp_ = sample_buffer_[i].second;
|
||||||
}
|
}
|
||||||
|
|
||||||
cc->Outputs().Index(0).Add(output_frame.release(),
|
cc->Outputs().Index(0).Add(output_frame.release(),
|
||||||
@@ -258,6 +290,7 @@ void TimeSeriesFramerCalculator::FrameOutput(CalculatorContext* cc) {
|
|||||||
cumulative_output_frames_ = 0;
|
cumulative_output_frames_ = 0;
|
||||||
samples_still_to_drop_ = 0;
|
samples_still_to_drop_ = 0;
|
||||||
initial_input_timestamp_ = Timestamp::Unstarted();
|
initial_input_timestamp_ = Timestamp::Unstarted();
|
||||||
|
current_timestamp_ = Timestamp::Unstarted();
|
||||||
|
|
||||||
std::vector<double> window_vector;
|
std::vector<double> window_vector;
|
||||||
use_window_ = false;
|
use_window_ = false;
|
||||||
@@ -282,6 +315,7 @@ void TimeSeriesFramerCalculator::FrameOutput(CalculatorContext* cc) {
|
|||||||
frame_duration_samples_)
|
frame_duration_samples_)
|
||||||
.cast<float>();
|
.cast<float>();
|
||||||
}
|
}
|
||||||
|
use_local_timestamp_ = framer_options.use_local_timestamp();
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -62,4 +62,11 @@ message TimeSeriesFramerCalculatorOptions {
|
|||||||
HANN = 2;
|
HANN = 2;
|
||||||
}
|
}
|
||||||
optional WindowFunction window_function = 4 [default = NONE];
|
optional WindowFunction window_function = 4 [default = NONE];
|
||||||
|
|
||||||
|
// If use_local_timestamp is true, the output packet's timestamp is based on
|
||||||
|
// the last sample of the packet and it's inferred from the latest input
|
||||||
|
// packet's timestamp. If false, the output packet's timestamp is based on
|
||||||
|
// the cumulative timestamping, which is inferred from the intial input
|
||||||
|
// timestamp and the cumulative number of samples.
|
||||||
|
optional bool use_local_timestamp = 6 [default = false];
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -35,6 +35,8 @@ namespace mediapipe {
|
|||||||
namespace {
|
namespace {
|
||||||
|
|
||||||
const int kInitialTimestampOffsetMicroseconds = 4;
|
const int kInitialTimestampOffsetMicroseconds = 4;
|
||||||
|
const int kGapBetweenPacketsInSeconds = 1;
|
||||||
|
const int kUniversalInputPacketSize = 50;
|
||||||
|
|
||||||
class TimeSeriesFramerCalculatorTest
|
class TimeSeriesFramerCalculatorTest
|
||||||
: public TimeSeriesCalculatorTest<TimeSeriesFramerCalculatorOptions> {
|
: public TimeSeriesCalculatorTest<TimeSeriesFramerCalculatorOptions> {
|
||||||
@@ -391,5 +393,93 @@ TEST_F(TimeSeriesFramerCalculatorWindowingSanityTest, HannWindowSanityCheck) {
|
|||||||
RunAndTestSinglePacketAverage(0.5f);
|
RunAndTestSinglePacketAverage(0.5f);
|
||||||
}
|
}
|
||||||
|
|
||||||
} // anonymous namespace
|
// A simple test class that checks the local packet time stamp. This class
|
||||||
|
// generate a series of packets with and without gaps between packets and tests
|
||||||
|
// the behavior with cumulative timestamping and local packet timestamping.
|
||||||
|
class TimeSeriesFramerCalculatorTimestampingTest
|
||||||
|
: public TimeSeriesFramerCalculatorTest {
|
||||||
|
protected:
|
||||||
|
// Creates test input and saves a reference copy.
|
||||||
|
void InitializeInputForTimeStampingTest() {
|
||||||
|
concatenated_input_samples_.resize(0, num_input_channels_);
|
||||||
|
num_input_samples_ = 0;
|
||||||
|
for (int i = 0; i < 10; ++i) {
|
||||||
|
// This range of packet sizes was chosen such that some input
|
||||||
|
// packets will be smaller than the output packet size and other
|
||||||
|
// input packets will be larger.
|
||||||
|
int packet_size = kUniversalInputPacketSize;
|
||||||
|
double timestamp_seconds = kInitialTimestampOffsetMicroseconds * 1.0e-6 +
|
||||||
|
num_input_samples_ / input_sample_rate_;
|
||||||
|
if (options_.use_local_timestamp()) {
|
||||||
|
timestamp_seconds += kGapBetweenPacketsInSeconds * i;
|
||||||
|
}
|
||||||
|
|
||||||
|
Matrix* data_frame =
|
||||||
|
NewTestFrame(num_input_channels_, packet_size, timestamp_seconds);
|
||||||
|
|
||||||
|
AppendInputPacket(data_frame, round(timestamp_seconds *
|
||||||
|
Timestamp::kTimestampUnitsPerSecond));
|
||||||
|
num_input_samples_ += packet_size;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
void CheckOutputTimestamps() {
|
||||||
|
int num_full_packets = output().packets.size();
|
||||||
|
if (options_.pad_final_packet()) {
|
||||||
|
num_full_packets -= 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
int64 num_samples = 0;
|
||||||
|
for (int packet_num = 0; packet_num < num_full_packets; ++packet_num) {
|
||||||
|
const Packet& packet = output().packets[packet_num];
|
||||||
|
num_samples += FrameDurationSamples();
|
||||||
|
double expected_timestamp =
|
||||||
|
options_.use_local_timestamp()
|
||||||
|
? GetExpectedLocalTimestampForSample(num_samples - 1)
|
||||||
|
: GetExpectedCumulativeTimestamp(num_samples - 1);
|
||||||
|
ASSERT_NEAR(packet.Timestamp().Seconds(), expected_timestamp, 1e-10);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status RunTimestampTest() {
|
||||||
|
InitializeGraph();
|
||||||
|
InitializeInputForTimeStampingTest();
|
||||||
|
FillInputHeader();
|
||||||
|
return RunGraph();
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
// Returns the timestamp in seconds based on local timestamping.
|
||||||
|
double GetExpectedLocalTimestampForSample(int sample_index) {
|
||||||
|
return kInitialTimestampOffsetMicroseconds * 1.0e-6 +
|
||||||
|
sample_index / input_sample_rate_ +
|
||||||
|
(sample_index / kUniversalInputPacketSize) *
|
||||||
|
kGapBetweenPacketsInSeconds;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Returns the timestamp inseconds based on cumulative timestamping.
|
||||||
|
double GetExpectedCumulativeTimestamp(int sample_index) {
|
||||||
|
return kInitialTimestampOffsetMicroseconds * 1.0e-6 +
|
||||||
|
sample_index / FrameDurationSamples() * FrameDurationSamples() /
|
||||||
|
input_sample_rate_;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
TEST_F(TimeSeriesFramerCalculatorTimestampingTest, UseLocalTimeStamp) {
|
||||||
|
options_.set_frame_duration_seconds(100.0 / input_sample_rate_);
|
||||||
|
options_.set_use_local_timestamp(true);
|
||||||
|
|
||||||
|
MP_ASSERT_OK(RunTimestampTest());
|
||||||
|
CheckOutputTimestamps();
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(TimeSeriesFramerCalculatorTimestampingTest, UseCumulativeTimeStamp) {
|
||||||
|
options_.set_frame_duration_seconds(100.0 / input_sample_rate_);
|
||||||
|
options_.set_use_local_timestamp(false);
|
||||||
|
|
||||||
|
MP_ASSERT_OK(RunTimestampTest());
|
||||||
|
CheckOutputTimestamps();
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -13,12 +13,12 @@
|
|||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
#
|
#
|
||||||
|
|
||||||
|
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||||
|
|
||||||
licenses(["notice"]) # Apache 2.0
|
licenses(["notice"]) # Apache 2.0
|
||||||
|
|
||||||
package(default_visibility = ["//visibility:private"])
|
package(default_visibility = ["//visibility:private"])
|
||||||
|
|
||||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
|
||||||
|
|
||||||
proto_library(
|
proto_library(
|
||||||
name = "concatenate_vector_calculator_proto",
|
name = "concatenate_vector_calculator_proto",
|
||||||
srcs = ["concatenate_vector_calculator.proto"],
|
srcs = ["concatenate_vector_calculator.proto"],
|
||||||
@@ -26,6 +26,13 @@ proto_library(
|
|||||||
deps = ["//mediapipe/framework:calculator_proto"],
|
deps = ["//mediapipe/framework:calculator_proto"],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
proto_library(
|
||||||
|
name = "dequantize_byte_array_calculator_proto",
|
||||||
|
srcs = ["dequantize_byte_array_calculator.proto"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = ["//mediapipe/framework:calculator_proto"],
|
||||||
|
)
|
||||||
|
|
||||||
proto_library(
|
proto_library(
|
||||||
name = "packet_cloner_calculator_proto",
|
name = "packet_cloner_calculator_proto",
|
||||||
srcs = ["packet_cloner_calculator.proto"],
|
srcs = ["packet_cloner_calculator.proto"],
|
||||||
@@ -40,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"],
|
||||||
@@ -72,6 +86,22 @@ 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(
|
||||||
|
name = "clip_vector_size_calculator_proto",
|
||||||
|
srcs = ["clip_vector_size_calculator.proto"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = ["//mediapipe/framework:calculator_proto"],
|
||||||
|
)
|
||||||
|
|
||||||
mediapipe_cc_proto_library(
|
mediapipe_cc_proto_library(
|
||||||
name = "packet_cloner_calculator_cc_proto",
|
name = "packet_cloner_calculator_cc_proto",
|
||||||
srcs = ["packet_cloner_calculator.proto"],
|
srcs = ["packet_cloner_calculator.proto"],
|
||||||
@@ -88,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"],
|
||||||
@@ -104,6 +142,22 @@ mediapipe_cc_proto_library(
|
|||||||
deps = [":concatenate_vector_calculator_proto"],
|
deps = [":concatenate_vector_calculator_proto"],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
mediapipe_cc_proto_library(
|
||||||
|
name = "clip_vector_size_calculator_cc_proto",
|
||||||
|
srcs = ["clip_vector_size_calculator.proto"],
|
||||||
|
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [":clip_vector_size_calculator_proto"],
|
||||||
|
)
|
||||||
|
|
||||||
|
mediapipe_cc_proto_library(
|
||||||
|
name = "dequantize_byte_array_calculator_cc_proto",
|
||||||
|
srcs = ["dequantize_byte_array_calculator.proto"],
|
||||||
|
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [":dequantize_byte_array_calculator_proto"],
|
||||||
|
)
|
||||||
|
|
||||||
mediapipe_cc_proto_library(
|
mediapipe_cc_proto_library(
|
||||||
name = "quantize_float_vector_calculator_cc_proto",
|
name = "quantize_float_vector_calculator_cc_proto",
|
||||||
srcs = ["quantize_float_vector_calculator.proto"],
|
srcs = ["quantize_float_vector_calculator.proto"],
|
||||||
@@ -128,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"],
|
||||||
@@ -135,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,
|
||||||
)
|
)
|
||||||
@@ -154,6 +217,66 @@ cc_test(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "begin_loop_calculator",
|
||||||
|
srcs = ["begin_loop_calculator.cc"],
|
||||||
|
hdrs = ["begin_loop_calculator.h"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_context",
|
||||||
|
"//mediapipe/framework:calculator_contract",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:collection_item_id",
|
||||||
|
"//mediapipe/framework:packet",
|
||||||
|
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||||
|
"//mediapipe/framework/formats:rect_cc_proto",
|
||||||
|
"//mediapipe/framework/port:integral_types",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"@com_google_absl//absl/memory",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "end_loop_calculator",
|
||||||
|
srcs = ["end_loop_calculator.cc"],
|
||||||
|
hdrs = ["end_loop_calculator.h"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_context",
|
||||||
|
"//mediapipe/framework:calculator_contract",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:collection_item_id",
|
||||||
|
"//mediapipe/framework:packet",
|
||||||
|
"//mediapipe/framework/formats:landmark_cc_proto",
|
||||||
|
"//mediapipe/framework/formats:rect_cc_proto",
|
||||||
|
"//mediapipe/framework/port:integral_types",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"//mediapipe/util:render_data_cc_proto",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_test(
|
||||||
|
name = "begin_end_loop_calculator_graph_test",
|
||||||
|
srcs = ["begin_end_loop_calculator_graph_test.cc"],
|
||||||
|
deps = [
|
||||||
|
":begin_loop_calculator",
|
||||||
|
":end_loop_calculator",
|
||||||
|
"//mediapipe/calculators/core:packet_cloner_calculator",
|
||||||
|
"//mediapipe/framework:calculator_context",
|
||||||
|
"//mediapipe/framework:calculator_contract",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/port:gtest_main",
|
||||||
|
"//mediapipe/framework/port:integral_types",
|
||||||
|
"//mediapipe/framework/port:parse_text_proto",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"@com_google_absl//absl/memory",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
cc_library(
|
cc_library(
|
||||||
name = "concatenate_vector_calculator",
|
name = "concatenate_vector_calculator",
|
||||||
srcs = ["concatenate_vector_calculator.cc"],
|
srcs = ["concatenate_vector_calculator.cc"],
|
||||||
@@ -166,7 +289,13 @@ cc_library(
|
|||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
"@org_tensorflow//tensorflow/lite:framework",
|
"@org_tensorflow//tensorflow/lite:framework",
|
||||||
],
|
] + select({
|
||||||
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
|
"//mediapipe:ios": [],
|
||||||
|
"//conditions:default": [
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
|
||||||
|
],
|
||||||
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -187,6 +316,49 @@ cc_test(
|
|||||||
srcs = ["concatenate_vector_calculator_test.cc"],
|
srcs = ["concatenate_vector_calculator_test.cc"],
|
||||||
deps = [
|
deps = [
|
||||||
":concatenate_vector_calculator",
|
":concatenate_vector_calculator",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:calculator_runner",
|
||||||
|
"//mediapipe/framework:timestamp",
|
||||||
|
"//mediapipe/framework/port:gtest_main",
|
||||||
|
"//mediapipe/framework/port:parse_text_proto",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"@com_google_absl//absl/strings",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "clip_vector_size_calculator",
|
||||||
|
srcs = ["clip_vector_size_calculator.cc"],
|
||||||
|
hdrs = ["clip_vector_size_calculator.h"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
":clip_vector_size_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/formats:rect_cc_proto",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"@org_tensorflow//tensorflow/lite:framework",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "clip_detection_vector_size_calculator",
|
||||||
|
srcs = ["clip_detection_vector_size_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
":clip_vector_size_calculator",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/formats:detection_cc_proto",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_test(
|
||||||
|
name = "clip_vector_size_calculator_test",
|
||||||
|
srcs = ["clip_vector_size_calculator_test.cc"],
|
||||||
|
deps = [
|
||||||
|
":clip_vector_size_calculator",
|
||||||
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
|
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
@@ -279,7 +451,7 @@ cc_library(
|
|||||||
"//visibility:public",
|
"//visibility:public",
|
||||||
],
|
],
|
||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/calculators/core:packet_cloner_calculator_cc_proto",
|
":packet_cloner_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
],
|
],
|
||||||
@@ -310,6 +482,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"],
|
||||||
@@ -381,6 +584,48 @@ cc_library(
|
|||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "string_to_int_calculator",
|
||||||
|
srcs = ["string_to_int_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/port:integral_types",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"@com_google_absl//absl/strings",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "side_packet_to_stream_calculator",
|
||||||
|
srcs = ["side_packet_to_stream_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/port:logging",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_test(
|
||||||
|
name = "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"],
|
||||||
@@ -405,6 +650,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",
|
||||||
],
|
],
|
||||||
@@ -428,17 +674,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",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
@@ -525,12 +771,19 @@ cc_library(
|
|||||||
":split_vector_calculator_cc_proto",
|
":split_vector_calculator_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/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,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -552,6 +805,32 @@ cc_test(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "dequantize_byte_array_calculator",
|
||||||
|
srcs = ["dequantize_byte_array_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
":dequantize_byte_array_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework:calculator_context",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_test(
|
||||||
|
name = "dequantize_byte_array_calculator_test",
|
||||||
|
srcs = ["dequantize_byte_array_calculator_test.cc"],
|
||||||
|
deps = [
|
||||||
|
":dequantize_byte_array_calculator",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:calculator_runner",
|
||||||
|
"//mediapipe/framework/port:gtest_main",
|
||||||
|
"//mediapipe/framework/port:parse_text_proto",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
cc_library(
|
cc_library(
|
||||||
name = "quantize_float_vector_calculator",
|
name = "quantize_float_vector_calculator",
|
||||||
srcs = ["quantize_float_vector_calculator.cc"],
|
srcs = ["quantize_float_vector_calculator.cc"],
|
||||||
@@ -688,3 +967,56 @@ cc_test(
|
|||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "stream_to_side_packet_calculator",
|
||||||
|
srcs = ["stream_to_side_packet_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:timestamp",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
|
cc_test(
|
||||||
|
name = "stream_to_side_packet_calculator_test",
|
||||||
|
srcs = ["stream_to_side_packet_calculator_test.cc"],
|
||||||
|
deps = [
|
||||||
|
":stream_to_side_packet_calculator",
|
||||||
|
"//mediapipe/framework:calculator_runner",
|
||||||
|
"//mediapipe/framework:packet",
|
||||||
|
"//mediapipe/framework:timestamp",
|
||||||
|
"//mediapipe/framework/port:gtest_main",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"@com_google_absl//absl/memory",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
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
|
||||||
|
|||||||
@@ -0,0 +1,335 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "absl/memory/memory.h"
|
||||||
|
#include "mediapipe/calculators/core/begin_loop_calculator.h"
|
||||||
|
#include "mediapipe/calculators/core/end_loop_calculator.h"
|
||||||
|
#include "mediapipe/framework/calculator_contract.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "mediapipe/framework/port/integral_types.h"
|
||||||
|
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||||
|
#include "mediapipe/framework/port/status_matchers.h" // NOLINT
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
typedef BeginLoopCalculator<std::vector<int>> BeginLoopIntegerCalculator;
|
||||||
|
REGISTER_CALCULATOR(BeginLoopIntegerCalculator);
|
||||||
|
|
||||||
|
class IncrementCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
cc->Inputs().Index(0).Set<int>();
|
||||||
|
cc->Outputs().Index(0).Set<int>();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
|
cc->SetOffset(TimestampDiff(0));
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
|
const int& input_int = cc->Inputs().Index(0).Get<int>();
|
||||||
|
auto output_int = absl::make_unique<int>(input_int + 1);
|
||||||
|
cc->Outputs().Index(0).Add(output_int.release(), cc->InputTimestamp());
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
REGISTER_CALCULATOR(IncrementCalculator);
|
||||||
|
|
||||||
|
typedef EndLoopCalculator<std::vector<int>> EndLoopIntegersCalculator;
|
||||||
|
REGISTER_CALCULATOR(EndLoopIntegersCalculator);
|
||||||
|
|
||||||
|
class BeginEndLoopCalculatorGraphTest : public ::testing::Test {
|
||||||
|
protected:
|
||||||
|
BeginEndLoopCalculatorGraphTest() {
|
||||||
|
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
num_threads: 4
|
||||||
|
input_stream: "ints"
|
||||||
|
node {
|
||||||
|
calculator: "BeginLoopIntegerCalculator"
|
||||||
|
input_stream: "ITERABLE:ints"
|
||||||
|
output_stream: "ITEM:int"
|
||||||
|
output_stream: "BATCH_END:timestamp"
|
||||||
|
}
|
||||||
|
node {
|
||||||
|
calculator: "IncrementCalculator"
|
||||||
|
input_stream: "int"
|
||||||
|
output_stream: "int_plus_one"
|
||||||
|
}
|
||||||
|
node {
|
||||||
|
calculator: "EndLoopIntegersCalculator"
|
||||||
|
input_stream: "ITEM:int_plus_one"
|
||||||
|
input_stream: "BATCH_END:timestamp"
|
||||||
|
output_stream: "ITERABLE:ints_plus_one"
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
tool::AddVectorSink("ints_plus_one", &graph_config_, &output_packets_);
|
||||||
|
}
|
||||||
|
|
||||||
|
CalculatorGraphConfig graph_config_;
|
||||||
|
std::vector<Packet> output_packets_;
|
||||||
|
};
|
||||||
|
|
||||||
|
TEST_F(BeginEndLoopCalculatorGraphTest, SingleEmptyVector) {
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||||
|
MP_EXPECT_OK(graph.StartRun({}));
|
||||||
|
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||||
|
Timestamp input_timestamp = Timestamp(0);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
|
||||||
|
// EndLoopCalc will forward the timestamp bound because there are no elements
|
||||||
|
// in collection to output.
|
||||||
|
ASSERT_EQ(0, output_packets_.size());
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(BeginEndLoopCalculatorGraphTest, SingleNonEmptyVector) {
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||||
|
MP_EXPECT_OK(graph.StartRun({}));
|
||||||
|
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||||
|
input_vector->emplace_back(0);
|
||||||
|
input_vector->emplace_back(1);
|
||||||
|
input_vector->emplace_back(2);
|
||||||
|
Timestamp input_timestamp = Timestamp(0);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
|
||||||
|
ASSERT_EQ(1, output_packets_.size());
|
||||||
|
EXPECT_EQ(input_timestamp, output_packets_[0].Timestamp());
|
||||||
|
std::vector<int> expected_output_vector = {1, 2, 3};
|
||||||
|
EXPECT_EQ(expected_output_vector, output_packets_[0].Get<std::vector<int>>());
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(BeginEndLoopCalculatorGraphTest, MultipleVectors) {
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||||
|
MP_EXPECT_OK(graph.StartRun({}));
|
||||||
|
|
||||||
|
auto input_vector0 = absl::make_unique<std::vector<int>>();
|
||||||
|
input_vector0->emplace_back(0);
|
||||||
|
input_vector0->emplace_back(1);
|
||||||
|
Timestamp input_timestamp0 = Timestamp(0);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"ints", Adopt(input_vector0.release()).At(input_timestamp0)));
|
||||||
|
|
||||||
|
auto input_vector1 = absl::make_unique<std::vector<int>>();
|
||||||
|
Timestamp input_timestamp1 = Timestamp(1);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"ints", Adopt(input_vector1.release()).At(input_timestamp1)));
|
||||||
|
|
||||||
|
auto input_vector2 = absl::make_unique<std::vector<int>>();
|
||||||
|
input_vector2->emplace_back(2);
|
||||||
|
input_vector2->emplace_back(3);
|
||||||
|
Timestamp input_timestamp2 = Timestamp(2);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"ints", Adopt(input_vector2.release()).At(input_timestamp2)));
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
|
|
||||||
|
ASSERT_EQ(2, output_packets_.size());
|
||||||
|
|
||||||
|
EXPECT_EQ(input_timestamp0, output_packets_[0].Timestamp());
|
||||||
|
std::vector<int> expected_output_vector0 = {1, 2};
|
||||||
|
EXPECT_EQ(expected_output_vector0,
|
||||||
|
output_packets_[0].Get<std::vector<int>>());
|
||||||
|
|
||||||
|
// At input_timestamp1, EndLoopCalc will forward timestamp bound as there are
|
||||||
|
// no elements in vector to process.
|
||||||
|
|
||||||
|
EXPECT_EQ(input_timestamp2, output_packets_[1].Timestamp());
|
||||||
|
std::vector<int> expected_output_vector2 = {3, 4};
|
||||||
|
EXPECT_EQ(expected_output_vector2,
|
||||||
|
output_packets_[1].Get<std::vector<int>>());
|
||||||
|
}
|
||||||
|
|
||||||
|
class MultiplierCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
cc->Inputs().Index(0).Set<int>();
|
||||||
|
cc->Inputs().Index(1).Set<int>();
|
||||||
|
cc->Outputs().Index(0).Set<int>();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
|
cc->SetOffset(TimestampDiff(0));
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
|
const int& input_int = cc->Inputs().Index(0).Get<int>();
|
||||||
|
const int& multiplier_int = cc->Inputs().Index(1).Get<int>();
|
||||||
|
auto output_int = absl::make_unique<int>(input_int * multiplier_int);
|
||||||
|
cc->Outputs().Index(0).Add(output_int.release(), cc->InputTimestamp());
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
REGISTER_CALCULATOR(MultiplierCalculator);
|
||||||
|
|
||||||
|
class BeginEndLoopCalculatorGraphWithClonedInputsTest : public ::testing::Test {
|
||||||
|
protected:
|
||||||
|
BeginEndLoopCalculatorGraphWithClonedInputsTest() {
|
||||||
|
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
num_threads: 4
|
||||||
|
input_stream: "ints"
|
||||||
|
input_stream: "multiplier"
|
||||||
|
node {
|
||||||
|
calculator: "BeginLoopIntegerCalculator"
|
||||||
|
input_stream: "ITERABLE:ints"
|
||||||
|
input_stream: "CLONE:multiplier"
|
||||||
|
output_stream: "ITEM:int_at_loop"
|
||||||
|
output_stream: "CLONE:multiplier_cloned_at_loop"
|
||||||
|
output_stream: "BATCH_END:timestamp"
|
||||||
|
}
|
||||||
|
node {
|
||||||
|
calculator: "MultiplierCalculator"
|
||||||
|
input_stream: "int_at_loop"
|
||||||
|
input_stream: "multiplier_cloned_at_loop"
|
||||||
|
output_stream: "multiplied_int_at_loop"
|
||||||
|
}
|
||||||
|
node {
|
||||||
|
calculator: "EndLoopIntegersCalculator"
|
||||||
|
input_stream: "ITEM:multiplied_int_at_loop"
|
||||||
|
input_stream: "BATCH_END:timestamp"
|
||||||
|
output_stream: "ITERABLE:multiplied_ints"
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
tool::AddVectorSink("multiplied_ints", &graph_config_, &output_packets_);
|
||||||
|
}
|
||||||
|
|
||||||
|
CalculatorGraphConfig graph_config_;
|
||||||
|
std::vector<Packet> output_packets_;
|
||||||
|
};
|
||||||
|
|
||||||
|
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, SingleEmptyVector) {
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||||
|
MP_EXPECT_OK(graph.StartRun({}));
|
||||||
|
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||||
|
Timestamp input_timestamp = Timestamp(42);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||||
|
auto multiplier = absl::make_unique<int>(2);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"multiplier", Adopt(multiplier.release()).At(input_timestamp)));
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
|
||||||
|
// EndLoopCalc will forward the timestamp bound because there are no elements
|
||||||
|
// in collection to output.
|
||||||
|
ASSERT_EQ(0, output_packets_.size());
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, SingleNonEmptyVector) {
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||||
|
MP_EXPECT_OK(graph.StartRun({}));
|
||||||
|
auto input_vector = absl::make_unique<std::vector<int>>();
|
||||||
|
input_vector->emplace_back(0);
|
||||||
|
input_vector->emplace_back(1);
|
||||||
|
input_vector->emplace_back(2);
|
||||||
|
Timestamp input_timestamp = Timestamp(42);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"ints", Adopt(input_vector.release()).At(input_timestamp)));
|
||||||
|
auto multiplier = absl::make_unique<int>(2);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"multiplier", Adopt(multiplier.release()).At(input_timestamp)));
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
|
||||||
|
ASSERT_EQ(1, output_packets_.size());
|
||||||
|
EXPECT_EQ(input_timestamp, output_packets_[0].Timestamp());
|
||||||
|
std::vector<int> expected_output_vector = {0, 2, 4};
|
||||||
|
EXPECT_EQ(expected_output_vector, output_packets_[0].Get<std::vector<int>>());
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(BeginEndLoopCalculatorGraphWithClonedInputsTest, MultipleVectors) {
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_EXPECT_OK(graph.Initialize(graph_config_));
|
||||||
|
MP_EXPECT_OK(graph.StartRun({}));
|
||||||
|
|
||||||
|
auto input_vector0 = absl::make_unique<std::vector<int>>();
|
||||||
|
input_vector0->emplace_back(0);
|
||||||
|
input_vector0->emplace_back(1);
|
||||||
|
Timestamp input_timestamp0 = Timestamp(42);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"ints", Adopt(input_vector0.release()).At(input_timestamp0)));
|
||||||
|
auto multiplier0 = absl::make_unique<int>(2);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"multiplier", Adopt(multiplier0.release()).At(input_timestamp0)));
|
||||||
|
|
||||||
|
auto input_vector1 = absl::make_unique<std::vector<int>>();
|
||||||
|
Timestamp input_timestamp1 = Timestamp(43);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"ints", Adopt(input_vector1.release()).At(input_timestamp1)));
|
||||||
|
auto multiplier1 = absl::make_unique<int>(2);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"multiplier", Adopt(multiplier1.release()).At(input_timestamp1)));
|
||||||
|
|
||||||
|
auto input_vector2 = absl::make_unique<std::vector<int>>();
|
||||||
|
input_vector2->emplace_back(2);
|
||||||
|
input_vector2->emplace_back(3);
|
||||||
|
Timestamp input_timestamp2 = Timestamp(44);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"ints", Adopt(input_vector2.release()).At(input_timestamp2)));
|
||||||
|
auto multiplier2 = absl::make_unique<int>(3);
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"multiplier", Adopt(multiplier2.release()).At(input_timestamp2)));
|
||||||
|
|
||||||
|
MP_ASSERT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
|
|
||||||
|
ASSERT_EQ(2, output_packets_.size());
|
||||||
|
|
||||||
|
EXPECT_EQ(input_timestamp0, output_packets_[0].Timestamp());
|
||||||
|
std::vector<int> expected_output_vector0 = {0, 2};
|
||||||
|
EXPECT_EQ(expected_output_vector0,
|
||||||
|
output_packets_[0].Get<std::vector<int>>());
|
||||||
|
|
||||||
|
// At input_timestamp1, EndLoopCalc will forward timestamp bound as there are
|
||||||
|
// no elements in vector to process.
|
||||||
|
|
||||||
|
EXPECT_EQ(input_timestamp2, output_packets_[1].Timestamp());
|
||||||
|
std::vector<int> expected_output_vector2 = {6, 9};
|
||||||
|
EXPECT_EQ(expected_output_vector2,
|
||||||
|
output_packets_[1].Get<std::vector<int>>());
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,34 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/core/begin_loop_calculator.h"
|
||||||
|
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||||
|
#include "mediapipe/framework/formats/rect.pb.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
// A calculator to process std::vector<NormalizedLandmarkList>.
|
||||||
|
typedef BeginLoopCalculator<std::vector<::mediapipe::NormalizedLandmarkList>>
|
||||||
|
BeginLoopNormalizedLandmarkListVectorCalculator;
|
||||||
|
REGISTER_CALCULATOR(BeginLoopNormalizedLandmarkListVectorCalculator);
|
||||||
|
|
||||||
|
// A calculator to process std::vector<NormalizedRect>.
|
||||||
|
typedef BeginLoopCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||||
|
BeginLoopNormalizedRectCalculator;
|
||||||
|
REGISTER_CALCULATOR(BeginLoopNormalizedRectCalculator);
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,157 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#ifndef MEDIAPIPE_CALCULATORS_CORE_BEGIN_LOOP_CALCULATOR_H_
|
||||||
|
#define MEDIAPIPE_CALCULATORS_CORE_BEGIN_LOOP_CALCULATOR_H_
|
||||||
|
|
||||||
|
#include "absl/memory/memory.h"
|
||||||
|
#include "mediapipe/framework/calculator_context.h"
|
||||||
|
#include "mediapipe/framework/calculator_contract.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/collection_item_id.h"
|
||||||
|
#include "mediapipe/framework/packet.h"
|
||||||
|
#include "mediapipe/framework/port/integral_types.h"
|
||||||
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
// Calculator for implementing loops on iterable collections inside a MediaPipe
|
||||||
|
// graph.
|
||||||
|
//
|
||||||
|
// It is designed to be used like:
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "BeginLoopWithIterableCalculator"
|
||||||
|
// input_stream: "ITERABLE:input_iterable" # IterableT @ext_ts
|
||||||
|
// output_stream: "ITEM:input_element" # ItemT @loop_internal_ts
|
||||||
|
// output_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "ElementToBlaConverterSubgraph"
|
||||||
|
// input_stream: "ITEM:input_to_loop_body" # ItemT @loop_internal_ts
|
||||||
|
// output_stream: "BLA:output_of_loop_body" # ItemU @loop_internal_ts
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "EndLoopWithOutputCalculator"
|
||||||
|
// input_stream: "ITEM:output_of_loop_body" # ItemU @loop_internal_ts
|
||||||
|
// input_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||||
|
// output_stream: "OUTPUT:aggregated_result" # IterableU @ext_ts
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// BeginLoopCalculator accepts an optional input stream tagged with "TICK"
|
||||||
|
// which if non-empty, wakes up the calculator and calls
|
||||||
|
// BeginLoopCalculator::Process(). Input streams tagged with "CLONE" are cloned
|
||||||
|
// to the corresponding output streams at loop timestamps. This ensures that a
|
||||||
|
// MediaPipe graph or sub-graph can run multiple times, once per element in the
|
||||||
|
// "ITERABLE" for each pakcet clone of the packets in the "CLONE" input streams.
|
||||||
|
template <typename IterableT>
|
||||||
|
class BeginLoopCalculator : public CalculatorBase {
|
||||||
|
using ItemT = typename IterableT::value_type;
|
||||||
|
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
// A non-empty packet in the optional "TICK" input stream wakes up the
|
||||||
|
// calculator.
|
||||||
|
if (cc->Inputs().HasTag("TICK")) {
|
||||||
|
cc->Inputs().Tag("TICK").SetAny();
|
||||||
|
}
|
||||||
|
|
||||||
|
// An iterable collection in the input stream.
|
||||||
|
RET_CHECK(cc->Inputs().HasTag("ITERABLE"));
|
||||||
|
cc->Inputs().Tag("ITERABLE").Set<IterableT>();
|
||||||
|
|
||||||
|
// An element from the collection.
|
||||||
|
RET_CHECK(cc->Outputs().HasTag("ITEM"));
|
||||||
|
cc->Outputs().Tag("ITEM").Set<ItemT>();
|
||||||
|
|
||||||
|
RET_CHECK(cc->Outputs().HasTag("BATCH_END"));
|
||||||
|
cc->Outputs()
|
||||||
|
.Tag("BATCH_END")
|
||||||
|
.Set<Timestamp>(
|
||||||
|
// A flush signal to the corresponding EndLoopCalculator for it to
|
||||||
|
// emit the aggregated result with the timestamp contained in this
|
||||||
|
// flush signal packet.
|
||||||
|
);
|
||||||
|
|
||||||
|
// Input streams tagged with "CLONE" are cloned to the corresponding
|
||||||
|
// "CLONE" output streams at loop timestamps.
|
||||||
|
RET_CHECK(cc->Inputs().NumEntries("CLONE") ==
|
||||||
|
cc->Outputs().NumEntries("CLONE"));
|
||||||
|
if (cc->Inputs().NumEntries("CLONE") > 0) {
|
||||||
|
for (int i = 0; i < cc->Inputs().NumEntries("CLONE"); ++i) {
|
||||||
|
cc->Inputs().Get("CLONE", i).SetAny();
|
||||||
|
cc->Outputs().Get("CLONE", i).SetSameAs(&cc->Inputs().Get("CLONE", i));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||||
|
Timestamp last_timestamp = loop_internal_timestamp_;
|
||||||
|
if (!cc->Inputs().Tag("ITERABLE").IsEmpty()) {
|
||||||
|
const IterableT& collection =
|
||||||
|
cc->Inputs().Tag("ITERABLE").template Get<IterableT>();
|
||||||
|
for (const auto& item : collection) {
|
||||||
|
cc->Outputs().Tag("ITEM").AddPacket(
|
||||||
|
MakePacket<ItemT>(item).At(loop_internal_timestamp_));
|
||||||
|
ForwardClonePackets(cc, loop_internal_timestamp_);
|
||||||
|
++loop_internal_timestamp_;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// The collection was empty and nothing was processed.
|
||||||
|
if (last_timestamp == loop_internal_timestamp_) {
|
||||||
|
// Increment loop_internal_timestamp_ because it is used up now.
|
||||||
|
++loop_internal_timestamp_;
|
||||||
|
for (auto it = cc->Outputs().begin(); it < cc->Outputs().end(); ++it) {
|
||||||
|
it->SetNextTimestampBound(loop_internal_timestamp_);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// The for loop processing the input collection already incremented
|
||||||
|
// loop_internal_timestamp_. To emit BATCH_END packet along the last
|
||||||
|
// non-BATCH_END packet, decrement by one.
|
||||||
|
cc->Outputs()
|
||||||
|
.Tag("BATCH_END")
|
||||||
|
.AddPacket(MakePacket<Timestamp>(cc->InputTimestamp())
|
||||||
|
.At(Timestamp(loop_internal_timestamp_ - 1)));
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
void ForwardClonePackets(CalculatorContext* cc, Timestamp output_timestamp) {
|
||||||
|
if (cc->Inputs().NumEntries("CLONE") > 0) {
|
||||||
|
for (int i = 0; i < cc->Inputs().NumEntries("CLONE"); ++i) {
|
||||||
|
if (!cc->Inputs().Get("CLONE", i).IsEmpty()) {
|
||||||
|
auto input_packet = cc->Inputs().Get("CLONE", i).Value();
|
||||||
|
cc->Outputs()
|
||||||
|
.Get("CLONE", i)
|
||||||
|
.AddPacket(std::move(input_packet).At(output_timestamp));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Fake timestamps generated per element in collection.
|
||||||
|
Timestamp loop_internal_timestamp_ = Timestamp(0);
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
|
|
||||||
|
#endif // MEDIAPIPE_CALCULATORS_CORE_BEGIN_LOOP_CALCULATOR_H_
|
||||||
@@ -0,0 +1,26 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/core/clip_vector_size_calculator.h"
|
||||||
|
#include "mediapipe/framework/formats/detection.pb.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
typedef ClipVectorSizeCalculator<::mediapipe::Detection>
|
||||||
|
ClipDetectionVectorSizeCalculator;
|
||||||
|
REGISTER_CALCULATOR(ClipDetectionVectorSizeCalculator);
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/core/clip_vector_size_calculator.h"
|
||||||
|
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/formats/rect.pb.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
typedef ClipVectorSizeCalculator<::mediapipe::NormalizedRect>
|
||||||
|
ClipNormalizedRectVectorSizeCalculator;
|
||||||
|
REGISTER_CALCULATOR(ClipNormalizedRectVectorSizeCalculator);
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,148 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#ifndef MEDIAPIPE_CALCULATORS_CORE_CLIP_VECTOR_SIZE_CALCULATOR_H_
|
||||||
|
#define MEDIAPIPE_CALCULATORS_CORE_CLIP_VECTOR_SIZE_CALCULATOR_H_
|
||||||
|
|
||||||
|
#include <type_traits>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/core/clip_vector_size_calculator.pb.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/canonical_errors.h"
|
||||||
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
// Clips the size of the input vector of type T to a specified max_vec_size.
|
||||||
|
// In a graph it will be used as:
|
||||||
|
// node {
|
||||||
|
// calculator: "ClipIntVectorSizeCalculator"
|
||||||
|
// input_stream: "input_vector"
|
||||||
|
// output_stream: "output_vector"
|
||||||
|
// options {
|
||||||
|
// [mediapipe.ClipIntVectorSizeCalculatorOptions.ext] {
|
||||||
|
// max_vec_size: 5
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
// Optionally, you can pass in a side packet that will override `max_vec_size`
|
||||||
|
// that is specified in the options.
|
||||||
|
template <typename T>
|
||||||
|
class ClipVectorSizeCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
RET_CHECK(cc->Inputs().NumEntries() == 1);
|
||||||
|
RET_CHECK(cc->Outputs().NumEntries() == 1);
|
||||||
|
|
||||||
|
if (cc->Options<::mediapipe::ClipVectorSizeCalculatorOptions>()
|
||||||
|
.max_vec_size() < 1) {
|
||||||
|
return ::mediapipe::InternalError(
|
||||||
|
"max_vec_size should be greater than or equal to 1.");
|
||||||
|
}
|
||||||
|
|
||||||
|
cc->Inputs().Index(0).Set<std::vector<T>>();
|
||||||
|
cc->Outputs().Index(0).Set<std::vector<T>>();
|
||||||
|
// Optional input side packet that determines `max_vec_size`.
|
||||||
|
if (cc->InputSidePackets().NumEntries() > 0) {
|
||||||
|
cc->InputSidePackets().Index(0).Set<int>();
|
||||||
|
}
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
|
cc->SetOffset(TimestampDiff(0));
|
||||||
|
max_vec_size_ = cc->Options<::mediapipe::ClipVectorSizeCalculatorOptions>()
|
||||||
|
.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();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
|
if (max_vec_size_ < 1) {
|
||||||
|
return ::mediapipe::InternalError(
|
||||||
|
"max_vec_size should be greater than or equal to 1.");
|
||||||
|
}
|
||||||
|
if (cc->Inputs().Index(0).IsEmpty()) {
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
return ClipVectorSize<T>(std::is_copy_constructible<T>(), cc);
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename U>
|
||||||
|
::mediapipe::Status ClipVectorSize(std::true_type, CalculatorContext* cc) {
|
||||||
|
auto output = absl::make_unique<std::vector<U>>();
|
||||||
|
const std::vector<U>& input_vector =
|
||||||
|
cc->Inputs().Index(0).Get<std::vector<U>>();
|
||||||
|
if (max_vec_size_ >= input_vector.size()) {
|
||||||
|
output->insert(output->end(), input_vector.begin(), input_vector.end());
|
||||||
|
} else {
|
||||||
|
for (int i = 0; i < max_vec_size_; ++i) {
|
||||||
|
output->push_back(input_vector[i]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename U>
|
||||||
|
::mediapipe::Status ClipVectorSize(std::false_type, CalculatorContext* cc) {
|
||||||
|
return ConsumeAndClipVectorSize<T>(std::is_move_constructible<U>(), cc);
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename U>
|
||||||
|
::mediapipe::Status ConsumeAndClipVectorSize(std::true_type,
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
auto output = absl::make_unique<std::vector<U>>();
|
||||||
|
::mediapipe::StatusOr<std::unique_ptr<std::vector<U>>> input_status =
|
||||||
|
cc->Inputs().Index(0).Value().Consume<std::vector<U>>();
|
||||||
|
|
||||||
|
if (input_status.ok()) {
|
||||||
|
std::unique_ptr<std::vector<U>> input_vector =
|
||||||
|
std::move(input_status).ValueOrDie();
|
||||||
|
auto begin_it = input_vector->begin();
|
||||||
|
auto end_it = input_vector->end();
|
||||||
|
if (max_vec_size_ < input_vector->size()) {
|
||||||
|
end_it = input_vector->begin() + max_vec_size_;
|
||||||
|
}
|
||||||
|
output->insert(output->end(), std::make_move_iterator(begin_it),
|
||||||
|
std::make_move_iterator(end_it));
|
||||||
|
} else {
|
||||||
|
return input_status.status();
|
||||||
|
}
|
||||||
|
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename U>
|
||||||
|
::mediapipe::Status ConsumeAndClipVectorSize(std::false_type,
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
return ::mediapipe::InternalError(
|
||||||
|
"Cannot copy or move input vectors and clip their size.");
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
int max_vec_size_ = 0;
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
|
|
||||||
|
#endif // MEDIAPIPE_CALCULATORS_CORE_CLIP_VECTOR_SIZE_CALCULATOR_H_
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
syntax = "proto2";
|
||||||
|
|
||||||
|
package mediapipe;
|
||||||
|
|
||||||
|
import "mediapipe/framework/calculator.proto";
|
||||||
|
|
||||||
|
message ClipVectorSizeCalculatorOptions {
|
||||||
|
extend CalculatorOptions {
|
||||||
|
optional ClipVectorSizeCalculatorOptions ext = 274674998;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Maximum size of output vector.
|
||||||
|
optional int32 max_vec_size = 1 [default = 1];
|
||||||
|
}
|
||||||
@@ -0,0 +1,206 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/core/clip_vector_size_calculator.h"
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/calculator_runner.h"
|
||||||
|
#include "mediapipe/framework/port/gmock.h"
|
||||||
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||||
|
#include "mediapipe/framework/port/status_matchers.h" // NOLINT
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
typedef ClipVectorSizeCalculator<int> TestClipIntVectorSizeCalculator;
|
||||||
|
REGISTER_CALCULATOR(TestClipIntVectorSizeCalculator);
|
||||||
|
|
||||||
|
void AddInputVector(const std::vector<int>& input, int64 timestamp,
|
||||||
|
CalculatorRunner* runner) {
|
||||||
|
runner->MutableInputs()->Index(0).packets.push_back(
|
||||||
|
MakePacket<std::vector<int>>(input).At(Timestamp(timestamp)));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(TestClipIntVectorSizeCalculatorTest, EmptyVectorInput) {
|
||||||
|
CalculatorGraphConfig::Node node_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||||
|
calculator: "TestClipIntVectorSizeCalculator"
|
||||||
|
input_stream: "input_vector"
|
||||||
|
output_stream: "output_vector"
|
||||||
|
options {
|
||||||
|
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 1 }
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorRunner runner(node_config);
|
||||||
|
|
||||||
|
std::vector<int> input = {};
|
||||||
|
AddInputVector(input, /*timestamp=*/1, &runner);
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||||
|
EXPECT_EQ(1, outputs.size());
|
||||||
|
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||||
|
EXPECT_TRUE(outputs[0].Get<std::vector<int>>().empty());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(TestClipIntVectorSizeCalculatorTest, OneTimestamp) {
|
||||||
|
CalculatorGraphConfig::Node node_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||||
|
calculator: "TestClipIntVectorSizeCalculator"
|
||||||
|
input_stream: "input_vector"
|
||||||
|
output_stream: "output_vector"
|
||||||
|
options {
|
||||||
|
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 2 }
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorRunner runner(node_config);
|
||||||
|
|
||||||
|
std::vector<int> input = {0, 1, 2, 3};
|
||||||
|
AddInputVector(input, /*timestamp=*/1, &runner);
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||||
|
EXPECT_EQ(1, outputs.size());
|
||||||
|
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||||
|
const std::vector<int>& output = outputs[0].Get<std::vector<int>>();
|
||||||
|
EXPECT_EQ(2, output.size());
|
||||||
|
std::vector<int> expected_vector = {0, 1};
|
||||||
|
EXPECT_EQ(expected_vector, output);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(TestClipIntVectorSizeCalculatorTest, TwoInputsAtTwoTimestamps) {
|
||||||
|
CalculatorGraphConfig::Node node_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||||
|
calculator: "TestClipIntVectorSizeCalculator"
|
||||||
|
input_stream: "input_vector"
|
||||||
|
output_stream: "output_vector"
|
||||||
|
options {
|
||||||
|
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 3 }
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorRunner runner(node_config);
|
||||||
|
|
||||||
|
{
|
||||||
|
std::vector<int> input = {0, 1, 2, 3};
|
||||||
|
AddInputVector(input, /*timestamp=*/1, &runner);
|
||||||
|
}
|
||||||
|
{
|
||||||
|
std::vector<int> input = {2, 3, 4, 5};
|
||||||
|
AddInputVector(input, /*timestamp=*/2, &runner);
|
||||||
|
}
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
|
||||||
|
const std::vector<Packet>& outputs = runner.Outputs().Index(0).packets;
|
||||||
|
EXPECT_EQ(2, outputs.size());
|
||||||
|
{
|
||||||
|
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||||
|
const std::vector<int>& output = outputs[0].Get<std::vector<int>>();
|
||||||
|
EXPECT_EQ(3, output.size());
|
||||||
|
std::vector<int> expected_vector = {0, 1, 2};
|
||||||
|
EXPECT_EQ(expected_vector, output);
|
||||||
|
}
|
||||||
|
{
|
||||||
|
EXPECT_EQ(Timestamp(2), outputs[1].Timestamp());
|
||||||
|
const std::vector<int>& output = outputs[1].Get<std::vector<int>>();
|
||||||
|
EXPECT_EQ(3, output.size());
|
||||||
|
std::vector<int> expected_vector = {2, 3, 4};
|
||||||
|
EXPECT_EQ(expected_vector, output);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
typedef ClipVectorSizeCalculator<std::unique_ptr<int>>
|
||||||
|
TestClipUniqueIntPtrVectorSizeCalculator;
|
||||||
|
REGISTER_CALCULATOR(TestClipUniqueIntPtrVectorSizeCalculator);
|
||||||
|
|
||||||
|
TEST(TestClipUniqueIntPtrVectorSizeCalculatorTest, ConsumeOneTimestamp) {
|
||||||
|
/* Note: We don't use CalculatorRunner for this test because it keeps copies
|
||||||
|
* of input packets, so packets sent to the graph don't have sole ownership.
|
||||||
|
* The test needs to send packets that own the data.
|
||||||
|
*/
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||||
|
input_stream: "input_vector"
|
||||||
|
node {
|
||||||
|
calculator: "TestClipUniqueIntPtrVectorSizeCalculator"
|
||||||
|
input_stream: "input_vector"
|
||||||
|
output_stream: "output_vector"
|
||||||
|
options {
|
||||||
|
[mediapipe.ClipVectorSizeCalculatorOptions.ext] { max_vec_size: 3 }
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
std::vector<Packet> outputs;
|
||||||
|
tool::AddVectorSink("output_vector", &graph_config, &outputs);
|
||||||
|
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_EXPECT_OK(graph.Initialize(graph_config));
|
||||||
|
MP_EXPECT_OK(graph.StartRun({}));
|
||||||
|
|
||||||
|
// input1 : {0, 1, 2, 3, 4, 5}
|
||||||
|
auto input_vector = absl::make_unique<std::vector<std::unique_ptr<int>>>(6);
|
||||||
|
for (int i = 0; i < 6; ++i) {
|
||||||
|
input_vector->at(i) = absl::make_unique<int>(i);
|
||||||
|
}
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph.AddPacketToInputStream(
|
||||||
|
"input_vector", Adopt(input_vector.release()).At(Timestamp(1))));
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph.WaitUntilIdle());
|
||||||
|
MP_EXPECT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_EXPECT_OK(graph.WaitUntilDone());
|
||||||
|
|
||||||
|
EXPECT_EQ(1, outputs.size());
|
||||||
|
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||||
|
const std::vector<std::unique_ptr<int>>& result =
|
||||||
|
outputs[0].Get<std::vector<std::unique_ptr<int>>>();
|
||||||
|
EXPECT_EQ(3, result.size());
|
||||||
|
for (int i = 0; i < 3; ++i) {
|
||||||
|
const std::unique_ptr<int>& v = result[i];
|
||||||
|
EXPECT_EQ(i, *v);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
@@ -19,6 +19,10 @@
|
|||||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
#include "mediapipe/framework/formats/landmark.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:
|
||||||
@@ -31,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"
|
||||||
@@ -45,4 +59,11 @@ REGISTER_CALCULATOR(ConcatenateTfLiteTensorVectorCalculator);
|
|||||||
typedef ConcatenateVectorCalculator<::mediapipe::NormalizedLandmark>
|
typedef ConcatenateVectorCalculator<::mediapipe::NormalizedLandmark>
|
||||||
ConcatenateLandmarkVectorCalculator;
|
ConcatenateLandmarkVectorCalculator;
|
||||||
REGISTER_CALCULATOR(ConcatenateLandmarkVectorCalculator);
|
REGISTER_CALCULATOR(ConcatenateLandmarkVectorCalculator);
|
||||||
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
|
typedef ConcatenateVectorCalculator<::tflite::gpu::gl::GlBuffer>
|
||||||
|
ConcatenateGlBufferVectorCalculator;
|
||||||
|
REGISTER_CALCULATOR(ConcatenateGlBufferVectorCalculator);
|
||||||
|
#endif
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -15,6 +15,7 @@
|
|||||||
#ifndef MEDIAPIPE_CALCULATORS_CORE_CONCATENATE_VECTOR_CALCULATOR_H_
|
#ifndef MEDIAPIPE_CALCULATORS_CORE_CONCATENATE_VECTOR_CALCULATOR_H_
|
||||||
#define MEDIAPIPE_CALCULATORS_CORE_CONCATENATE_VECTOR_CALCULATOR_H_
|
#define MEDIAPIPE_CALCULATORS_CORE_CONCATENATE_VECTOR_CALCULATOR_H_
|
||||||
|
|
||||||
|
#include <type_traits>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
#include "mediapipe/calculators/core/concatenate_vector_calculator.pb.h"
|
#include "mediapipe/calculators/core/concatenate_vector_calculator.pb.h"
|
||||||
@@ -59,16 +60,58 @@ class ConcatenateVectorCalculator : public CalculatorBase {
|
|||||||
if (cc->Inputs().Index(i).IsEmpty()) return ::mediapipe::OkStatus();
|
if (cc->Inputs().Index(i).IsEmpty()) return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
auto output = absl::make_unique<std::vector<T>>();
|
|
||||||
|
return ConcatenateVectors<T>(std::is_copy_constructible<T>(), cc);
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename U>
|
||||||
|
::mediapipe::Status ConcatenateVectors(std::true_type,
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
auto output = absl::make_unique<std::vector<U>>();
|
||||||
for (int i = 0; i < cc->Inputs().NumEntries(); ++i) {
|
for (int i = 0; i < cc->Inputs().NumEntries(); ++i) {
|
||||||
if (cc->Inputs().Index(i).IsEmpty()) continue;
|
if (cc->Inputs().Index(i).IsEmpty()) continue;
|
||||||
const std::vector<T>& input = cc->Inputs().Index(i).Get<std::vector<T>>();
|
const std::vector<U>& input = cc->Inputs().Index(i).Get<std::vector<U>>();
|
||||||
output->insert(output->end(), input.begin(), input.end());
|
output->insert(output->end(), input.begin(), input.end());
|
||||||
}
|
}
|
||||||
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
template <typename U>
|
||||||
|
::mediapipe::Status ConcatenateVectors(std::false_type,
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
return ConsumeAndConcatenateVectors<T>(std::is_move_constructible<U>(), cc);
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename U>
|
||||||
|
::mediapipe::Status ConsumeAndConcatenateVectors(std::true_type,
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
auto output = absl::make_unique<std::vector<U>>();
|
||||||
|
for (int i = 0; i < cc->Inputs().NumEntries(); ++i) {
|
||||||
|
if (cc->Inputs().Index(i).IsEmpty()) continue;
|
||||||
|
::mediapipe::StatusOr<std::unique_ptr<std::vector<U>>> input_status =
|
||||||
|
cc->Inputs().Index(i).Value().Consume<std::vector<U>>();
|
||||||
|
if (input_status.ok()) {
|
||||||
|
std::unique_ptr<std::vector<U>> input_vector =
|
||||||
|
std::move(input_status).ValueOrDie();
|
||||||
|
output->insert(output->end(),
|
||||||
|
std::make_move_iterator(input_vector->begin()),
|
||||||
|
std::make_move_iterator(input_vector->end()));
|
||||||
|
} else {
|
||||||
|
return input_status.status();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename U>
|
||||||
|
::mediapipe::Status ConsumeAndConcatenateVectors(std::false_type,
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
return ::mediapipe::InternalError(
|
||||||
|
"Cannot copy or move input vectors to concatenate them");
|
||||||
|
}
|
||||||
|
|
||||||
private:
|
private:
|
||||||
bool only_emit_if_all_present_;
|
bool only_emit_if_all_present_;
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -235,4 +235,167 @@ TEST(ConcatenateFloatVectorCalculatorTest, OneEmptyStreamNoOutput) {
|
|||||||
EXPECT_EQ(0, outputs.size());
|
EXPECT_EQ(0, outputs.size());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
typedef ConcatenateVectorCalculator<std::unique_ptr<int>>
|
||||||
|
TestConcatenateUniqueIntPtrCalculator;
|
||||||
|
REGISTER_CALCULATOR(TestConcatenateUniqueIntPtrCalculator);
|
||||||
|
|
||||||
|
TEST(TestConcatenateUniqueIntVectorCalculatorTest, ConsumeOneTimestamp) {
|
||||||
|
/* Note: We don't use CalculatorRunner for this test because it keeps copies
|
||||||
|
* of input packets, so packets sent to the graph don't have sole ownership.
|
||||||
|
* The test needs to send packets that own the data.
|
||||||
|
*/
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||||
|
input_stream: "in_1"
|
||||||
|
input_stream: "in_2"
|
||||||
|
input_stream: "in_3"
|
||||||
|
node {
|
||||||
|
calculator: "TestConcatenateUniqueIntPtrCalculator"
|
||||||
|
input_stream: "in_1"
|
||||||
|
input_stream: "in_2"
|
||||||
|
input_stream: "in_3"
|
||||||
|
output_stream: "out"
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
std::vector<Packet> outputs;
|
||||||
|
tool::AddVectorSink("out", &graph_config, &outputs);
|
||||||
|
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_EXPECT_OK(graph.Initialize(graph_config));
|
||||||
|
MP_EXPECT_OK(graph.StartRun({}));
|
||||||
|
|
||||||
|
// input1 : {0, 1, 2}
|
||||||
|
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_1 =
|
||||||
|
absl::make_unique<std::vector<std::unique_ptr<int>>>(3);
|
||||||
|
for (int i = 0; i < 3; ++i) {
|
||||||
|
input_1->at(i) = absl::make_unique<int>(i);
|
||||||
|
}
|
||||||
|
// input2: {3}
|
||||||
|
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_2 =
|
||||||
|
absl::make_unique<std::vector<std::unique_ptr<int>>>(1);
|
||||||
|
input_2->at(0) = absl::make_unique<int>(3);
|
||||||
|
// input3: {4, 5}
|
||||||
|
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_3 =
|
||||||
|
absl::make_unique<std::vector<std::unique_ptr<int>>>(2);
|
||||||
|
input_3->at(0) = absl::make_unique<int>(4);
|
||||||
|
input_3->at(1) = absl::make_unique<int>(5);
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph.AddPacketToInputStream(
|
||||||
|
"in_1", Adopt(input_1.release()).At(Timestamp(1))));
|
||||||
|
MP_EXPECT_OK(graph.AddPacketToInputStream(
|
||||||
|
"in_2", Adopt(input_2.release()).At(Timestamp(1))));
|
||||||
|
MP_EXPECT_OK(graph.AddPacketToInputStream(
|
||||||
|
"in_3", Adopt(input_3.release()).At(Timestamp(1))));
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph.WaitUntilIdle());
|
||||||
|
MP_EXPECT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_EXPECT_OK(graph.WaitUntilDone());
|
||||||
|
|
||||||
|
EXPECT_EQ(1, outputs.size());
|
||||||
|
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||||
|
const std::vector<std::unique_ptr<int>>& result =
|
||||||
|
outputs[0].Get<std::vector<std::unique_ptr<int>>>();
|
||||||
|
EXPECT_EQ(6, result.size());
|
||||||
|
for (int i = 0; i < 6; ++i) {
|
||||||
|
const std::unique_ptr<int>& v = result[i];
|
||||||
|
EXPECT_EQ(i, *v);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(TestConcatenateUniqueIntVectorCalculatorTest, OneEmptyStreamStillOutput) {
|
||||||
|
/* Note: We don't use CalculatorRunner for this test because it keeps copies
|
||||||
|
* of input packets, so packets sent to the graph don't have sole ownership.
|
||||||
|
* The test needs to send packets that own the data.
|
||||||
|
*/
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||||
|
input_stream: "in_1"
|
||||||
|
input_stream: "in_2"
|
||||||
|
node {
|
||||||
|
calculator: "TestConcatenateUniqueIntPtrCalculator"
|
||||||
|
input_stream: "in_1"
|
||||||
|
input_stream: "in_2"
|
||||||
|
output_stream: "out"
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
std::vector<Packet> outputs;
|
||||||
|
tool::AddVectorSink("out", &graph_config, &outputs);
|
||||||
|
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_EXPECT_OK(graph.Initialize(graph_config));
|
||||||
|
MP_EXPECT_OK(graph.StartRun({}));
|
||||||
|
|
||||||
|
// input1 : {0, 1, 2}
|
||||||
|
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_1 =
|
||||||
|
absl::make_unique<std::vector<std::unique_ptr<int>>>(3);
|
||||||
|
for (int i = 0; i < 3; ++i) {
|
||||||
|
input_1->at(i) = absl::make_unique<int>(i);
|
||||||
|
}
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph.AddPacketToInputStream(
|
||||||
|
"in_1", Adopt(input_1.release()).At(Timestamp(1))));
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph.WaitUntilIdle());
|
||||||
|
MP_EXPECT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_EXPECT_OK(graph.WaitUntilDone());
|
||||||
|
|
||||||
|
EXPECT_EQ(1, outputs.size());
|
||||||
|
EXPECT_EQ(Timestamp(1), outputs[0].Timestamp());
|
||||||
|
const std::vector<std::unique_ptr<int>>& result =
|
||||||
|
outputs[0].Get<std::vector<std::unique_ptr<int>>>();
|
||||||
|
EXPECT_EQ(3, result.size());
|
||||||
|
for (int i = 0; i < 3; ++i) {
|
||||||
|
const std::unique_ptr<int>& v = result[i];
|
||||||
|
EXPECT_EQ(i, *v);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(TestConcatenateUniqueIntVectorCalculatorTest, OneEmptyStreamNoOutput) {
|
||||||
|
/* Note: We don't use CalculatorRunner for this test because it keeps copies
|
||||||
|
* of input packets, so packets sent to the graph don't have sole ownership.
|
||||||
|
* The test needs to send packets that own the data.
|
||||||
|
*/
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||||
|
input_stream: "in_1"
|
||||||
|
input_stream: "in_2"
|
||||||
|
node {
|
||||||
|
calculator: "TestConcatenateUniqueIntPtrCalculator"
|
||||||
|
input_stream: "in_1"
|
||||||
|
input_stream: "in_2"
|
||||||
|
output_stream: "out"
|
||||||
|
options {
|
||||||
|
[mediapipe.ConcatenateVectorCalculatorOptions.ext] {
|
||||||
|
only_emit_if_all_present: true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
std::vector<Packet> outputs;
|
||||||
|
tool::AddVectorSink("out", &graph_config, &outputs);
|
||||||
|
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_EXPECT_OK(graph.Initialize(graph_config));
|
||||||
|
MP_EXPECT_OK(graph.StartRun({}));
|
||||||
|
|
||||||
|
// input1 : {0, 1, 2}
|
||||||
|
std::unique_ptr<std::vector<std::unique_ptr<int>>> input_1 =
|
||||||
|
absl::make_unique<std::vector<std::unique_ptr<int>>>(3);
|
||||||
|
for (int i = 0; i < 3; ++i) {
|
||||||
|
input_1->at(i) = absl::make_unique<int>(i);
|
||||||
|
}
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph.AddPacketToInputStream(
|
||||||
|
"in_1", Adopt(input_1.release()).At(Timestamp(1))));
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph.WaitUntilIdle());
|
||||||
|
MP_EXPECT_OK(graph.CloseAllPacketSources());
|
||||||
|
MP_EXPECT_OK(graph.WaitUntilDone());
|
||||||
|
|
||||||
|
EXPECT_EQ(0, outputs.size());
|
||||||
|
}
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -0,0 +1,90 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include <cfloat>
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/core/dequantize_byte_array_calculator.pb.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/canonical_errors.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
|
||||||
|
// Dequantizes a byte array to a vector of floats.
|
||||||
|
//
|
||||||
|
// Example config:
|
||||||
|
// node {
|
||||||
|
// calculator: "DequantizeByteArrayCalculator"
|
||||||
|
// input_stream: "ENCODED:encoded"
|
||||||
|
// output_stream: "FLOAT_VECTOR:float_vector"
|
||||||
|
// options {
|
||||||
|
// [mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||||
|
// max_quantized_value: 2
|
||||||
|
// min_quantized_value: -2
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
// }
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
class DequantizeByteArrayCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
cc->Inputs().Tag("ENCODED").Set<std::string>();
|
||||||
|
cc->Outputs().Tag("FLOAT_VECTOR").Set<std::vector<float>>();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) final {
|
||||||
|
const auto options =
|
||||||
|
cc->Options<::mediapipe::DequantizeByteArrayCalculatorOptions>();
|
||||||
|
if (!options.has_max_quantized_value() ||
|
||||||
|
!options.has_min_quantized_value()) {
|
||||||
|
return ::mediapipe::InvalidArgumentError(
|
||||||
|
"Both max_quantized_value and min_quantized_value must be provided "
|
||||||
|
"in DequantizeByteArrayCalculatorOptions.");
|
||||||
|
}
|
||||||
|
float max_quantized_value = options.max_quantized_value();
|
||||||
|
float min_quantized_value = options.min_quantized_value();
|
||||||
|
if (max_quantized_value < min_quantized_value + FLT_EPSILON) {
|
||||||
|
return ::mediapipe::InvalidArgumentError(
|
||||||
|
"max_quantized_value must be greater than min_quantized_value.");
|
||||||
|
}
|
||||||
|
float range = max_quantized_value - min_quantized_value;
|
||||||
|
scalar_ = range / 255.0;
|
||||||
|
bias_ = (range / 512.0) + min_quantized_value;
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||||
|
const std::string& encoded =
|
||||||
|
cc->Inputs().Tag("ENCODED").Value().Get<std::string>();
|
||||||
|
std::vector<float> float_vector;
|
||||||
|
float_vector.reserve(encoded.length());
|
||||||
|
for (int i = 0; i < encoded.length(); ++i) {
|
||||||
|
float_vector.push_back(
|
||||||
|
static_cast<unsigned char>(encoded.at(i)) * scalar_ + bias_);
|
||||||
|
}
|
||||||
|
cc->Outputs()
|
||||||
|
.Tag("FLOAT_VECTOR")
|
||||||
|
.AddPacket(MakePacket<std::vector<float>>(float_vector)
|
||||||
|
.At(cc->InputTimestamp()));
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
float scalar_;
|
||||||
|
float bias_;
|
||||||
|
};
|
||||||
|
|
||||||
|
REGISTER_CALCULATOR(DequantizeByteArrayCalculator);
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
syntax = "proto2";
|
||||||
|
|
||||||
|
package mediapipe;
|
||||||
|
|
||||||
|
import "mediapipe/framework/calculator.proto";
|
||||||
|
|
||||||
|
message DequantizeByteArrayCalculatorOptions {
|
||||||
|
extend CalculatorOptions {
|
||||||
|
optional DequantizeByteArrayCalculatorOptions ext = 272316343;
|
||||||
|
}
|
||||||
|
|
||||||
|
optional float max_quantized_value = 1;
|
||||||
|
optional float min_quantized_value = 2;
|
||||||
|
}
|
||||||
@@ -0,0 +1,137 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include <string>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/calculator_runner.h"
|
||||||
|
#include "mediapipe/framework/port/gmock.h"
|
||||||
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
#include "mediapipe/framework/port/status_matchers.h" // NOLINT
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
TEST(QuantizeFloatVectorCalculatorTest, WrongConfig) {
|
||||||
|
CalculatorGraphConfig::Node node_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||||
|
calculator: "DequantizeByteArrayCalculator"
|
||||||
|
input_stream: "ENCODED:encoded"
|
||||||
|
output_stream: "FLOAT_VECTOR:float_vector"
|
||||||
|
options {
|
||||||
|
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||||
|
max_quantized_value: 2
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorRunner runner(node_config);
|
||||||
|
std::string empty_string;
|
||||||
|
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||||
|
MakePacket<std::string>(empty_string).At(Timestamp(0)));
|
||||||
|
auto status = runner.Run();
|
||||||
|
EXPECT_FALSE(status.ok());
|
||||||
|
EXPECT_THAT(
|
||||||
|
status.message(),
|
||||||
|
testing::HasSubstr(
|
||||||
|
"Both max_quantized_value and min_quantized_value must be provided"));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(QuantizeFloatVectorCalculatorTest, WrongConfig2) {
|
||||||
|
CalculatorGraphConfig::Node node_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||||
|
calculator: "DequantizeByteArrayCalculator"
|
||||||
|
input_stream: "ENCODED:encoded"
|
||||||
|
output_stream: "FLOAT_VECTOR:float_vector"
|
||||||
|
options {
|
||||||
|
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||||
|
max_quantized_value: -2
|
||||||
|
min_quantized_value: 2
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorRunner runner(node_config);
|
||||||
|
std::string empty_string;
|
||||||
|
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||||
|
MakePacket<std::string>(empty_string).At(Timestamp(0)));
|
||||||
|
auto status = runner.Run();
|
||||||
|
EXPECT_FALSE(status.ok());
|
||||||
|
EXPECT_THAT(
|
||||||
|
status.message(),
|
||||||
|
testing::HasSubstr(
|
||||||
|
"max_quantized_value must be greater than min_quantized_value"));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(QuantizeFloatVectorCalculatorTest, WrongConfig3) {
|
||||||
|
CalculatorGraphConfig::Node node_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||||
|
calculator: "DequantizeByteArrayCalculator"
|
||||||
|
input_stream: "ENCODED:encoded"
|
||||||
|
output_stream: "FLOAT_VECTOR:float_vector"
|
||||||
|
options {
|
||||||
|
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||||
|
max_quantized_value: 1
|
||||||
|
min_quantized_value: 1
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorRunner runner(node_config);
|
||||||
|
std::string empty_string;
|
||||||
|
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||||
|
MakePacket<std::string>(empty_string).At(Timestamp(0)));
|
||||||
|
auto status = runner.Run();
|
||||||
|
EXPECT_FALSE(status.ok());
|
||||||
|
EXPECT_THAT(
|
||||||
|
status.message(),
|
||||||
|
testing::HasSubstr(
|
||||||
|
"max_quantized_value must be greater than min_quantized_value"));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(DequantizeByteArrayCalculatorTest, TestDequantization) {
|
||||||
|
CalculatorGraphConfig::Node node_config =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"(
|
||||||
|
calculator: "DequantizeByteArrayCalculator"
|
||||||
|
input_stream: "ENCODED:encoded"
|
||||||
|
output_stream: "FLOAT_VECTOR:float_vector"
|
||||||
|
options {
|
||||||
|
[mediapipe.DequantizeByteArrayCalculatorOptions.ext]: {
|
||||||
|
max_quantized_value: 2
|
||||||
|
min_quantized_value: -2
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
CalculatorRunner runner(node_config);
|
||||||
|
unsigned char input[4] = {0x7F, 0xFF, 0x00, 0x01};
|
||||||
|
runner.MutableInputs()->Tag("ENCODED").packets.push_back(
|
||||||
|
MakePacket<std::string>(
|
||||||
|
std::string(reinterpret_cast<char const*>(input), 4))
|
||||||
|
.At(Timestamp(0)));
|
||||||
|
auto status = runner.Run();
|
||||||
|
MP_ASSERT_OK(runner.Run());
|
||||||
|
const std::vector<Packet>& outputs =
|
||||||
|
runner.Outputs().Tag("FLOAT_VECTOR").packets;
|
||||||
|
EXPECT_EQ(1, outputs.size());
|
||||||
|
const std::vector<float>& result = outputs[0].Get<std::vector<float>>();
|
||||||
|
ASSERT_FALSE(result.empty());
|
||||||
|
EXPECT_EQ(4, result.size());
|
||||||
|
EXPECT_NEAR(0, result[0], 0.01);
|
||||||
|
EXPECT_NEAR(2, result[1], 0.01);
|
||||||
|
EXPECT_NEAR(-2, result[2], 0.01);
|
||||||
|
EXPECT_NEAR(-1.976, result[3], 0.01);
|
||||||
|
|
||||||
|
EXPECT_EQ(Timestamp(0), outputs[0].Timestamp());
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,40 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/core/end_loop_calculator.h"
|
||||||
|
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/formats/landmark.pb.h"
|
||||||
|
#include "mediapipe/framework/formats/rect.pb.h"
|
||||||
|
#include "mediapipe/util/render_data.pb.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
typedef EndLoopCalculator<std::vector<::mediapipe::NormalizedRect>>
|
||||||
|
EndLoopNormalizedRectCalculator;
|
||||||
|
REGISTER_CALCULATOR(EndLoopNormalizedRectCalculator);
|
||||||
|
|
||||||
|
typedef EndLoopCalculator<std::vector<::mediapipe::NormalizedLandmarkList>>
|
||||||
|
EndLoopNormalizedLandmarkListVectorCalculator;
|
||||||
|
REGISTER_CALCULATOR(EndLoopNormalizedLandmarkListVectorCalculator);
|
||||||
|
|
||||||
|
typedef EndLoopCalculator<std::vector<bool>> EndLoopBooleanCalculator;
|
||||||
|
REGISTER_CALCULATOR(EndLoopBooleanCalculator);
|
||||||
|
|
||||||
|
typedef EndLoopCalculator<std::vector<::mediapipe::RenderData>>
|
||||||
|
EndLoopRenderDataCalculator;
|
||||||
|
REGISTER_CALCULATOR(EndLoopRenderDataCalculator);
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,106 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#ifndef MEDIAPIPE_CALCULATORS_CORE_END_LOOP_CALCULATOR_H_
|
||||||
|
#define MEDIAPIPE_CALCULATORS_CORE_END_LOOP_CALCULATOR_H_
|
||||||
|
|
||||||
|
#include "mediapipe/framework/calculator_context.h"
|
||||||
|
#include "mediapipe/framework/calculator_contract.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/collection_item_id.h"
|
||||||
|
#include "mediapipe/framework/port/integral_types.h"
|
||||||
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
// Calculator for completing the processing of loops on iterable collections
|
||||||
|
// inside a MediaPipe graph. The EndLoopCalculator collects all input packets
|
||||||
|
// from ITEM input_stream into a collection and upon receiving the flush signal
|
||||||
|
// from the "BATCH_END" tagged input stream, it emits the aggregated results
|
||||||
|
// at the original timestamp contained in the "BATCH_END" input stream.
|
||||||
|
//
|
||||||
|
// It is designed to be used like:
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "BeginLoopWithIterableCalculator"
|
||||||
|
// input_stream: "ITERABLE:input_iterable" # IterableT @ext_ts
|
||||||
|
// output_stream: "ITEM:input_element" # ItemT @loop_internal_ts
|
||||||
|
// output_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "ElementToBlaConverterSubgraph"
|
||||||
|
// input_stream: "ITEM:input_to_loop_body" # ItemT @loop_internal_ts
|
||||||
|
// output_stream: "BLA:output_of_loop_body" # ItemU @loop_internal_ts
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "EndLoopWithOutputCalculator"
|
||||||
|
// input_stream: "ITEM:output_of_loop_body" # ItemU @loop_internal_ts
|
||||||
|
// input_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
|
||||||
|
// output_stream: "OUTPUT:aggregated_result" # IterableU @ext_ts
|
||||||
|
// }
|
||||||
|
template <typename IterableT>
|
||||||
|
class EndLoopCalculator : public CalculatorBase {
|
||||||
|
using ItemT = typename IterableT::value_type;
|
||||||
|
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
RET_CHECK(cc->Inputs().HasTag("BATCH_END"))
|
||||||
|
<< "Missing BATCH_END tagged input_stream.";
|
||||||
|
cc->Inputs().Tag("BATCH_END").Set<Timestamp>();
|
||||||
|
|
||||||
|
RET_CHECK(cc->Inputs().HasTag("ITEM"));
|
||||||
|
cc->Inputs().Tag("ITEM").Set<ItemT>();
|
||||||
|
|
||||||
|
RET_CHECK(cc->Outputs().HasTag("ITERABLE"));
|
||||||
|
cc->Outputs().Tag("ITERABLE").Set<IterableT>();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
|
if (!cc->Inputs().Tag("ITEM").IsEmpty()) {
|
||||||
|
if (!input_stream_collection_) {
|
||||||
|
input_stream_collection_.reset(new IterableT);
|
||||||
|
}
|
||||||
|
input_stream_collection_->push_back(
|
||||||
|
cc->Inputs().Tag("ITEM").template Get<ItemT>());
|
||||||
|
}
|
||||||
|
|
||||||
|
if (!cc->Inputs().Tag("BATCH_END").Value().IsEmpty()) { // flush signal
|
||||||
|
Timestamp loop_control_ts =
|
||||||
|
cc->Inputs().Tag("BATCH_END").template Get<Timestamp>();
|
||||||
|
if (input_stream_collection_) {
|
||||||
|
cc->Outputs()
|
||||||
|
.Tag("ITERABLE")
|
||||||
|
.Add(input_stream_collection_.release(), loop_control_ts);
|
||||||
|
} else {
|
||||||
|
// Since there is no collection, inform downstream calculators to not
|
||||||
|
// expect any packet by updating the timestamp bounds.
|
||||||
|
cc->Outputs()
|
||||||
|
.Tag("ITERABLE")
|
||||||
|
.SetNextTimestampBound(Timestamp(loop_control_ts.Value() + 1));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
std::unique_ptr<IterableT> input_stream_collection_;
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
|
|
||||||
|
#endif // MEDIAPIPE_CALCULATORS_CORE_END_LOOP_CALCULATOR_H_
|
||||||
@@ -12,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,114 +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:
|
|
||||||
// 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.
|
||||||
@@ -203,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()) {
|
||||||
@@ -233,6 +119,9 @@ 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_;
|
||||||
|
|
||||||
if (resampler_options.output_header() !=
|
if (resampler_options.output_header() !=
|
||||||
@@ -272,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();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -287,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));
|
||||||
@@ -295,6 +195,40 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
|||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
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_ =
|
||||||
|
first_timestamp_ + frame_time_usec_ * random_->RandFloat();
|
||||||
|
}
|
||||||
|
|
||||||
|
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_ +=
|
||||||
|
frame_time_usec_ *
|
||||||
|
((1.0 - jitter_) + 2.0 * jitter_ * random_->RandFloat());
|
||||||
|
}
|
||||||
|
|
||||||
::mediapipe::Status PacketResamplerCalculator::ProcessWithJitter(
|
::mediapipe::Status PacketResamplerCalculator::ProcessWithJitter(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
RET_CHECK_GT(cc->InputTimestamp(), Timestamp::PreStream());
|
RET_CHECK_GT(cc->InputTimestamp(), Timestamp::PreStream());
|
||||||
@@ -302,29 +236,37 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
|||||||
|
|
||||||
if (first_timestamp_ == Timestamp::Unset()) {
|
if (first_timestamp_ == Timestamp::Unset()) {
|
||||||
first_timestamp_ = cc->InputTimestamp();
|
first_timestamp_ = cc->InputTimestamp();
|
||||||
next_output_timestamp_ =
|
InitializeNextOutputTimestampWithJitter();
|
||||||
first_timestamp_ + frame_time_usec_ * random_->RandFloat();
|
if (first_timestamp_ == next_output_timestamp_) {
|
||||||
|
OutputWithinLimits(
|
||||||
|
cc,
|
||||||
|
cc->Inputs().Get(input_data_id_).Value().At(next_output_timestamp_));
|
||||||
|
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_));
|
|
||||||
next_output_timestamp_ +=
|
|
||||||
frame_time_usec_ *
|
|
||||||
((1.0 - jitter_) + 2.0 * jitter_ * random_->RandFloat());
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -405,6 +347,9 @@ TimestampDiff TimestampDiffFromSeconds(double seconds) {
|
|||||||
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 {
|
||||||
|
|
||||||
@@ -86,6 +87,7 @@ class PreviousLoopbackCalculator : public CalculatorBase {
|
|||||||
main_ts_.pop_front();
|
main_ts_.pop_front();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
auto& loop_out = cc->Outputs().Get(loop_out_id_);
|
||||||
|
|
||||||
while (!main_ts_.empty() && !loopback_packets_.empty()) {
|
while (!main_ts_.empty() && !loopback_packets_.empty()) {
|
||||||
Timestamp main_timestamp = main_ts_.front();
|
Timestamp main_timestamp = main_ts_.front();
|
||||||
@@ -95,13 +97,32 @@ class PreviousLoopbackCalculator : public CalculatorBase {
|
|||||||
|
|
||||||
if (previous_loopback.IsEmpty()) {
|
if (previous_loopback.IsEmpty()) {
|
||||||
// TODO: SetCompleteTimestampBound would be more useful.
|
// TODO: SetCompleteTimestampBound would be more useful.
|
||||||
cc->Outputs()
|
loop_out.SetNextTimestampBound(main_timestamp + 1);
|
||||||
.Get(loop_out_id_)
|
|
||||||
.SetNextTimestampBound(main_timestamp + 1);
|
|
||||||
} else {
|
} else {
|
||||||
cc->Outputs().Get(loop_out_id_).AddPacket(std::move(previous_loopback));
|
loop_out.AddPacket(std::move(previous_loopback));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// In case of an empty loopback input, the next timestamp bound for
|
||||||
|
// loopback input is the loopback timestamp + 1. The next timestamp bound
|
||||||
|
// for output is set and the main_ts_ vector is truncated accordingly.
|
||||||
|
if (loopback_packet.IsEmpty() &&
|
||||||
|
loopback_packet.Timestamp() != Timestamp::Unstarted()) {
|
||||||
|
Timestamp loopback_bound =
|
||||||
|
loopback_packet.Timestamp().NextAllowedInStream();
|
||||||
|
while (!main_ts_.empty() && main_ts_.front() <= loopback_bound) {
|
||||||
|
main_ts_.pop_front();
|
||||||
|
}
|
||||||
|
if (main_ts_.empty()) {
|
||||||
|
loop_out.SetNextTimestampBound(loopback_bound.NextAllowedInStream());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (!main_ts_.empty()) {
|
||||||
|
loop_out.SetNextTimestampBound(main_ts_.front());
|
||||||
|
}
|
||||||
|
if (cc->Inputs().Get(main_id_).IsDone() && main_ts_.empty()) {
|
||||||
|
loop_out.Close();
|
||||||
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -93,19 +93,178 @@ TEST(PreviousLoopbackCalculator, CorrectTimestamps) {
|
|||||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1}));
|
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1}));
|
||||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(1, -1));
|
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(1, -1));
|
||||||
|
|
||||||
|
send_packet("in", 2);
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
|
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 2}));
|
||||||
|
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(2, 1));
|
||||||
|
|
||||||
send_packet("in", 5);
|
send_packet("in", 5);
|
||||||
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 5}));
|
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 2, 5}));
|
||||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(5, 1));
|
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(5, 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, 5, 15}));
|
EXPECT_EQ(TimestampValues(in_prev), (std::vector<int64>{1, 2, 5, 15}));
|
||||||
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(15, 5));
|
EXPECT_EQ(pair_values(in_prev.back()), std::make_pair(15, 5));
|
||||||
|
|
||||||
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||||
MP_EXPECT_OK(graph_.WaitUntilDone());
|
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// A Calculator that outputs a summary packet in CalculatorBase::Close().
|
||||||
|
class PacketOnCloseCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
cc->Inputs().Index(0).Set<int>();
|
||||||
|
cc->Outputs().Index(0).Set<int>();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) final {
|
||||||
|
cc->SetOffset(TimestampDiff(0));
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||||
|
sum_ += cc->Inputs().Index(0).Value().Get<int>();
|
||||||
|
cc->Outputs().Index(0).AddPacket(cc->Inputs().Index(0).Value());
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Close(CalculatorContext* cc) final {
|
||||||
|
cc->Outputs().Index(0).AddPacket(
|
||||||
|
MakePacket<int>(sum_).At(Timestamp::Max()));
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
int sum_ = 0;
|
||||||
|
};
|
||||||
|
REGISTER_CALCULATOR(PacketOnCloseCalculator);
|
||||||
|
|
||||||
|
// Demonstrates that all ouput and input streams in PreviousLoopbackCalculator
|
||||||
|
// will close as expected when all graph input streams are closed.
|
||||||
|
TEST(PreviousLoopbackCalculator, ClosesCorrectly) {
|
||||||
|
std::vector<Packet> outputs;
|
||||||
|
CalculatorGraphConfig graph_config_ =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||||
|
input_stream: 'in'
|
||||||
|
node {
|
||||||
|
calculator: 'PreviousLoopbackCalculator'
|
||||||
|
input_stream: 'MAIN:in'
|
||||||
|
input_stream: 'LOOP:out'
|
||||||
|
input_stream_info: { tag_index: 'LOOP' back_edge: true }
|
||||||
|
output_stream: 'PREV_LOOP:previous'
|
||||||
|
}
|
||||||
|
# This calculator synchronizes its inputs as normal, so it is used
|
||||||
|
# to check that both "in" and "previous" are ready.
|
||||||
|
node {
|
||||||
|
calculator: 'PassThroughCalculator'
|
||||||
|
input_stream: 'in'
|
||||||
|
input_stream: 'previous'
|
||||||
|
output_stream: 'out'
|
||||||
|
output_stream: 'previous2'
|
||||||
|
}
|
||||||
|
node {
|
||||||
|
calculator: 'PacketOnCloseCalculator'
|
||||||
|
input_stream: 'out'
|
||||||
|
output_stream: 'close_out'
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
tool::AddVectorSink("close_out", &graph_config_, &outputs);
|
||||||
|
|
||||||
|
CalculatorGraph graph_;
|
||||||
|
MP_ASSERT_OK(graph_.Initialize(graph_config_, {}));
|
||||||
|
MP_ASSERT_OK(graph_.StartRun({}));
|
||||||
|
|
||||||
|
auto send_packet = [&graph_](const std::string& input_name, int n) {
|
||||||
|
MP_EXPECT_OK(graph_.AddPacketToInputStream(
|
||||||
|
input_name, MakePacket<int>(n).At(Timestamp(n))));
|
||||||
|
};
|
||||||
|
|
||||||
|
send_packet("in", 1);
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
|
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1}));
|
||||||
|
|
||||||
|
send_packet("in", 2);
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
|
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1, 2}));
|
||||||
|
|
||||||
|
send_packet("in", 5);
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
|
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1, 2, 5}));
|
||||||
|
|
||||||
|
send_packet("in", 15);
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
|
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{1, 2, 5, 15}));
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
|
EXPECT_EQ(TimestampValues(outputs),
|
||||||
|
(std::vector<int64>{1, 2, 5, 15, Timestamp::Max().Value()}));
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||||
|
}
|
||||||
|
|
||||||
|
// Demonstrates that downstream calculators won't be blocked by
|
||||||
|
// always-empty-LOOP-stream.
|
||||||
|
TEST(PreviousLoopbackCalculator, EmptyLoopForever) {
|
||||||
|
std::vector<Packet> outputs;
|
||||||
|
CalculatorGraphConfig graph_config_ =
|
||||||
|
ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||||
|
input_stream: 'in'
|
||||||
|
node {
|
||||||
|
calculator: 'PreviousLoopbackCalculator'
|
||||||
|
input_stream: 'MAIN:in'
|
||||||
|
input_stream: 'LOOP:previous'
|
||||||
|
input_stream_info: { tag_index: 'LOOP' back_edge: true }
|
||||||
|
output_stream: 'PREV_LOOP:previous'
|
||||||
|
}
|
||||||
|
# This calculator synchronizes its inputs as normal, so it is used
|
||||||
|
# to check that both "in" and "previous" are ready.
|
||||||
|
node {
|
||||||
|
calculator: 'PassThroughCalculator'
|
||||||
|
input_stream: 'in'
|
||||||
|
input_stream: 'previous'
|
||||||
|
output_stream: 'out'
|
||||||
|
output_stream: 'previous2'
|
||||||
|
}
|
||||||
|
node {
|
||||||
|
calculator: 'PacketOnCloseCalculator'
|
||||||
|
input_stream: 'out'
|
||||||
|
output_stream: 'close_out'
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
tool::AddVectorSink("close_out", &graph_config_, &outputs);
|
||||||
|
|
||||||
|
CalculatorGraph graph_;
|
||||||
|
MP_ASSERT_OK(graph_.Initialize(graph_config_, {}));
|
||||||
|
MP_ASSERT_OK(graph_.StartRun({}));
|
||||||
|
|
||||||
|
auto send_packet = [&graph_](const std::string& input_name, int n) {
|
||||||
|
MP_EXPECT_OK(graph_.AddPacketToInputStream(
|
||||||
|
input_name, MakePacket<int>(n).At(Timestamp(n))));
|
||||||
|
};
|
||||||
|
|
||||||
|
send_packet("in", 0);
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
|
EXPECT_EQ(TimestampValues(outputs), (std::vector<int64>{0}));
|
||||||
|
|
||||||
|
for (int main_ts = 1; main_ts < 50; ++main_ts) {
|
||||||
|
send_packet("in", main_ts);
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
|
std::vector<int64> ts_values = TimestampValues(outputs);
|
||||||
|
EXPECT_EQ(ts_values.size(), main_ts + 1);
|
||||||
|
for (int j = 0; j < main_ts; ++j) {
|
||||||
|
EXPECT_EQ(ts_values[j], j);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
MP_EXPECT_OK(graph_.CloseAllInputStreams());
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilIdle());
|
||||||
|
MP_EXPECT_OK(graph_.WaitUntilDone());
|
||||||
|
}
|
||||||
|
|
||||||
} // anonymous namespace
|
} // anonymous namespace
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -0,0 +1,161 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include <map>
|
||||||
|
#include <memory>
|
||||||
|
#include <set>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/logging.h"
|
||||||
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
using mediapipe::PacketTypeSet;
|
||||||
|
using mediapipe::Timestamp;
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
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 = []() {
|
||||||
|
auto* res = new std::map<std::string, Timestamp>();
|
||||||
|
res->emplace(kTagAtPreStream, Timestamp::PreStream());
|
||||||
|
res->emplace(kTagAtPostStream, Timestamp::PostStream());
|
||||||
|
res->emplace(kTagAtZero, Timestamp(0));
|
||||||
|
res->emplace(kTagAtTick, Timestamp::Unset());
|
||||||
|
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
|
||||||
|
|
||||||
|
// Outputs side packet(s) in corresponding output stream(s) with a particular
|
||||||
|
// 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 {
|
||||||
|
public:
|
||||||
|
SidePacketToStreamCalculator() = default;
|
||||||
|
~SidePacketToStreamCalculator() 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:
|
||||||
|
bool is_tick_processing_ = false;
|
||||||
|
std::string output_tag_;
|
||||||
|
};
|
||||||
|
REGISTER_CALCULATOR(SidePacketToStreamCalculator);
|
||||||
|
|
||||||
|
::mediapipe::Status SidePacketToStreamCalculator::GetContract(
|
||||||
|
CalculatorContract* cc) {
|
||||||
|
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());
|
||||||
|
}
|
||||||
|
|
||||||
|
if (cc->Inputs().HasTag(kTagTick)) {
|
||||||
|
cc->Inputs().Tag(kTagTick).SetAny();
|
||||||
|
}
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::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();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status SidePacketToStreamCalculator::Process(
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
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) {
|
||||||
|
if (!cc->Outputs().HasTag(kTagAtTick)) {
|
||||||
|
const auto& timestamp = kTimestampMap->at(output_tag_);
|
||||||
|
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();
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -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
|
||||||
@@ -17,8 +17,13 @@
|
|||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
#include "mediapipe/framework/formats/landmark.pb.h"
|
#include "mediapipe/framework/formats/landmark.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:
|
||||||
@@ -35,10 +40,21 @@ 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, false>
|
||||||
|
SplitNormalizedRectVectorCalculator;
|
||||||
|
REGISTER_CALCULATOR(SplitNormalizedRectVectorCalculator);
|
||||||
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
|
typedef SplitVectorCalculator<::tflite::gpu::gl::GlBuffer, true>
|
||||||
|
MovableSplitGlBufferVectorCalculator;
|
||||||
|
REGISTER_CALCULATOR(MovableSplitGlBufferVectorCalculator);
|
||||||
|
#endif
|
||||||
} // namespace mediapipe
|
} // 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,15 +31,31 @@
|
|||||||
|
|
||||||
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,
|
||||||
// all ranges should be of size 1 and all outputs will be elements of type T. If
|
// all ranges should be of size 1 and all outputs will be elements of type T. If
|
||||||
// "element_only" is false, ranges can be non-zero in size and all outputs will
|
// "element_only" is false, ranges can be non-zero in size and all outputs will
|
||||||
// be of type std::vector<T>.
|
// be of type std::vector<T>. If the option "combine_outputs" is set to true,
|
||||||
|
// only one output stream can be specified and all ranges of elements will be
|
||||||
|
// combined into one vector.
|
||||||
// To use this class for a particular type T, register a calculator using
|
// 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) {
|
||||||
@@ -49,28 +67,40 @@ class SplitVectorCalculator : public CalculatorBase {
|
|||||||
const auto& options =
|
const auto& options =
|
||||||
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
|
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
|
||||||
|
|
||||||
if (cc->Outputs().NumEntries() != options.ranges_size()) {
|
if (!std::is_copy_constructible<T>::value || move_elements) {
|
||||||
return ::mediapipe::InvalidArgumentError(
|
// Ranges of elements shouldn't overlap when the vector contains
|
||||||
"The number of output streams should match the number of ranges "
|
// non-copyable elements.
|
||||||
"specified in the CalculatorOptions.");
|
RET_CHECK_OK(checkRangesDontOverlap(options));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Set the output types for each output stream.
|
if (options.combine_outputs()) {
|
||||||
for (int i = 0; i < cc->Outputs().NumEntries(); ++i) {
|
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1);
|
||||||
if (options.ranges(i).begin() < 0 || options.ranges(i).end() < 0 ||
|
cc->Outputs().Index(0).Set<std::vector<T>>();
|
||||||
options.ranges(i).begin() >= options.ranges(i).end()) {
|
RET_CHECK_OK(checkRangesDontOverlap(options));
|
||||||
|
} else {
|
||||||
|
if (cc->Outputs().NumEntries() != options.ranges_size()) {
|
||||||
return ::mediapipe::InvalidArgumentError(
|
return ::mediapipe::InvalidArgumentError(
|
||||||
"Indices should be non-negative and begin index should be less "
|
"The number of output streams should match the number of ranges "
|
||||||
"than the end index.");
|
"specified in the CalculatorOptions.");
|
||||||
}
|
}
|
||||||
if (options.element_only()) {
|
|
||||||
if (options.ranges(i).end() - options.ranges(i).begin() != 1) {
|
// Set the output types for each output stream.
|
||||||
|
for (int i = 0; i < cc->Outputs().NumEntries(); ++i) {
|
||||||
|
if (options.ranges(i).begin() < 0 || options.ranges(i).end() < 0 ||
|
||||||
|
options.ranges(i).begin() >= options.ranges(i).end()) {
|
||||||
return ::mediapipe::InvalidArgumentError(
|
return ::mediapipe::InvalidArgumentError(
|
||||||
"Since element_only is true, all ranges should be of size 1.");
|
"Indices should be non-negative and begin index should be less "
|
||||||
|
"than the end index.");
|
||||||
|
}
|
||||||
|
if (options.element_only()) {
|
||||||
|
if (options.ranges(i).end() - options.ranges(i).begin() != 1) {
|
||||||
|
return ::mediapipe::InvalidArgumentError(
|
||||||
|
"Since element_only is true, all ranges should be of size 1.");
|
||||||
|
}
|
||||||
|
cc->Outputs().Index(i).Set<T>();
|
||||||
|
} else {
|
||||||
|
cc->Outputs().Index(i).Set<std::vector<T>>();
|
||||||
}
|
}
|
||||||
cc->Outputs().Index(i).Set<T>();
|
|
||||||
} else {
|
|
||||||
cc->Outputs().Index(i).Set<std::vector<T>>();
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -83,41 +113,140 @@ class SplitVectorCalculator : public CalculatorBase {
|
|||||||
const auto& options =
|
const auto& options =
|
||||||
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
|
cc->Options<::mediapipe::SplitVectorCalculatorOptions>();
|
||||||
|
|
||||||
|
element_only_ = options.element_only();
|
||||||
|
combine_outputs_ = options.combine_outputs();
|
||||||
|
|
||||||
for (const auto& range : options.ranges()) {
|
for (const auto& range : options.ranges()) {
|
||||||
ranges_.push_back({range.begin(), range.end()});
|
ranges_.push_back({range.begin(), range.end()});
|
||||||
max_range_end_ = std::max(max_range_end_, range.end());
|
max_range_end_ = std::max(max_range_end_, range.end());
|
||||||
|
total_elements_ += range.end() - range.begin();
|
||||||
}
|
}
|
||||||
|
|
||||||
element_only_ = options.element_only();
|
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::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 (element_only_) {
|
if (move_elements) {
|
||||||
for (int i = 0; i < ranges_.size(); ++i) {
|
return ProcessMovableElements<T>(cc);
|
||||||
cc->Outputs().Index(i).AddPacket(
|
|
||||||
MakePacket<T>(input[ranges_[i].first]).At(cc->InputTimestamp()));
|
|
||||||
}
|
|
||||||
} else {
|
} else {
|
||||||
|
return ProcessCopyableElements<T>(cc);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename U, IsCopyable<U> = true>
|
||||||
|
::mediapipe::Status ProcessCopyableElements(CalculatorContext* cc) {
|
||||||
|
// static_assert(std::is_copy_constructible<U>::value,
|
||||||
|
// "Cannot copy non-copyable elements");
|
||||||
|
const auto& input = cc->Inputs().Index(0).Get<std::vector<U>>();
|
||||||
|
RET_CHECK_GE(input.size(), max_range_end_);
|
||||||
|
if (combine_outputs_) {
|
||||||
|
auto output = absl::make_unique<std::vector<U>>();
|
||||||
|
output->reserve(total_elements_);
|
||||||
for (int i = 0; i < ranges_.size(); ++i) {
|
for (int i = 0; i < ranges_.size(); ++i) {
|
||||||
auto output = absl::make_unique<std::vector<T>>(
|
auto elements = absl::make_unique<std::vector<U>>(
|
||||||
input.begin() + ranges_[i].first,
|
input.begin() + ranges_[i].first,
|
||||||
input.begin() + ranges_[i].second);
|
input.begin() + ranges_[i].second);
|
||||||
cc->Outputs().Index(i).Add(output.release(), cc->InputTimestamp());
|
output->insert(output->end(), elements->begin(), elements->end());
|
||||||
|
}
|
||||||
|
cc->Outputs().Index(0).Add(output.release(), cc->InputTimestamp());
|
||||||
|
} else {
|
||||||
|
if (element_only_) {
|
||||||
|
for (int i = 0; i < ranges_.size(); ++i) {
|
||||||
|
cc->Outputs().Index(i).AddPacket(
|
||||||
|
MakePacket<U>(input[ranges_[i].first]).At(cc->InputTimestamp()));
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
for (int i = 0; i < ranges_.size(); ++i) {
|
||||||
|
auto output = absl::make_unique<std::vector<T>>(
|
||||||
|
input.begin() + ranges_[i].first,
|
||||||
|
input.begin() + ranges_[i].second);
|
||||||
|
cc->Outputs().Index(i).Add(output.release(), cc->InputTimestamp());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
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;
|
||||||
bool element_only_ = false;
|
bool element_only_ = false;
|
||||||
|
bool combine_outputs_ = false;
|
||||||
};
|
};
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -37,4 +37,7 @@ message SplitVectorCalculatorOptions {
|
|||||||
// just element of type T. By default, if a range specifies only one element,
|
// just element of type T. By default, if a range specifies only one element,
|
||||||
// it is outputted as an std::vector<T>.
|
// it is outputted as an std::vector<T>.
|
||||||
optional bool element_only = 2 [default = false];
|
optional bool element_only = 2 [default = false];
|
||||||
|
|
||||||
|
// Combines output elements to one vector.
|
||||||
|
optional bool combine_outputs = 3 [default = false];
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -105,6 +105,34 @@ class SplitTfLiteTensorVectorCalculatorTest : public ::testing::Test {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
void ValidateCombinedVectorOutput(std::vector<Packet>& output_packets,
|
||||||
|
int expected_elements,
|
||||||
|
std::vector<int>& input_begin_indices,
|
||||||
|
std::vector<int>& input_end_indices) {
|
||||||
|
ASSERT_EQ(1, output_packets.size());
|
||||||
|
ASSERT_EQ(input_begin_indices.size(), input_end_indices.size());
|
||||||
|
const std::vector<TfLiteTensor>& output_vec =
|
||||||
|
output_packets[0].Get<std::vector<TfLiteTensor>>();
|
||||||
|
ASSERT_EQ(expected_elements, output_vec.size());
|
||||||
|
const int num_ranges = input_begin_indices.size();
|
||||||
|
|
||||||
|
int element_id = 0;
|
||||||
|
for (int range_id = 0; range_id < num_ranges; ++range_id) {
|
||||||
|
for (int i = input_begin_indices[range_id];
|
||||||
|
i < input_end_indices[range_id]; ++i) {
|
||||||
|
const int expected_value = i;
|
||||||
|
const TfLiteTensor* result = &output_vec[element_id];
|
||||||
|
float* result_buffer = result->data.f;
|
||||||
|
ASSERT_NE(result_buffer, nullptr);
|
||||||
|
ASSERT_EQ(result_buffer, input_buffers_[i]);
|
||||||
|
for (int j = 0; j < width * height * channels; ++j) {
|
||||||
|
ASSERT_EQ(expected_value, result_buffer[j]);
|
||||||
|
}
|
||||||
|
element_id++;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
void ValidateElementOutput(std::vector<Packet>& output_packets,
|
void ValidateElementOutput(std::vector<Packet>& output_packets,
|
||||||
int input_begin_index) {
|
int input_begin_index) {
|
||||||
ASSERT_EQ(1, output_packets.size());
|
ASSERT_EQ(1, output_packets.size());
|
||||||
@@ -234,6 +262,65 @@ TEST_F(SplitTfLiteTensorVectorCalculatorTest, InvalidOutputStreamCountTest) {
|
|||||||
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
TEST_F(SplitTfLiteTensorVectorCalculatorTest,
|
||||||
|
InvalidCombineOutputsMultipleOutputsTest) {
|
||||||
|
ASSERT_NE(interpreter_, nullptr);
|
||||||
|
|
||||||
|
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "tensor_in"
|
||||||
|
node {
|
||||||
|
calculator: "SplitTfLiteTensorVectorCalculator"
|
||||||
|
input_stream: "tensor_in"
|
||||||
|
output_stream: "range_0"
|
||||||
|
output_stream: "range_1"
|
||||||
|
options {
|
||||||
|
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||||
|
ranges: { begin: 0 end: 1 }
|
||||||
|
ranges: { begin: 2 end: 3 }
|
||||||
|
combine_outputs: true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
// Run the graph.
|
||||||
|
CalculatorGraph graph;
|
||||||
|
// The graph should fail running because the number of output streams does not
|
||||||
|
// match the number of range elements in the options.
|
||||||
|
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(SplitTfLiteTensorVectorCalculatorTest, InvalidOverlappingRangesTest) {
|
||||||
|
ASSERT_NE(interpreter_, nullptr);
|
||||||
|
|
||||||
|
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "tensor_in"
|
||||||
|
node {
|
||||||
|
calculator: "SplitTfLiteTensorVectorCalculator"
|
||||||
|
input_stream: "tensor_in"
|
||||||
|
output_stream: "range_0"
|
||||||
|
options {
|
||||||
|
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||||
|
ranges: { begin: 0 end: 3 }
|
||||||
|
ranges: { begin: 1 end: 4 }
|
||||||
|
combine_outputs: true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
|
||||||
|
// Run the graph.
|
||||||
|
CalculatorGraph graph;
|
||||||
|
// The graph should fail running because there are overlapping ranges.
|
||||||
|
ASSERT_FALSE(graph.Initialize(graph_config).ok());
|
||||||
|
}
|
||||||
|
|
||||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest, SmokeTestElementOnly) {
|
TEST_F(SplitTfLiteTensorVectorCalculatorTest, SmokeTestElementOnly) {
|
||||||
ASSERT_NE(interpreter_, nullptr);
|
ASSERT_NE(interpreter_, nullptr);
|
||||||
|
|
||||||
@@ -289,6 +376,53 @@ TEST_F(SplitTfLiteTensorVectorCalculatorTest, SmokeTestElementOnly) {
|
|||||||
MP_ASSERT_OK(graph.WaitUntilDone());
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
TEST_F(SplitTfLiteTensorVectorCalculatorTest, SmokeTestCombiningOutputs) {
|
||||||
|
ASSERT_NE(interpreter_, nullptr);
|
||||||
|
|
||||||
|
PrepareTfLiteTensorVector(/*vector_size=*/5);
|
||||||
|
ASSERT_NE(input_vec_, nullptr);
|
||||||
|
|
||||||
|
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||||
|
CalculatorGraphConfig graph_config =
|
||||||
|
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||||
|
R"(
|
||||||
|
input_stream: "tensor_in"
|
||||||
|
node {
|
||||||
|
calculator: "SplitTfLiteTensorVectorCalculator"
|
||||||
|
input_stream: "tensor_in"
|
||||||
|
output_stream: "range_0"
|
||||||
|
options {
|
||||||
|
[mediapipe.SplitVectorCalculatorOptions.ext] {
|
||||||
|
ranges: { begin: 0 end: 1 }
|
||||||
|
ranges: { begin: 2 end: 3 }
|
||||||
|
ranges: { begin: 4 end: 5 }
|
||||||
|
combine_outputs: true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)");
|
||||||
|
std::vector<Packet> range_0_packets;
|
||||||
|
tool::AddVectorSink("range_0", &graph_config, &range_0_packets);
|
||||||
|
|
||||||
|
// Run the graph.
|
||||||
|
CalculatorGraph graph;
|
||||||
|
MP_ASSERT_OK(graph.Initialize(graph_config));
|
||||||
|
MP_ASSERT_OK(graph.StartRun({}));
|
||||||
|
MP_ASSERT_OK(graph.AddPacketToInputStream(
|
||||||
|
"tensor_in", Adopt(input_vec_.release()).At(Timestamp(0))));
|
||||||
|
// Wait until the calculator finishes processing.
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilIdle());
|
||||||
|
|
||||||
|
std::vector<int> input_begin_indices = {0, 2, 4};
|
||||||
|
std::vector<int> input_end_indices = {1, 3, 5};
|
||||||
|
ValidateCombinedVectorOutput(range_0_packets, /*expected_elements=*/3,
|
||||||
|
input_begin_indices, input_end_indices);
|
||||||
|
|
||||||
|
// Fully close the graph at the end.
|
||||||
|
MP_ASSERT_OK(graph.CloseInputStream("tensor_in"));
|
||||||
|
MP_ASSERT_OK(graph.WaitUntilDone());
|
||||||
|
}
|
||||||
|
|
||||||
TEST_F(SplitTfLiteTensorVectorCalculatorTest,
|
TEST_F(SplitTfLiteTensorVectorCalculatorTest,
|
||||||
ElementOnlyDisablesVectorOutputs) {
|
ElementOnlyDisablesVectorOutputs) {
|
||||||
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
// Prepare a graph to use the SplitTfLiteTensorVectorCalculator.
|
||||||
@@ -318,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
|
||||||
|
|||||||
@@ -0,0 +1,48 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
#include "mediapipe/framework/timestamp.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
// A calculator that takes a packet of an input stream and converts it to an
|
||||||
|
// output side packet. This calculator only works under the assumption that the
|
||||||
|
// input stream only has a single packet passing through.
|
||||||
|
//
|
||||||
|
// Example config:
|
||||||
|
// node {
|
||||||
|
// calculator: "StreamToSidePacketCalculator"
|
||||||
|
// input_stream: "stream"
|
||||||
|
// output_side_packet: "side_packet"
|
||||||
|
// }
|
||||||
|
class StreamToSidePacketCalculator : public mediapipe::CalculatorBase {
|
||||||
|
public:
|
||||||
|
static mediapipe::Status GetContract(mediapipe::CalculatorContract* cc) {
|
||||||
|
cc->Inputs().Index(0).SetAny();
|
||||||
|
cc->OutputSidePackets().Index(0).SetAny();
|
||||||
|
return mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
mediapipe::Status Process(mediapipe::CalculatorContext* cc) override {
|
||||||
|
mediapipe::Packet& packet = cc->Inputs().Index(0).Value();
|
||||||
|
cc->OutputSidePackets().Index(0).Set(
|
||||||
|
packet.At(mediapipe::Timestamp::Unset()));
|
||||||
|
return mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
};
|
||||||
|
REGISTER_CALCULATOR(StreamToSidePacketCalculator);
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,67 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
#include "absl/memory/memory.h"
|
||||||
|
#include "mediapipe/framework/calculator_runner.h"
|
||||||
|
#include "mediapipe/framework/packet.h"
|
||||||
|
#include "mediapipe/framework/port/gmock.h"
|
||||||
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
#include "mediapipe/framework/port/status_matchers.h"
|
||||||
|
#include "mediapipe/framework/timestamp.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
using ::testing::Test;
|
||||||
|
|
||||||
|
class StreamToSidePacketCalculatorTest : public Test {
|
||||||
|
protected:
|
||||||
|
StreamToSidePacketCalculatorTest() {
|
||||||
|
const char kConfig[] = R"(
|
||||||
|
calculator: "StreamToSidePacketCalculator"
|
||||||
|
input_stream: "stream"
|
||||||
|
output_side_packet: "side_packet"
|
||||||
|
)";
|
||||||
|
runner_ = absl::make_unique<CalculatorRunner>(kConfig);
|
||||||
|
}
|
||||||
|
|
||||||
|
std::unique_ptr<CalculatorRunner> runner_;
|
||||||
|
};
|
||||||
|
|
||||||
|
TEST_F(StreamToSidePacketCalculatorTest,
|
||||||
|
StreamToSidePacketCalculatorWithEmptyStreamFails) {
|
||||||
|
EXPECT_EQ(runner_->Run().code(), mediapipe::StatusCode::kUnavailable);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(StreamToSidePacketCalculatorTest,
|
||||||
|
StreamToSidePacketCalculatorWithSinglePacketCreatesSidePacket) {
|
||||||
|
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||||
|
Adopt(new std::string("test")).At(Timestamp(1)));
|
||||||
|
MP_ASSERT_OK(runner_->Run());
|
||||||
|
EXPECT_EQ(runner_->OutputSidePackets().Index(0).Get<std::string>(), "test");
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(StreamToSidePacketCalculatorTest,
|
||||||
|
StreamToSidePacketCalculatorWithMultiplePacketsFails) {
|
||||||
|
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||||
|
Adopt(new std::string("test1")).At(Timestamp(1)));
|
||||||
|
runner_->MutableInputs()->Index(0).packets.push_back(
|
||||||
|
Adopt(new std::string("test2")).At(Timestamp(2)));
|
||||||
|
EXPECT_EQ(runner_->Run().code(), mediapipe::StatusCode::kAlreadyExists);
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,79 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include <sys/types.h>
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
#include "absl/strings/numbers.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/integral_types.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
// Calculator that converts a std::string into an integer type, or fails if the
|
||||||
|
// conversion is not possible.
|
||||||
|
//
|
||||||
|
// Example config:
|
||||||
|
// node {
|
||||||
|
// calculator: "StringToIntCalculator"
|
||||||
|
// input_side_packet: "string"
|
||||||
|
// output_side_packet: "index"
|
||||||
|
// }
|
||||||
|
template <typename IntType>
|
||||||
|
class StringToIntCalculatorTemplate : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
cc->InputSidePackets().Index(0).Set<std::string>();
|
||||||
|
cc->OutputSidePackets().Index(0).Set<IntType>();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
|
IntType number;
|
||||||
|
if (!absl::SimpleAtoi(cc->InputSidePackets().Index(0).Get<std::string>(),
|
||||||
|
&number)) {
|
||||||
|
return ::mediapipe::InvalidArgumentError(
|
||||||
|
"The std::string could not be parsed as an integer.");
|
||||||
|
}
|
||||||
|
cc->OutputSidePackets().Index(0).Set(MakePacket<IntType>(number));
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
using StringToIntCalculator = StringToIntCalculatorTemplate<int>;
|
||||||
|
REGISTER_CALCULATOR(StringToIntCalculator);
|
||||||
|
|
||||||
|
using StringToUintCalculator = StringToIntCalculatorTemplate<uint>;
|
||||||
|
REGISTER_CALCULATOR(StringToUintCalculator);
|
||||||
|
|
||||||
|
using StringToInt32Calculator = StringToIntCalculatorTemplate<int32>;
|
||||||
|
REGISTER_CALCULATOR(StringToInt32Calculator);
|
||||||
|
|
||||||
|
using StringToUint32Calculator = StringToIntCalculatorTemplate<uint32>;
|
||||||
|
REGISTER_CALCULATOR(StringToUint32Calculator);
|
||||||
|
|
||||||
|
using StringToInt64Calculator = StringToIntCalculatorTemplate<int64>;
|
||||||
|
REGISTER_CALCULATOR(StringToInt64Calculator);
|
||||||
|
|
||||||
|
using StringToUint64Calculator = StringToIntCalculatorTemplate<uint64>;
|
||||||
|
REGISTER_CALCULATOR(StringToUint64Calculator);
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -12,15 +12,14 @@
|
|||||||
# 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.
|
||||||
|
|
||||||
|
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||||
|
|
||||||
licenses(["notice"]) # Apache 2.0
|
licenses(["notice"]) # Apache 2.0
|
||||||
|
|
||||||
package(default_visibility = ["//visibility:private"])
|
package(default_visibility = ["//visibility:private"])
|
||||||
|
|
||||||
exports_files(["LICENSE"])
|
exports_files(["LICENSE"])
|
||||||
|
|
||||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
|
||||||
load("@bazel_skylib//lib:selects.bzl", "selects")
|
|
||||||
|
|
||||||
proto_library(
|
proto_library(
|
||||||
name = "opencv_image_encoder_calculator_proto",
|
name = "opencv_image_encoder_calculator_proto",
|
||||||
srcs = ["opencv_image_encoder_calculator.proto"],
|
srcs = ["opencv_image_encoder_calculator.proto"],
|
||||||
@@ -81,7 +80,9 @@ 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"],
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -227,19 +228,13 @@ cc_library(
|
|||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
"//mediapipe/framework/port:vector",
|
"//mediapipe/framework/port:vector",
|
||||||
] + select({
|
] + select({
|
||||||
"//mediapipe:android": [
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
|
"//conditions:default": [
|
||||||
"//mediapipe/gpu:gl_calculator_helper",
|
"//mediapipe/gpu:gl_calculator_helper",
|
||||||
"//mediapipe/gpu:gl_simple_shaders",
|
"//mediapipe/gpu:gl_simple_shaders",
|
||||||
"//mediapipe/gpu:gpu_buffer",
|
"//mediapipe/gpu:gl_quad_renderer",
|
||||||
"//mediapipe/gpu:shader_util",
|
"//mediapipe/gpu:shader_util",
|
||||||
],
|
],
|
||||||
"//mediapipe:ios": [
|
|
||||||
"//mediapipe/gpu:gl_calculator_helper",
|
|
||||||
"//mediapipe/gpu:gl_simple_shaders",
|
|
||||||
"//mediapipe/gpu:gpu_buffer",
|
|
||||||
"//mediapipe/gpu:shader_util",
|
|
||||||
],
|
|
||||||
"//conditions:default": [],
|
|
||||||
}),
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -263,13 +258,13 @@ cc_library(
|
|||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
"//mediapipe/framework/port:vector",
|
"//mediapipe/framework/port:vector",
|
||||||
] + select({
|
] + select({
|
||||||
"//mediapipe:android": [
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
|
"//conditions:default": [
|
||||||
"//mediapipe/gpu:gl_calculator_helper",
|
"//mediapipe/gpu:gl_calculator_helper",
|
||||||
"//mediapipe/gpu:gl_simple_shaders",
|
"//mediapipe/gpu:gl_simple_shaders",
|
||||||
"//mediapipe/gpu:gpu_buffer",
|
"//mediapipe/gpu:gl_quad_renderer",
|
||||||
"//mediapipe/gpu:shader_util",
|
"//mediapipe/gpu:shader_util",
|
||||||
],
|
],
|
||||||
"//conditions:default": [],
|
|
||||||
}),
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -322,14 +317,14 @@ cc_library(
|
|||||||
"//mediapipe/framework/port:opencv_imgproc",
|
"//mediapipe/framework/port:opencv_imgproc",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
] + selects.with_or({
|
] + select({
|
||||||
("//mediapipe:android", "//mediapipe:ios"): [
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
|
"//conditions:default": [
|
||||||
"//mediapipe/gpu:gl_calculator_helper",
|
"//mediapipe/gpu:gl_calculator_helper",
|
||||||
"//mediapipe/gpu:gl_simple_shaders",
|
"//mediapipe/gpu:gl_simple_shaders",
|
||||||
"//mediapipe/gpu:gl_quad_renderer",
|
"//mediapipe/gpu:gl_quad_renderer",
|
||||||
"//mediapipe/gpu:shader_util",
|
"//mediapipe/gpu:shader_util",
|
||||||
],
|
],
|
||||||
"//conditions:default": [],
|
|
||||||
}),
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -337,6 +332,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++",
|
||||||
@@ -350,9 +346,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",
|
||||||
@@ -363,18 +357,35 @@ cc_library(
|
|||||||
"//mediapipe/framework/port:opencv_imgproc",
|
"//mediapipe/framework/port:opencv_imgproc",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
] + selects.with_or({
|
] + select({
|
||||||
("//mediapipe:android", "//mediapipe:ios"): [
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
|
"//conditions:default": [
|
||||||
"//mediapipe/gpu:gl_calculator_helper",
|
"//mediapipe/gpu:gl_calculator_helper",
|
||||||
"//mediapipe/gpu:gl_simple_shaders",
|
"//mediapipe/gpu:gl_simple_shaders",
|
||||||
"//mediapipe/gpu:gl_quad_renderer",
|
"//mediapipe/gpu:gl_quad_renderer",
|
||||||
|
"//mediapipe/gpu:gpu_buffer",
|
||||||
"//mediapipe/gpu:shader_util",
|
"//mediapipe/gpu:shader_util",
|
||||||
],
|
],
|
||||||
"//conditions:default": [],
|
|
||||||
}),
|
}),
|
||||||
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"],
|
||||||
@@ -415,19 +426,13 @@ cc_library(
|
|||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/util:color_cc_proto",
|
"//mediapipe/util:color_cc_proto",
|
||||||
] + select({
|
] + select({
|
||||||
"//mediapipe:android": [
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
|
"//conditions:default": [
|
||||||
"//mediapipe/gpu:gl_calculator_helper",
|
"//mediapipe/gpu:gl_calculator_helper",
|
||||||
"//mediapipe/gpu:gl_simple_shaders",
|
"//mediapipe/gpu:gl_simple_shaders",
|
||||||
"//mediapipe/gpu:gpu_buffer",
|
"//mediapipe/gpu:gl_quad_renderer",
|
||||||
"//mediapipe/gpu:shader_util",
|
"//mediapipe/gpu:shader_util",
|
||||||
],
|
],
|
||||||
"//mediapipe:ios": [
|
|
||||||
"//mediapipe/gpu:gl_calculator_helper",
|
|
||||||
"//mediapipe/gpu:gl_simple_shaders",
|
|
||||||
"//mediapipe/gpu:gpu_buffer",
|
|
||||||
"//mediapipe/gpu:shader_util",
|
|
||||||
],
|
|
||||||
"//conditions:default": [],
|
|
||||||
}),
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -486,11 +491,11 @@ cc_library(
|
|||||||
"//mediapipe/framework/formats:image_frame",
|
"//mediapipe/framework/formats:image_frame",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
] + selects.with_or({
|
] + select({
|
||||||
("//mediapipe:android", "//mediapipe:ios"): [
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
|
"//conditions:default": [
|
||||||
"//mediapipe/gpu:gpu_buffer",
|
"//mediapipe/gpu:gpu_buffer",
|
||||||
],
|
],
|
||||||
"//conditions:default": [],
|
|
||||||
}),
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -27,11 +27,11 @@
|
|||||||
#include "mediapipe/framework/port/status.h"
|
#include "mediapipe/framework/port/status.h"
|
||||||
#include "mediapipe/framework/port/vector.h"
|
#include "mediapipe/framework/port/vector.h"
|
||||||
|
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||||
#include "mediapipe/gpu/gl_simple_shaders.h"
|
#include "mediapipe/gpu/gl_simple_shaders.h"
|
||||||
#include "mediapipe/gpu/shader_util.h"
|
#include "mediapipe/gpu/shader_util.h"
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
@@ -101,11 +101,11 @@ class BilateralFilterCalculator : public CalculatorBase {
|
|||||||
|
|
||||||
bool use_gpu_ = false;
|
bool use_gpu_ = false;
|
||||||
bool gpu_initialized_ = false;
|
bool gpu_initialized_ = false;
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||||
GLuint program_ = 0;
|
GLuint program_ = 0;
|
||||||
GLuint program_joint_ = 0;
|
GLuint program_joint_ = 0;
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
};
|
};
|
||||||
REGISTER_CALCULATOR(BilateralFilterCalculator);
|
REGISTER_CALCULATOR(BilateralFilterCalculator);
|
||||||
|
|
||||||
@@ -122,39 +122,46 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
|
|||||||
return ::mediapipe::InternalError("GPU output must have GPU input.");
|
return ::mediapipe::InternalError("GPU output must have GPU input.");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
bool use_gpu = false;
|
||||||
|
|
||||||
// Input image to filter.
|
// Input image to filter.
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Inputs().HasTag(kInputFrameTagGpu)) {
|
if (cc->Inputs().HasTag(kInputFrameTagGpu)) {
|
||||||
cc->Inputs().Tag(kInputFrameTagGpu).Set<mediapipe::GpuBuffer>();
|
cc->Inputs().Tag(kInputFrameTagGpu).Set<mediapipe::GpuBuffer>();
|
||||||
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
if (cc->Inputs().HasTag(kInputFrameTag)) {
|
if (cc->Inputs().HasTag(kInputFrameTag)) {
|
||||||
cc->Inputs().Tag(kInputFrameTag).Set<ImageFrame>();
|
cc->Inputs().Tag(kInputFrameTag).Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
|
|
||||||
// Input guide image mask (optional)
|
// Input guide image mask (optional)
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Inputs().HasTag(kInputGuideTagGpu)) {
|
if (cc->Inputs().HasTag(kInputGuideTagGpu)) {
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
|
||||||
cc->Inputs().Tag(kInputGuideTagGpu).Set<mediapipe::GpuBuffer>();
|
cc->Inputs().Tag(kInputGuideTagGpu).Set<mediapipe::GpuBuffer>();
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
if (cc->Inputs().HasTag(kInputGuideTag)) {
|
if (cc->Inputs().HasTag(kInputGuideTag)) {
|
||||||
cc->Inputs().Tag(kInputGuideTag).Set<ImageFrame>();
|
cc->Inputs().Tag(kInputGuideTag).Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
|
|
||||||
// Output image.
|
// Output image.
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Outputs().HasTag(kOutputFrameTagGpu)) {
|
if (cc->Outputs().HasTag(kOutputFrameTagGpu)) {
|
||||||
cc->Outputs().Tag(kOutputFrameTagGpu).Set<mediapipe::GpuBuffer>();
|
cc->Outputs().Tag(kOutputFrameTagGpu).Set<mediapipe::GpuBuffer>();
|
||||||
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
if (cc->Outputs().HasTag(kOutputFrameTag)) {
|
if (cc->Outputs().HasTag(kOutputFrameTag)) {
|
||||||
cc->Outputs().Tag(kOutputFrameTag).Set<ImageFrame>();
|
cc->Outputs().Tag(kOutputFrameTag).Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
|
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
if (use_gpu) {
|
||||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -166,11 +173,11 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
|
|||||||
|
|
||||||
if (cc->Inputs().HasTag(kInputFrameTagGpu) &&
|
if (cc->Inputs().HasTag(kInputFrameTagGpu) &&
|
||||||
cc->Outputs().HasTag(kOutputFrameTagGpu)) {
|
cc->Outputs().HasTag(kOutputFrameTagGpu)) {
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
use_gpu_ = true;
|
use_gpu_ = true;
|
||||||
#else
|
#else
|
||||||
RET_CHECK_FAIL() << "GPU processing on non-Android not supported yet.";
|
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
sigma_color_ = options_.sigma_color();
|
sigma_color_ = options_.sigma_color();
|
||||||
@@ -180,9 +187,9 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
|
|||||||
if (!use_gpu_) sigma_color_ *= 255.0;
|
if (!use_gpu_) sigma_color_ *= 255.0;
|
||||||
|
|
||||||
if (use_gpu_) {
|
if (use_gpu_) {
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||||
#endif
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
@@ -190,7 +197,7 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
|
|||||||
|
|
||||||
::mediapipe::Status BilateralFilterCalculator::Process(CalculatorContext* cc) {
|
::mediapipe::Status BilateralFilterCalculator::Process(CalculatorContext* cc) {
|
||||||
if (use_gpu_) {
|
if (use_gpu_) {
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
MP_RETURN_IF_ERROR(
|
MP_RETURN_IF_ERROR(
|
||||||
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
|
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
|
||||||
if (!gpu_initialized_) {
|
if (!gpu_initialized_) {
|
||||||
@@ -200,7 +207,7 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
|
|||||||
MP_RETURN_IF_ERROR(RenderGpu(cc));
|
MP_RETURN_IF_ERROR(RenderGpu(cc));
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}));
|
}));
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
} else {
|
} else {
|
||||||
MP_RETURN_IF_ERROR(RenderCpu(cc));
|
MP_RETURN_IF_ERROR(RenderCpu(cc));
|
||||||
}
|
}
|
||||||
@@ -209,14 +216,14 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status BilateralFilterCalculator::Close(CalculatorContext* cc) {
|
::mediapipe::Status BilateralFilterCalculator::Close(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
gpu_helper_.RunInGlContext([this] {
|
gpu_helper_.RunInGlContext([this] {
|
||||||
if (program_) glDeleteProgram(program_);
|
if (program_) glDeleteProgram(program_);
|
||||||
program_ = 0;
|
program_ = 0;
|
||||||
if (program_joint_) glDeleteProgram(program_joint_);
|
if (program_joint_) glDeleteProgram(program_joint_);
|
||||||
program_joint_ = 0;
|
program_joint_ = 0;
|
||||||
});
|
});
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -263,7 +270,7 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
|
|||||||
if (cc->Inputs().Tag(kInputFrameTagGpu).IsEmpty()) {
|
if (cc->Inputs().Tag(kInputFrameTagGpu).IsEmpty()) {
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
const auto& input_frame =
|
const auto& input_frame =
|
||||||
cc->Inputs().Tag(kInputFrameTagGpu).Get<mediapipe::GpuBuffer>();
|
cc->Inputs().Tag(kInputFrameTagGpu).Get<mediapipe::GpuBuffer>();
|
||||||
auto input_texture = gpu_helper_.CreateSourceTexture(input_frame);
|
auto input_texture = gpu_helper_.CreateSourceTexture(input_frame);
|
||||||
@@ -321,13 +328,13 @@ REGISTER_CALCULATOR(BilateralFilterCalculator);
|
|||||||
// Cleanup
|
// Cleanup
|
||||||
input_texture.Release();
|
input_texture.Release();
|
||||||
output_texture.Release();
|
output_texture.Release();
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
void BilateralFilterCalculator::GlRender(CalculatorContext* cc) {
|
void BilateralFilterCalculator::GlRender(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
static const GLfloat square_vertices[] = {
|
static const GLfloat square_vertices[] = {
|
||||||
-1.0f, -1.0f, // bottom left
|
-1.0f, -1.0f, // bottom left
|
||||||
1.0f, -1.0f, // bottom right
|
1.0f, -1.0f, // bottom right
|
||||||
@@ -373,11 +380,11 @@ void BilateralFilterCalculator::GlRender(CalculatorContext* cc) {
|
|||||||
glDeleteVertexArrays(1, &vao);
|
glDeleteVertexArrays(1, &vao);
|
||||||
glDeleteBuffers(2, vbo);
|
glDeleteBuffers(2, vbo);
|
||||||
|
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status BilateralFilterCalculator::GlSetup(CalculatorContext* cc) {
|
::mediapipe::Status BilateralFilterCalculator::GlSetup(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || defined(__EMSCRIPTEN__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
const GLint attr_location[NUM_ATTRIBUTES] = {
|
const GLint attr_location[NUM_ATTRIBUTES] = {
|
||||||
ATTRIB_VERTEX,
|
ATTRIB_VERTEX,
|
||||||
ATTRIB_TEXTURE_POSITION,
|
ATTRIB_TEXTURE_POSITION,
|
||||||
@@ -545,7 +552,7 @@ void BilateralFilterCalculator::GlRender(CalculatorContext* cc) {
|
|||||||
glUniform1i(glGetUniformLocation(program_joint_, "input_frame"), 1);
|
glUniform1i(glGetUniformLocation(program_joint_, "input_frame"), 1);
|
||||||
glUniform1i(glGetUniformLocation(program_joint_, "guide_frame"), 2);
|
glUniform1i(glGetUniformLocation(program_joint_, "guide_frame"), 2);
|
||||||
|
|
||||||
#endif // __ANDROID__ || __EMSCRIPTEN__
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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"
|
||||||
@@ -24,12 +24,11 @@
|
|||||||
#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"
|
||||||
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#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"
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
|
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
|
||||||
@@ -37,102 +36,85 @@ enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
|
|||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
namespace {
|
||||||
|
|
||||||
#endif // __ANDROID__ or iOS
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
|
|
||||||
// Crops the input texture to the given rectangle region. The rectangle can
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
// 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);
|
constexpr char kRectTag[] = "RECT";
|
||||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
constexpr char kNormRectTag[] = "NORM_RECT";
|
||||||
::mediapipe::Status Process(CalculatorContext* cc) override;
|
constexpr char kHeightTag[] = "HEIGHT";
|
||||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
constexpr char kImageTag[] = "IMAGE";
|
||||||
|
constexpr char kImageGpuTag[] = "IMAGE_GPU";
|
||||||
|
constexpr char kWidthTag[] = "WIDTH";
|
||||||
|
|
||||||
private:
|
} // namespace
|
||||||
::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(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
|
||||||
bool gpu_initialized_ = false;
|
|
||||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
|
||||||
GLuint program_ = 0;
|
|
||||||
#endif // __ANDROID__ or iOS
|
|
||||||
};
|
|
||||||
REGISTER_CALCULATOR(ImageCroppingCalculator);
|
REGISTER_CALCULATOR(ImageCroppingCalculator);
|
||||||
|
|
||||||
::mediapipe::Status ImageCroppingCalculator::GetContract(
|
::mediapipe::Status ImageCroppingCalculator::GetContract(
|
||||||
CalculatorContract* cc) {
|
CalculatorContract* cc) {
|
||||||
RET_CHECK(cc->Inputs().HasTag("IMAGE") ^ cc->Inputs().HasTag("IMAGE_GPU"));
|
RET_CHECK(cc->Inputs().HasTag(kImageTag) ^ cc->Inputs().HasTag(kImageGpuTag));
|
||||||
RET_CHECK(cc->Outputs().HasTag("IMAGE") ^ cc->Outputs().HasTag("IMAGE_GPU"));
|
RET_CHECK(cc->Outputs().HasTag(kImageTag) ^
|
||||||
|
cc->Outputs().HasTag(kImageGpuTag));
|
||||||
|
|
||||||
if (cc->Inputs().HasTag("IMAGE")) {
|
bool use_gpu = false;
|
||||||
RET_CHECK(cc->Outputs().HasTag("IMAGE"));
|
|
||||||
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
|
||||||
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
|
|
||||||
}
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
|
||||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
|
||||||
RET_CHECK(cc->Outputs().HasTag("IMAGE_GPU"));
|
|
||||||
cc->Inputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
|
|
||||||
cc->Outputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
|
|
||||||
}
|
|
||||||
#endif // __ANDROID__ or iOS
|
|
||||||
|
|
||||||
if (cc->Inputs().HasTag("RECT")) {
|
if (cc->Inputs().HasTag(kImageTag)) {
|
||||||
cc->Inputs().Tag("RECT").Set<Rect>();
|
RET_CHECK(cc->Outputs().HasTag(kImageTag));
|
||||||
|
cc->Inputs().Tag(kImageTag).Set<ImageFrame>();
|
||||||
|
cc->Outputs().Tag(kImageTag).Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
if (cc->Inputs().HasTag("NORM_RECT")) {
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
cc->Inputs().Tag("NORM_RECT").Set<NormalizedRect>();
|
if (cc->Inputs().HasTag(kImageGpuTag)) {
|
||||||
|
RET_CHECK(cc->Outputs().HasTag(kImageGpuTag));
|
||||||
|
cc->Inputs().Tag(kImageGpuTag).Set<GpuBuffer>();
|
||||||
|
cc->Outputs().Tag(kImageGpuTag).Set<GpuBuffer>();
|
||||||
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
if (cc->Inputs().HasTag("WIDTH")) {
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
cc->Inputs().Tag("WIDTH").Set<int>();
|
|
||||||
|
int flags = 0;
|
||||||
|
if (cc->Inputs().HasTag(kRectTag)) {
|
||||||
|
++flags;
|
||||||
}
|
}
|
||||||
if (cc->Inputs().HasTag("HEIGHT")) {
|
if (cc->Inputs().HasTag(kWidthTag) && cc->Inputs().HasTag(kHeightTag)) {
|
||||||
cc->Inputs().Tag("HEIGHT").Set<int>();
|
++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)) {
|
||||||
|
cc->Inputs().Tag(kRectTag).Set<Rect>();
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag(kNormRectTag)) {
|
||||||
|
cc->Inputs().Tag(kNormRectTag).Set<NormalizedRect>();
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag(kWidthTag)) {
|
||||||
|
cc->Inputs().Tag(kWidthTag).Set<int>();
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag(kHeightTag)) {
|
||||||
|
cc->Inputs().Tag(kHeightTag).Set<int>();
|
||||||
}
|
}
|
||||||
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
if (use_gpu) {
|
||||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
#endif // __ANDROID__ or iOS
|
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -140,26 +122,42 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
|
|||||||
::mediapipe::Status ImageCroppingCalculator::Open(CalculatorContext* cc) {
|
::mediapipe::Status ImageCroppingCalculator::Open(CalculatorContext* cc) {
|
||||||
cc->SetOffset(TimestampDiff(0));
|
cc->SetOffset(TimestampDiff(0));
|
||||||
|
|
||||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
if (cc->Inputs().HasTag(kImageGpuTag)) {
|
||||||
use_gpu_ = true;
|
use_gpu_ = true;
|
||||||
}
|
}
|
||||||
|
|
||||||
options_ = cc->Options<mediapipe::ImageCroppingCalculatorOptions>();
|
options_ = cc->Options<mediapipe::ImageCroppingCalculatorOptions>();
|
||||||
|
|
||||||
if (use_gpu_) {
|
if (use_gpu_) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||||
#else
|
#else
|
||||||
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
|
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
|
||||||
#endif // __ANDROID__ or iOS
|
#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();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status ImageCroppingCalculator::Process(CalculatorContext* cc) {
|
::mediapipe::Status ImageCroppingCalculator::Process(CalculatorContext* cc) {
|
||||||
|
if (cc->Inputs().HasTag(kRectTag) && cc->Inputs().Tag(kRectTag).IsEmpty()) {
|
||||||
|
VLOG(1) << "RECT is empty for timestamp: " << cc->InputTimestamp();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
if (cc->Inputs().HasTag(kNormRectTag) &&
|
||||||
|
cc->Inputs().Tag(kNormRectTag).IsEmpty()) {
|
||||||
|
VLOG(1) << "NORM_RECT is empty for timestamp: " << cc->InputTimestamp();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
if (use_gpu_) {
|
if (use_gpu_) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
MP_RETURN_IF_ERROR(
|
MP_RETURN_IF_ERROR(
|
||||||
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
|
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
|
||||||
if (!gpu_initialized_) {
|
if (!gpu_initialized_) {
|
||||||
@@ -169,7 +167,7 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
|
|||||||
MP_RETURN_IF_ERROR(RenderGpu(cc));
|
MP_RETURN_IF_ERROR(RenderGpu(cc));
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}));
|
}));
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
} else {
|
} else {
|
||||||
MP_RETURN_IF_ERROR(RenderCpu(cc));
|
MP_RETURN_IF_ERROR(RenderCpu(cc));
|
||||||
}
|
}
|
||||||
@@ -177,56 +175,56 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status ImageCroppingCalculator::Close(CalculatorContext* cc) {
|
::mediapipe::Status ImageCroppingCalculator::Close(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
gpu_helper_.RunInGlContext([this] {
|
gpu_helper_.RunInGlContext([this] {
|
||||||
if (program_) glDeleteProgram(program_);
|
if (program_) glDeleteProgram(program_);
|
||||||
program_ = 0;
|
program_ = 0;
|
||||||
});
|
});
|
||||||
gpu_initialized_ = false;
|
gpu_initialized_ = false;
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
|
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();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status ImageCroppingCalculator::RenderCpu(CalculatorContext* cc) {
|
::mediapipe::Status ImageCroppingCalculator::RenderCpu(CalculatorContext* cc) {
|
||||||
const auto& input_img = cc->Inputs().Tag("IMAGE").Get<ImageFrame>();
|
if (cc->Inputs().Tag(kImageTag).IsEmpty()) {
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
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;
|
auto [target_width, target_height, rect_center_x, rect_center_y, rotation] =
|
||||||
float rect_center_y = input_img.Height() / 2.0f;
|
GetCropSpecs(cc, input_img.Width(), input_img.Height());
|
||||||
float rotation = 0.0f;
|
|
||||||
int target_width = input_img.Width();
|
// Get border mode and value for OpenCV.
|
||||||
int target_height = input_img.Height();
|
int border_mode;
|
||||||
if (cc->Inputs().HasTag("RECT")) {
|
MP_RETURN_IF_ERROR(GetBorderModeForOpenCV(cc, &border_mode));
|
||||||
const auto& rect = cc->Inputs().Tag("RECT").Get<Rect>();
|
|
||||||
if (rect.width() > 0 && rect.height() > 0 && rect.x_center() >= 0 &&
|
|
||||||
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("NORM_RECT")) {
|
|
||||||
const auto& rect = cc->Inputs().Tag("NORM_RECT").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("WIDTH") && cc->Inputs().HasTag("HEIGHT")) {
|
|
||||||
target_width = cc->Inputs().Tag("WIDTH").Get<int>();
|
|
||||||
target_height = cc->Inputs().Tag("HEIGHT").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),
|
||||||
@@ -247,22 +245,25 @@ 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));
|
||||||
cv::Mat output_mat = formats::MatView(output_frame.get());
|
cv::Mat output_mat = formats::MatView(output_frame.get());
|
||||||
cropped_image.copyTo(output_mat);
|
cropped_image.copyTo(output_mat);
|
||||||
cc->Outputs().Tag("IMAGE").Add(output_frame.release(), cc->InputTimestamp());
|
cc->Outputs().Tag(kImageTag).Add(output_frame.release(),
|
||||||
|
cc->InputTimestamp());
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status ImageCroppingCalculator::RenderGpu(CalculatorContext* cc) {
|
::mediapipe::Status ImageCroppingCalculator::RenderGpu(CalculatorContext* cc) {
|
||||||
if (cc->Inputs().Tag("IMAGE_GPU").IsEmpty()) {
|
if (cc->Inputs().Tag(kImageGpuTag).IsEmpty()) {
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
const Packet& input_packet = cc->Inputs().Tag("IMAGE_GPU").Value();
|
const Packet& input_packet = cc->Inputs().Tag(kImageGpuTag).Value();
|
||||||
const auto& input_buffer = input_packet.Get<mediapipe::GpuBuffer>();
|
const auto& input_buffer = input_packet.Get<mediapipe::GpuBuffer>();
|
||||||
auto src_tex = gpu_helper_.CreateSourceTexture(input_buffer);
|
auto src_tex = gpu_helper_.CreateSourceTexture(input_buffer);
|
||||||
|
|
||||||
@@ -287,18 +288,18 @@ REGISTER_CALCULATOR(ImageCroppingCalculator);
|
|||||||
|
|
||||||
// Send result image in GPU packet.
|
// Send result image in GPU packet.
|
||||||
auto output = dst_tex.GetFrame<mediapipe::GpuBuffer>();
|
auto output = dst_tex.GetFrame<mediapipe::GpuBuffer>();
|
||||||
cc->Outputs().Tag("IMAGE_GPU").Add(output.release(), cc->InputTimestamp());
|
cc->Outputs().Tag(kImageGpuTag).Add(output.release(), cc->InputTimestamp());
|
||||||
|
|
||||||
// Cleanup
|
// Cleanup
|
||||||
src_tex.Release();
|
src_tex.Release();
|
||||||
dst_tex.Release();
|
dst_tex.Release();
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
void ImageCroppingCalculator::GlRender() {
|
void ImageCroppingCalculator::GlRender() {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
static const GLfloat square_vertices[] = {
|
static const GLfloat square_vertices[] = {
|
||||||
-1.0f, -1.0f, // bottom left
|
-1.0f, -1.0f, // bottom left
|
||||||
1.0f, -1.0f, // bottom right
|
1.0f, -1.0f, // bottom right
|
||||||
@@ -342,11 +343,11 @@ void ImageCroppingCalculator::GlRender() {
|
|||||||
glDeleteVertexArrays(1, &vao);
|
glDeleteVertexArrays(1, &vao);
|
||||||
glDeleteBuffers(2, vbo);
|
glDeleteBuffers(2, vbo);
|
||||||
|
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status ImageCroppingCalculator::InitGpu(CalculatorContext* cc) {
|
::mediapipe::Status ImageCroppingCalculator::InitGpu(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
const GLint attr_location[NUM_ATTRIBUTES] = {
|
const GLint attr_location[NUM_ATTRIBUTES] = {
|
||||||
ATTRIB_VERTEX,
|
ATTRIB_VERTEX,
|
||||||
ATTRIB_TEXTURE_POSITION,
|
ATTRIB_TEXTURE_POSITION,
|
||||||
@@ -392,7 +393,7 @@ void ImageCroppingCalculator::GlRender() {
|
|||||||
// Parameters
|
// Parameters
|
||||||
glUseProgram(program_);
|
glUseProgram(program_);
|
||||||
glUniform1i(glGetUniformLocation(program_, "input_frame"), 1);
|
glUniform1i(glGetUniformLocation(program_, "input_frame"), 1);
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -402,46 +403,8 @@ 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.
|
auto [crop_width, crop_height, x_center, y_center, rotation] =
|
||||||
int crop_width = src_width;
|
GetCropSpecs(cc, src_width, src_height);
|
||||||
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;
|
|
||||||
if (cc->Inputs().HasTag("RECT")) {
|
|
||||||
const auto& rect = cc->Inputs().Tag("RECT").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("NORM_RECT")) {
|
|
||||||
const auto& rect = cc->Inputs().Tag("NORM_RECT").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("WIDTH") && cc->Inputs().HasTag("HEIGHT")) {
|
|
||||||
crop_width = cc->Inputs().Tag("WIDTH").Get<int>();
|
|
||||||
crop_height = cc->Inputs().Tag("HEIGHT").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;
|
||||||
@@ -470,8 +433,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
|
||||||
@@ -15,9 +15,9 @@
|
|||||||
#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"
|
||||||
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
#include "mediapipe/gpu/gpu_buffer.h"
|
#include "mediapipe/gpu/gpu_buffer.h"
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
@@ -44,11 +44,11 @@ class ImagePropertiesCalculator : public CalculatorBase {
|
|||||||
if (cc->Inputs().HasTag("IMAGE")) {
|
if (cc->Inputs().HasTag("IMAGE")) {
|
||||||
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||||
cc->Inputs().Tag("IMAGE_GPU").Set<::mediapipe::GpuBuffer>();
|
cc->Inputs().Tag("IMAGE_GPU").Set<::mediapipe::GpuBuffer>();
|
||||||
}
|
}
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
if (cc->Outputs().HasTag("SIZE")) {
|
if (cc->Outputs().HasTag("SIZE")) {
|
||||||
cc->Outputs().Tag("SIZE").Set<std::pair<int, int>>();
|
cc->Outputs().Tag("SIZE").Set<std::pair<int, int>>();
|
||||||
@@ -71,7 +71,7 @@ class ImagePropertiesCalculator : public CalculatorBase {
|
|||||||
width = image.Width();
|
width = image.Width();
|
||||||
height = image.Height();
|
height = image.Height();
|
||||||
}
|
}
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Inputs().HasTag("IMAGE_GPU") &&
|
if (cc->Inputs().HasTag("IMAGE_GPU") &&
|
||||||
!cc->Inputs().Tag("IMAGE_GPU").IsEmpty()) {
|
!cc->Inputs().Tag("IMAGE_GPU").IsEmpty()) {
|
||||||
const auto& image =
|
const auto& image =
|
||||||
@@ -79,7 +79,7 @@ class ImagePropertiesCalculator : public CalculatorBase {
|
|||||||
width = image.width();
|
width = image.width();
|
||||||
height = image.height();
|
height = image.height();
|
||||||
}
|
}
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
cc->Outputs().Tag("SIZE").AddPacket(
|
cc->Outputs().Tag("SIZE").AddPacket(
|
||||||
MakePacket<std::pair<int, int>>(width, height)
|
MakePacket<std::pair<int, int>>(width, height)
|
||||||
|
|||||||
@@ -22,12 +22,12 @@
|
|||||||
#include "mediapipe/framework/port/status.h"
|
#include "mediapipe/framework/port/status.h"
|
||||||
#include "mediapipe/gpu/scale_mode.pb.h"
|
#include "mediapipe/gpu/scale_mode.pb.h"
|
||||||
|
|
||||||
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||||
#include "mediapipe/gpu/gl_quad_renderer.h"
|
#include "mediapipe/gpu/gl_quad_renderer.h"
|
||||||
#include "mediapipe/gpu/gl_simple_shaders.h"
|
#include "mediapipe/gpu/gl_simple_shaders.h"
|
||||||
#include "mediapipe/gpu/shader_util.h"
|
#include "mediapipe/gpu/shader_util.h"
|
||||||
#endif // __ANDROID__ || iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if defined(__ANDROID__)
|
||||||
// The size of Java arrays is dynamic, which makes it difficult to
|
// The size of Java arrays is dynamic, which makes it difficult to
|
||||||
@@ -42,9 +42,9 @@ typedef int DimensionsPacketType[2];
|
|||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
|
|
||||||
#endif // __ANDROID__ || iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
int RotationModeToDegrees(mediapipe::RotationMode_Mode rotation) {
|
int RotationModeToDegrees(mediapipe::RotationMode_Mode rotation) {
|
||||||
@@ -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,14 +181,16 @@ 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(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
GlCalculatorHelper helper_;
|
GlCalculatorHelper helper_;
|
||||||
std::unique_ptr<QuadRenderer> rgb_renderer_;
|
std::unique_ptr<QuadRenderer> rgb_renderer_;
|
||||||
std::unique_ptr<QuadRenderer> yuv_renderer_;
|
std::unique_ptr<QuadRenderer> yuv_renderer_;
|
||||||
std::unique_ptr<QuadRenderer> ext_rgb_renderer_;
|
std::unique_ptr<QuadRenderer> ext_rgb_renderer_;
|
||||||
#endif // __ANDROID__ || iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
};
|
};
|
||||||
REGISTER_CALCULATOR(ImageTransformationCalculator);
|
REGISTER_CALCULATOR(ImageTransformationCalculator);
|
||||||
|
|
||||||
@@ -185,21 +200,31 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
RET_CHECK(cc->Inputs().HasTag("IMAGE") ^ cc->Inputs().HasTag("IMAGE_GPU"));
|
RET_CHECK(cc->Inputs().HasTag("IMAGE") ^ cc->Inputs().HasTag("IMAGE_GPU"));
|
||||||
RET_CHECK(cc->Outputs().HasTag("IMAGE") ^ cc->Outputs().HasTag("IMAGE_GPU"));
|
RET_CHECK(cc->Outputs().HasTag("IMAGE") ^ cc->Outputs().HasTag("IMAGE_GPU"));
|
||||||
|
|
||||||
|
bool use_gpu = false;
|
||||||
|
|
||||||
if (cc->Inputs().HasTag("IMAGE")) {
|
if (cc->Inputs().HasTag("IMAGE")) {
|
||||||
RET_CHECK(cc->Outputs().HasTag("IMAGE"));
|
RET_CHECK(cc->Outputs().HasTag("IMAGE"));
|
||||||
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
||||||
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
|
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||||
RET_CHECK(cc->Outputs().HasTag("IMAGE_GPU"));
|
RET_CHECK(cc->Outputs().HasTag("IMAGE_GPU"));
|
||||||
cc->Inputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
|
cc->Inputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
|
||||||
cc->Outputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
|
cc->Outputs().Tag("IMAGE_GPU").Set<GpuBuffer>();
|
||||||
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
#endif // __ANDROID__ || iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
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>();
|
||||||
@@ -207,14 +232,22 @@ 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>>();
|
||||||
}
|
}
|
||||||
|
|
||||||
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
|
if (use_gpu) {
|
||||||
MP_RETURN_IF_ERROR(GlCalculatorHelper::UpdateContract(cc));
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
#endif // __ANDROID__ || iOS
|
MP_RETURN_IF_ERROR(GlCalculatorHelper::UpdateContract(cc));
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -240,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>());
|
||||||
@@ -247,15 +281,29 @@ 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_) {
|
||||||
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
// Let the helper access the GL context information.
|
// Let the helper access the GL context information.
|
||||||
MP_RETURN_IF_ERROR(helper_.Open(cc));
|
MP_RETURN_IF_ERROR(helper_.Open(cc));
|
||||||
#else
|
#else
|
||||||
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
|
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
||||||
#endif // __ANDROID__ || iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
@@ -263,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(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
|
#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 // __ANDROID__ || iOS
|
#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();
|
||||||
@@ -277,7 +350,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
::mediapipe::Status ImageTransformationCalculator::Close(
|
::mediapipe::Status ImageTransformationCalculator::Close(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
if (use_gpu_) {
|
if (use_gpu_) {
|
||||||
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
QuadRenderer* rgb_renderer = rgb_renderer_.release();
|
QuadRenderer* rgb_renderer = rgb_renderer_.release();
|
||||||
QuadRenderer* yuv_renderer = yuv_renderer_.release();
|
QuadRenderer* yuv_renderer = yuv_renderer_.release();
|
||||||
QuadRenderer* ext_rgb_renderer = ext_rgb_renderer_.release();
|
QuadRenderer* ext_rgb_renderer = ext_rgb_renderer_.release();
|
||||||
@@ -295,8 +368,9 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
delete yuv_renderer;
|
delete yuv_renderer;
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
#endif // __ANDROID__ || iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -309,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 {
|
||||||
@@ -349,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();
|
||||||
@@ -371,7 +454,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
|
|
||||||
::mediapipe::Status ImageTransformationCalculator::RenderGpu(
|
::mediapipe::Status ImageTransformationCalculator::RenderGpu(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || defined(__APPLE__) && !TARGET_OS_OSX
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
int input_width = cc->Inputs().Tag("IMAGE_GPU").Get<GpuBuffer>().width();
|
int input_width = cc->Inputs().Tag("IMAGE_GPU").Get<GpuBuffer>().width();
|
||||||
int input_height = cc->Inputs().Tag("IMAGE_GPU").Get<GpuBuffer>().height();
|
int input_height = cc->Inputs().Tag("IMAGE_GPU").Get<GpuBuffer>().height();
|
||||||
|
|
||||||
@@ -393,7 +476,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
QuadRenderer* renderer = nullptr;
|
QuadRenderer* renderer = nullptr;
|
||||||
GlTexture src1;
|
GlTexture src1;
|
||||||
|
|
||||||
#if defined(__APPLE__) && !TARGET_OS_OSX
|
#if defined(MEDIAPIPE_IOS)
|
||||||
if (input.format() == GpuBufferFormat::kBiPlanar420YpCbCr8VideoRange ||
|
if (input.format() == GpuBufferFormat::kBiPlanar420YpCbCr8VideoRange ||
|
||||||
input.format() == GpuBufferFormat::kBiPlanar420YpCbCr8FullRange) {
|
input.format() == GpuBufferFormat::kBiPlanar420YpCbCr8FullRange) {
|
||||||
if (!yuv_renderer_) {
|
if (!yuv_renderer_) {
|
||||||
@@ -408,7 +491,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
#endif // iOS
|
#endif // iOS
|
||||||
{
|
{
|
||||||
src1 = helper_.CreateSourceTexture(input);
|
src1 = helper_.CreateSourceTexture(input);
|
||||||
#if defined(__ANDROID__)
|
#if defined(TEXTURE_EXTERNAL_OES)
|
||||||
if (src1.target() == GL_TEXTURE_EXTERNAL_OES) {
|
if (src1.target() == GL_TEXTURE_EXTERNAL_OES) {
|
||||||
if (!ext_rgb_renderer_) {
|
if (!ext_rgb_renderer_) {
|
||||||
ext_rgb_renderer_ = absl::make_unique<QuadRenderer>();
|
ext_rgb_renderer_ = absl::make_unique<QuadRenderer>();
|
||||||
@@ -417,7 +500,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
}
|
}
|
||||||
renderer = ext_rgb_renderer_.get();
|
renderer = ext_rgb_renderer_.get();
|
||||||
} else // NOLINT(readability/braces)
|
} else // NOLINT(readability/braces)
|
||||||
#endif // __ANDROID__
|
#endif // TEXTURE_EXTERNAL_OES
|
||||||
{
|
{
|
||||||
if (!rgb_renderer_) {
|
if (!rgb_renderer_) {
|
||||||
rgb_renderer_ = absl::make_unique<QuadRenderer>();
|
rgb_renderer_ = absl::make_unique<QuadRenderer>();
|
||||||
@@ -428,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_));
|
||||||
|
|
||||||
@@ -448,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);
|
||||||
@@ -460,7 +537,7 @@ REGISTER_CALCULATOR(ImageTransformationCalculator);
|
|||||||
auto output = dst.GetFrame<GpuBuffer>();
|
auto output = dst.GetFrame<GpuBuffer>();
|
||||||
cc->Outputs().Tag("IMAGE_GPU").Add(output.release(), cc->InputTimestamp());
|
cc->Outputs().Tag("IMAGE_GPU").Add(output.release(), cc->InputTimestamp());
|
||||||
|
|
||||||
#endif // __ANDROID__ || iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -21,12 +21,11 @@
|
|||||||
#include "mediapipe/framework/port/status.h"
|
#include "mediapipe/framework/port/status.h"
|
||||||
#include "mediapipe/util/color.pb.h"
|
#include "mediapipe/util/color.pb.h"
|
||||||
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
#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/shader_util.h"
|
#include "mediapipe/gpu/shader_util.h"
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
|
enum { ATTRIB_VERTEX, ATTRIB_TEXTURE_POSITION, NUM_ATTRIBUTES };
|
||||||
@@ -95,10 +94,10 @@ class RecolorCalculator : public CalculatorBase {
|
|||||||
mediapipe::RecolorCalculatorOptions::MaskChannel mask_channel_;
|
mediapipe::RecolorCalculatorOptions::MaskChannel mask_channel_;
|
||||||
|
|
||||||
bool use_gpu_ = false;
|
bool use_gpu_ = false;
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||||
GLuint program_ = 0;
|
GLuint program_ = 0;
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
};
|
};
|
||||||
REGISTER_CALCULATOR(RecolorCalculator);
|
REGISTER_CALCULATOR(RecolorCalculator);
|
||||||
|
|
||||||
@@ -107,36 +106,43 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
|||||||
RET_CHECK(!cc->Inputs().GetTags().empty());
|
RET_CHECK(!cc->Inputs().GetTags().empty());
|
||||||
RET_CHECK(!cc->Outputs().GetTags().empty());
|
RET_CHECK(!cc->Outputs().GetTags().empty());
|
||||||
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
bool use_gpu = false;
|
||||||
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||||
cc->Inputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
|
cc->Inputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
|
||||||
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
if (cc->Inputs().HasTag("IMAGE")) {
|
if (cc->Inputs().HasTag("IMAGE")) {
|
||||||
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Inputs().HasTag("MASK_GPU")) {
|
if (cc->Inputs().HasTag("MASK_GPU")) {
|
||||||
cc->Inputs().Tag("MASK_GPU").Set<mediapipe::GpuBuffer>();
|
cc->Inputs().Tag("MASK_GPU").Set<mediapipe::GpuBuffer>();
|
||||||
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
if (cc->Inputs().HasTag("MASK")) {
|
if (cc->Inputs().HasTag("MASK")) {
|
||||||
cc->Inputs().Tag("MASK").Set<ImageFrame>();
|
cc->Inputs().Tag("MASK").Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Outputs().HasTag("IMAGE_GPU")) {
|
if (cc->Outputs().HasTag("IMAGE_GPU")) {
|
||||||
cc->Outputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
|
cc->Outputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
|
||||||
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
if (cc->Outputs().HasTag("IMAGE")) {
|
if (cc->Outputs().HasTag("IMAGE")) {
|
||||||
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
|
cc->Outputs().Tag("IMAGE").Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
if (use_gpu) {
|
||||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
#endif // __ANDROID__ or iOS
|
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -146,9 +152,9 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
|||||||
|
|
||||||
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||||
use_gpu_ = true;
|
use_gpu_ = true;
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
MP_RETURN_IF_ERROR(LoadOptions(cc));
|
MP_RETURN_IF_ERROR(LoadOptions(cc));
|
||||||
@@ -158,7 +164,7 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
|||||||
|
|
||||||
::mediapipe::Status RecolorCalculator::Process(CalculatorContext* cc) {
|
::mediapipe::Status RecolorCalculator::Process(CalculatorContext* cc) {
|
||||||
if (use_gpu_) {
|
if (use_gpu_) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
MP_RETURN_IF_ERROR(
|
MP_RETURN_IF_ERROR(
|
||||||
gpu_helper_.RunInGlContext([this, &cc]() -> ::mediapipe::Status {
|
gpu_helper_.RunInGlContext([this, &cc]() -> ::mediapipe::Status {
|
||||||
if (!initialized_) {
|
if (!initialized_) {
|
||||||
@@ -168,7 +174,7 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
|||||||
MP_RETURN_IF_ERROR(RenderGpu(cc));
|
MP_RETURN_IF_ERROR(RenderGpu(cc));
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}));
|
}));
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
} else {
|
} else {
|
||||||
MP_RETURN_IF_ERROR(RenderCpu(cc));
|
MP_RETURN_IF_ERROR(RenderCpu(cc));
|
||||||
}
|
}
|
||||||
@@ -176,12 +182,12 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status RecolorCalculator::Close(CalculatorContext* cc) {
|
::mediapipe::Status RecolorCalculator::Close(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
gpu_helper_.RunInGlContext([this] {
|
gpu_helper_.RunInGlContext([this] {
|
||||||
if (program_) glDeleteProgram(program_);
|
if (program_) glDeleteProgram(program_);
|
||||||
program_ = 0;
|
program_ = 0;
|
||||||
});
|
});
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -194,7 +200,7 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
|||||||
if (cc->Inputs().Tag("MASK_GPU").IsEmpty()) {
|
if (cc->Inputs().Tag("MASK_GPU").IsEmpty()) {
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
// Get inputs and setup output.
|
// Get inputs and setup output.
|
||||||
const Packet& input_packet = cc->Inputs().Tag("IMAGE_GPU").Value();
|
const Packet& input_packet = cc->Inputs().Tag("IMAGE_GPU").Value();
|
||||||
const Packet& mask_packet = cc->Inputs().Tag("MASK_GPU").Value();
|
const Packet& mask_packet = cc->Inputs().Tag("MASK_GPU").Value();
|
||||||
@@ -233,13 +239,13 @@ REGISTER_CALCULATOR(RecolorCalculator);
|
|||||||
img_tex.Release();
|
img_tex.Release();
|
||||||
mask_tex.Release();
|
mask_tex.Release();
|
||||||
dst_tex.Release();
|
dst_tex.Release();
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
void RecolorCalculator::GlRender() {
|
void RecolorCalculator::GlRender() {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
static const GLfloat square_vertices[] = {
|
static const GLfloat square_vertices[] = {
|
||||||
-1.0f, -1.0f, // bottom left
|
-1.0f, -1.0f, // bottom left
|
||||||
1.0f, -1.0f, // bottom right
|
1.0f, -1.0f, // bottom right
|
||||||
@@ -287,7 +293,7 @@ void RecolorCalculator::GlRender() {
|
|||||||
glBindVertexArray(0);
|
glBindVertexArray(0);
|
||||||
glDeleteVertexArrays(1, &vao);
|
glDeleteVertexArrays(1, &vao);
|
||||||
glDeleteBuffers(2, vbo);
|
glDeleteBuffers(2, vbo);
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status RecolorCalculator::LoadOptions(CalculatorContext* cc) {
|
::mediapipe::Status RecolorCalculator::LoadOptions(CalculatorContext* cc) {
|
||||||
@@ -305,7 +311,7 @@ void RecolorCalculator::GlRender() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status RecolorCalculator::InitGpu(CalculatorContext* cc) {
|
::mediapipe::Status RecolorCalculator::InitGpu(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
const GLint attr_location[NUM_ATTRIBUTES] = {
|
const GLint attr_location[NUM_ATTRIBUTES] = {
|
||||||
ATTRIB_VERTEX,
|
ATTRIB_VERTEX,
|
||||||
ATTRIB_TEXTURE_POSITION,
|
ATTRIB_TEXTURE_POSITION,
|
||||||
@@ -374,7 +380,7 @@ void RecolorCalculator::GlRender() {
|
|||||||
glUniform1i(glGetUniformLocation(program_, "mask"), 2);
|
glUniform1i(glGetUniformLocation(program_, "mask"), 2);
|
||||||
glUniform3f(glGetUniformLocation(program_, "recolor"), color_[0], color_[1],
|
glUniform3f(glGetUniformLocation(program_, "recolor"), color_[0], color_[1],
|
||||||
color_[2]);
|
color_[2]);
|
||||||
#endif // __ANDROID__ or iOS
|
#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,12 +32,14 @@ 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
|
||||||
// aspect ratio (preserving the center of the frame). This is done
|
// aspect ratio (preserving the center of the frame). This is done
|
||||||
// before scaling.
|
// before scaling. The string must contain "/", so to disable cropping,
|
||||||
|
// set both to "0/1".
|
||||||
// For example, for a min_aspect_ratio of "9/16" and max of "16/9" the
|
// For example, for a min_aspect_ratio of "9/16" and max of "16/9" the
|
||||||
// following cropping will occur:
|
// following cropping will occur:
|
||||||
// 1920x1080 (which is 16:9) is not cropped
|
// 1920x1080 (which is 16:9) is not cropped
|
||||||
@@ -94,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
|
||||||
|
|||||||
@@ -25,12 +25,11 @@
|
|||||||
#include "mediapipe/framework/port/status.h"
|
#include "mediapipe/framework/port/status.h"
|
||||||
#include "mediapipe/framework/port/vector.h"
|
#include "mediapipe/framework/port/vector.h"
|
||||||
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
#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/shader_util.h"
|
#include "mediapipe/gpu/shader_util.h"
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
@@ -107,16 +106,18 @@ class SetAlphaCalculator : public CalculatorBase {
|
|||||||
|
|
||||||
bool use_gpu_ = false;
|
bool use_gpu_ = false;
|
||||||
bool gpu_initialized_ = false;
|
bool gpu_initialized_ = false;
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||||
GLuint program_ = 0;
|
GLuint program_ = 0;
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
};
|
};
|
||||||
REGISTER_CALCULATOR(SetAlphaCalculator);
|
REGISTER_CALCULATOR(SetAlphaCalculator);
|
||||||
|
|
||||||
::mediapipe::Status SetAlphaCalculator::GetContract(CalculatorContract* cc) {
|
::mediapipe::Status SetAlphaCalculator::GetContract(CalculatorContract* cc) {
|
||||||
CHECK_GE(cc->Inputs().NumEntries(), 1);
|
CHECK_GE(cc->Inputs().NumEntries(), 1);
|
||||||
|
|
||||||
|
bool use_gpu = false;
|
||||||
|
|
||||||
if (cc->Inputs().HasTag(kInputFrameTag) &&
|
if (cc->Inputs().HasTag(kInputFrameTag) &&
|
||||||
cc->Inputs().HasTag(kInputFrameTagGpu)) {
|
cc->Inputs().HasTag(kInputFrameTagGpu)) {
|
||||||
return ::mediapipe::InternalError("Cannot have multiple input images.");
|
return ::mediapipe::InternalError("Cannot have multiple input images.");
|
||||||
@@ -127,38 +128,43 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Input image to add/edit alpha channel.
|
// Input image to add/edit alpha channel.
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Inputs().HasTag(kInputFrameTagGpu)) {
|
if (cc->Inputs().HasTag(kInputFrameTagGpu)) {
|
||||||
cc->Inputs().Tag(kInputFrameTagGpu).Set<mediapipe::GpuBuffer>();
|
cc->Inputs().Tag(kInputFrameTagGpu).Set<mediapipe::GpuBuffer>();
|
||||||
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
if (cc->Inputs().HasTag(kInputFrameTag)) {
|
if (cc->Inputs().HasTag(kInputFrameTag)) {
|
||||||
cc->Inputs().Tag(kInputFrameTag).Set<ImageFrame>();
|
cc->Inputs().Tag(kInputFrameTag).Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
|
|
||||||
// Input alpha image mask (optional)
|
// Input alpha image mask (optional)
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Inputs().HasTag(kInputAlphaTagGpu)) {
|
if (cc->Inputs().HasTag(kInputAlphaTagGpu)) {
|
||||||
cc->Inputs().Tag(kInputAlphaTagGpu).Set<mediapipe::GpuBuffer>();
|
cc->Inputs().Tag(kInputAlphaTagGpu).Set<mediapipe::GpuBuffer>();
|
||||||
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
if (cc->Inputs().HasTag(kInputAlphaTag)) {
|
if (cc->Inputs().HasTag(kInputAlphaTag)) {
|
||||||
cc->Inputs().Tag(kInputAlphaTag).Set<ImageFrame>();
|
cc->Inputs().Tag(kInputAlphaTag).Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
|
|
||||||
// RGBA output image.
|
// RGBA output image.
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
if (cc->Outputs().HasTag(kOutputFrameTagGpu)) {
|
if (cc->Outputs().HasTag(kOutputFrameTagGpu)) {
|
||||||
cc->Outputs().Tag(kOutputFrameTagGpu).Set<mediapipe::GpuBuffer>();
|
cc->Outputs().Tag(kOutputFrameTagGpu).Set<mediapipe::GpuBuffer>();
|
||||||
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
if (cc->Outputs().HasTag(kOutputFrameTag)) {
|
if (cc->Outputs().HasTag(kOutputFrameTag)) {
|
||||||
cc->Outputs().Tag(kOutputFrameTag).Set<ImageFrame>();
|
cc->Outputs().Tag(kOutputFrameTag).Set<ImageFrame>();
|
||||||
}
|
}
|
||||||
|
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
if (use_gpu) {
|
||||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
#endif // __ANDROID__ or iOS
|
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -170,11 +176,11 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
|
|||||||
|
|
||||||
if (cc->Inputs().HasTag(kInputFrameTagGpu) &&
|
if (cc->Inputs().HasTag(kInputFrameTagGpu) &&
|
||||||
cc->Outputs().HasTag(kOutputFrameTagGpu)) {
|
cc->Outputs().HasTag(kOutputFrameTagGpu)) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
use_gpu_ = true;
|
use_gpu_ = true;
|
||||||
#else
|
#else
|
||||||
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
|
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
// Get global value from options (-1 if not set).
|
// Get global value from options (-1 if not set).
|
||||||
@@ -187,17 +193,17 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
|
|||||||
RET_CHECK_FAIL() << "Must use either image mask or options alpha value.";
|
RET_CHECK_FAIL() << "Must use either image mask or options alpha value.";
|
||||||
|
|
||||||
if (use_gpu_) {
|
if (use_gpu_) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||||
#endif
|
#endif
|
||||||
}
|
} // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status SetAlphaCalculator::Process(CalculatorContext* cc) {
|
::mediapipe::Status SetAlphaCalculator::Process(CalculatorContext* cc) {
|
||||||
if (use_gpu_) {
|
if (use_gpu_) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
MP_RETURN_IF_ERROR(
|
MP_RETURN_IF_ERROR(
|
||||||
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
|
gpu_helper_.RunInGlContext([this, cc]() -> ::mediapipe::Status {
|
||||||
if (!gpu_initialized_) {
|
if (!gpu_initialized_) {
|
||||||
@@ -207,7 +213,7 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
|
|||||||
MP_RETURN_IF_ERROR(RenderGpu(cc));
|
MP_RETURN_IF_ERROR(RenderGpu(cc));
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}));
|
}));
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
} else {
|
} else {
|
||||||
MP_RETURN_IF_ERROR(RenderCpu(cc));
|
MP_RETURN_IF_ERROR(RenderCpu(cc));
|
||||||
}
|
}
|
||||||
@@ -216,12 +222,12 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status SetAlphaCalculator::Close(CalculatorContext* cc) {
|
::mediapipe::Status SetAlphaCalculator::Close(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
gpu_helper_.RunInGlContext([this] {
|
gpu_helper_.RunInGlContext([this] {
|
||||||
if (program_) glDeleteProgram(program_);
|
if (program_) glDeleteProgram(program_);
|
||||||
program_ = 0;
|
program_ = 0;
|
||||||
});
|
});
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -295,7 +301,7 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
|
|||||||
if (cc->Inputs().Tag(kInputFrameTagGpu).IsEmpty()) {
|
if (cc->Inputs().Tag(kInputFrameTagGpu).IsEmpty()) {
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
// Setup source texture.
|
// Setup source texture.
|
||||||
const auto& input_frame =
|
const auto& input_frame =
|
||||||
cc->Inputs().Tag(kInputFrameTagGpu).Get<mediapipe::GpuBuffer>();
|
cc->Inputs().Tag(kInputFrameTagGpu).Get<mediapipe::GpuBuffer>();
|
||||||
@@ -348,13 +354,13 @@ REGISTER_CALCULATOR(SetAlphaCalculator);
|
|||||||
// Cleanup
|
// Cleanup
|
||||||
input_texture.Release();
|
input_texture.Release();
|
||||||
output_texture.Release();
|
output_texture.Release();
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
void SetAlphaCalculator::GlRender(CalculatorContext* cc) {
|
void SetAlphaCalculator::GlRender(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
static const GLfloat square_vertices[] = {
|
static const GLfloat square_vertices[] = {
|
||||||
-1.0f, -1.0f, // bottom left
|
-1.0f, -1.0f, // bottom left
|
||||||
1.0f, -1.0f, // bottom right
|
1.0f, -1.0f, // bottom right
|
||||||
@@ -403,11 +409,11 @@ void SetAlphaCalculator::GlRender(CalculatorContext* cc) {
|
|||||||
glDeleteVertexArrays(1, &vao);
|
glDeleteVertexArrays(1, &vao);
|
||||||
glDeleteBuffers(2, vbo);
|
glDeleteBuffers(2, vbo);
|
||||||
|
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status SetAlphaCalculator::GlSetup(CalculatorContext* cc) {
|
::mediapipe::Status SetAlphaCalculator::GlSetup(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU)
|
||||||
const GLint attr_location[NUM_ATTRIBUTES] = {
|
const GLint attr_location[NUM_ATTRIBUTES] = {
|
||||||
ATTRIB_VERTEX,
|
ATTRIB_VERTEX,
|
||||||
ATTRIB_TEXTURE_POSITION,
|
ATTRIB_TEXTURE_POSITION,
|
||||||
@@ -460,7 +466,7 @@ void SetAlphaCalculator::GlRender(CalculatorContext* cc) {
|
|||||||
glUniform1i(glGetUniformLocation(program_, "alpha_mask"), 2);
|
glUniform1i(glGetUniformLocation(program_, "alpha_mask"), 2);
|
||||||
glUniform1f(glGetUniformLocation(program_, "alpha_value"), alpha_value_);
|
glUniform1f(glGetUniformLocation(program_, "alpha_value"), alpha_value_);
|
||||||
|
|
||||||
#endif // __ANDROID__ or iOS
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|||||||
+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 |
@@ -13,12 +13,12 @@
|
|||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
#
|
#
|
||||||
|
|
||||||
|
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||||
|
|
||||||
licenses(["notice"]) # Apache 2.0
|
licenses(["notice"]) # Apache 2.0
|
||||||
|
|
||||||
package(default_visibility = ["//visibility:private"])
|
package(default_visibility = ["//visibility:private"])
|
||||||
|
|
||||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
|
||||||
|
|
||||||
proto_library(
|
proto_library(
|
||||||
name = "graph_tensors_packet_generator_proto",
|
name = "graph_tensors_packet_generator_proto",
|
||||||
srcs = ["graph_tensors_packet_generator.proto"],
|
srcs = ["graph_tensors_packet_generator.proto"],
|
||||||
@@ -104,6 +104,17 @@ proto_library(
|
|||||||
deps = ["//mediapipe/framework:calculator_proto"],
|
deps = ["//mediapipe/framework:calculator_proto"],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
proto_library(
|
||||||
|
name = "unpack_media_sequence_calculator_proto",
|
||||||
|
srcs = ["unpack_media_sequence_calculator.proto"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/calculators/core:packet_resampler_calculator_proto",
|
||||||
|
"//mediapipe/framework:calculator_proto",
|
||||||
|
"//mediapipe/util:audio_decoder_proto",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
proto_library(
|
proto_library(
|
||||||
name = "vector_float_to_tensor_calculator_options_proto",
|
name = "vector_float_to_tensor_calculator_options_proto",
|
||||||
srcs = ["vector_float_to_tensor_calculator_options.proto"],
|
srcs = ["vector_float_to_tensor_calculator_options.proto"],
|
||||||
@@ -255,11 +266,23 @@ mediapipe_cc_proto_library(
|
|||||||
cc_deps = [
|
cc_deps = [
|
||||||
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
|
"//mediapipe/calculators/core:packet_resampler_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_cc_proto",
|
"//mediapipe/framework:calculator_cc_proto",
|
||||||
|
"//mediapipe/util:audio_decoder_cc_proto",
|
||||||
],
|
],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [":unpack_media_sequence_calculator_proto"],
|
deps = [":unpack_media_sequence_calculator_proto"],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
mediapipe_cc_proto_library(
|
||||||
|
name = "vector_int_to_tensor_calculator_options_cc_proto",
|
||||||
|
srcs = ["vector_int_to_tensor_calculator_options.proto"],
|
||||||
|
cc_deps = [
|
||||||
|
"//mediapipe/framework:calculator_cc_proto",
|
||||||
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
|
],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [":vector_int_to_tensor_calculator_options_proto"],
|
||||||
|
)
|
||||||
|
|
||||||
mediapipe_cc_proto_library(
|
mediapipe_cc_proto_library(
|
||||||
name = "vector_float_to_tensor_calculator_options_cc_proto",
|
name = "vector_float_to_tensor_calculator_options_cc_proto",
|
||||||
srcs = ["vector_float_to_tensor_calculator_options.proto"],
|
srcs = ["vector_float_to_tensor_calculator_options.proto"],
|
||||||
@@ -273,7 +296,7 @@ cc_library(
|
|||||||
srcs = ["graph_tensors_packet_generator.cc"],
|
srcs = ["graph_tensors_packet_generator.cc"],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/calculators/tensorflow:graph_tensors_packet_generator_cc_proto",
|
":graph_tensors_packet_generator_cc_proto",
|
||||||
"//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",
|
||||||
@@ -288,7 +311,7 @@ cc_library(
|
|||||||
srcs = ["image_frame_to_tensor_calculator.cc"],
|
srcs = ["image_frame_to_tensor_calculator.cc"],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/calculators/tensorflow:image_frame_to_tensor_calculator_cc_proto",
|
":image_frame_to_tensor_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/formats:image_frame",
|
"//mediapipe/framework/formats:image_frame",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
@@ -310,7 +333,7 @@ cc_library(
|
|||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
||||||
"//mediapipe/calculators/tensorflow:matrix_to_tensor_calculator_options_cc_proto",
|
":matrix_to_tensor_calculator_options_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/formats:matrix",
|
"//mediapipe/framework/formats:matrix",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
@@ -331,7 +354,7 @@ cc_library(
|
|||||||
srcs = ["lapped_tensor_buffer_calculator.cc"],
|
srcs = ["lapped_tensor_buffer_calculator.cc"],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/calculators/tensorflow:lapped_tensor_buffer_calculator_cc_proto",
|
":lapped_tensor_buffer_calculator_cc_proto",
|
||||||
"//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",
|
||||||
@@ -413,7 +436,7 @@ cc_library(
|
|||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
":tensorflow_session",
|
":tensorflow_session",
|
||||||
"//mediapipe/calculators/tensorflow:tensorflow_inference_calculator_cc_proto",
|
":tensorflow_inference_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/tool:status_util",
|
"//mediapipe/framework/tool:status_util",
|
||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
@@ -491,7 +514,7 @@ cc_library(
|
|||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
":tensorflow_session",
|
":tensorflow_session",
|
||||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_frozen_graph_generator_cc_proto",
|
":tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/tool:status_util",
|
"//mediapipe/framework/tool:status_util",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
@@ -550,7 +573,7 @@ cc_library(
|
|||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
":tensorflow_session",
|
":tensorflow_session",
|
||||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_saved_model_generator_cc_proto",
|
":tensorflow_session_from_saved_model_generator_cc_proto",
|
||||||
"//mediapipe/framework:packet_generator",
|
"//mediapipe/framework:packet_generator",
|
||||||
"//mediapipe/framework:packet_type",
|
"//mediapipe/framework:packet_type",
|
||||||
"//mediapipe/framework/tool:status_util",
|
"//mediapipe/framework/tool:status_util",
|
||||||
@@ -574,7 +597,7 @@ cc_library(
|
|||||||
srcs = ["tensor_squeeze_dimensions_calculator.cc"],
|
srcs = ["tensor_squeeze_dimensions_calculator.cc"],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/calculators/tensorflow:tensor_squeeze_dimensions_calculator_cc_proto",
|
":tensor_squeeze_dimensions_calculator_cc_proto",
|
||||||
"//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",
|
||||||
@@ -588,7 +611,7 @@ cc_library(
|
|||||||
srcs = ["tensor_to_image_frame_calculator.cc"],
|
srcs = ["tensor_to_image_frame_calculator.cc"],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/calculators/tensorflow:tensor_to_image_frame_calculator_cc_proto",
|
":tensor_to_image_frame_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/formats:image_frame",
|
"//mediapipe/framework/formats:image_frame",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
@@ -604,7 +627,7 @@ cc_library(
|
|||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
"//mediapipe/framework/formats:time_series_header_cc_proto",
|
||||||
"//mediapipe/calculators/tensorflow:tensor_to_matrix_calculator_cc_proto",
|
":tensor_to_matrix_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/formats:matrix",
|
"//mediapipe/framework/formats:matrix",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
@@ -620,6 +643,22 @@ cc_library(
|
|||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "tfrecord_reader_calculator",
|
||||||
|
srcs = ["tfrecord_reader_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/port:integral_types",
|
||||||
|
"//mediapipe/framework/port:logging",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"@org_tensorflow//tensorflow/core:lib",
|
||||||
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
cc_library(
|
cc_library(
|
||||||
name = "tensor_to_vector_float_calculator",
|
name = "tensor_to_vector_float_calculator",
|
||||||
srcs = ["tensor_to_vector_float_calculator.cc"],
|
srcs = ["tensor_to_vector_float_calculator.cc"],
|
||||||
@@ -628,7 +667,7 @@ cc_library(
|
|||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/calculators/tensorflow:tensor_to_vector_float_calculator_options_cc_proto",
|
":tensor_to_vector_float_calculator_options_cc_proto",
|
||||||
] + select({
|
] + select({
|
||||||
"//conditions:default": [
|
"//conditions:default": [
|
||||||
"@org_tensorflow//tensorflow/core:framework",
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
@@ -653,6 +692,7 @@ cc_library(
|
|||||||
"//mediapipe/framework/formats:location",
|
"//mediapipe/framework/formats:location",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
|
"//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_cc",
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
@@ -660,12 +700,26 @@ cc_library(
|
|||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "vector_int_to_tensor_calculator",
|
||||||
|
srcs = ["vector_int_to_tensor_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
":vector_int_to_tensor_calculator_options_cc_proto",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework/port:ret_check",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
cc_library(
|
cc_library(
|
||||||
name = "vector_float_to_tensor_calculator",
|
name = "vector_float_to_tensor_calculator",
|
||||||
srcs = ["vector_float_to_tensor_calculator.cc"],
|
srcs = ["vector_float_to_tensor_calculator.cc"],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
"//mediapipe/calculators/tensorflow:vector_float_to_tensor_calculator_options_cc_proto",
|
":vector_float_to_tensor_calculator_options_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"//mediapipe/framework/port:status",
|
"//mediapipe/framework/port:status",
|
||||||
@@ -674,12 +728,27 @@ cc_library(
|
|||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "unpack_yt8m_sequence_example_calculator",
|
||||||
|
srcs = ["unpack_yt8m_sequence_example_calculator.cc"],
|
||||||
|
visibility = ["//visibility:public"],
|
||||||
|
deps = [
|
||||||
|
":lapped_tensor_buffer_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:packet",
|
||||||
|
"//mediapipe/framework/port:status",
|
||||||
|
"@org_tensorflow//tensorflow/core:protos_all_cc",
|
||||||
|
],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
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",
|
||||||
"//mediapipe/calculators/tensorflow:graph_tensors_packet_generator_cc_proto",
|
":graph_tensors_packet_generator_cc_proto",
|
||||||
"//mediapipe/framework:packet",
|
"//mediapipe/framework:packet",
|
||||||
"//mediapipe/framework:packet_generator_cc_proto",
|
"//mediapipe/framework:packet_generator_cc_proto",
|
||||||
"//mediapipe/framework:packet_set",
|
"//mediapipe/framework:packet_set",
|
||||||
@@ -693,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",
|
||||||
@@ -709,9 +779,10 @@ 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",
|
||||||
"//mediapipe/calculators/tensorflow:matrix_to_tensor_calculator_options_cc_proto",
|
":matrix_to_tensor_calculator_options_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/formats:matrix",
|
"//mediapipe/framework/formats:matrix",
|
||||||
@@ -725,9 +796,10 @@ 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",
|
||||||
"//mediapipe/calculators/tensorflow:lapped_tensor_buffer_calculator_cc_proto",
|
":lapped_tensor_buffer_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
@@ -769,7 +841,6 @@ cc_test(
|
|||||||
"//mediapipe/framework/formats:location",
|
"//mediapipe/framework/formats:location",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"//mediapipe/framework/port:opencv_imgcodecs",
|
"//mediapipe/framework/port:opencv_imgcodecs",
|
||||||
"//mediapipe/framework/port:status",
|
|
||||||
"//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",
|
||||||
@@ -816,7 +887,7 @@ cc_test(
|
|||||||
":tensorflow_inference_calculator",
|
":tensorflow_inference_calculator",
|
||||||
":tensorflow_session",
|
":tensorflow_session",
|
||||||
":tensorflow_session_from_frozen_graph_generator",
|
":tensorflow_session_from_frozen_graph_generator",
|
||||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_frozen_graph_generator_cc_proto",
|
":tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:packet",
|
"//mediapipe/framework:packet",
|
||||||
"//mediapipe/framework:packet_generator_cc_proto",
|
"//mediapipe/framework:packet_generator_cc_proto",
|
||||||
@@ -846,7 +917,7 @@ cc_test(
|
|||||||
":tensorflow_inference_calculator",
|
":tensorflow_inference_calculator",
|
||||||
":tensorflow_session",
|
":tensorflow_session",
|
||||||
":tensorflow_session_from_saved_model_generator",
|
":tensorflow_session_from_saved_model_generator",
|
||||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_saved_model_generator_cc_proto",
|
":tensorflow_session_from_saved_model_generator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:packet",
|
"//mediapipe/framework:packet",
|
||||||
"//mediapipe/framework:packet_generator_cc_proto",
|
"//mediapipe/framework:packet_generator_cc_proto",
|
||||||
@@ -856,14 +927,8 @@ cc_test(
|
|||||||
"//mediapipe/framework/tool:tag_map_helper",
|
"//mediapipe/framework/tool:tag_map_helper",
|
||||||
"//mediapipe/framework/tool:validate_type",
|
"//mediapipe/framework/tool:validate_type",
|
||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
|
"@org_tensorflow//tensorflow/core:all_kernels",
|
||||||
"@org_tensorflow//tensorflow/core:direct_session",
|
"@org_tensorflow//tensorflow/core:direct_session",
|
||||||
"@org_tensorflow//tensorflow/core/kernels:array",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels:bitcast_op",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels:io",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels:state",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels:string",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels/data:tensor_dataset_op",
|
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -887,23 +952,18 @@ cc_test(
|
|||||||
"//mediapipe/framework/tool:tag_map_helper",
|
"//mediapipe/framework/tool:tag_map_helper",
|
||||||
"//mediapipe/framework/tool:validate_type",
|
"//mediapipe/framework/tool:validate_type",
|
||||||
"@com_google_absl//absl/strings",
|
"@com_google_absl//absl/strings",
|
||||||
|
"@org_tensorflow//tensorflow/core:all_kernels",
|
||||||
"@org_tensorflow//tensorflow/core:direct_session",
|
"@org_tensorflow//tensorflow/core:direct_session",
|
||||||
"@org_tensorflow//tensorflow/core/kernels:array",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels:bitcast_op",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels:conv_ops",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels:io",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels:state",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels:string",
|
|
||||||
"@org_tensorflow//tensorflow/core/kernels/data:tensor_dataset_op",
|
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
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",
|
||||||
"//mediapipe/calculators/tensorflow:tensor_squeeze_dimensions_calculator_cc_proto",
|
":tensor_squeeze_dimensions_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
@@ -916,9 +976,10 @@ 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",
|
||||||
"//mediapipe/calculators/tensorflow:tensor_to_image_frame_calculator_cc_proto",
|
":tensor_to_image_frame_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/formats:image_frame",
|
"//mediapipe/framework/formats:image_frame",
|
||||||
@@ -932,9 +993,10 @@ 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",
|
||||||
"//mediapipe/calculators/tensorflow:tensor_to_matrix_calculator_cc_proto",
|
":tensor_to_matrix_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/formats:matrix",
|
"//mediapipe/framework/formats:matrix",
|
||||||
@@ -948,9 +1010,10 @@ cc_test(
|
|||||||
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",
|
||||||
"//mediapipe/calculators/tensorflow:tensor_to_vector_float_calculator_options_cc_proto",
|
":tensor_to_vector_float_calculator_options_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
@@ -971,6 +1034,7 @@ cc_test(
|
|||||||
"//mediapipe/framework/formats:location",
|
"//mediapipe/framework/formats:location",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
"//mediapipe/framework/port:rectangle",
|
"//mediapipe/framework/port:rectangle",
|
||||||
|
"//mediapipe/util:audio_decoder_cc_proto",
|
||||||
"//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",
|
||||||
@@ -978,12 +1042,28 @@ cc_test(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_test(
|
||||||
|
name = "vector_int_to_tensor_calculator_test",
|
||||||
|
srcs = ["vector_int_to_tensor_calculator_test.cc"],
|
||||||
|
linkstatic = 1,
|
||||||
|
deps = [
|
||||||
|
":vector_int_to_tensor_calculator",
|
||||||
|
":vector_int_to_tensor_calculator_options_cc_proto",
|
||||||
|
"//mediapipe/framework:calculator_framework",
|
||||||
|
"//mediapipe/framework:calculator_runner",
|
||||||
|
"//mediapipe/framework/port:gtest_main",
|
||||||
|
"@org_tensorflow//tensorflow/core:framework",
|
||||||
|
"@org_tensorflow//tensorflow/core:protos_all_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",
|
||||||
"//mediapipe/calculators/tensorflow:vector_float_to_tensor_calculator_options_cc_proto",
|
":vector_float_to_tensor_calculator_options_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
"//mediapipe/framework/port:gtest_main",
|
"//mediapipe/framework/port:gtest_main",
|
||||||
@@ -1012,7 +1092,7 @@ cc_test(
|
|||||||
":tensorflow_session",
|
":tensorflow_session",
|
||||||
":tensorflow_inference_calculator",
|
":tensorflow_inference_calculator",
|
||||||
":tensorflow_session_from_frozen_graph_generator",
|
":tensorflow_session_from_frozen_graph_generator",
|
||||||
"//mediapipe/calculators/tensorflow:tensorflow_session_from_frozen_graph_generator_cc_proto",
|
":tensorflow_session_from_frozen_graph_generator_cc_proto",
|
||||||
"//mediapipe/framework/deps:file_path",
|
"//mediapipe/framework/deps:file_path",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework:calculator_runner",
|
"//mediapipe/framework:calculator_runner",
|
||||||
|
|||||||
@@ -29,6 +29,11 @@
|
|||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
|
const char kBufferSize[] = "BUFFER_SIZE";
|
||||||
|
const char kOverlap[] = "OVERLAP";
|
||||||
|
const char kTimestampOffset[] = "TIMESTAMP_OFFSET";
|
||||||
|
const char kCalculatorOptions[] = "CALCULATOR_OPTIONS";
|
||||||
|
|
||||||
namespace tf = tensorflow;
|
namespace tf = tensorflow;
|
||||||
|
|
||||||
// Given an input stream of tensors, concatenates the tensors over timesteps.
|
// Given an input stream of tensors, concatenates the tensors over timesteps.
|
||||||
@@ -72,6 +77,9 @@ class LappedTensorBufferCalculator : public CalculatorBase {
|
|||||||
::mediapipe::Status AddBatchDimension(tf::Tensor* input_tensor);
|
::mediapipe::Status AddBatchDimension(tf::Tensor* input_tensor);
|
||||||
|
|
||||||
int steps_until_output_;
|
int steps_until_output_;
|
||||||
|
int buffer_size_;
|
||||||
|
int overlap_;
|
||||||
|
int timestamp_offset_;
|
||||||
std::unique_ptr<CircularBuffer<Timestamp>> timestamp_buffer_;
|
std::unique_ptr<CircularBuffer<Timestamp>> timestamp_buffer_;
|
||||||
std::unique_ptr<CircularBuffer<tf::Tensor>> buffer_;
|
std::unique_ptr<CircularBuffer<tf::Tensor>> buffer_;
|
||||||
LappedTensorBufferCalculatorOptions options_;
|
LappedTensorBufferCalculatorOptions options_;
|
||||||
@@ -87,6 +95,21 @@ REGISTER_CALCULATOR(LappedTensorBufferCalculator);
|
|||||||
);
|
);
|
||||||
RET_CHECK_EQ(cc->Inputs().NumEntries(), 1)
|
RET_CHECK_EQ(cc->Inputs().NumEntries(), 1)
|
||||||
<< "Only one output stream is supported.";
|
<< "Only one output stream is supported.";
|
||||||
|
|
||||||
|
if (cc->InputSidePackets().HasTag(kBufferSize)) {
|
||||||
|
cc->InputSidePackets().Tag(kBufferSize).Set<int>();
|
||||||
|
}
|
||||||
|
if (cc->InputSidePackets().HasTag(kOverlap)) {
|
||||||
|
cc->InputSidePackets().Tag(kOverlap).Set<int>();
|
||||||
|
}
|
||||||
|
if (cc->InputSidePackets().HasTag(kTimestampOffset)) {
|
||||||
|
cc->InputSidePackets().Tag(kTimestampOffset).Set<int>();
|
||||||
|
}
|
||||||
|
if (cc->InputSidePackets().HasTag(kCalculatorOptions)) {
|
||||||
|
cc->InputSidePackets()
|
||||||
|
.Tag(kCalculatorOptions)
|
||||||
|
.Set<LappedTensorBufferCalculatorOptions>();
|
||||||
|
}
|
||||||
cc->Outputs().Index(0).Set<tf::Tensor>(
|
cc->Outputs().Index(0).Set<tf::Tensor>(
|
||||||
// Output tensorflow::Tensor stream with possibly overlapping steps.
|
// Output tensorflow::Tensor stream with possibly overlapping steps.
|
||||||
);
|
);
|
||||||
@@ -95,16 +118,33 @@ REGISTER_CALCULATOR(LappedTensorBufferCalculator);
|
|||||||
|
|
||||||
::mediapipe::Status LappedTensorBufferCalculator::Open(CalculatorContext* cc) {
|
::mediapipe::Status LappedTensorBufferCalculator::Open(CalculatorContext* cc) {
|
||||||
options_ = cc->Options<LappedTensorBufferCalculatorOptions>();
|
options_ = cc->Options<LappedTensorBufferCalculatorOptions>();
|
||||||
RET_CHECK_LT(options_.overlap(), options_.buffer_size());
|
if (cc->InputSidePackets().HasTag(kCalculatorOptions)) {
|
||||||
RET_CHECK_GE(options_.timestamp_offset(), 0)
|
options_ = cc->InputSidePackets()
|
||||||
|
.Tag(kCalculatorOptions)
|
||||||
|
.Get<LappedTensorBufferCalculatorOptions>();
|
||||||
|
}
|
||||||
|
buffer_size_ = options_.buffer_size();
|
||||||
|
if (cc->InputSidePackets().HasTag(kBufferSize)) {
|
||||||
|
buffer_size_ = cc->InputSidePackets().Tag(kBufferSize).Get<int>();
|
||||||
|
}
|
||||||
|
overlap_ = options_.overlap();
|
||||||
|
if (cc->InputSidePackets().HasTag(kOverlap)) {
|
||||||
|
overlap_ = cc->InputSidePackets().Tag(kOverlap).Get<int>();
|
||||||
|
}
|
||||||
|
timestamp_offset_ = options_.timestamp_offset();
|
||||||
|
if (cc->InputSidePackets().HasTag(kTimestampOffset)) {
|
||||||
|
timestamp_offset_ = cc->InputSidePackets().Tag(kTimestampOffset).Get<int>();
|
||||||
|
}
|
||||||
|
|
||||||
|
RET_CHECK_LT(overlap_, buffer_size_);
|
||||||
|
RET_CHECK_GE(timestamp_offset_, 0)
|
||||||
<< "Negative timestamp_offset is not allowed.";
|
<< "Negative timestamp_offset is not allowed.";
|
||||||
RET_CHECK_LT(options_.timestamp_offset(), options_.buffer_size())
|
RET_CHECK_LT(timestamp_offset_, buffer_size_)
|
||||||
<< "output_frame_num_offset has to be less than buffer_size.";
|
<< "output_frame_num_offset has to be less than buffer_size.";
|
||||||
timestamp_buffer_ =
|
timestamp_buffer_ =
|
||||||
absl::make_unique<CircularBuffer<Timestamp>>(options_.buffer_size());
|
absl::make_unique<CircularBuffer<Timestamp>>(buffer_size_);
|
||||||
buffer_ =
|
buffer_ = absl::make_unique<CircularBuffer<tf::Tensor>>(buffer_size_);
|
||||||
absl::make_unique<CircularBuffer<tf::Tensor>>(options_.buffer_size());
|
steps_until_output_ = buffer_size_;
|
||||||
steps_until_output_ = options_.buffer_size();
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -128,11 +168,10 @@ REGISTER_CALCULATOR(LappedTensorBufferCalculator);
|
|||||||
concatenated.get());
|
concatenated.get());
|
||||||
RET_CHECK(concat_status.ok()) << concat_status.ToString();
|
RET_CHECK(concat_status.ok()) << concat_status.ToString();
|
||||||
|
|
||||||
cc->Outputs().Index(0).Add(
|
cc->Outputs().Index(0).Add(concatenated.release(),
|
||||||
concatenated.release(),
|
timestamp_buffer_->Get(timestamp_offset_));
|
||||||
timestamp_buffer_->Get(options_.timestamp_offset()));
|
|
||||||
|
|
||||||
steps_until_output_ = options_.buffer_size() - options_.overlap();
|
steps_until_output_ = buffer_size_ - overlap_;
|
||||||
}
|
}
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -285,6 +289,10 @@ class PackMediaSequenceCalculator : public CalculatorBase {
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
|
int image_height = -1;
|
||||||
|
int image_width = -1;
|
||||||
|
// Because the tag order may vary, we need to loop through tags to get
|
||||||
|
// image information before processing other tag types.
|
||||||
for (const auto& tag : cc->Inputs().GetTags()) {
|
for (const auto& tag : cc->Inputs().GetTags()) {
|
||||||
if (!cc->Inputs().Tag(tag).IsEmpty()) {
|
if (!cc->Inputs().Tag(tag).IsEmpty()) {
|
||||||
features_present_[tag] = true;
|
features_present_[tag] = true;
|
||||||
@@ -306,14 +314,21 @@ class PackMediaSequenceCalculator : public CalculatorBase {
|
|||||||
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||||
<< "No encoded image";
|
<< "No encoded image";
|
||||||
}
|
}
|
||||||
|
image_height = image.height();
|
||||||
|
image_width = image.width();
|
||||||
mpms::AddImageTimestamp(key, cc->InputTimestamp().Value(),
|
mpms::AddImageTimestamp(key, cc->InputTimestamp().Value(),
|
||||||
sequence_.get());
|
sequence_.get());
|
||||||
mpms::AddImageEncoded(key, image.encoded_image(), sequence_.get());
|
mpms::AddImageEncoded(key, image.encoded_image(), sequence_.get());
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
for (const auto& tag : cc->Inputs().GetTags()) {
|
||||||
|
if (!cc->Inputs().Tag(tag).IsEmpty()) {
|
||||||
|
features_present_[tag] = true;
|
||||||
|
}
|
||||||
if (absl::StartsWith(tag, kKeypointsTag) &&
|
if (absl::StartsWith(tag, kKeypointsTag) &&
|
||||||
!cc->Inputs().Tag(tag).IsEmpty()) {
|
!cc->Inputs().Tag(tag).IsEmpty()) {
|
||||||
std::string key = "";
|
std::string key = "";
|
||||||
if (tag != kImageTag) {
|
if (tag != kKeypointsTag) {
|
||||||
int tag_length = sizeof(kKeypointsTag) / sizeof(*kKeypointsTag) - 1;
|
int tag_length = sizeof(kKeypointsTag) / sizeof(*kKeypointsTag) - 1;
|
||||||
if (tag[tag_length] == '_') {
|
if (tag[tag_length] == '_') {
|
||||||
key = tag.substr(tag_length + 1);
|
key = tag.substr(tag_length + 1);
|
||||||
@@ -333,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) /
|
||||||
@@ -363,11 +389,20 @@ class PackMediaSequenceCalculator : public CalculatorBase {
|
|||||||
LocationData::BOUNDING_BOX ||
|
LocationData::BOUNDING_BOX ||
|
||||||
detection.location_data().format() ==
|
detection.location_data().format() ==
|
||||||
LocationData::RELATIVE_BOUNDING_BOX) {
|
LocationData::RELATIVE_BOUNDING_BOX) {
|
||||||
int height = mpms::GetImageHeight(*sequence_);
|
if (mpms::HasImageHeight(*sequence_) &&
|
||||||
int width = mpms::GetImageWidth(*sequence_);
|
mpms::HasImageWidth(*sequence_)) {
|
||||||
|
image_height = mpms::GetImageHeight(*sequence_);
|
||||||
|
image_width = mpms::GetImageWidth(*sequence_);
|
||||||
|
}
|
||||||
|
if (image_height == -1 || image_width == -1) {
|
||||||
|
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||||
|
<< "Images must be provided with bounding boxes or the "
|
||||||
|
"image "
|
||||||
|
<< "height and width must already be in the example.";
|
||||||
|
}
|
||||||
Location relative_bbox = Location::CreateRelativeBBoxLocation(
|
Location relative_bbox = Location::CreateRelativeBBoxLocation(
|
||||||
Location(detection.location_data())
|
Location(detection.location_data())
|
||||||
.ConvertToRelativeBBox(width, height));
|
.ConvertToRelativeBBox(image_width, image_height));
|
||||||
predicted_locations.push_back(relative_bbox);
|
predicted_locations.push_back(relative_bbox);
|
||||||
if (detection.label_size() > 0) {
|
if (detection.label_size() > 0) {
|
||||||
predicted_class_strings.push_back(detection.label(0));
|
predicted_class_strings.push_back(detection.label(0));
|
||||||
|
|||||||
@@ -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");
|
||||||
@@ -357,6 +389,148 @@ TEST_F(PackMediaSequenceCalculatorTest, PacksTwoBBoxDetections) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
TEST_F(PackMediaSequenceCalculatorTest, PacksBBoxWithoutImageDims) {
|
||||||
|
SetUpCalculator({"BBOX_PREDICTED:detections"}, {}, false, true);
|
||||||
|
auto input_sequence = ::absl::make_unique<tf::SequenceExample>();
|
||||||
|
std::string test_video_id = "test_video_id";
|
||||||
|
mpms::SetClipMediaId(test_video_id, input_sequence.get());
|
||||||
|
int height = 480;
|
||||||
|
int width = 640;
|
||||||
|
int num_vectors = 2;
|
||||||
|
for (int i = 0; i < num_vectors; ++i) {
|
||||||
|
auto detections = ::absl::make_unique<::std::vector<Detection>>();
|
||||||
|
Detection detection;
|
||||||
|
detection.add_label("absolute bbox");
|
||||||
|
detection.add_label_id(0);
|
||||||
|
detection.add_score(0.5);
|
||||||
|
Location::CreateBBoxLocation(0, height / 2, width / 2, height / 2)
|
||||||
|
.ConvertToProto(detection.mutable_location_data());
|
||||||
|
detections->push_back(detection);
|
||||||
|
|
||||||
|
detection = Detection();
|
||||||
|
detection.add_label("relative bbox");
|
||||||
|
detection.add_label_id(1);
|
||||||
|
detection.add_score(0.75);
|
||||||
|
Location::CreateRelativeBBoxLocation(0, 0.5, 0.5, 0.5)
|
||||||
|
.ConvertToProto(detection.mutable_location_data());
|
||||||
|
detections->push_back(detection);
|
||||||
|
|
||||||
|
// The mask detection should be ignored in the output.
|
||||||
|
detection = Detection();
|
||||||
|
detection.add_label("mask");
|
||||||
|
detection.add_score(1.0);
|
||||||
|
cv::Mat image(2, 3, CV_8UC1, cv::Scalar(0));
|
||||||
|
Location::CreateCvMaskLocation<uint8>(image).ConvertToProto(
|
||||||
|
detection.mutable_location_data());
|
||||||
|
detections->push_back(detection);
|
||||||
|
|
||||||
|
runner_->MutableInputs()
|
||||||
|
->Tag("BBOX_PREDICTED")
|
||||||
|
.packets.push_back(Adopt(detections.release()).At(Timestamp(i)));
|
||||||
|
}
|
||||||
|
|
||||||
|
runner_->MutableSidePackets()->Tag("SEQUENCE_EXAMPLE") =
|
||||||
|
Adopt(input_sequence.release());
|
||||||
|
|
||||||
|
auto status = runner_->Run();
|
||||||
|
EXPECT_EQ(::mediapipe::StatusCode::kInvalidArgument, status.code());
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(PackMediaSequenceCalculatorTest, PacksBBoxWithImages) {
|
||||||
|
SetUpCalculator({"BBOX_PREDICTED:detections", "IMAGE:images"}, {}, false,
|
||||||
|
true);
|
||||||
|
auto input_sequence = ::absl::make_unique<tf::SequenceExample>();
|
||||||
|
std::string test_video_id = "test_video_id";
|
||||||
|
mpms::SetClipMediaId(test_video_id, input_sequence.get());
|
||||||
|
int height = 480;
|
||||||
|
int width = 640;
|
||||||
|
int num_vectors = 2;
|
||||||
|
for (int i = 0; i < num_vectors; ++i) {
|
||||||
|
auto detections = ::absl::make_unique<::std::vector<Detection>>();
|
||||||
|
Detection detection;
|
||||||
|
detection.add_label("absolute bbox");
|
||||||
|
detection.add_label_id(0);
|
||||||
|
detection.add_score(0.5);
|
||||||
|
Location::CreateBBoxLocation(0, height / 2, width / 2, height / 2)
|
||||||
|
.ConvertToProto(detection.mutable_location_data());
|
||||||
|
detections->push_back(detection);
|
||||||
|
|
||||||
|
detection = Detection();
|
||||||
|
detection.add_label("relative bbox");
|
||||||
|
detection.add_label_id(1);
|
||||||
|
detection.add_score(0.75);
|
||||||
|
Location::CreateRelativeBBoxLocation(0, 0.5, 0.5, 0.5)
|
||||||
|
.ConvertToProto(detection.mutable_location_data());
|
||||||
|
detections->push_back(detection);
|
||||||
|
|
||||||
|
// The mask detection should be ignored in the output.
|
||||||
|
detection = Detection();
|
||||||
|
detection.add_label("mask");
|
||||||
|
detection.add_score(1.0);
|
||||||
|
cv::Mat image(2, 3, CV_8UC1, cv::Scalar(0));
|
||||||
|
Location::CreateCvMaskLocation<uint8>(image).ConvertToProto(
|
||||||
|
detection.mutable_location_data());
|
||||||
|
detections->push_back(detection);
|
||||||
|
|
||||||
|
runner_->MutableInputs()
|
||||||
|
->Tag("BBOX_PREDICTED")
|
||||||
|
.packets.push_back(Adopt(detections.release()).At(Timestamp(i)));
|
||||||
|
}
|
||||||
|
cv::Mat image(height, width, CV_8UC3, cv::Scalar(0, 0, 255));
|
||||||
|
std::vector<uchar> bytes;
|
||||||
|
ASSERT_TRUE(cv::imencode(".jpg", image, bytes, {80}));
|
||||||
|
std::string test_image_string(bytes.begin(), bytes.end());
|
||||||
|
OpenCvImageEncoderCalculatorResults encoded_image;
|
||||||
|
encoded_image.set_encoded_image(test_image_string);
|
||||||
|
encoded_image.set_width(width);
|
||||||
|
encoded_image.set_height(height);
|
||||||
|
|
||||||
|
int num_images = 2;
|
||||||
|
for (int i = 0; i < num_images; ++i) {
|
||||||
|
auto image_ptr =
|
||||||
|
::absl::make_unique<OpenCvImageEncoderCalculatorResults>(encoded_image);
|
||||||
|
runner_->MutableInputs()->Tag("IMAGE").packets.push_back(
|
||||||
|
Adopt(image_ptr.release()).At(Timestamp(i)));
|
||||||
|
}
|
||||||
|
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_EQ(test_video_id, mpms::GetClipMediaId(output_sequence));
|
||||||
|
ASSERT_EQ(height, mpms::GetImageHeight(output_sequence));
|
||||||
|
ASSERT_EQ(width, mpms::GetImageWidth(output_sequence));
|
||||||
|
ASSERT_EQ(num_vectors, mpms::GetPredictedBBoxSize(output_sequence));
|
||||||
|
ASSERT_EQ(num_vectors, mpms::GetPredictedBBoxTimestampSize(output_sequence));
|
||||||
|
ASSERT_EQ(0, mpms::GetClassSegmentationEncodedSize(output_sequence));
|
||||||
|
ASSERT_EQ(0, mpms::GetClassSegmentationTimestampSize(output_sequence));
|
||||||
|
for (int i = 0; i < num_vectors; ++i) {
|
||||||
|
ASSERT_EQ(i, mpms::GetPredictedBBoxTimestampAt(output_sequence, i));
|
||||||
|
auto bboxes = mpms::GetPredictedBBoxAt(output_sequence, i);
|
||||||
|
ASSERT_EQ(2, bboxes.size());
|
||||||
|
for (int j = 0; j < bboxes.size(); ++j) {
|
||||||
|
auto rect = bboxes[j].GetRelativeBBox();
|
||||||
|
ASSERT_NEAR(0, rect.xmin(), 0.001);
|
||||||
|
ASSERT_NEAR(0.5, rect.ymin(), 0.001);
|
||||||
|
ASSERT_NEAR(0.5, rect.xmax(), 0.001);
|
||||||
|
ASSERT_NEAR(1.0, rect.ymax(), 0.001);
|
||||||
|
}
|
||||||
|
auto class_strings =
|
||||||
|
mpms::GetPredictedBBoxLabelStringAt(output_sequence, i);
|
||||||
|
ASSERT_EQ("absolute bbox", class_strings[0]);
|
||||||
|
ASSERT_EQ("relative bbox", class_strings[1]);
|
||||||
|
auto class_indices = mpms::GetPredictedBBoxLabelIndexAt(output_sequence, i);
|
||||||
|
ASSERT_EQ(0, class_indices[0]);
|
||||||
|
ASSERT_EQ(1, class_indices[1]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
TEST_F(PackMediaSequenceCalculatorTest, PacksTwoKeypoints) {
|
TEST_F(PackMediaSequenceCalculatorTest, PacksTwoKeypoints) {
|
||||||
SetUpCalculator({"KEYPOINTS_TEST:keypoints"}, {}, false, true);
|
SetUpCalculator({"KEYPOINTS_TEST:keypoints"}, {}, false, true);
|
||||||
auto input_sequence = ::absl::make_unique<tf::SequenceExample>();
|
auto input_sequence = ::absl::make_unique<tf::SequenceExample>();
|
||||||
|
|||||||
@@ -34,7 +34,7 @@
|
|||||||
#include "tensorflow/core/framework/tensor_shape.h"
|
#include "tensorflow/core/framework/tensor_shape.h"
|
||||||
#include "tensorflow/core/framework/tensor_util.h"
|
#include "tensorflow/core/framework/tensor_util.h"
|
||||||
|
|
||||||
#if !defined(__ANDROID__) && !defined(__APPLE__)
|
#if !defined(MEDIAPIPE_MOBILE) && !defined(__APPLE__)
|
||||||
#include "tensorflow/core/profiler/lib/traceme.h"
|
#include "tensorflow/core/profiler/lib/traceme.h"
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
@@ -441,7 +441,7 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
|
|||||||
const int64 run_start_time = absl::ToUnixMicros(clock_->TimeNow());
|
const int64 run_start_time = absl::ToUnixMicros(clock_->TimeNow());
|
||||||
tf::Status tf_status;
|
tf::Status tf_status;
|
||||||
{
|
{
|
||||||
#if !defined(__ANDROID__) && !defined(__APPLE__)
|
#if !defined(MEDIAPIPE_MOBILE) && !defined(__APPLE__)
|
||||||
tensorflow::profiler::TraceMe trace(absl::string_view(cc->NodeName()));
|
tensorflow::profiler::TraceMe trace(absl::string_view(cc->NodeName()));
|
||||||
#endif
|
#endif
|
||||||
tf_status = session_->Run(input_tensors, output_tensor_names,
|
tf_status = session_->Run(input_tensors, output_tensor_names,
|
||||||
@@ -454,7 +454,7 @@ class TensorFlowInferenceCalculator : public CalculatorBase {
|
|||||||
|
|
||||||
// RET_CHECK on the tf::Status object itself in order to print an
|
// 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)
|
||||||
|
|||||||
@@ -31,8 +31,7 @@
|
|||||||
#include "mediapipe/framework/tool/status_util.h"
|
#include "mediapipe/framework/tool/status_util.h"
|
||||||
#include "tensorflow/core/public/session_options.h"
|
#include "tensorflow/core/public/session_options.h"
|
||||||
|
|
||||||
#if defined(MEDIAPIPE_LITE) || defined(__ANDROID__) || \
|
#if defined(MEDIAPIPE_MOBILE)
|
||||||
defined(__APPLE__) && !TARGET_OS_OSX
|
|
||||||
#include "mediapipe/util/android/file/base/helpers.h"
|
#include "mediapipe/util/android/file/base/helpers.h"
|
||||||
#else
|
#else
|
||||||
#include "mediapipe/framework/port/file_helpers.h"
|
#include "mediapipe/framework/port/file_helpers.h"
|
||||||
@@ -110,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;
|
||||||
@@ -120,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"
|
||||||
@@ -63,7 +63,7 @@ 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, {{"/", "_"}});
|
||||||
return output;
|
return output;
|
||||||
} else {
|
} else {
|
||||||
return name;
|
return name;
|
||||||
@@ -140,7 +140,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>();
|
||||||
|
|||||||
+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"
|
||||||
@@ -65,7 +65,7 @@ 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, {{"/", "_"}});
|
||||||
return output;
|
return output;
|
||||||
} else {
|
} else {
|
||||||
return name;
|
return name;
|
||||||
@@ -135,7 +135,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>();
|
||||||
|
|||||||
+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"
|
||||||
|
|||||||
@@ -0,0 +1,126 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
#include <utility>
|
||||||
|
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/integral_types.h"
|
||||||
|
#include "mediapipe/framework/port/logging.h"
|
||||||
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
#include "tensorflow/core/example/example.pb.h"
|
||||||
|
#include "tensorflow/core/lib/core/status.h"
|
||||||
|
#include "tensorflow/core/lib/io/record_reader.h"
|
||||||
|
#include "tensorflow/core/platform/env.h"
|
||||||
|
#include "tensorflow/core/platform/file_system.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
const char kTFRecordPath[] = "TFRECORD_PATH";
|
||||||
|
const char kRecordIndex[] = "RECORD_INDEX";
|
||||||
|
const char kExampleTag[] = "EXAMPLE";
|
||||||
|
const char kSequenceExampleTag[] = "SEQUENCE_EXAMPLE";
|
||||||
|
|
||||||
|
// Reads a tensorflow example/sequence example from a tfrecord file.
|
||||||
|
// If the "RECORD_INDEX" input side packet is provided, the calculator is going
|
||||||
|
// to fetch the example/sequence example of the tfrecord file at the target
|
||||||
|
// record index. Otherwise, the reader always reads the first example/sequence
|
||||||
|
// example of the tfrecord file.
|
||||||
|
//
|
||||||
|
// Example config:
|
||||||
|
// node {
|
||||||
|
// calculator: "TFRecordReaderCalculator"
|
||||||
|
// input_side_packet: "TFRECORD_PATH:tfrecord_path"
|
||||||
|
// input_side_packet: "RECORD_INDEX:record_index"
|
||||||
|
// output_side_packet: "SEQUENCE_EXAMPLE:sequence_example"
|
||||||
|
// }
|
||||||
|
class TFRecordReaderCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||||
|
};
|
||||||
|
|
||||||
|
::mediapipe::Status TFRecordReaderCalculator::GetContract(
|
||||||
|
CalculatorContract* cc) {
|
||||||
|
cc->InputSidePackets().Tag(kTFRecordPath).Set<std::string>();
|
||||||
|
if (cc->InputSidePackets().HasTag(kRecordIndex)) {
|
||||||
|
cc->InputSidePackets().Tag(kRecordIndex).Set<int>();
|
||||||
|
}
|
||||||
|
|
||||||
|
RET_CHECK(cc->OutputSidePackets().HasTag(kExampleTag) ||
|
||||||
|
cc->OutputSidePackets().HasTag(kSequenceExampleTag))
|
||||||
|
<< "TFRecordReaderCalculator must output either Tensorflow example or "
|
||||||
|
"sequence example.";
|
||||||
|
if (cc->OutputSidePackets().HasTag(kExampleTag)) {
|
||||||
|
cc->OutputSidePackets().Tag(kExampleTag).Set<tensorflow::Example>();
|
||||||
|
} else {
|
||||||
|
cc->OutputSidePackets()
|
||||||
|
.Tag(kSequenceExampleTag)
|
||||||
|
.Set<tensorflow::SequenceExample>();
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status TFRecordReaderCalculator::Open(CalculatorContext* cc) {
|
||||||
|
std::unique_ptr<tensorflow::RandomAccessFile> file;
|
||||||
|
auto tf_status = tensorflow::Env::Default()->NewRandomAccessFile(
|
||||||
|
cc->InputSidePackets().Tag(kTFRecordPath).Get<std::string>(), &file);
|
||||||
|
RET_CHECK(tf_status.ok())
|
||||||
|
<< "Failed to open tfrecord file: " << tf_status.ToString();
|
||||||
|
tensorflow::io::RecordReader reader(file.get(),
|
||||||
|
tensorflow::io::RecordReaderOptions());
|
||||||
|
tensorflow::uint64 offset = 0;
|
||||||
|
tensorflow::tstring example_str;
|
||||||
|
const int target_idx =
|
||||||
|
cc->InputSidePackets().HasTag(kRecordIndex)
|
||||||
|
? cc->InputSidePackets().Tag(kRecordIndex).Get<int>()
|
||||||
|
: 0;
|
||||||
|
int current_idx = 0;
|
||||||
|
while (current_idx <= target_idx) {
|
||||||
|
tf_status = reader.ReadRecord(&offset, &example_str);
|
||||||
|
RET_CHECK(tf_status.ok())
|
||||||
|
<< "Failed to read tfrecord: " << tf_status.ToString();
|
||||||
|
if (current_idx == target_idx) {
|
||||||
|
if (cc->OutputSidePackets().HasTag(kExampleTag)) {
|
||||||
|
tensorflow::Example tf_example;
|
||||||
|
tf_example.ParseFromArray(example_str.data(), example_str.size());
|
||||||
|
cc->OutputSidePackets()
|
||||||
|
.Tag(kExampleTag)
|
||||||
|
.Set(MakePacket<tensorflow::Example>(std::move(tf_example)));
|
||||||
|
} else {
|
||||||
|
tensorflow::SequenceExample tf_sequence_example;
|
||||||
|
tf_sequence_example.ParseFromString(example_str);
|
||||||
|
cc->OutputSidePackets()
|
||||||
|
.Tag(kSequenceExampleTag)
|
||||||
|
.Set(MakePacket<tensorflow::SequenceExample>(
|
||||||
|
std::move(tf_sequence_example)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
++current_idx;
|
||||||
|
}
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status TFRecordReaderCalculator::Process(CalculatorContext* cc) {
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
REGISTER_CALCULATOR(TFRecordReaderCalculator);
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -19,6 +19,7 @@
|
|||||||
#include "mediapipe/framework/formats/location.h"
|
#include "mediapipe/framework/formats/location.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/audio_decoder.pb.h"
|
||||||
#include "mediapipe/util/sequence/media_sequence.h"
|
#include "mediapipe/util/sequence/media_sequence.h"
|
||||||
#include "tensorflow/core/example/example.pb.h"
|
#include "tensorflow/core/example/example.pb.h"
|
||||||
#include "tensorflow/core/example/feature.pb.h"
|
#include "tensorflow/core/example/feature.pb.h"
|
||||||
@@ -37,6 +38,7 @@ const char kDatasetRootDirTag[] = "DATASET_ROOT";
|
|||||||
const char kDataPath[] = "DATA_PATH";
|
const char kDataPath[] = "DATA_PATH";
|
||||||
const char kPacketResamplerOptions[] = "RESAMPLER_OPTIONS";
|
const char kPacketResamplerOptions[] = "RESAMPLER_OPTIONS";
|
||||||
const char kImagesFrameRateTag[] = "IMAGE_FRAME_RATE";
|
const char kImagesFrameRateTag[] = "IMAGE_FRAME_RATE";
|
||||||
|
const char kAudioDecoderOptions[] = "AUDIO_DECODER_OPTIONS";
|
||||||
|
|
||||||
namespace tf = ::tensorflow;
|
namespace tf = ::tensorflow;
|
||||||
namespace mpms = ::mediapipe::mediasequence;
|
namespace mpms = ::mediapipe::mediasequence;
|
||||||
@@ -126,6 +128,11 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
|
|||||||
if (cc->OutputSidePackets().HasTag(kDataPath)) {
|
if (cc->OutputSidePackets().HasTag(kDataPath)) {
|
||||||
cc->OutputSidePackets().Tag(kDataPath).Set<std::string>();
|
cc->OutputSidePackets().Tag(kDataPath).Set<std::string>();
|
||||||
}
|
}
|
||||||
|
if (cc->OutputSidePackets().HasTag(kAudioDecoderOptions)) {
|
||||||
|
cc->OutputSidePackets()
|
||||||
|
.Tag(kAudioDecoderOptions)
|
||||||
|
.Set<AudioDecoderOptions>();
|
||||||
|
}
|
||||||
if (cc->OutputSidePackets().HasTag(kImagesFrameRateTag)) {
|
if (cc->OutputSidePackets().HasTag(kImagesFrameRateTag)) {
|
||||||
cc->OutputSidePackets().Tag(kImagesFrameRateTag).Set<double>();
|
cc->OutputSidePackets().Tag(kImagesFrameRateTag).Set<double>();
|
||||||
}
|
}
|
||||||
@@ -136,10 +143,11 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
|
|||||||
}
|
}
|
||||||
if ((options.has_padding_before_label() ||
|
if ((options.has_padding_before_label() ||
|
||||||
options.has_padding_after_label()) &&
|
options.has_padding_after_label()) &&
|
||||||
!(cc->OutputSidePackets().HasTag(kPacketResamplerOptions))) {
|
!(cc->OutputSidePackets().HasTag(kAudioDecoderOptions) ||
|
||||||
|
cc->OutputSidePackets().HasTag(kPacketResamplerOptions))) {
|
||||||
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
return ::mediapipe::InvalidArgumentErrorBuilder(MEDIAPIPE_LOC)
|
||||||
<< "If specifying padding, must output "
|
<< "If specifying padding, must output " << kPacketResamplerOptions
|
||||||
<< kPacketResamplerOptions;
|
<< "or" << kAudioDecoderOptions;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Optional streams.
|
// Optional streams.
|
||||||
@@ -260,7 +268,8 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
|
|||||||
// Set the start and end of the clip in the appropriate options protos.
|
// Set the start and end of the clip in the appropriate options protos.
|
||||||
double start_time = 0;
|
double start_time = 0;
|
||||||
double end_time = 0;
|
double end_time = 0;
|
||||||
if (cc->OutputSidePackets().HasTag(kPacketResamplerOptions)) {
|
if (cc->OutputSidePackets().HasTag(kAudioDecoderOptions) ||
|
||||||
|
cc->OutputSidePackets().HasTag(kPacketResamplerOptions)) {
|
||||||
if (mpms::HasClipStartTimestamp(sequence)) {
|
if (mpms::HasClipStartTimestamp(sequence)) {
|
||||||
start_time =
|
start_time =
|
||||||
Timestamp(mpms::GetClipStartTimestamp(sequence)).Seconds() -
|
Timestamp(mpms::GetClipStartTimestamp(sequence)).Seconds() -
|
||||||
@@ -271,6 +280,27 @@ class UnpackMediaSequenceCalculator : public CalculatorBase {
|
|||||||
options.padding_after_label();
|
options.padding_after_label();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if (cc->OutputSidePackets().HasTag(kAudioDecoderOptions)) {
|
||||||
|
auto audio_decoder_options = absl::make_unique<AudioDecoderOptions>(
|
||||||
|
options.base_audio_decoder_options());
|
||||||
|
if (mpms::HasClipStartTimestamp(sequence)) {
|
||||||
|
if (options.force_decoding_from_start_of_media()) {
|
||||||
|
audio_decoder_options->set_start_time(0);
|
||||||
|
} else {
|
||||||
|
audio_decoder_options->set_start_time(
|
||||||
|
start_time - options.extra_padding_from_media_decoder());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (mpms::HasClipEndTimestamp(sequence)) {
|
||||||
|
audio_decoder_options->set_end_time(
|
||||||
|
end_time + options.extra_padding_from_media_decoder());
|
||||||
|
}
|
||||||
|
LOG(INFO) << "Created AudioDecoderOptions:\n"
|
||||||
|
<< audio_decoder_options->DebugString();
|
||||||
|
cc->OutputSidePackets()
|
||||||
|
.Tag(kAudioDecoderOptions)
|
||||||
|
.Set(Adopt(audio_decoder_options.release()));
|
||||||
|
}
|
||||||
if (cc->OutputSidePackets().HasTag(kPacketResamplerOptions)) {
|
if (cc->OutputSidePackets().HasTag(kPacketResamplerOptions)) {
|
||||||
auto resampler_options = absl::make_unique<CalculatorOptions>();
|
auto resampler_options = absl::make_unique<CalculatorOptions>();
|
||||||
*(resampler_options->MutableExtension(
|
*(resampler_options->MutableExtension(
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ package mediapipe;
|
|||||||
|
|
||||||
import "mediapipe/calculators/core/packet_resampler_calculator.proto";
|
import "mediapipe/calculators/core/packet_resampler_calculator.proto";
|
||||||
import "mediapipe/framework/calculator.proto";
|
import "mediapipe/framework/calculator.proto";
|
||||||
|
import "mediapipe/util/audio_decoder.proto";
|
||||||
|
|
||||||
message UnpackMediaSequenceCalculatorOptions {
|
message UnpackMediaSequenceCalculatorOptions {
|
||||||
extend mediapipe.CalculatorOptions {
|
extend mediapipe.CalculatorOptions {
|
||||||
@@ -49,4 +50,10 @@ message UnpackMediaSequenceCalculatorOptions {
|
|||||||
// parameters for the MediaDecoderCalculator. End time parameters are still
|
// parameters for the MediaDecoderCalculator. End time parameters are still
|
||||||
// respected.
|
// respected.
|
||||||
optional bool force_decoding_from_start_of_media = 7;
|
optional bool force_decoding_from_start_of_media = 7;
|
||||||
|
|
||||||
|
// Stores the audio decoder settings for the graph. (e.g. which audio
|
||||||
|
// stream to pull from the video.) The sequence's metadata overrides
|
||||||
|
// the clip start and end times and outputs these for the
|
||||||
|
// AudioDecoderCalculator to consume.
|
||||||
|
optional AudioDecoderOptions base_audio_decoder_options = 9;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -23,6 +23,7 @@
|
|||||||
#include "mediapipe/framework/port/gtest.h"
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
#include "mediapipe/framework/port/rectangle.h"
|
#include "mediapipe/framework/port/rectangle.h"
|
||||||
#include "mediapipe/framework/port/status_matchers.h"
|
#include "mediapipe/framework/port/status_matchers.h"
|
||||||
|
#include "mediapipe/util/audio_decoder.pb.h"
|
||||||
#include "mediapipe/util/sequence/media_sequence.h"
|
#include "mediapipe/util/sequence/media_sequence.h"
|
||||||
#include "tensorflow/core/example/example.pb.h"
|
#include "tensorflow/core/example/example.pb.h"
|
||||||
|
|
||||||
@@ -459,6 +460,62 @@ TEST_F(UnpackMediaSequenceCalculatorTest, GetDatasetFromExample) {
|
|||||||
data_path_);
|
data_path_);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
TEST_F(UnpackMediaSequenceCalculatorTest, GetAudioDecoderOptions) {
|
||||||
|
CalculatorOptions options;
|
||||||
|
options.MutableExtension(UnpackMediaSequenceCalculatorOptions::ext)
|
||||||
|
->set_padding_before_label(1);
|
||||||
|
options.MutableExtension(UnpackMediaSequenceCalculatorOptions::ext)
|
||||||
|
->set_padding_after_label(2);
|
||||||
|
SetUpCalculator({}, {"AUDIO_DECODER_OPTIONS:audio_decoder_options"}, {},
|
||||||
|
&options);
|
||||||
|
runner_->MutableSidePackets()->Tag("SEQUENCE_EXAMPLE") =
|
||||||
|
Adopt(sequence_.release());
|
||||||
|
MP_ASSERT_OK(runner_->Run());
|
||||||
|
|
||||||
|
MP_EXPECT_OK(runner_->OutputSidePackets()
|
||||||
|
.Tag("AUDIO_DECODER_OPTIONS")
|
||||||
|
.ValidateAsType<AudioDecoderOptions>());
|
||||||
|
EXPECT_NEAR(runner_->OutputSidePackets()
|
||||||
|
.Tag("AUDIO_DECODER_OPTIONS")
|
||||||
|
.Get<AudioDecoderOptions>()
|
||||||
|
.start_time(),
|
||||||
|
2.0, 1e-5);
|
||||||
|
EXPECT_NEAR(runner_->OutputSidePackets()
|
||||||
|
.Tag("AUDIO_DECODER_OPTIONS")
|
||||||
|
.Get<AudioDecoderOptions>()
|
||||||
|
.end_time(),
|
||||||
|
7.0, 1e-5);
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(UnpackMediaSequenceCalculatorTest, GetAudioDecoderOptionsOverride) {
|
||||||
|
CalculatorOptions options;
|
||||||
|
options.MutableExtension(UnpackMediaSequenceCalculatorOptions::ext)
|
||||||
|
->set_padding_before_label(1);
|
||||||
|
options.MutableExtension(UnpackMediaSequenceCalculatorOptions::ext)
|
||||||
|
->set_padding_after_label(2);
|
||||||
|
options.MutableExtension(UnpackMediaSequenceCalculatorOptions::ext)
|
||||||
|
->set_force_decoding_from_start_of_media(true);
|
||||||
|
SetUpCalculator({}, {"AUDIO_DECODER_OPTIONS:audio_decoder_options"}, {},
|
||||||
|
&options);
|
||||||
|
runner_->MutableSidePackets()->Tag("SEQUENCE_EXAMPLE") =
|
||||||
|
Adopt(sequence_.release());
|
||||||
|
MP_ASSERT_OK(runner_->Run());
|
||||||
|
|
||||||
|
MP_EXPECT_OK(runner_->OutputSidePackets()
|
||||||
|
.Tag("AUDIO_DECODER_OPTIONS")
|
||||||
|
.ValidateAsType<AudioDecoderOptions>());
|
||||||
|
EXPECT_NEAR(runner_->OutputSidePackets()
|
||||||
|
.Tag("AUDIO_DECODER_OPTIONS")
|
||||||
|
.Get<AudioDecoderOptions>()
|
||||||
|
.start_time(),
|
||||||
|
0.0, 1e-5);
|
||||||
|
EXPECT_NEAR(runner_->OutputSidePackets()
|
||||||
|
.Tag("AUDIO_DECODER_OPTIONS")
|
||||||
|
.Get<AudioDecoderOptions>()
|
||||||
|
.end_time(),
|
||||||
|
7.0, 1e-5);
|
||||||
|
}
|
||||||
|
|
||||||
TEST_F(UnpackMediaSequenceCalculatorTest, GetPacketResamplingOptions) {
|
TEST_F(UnpackMediaSequenceCalculatorTest, GetPacketResamplingOptions) {
|
||||||
// TODO: Suport proto3 proto.Any in CalculatorOptions.
|
// TODO: Suport proto3 proto.Any in CalculatorOptions.
|
||||||
// TODO: Avoid proto2 extensions in "RESAMPLER_OPTIONS".
|
// TODO: Avoid proto2 extensions in "RESAMPLER_OPTIONS".
|
||||||
|
|||||||
@@ -0,0 +1,192 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include <iterator>
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/tensorflow/lapped_tensor_buffer_calculator.pb.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/packet.h"
|
||||||
|
#include "mediapipe/framework/port/canonical_errors.h"
|
||||||
|
#include "tensorflow/core/example/example.pb.h"
|
||||||
|
#include "tensorflow/core/example/feature.pb.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
const char kId[] = "id";
|
||||||
|
const char kRgb[] = "rgb";
|
||||||
|
const char kAudio[] = "audio";
|
||||||
|
const char kDesiredSegmentSize[] = "DESIRED_SEGMENT_SIZE";
|
||||||
|
const char kYt8mId[] = "YT8M_ID";
|
||||||
|
const char kYt8mSequenceExample[] = "YT8M_SEQUENCE_EXAMPLE";
|
||||||
|
const char kQuantizedRgbFeature[] = "QUANTIZED_RGB_FEATURE";
|
||||||
|
const char kQuantizedAudioFeature[] = "QUANTIZED_AUDIO_FEATURE";
|
||||||
|
const char kSegmentSize[] = "SEGMENT_SIZE";
|
||||||
|
const char kLappedTensorBufferCalculatorOptions[] =
|
||||||
|
"LAPPED_TENSOR_BUFFER_CALCULATOR_OPTIONS";
|
||||||
|
|
||||||
|
std::string GetQuantizedFeature(
|
||||||
|
const tensorflow::SequenceExample& sequence_example, const std::string& key,
|
||||||
|
int index) {
|
||||||
|
const auto& bytes_list = sequence_example.feature_lists()
|
||||||
|
.feature_list()
|
||||||
|
.at(key)
|
||||||
|
.feature()
|
||||||
|
.Get(index)
|
||||||
|
.bytes_list()
|
||||||
|
.value();
|
||||||
|
CHECK_EQ(1, bytes_list.size());
|
||||||
|
return bytes_list.Get(0);
|
||||||
|
}
|
||||||
|
} // namespace
|
||||||
|
|
||||||
|
// Unpacks YT8M Sequence Example. Note that the audio feature and rgb feature
|
||||||
|
// output are quantized. DequantizeByteArrayCalculator can do the dequantization
|
||||||
|
// for you.
|
||||||
|
//
|
||||||
|
// Example config:
|
||||||
|
// node {
|
||||||
|
// calculator: "UnpackYt8mSequenceExampleCalculator"
|
||||||
|
// input_side_packet: "YT8M_SEQUENCE_EXAMPLE:yt8m_sequence_example"
|
||||||
|
// output_stream: "QUANTIZED_RGB_FEATURE:quantized_rgb_feature"
|
||||||
|
// output_stream: "QUANTIZED_AUDIO_FEATURE:quantized_audio_feature"
|
||||||
|
// }
|
||||||
|
class UnpackYt8mSequenceExampleCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||||
|
cc->InputSidePackets()
|
||||||
|
.Tag(kYt8mSequenceExample)
|
||||||
|
.Set<tensorflow::SequenceExample>();
|
||||||
|
if (cc->InputSidePackets().HasTag(kDesiredSegmentSize)) {
|
||||||
|
cc->InputSidePackets().Tag(kDesiredSegmentSize).Set<int>();
|
||||||
|
}
|
||||||
|
cc->Outputs().Tag(kQuantizedRgbFeature).Set<std::string>();
|
||||||
|
cc->Outputs().Tag(kQuantizedAudioFeature).Set<std::string>();
|
||||||
|
if (cc->OutputSidePackets().HasTag(kYt8mId)) {
|
||||||
|
cc->OutputSidePackets().Tag(kYt8mId).Set<std::string>();
|
||||||
|
}
|
||||||
|
if (cc->OutputSidePackets().HasTag(kLappedTensorBufferCalculatorOptions)) {
|
||||||
|
cc->OutputSidePackets()
|
||||||
|
.Tag(kLappedTensorBufferCalculatorOptions)
|
||||||
|
.Set<::mediapipe::LappedTensorBufferCalculatorOptions>();
|
||||||
|
}
|
||||||
|
if (cc->OutputSidePackets().HasTag(kSegmentSize)) {
|
||||||
|
cc->OutputSidePackets().Tag(kSegmentSize).Set<int>();
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||||
|
const tensorflow::SequenceExample& sequence_example =
|
||||||
|
cc->InputSidePackets()
|
||||||
|
.Tag(kYt8mSequenceExample)
|
||||||
|
.Get<tensorflow::SequenceExample>();
|
||||||
|
const std::string& yt8m_id =
|
||||||
|
sequence_example.context().feature().at(kId).bytes_list().value().Get(
|
||||||
|
0);
|
||||||
|
if (cc->OutputSidePackets().HasTag(kYt8mId)) {
|
||||||
|
cc->OutputSidePackets().Tag(kYt8mId).Set(
|
||||||
|
MakePacket<std::string>(yt8m_id));
|
||||||
|
}
|
||||||
|
|
||||||
|
int rgb_feature_list_length =
|
||||||
|
sequence_example.feature_lists().feature_list().at(kRgb).feature_size();
|
||||||
|
int audio_feature_list_length = sequence_example.feature_lists()
|
||||||
|
.feature_list()
|
||||||
|
.at(kAudio)
|
||||||
|
.feature_size();
|
||||||
|
|
||||||
|
if (rgb_feature_list_length != audio_feature_list_length) {
|
||||||
|
return ::mediapipe::FailedPreconditionError(absl::StrCat(
|
||||||
|
"Data corruption: the length of audio features and rgb features are "
|
||||||
|
"not equal. Please check the sequence example that contains yt8m "
|
||||||
|
"id: ",
|
||||||
|
yt8m_id));
|
||||||
|
}
|
||||||
|
feature_list_length_ = rgb_feature_list_length;
|
||||||
|
if (cc->OutputSidePackets().HasTag(kLappedTensorBufferCalculatorOptions) ||
|
||||||
|
cc->OutputSidePackets().HasTag(kSegmentSize)) {
|
||||||
|
// If the desired segment size is specified, take the min of the length of
|
||||||
|
// the feature list and the desired size to be the output segment size.
|
||||||
|
int segment_size = feature_list_length_;
|
||||||
|
if (cc->InputSidePackets().HasTag(kDesiredSegmentSize)) {
|
||||||
|
int desired_segment_size =
|
||||||
|
cc->InputSidePackets().Tag(kDesiredSegmentSize).Get<int>();
|
||||||
|
RET_CHECK(desired_segment_size > 0)
|
||||||
|
<< "The desired segment size must be greater than zero.";
|
||||||
|
segment_size = std::min(
|
||||||
|
feature_list_length_,
|
||||||
|
cc->InputSidePackets().Tag(kDesiredSegmentSize).Get<int>());
|
||||||
|
}
|
||||||
|
if (cc->OutputSidePackets().HasTag(
|
||||||
|
kLappedTensorBufferCalculatorOptions)) {
|
||||||
|
auto lapped_tensor_buffer_calculator_options = absl::make_unique<
|
||||||
|
::mediapipe::LappedTensorBufferCalculatorOptions>();
|
||||||
|
lapped_tensor_buffer_calculator_options->set_add_batch_dim_to_tensors(
|
||||||
|
true);
|
||||||
|
lapped_tensor_buffer_calculator_options->set_buffer_size(segment_size);
|
||||||
|
lapped_tensor_buffer_calculator_options->set_overlap(segment_size - 1);
|
||||||
|
lapped_tensor_buffer_calculator_options->set_timestamp_offset(
|
||||||
|
segment_size - 1);
|
||||||
|
cc->OutputSidePackets()
|
||||||
|
.Tag(kLappedTensorBufferCalculatorOptions)
|
||||||
|
.Set(Adopt(lapped_tensor_buffer_calculator_options.release()));
|
||||||
|
}
|
||||||
|
if (cc->OutputSidePackets().HasTag(kSegmentSize)) {
|
||||||
|
cc->OutputSidePackets()
|
||||||
|
.Tag(kSegmentSize)
|
||||||
|
.Set(MakePacket<int>(segment_size));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
LOG(INFO) << "Reading the sequence example that contains yt8m id: "
|
||||||
|
<< yt8m_id << ". Feature list length: " << feature_list_length_;
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||||
|
if (current_index_ >= feature_list_length_) {
|
||||||
|
return ::mediapipe::tool::StatusStop();
|
||||||
|
}
|
||||||
|
const tensorflow::SequenceExample& sequence_example =
|
||||||
|
cc->InputSidePackets()
|
||||||
|
.Tag(kYt8mSequenceExample)
|
||||||
|
.Get<tensorflow::SequenceExample>();
|
||||||
|
|
||||||
|
// Uses microsecond as the unit of time. In the YT8M dataset, each feature
|
||||||
|
// represents a second.
|
||||||
|
const Timestamp timestamp = Timestamp(current_index_ * 1000000);
|
||||||
|
cc->Outputs()
|
||||||
|
.Tag(kQuantizedRgbFeature)
|
||||||
|
.AddPacket(
|
||||||
|
MakePacket<std::string>(
|
||||||
|
GetQuantizedFeature(sequence_example, kRgb, current_index_))
|
||||||
|
.At(timestamp));
|
||||||
|
cc->Outputs()
|
||||||
|
.Tag(kQuantizedAudioFeature)
|
||||||
|
.AddPacket(
|
||||||
|
MakePacket<std::string>(
|
||||||
|
GetQuantizedFeature(sequence_example, kAudio, current_index_))
|
||||||
|
.At(timestamp));
|
||||||
|
++current_index_;
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
int current_index_ = 0;
|
||||||
|
int feature_list_length_ = 0;
|
||||||
|
};
|
||||||
|
|
||||||
|
REGISTER_CALCULATOR(UnpackYt8mSequenceExampleCalculator);
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -23,10 +23,12 @@
|
|||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
namespace tf = ::tensorflow;
|
namespace {
|
||||||
|
|
||||||
auto& INPUT_1D = VectorFloatToTensorCalculatorOptions::INPUT_1D;
|
auto& INPUT_1D = VectorFloatToTensorCalculatorOptions::INPUT_1D;
|
||||||
auto& INPUT_2D = VectorFloatToTensorCalculatorOptions::INPUT_2D;
|
auto& INPUT_2D = VectorFloatToTensorCalculatorOptions::INPUT_2D;
|
||||||
|
} // namespace
|
||||||
|
|
||||||
|
namespace tf = ::tensorflow;
|
||||||
|
|
||||||
// The calculator expects one input (a packet containing a vector<float> or
|
// The calculator expects one input (a packet containing a vector<float> or
|
||||||
// vector<vector<float>>) and generates one output (a packet containing a
|
// vector<vector<float>>) and generates one output (a packet containing a
|
||||||
|
|||||||
@@ -0,0 +1,203 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
//
|
||||||
|
// Converts a single int or vector<int> or vector<vector<int>> to 1D (or 2D)
|
||||||
|
// tf::Tensor.
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/tensorflow/vector_int_to_tensor_calculator_options.pb.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
|
#include "mediapipe/framework/port/status.h"
|
||||||
|
#include "tensorflow/core/framework/tensor.h"
|
||||||
|
#include "tensorflow/core/framework/types.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
const char kVectorInt[] = "VECTOR_INT";
|
||||||
|
const char kSingleInt[] = "SINGLE_INT";
|
||||||
|
const char kTensorOut[] = "TENSOR_OUT";
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
auto& INPUT_1D = VectorIntToTensorCalculatorOptions::INPUT_1D;
|
||||||
|
auto& INPUT_2D = VectorIntToTensorCalculatorOptions::INPUT_2D;
|
||||||
|
} // namespace
|
||||||
|
|
||||||
|
namespace tf = ::tensorflow;
|
||||||
|
|
||||||
|
template <typename TensorType>
|
||||||
|
void AssignMatrixValue(int r, int c, int value, tf::Tensor* output_tensor) {
|
||||||
|
output_tensor->tensor<TensorType, 2>()(r, c) = value;
|
||||||
|
}
|
||||||
|
|
||||||
|
// The calculator expects one input (a packet containing a single int or
|
||||||
|
// vector<int> or vector<vector<int>>) and generates one output (a packet
|
||||||
|
// containing a tf::Tensor containing the same data). The output tensor will be
|
||||||
|
// either 1D or 2D with dimensions corresponding to the input vector int. It
|
||||||
|
// will hold DT_INT32 or DT_UINT8 or DT_INT64 values.
|
||||||
|
//
|
||||||
|
// Example config:
|
||||||
|
// node {
|
||||||
|
// calculator: "VectorIntToTensorCalculator"
|
||||||
|
// input_stream: "SINGLE_INT:segment_size_int_stream"
|
||||||
|
// output_stream: "TENSOR_OUT:segment_size_tensor"
|
||||||
|
// }
|
||||||
|
//
|
||||||
|
// or
|
||||||
|
//
|
||||||
|
// node {
|
||||||
|
// calculator: "VectorIntToTensorCalculator"
|
||||||
|
// input_stream: "VECTOR_INT:vector_int_features"
|
||||||
|
// output_stream: "TENSOR_OUT:tensor_features"
|
||||||
|
// }
|
||||||
|
class VectorIntToTensorCalculator : public CalculatorBase {
|
||||||
|
public:
|
||||||
|
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||||
|
|
||||||
|
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||||
|
::mediapipe::Status Process(CalculatorContext* cc) override;
|
||||||
|
|
||||||
|
private:
|
||||||
|
VectorIntToTensorCalculatorOptions options_;
|
||||||
|
};
|
||||||
|
REGISTER_CALCULATOR(VectorIntToTensorCalculator);
|
||||||
|
|
||||||
|
::mediapipe::Status VectorIntToTensorCalculator::GetContract(
|
||||||
|
CalculatorContract* cc) {
|
||||||
|
const auto& options = cc->Options<VectorIntToTensorCalculatorOptions>();
|
||||||
|
// Start with only one input packet.
|
||||||
|
RET_CHECK_EQ(cc->Inputs().NumEntries(), 1)
|
||||||
|
<< "Only one input stream is supported.";
|
||||||
|
if (options.input_size() == INPUT_2D) {
|
||||||
|
cc->Inputs().Tag(kVectorInt).Set<std::vector<std::vector<int>>>();
|
||||||
|
} else if (options.input_size() == INPUT_1D) {
|
||||||
|
if (cc->Inputs().HasTag(kSingleInt)) {
|
||||||
|
cc->Inputs().Tag(kSingleInt).Set<int>();
|
||||||
|
} else {
|
||||||
|
cc->Inputs().Tag(kVectorInt).Set<std::vector<int>>();
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
LOG(FATAL) << "input size not supported";
|
||||||
|
}
|
||||||
|
RET_CHECK_EQ(cc->Outputs().NumEntries(), 1)
|
||||||
|
<< "Only one output stream is supported.";
|
||||||
|
cc->Outputs().Tag(kTensorOut).Set<tf::Tensor>();
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status VectorIntToTensorCalculator::Open(CalculatorContext* cc) {
|
||||||
|
options_ = cc->Options<VectorIntToTensorCalculatorOptions>();
|
||||||
|
RET_CHECK(options_.tensor_data_type() == tf::DT_UINT8 ||
|
||||||
|
options_.tensor_data_type() == tf::DT_INT32 ||
|
||||||
|
options_.tensor_data_type() == tf::DT_INT64)
|
||||||
|
<< "Output tensor data type is not supported.";
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
::mediapipe::Status VectorIntToTensorCalculator::Process(
|
||||||
|
CalculatorContext* cc) {
|
||||||
|
tf::TensorShape tensor_shape;
|
||||||
|
if (options_.input_size() == INPUT_2D) {
|
||||||
|
const std::vector<std::vector<int>>& input =
|
||||||
|
cc->Inputs()
|
||||||
|
.Tag(kVectorInt)
|
||||||
|
.Value()
|
||||||
|
.Get<std::vector<std::vector<int>>>();
|
||||||
|
|
||||||
|
const int32 rows = input.size();
|
||||||
|
CHECK_GE(rows, 1);
|
||||||
|
const int32 cols = input[0].size();
|
||||||
|
CHECK_GE(cols, 1);
|
||||||
|
for (int i = 1; i < rows; ++i) {
|
||||||
|
CHECK_EQ(input[i].size(), cols);
|
||||||
|
}
|
||||||
|
if (options_.transpose()) {
|
||||||
|
tensor_shape = tf::TensorShape({cols, rows});
|
||||||
|
} else {
|
||||||
|
tensor_shape = tf::TensorShape({rows, cols});
|
||||||
|
}
|
||||||
|
auto output = ::absl::make_unique<tf::Tensor>(options_.tensor_data_type(),
|
||||||
|
tensor_shape);
|
||||||
|
if (options_.transpose()) {
|
||||||
|
for (int r = 0; r < rows; ++r) {
|
||||||
|
for (int c = 0; c < cols; ++c) {
|
||||||
|
switch (options_.tensor_data_type()) {
|
||||||
|
case tf::DT_INT64:
|
||||||
|
AssignMatrixValue<tf::int64>(c, r, input[r][c], output.get());
|
||||||
|
break;
|
||||||
|
case tf::DT_UINT8:
|
||||||
|
AssignMatrixValue<uint8>(c, r, input[r][c], output.get());
|
||||||
|
break;
|
||||||
|
case tf::DT_INT32:
|
||||||
|
AssignMatrixValue<int>(c, r, input[r][c], output.get());
|
||||||
|
break;
|
||||||
|
default:
|
||||||
|
LOG(FATAL) << "tensor data type is not supported.";
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
for (int r = 0; r < rows; ++r) {
|
||||||
|
for (int c = 0; c < cols; ++c) {
|
||||||
|
switch (options_.tensor_data_type()) {
|
||||||
|
case tf::DT_INT64:
|
||||||
|
AssignMatrixValue<tf::int64>(r, c, input[r][c], output.get());
|
||||||
|
break;
|
||||||
|
case tf::DT_UINT8:
|
||||||
|
AssignMatrixValue<uint8>(r, c, input[r][c], output.get());
|
||||||
|
break;
|
||||||
|
case tf::DT_INT32:
|
||||||
|
AssignMatrixValue<int>(r, c, input[r][c], output.get());
|
||||||
|
break;
|
||||||
|
default:
|
||||||
|
LOG(FATAL) << "tensor data type is not supported.";
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
cc->Outputs().Tag(kTensorOut).Add(output.release(), cc->InputTimestamp());
|
||||||
|
} else if (options_.input_size() == INPUT_1D) {
|
||||||
|
std::vector<int> input;
|
||||||
|
if (cc->Inputs().HasTag(kSingleInt)) {
|
||||||
|
input.push_back(cc->Inputs().Tag(kSingleInt).Get<int>());
|
||||||
|
} else {
|
||||||
|
input = cc->Inputs().Tag(kVectorInt).Value().Get<std::vector<int>>();
|
||||||
|
}
|
||||||
|
CHECK_GE(input.size(), 1);
|
||||||
|
const int32 length = input.size();
|
||||||
|
tensor_shape = tf::TensorShape({length});
|
||||||
|
auto output = ::absl::make_unique<tf::Tensor>(options_.tensor_data_type(),
|
||||||
|
tensor_shape);
|
||||||
|
for (int i = 0; i < length; ++i) {
|
||||||
|
switch (options_.tensor_data_type()) {
|
||||||
|
case tf::DT_INT64:
|
||||||
|
output->tensor<tf::int64, 1>()(i) = input.at(i);
|
||||||
|
break;
|
||||||
|
case tf::DT_UINT8:
|
||||||
|
output->tensor<uint8, 1>()(i) = input.at(i);
|
||||||
|
break;
|
||||||
|
case tf::DT_INT32:
|
||||||
|
output->tensor<int, 1>()(i) = input.at(i);
|
||||||
|
break;
|
||||||
|
default:
|
||||||
|
LOG(FATAL) << "tensor data type is not supported.";
|
||||||
|
}
|
||||||
|
}
|
||||||
|
cc->Outputs().Tag(kTensorOut).Add(output.release(), cc->InputTimestamp());
|
||||||
|
} else {
|
||||||
|
LOG(FATAL) << "input size not supported";
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -0,0 +1,43 @@
|
|||||||
|
// Copyright 2019 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
syntax = "proto2";
|
||||||
|
|
||||||
|
package mediapipe;
|
||||||
|
|
||||||
|
import "mediapipe/framework/calculator.proto";
|
||||||
|
import "tensorflow/core/framework/types.proto";
|
||||||
|
|
||||||
|
message VectorIntToTensorCalculatorOptions {
|
||||||
|
extend mediapipe.CalculatorOptions {
|
||||||
|
optional VectorIntToTensorCalculatorOptions ext = 275364184;
|
||||||
|
}
|
||||||
|
enum InputSize {
|
||||||
|
UNKNOWN = 0;
|
||||||
|
INPUT_1D = 1;
|
||||||
|
INPUT_2D = 2;
|
||||||
|
}
|
||||||
|
|
||||||
|
// If input_size is INPUT_2D, unpack a vector<vector<int>> to a
|
||||||
|
// 2d tensor (matrix). If INPUT_1D, convert a single int or vector<int>
|
||||||
|
// into a 1d tensor (vector).
|
||||||
|
optional InputSize input_size = 1 [default = INPUT_1D];
|
||||||
|
|
||||||
|
// If true, the output tensor is transposed.
|
||||||
|
// Otherwise, the output tensor is not transposed.
|
||||||
|
// It will be ignored if tensor_is_2d is INPUT_1D.
|
||||||
|
optional bool transpose = 2 [default = false];
|
||||||
|
|
||||||
|
optional tensorflow.DataType tensor_data_type = 3 [default = DT_INT32];
|
||||||
|
}
|
||||||
@@ -0,0 +1,202 @@
|
|||||||
|
// Copyright 2018 The MediaPipe Authors.
|
||||||
|
//
|
||||||
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
// you may not use this file except in compliance with the License.
|
||||||
|
// You may obtain a copy of the License at
|
||||||
|
//
|
||||||
|
// http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
//
|
||||||
|
// Unless required by applicable law or agreed to in writing, software
|
||||||
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
// See the License for the specific language governing permissions and
|
||||||
|
// limitations under the License.
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/tensorflow/vector_int_to_tensor_calculator_options.pb.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/calculator_runner.h"
|
||||||
|
#include "mediapipe/framework/port/gmock.h"
|
||||||
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "tensorflow/core/framework/tensor.h"
|
||||||
|
#include "tensorflow/core/framework/types.pb.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
namespace tf = ::tensorflow;
|
||||||
|
|
||||||
|
class VectorIntToTensorCalculatorTest : public ::testing::Test {
|
||||||
|
protected:
|
||||||
|
void SetUpRunner(
|
||||||
|
const VectorIntToTensorCalculatorOptions::InputSize input_size,
|
||||||
|
const tensorflow::DataType tensor_data_type, const bool transpose,
|
||||||
|
const bool single_value) {
|
||||||
|
CalculatorGraphConfig::Node config;
|
||||||
|
config.set_calculator("VectorIntToTensorCalculator");
|
||||||
|
if (single_value) {
|
||||||
|
config.add_input_stream("SINGLE_INT:input_int");
|
||||||
|
} else {
|
||||||
|
config.add_input_stream("VECTOR_INT:input_int");
|
||||||
|
}
|
||||||
|
config.add_output_stream("TENSOR_OUT:output_tensor");
|
||||||
|
auto options = config.mutable_options()->MutableExtension(
|
||||||
|
VectorIntToTensorCalculatorOptions::ext);
|
||||||
|
options->set_input_size(input_size);
|
||||||
|
options->set_transpose(transpose);
|
||||||
|
options->set_tensor_data_type(tensor_data_type);
|
||||||
|
runner_ = ::absl::make_unique<CalculatorRunner>(config);
|
||||||
|
}
|
||||||
|
|
||||||
|
void TestConvertFromVectoVectorInt(const bool transpose) {
|
||||||
|
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_2D,
|
||||||
|
tensorflow::DT_INT32, transpose, false);
|
||||||
|
auto input = ::absl::make_unique<std::vector<std::vector<int>>>(
|
||||||
|
2, std::vector<int>(2));
|
||||||
|
for (int i = 0; i < 2; ++i) {
|
||||||
|
for (int j = 0; j < 2; ++j) {
|
||||||
|
input->at(i).at(j) = i * 2 + j;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const int64 time = 1234;
|
||||||
|
runner_->MutableInputs()
|
||||||
|
->Tag("VECTOR_INT")
|
||||||
|
.packets.push_back(Adopt(input.release()).At(Timestamp(time)));
|
||||||
|
|
||||||
|
EXPECT_TRUE(runner_->Run().ok());
|
||||||
|
|
||||||
|
const std::vector<Packet>& output_packets =
|
||||||
|
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||||
|
EXPECT_EQ(1, output_packets.size());
|
||||||
|
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||||
|
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||||
|
|
||||||
|
EXPECT_EQ(2, output_tensor.dims());
|
||||||
|
EXPECT_EQ(tf::DT_INT32, output_tensor.dtype());
|
||||||
|
const auto matrix = output_tensor.matrix<int>();
|
||||||
|
|
||||||
|
for (int i = 0; i < 2; ++i) {
|
||||||
|
for (int j = 0; j < 2; ++j) {
|
||||||
|
if (!transpose) {
|
||||||
|
EXPECT_EQ(i * 2 + j, matrix(i, j));
|
||||||
|
} else {
|
||||||
|
EXPECT_EQ(j * 2 + i, matrix(i, j));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
std::unique_ptr<CalculatorRunner> runner_;
|
||||||
|
};
|
||||||
|
|
||||||
|
TEST_F(VectorIntToTensorCalculatorTest, TestSingleValue) {
|
||||||
|
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||||
|
tensorflow::DT_INT32, false, true);
|
||||||
|
const int64 time = 1234;
|
||||||
|
runner_->MutableInputs()
|
||||||
|
->Tag("SINGLE_INT")
|
||||||
|
.packets.push_back(MakePacket<int>(1).At(Timestamp(time)));
|
||||||
|
|
||||||
|
EXPECT_TRUE(runner_->Run().ok());
|
||||||
|
|
||||||
|
const std::vector<Packet>& output_packets =
|
||||||
|
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||||
|
EXPECT_EQ(1, output_packets.size());
|
||||||
|
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||||
|
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||||
|
|
||||||
|
EXPECT_EQ(1, output_tensor.dims());
|
||||||
|
EXPECT_EQ(tf::DT_INT32, output_tensor.dtype());
|
||||||
|
const auto vec = output_tensor.vec<int32>();
|
||||||
|
EXPECT_EQ(1, vec(0));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(VectorIntToTensorCalculatorTest, TesOneDim) {
|
||||||
|
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||||
|
tensorflow::DT_INT32, false, false);
|
||||||
|
auto input = ::absl::make_unique<std::vector<int>>(5);
|
||||||
|
for (int i = 0; i < 5; ++i) {
|
||||||
|
input->at(i) = i;
|
||||||
|
}
|
||||||
|
const int64 time = 1234;
|
||||||
|
runner_->MutableInputs()
|
||||||
|
->Tag("VECTOR_INT")
|
||||||
|
.packets.push_back(Adopt(input.release()).At(Timestamp(time)));
|
||||||
|
|
||||||
|
EXPECT_TRUE(runner_->Run().ok());
|
||||||
|
|
||||||
|
const std::vector<Packet>& output_packets =
|
||||||
|
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||||
|
EXPECT_EQ(1, output_packets.size());
|
||||||
|
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||||
|
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||||
|
|
||||||
|
EXPECT_EQ(1, output_tensor.dims());
|
||||||
|
EXPECT_EQ(tf::DT_INT32, output_tensor.dtype());
|
||||||
|
const auto vec = output_tensor.vec<int32>();
|
||||||
|
|
||||||
|
for (int i = 0; i < 5; ++i) {
|
||||||
|
EXPECT_EQ(i, vec(i));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(VectorIntToTensorCalculatorTest, TestTwoDims) {
|
||||||
|
for (bool transpose : {false, true}) {
|
||||||
|
TestConvertFromVectoVectorInt(transpose);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(VectorIntToTensorCalculatorTest, TestInt64) {
|
||||||
|
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||||
|
tensorflow::DT_INT64, false, true);
|
||||||
|
const int64 time = 1234;
|
||||||
|
runner_->MutableInputs()
|
||||||
|
->Tag("SINGLE_INT")
|
||||||
|
.packets.push_back(MakePacket<int>(2 ^ 31).At(Timestamp(time)));
|
||||||
|
|
||||||
|
EXPECT_TRUE(runner_->Run().ok());
|
||||||
|
|
||||||
|
const std::vector<Packet>& output_packets =
|
||||||
|
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||||
|
EXPECT_EQ(1, output_packets.size());
|
||||||
|
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||||
|
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||||
|
|
||||||
|
EXPECT_EQ(1, output_tensor.dims());
|
||||||
|
EXPECT_EQ(tf::DT_INT64, output_tensor.dtype());
|
||||||
|
const auto vec = output_tensor.vec<tf::int64>();
|
||||||
|
EXPECT_EQ(2 ^ 31, vec(0));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST_F(VectorIntToTensorCalculatorTest, TestUint8) {
|
||||||
|
SetUpRunner(VectorIntToTensorCalculatorOptions::INPUT_1D,
|
||||||
|
tensorflow::DT_UINT8, false, false);
|
||||||
|
auto input = ::absl::make_unique<std::vector<int>>(5);
|
||||||
|
for (int i = 0; i < 5; ++i) {
|
||||||
|
input->at(i) = i;
|
||||||
|
}
|
||||||
|
const int64 time = 1234;
|
||||||
|
runner_->MutableInputs()
|
||||||
|
->Tag("VECTOR_INT")
|
||||||
|
.packets.push_back(Adopt(input.release()).At(Timestamp(time)));
|
||||||
|
|
||||||
|
EXPECT_TRUE(runner_->Run().ok());
|
||||||
|
|
||||||
|
const std::vector<Packet>& output_packets =
|
||||||
|
runner_->Outputs().Tag("TENSOR_OUT").packets;
|
||||||
|
EXPECT_EQ(1, output_packets.size());
|
||||||
|
EXPECT_EQ(time, output_packets[0].Timestamp().Value());
|
||||||
|
const tf::Tensor& output_tensor = output_packets[0].Get<tf::Tensor>();
|
||||||
|
|
||||||
|
EXPECT_EQ(1, output_tensor.dims());
|
||||||
|
EXPECT_EQ(tf::DT_UINT8, output_tensor.dtype());
|
||||||
|
const auto vec = output_tensor.vec<uint8>();
|
||||||
|
|
||||||
|
for (int i = 0; i < 5; ++i) {
|
||||||
|
EXPECT_EQ(i, vec(i));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace
|
||||||
|
} // namespace mediapipe
|
||||||
@@ -13,12 +13,12 @@
|
|||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
#
|
#
|
||||||
|
|
||||||
|
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||||
|
|
||||||
licenses(["notice"]) # Apache 2.0
|
licenses(["notice"]) # Apache 2.0
|
||||||
|
|
||||||
package(default_visibility = ["//visibility:private"])
|
package(default_visibility = ["//visibility:private"])
|
||||||
|
|
||||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
|
||||||
|
|
||||||
proto_library(
|
proto_library(
|
||||||
name = "ssd_anchors_calculator_proto",
|
name = "ssd_anchors_calculator_proto",
|
||||||
srcs = ["ssd_anchors_calculator.proto"],
|
srcs = ["ssd_anchors_calculator.proto"],
|
||||||
@@ -195,6 +195,12 @@ cc_test(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "util",
|
||||||
|
hdrs = ["util.h"],
|
||||||
|
alwayslink = 1,
|
||||||
|
)
|
||||||
|
|
||||||
cc_library(
|
cc_library(
|
||||||
name = "tflite_inference_calculator",
|
name = "tflite_inference_calculator",
|
||||||
srcs = ["tflite_inference_calculator.cc"],
|
srcs = ["tflite_inference_calculator.cc"],
|
||||||
@@ -214,6 +220,7 @@ cc_library(
|
|||||||
}),
|
}),
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
|
":util",
|
||||||
":tflite_inference_calculator_cc_proto",
|
":tflite_inference_calculator_cc_proto",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/util:resource_util",
|
"//mediapipe/util:resource_util",
|
||||||
@@ -222,20 +229,31 @@ cc_library(
|
|||||||
"//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",
|
||||||
] + select({
|
] + select({
|
||||||
"//mediapipe:android": [
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
|
"//mediapipe:ios": [
|
||||||
|
"//mediapipe/gpu:MPPMetalHelper",
|
||||||
|
"//mediapipe/gpu:MPPMetalUtil",
|
||||||
|
"//mediapipe/gpu:gpu_buffer",
|
||||||
|
"//mediapipe/objc:mediapipe_framework_ios",
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/metal:buffer_convert",
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate_internal",
|
||||||
|
],
|
||||||
|
"//conditions:default": [
|
||||||
"//mediapipe/gpu:gl_calculator_helper",
|
"//mediapipe/gpu:gl_calculator_helper",
|
||||||
"//mediapipe/gpu:gpu_buffer",
|
"//mediapipe/gpu:gpu_buffer",
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/common:shape",
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
|
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
|
||||||
],
|
],
|
||||||
"//mediapipe:ios": [
|
}) + select({
|
||||||
"//mediapipe/gpu:MPPMetalHelper",
|
|
||||||
"//mediapipe/objc:mediapipe_framework_ios",
|
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
|
|
||||||
],
|
|
||||||
"//conditions:default": [],
|
"//conditions:default": [],
|
||||||
|
"//mediapipe:android": [
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/nnapi:nnapi_delegate",
|
||||||
|
],
|
||||||
}),
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -259,33 +277,33 @@ cc_library(
|
|||||||
}),
|
}),
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
|
":util",
|
||||||
":tflite_converter_calculator_cc_proto",
|
":tflite_converter_calculator_cc_proto",
|
||||||
"//mediapipe/util:resource_util",
|
"//mediapipe/util:resource_util",
|
||||||
"//mediapipe/framework:calculator_framework",
|
"//mediapipe/framework:calculator_framework",
|
||||||
"//mediapipe/framework/formats:image_frame",
|
"//mediapipe/framework/formats:image_frame",
|
||||||
"//mediapipe/framework/formats:matrix",
|
"//mediapipe/framework/formats:matrix",
|
||||||
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
|
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
|
||||||
"//mediapipe/framework/tool:status_util",
|
|
||||||
"//mediapipe/framework/port:status",
|
|
||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"@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({
|
] + select({
|
||||||
"//mediapipe:android": [
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
"//mediapipe/gpu:gl_calculator_helper",
|
|
||||||
"//mediapipe/gpu:gpu_buffer",
|
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
|
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
|
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
|
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
|
|
||||||
],
|
|
||||||
"//mediapipe:ios": [
|
"//mediapipe:ios": [
|
||||||
|
"//mediapipe/gpu:MPPMetalUtil",
|
||||||
"//mediapipe/gpu:gpu_buffer",
|
"//mediapipe/gpu:gpu_buffer",
|
||||||
"//mediapipe/gpu:MPPMetalHelper",
|
"//mediapipe/gpu:MPPMetalHelper",
|
||||||
"//mediapipe/objc:mediapipe_framework_ios",
|
"//mediapipe/objc:mediapipe_framework_ios",
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
|
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
|
||||||
],
|
],
|
||||||
"//conditions:default": [],
|
"//conditions:default": [
|
||||||
|
"//mediapipe/gpu:gpu_buffer",
|
||||||
|
"//mediapipe/gpu:gl_calculator_helper",
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
|
||||||
|
],
|
||||||
}),
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -295,6 +313,7 @@ cc_library(
|
|||||||
srcs = ["tflite_tensors_to_segmentation_calculator.cc"],
|
srcs = ["tflite_tensors_to_segmentation_calculator.cc"],
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
|
":util",
|
||||||
":tflite_tensors_to_segmentation_calculator_cc_proto",
|
":tflite_tensors_to_segmentation_calculator_cc_proto",
|
||||||
"@com_google_absl//absl/strings:str_format",
|
"@com_google_absl//absl/strings:str_format",
|
||||||
"@com_google_absl//absl/types:span",
|
"@com_google_absl//absl/types:span",
|
||||||
@@ -308,7 +327,9 @@ cc_library(
|
|||||||
"//mediapipe/util:resource_util",
|
"//mediapipe/util:resource_util",
|
||||||
"@org_tensorflow//tensorflow/lite:framework",
|
"@org_tensorflow//tensorflow/lite:framework",
|
||||||
] + select({
|
] + select({
|
||||||
"//mediapipe:android": [
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
|
"//mediapipe:ios": [],
|
||||||
|
"//conditions:default": [
|
||||||
"//mediapipe/gpu:gl_calculator_helper",
|
"//mediapipe/gpu:gl_calculator_helper",
|
||||||
"//mediapipe/gpu:gl_simple_shaders",
|
"//mediapipe/gpu:gl_simple_shaders",
|
||||||
"//mediapipe/gpu:gpu_buffer",
|
"//mediapipe/gpu:gpu_buffer",
|
||||||
@@ -319,7 +340,6 @@ cc_library(
|
|||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_texture",
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_texture",
|
||||||
],
|
],
|
||||||
"//conditions:default": [],
|
|
||||||
}),
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -346,8 +366,23 @@ cc_test(
|
|||||||
cc_library(
|
cc_library(
|
||||||
name = "tflite_tensors_to_detections_calculator",
|
name = "tflite_tensors_to_detections_calculator",
|
||||||
srcs = ["tflite_tensors_to_detections_calculator.cc"],
|
srcs = ["tflite_tensors_to_detections_calculator.cc"],
|
||||||
|
copts = select({
|
||||||
|
"//mediapipe:ios": [
|
||||||
|
"-x objective-c++",
|
||||||
|
"-fobjc-arc", # enable reference-counting
|
||||||
|
],
|
||||||
|
"//conditions:default": [],
|
||||||
|
}),
|
||||||
|
linkopts = select({
|
||||||
|
"//mediapipe:ios": [
|
||||||
|
"-framework CoreVideo",
|
||||||
|
"-framework MetalKit",
|
||||||
|
],
|
||||||
|
"//conditions:default": [],
|
||||||
|
}),
|
||||||
visibility = ["//visibility:public"],
|
visibility = ["//visibility:public"],
|
||||||
deps = [
|
deps = [
|
||||||
|
":util",
|
||||||
":tflite_tensors_to_detections_calculator_cc_proto",
|
":tflite_tensors_to_detections_calculator_cc_proto",
|
||||||
"//mediapipe/framework/formats:detection_cc_proto",
|
"//mediapipe/framework/formats:detection_cc_proto",
|
||||||
"@com_google_absl//absl/strings:str_format",
|
"@com_google_absl//absl/strings:str_format",
|
||||||
@@ -359,14 +394,21 @@ cc_library(
|
|||||||
"//mediapipe/framework/port:ret_check",
|
"//mediapipe/framework/port:ret_check",
|
||||||
"@org_tensorflow//tensorflow/lite:framework",
|
"@org_tensorflow//tensorflow/lite:framework",
|
||||||
] + select({
|
] + select({
|
||||||
"//mediapipe:android": [
|
"//mediapipe/gpu:disable_gpu": [],
|
||||||
|
"//mediapipe:ios": [
|
||||||
|
"//mediapipe/gpu:MPPMetalUtil",
|
||||||
|
"//mediapipe/gpu:gpu_buffer",
|
||||||
|
"//mediapipe/gpu:MPPMetalHelper",
|
||||||
|
"//mediapipe/objc:mediapipe_framework_ios",
|
||||||
|
"@org_tensorflow//tensorflow/lite/delegates/gpu:metal_delegate",
|
||||||
|
],
|
||||||
|
"//conditions:default": [
|
||||||
"//mediapipe/gpu:gl_calculator_helper",
|
"//mediapipe/gpu:gl_calculator_helper",
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
|
"@org_tensorflow//tensorflow/lite/delegates/gpu:gl_delegate",
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_buffer",
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_program",
|
||||||
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
|
"@org_tensorflow//tensorflow/lite/delegates/gpu/gl:gl_shader",
|
||||||
],
|
],
|
||||||
"//conditions:default": [],
|
|
||||||
}),
|
}),
|
||||||
alwayslink = 1,
|
alwayslink = 1,
|
||||||
)
|
)
|
||||||
@@ -443,6 +485,7 @@ 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/strings",
|
||||||
"@org_tensorflow//tensorflow/lite:framework",
|
"@org_tensorflow//tensorflow/lite:framework",
|
||||||
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
"@org_tensorflow//tensorflow/lite/kernels:builtin_ops",
|
||||||
],
|
],
|
||||||
|
|||||||
@@ -16,37 +16,38 @@
|
|||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
#include "mediapipe/calculators/tflite/tflite_converter_calculator.pb.h"
|
#include "mediapipe/calculators/tflite/tflite_converter_calculator.pb.h"
|
||||||
|
#include "mediapipe/calculators/tflite/util.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/matrix.h"
|
#include "mediapipe/framework/formats/matrix.h"
|
||||||
#include "mediapipe/framework/port/canonical_errors.h"
|
|
||||||
#include "mediapipe/framework/port/ret_check.h"
|
#include "mediapipe/framework/port/ret_check.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"
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||||
#include "mediapipe/gpu/gpu_buffer.h"
|
#include "mediapipe/gpu/gpu_buffer.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
|
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
|
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
|
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
|
||||||
#endif // __ANDROID__
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#if defined(MEDIAPIPE_IOS)
|
||||||
#import <CoreVideo/CoreVideo.h>
|
#import <CoreVideo/CoreVideo.h>
|
||||||
#import <Metal/Metal.h>
|
#import <Metal/Metal.h>
|
||||||
#import <MetalKit/MetalKit.h>
|
#import <MetalKit/MetalKit.h>
|
||||||
|
|
||||||
#import "mediapipe/gpu/MPPMetalHelper.h"
|
#import "mediapipe/gpu/MPPMetalHelper.h"
|
||||||
|
#include "mediapipe/gpu/MPPMetalUtil.h"
|
||||||
#include "mediapipe/gpu/gpu_buffer.h"
|
#include "mediapipe/gpu/gpu_buffer.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
|
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
|
||||||
#endif // iOS
|
#endif // iOS
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
|
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
typedef id<MTLBuffer> GpuTensor;
|
typedef id<MTLBuffer> GpuTensor;
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
@@ -66,26 +67,27 @@ typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>
|
|||||||
|
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
using ::tflite::gpu::gl::GlBuffer;
|
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
||||||
using ::tflite::gpu::gl::GlProgram;
|
using ::tflite::gpu::gl::GlProgram;
|
||||||
using ::tflite::gpu::gl::GlShader;
|
using ::tflite::gpu::gl::GlShader;
|
||||||
struct GPUData {
|
struct GPUData {
|
||||||
int elements = 1;
|
int elements = 1;
|
||||||
GlBuffer buffer;
|
GpuTensor buffer;
|
||||||
GlShader shader;
|
GlShader shader;
|
||||||
GlProgram program;
|
GlProgram program;
|
||||||
};
|
};
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
struct GPUData {
|
struct GPUData {
|
||||||
int elements = 1;
|
int elements = 1;
|
||||||
id<MTLBuffer> buffer;
|
GpuTensor buffer;
|
||||||
id<MTLComputePipelineState> pipeline_state;
|
id<MTLComputePipelineState> pipeline_state;
|
||||||
};
|
};
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
// Calculator for normalizing and converting an ImageFrame or Matrix
|
// Calculator for normalizing and converting an ImageFrame or Matrix
|
||||||
// into a TfLiteTensor (float 32) or a GpuBuffer to a tflite::gpu::GlBuffer.
|
// into a TfLiteTensor (float 32) or a GpuBuffer to a tflite::gpu::GlBuffer
|
||||||
|
// or MTLBuffer.
|
||||||
//
|
//
|
||||||
// This calculator is designed to be used with the TfLiteInferenceCalcualtor,
|
// This calculator is designed to be used with the TfLiteInferenceCalcualtor,
|
||||||
// as a pre-processing step for calculator inputs.
|
// as a pre-processing step for calculator inputs.
|
||||||
@@ -102,7 +104,7 @@ struct GPUData {
|
|||||||
// Output:
|
// Output:
|
||||||
// One of the following tags:
|
// One of the following tags:
|
||||||
// TENSORS - Vector of TfLiteTensor of type kTfLiteFloat32, or kTfLiteUint8.
|
// TENSORS - Vector of TfLiteTensor of type kTfLiteFloat32, or kTfLiteUint8.
|
||||||
// TENSORS_GPU - vector of GlBuffer.
|
// TENSORS_GPU - vector of GlBuffer or MTLBuffer.
|
||||||
//
|
//
|
||||||
// Example use:
|
// Example use:
|
||||||
// node {
|
// node {
|
||||||
@@ -144,10 +146,10 @@ class TfLiteConverterCalculator : public CalculatorBase {
|
|||||||
|
|
||||||
std::unique_ptr<tflite::Interpreter> interpreter_ = nullptr;
|
std::unique_ptr<tflite::Interpreter> interpreter_ = nullptr;
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||||
std::unique_ptr<GPUData> gpu_data_out_;
|
std::unique_ptr<GPUData> gpu_data_out_;
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
MPPMetalHelper* gpu_helper_ = nullptr;
|
MPPMetalHelper* gpu_helper_ = nullptr;
|
||||||
std::unique_ptr<GPUData> gpu_data_out_;
|
std::unique_ptr<GPUData> gpu_data_out_;
|
||||||
#endif
|
#endif
|
||||||
@@ -175,25 +177,33 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
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;
|
||||||
|
|
||||||
if (cc->Inputs().HasTag("IMAGE")) cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
if (cc->Inputs().HasTag("IMAGE")) cc->Inputs().Tag("IMAGE").Set<ImageFrame>();
|
||||||
if (cc->Inputs().HasTag("MATRIX")) cc->Inputs().Tag("MATRIX").Set<Matrix>();
|
if (cc->Inputs().HasTag("MATRIX")) cc->Inputs().Tag("MATRIX").Set<Matrix>();
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
if (cc->Inputs().HasTag("IMAGE_GPU"))
|
if (cc->Inputs().HasTag("IMAGE_GPU")) {
|
||||||
cc->Inputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
|
cc->Inputs().Tag("IMAGE_GPU").Set<mediapipe::GpuBuffer>();
|
||||||
#endif
|
use_gpu |= true;
|
||||||
|
}
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
if (cc->Outputs().HasTag("TENSORS"))
|
if (cc->Outputs().HasTag("TENSORS"))
|
||||||
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
if (cc->Outputs().HasTag("TENSORS_GPU"))
|
if (cc->Outputs().HasTag("TENSORS_GPU")) {
|
||||||
cc->Outputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
cc->Outputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||||
#endif
|
use_gpu |= true;
|
||||||
|
}
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
if (use_gpu) {
|
||||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||||
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
|
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
|
||||||
#endif
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
// Assign this calculator's default InputStreamHandler.
|
// Assign this calculator's default InputStreamHandler.
|
||||||
cc->SetInputStreamHandler("FixedSizeInputStreamHandler");
|
cc->SetInputStreamHandler("FixedSizeInputStreamHandler");
|
||||||
@@ -208,10 +218,10 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
|
|
||||||
if (cc->Inputs().HasTag("IMAGE_GPU") ||
|
if (cc->Inputs().HasTag("IMAGE_GPU") ||
|
||||||
cc->Outputs().HasTag("IMAGE_OUT_GPU")) {
|
cc->Outputs().HasTag("IMAGE_OUT_GPU")) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
use_gpu_ = true;
|
use_gpu_ = true;
|
||||||
#else
|
#else
|
||||||
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
|
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
||||||
#endif
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -221,9 +231,9 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
cc->Outputs().HasTag("TENSORS_GPU"));
|
cc->Outputs().HasTag("TENSORS_GPU"));
|
||||||
// Cannot use quantization.
|
// Cannot use quantization.
|
||||||
use_quantized_tensors_ = false;
|
use_quantized_tensors_ = false;
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
|
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
|
||||||
RET_CHECK(gpu_helper_);
|
RET_CHECK(gpu_helper_);
|
||||||
#endif
|
#endif
|
||||||
@@ -238,6 +248,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
|
|
||||||
::mediapipe::Status TfLiteConverterCalculator::Process(CalculatorContext* cc) {
|
::mediapipe::Status TfLiteConverterCalculator::Process(CalculatorContext* cc) {
|
||||||
if (use_gpu_) {
|
if (use_gpu_) {
|
||||||
|
// GpuBuffer to tflite::gpu::GlBuffer conversion.
|
||||||
if (!initialized_) {
|
if (!initialized_) {
|
||||||
MP_RETURN_IF_ERROR(InitGpu(cc));
|
MP_RETURN_IF_ERROR(InitGpu(cc));
|
||||||
initialized_ = true;
|
initialized_ = true;
|
||||||
@@ -253,10 +264,10 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status TfLiteConverterCalculator::Close(CalculatorContext* cc) {
|
::mediapipe::Status TfLiteConverterCalculator::Close(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
gpu_helper_.RunInGlContext([this] { gpu_data_out_.reset(); });
|
gpu_helper_.RunInGlContext([this] { gpu_data_out_.reset(); });
|
||||||
#endif
|
#endif
|
||||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#if defined(MEDIAPIPE_IOS)
|
||||||
gpu_data_out_.reset();
|
gpu_data_out_.reset();
|
||||||
#endif
|
#endif
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
@@ -283,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);
|
||||||
}
|
}
|
||||||
@@ -372,7 +387,7 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
|
|
||||||
::mediapipe::Status TfLiteConverterCalculator::ProcessGPU(
|
::mediapipe::Status TfLiteConverterCalculator::ProcessGPU(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
// GpuBuffer to tflite::gpu::GlBuffer conversion.
|
// GpuBuffer to tflite::gpu::GlBuffer conversion.
|
||||||
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
||||||
MP_RETURN_IF_ERROR(
|
MP_RETURN_IF_ERROR(
|
||||||
@@ -381,17 +396,11 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
auto src = gpu_helper_.CreateSourceTexture(input);
|
auto src = gpu_helper_.CreateSourceTexture(input);
|
||||||
glActiveTexture(GL_TEXTURE0 + 0);
|
glActiveTexture(GL_TEXTURE0 + 0);
|
||||||
glBindTexture(GL_TEXTURE_2D, src.name());
|
glBindTexture(GL_TEXTURE_2D, src.name());
|
||||||
auto status = gpu_data_out_->buffer.BindToIndex(1);
|
RET_CHECK_CALL(gpu_data_out_->buffer.BindToIndex(1));
|
||||||
if (!status.ok()) {
|
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
|
||||||
}
|
|
||||||
const tflite::gpu::uint3 workgroups = {
|
const tflite::gpu::uint3 workgroups = {
|
||||||
NumGroups(input.width(), kWorkgroupSize),
|
NumGroups(input.width(), kWorkgroupSize),
|
||||||
NumGroups(input.height(), kWorkgroupSize), 1};
|
NumGroups(input.height(), kWorkgroupSize), 1};
|
||||||
status = gpu_data_out_->program.Dispatch(workgroups);
|
RET_CHECK_CALL(gpu_data_out_->program.Dispatch(workgroups));
|
||||||
if (!status.ok()) {
|
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
|
||||||
}
|
|
||||||
glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
|
glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
|
||||||
glBindTexture(GL_TEXTURE_2D, 0);
|
glBindTexture(GL_TEXTURE_2D, 0);
|
||||||
src.Release();
|
src.Release();
|
||||||
@@ -400,104 +409,93 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
|
|
||||||
// Copy into outputs.
|
// Copy into outputs.
|
||||||
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
||||||
output_tensors->resize(1);
|
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||||
{
|
[this, &output_tensors]() -> ::mediapipe::Status {
|
||||||
GlBuffer& tensor = output_tensors->at(0);
|
output_tensors->resize(1);
|
||||||
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
{
|
||||||
auto status = CreateReadWriteShaderStorageBuffer<float>(
|
GpuTensor& tensor = output_tensors->at(0);
|
||||||
gpu_data_out_->elements, &tensor);
|
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
|
||||||
if (!status.ok()) {
|
gpu_data_out_->elements, &tensor));
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
RET_CHECK_CALL(CopyBuffer(gpu_data_out_->buffer, tensor));
|
||||||
}
|
}
|
||||||
tflite::gpu::gl::CopyBuffer(gpu_data_out_->buffer, tensor);
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}));
|
||||||
cc->Outputs()
|
cc->Outputs()
|
||||||
.Tag("TENSORS_GPU")
|
.Tag("TENSORS_GPU")
|
||||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // 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.
|
||||||
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
||||||
{
|
output_tensors->resize(1);
|
||||||
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
||||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
output_tensors->at(0) =
|
||||||
command_buffer.label = @"TfLiteConverterCalculatorCopy";
|
[device newBufferWithLength:gpu_data_out_->elements * sizeof(float)
|
||||||
id<MTLBuffer> tensor =
|
options:MTLResourceStorageModeShared];
|
||||||
[device newBufferWithLength:gpu_data_out_->elements * sizeof(float)
|
[MPPMetalUtil blitMetalBufferTo:output_tensors->at(0)
|
||||||
options:MTLResourceStorageModeShared];
|
from:gpu_data_out_->buffer
|
||||||
id<MTLBlitCommandEncoder> blit_command =
|
blocking:false
|
||||||
[command_buffer blitCommandEncoder];
|
commandBuffer:command_buffer];
|
||||||
[blit_command copyFromBuffer:gpu_data_out_->buffer
|
|
||||||
sourceOffset:0
|
|
||||||
toBuffer:tensor
|
|
||||||
destinationOffset:0
|
|
||||||
size:gpu_data_out_->elements * sizeof(float)];
|
|
||||||
[blit_command endEncoding];
|
|
||||||
[command_buffer commit];
|
|
||||||
[command_buffer waitUntilCompleted];
|
|
||||||
|
|
||||||
output_tensors->push_back(tensor);
|
|
||||||
}
|
|
||||||
|
|
||||||
cc->Outputs()
|
cc->Outputs()
|
||||||
.Tag("TENSORS_GPU")
|
.Tag("TENSORS_GPU")
|
||||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||||
#else
|
#else
|
||||||
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
|
RET_CHECK_FAIL() << "GPU processing is not enabled.";
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status TfLiteConverterCalculator::InitGpu(CalculatorContext* cc) {
|
::mediapipe::Status TfLiteConverterCalculator::InitGpu(CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
// Configure inputs.
|
// Get input image sizes.
|
||||||
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
const auto& input = cc->Inputs().Tag("IMAGE_GPU").Get<mediapipe::GpuBuffer>();
|
||||||
mediapipe::ImageFormat::Format format =
|
mediapipe::ImageFormat::Format format =
|
||||||
mediapipe::ImageFormatForGpuBufferFormat(input.format());
|
mediapipe::ImageFormatForGpuBufferFormat(input.format());
|
||||||
gpu_data_out_ = absl::make_unique<GPUData>();
|
gpu_data_out_ = absl::make_unique<GPUData>();
|
||||||
gpu_data_out_->elements = input.height() * input.width() * max_num_channels_;
|
gpu_data_out_->elements = input.height() * input.width() * max_num_channels_;
|
||||||
const bool include_alpha = (max_num_channels_ == 4);
|
const bool include_alpha = (max_num_channels_ == 4);
|
||||||
if (!(format == mediapipe::ImageFormat::SRGB ||
|
const bool single_channel = (max_num_channels_ == 1);
|
||||||
|
if (!(format == mediapipe::ImageFormat::GRAY8 ||
|
||||||
|
format == mediapipe::ImageFormat::SRGB ||
|
||||||
format == mediapipe::ImageFormat::SRGBA))
|
format == mediapipe::ImageFormat::SRGBA))
|
||||||
RET_CHECK_FAIL() << "Unsupported GPU input format.";
|
RET_CHECK_FAIL() << "Unsupported GPU input format.";
|
||||||
if (include_alpha && (format != mediapipe::ImageFormat::SRGBA))
|
if (include_alpha && (format != mediapipe::ImageFormat::SRGBA))
|
||||||
RET_CHECK_FAIL() << "Num input channels is less than desired output.";
|
RET_CHECK_FAIL() << "Num input channels is less than desired output.";
|
||||||
#endif
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
// Device memory.
|
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||||
auto status = ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
|
[this, &include_alpha, &input, &single_channel]() -> ::mediapipe::Status {
|
||||||
gpu_data_out_->elements, &gpu_data_out_->buffer);
|
// Device memory.
|
||||||
if (!status.ok()) {
|
RET_CHECK_CALL(
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
|
||||||
}
|
gpu_data_out_->elements, &gpu_data_out_->buffer));
|
||||||
|
|
||||||
// Shader to convert GL Texture to Shader Storage Buffer Object (SSBO),
|
// Shader to convert GL Texture to Shader Storage Buffer Object (SSBO),
|
||||||
// with normalization to either: [0,1] or [-1,1].
|
// with normalization to either: [0,1] or [-1,1].
|
||||||
const std::string shader_source = absl::Substitute(
|
const std::string shader_source = absl::Substitute(
|
||||||
R"( #version 310 es
|
R"( #version 310 es
|
||||||
layout(local_size_x = $0, local_size_y = $0) in;
|
layout(local_size_x = $0, local_size_y = $0) in;
|
||||||
layout(binding = 0) uniform sampler2D input_texture;
|
layout(binding = 0) uniform sampler2D input_texture;
|
||||||
layout(std430, binding = 1) buffer Output {float elements[];} output_data;
|
layout(std430, binding = 1) buffer Output {float elements[];} output_data;
|
||||||
@@ -505,34 +503,34 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
void main() {
|
void main() {
|
||||||
ivec2 gid = ivec2(gl_GlobalInvocationID.xy);
|
ivec2 gid = ivec2(gl_GlobalInvocationID.xy);
|
||||||
if (gid.x >= width_height.x || gid.y >= width_height.y) return;
|
if (gid.x >= width_height.x || gid.y >= width_height.y) return;
|
||||||
$5 // pixel fetch
|
vec4 pixel = texelFetch(input_texture, gid, 0);
|
||||||
$3 // normalize [-1,1]
|
$3 // normalize [-1,1]
|
||||||
int linear_index = $7 * ($4 * width_height.x + gid.x);
|
int linear_index = $7 * ($4 * width_height.x + gid.x);
|
||||||
output_data.elements[linear_index + 0] = pixel.x;
|
output_data.elements[linear_index + 0] = pixel.x; // r channel
|
||||||
output_data.elements[linear_index + 1] = pixel.y;
|
$5 // g & b channels
|
||||||
output_data.elements[linear_index + 2] = pixel.z;
|
|
||||||
$6 // alpha channel
|
$6 // alpha channel
|
||||||
})",
|
})",
|
||||||
/*$0=*/kWorkgroupSize, /*$1=*/input.width(), /*$2=*/input.height(),
|
/*$0=*/kWorkgroupSize, /*$1=*/input.width(), /*$2=*/input.height(),
|
||||||
/*$3=*/zero_center_ ? "pixel = (pixel - 0.5) * 2.0;" : "",
|
/*$3=*/zero_center_ ? "pixel = (pixel - 0.5) * 2.0;" : "",
|
||||||
/*$4=*/flip_vertically_ ? "(width_height.y - 1 - gid.y)" : "gid.y",
|
/*$4=*/flip_vertically_ ? "(width_height.y - 1 - gid.y)" : "gid.y",
|
||||||
/*$5=*/
|
/*$5=*/
|
||||||
include_alpha ? "vec4 pixel = texelFetch(input_texture, gid, 0);"
|
single_channel
|
||||||
: "vec3 pixel = texelFetch(input_texture, gid, 0).xyz;",
|
? ""
|
||||||
/*$6=*/
|
: R"(output_data.elements[linear_index + 1] = pixel.y;
|
||||||
include_alpha ? "output_data.elements[linear_index + 3] = pixel.w;" : "",
|
output_data.elements[linear_index + 2] = pixel.z;)",
|
||||||
/*$7=*/include_alpha ? 4 : 3);
|
/*$6=*/
|
||||||
status = GlShader::CompileShader(GL_COMPUTE_SHADER, shader_source,
|
include_alpha ? "output_data.elements[linear_index + 3] = pixel.w;"
|
||||||
&gpu_data_out_->shader);
|
: "",
|
||||||
if (!status.ok()) {
|
/*$7=*/max_num_channels_);
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
RET_CHECK_CALL(GlShader::CompileShader(GL_COMPUTE_SHADER, shader_source,
|
||||||
}
|
&gpu_data_out_->shader));
|
||||||
status = GlProgram::CreateWithShader(gpu_data_out_->shader,
|
RET_CHECK_CALL(GlProgram::CreateWithShader(gpu_data_out_->shader,
|
||||||
&gpu_data_out_->program);
|
&gpu_data_out_->program));
|
||||||
if (!status.ok()) {
|
return ::mediapipe::OkStatus();
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
}));
|
||||||
}
|
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
|
|
||||||
RET_CHECK(include_alpha)
|
RET_CHECK(include_alpha)
|
||||||
<< "iOS GPU inference currently accepts only RGBA input.";
|
<< "iOS GPU inference currently accepts only RGBA input.";
|
||||||
|
|
||||||
@@ -546,8 +544,6 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
// with normalization to either: [0,1] or [-1,1].
|
// with normalization to either: [0,1] or [-1,1].
|
||||||
const std::string shader_source = absl::Substitute(
|
const std::string shader_source = absl::Substitute(
|
||||||
R"(
|
R"(
|
||||||
#include <simd/simd.h>
|
|
||||||
|
|
||||||
#include <metal_stdlib>
|
#include <metal_stdlib>
|
||||||
|
|
||||||
using namespace metal;
|
using namespace metal;
|
||||||
@@ -612,10 +608,10 @@ REGISTER_CALCULATOR(TfLiteConverterCalculator);
|
|||||||
|
|
||||||
// Get desired way to handle input channels.
|
// Get desired way to handle input channels.
|
||||||
max_num_channels_ = options.max_num_channels();
|
max_num_channels_ = options.max_num_channels();
|
||||||
// Currently only alpha channel toggling is suppored.
|
CHECK_GE(max_num_channels_, 1);
|
||||||
CHECK_GE(max_num_channels_, 3);
|
|
||||||
CHECK_LE(max_num_channels_, 4);
|
CHECK_LE(max_num_channels_, 4);
|
||||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
CHECK_NE(max_num_channels_, 2);
|
||||||
|
#if defined(MEDIAPIPE_IOS)
|
||||||
if (cc->Inputs().HasTag("IMAGE_GPU"))
|
if (cc->Inputs().HasTag("IMAGE_GPU"))
|
||||||
// Currently on iOS, tflite gpu input tensor must be 4 channels,
|
// Currently on iOS, tflite gpu input tensor must be 4 channels,
|
||||||
// so input image must be 4 channels also (checked in InitGpu).
|
// so input image must be 4 channels also (checked in InitGpu).
|
||||||
|
|||||||
@@ -36,8 +36,7 @@ message TfLiteConverterCalculatorOptions {
|
|||||||
optional bool flip_vertically = 2 [default = false];
|
optional bool flip_vertically = 2 [default = false];
|
||||||
|
|
||||||
// Controls how many channels of the input image get passed through to the
|
// Controls how many channels of the input image get passed through to the
|
||||||
// tensor. Currently this only controls whether or not to ignore alpha
|
// tensor. Valid values are 1,3,4 only. Ignored for iOS GPU.
|
||||||
// channel, so it must be 3 or 4.
|
|
||||||
optional int32 max_num_channels = 3 [default = 3];
|
optional int32 max_num_channels = 3 [default = 3];
|
||||||
|
|
||||||
// The calculator expects Matrix inputs to be in column-major order. Set
|
// The calculator expects Matrix inputs to be in column-major order. Set
|
||||||
|
|||||||
@@ -12,10 +12,13 @@
|
|||||||
// 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 <cstring>
|
||||||
|
#include <memory>
|
||||||
#include <string>
|
#include <string>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
#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/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/util/resource_util.h"
|
#include "mediapipe/util/resource_util.h"
|
||||||
@@ -24,48 +27,89 @@
|
|||||||
#include "tensorflow/lite/kernels/register.h"
|
#include "tensorflow/lite/kernels/register.h"
|
||||||
#include "tensorflow/lite/model.h"
|
#include "tensorflow/lite/model.h"
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||||
#include "mediapipe/gpu/gpu_buffer.h"
|
#include "mediapipe/gpu/gpu_buffer.h"
|
||||||
|
#include "tensorflow/lite/delegates/gpu/common/shape.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
|
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
|
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
|
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
|
||||||
#endif // __ANDROID__
|
#endif // !MEDIAPIPE_DISABLE_GL_COMPUTE
|
||||||
|
|
||||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#if defined(MEDIAPIPE_IOS)
|
||||||
#import <CoreVideo/CoreVideo.h>
|
#import <CoreVideo/CoreVideo.h>
|
||||||
#import <Metal/Metal.h>
|
#import <Metal/Metal.h>
|
||||||
#import <MetalKit/MetalKit.h>
|
#import <MetalKit/MetalKit.h>
|
||||||
|
|
||||||
#import "mediapipe/gpu/MPPMetalHelper.h"
|
#import "mediapipe/gpu/MPPMetalHelper.h"
|
||||||
|
#include "mediapipe/gpu/MPPMetalUtil.h"
|
||||||
|
#include "mediapipe/gpu/gpu_buffer.h"
|
||||||
|
#include "tensorflow/lite/delegates/gpu/common/shape.h"
|
||||||
|
#include "tensorflow/lite/delegates/gpu/metal/buffer_convert.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
|
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
|
||||||
|
#include "tensorflow/lite/delegates/gpu/metal_delegate_internal.h"
|
||||||
#endif // iOS
|
#endif // iOS
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if defined(MEDIAPIPE_ANDROID)
|
||||||
|
#include "tensorflow/lite/delegates/nnapi/nnapi_delegate.h"
|
||||||
|
#endif // ANDROID
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
// Commonly used to compute the number of blocks to launch in a kernel.
|
||||||
|
int NumGroups(const int size, const int group_size) { // NOLINT
|
||||||
|
return (size + group_size - 1) / group_size;
|
||||||
|
}
|
||||||
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
|
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
typedef id<MTLBuffer> GpuTensor;
|
typedef id<MTLBuffer> GpuTensor;
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
|
// Round up n to next multiple of m.
|
||||||
|
size_t RoundUp(size_t n, size_t m) { return ((n + m - 1) / m) * m; } // NOLINT
|
||||||
|
} // namespace
|
||||||
|
|
||||||
|
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||||
|
#include "edgetpu.h"
|
||||||
|
|
||||||
|
// Creates and returns an Edge TPU interpreter to run the given edgetpu model.
|
||||||
|
std::unique_ptr<tflite::Interpreter> BuildEdgeTpuInterpreter(
|
||||||
|
const tflite::FlatBufferModel& model,
|
||||||
|
tflite::ops::builtin::BuiltinOpResolver* resolver,
|
||||||
|
edgetpu::EdgeTpuContext* edgetpu_context) {
|
||||||
|
resolver->AddCustom(edgetpu::kCustomOp, edgetpu::RegisterCustomOp());
|
||||||
|
std::unique_ptr<tflite::Interpreter> interpreter;
|
||||||
|
if (tflite::InterpreterBuilder(model, *resolver)(&interpreter) != kTfLiteOk) {
|
||||||
|
std::cerr << "Failed to build edge TPU interpreter." << std::endl;
|
||||||
|
}
|
||||||
|
interpreter->SetExternalContext(kTfLiteEdgeTpuContext, edgetpu_context);
|
||||||
|
interpreter->SetNumThreads(1);
|
||||||
|
if (interpreter->AllocateTensors() != kTfLiteOk) {
|
||||||
|
std::cerr << "Failed to allocate edge TPU tensors." << std::endl;
|
||||||
|
}
|
||||||
|
return interpreter;
|
||||||
|
}
|
||||||
|
#endif // MEDIAPIPE_EDGE_TPU
|
||||||
|
|
||||||
// TfLiteInferenceCalculator File Layout:
|
// TfLiteInferenceCalculator File Layout:
|
||||||
// * Header
|
// * Header
|
||||||
// * Core
|
// * Core
|
||||||
// * Aux
|
// * Aux
|
||||||
namespace mediapipe {
|
namespace mediapipe {
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
|
using ::tflite::gpu::gl::CopyBuffer;
|
||||||
|
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
||||||
using ::tflite::gpu::gl::GlBuffer;
|
using ::tflite::gpu::gl::GlBuffer;
|
||||||
using ::tflite::gpu::gl::GlProgram;
|
#endif
|
||||||
using ::tflite::gpu::gl::GlShader;
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
struct GPUData {
|
struct GPUData {
|
||||||
int elements = 1;
|
int elements = 1;
|
||||||
GlBuffer buffer;
|
GpuTensor buffer;
|
||||||
};
|
::tflite::gpu::BHWC shape;
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
|
||||||
struct GPUData {
|
|
||||||
int elements = 1;
|
|
||||||
id<MTLBuffer> buffer;
|
|
||||||
};
|
};
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
@@ -104,7 +148,7 @@ struct GPUData {
|
|||||||
// options: {
|
// options: {
|
||||||
// [mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
// [mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
||||||
// model_path: "modelname.tflite"
|
// model_path: "modelname.tflite"
|
||||||
// use_gpu: true
|
// delegate { gpu {} }
|
||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
// }
|
// }
|
||||||
@@ -119,6 +163,9 @@ struct GPUData {
|
|||||||
//
|
//
|
||||||
class TfLiteInferenceCalculator : public CalculatorBase {
|
class TfLiteInferenceCalculator : public CalculatorBase {
|
||||||
public:
|
public:
|
||||||
|
using TfLiteDelegatePtr =
|
||||||
|
std::unique_ptr<TfLiteDelegate, std::function<void(TfLiteDelegate*)>>;
|
||||||
|
|
||||||
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
static ::mediapipe::Status GetContract(CalculatorContract* cc);
|
||||||
|
|
||||||
::mediapipe::Status Open(CalculatorContext* cc) override;
|
::mediapipe::Status Open(CalculatorContext* cc) override;
|
||||||
@@ -132,16 +179,23 @@ class TfLiteInferenceCalculator : public CalculatorBase {
|
|||||||
|
|
||||||
std::unique_ptr<tflite::Interpreter> interpreter_;
|
std::unique_ptr<tflite::Interpreter> interpreter_;
|
||||||
std::unique_ptr<tflite::FlatBufferModel> model_;
|
std::unique_ptr<tflite::FlatBufferModel> model_;
|
||||||
TfLiteDelegate* delegate_ = nullptr;
|
TfLiteDelegatePtr delegate_;
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#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(__APPLE__) && !TARGET_OS_OSX // 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;
|
||||||
|
#endif
|
||||||
|
|
||||||
|
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||||
|
std::shared_ptr<edgetpu::EdgeTpuContext> edgetpu_context_ =
|
||||||
|
edgetpu::EdgeTpuManager::GetSingleton()->OpenDevice();
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
std::string model_path_ = "";
|
std::string model_path_ = "";
|
||||||
@@ -161,19 +215,34 @@ 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"));
|
||||||
|
|
||||||
|
const auto& options =
|
||||||
|
cc->Options<::mediapipe::TfLiteInferenceCalculatorOptions>();
|
||||||
|
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(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#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>>();
|
||||||
#endif
|
use_gpu |= true;
|
||||||
|
}
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
if (cc->Outputs().HasTag("TENSORS"))
|
if (cc->Outputs().HasTag("TENSORS"))
|
||||||
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
cc->Outputs().Tag("TENSORS").Set<std::vector<TfLiteTensor>>();
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#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>>();
|
||||||
#endif
|
use_gpu |= true;
|
||||||
|
}
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
|
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
|
||||||
cc->InputSidePackets()
|
cc->InputSidePackets()
|
||||||
@@ -181,11 +250,13 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
.Set<tflite::ops::builtin::BuiltinOpResolver>();
|
.Set<tflite::ops::builtin::BuiltinOpResolver>();
|
||||||
}
|
}
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
if (use_gpu) {
|
||||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||||
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
|
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
|
||||||
#endif
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
// Assign this calculator's default InputStreamHandler.
|
// Assign this calculator's default InputStreamHandler.
|
||||||
cc->SetInputStreamHandler("FixedSizeInputStreamHandler");
|
cc->SetInputStreamHandler("FixedSizeInputStreamHandler");
|
||||||
@@ -199,37 +270,47 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
MP_RETURN_IF_ERROR(LoadOptions(cc));
|
MP_RETURN_IF_ERROR(LoadOptions(cc));
|
||||||
|
|
||||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
gpu_input_ = true;
|
gpu_input_ = true;
|
||||||
gpu_inference_ = true; // Inference must be on GPU also.
|
gpu_inference_ = true; // Inference must be on GPU also.
|
||||||
#else
|
#else
|
||||||
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
|
RET_CHECK(!cc->Inputs().HasTag("TENSORS_GPU"))
|
||||||
#endif
|
<< "GPU processing not enabled.";
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
if (cc->Outputs().HasTag("TENSORS_GPU")) {
|
if (cc->Outputs().HasTag("TENSORS_GPU")) {
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
gpu_output_ = true;
|
gpu_output_ = true;
|
||||||
RET_CHECK(cc->Inputs().HasTag("TENSORS_GPU"))
|
RET_CHECK(cc->Inputs().HasTag("TENSORS_GPU"))
|
||||||
<< "GPU output must also have GPU Input.";
|
<< "GPU output must also have GPU Input.";
|
||||||
#else
|
#else
|
||||||
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
|
RET_CHECK(!cc->Inputs().HasTag("TENSORS_GPU"))
|
||||||
#endif
|
<< "GPU processing not enabled.";
|
||||||
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
}
|
}
|
||||||
|
|
||||||
MP_RETURN_IF_ERROR(LoadModel(cc));
|
MP_RETURN_IF_ERROR(LoadModel(cc));
|
||||||
|
|
||||||
if (gpu_inference_) {
|
if (gpu_inference_) {
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
|
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
|
||||||
RET_CHECK(gpu_helper_);
|
RET_CHECK(gpu_helper_);
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
|
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||||
|
[this, &cc]() -> ::mediapipe::Status { return LoadDelegate(cc); }));
|
||||||
|
#else
|
||||||
MP_RETURN_IF_ERROR(LoadDelegate(cc));
|
MP_RETURN_IF_ERROR(LoadDelegate(cc));
|
||||||
|
#endif
|
||||||
|
} else {
|
||||||
|
#if defined(__EMSCRIPTEN__) || defined(MEDIAPIPE_ANDROID)
|
||||||
|
MP_RETURN_IF_ERROR(LoadDelegate(cc));
|
||||||
|
#endif // __EMSCRIPTEN__ || ANDROID
|
||||||
}
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -237,35 +318,46 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
// 1. Receive pre-processed tensor inputs.
|
// 1. Receive pre-processed tensor inputs.
|
||||||
if (gpu_input_) {
|
if (gpu_input_) {
|
||||||
// Read GPU input into SSBO.
|
// Read GPU input into SSBO.
|
||||||
#if defined(__ANDROID__)
|
#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.
|
||||||
tflite::gpu::gl::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(__APPLE__) && !TARGET_OS_OSX // 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 with conversion float 32 bits to 16 bits.
|
||||||
|
gpu_data_in_.resize(input_tensors.size());
|
||||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||||
command_buffer.label = @"TfLiteInferenceCalculatorInput";
|
command_buffer.label = @"TfLiteInferenceCalculatorConvert";
|
||||||
id<MTLBlitCommandEncoder> blit_command =
|
id<MTLComputeCommandEncoder> compute_encoder =
|
||||||
[command_buffer blitCommandEncoder];
|
[command_buffer computeCommandEncoder];
|
||||||
// Explicit copy input.
|
[compute_encoder setComputePipelineState:fp32_to_fp16_program_];
|
||||||
[blit_command copyFromBuffer:input_tensors[0]
|
for (int i = 0; i < input_tensors.size(); ++i) {
|
||||||
sourceOffset:0
|
[compute_encoder setBuffer:input_tensors[i] offset:0 atIndex:0];
|
||||||
toBuffer:gpu_data_in_->buffer
|
[compute_encoder setBuffer:gpu_data_in_[i]->buffer offset:0 atIndex:1];
|
||||||
destinationOffset:0
|
constexpr int kWorkgroupSize = 64; // Block size for GPU shader.
|
||||||
size:gpu_data_in_->elements * sizeof(float)];
|
MTLSize threads_per_group = MTLSizeMake(kWorkgroupSize, 1, 1);
|
||||||
[blit_command endEncoding];
|
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];
|
[command_buffer commit];
|
||||||
[command_buffer waitUntilCompleted];
|
|
||||||
#else
|
#else
|
||||||
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
|
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
||||||
#endif
|
#endif
|
||||||
} else {
|
} else {
|
||||||
// Read CPU input into tensors.
|
// Read CPU input into tensors.
|
||||||
@@ -278,24 +370,26 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
if (use_quantized_tensors_) {
|
if (use_quantized_tensors_) {
|
||||||
const uint8* input_tensor_buffer = input_tensor->data.uint8;
|
const uint8* input_tensor_buffer = input_tensor->data.uint8;
|
||||||
uint8* local_tensor_buffer = interpreter_->typed_input_tensor<uint8>(i);
|
uint8* local_tensor_buffer = interpreter_->typed_input_tensor<uint8>(i);
|
||||||
memcpy(local_tensor_buffer, input_tensor_buffer, input_tensor->bytes);
|
std::memcpy(local_tensor_buffer, input_tensor_buffer,
|
||||||
|
input_tensor->bytes);
|
||||||
} else {
|
} else {
|
||||||
const float* input_tensor_buffer = input_tensor->data.f;
|
const float* input_tensor_buffer = input_tensor->data.f;
|
||||||
float* local_tensor_buffer = interpreter_->typed_input_tensor<float>(i);
|
float* local_tensor_buffer = interpreter_->typed_input_tensor<float>(i);
|
||||||
memcpy(local_tensor_buffer, input_tensor_buffer, input_tensor->bytes);
|
std::memcpy(local_tensor_buffer, input_tensor_buffer,
|
||||||
|
input_tensor->bytes);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// 2. Run inference.
|
// 2. Run inference.
|
||||||
if (gpu_inference_) {
|
if (gpu_inference_) {
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
MP_RETURN_IF_ERROR(
|
MP_RETURN_IF_ERROR(
|
||||||
gpu_helper_.RunInGlContext([this]() -> ::mediapipe::Status {
|
gpu_helper_.RunInGlContext([this]() -> ::mediapipe::Status {
|
||||||
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
|
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}));
|
}));
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
|
RET_CHECK_EQ(interpreter_->Invoke(), kTfLiteOk);
|
||||||
#endif
|
#endif
|
||||||
} else {
|
} else {
|
||||||
@@ -304,52 +398,50 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
|
|
||||||
// 3. Output processed tensors.
|
// 3. Output processed tensors.
|
||||||
if (gpu_output_) {
|
if (gpu_output_) {
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
// Output result tensors (GPU).
|
// Output result tensors (GPU).
|
||||||
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
||||||
output_tensors->resize(gpu_data_out_.size());
|
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
||||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
[this, &output_tensors]() -> ::mediapipe::Status {
|
||||||
GlBuffer& tensor = output_tensors->at(i);
|
output_tensors->resize(gpu_data_out_.size());
|
||||||
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||||
auto status = CreateReadWriteShaderStorageBuffer<float>(
|
GpuTensor& tensor = output_tensors->at(i);
|
||||||
gpu_data_out_[i]->elements, &tensor);
|
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
|
||||||
if (!status.ok()) {
|
gpu_data_out_[i]->elements, &tensor));
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
RET_CHECK_CALL(CopyBuffer(gpu_data_out_[i]->buffer, tensor));
|
||||||
}
|
}
|
||||||
tflite::gpu::gl::CopyBuffer(gpu_data_out_[i]->buffer, tensor);
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}));
|
||||||
cc->Outputs()
|
cc->Outputs()
|
||||||
.Tag("TENSORS_GPU")
|
.Tag("TENSORS_GPU")
|
||||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
// Output result tensors (GPU).
|
// Output result tensors (GPU).
|
||||||
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
auto output_tensors = absl::make_unique<std::vector<GpuTensor>>();
|
||||||
|
output_tensors->resize(gpu_data_out_.size());
|
||||||
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
||||||
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||||
command_buffer.label = @"TfLiteInferenceCalculatorOutput";
|
command_buffer.label = @"TfLiteInferenceBPHWC4Convert";
|
||||||
|
id<MTLComputeCommandEncoder> convert_command =
|
||||||
|
[command_buffer computeCommandEncoder];
|
||||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||||
id<MTLBuffer> tensor =
|
output_tensors->at(i) =
|
||||||
[device newBufferWithLength:gpu_data_out_[i]->elements * sizeof(float)
|
[device newBufferWithLength:gpu_data_out_[i]->elements * sizeof(float)
|
||||||
options:MTLResourceStorageModeShared];
|
options:MTLResourceStorageModeShared];
|
||||||
id<MTLBlitCommandEncoder> blit_command =
|
// Reshape tensor.
|
||||||
[command_buffer blitCommandEncoder];
|
[converter_from_BPHWC4_ convertWithEncoder:convert_command
|
||||||
// Explicit copy input.
|
shape:gpu_data_out_[i]->shape
|
||||||
[blit_command copyFromBuffer:gpu_data_out_[i]->buffer
|
sourceBuffer:gpu_data_out_[i]->buffer
|
||||||
sourceOffset:0
|
convertedBuffer:output_tensors->at(i)];
|
||||||
toBuffer:tensor
|
|
||||||
destinationOffset:0
|
|
||||||
size:gpu_data_out_[i]->elements * sizeof(float)];
|
|
||||||
[blit_command endEncoding];
|
|
||||||
[command_buffer commit];
|
|
||||||
[command_buffer waitUntilCompleted];
|
|
||||||
output_tensors->push_back(tensor);
|
|
||||||
}
|
}
|
||||||
|
[convert_command endEncoding];
|
||||||
|
[command_buffer commit];
|
||||||
cc->Outputs()
|
cc->Outputs()
|
||||||
.Tag("TENSORS_GPU")
|
.Tag("TENSORS_GPU")
|
||||||
.Add(output_tensors.release(), cc->InputTimestamp());
|
.Add(output_tensors.release(), cc->InputTimestamp());
|
||||||
#else
|
#else
|
||||||
RET_CHECK_FAIL() << "GPU processing is for Android and iOS only.";
|
RET_CHECK_FAIL() << "GPU processing not enabled.";
|
||||||
#endif
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
} else {
|
} else {
|
||||||
// Output result tensors (CPU).
|
// Output result tensors (CPU).
|
||||||
const auto& tensor_indexes = interpreter_->outputs();
|
const auto& tensor_indexes = interpreter_->outputs();
|
||||||
@@ -367,24 +459,34 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
|
|
||||||
::mediapipe::Status TfLiteInferenceCalculator::Close(CalculatorContext* cc) {
|
::mediapipe::Status TfLiteInferenceCalculator::Close(CalculatorContext* cc) {
|
||||||
if (delegate_) {
|
if (delegate_) {
|
||||||
#if defined(__ANDROID__)
|
if (gpu_inference_) {
|
||||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
TfLiteGpuDelegateDelete(delegate_);
|
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]() -> Status {
|
||||||
gpu_data_in_.reset();
|
delegate_ = nullptr;
|
||||||
|
for (int i = 0; i < gpu_data_in_.size(); ++i) {
|
||||||
|
gpu_data_in_[i].reset();
|
||||||
|
}
|
||||||
|
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
||||||
|
gpu_data_out_[i].reset();
|
||||||
|
}
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}));
|
||||||
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
|
delegate_ = nullptr;
|
||||||
|
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();
|
|
||||||
}));
|
|
||||||
#elif defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
|
||||||
TFLGpuDelegateDelete(delegate_);
|
|
||||||
gpu_data_in_.reset();
|
|
||||||
for (int i = 0; i < gpu_data_out_.size(); ++i) {
|
|
||||||
gpu_data_out_[i].reset();
|
|
||||||
}
|
|
||||||
#endif
|
#endif
|
||||||
delegate_ = nullptr;
|
} else {
|
||||||
|
delegate_ = nullptr;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||||
|
edgetpu_context_.reset();
|
||||||
|
#endif
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -398,7 +500,7 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
|
|
||||||
// Get model name.
|
// Get model name.
|
||||||
if (!options.model_path().empty()) {
|
if (!options.model_path().empty()) {
|
||||||
auto model_path = options.model_path();
|
std::string model_path = options.model_path();
|
||||||
|
|
||||||
ASSIGN_OR_RETURN(model_path_, mediapipe::PathToResourceAsFile(model_path));
|
ASSIGN_OR_RETURN(model_path_, mediapipe::PathToResourceAsFile(model_path));
|
||||||
} else {
|
} else {
|
||||||
@@ -408,7 +510,8 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Get execution modes.
|
// Get execution modes.
|
||||||
gpu_inference_ = options.use_gpu();
|
gpu_inference_ =
|
||||||
|
options.has_delegate() ? options.delegate().has_gpu() : options.use_gpu();
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -418,19 +521,29 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
model_ = tflite::FlatBufferModel::BuildFromFile(model_path_.c_str());
|
model_ = tflite::FlatBufferModel::BuildFromFile(model_path_.c_str());
|
||||||
RET_CHECK(model_);
|
RET_CHECK(model_);
|
||||||
|
|
||||||
|
tflite::ops::builtin::BuiltinOpResolver op_resolver;
|
||||||
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
|
if (cc->InputSidePackets().HasTag("CUSTOM_OP_RESOLVER")) {
|
||||||
const auto& op_resolver =
|
op_resolver = cc->InputSidePackets()
|
||||||
cc->InputSidePackets()
|
.Tag("CUSTOM_OP_RESOLVER")
|
||||||
.Tag("CUSTOM_OP_RESOLVER")
|
.Get<tflite::ops::builtin::BuiltinOpResolver>();
|
||||||
.Get<tflite::ops::builtin::BuiltinOpResolver>();
|
|
||||||
tflite::InterpreterBuilder(*model_, op_resolver)(&interpreter_);
|
|
||||||
} else {
|
|
||||||
const tflite::ops::builtin::BuiltinOpResolver op_resolver;
|
|
||||||
tflite::InterpreterBuilder(*model_, op_resolver)(&interpreter_);
|
|
||||||
}
|
}
|
||||||
|
#if defined(MEDIAPIPE_EDGE_TPU)
|
||||||
|
interpreter_ =
|
||||||
|
BuildEdgeTpuInterpreter(*model_, &op_resolver, edgetpu_context_.get());
|
||||||
|
#else
|
||||||
|
tflite::InterpreterBuilder(*model_, op_resolver)(&interpreter_);
|
||||||
|
#endif // MEDIAPIPE_EDGE_TPU
|
||||||
|
|
||||||
RET_CHECK(interpreter_);
|
RET_CHECK(interpreter_);
|
||||||
|
|
||||||
|
#if defined(__EMSCRIPTEN__) || defined(MEDIAPIPE_EDGE_TPU)
|
||||||
|
interpreter_->SetNumThreads(1);
|
||||||
|
#else
|
||||||
|
interpreter_->SetNumThreads(
|
||||||
|
cc->Options<mediapipe::TfLiteInferenceCalculatorOptions>()
|
||||||
|
.cpu_num_thread());
|
||||||
|
#endif // __EMSCRIPTEN__
|
||||||
|
|
||||||
if (gpu_output_) {
|
if (gpu_output_) {
|
||||||
use_quantized_tensors_ = false;
|
use_quantized_tensors_ = false;
|
||||||
} else {
|
} else {
|
||||||
@@ -446,7 +559,39 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
|
|
||||||
::mediapipe::Status TfLiteInferenceCalculator::LoadDelegate(
|
::mediapipe::Status TfLiteInferenceCalculator::LoadDelegate(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
#if defined(__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 defined(MEDIAPIPE_ANDROID)
|
||||||
|
const bool nnapi_requested = calculator_opts.has_delegate()
|
||||||
|
? calculator_opts.delegate().has_nnapi()
|
||||||
|
: calculator_opts.use_nnapi();
|
||||||
|
if (nnapi_requested) {
|
||||||
|
// Attempt to use NNAPI.
|
||||||
|
// If not supported, the default CPU delegate will be created and used.
|
||||||
|
interpreter_->SetAllowFp16PrecisionForFp32(1);
|
||||||
|
delegate_ =
|
||||||
|
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
|
||||||
|
|
||||||
|
// Return, no need for GPU delegate below.
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}
|
||||||
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
// Configure and create the delegate.
|
// Configure and create the delegate.
|
||||||
TfLiteGpuDelegateOptions options = TfLiteGpuDelegateOptionsDefault();
|
TfLiteGpuDelegateOptions options = TfLiteGpuDelegateOptionsDefault();
|
||||||
options.compile_options.precision_loss_allowed = 1;
|
options.compile_options.precision_loss_allowed = 1;
|
||||||
@@ -454,31 +599,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[0]);
|
||||||
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);
|
||||||
}
|
}
|
||||||
// Input to model can be either RGB/RGBA only.
|
|
||||||
RET_CHECK_GE(tensor->dims->data[3], 3);
|
|
||||||
RET_CHECK_LE(tensor->dims->data[3], 4);
|
|
||||||
// Create and bind input buffer.
|
|
||||||
auto status = ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer<float>(
|
|
||||||
gpu_data_in_->elements, &gpu_data_in_->buffer);
|
|
||||||
if (!status.ok()) {
|
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
|
||||||
}
|
|
||||||
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.
|
||||||
@@ -496,60 +640,87 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
// Create and bind output buffers.
|
// Create and bind output buffers.
|
||||||
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) {
|
||||||
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
|
||||||
auto status = 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(TfLiteGpuDelegateBindBufferToTensor(
|
||||||
if (!status.ok()) {
|
delegate_.get(), gpu_data_out_[i]->buffer.id(),
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
output_indices[i]),
|
||||||
}
|
kTfLiteOk);
|
||||||
RET_CHECK_EQ(
|
|
||||||
TfLiteGpuDelegateBindBufferToTensor(
|
|
||||||
delegate_, gpu_data_out_[i]->buffer.id(), output_indices[i]),
|
|
||||||
kTfLiteOk);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Must call this last.
|
// Must call this last.
|
||||||
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_), kTfLiteOk);
|
RET_CHECK_EQ(interpreter_->ModifyGraphWithDelegate(delegate_.get()),
|
||||||
#endif // __ANDROID__
|
kTfLiteOk);
|
||||||
|
#endif // OpenGL
|
||||||
|
|
||||||
#if defined(__APPLE__) && !TARGET_OS_OSX // iOS
|
#if defined(MEDIAPIPE_IOS)
|
||||||
|
const int kHalfSize = 2; // sizeof(half)
|
||||||
// Configure and create the delegate.
|
// Configure and create the delegate.
|
||||||
GpuDelegateOptions options;
|
TFLGpuDelegateOptions options;
|
||||||
options.allow_precision_loss = false; // Must match converter, F=float/T=half
|
options.allow_precision_loss = true;
|
||||||
options.wait_type = GpuDelegateOptions::WaitType::kActive;
|
options.wait_type = TFLGpuDelegateWaitType::TFLGpuDelegateWaitTypePassive;
|
||||||
if (!delegate_) delegate_ = TFLGpuDelegateCreate(&options);
|
if (!delegate_)
|
||||||
|
delegate_ = TfLiteDelegatePtr(TFLGpuDelegateCreate(&options),
|
||||||
|
&TFLGpuDelegateDelete);
|
||||||
|
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.
|
|
||||||
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
|
||||||
gpu_data_in_->buffer =
|
|
||||||
[device newBufferWithLength:gpu_data_in_->elements * sizeof(float)
|
|
||||||
options:MTLResourceStorageModeShared];
|
|
||||||
// Must call this before TFLGpuDelegateBindMetalBufferToTensor.
|
|
||||||
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.
|
||||||
@@ -561,19 +732,50 @@ REGISTER_CALCULATOR(TfLiteInferenceCalculator);
|
|||||||
gpu_data_out_[i]->elements = 1;
|
gpu_data_out_[i]->elements = 1;
|
||||||
// TODO handle *2 properly on some dialated models
|
// TODO handle *2 properly on some dialated models
|
||||||
for (int d = 0; d < tensor->dims->size; ++d) {
|
for (int d = 0; d < tensor->dims->size; ++d) {
|
||||||
gpu_data_out_[i]->elements *= tensor->dims->data[d];
|
// Pad each dim for BHWC4 conversion inside delegate.
|
||||||
|
gpu_data_out_[i]->elements *= RoundUp(tensor->dims->data[d], 4);
|
||||||
|
}
|
||||||
|
// Save dimensions for reshaping back later.
|
||||||
|
gpu_data_out_[i]->shape.b = tensor->dims->data[0];
|
||||||
|
switch (tensor->dims->size) {
|
||||||
|
case 2:
|
||||||
|
gpu_data_out_[i]->shape.h = 1;
|
||||||
|
gpu_data_out_[i]->shape.w = 1;
|
||||||
|
gpu_data_out_[i]->shape.c = tensor->dims->data[1];
|
||||||
|
break;
|
||||||
|
case 3:
|
||||||
|
gpu_data_out_[i]->shape.h = 1;
|
||||||
|
gpu_data_out_[i]->shape.w = tensor->dims->data[1];
|
||||||
|
gpu_data_out_[i]->shape.c = tensor->dims->data[2];
|
||||||
|
break;
|
||||||
|
case 4:
|
||||||
|
gpu_data_out_[i]->shape.h = tensor->dims->data[1];
|
||||||
|
gpu_data_out_[i]->shape.w = tensor->dims->data[2];
|
||||||
|
gpu_data_out_[i]->shape.c = tensor->dims->data[3];
|
||||||
|
break;
|
||||||
|
default:
|
||||||
|
return mediapipe::InternalError("Unsupported tensor shape.");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
// Create and bind output buffers.
|
// Create and bind output buffers.
|
||||||
interpreter_->SetAllowBufferHandleOutput(true);
|
interpreter_->SetAllowBufferHandleOutput(true);
|
||||||
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
|
||||||
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.
|
||||||
|
converter_from_BPHWC4_ = [[TFLBufferConvert alloc] initWithDevice:device
|
||||||
|
isFloat16:true
|
||||||
|
convertToPBHWC4:false];
|
||||||
|
if (converter_from_BPHWC4_ == nil) {
|
||||||
|
return mediapipe::InternalError(
|
||||||
|
"Error initializating output buffer converter");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
#endif // iOS
|
#endif // iOS
|
||||||
|
|||||||
@@ -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,22 @@ 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 {}
|
||||||
|
|
||||||
|
oneof delegate {
|
||||||
|
TfLite tflite = 1;
|
||||||
|
Gpu gpu = 2;
|
||||||
|
Nnapi nnapi = 3;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// 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,5 +60,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.
|
||||||
|
// If NNAPI is not available, then the default CPU delegate will be used
|
||||||
|
// automatically.
|
||||||
|
// 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(absl::string_view delegate) {
|
||||||
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;
|
||||||
@@ -73,23 +69,24 @@ TEST_F(TfLiteInferenceCalculatorTest, SmokeTest) {
|
|||||||
auto input_vec = absl::make_unique<std::vector<TfLiteTensor>>();
|
auto input_vec = absl::make_unique<std::vector<TfLiteTensor>>();
|
||||||
input_vec->emplace_back(*tensor);
|
input_vec->emplace_back(*tensor);
|
||||||
|
|
||||||
|
std::string graph_proto = R"(
|
||||||
|
input_stream: "tensor_in"
|
||||||
|
node {
|
||||||
|
calculator: "TfLiteInferenceCalculator"
|
||||||
|
input_stream: "TENSORS:tensor_in"
|
||||||
|
output_stream: "TENSORS:tensor_out"
|
||||||
|
options {
|
||||||
|
[mediapipe.TfLiteInferenceCalculatorOptions.ext] {
|
||||||
|
model_path: "mediapipe/calculators/tflite/testdata/add.bin"
|
||||||
|
$delegate
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
)";
|
||||||
|
ASSERT_EQ(absl::StrReplaceAll({{"$delegate", delegate}}, &graph_proto), 1);
|
||||||
// Prepare single calculator graph to and wait for packets.
|
// Prepare single calculator graph to and wait for packets.
|
||||||
CalculatorGraphConfig graph_config =
|
CalculatorGraphConfig graph_config =
|
||||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
::mediapipe::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 +117,10 @@ 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) {
|
||||||
|
DoSmokeTest(/*delegate=*/"");
|
||||||
|
DoSmokeTest(/*delegate=*/"delegate { tflite {} }");
|
||||||
|
}
|
||||||
|
|
||||||
} // namespace mediapipe
|
} // namespace mediapipe
|
||||||
|
|||||||
@@ -24,7 +24,7 @@
|
|||||||
#include "mediapipe/framework/port/ret_check.h"
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
#include "mediapipe/util/resource_util.h"
|
#include "mediapipe/util/resource_util.h"
|
||||||
#include "tensorflow/lite/interpreter.h"
|
#include "tensorflow/lite/interpreter.h"
|
||||||
#if defined(__ANDROID__) || (defined(__APPLE__) && !TARGET_OS_OSX)
|
#if defined(MEDIAPIPE_MOBILE)
|
||||||
#include "mediapipe/util/android/file/base/file.h"
|
#include "mediapipe/util/android/file/base/file.h"
|
||||||
#include "mediapipe/util/android/file/base/helpers.h"
|
#include "mediapipe/util/android/file/base/helpers.h"
|
||||||
#else
|
#else
|
||||||
@@ -66,8 +66,8 @@ class TfLiteTensorsToClassificationCalculator : public CalculatorBase {
|
|||||||
::mediapipe::Status Close(CalculatorContext* cc) override;
|
::mediapipe::Status Close(CalculatorContext* cc) override;
|
||||||
|
|
||||||
private:
|
private:
|
||||||
|
::mediapipe::TfLiteTensorsToClassificationCalculatorOptions options_;
|
||||||
int top_k_ = 0;
|
int top_k_ = 0;
|
||||||
double min_score_threshold_ = 0;
|
|
||||||
std::unordered_map<int, std::string> label_map_;
|
std::unordered_map<int, std::string> label_map_;
|
||||||
bool label_map_loaded_ = false;
|
bool label_map_loaded_ = false;
|
||||||
};
|
};
|
||||||
@@ -93,15 +93,14 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
|||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
cc->SetOffset(TimestampDiff(0));
|
cc->SetOffset(TimestampDiff(0));
|
||||||
|
|
||||||
auto options = cc->Options<
|
options_ = cc->Options<
|
||||||
::mediapipe::TfLiteTensorsToClassificationCalculatorOptions>();
|
::mediapipe::TfLiteTensorsToClassificationCalculatorOptions>();
|
||||||
|
|
||||||
top_k_ = options.top_k();
|
top_k_ = options_.top_k();
|
||||||
min_score_threshold_ = options.min_score_threshold();
|
if (options_.has_label_map_path()) {
|
||||||
if (options.has_label_map_path()) {
|
|
||||||
std::string string_path;
|
std::string string_path;
|
||||||
ASSIGN_OR_RETURN(string_path,
|
ASSIGN_OR_RETURN(string_path,
|
||||||
PathToResourceAsFile(options.label_map_path()));
|
PathToResourceAsFile(options_.label_map_path()));
|
||||||
std::string label_map_string;
|
std::string label_map_string;
|
||||||
MP_RETURN_IF_ERROR(file::GetContents(string_path, &label_map_string));
|
MP_RETURN_IF_ERROR(file::GetContents(string_path, &label_map_string));
|
||||||
|
|
||||||
@@ -125,9 +124,11 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
|||||||
RET_CHECK_EQ(input_tensors.size(), 1);
|
RET_CHECK_EQ(input_tensors.size(), 1);
|
||||||
|
|
||||||
const TfLiteTensor* raw_score_tensor = &input_tensors[0];
|
const TfLiteTensor* raw_score_tensor = &input_tensors[0];
|
||||||
RET_CHECK_EQ(raw_score_tensor->dims->size, 2);
|
int num_classes = 1;
|
||||||
RET_CHECK_EQ(raw_score_tensor->dims->data[0], 1);
|
for (int i = 0; i < raw_score_tensor->dims->size; ++i) {
|
||||||
int num_classes = raw_score_tensor->dims->data[1];
|
num_classes *= raw_score_tensor->dims->data[i];
|
||||||
|
}
|
||||||
|
|
||||||
if (label_map_loaded_) {
|
if (label_map_loaded_) {
|
||||||
RET_CHECK_EQ(num_classes, label_map_.size());
|
RET_CHECK_EQ(num_classes, label_map_.size());
|
||||||
}
|
}
|
||||||
@@ -135,7 +136,8 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
|||||||
|
|
||||||
auto classification_list = absl::make_unique<ClassificationList>();
|
auto classification_list = absl::make_unique<ClassificationList>();
|
||||||
for (int i = 0; i < num_classes; ++i) {
|
for (int i = 0; i < num_classes; ++i) {
|
||||||
if (raw_scores[i] < min_score_threshold_) {
|
if (options_.has_min_score_threshold() &&
|
||||||
|
raw_scores[i] < options_.min_score_threshold()) {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
Classification* classification = classification_list->add_classification();
|
Classification* classification = classification_list->add_classification();
|
||||||
@@ -148,6 +150,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToClassificationCalculator);
|
|||||||
|
|
||||||
// Note that partial_sort will raise error when top_k_ >
|
// Note that partial_sort will raise error when top_k_ >
|
||||||
// classification_list->classification_size().
|
// classification_list->classification_size().
|
||||||
|
CHECK_GE(classification_list->classification_size(), top_k_);
|
||||||
auto raw_classification_list = classification_list->mutable_classification();
|
auto raw_classification_list = classification_list->mutable_classification();
|
||||||
if (top_k_ > 0 && classification_list->classification_size() >= top_k_) {
|
if (top_k_ > 0 && classification_list->classification_size() >= top_k_) {
|
||||||
std::partial_sort(raw_classification_list->begin(),
|
std::partial_sort(raw_classification_list->begin(),
|
||||||
|
|||||||
@@ -18,6 +18,7 @@
|
|||||||
#include "absl/strings/str_format.h"
|
#include "absl/strings/str_format.h"
|
||||||
#include "absl/types/span.h"
|
#include "absl/types/span.h"
|
||||||
#include "mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.pb.h"
|
#include "mediapipe/calculators/tflite/tflite_tensors_to_detections_calculator.pb.h"
|
||||||
|
#include "mediapipe/calculators/tflite/util.h"
|
||||||
#include "mediapipe/framework/calculator_framework.h"
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
#include "mediapipe/framework/deps/file_path.h"
|
#include "mediapipe/framework/deps/file_path.h"
|
||||||
#include "mediapipe/framework/formats/detection.pb.h"
|
#include "mediapipe/framework/formats/detection.pb.h"
|
||||||
@@ -26,28 +27,61 @@
|
|||||||
#include "mediapipe/framework/port/ret_check.h"
|
#include "mediapipe/framework/port/ret_check.h"
|
||||||
#include "tensorflow/lite/interpreter.h"
|
#include "tensorflow/lite/interpreter.h"
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
#include "mediapipe/gpu/gl_calculator_helper.h"
|
#include "mediapipe/gpu/gl_calculator_helper.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
#include "tensorflow/lite/delegates/gpu/gl/gl_buffer.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
|
#include "tensorflow/lite/delegates/gpu/gl/gl_program.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
|
#include "tensorflow/lite/delegates/gpu/gl/gl_shader.h"
|
||||||
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
|
#include "tensorflow/lite/delegates/gpu/gl_delegate.h"
|
||||||
#endif // ANDROID
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if defined(MEDIAPIPE_IOS)
|
||||||
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
#import <CoreVideo/CoreVideo.h>
|
||||||
using ::tflite::gpu::gl::GlBuffer;
|
#import <Metal/Metal.h>
|
||||||
using ::tflite::gpu::gl::GlProgram;
|
#import <MetalKit/MetalKit.h>
|
||||||
using ::tflite::gpu::gl::GlShader;
|
|
||||||
#endif // ANDROID
|
|
||||||
|
|
||||||
namespace mediapipe {
|
#import "mediapipe/gpu/MPPMetalHelper.h"
|
||||||
|
#include "mediapipe/gpu/MPPMetalUtil.h"
|
||||||
|
#include "mediapipe/gpu/gpu_buffer.h"
|
||||||
|
#include "tensorflow/lite/delegates/gpu/metal_delegate.h"
|
||||||
|
#endif // iOS
|
||||||
|
|
||||||
namespace {
|
namespace {
|
||||||
|
|
||||||
constexpr int kNumInputTensorsWithAnchors = 3;
|
constexpr int kNumInputTensorsWithAnchors = 3;
|
||||||
constexpr int kNumCoordsPerBox = 4;
|
constexpr int kNumCoordsPerBox = 4;
|
||||||
|
|
||||||
|
} // namespace
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
|
using ::tflite::gpu::gl::CreateReadWriteShaderStorageBuffer;
|
||||||
|
using ::tflite::gpu::gl::GlShader;
|
||||||
|
#endif
|
||||||
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
|
typedef ::tflite::gpu::gl::GlBuffer GpuTensor;
|
||||||
|
typedef ::tflite::gpu::gl::GlProgram GpuProgram;
|
||||||
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
|
typedef id<MTLBuffer> GpuTensor;
|
||||||
|
typedef id<MTLComputePipelineState> GpuProgram;
|
||||||
|
#endif
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
|
struct GPUData {
|
||||||
|
GpuProgram decode_program;
|
||||||
|
GpuProgram score_program;
|
||||||
|
GpuTensor decoded_boxes_buffer;
|
||||||
|
GpuTensor raw_boxes_buffer;
|
||||||
|
GpuTensor raw_anchors_buffer;
|
||||||
|
GpuTensor scored_boxes_buffer;
|
||||||
|
GpuTensor raw_scores_buffer;
|
||||||
|
};
|
||||||
|
#endif
|
||||||
|
|
||||||
void ConvertRawValuesToAnchors(const float* raw_anchors, int num_boxes,
|
void ConvertRawValuesToAnchors(const float* raw_anchors, int num_boxes,
|
||||||
std::vector<Anchor>* anchors) {
|
std::vector<Anchor>* anchors) {
|
||||||
anchors->clear();
|
anchors->clear();
|
||||||
@@ -88,7 +122,7 @@ void ConvertAnchorsToRawValues(const std::vector<Anchor>& anchors,
|
|||||||
// optional to pass in a third tensor for anchors (e.g. for SSD
|
// optional to pass in a third tensor for anchors (e.g. for SSD
|
||||||
// models) depend on the outputs of the detection model. The size
|
// models) depend on the outputs of the detection model. The size
|
||||||
// of anchor tensor must be (num_boxes * 4).
|
// of anchor tensor must be (num_boxes * 4).
|
||||||
// TENSORS_GPU - vector of GlBuffer.
|
// TENSORS_GPU - vector of GlBuffer of MTLBuffer.
|
||||||
// Output:
|
// Output:
|
||||||
// DETECTIONS - Result MediaPipe detections.
|
// DETECTIONS - Result MediaPipe detections.
|
||||||
//
|
//
|
||||||
@@ -126,7 +160,7 @@ class TfLiteTensorsToDetectionsCalculator : public CalculatorBase {
|
|||||||
std::vector<Detection>* output_detections);
|
std::vector<Detection>* output_detections);
|
||||||
|
|
||||||
::mediapipe::Status LoadOptions(CalculatorContext* cc);
|
::mediapipe::Status LoadOptions(CalculatorContext* cc);
|
||||||
::mediapipe::Status GlSetup(CalculatorContext* cc);
|
::mediapipe::Status GpuInit(CalculatorContext* cc);
|
||||||
::mediapipe::Status DecodeBoxes(const float* raw_boxes,
|
::mediapipe::Status DecodeBoxes(const float* raw_boxes,
|
||||||
const std::vector<Anchor>& anchors,
|
const std::vector<Anchor>& anchors,
|
||||||
std::vector<float>* boxes);
|
std::vector<float>* boxes);
|
||||||
@@ -146,15 +180,12 @@ class TfLiteTensorsToDetectionsCalculator : public CalculatorBase {
|
|||||||
std::vector<Anchor> anchors_;
|
std::vector<Anchor> anchors_;
|
||||||
bool side_packet_anchors_{};
|
bool side_packet_anchors_{};
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
mediapipe::GlCalculatorHelper gpu_helper_;
|
mediapipe::GlCalculatorHelper gpu_helper_;
|
||||||
std::unique_ptr<GlProgram> decode_program_;
|
std::unique_ptr<GPUData> gpu_data_;
|
||||||
std::unique_ptr<GlProgram> score_program_;
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
std::unique_ptr<GlBuffer> decoded_boxes_buffer_;
|
MPPMetalHelper* gpu_helper_ = nullptr;
|
||||||
std::unique_ptr<GlBuffer> raw_boxes_buffer_;
|
std::unique_ptr<GPUData> gpu_data_;
|
||||||
std::unique_ptr<GlBuffer> raw_anchors_buffer_;
|
|
||||||
std::unique_ptr<GlBuffer> scored_boxes_buffer_;
|
|
||||||
std::unique_ptr<GlBuffer> raw_scores_buffer_;
|
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
bool gpu_input_ = false;
|
bool gpu_input_ = false;
|
||||||
@@ -167,15 +198,18 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
RET_CHECK(!cc->Inputs().GetTags().empty());
|
RET_CHECK(!cc->Inputs().GetTags().empty());
|
||||||
RET_CHECK(!cc->Outputs().GetTags().empty());
|
RET_CHECK(!cc->Outputs().GetTags().empty());
|
||||||
|
|
||||||
|
bool use_gpu = false;
|
||||||
|
|
||||||
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(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GPU) && !defined(__EMSCRIPTEN__)
|
||||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||||
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GlBuffer>>();
|
cc->Inputs().Tag("TENSORS_GPU").Set<std::vector<GpuTensor>>();
|
||||||
|
use_gpu |= true;
|
||||||
}
|
}
|
||||||
#endif
|
#endif // !MEDIAPIPE_DISABLE_GPU
|
||||||
|
|
||||||
if (cc->Outputs().HasTag("DETECTIONS")) {
|
if (cc->Outputs().HasTag("DETECTIONS")) {
|
||||||
cc->Outputs().Tag("DETECTIONS").Set<std::vector<Detection>>();
|
cc->Outputs().Tag("DETECTIONS").Set<std::vector<Detection>>();
|
||||||
@@ -187,9 +221,13 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#if defined(__ANDROID__)
|
if (use_gpu) {
|
||||||
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
|
MP_RETURN_IF_ERROR(mediapipe::GlCalculatorHelper::UpdateContract(cc));
|
||||||
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
|
MP_RETURN_IF_ERROR([MPPMetalHelper updateContract:cc]);
|
||||||
#endif
|
#endif
|
||||||
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -200,8 +238,11 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
|
|
||||||
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
if (cc->Inputs().HasTag("TENSORS_GPU")) {
|
||||||
gpu_input_ = true;
|
gpu_input_ = true;
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
MP_RETURN_IF_ERROR(gpu_helper_.Open(cc));
|
||||||
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
|
gpu_helper_ = [[MPPMetalHelper alloc] initWithCalculatorContext:cc];
|
||||||
|
RET_CHECK(gpu_helper_);
|
||||||
#endif
|
#endif
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -209,7 +250,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
side_packet_anchors_ = cc->InputSidePackets().HasTag("ANCHORS");
|
side_packet_anchors_ = cc->InputSidePackets().HasTag("ANCHORS");
|
||||||
|
|
||||||
if (gpu_input_) {
|
if (gpu_input_) {
|
||||||
MP_RETURN_IF_ERROR(GlSetup(cc));
|
MP_RETURN_IF_ERROR(GpuInit(cc));
|
||||||
}
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
@@ -228,7 +269,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
MP_RETURN_IF_ERROR(ProcessGPU(cc, output_detections.get()));
|
MP_RETURN_IF_ERROR(ProcessGPU(cc, output_detections.get()));
|
||||||
} else {
|
} else {
|
||||||
MP_RETURN_IF_ERROR(ProcessCPU(cc, output_detections.get()));
|
MP_RETURN_IF_ERROR(ProcessCPU(cc, output_detections.get()));
|
||||||
} // if gpu_input_
|
}
|
||||||
|
|
||||||
// Output
|
// Output
|
||||||
if (cc->Outputs().HasTag("DETECTIONS")) {
|
if (cc->Outputs().HasTag("DETECTIONS")) {
|
||||||
@@ -245,7 +286,8 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
const auto& input_tensors =
|
const auto& input_tensors =
|
||||||
cc->Inputs().Tag("TENSORS").Get<std::vector<TfLiteTensor>>();
|
cc->Inputs().Tag("TENSORS").Get<std::vector<TfLiteTensor>>();
|
||||||
|
|
||||||
if (input_tensors.size() == 2) {
|
if (input_tensors.size() == 2 ||
|
||||||
|
input_tensors.size() == kNumInputTensorsWithAnchors) {
|
||||||
// Postprocessing on CPU for model without postprocessing op. E.g. output
|
// Postprocessing on CPU for model without postprocessing op. E.g. output
|
||||||
// raw score tensor and box tensor. Anchor decoding will be handled below.
|
// raw score tensor and box tensor. Anchor decoding will be handled below.
|
||||||
const TfLiteTensor* raw_box_tensor = &input_tensors[0];
|
const TfLiteTensor* raw_box_tensor = &input_tensors[0];
|
||||||
@@ -358,13 +400,84 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
}
|
}
|
||||||
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::ProcessGPU(
|
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::ProcessGPU(
|
||||||
CalculatorContext* cc, std::vector<Detection>* output_detections) {
|
CalculatorContext* cc, std::vector<Detection>* output_detections) {
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
const auto& input_tensors =
|
const auto& input_tensors =
|
||||||
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GlBuffer>>();
|
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||||
|
RET_CHECK_GE(input_tensors.size(), 2);
|
||||||
|
|
||||||
|
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this, &input_tensors, &cc,
|
||||||
|
&output_detections]()
|
||||||
|
-> ::mediapipe::Status {
|
||||||
|
// Copy inputs.
|
||||||
|
RET_CHECK_CALL(CopyBuffer(input_tensors[0], gpu_data_->raw_boxes_buffer));
|
||||||
|
RET_CHECK_CALL(CopyBuffer(input_tensors[1], gpu_data_->raw_scores_buffer));
|
||||||
|
if (!anchors_init_) {
|
||||||
|
if (side_packet_anchors_) {
|
||||||
|
CHECK(!cc->InputSidePackets().Tag("ANCHORS").IsEmpty());
|
||||||
|
const auto& anchors =
|
||||||
|
cc->InputSidePackets().Tag("ANCHORS").Get<std::vector<Anchor>>();
|
||||||
|
std::vector<float> raw_anchors(num_boxes_ * kNumCoordsPerBox);
|
||||||
|
ConvertAnchorsToRawValues(anchors, num_boxes_, raw_anchors.data());
|
||||||
|
RET_CHECK_CALL(gpu_data_->raw_anchors_buffer.Write<float>(
|
||||||
|
absl::MakeSpan(raw_anchors)));
|
||||||
|
} else {
|
||||||
|
CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
|
||||||
|
RET_CHECK_CALL(
|
||||||
|
CopyBuffer(input_tensors[2], gpu_data_->raw_anchors_buffer));
|
||||||
|
}
|
||||||
|
anchors_init_ = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
// Run shaders.
|
||||||
|
// Decode boxes.
|
||||||
|
RET_CHECK_CALL(gpu_data_->decoded_boxes_buffer.BindToIndex(0));
|
||||||
|
RET_CHECK_CALL(gpu_data_->raw_boxes_buffer.BindToIndex(1));
|
||||||
|
RET_CHECK_CALL(gpu_data_->raw_anchors_buffer.BindToIndex(2));
|
||||||
|
const tflite::gpu::uint3 decode_workgroups = {num_boxes_, 1, 1};
|
||||||
|
RET_CHECK_CALL(gpu_data_->decode_program.Dispatch(decode_workgroups));
|
||||||
|
|
||||||
|
// Score boxes.
|
||||||
|
RET_CHECK_CALL(gpu_data_->scored_boxes_buffer.BindToIndex(0));
|
||||||
|
RET_CHECK_CALL(gpu_data_->raw_scores_buffer.BindToIndex(1));
|
||||||
|
const tflite::gpu::uint3 score_workgroups = {num_boxes_, 1, 1};
|
||||||
|
RET_CHECK_CALL(gpu_data_->score_program.Dispatch(score_workgroups));
|
||||||
|
|
||||||
|
// Copy decoded boxes from GPU to CPU.
|
||||||
|
std::vector<float> boxes(num_boxes_ * num_coords_);
|
||||||
|
RET_CHECK_CALL(gpu_data_->decoded_boxes_buffer.Read(absl::MakeSpan(boxes)));
|
||||||
|
std::vector<float> score_class_id_pairs(num_boxes_ * 2);
|
||||||
|
RET_CHECK_CALL(gpu_data_->scored_boxes_buffer.Read(
|
||||||
|
absl::MakeSpan(score_class_id_pairs)));
|
||||||
|
|
||||||
|
// TODO: b/138851969. Is it possible to output a float vector
|
||||||
|
// for score and an int vector for class so that we can avoid copying twice?
|
||||||
|
std::vector<float> detection_scores(num_boxes_);
|
||||||
|
std::vector<int> detection_classes(num_boxes_);
|
||||||
|
for (int i = 0; i < num_boxes_; ++i) {
|
||||||
|
detection_scores[i] = score_class_id_pairs[i * 2];
|
||||||
|
detection_classes[i] = static_cast<int>(score_class_id_pairs[i * 2 + 1]);
|
||||||
|
}
|
||||||
|
MP_RETURN_IF_ERROR(
|
||||||
|
ConvertToDetections(boxes.data(), detection_scores.data(),
|
||||||
|
detection_classes.data(), output_detections));
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}));
|
||||||
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
|
|
||||||
|
const auto& input_tensors =
|
||||||
|
cc->Inputs().Tag("TENSORS_GPU").Get<std::vector<GpuTensor>>();
|
||||||
|
RET_CHECK_GE(input_tensors.size(), 2);
|
||||||
|
|
||||||
// Copy inputs.
|
// Copy inputs.
|
||||||
tflite::gpu::gl::CopyBuffer(input_tensors[0], *raw_boxes_buffer_.get());
|
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_boxes_buffer
|
||||||
tflite::gpu::gl::CopyBuffer(input_tensors[1], *raw_scores_buffer_.get());
|
from:input_tensors[0]
|
||||||
|
blocking:false
|
||||||
|
commandBuffer:[gpu_helper_ commandBuffer]];
|
||||||
|
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_scores_buffer
|
||||||
|
from:input_tensors[1]
|
||||||
|
blocking:false
|
||||||
|
commandBuffer:[gpu_helper_ commandBuffer]];
|
||||||
if (!anchors_init_) {
|
if (!anchors_init_) {
|
||||||
if (side_packet_anchors_) {
|
if (side_packet_anchors_) {
|
||||||
CHECK(!cc->InputSidePackets().Tag("ANCHORS").IsEmpty());
|
CHECK(!cc->InputSidePackets().Tag("ANCHORS").IsEmpty());
|
||||||
@@ -372,47 +485,54 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
cc->InputSidePackets().Tag("ANCHORS").Get<std::vector<Anchor>>();
|
cc->InputSidePackets().Tag("ANCHORS").Get<std::vector<Anchor>>();
|
||||||
std::vector<float> raw_anchors(num_boxes_ * kNumCoordsPerBox);
|
std::vector<float> raw_anchors(num_boxes_ * kNumCoordsPerBox);
|
||||||
ConvertAnchorsToRawValues(anchors, num_boxes_, raw_anchors.data());
|
ConvertAnchorsToRawValues(anchors, num_boxes_, raw_anchors.data());
|
||||||
raw_anchors_buffer_->Write<float>(absl::MakeSpan(raw_anchors));
|
memcpy([gpu_data_->raw_anchors_buffer contents], raw_anchors.data(),
|
||||||
|
raw_anchors.size() * sizeof(float));
|
||||||
} else {
|
} else {
|
||||||
CHECK_EQ(input_tensors.size(), 3);
|
RET_CHECK_EQ(input_tensors.size(), kNumInputTensorsWithAnchors);
|
||||||
tflite::gpu::gl::CopyBuffer(input_tensors[2], *raw_anchors_buffer_.get());
|
[MPPMetalUtil blitMetalBufferTo:gpu_data_->raw_anchors_buffer
|
||||||
|
from:input_tensors[2]
|
||||||
|
blocking:false
|
||||||
|
commandBuffer:[gpu_helper_ commandBuffer]];
|
||||||
}
|
}
|
||||||
anchors_init_ = true;
|
anchors_init_ = true;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Run shaders.
|
// Run shaders.
|
||||||
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext(
|
id<MTLCommandBuffer> command_buffer = [gpu_helper_ commandBuffer];
|
||||||
[this, &input_tensors]() -> ::mediapipe::Status {
|
command_buffer.label = @"TfLiteDecodeAndScoreBoxes";
|
||||||
// Decode boxes.
|
id<MTLComputeCommandEncoder> command_encoder =
|
||||||
decoded_boxes_buffer_->BindToIndex(0);
|
[command_buffer computeCommandEncoder];
|
||||||
raw_boxes_buffer_->BindToIndex(1);
|
[command_encoder setComputePipelineState:gpu_data_->decode_program];
|
||||||
raw_anchors_buffer_->BindToIndex(2);
|
[command_encoder setBuffer:gpu_data_->decoded_boxes_buffer
|
||||||
const tflite::gpu::uint3 decode_workgroups = {num_boxes_, 1, 1};
|
offset:0
|
||||||
decode_program_->Dispatch(decode_workgroups);
|
atIndex:0];
|
||||||
|
[command_encoder setBuffer:gpu_data_->raw_boxes_buffer offset:0 atIndex:1];
|
||||||
|
[command_encoder setBuffer:gpu_data_->raw_anchors_buffer offset:0 atIndex:2];
|
||||||
|
MTLSize decode_threads_per_group = MTLSizeMake(1, 1, 1);
|
||||||
|
MTLSize decode_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
|
||||||
|
[command_encoder dispatchThreadgroups:decode_threadgroups
|
||||||
|
threadsPerThreadgroup:decode_threads_per_group];
|
||||||
|
|
||||||
// Score boxes.
|
[command_encoder setComputePipelineState:gpu_data_->score_program];
|
||||||
scored_boxes_buffer_->BindToIndex(0);
|
[command_encoder setBuffer:gpu_data_->scored_boxes_buffer offset:0 atIndex:0];
|
||||||
raw_scores_buffer_->BindToIndex(1);
|
[command_encoder setBuffer:gpu_data_->raw_scores_buffer offset:0 atIndex:1];
|
||||||
const tflite::gpu::uint3 score_workgroups = {num_boxes_, 1, 1};
|
MTLSize score_threads_per_group = MTLSizeMake(1, num_classes_, 1);
|
||||||
score_program_->Dispatch(score_workgroups);
|
MTLSize score_threadgroups = MTLSizeMake(num_boxes_, 1, 1);
|
||||||
|
[command_encoder dispatchThreadgroups:score_threadgroups
|
||||||
return ::mediapipe::OkStatus();
|
threadsPerThreadgroup:score_threads_per_group];
|
||||||
}));
|
[command_encoder endEncoding];
|
||||||
|
[MPPMetalUtil commitCommandBufferAndWait:command_buffer];
|
||||||
|
|
||||||
// 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_);
|
||||||
auto status = decoded_boxes_buffer_->Read(absl::MakeSpan(boxes));
|
memcpy(boxes.data(), [gpu_data_->decoded_boxes_buffer contents],
|
||||||
if (!status.ok()) {
|
num_boxes_ * num_coords_ * sizeof(float));
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
|
||||||
}
|
|
||||||
std::vector<float> score_class_id_pairs(num_boxes_ * 2);
|
std::vector<float> score_class_id_pairs(num_boxes_ * 2);
|
||||||
status = scored_boxes_buffer_->Read(absl::MakeSpan(score_class_id_pairs));
|
memcpy(score_class_id_pairs.data(), [gpu_data_->scored_boxes_buffer contents],
|
||||||
if (!status.ok()) {
|
num_boxes_ * 2 * sizeof(float));
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
|
||||||
}
|
|
||||||
|
|
||||||
// TODO: b/138851969. Is it possible to output a float vector
|
// Output detections.
|
||||||
// for score and an int vector for class so that we can avoid copying twice?
|
// TODO Adjust shader to avoid copying shader output twice.
|
||||||
std::vector<float> detection_scores(num_boxes_);
|
std::vector<float> detection_scores(num_boxes_);
|
||||||
std::vector<int> detection_classes(num_boxes_);
|
std::vector<int> detection_classes(num_boxes_);
|
||||||
for (int i = 0; i < num_boxes_; ++i) {
|
for (int i = 0; i < num_boxes_; ++i) {
|
||||||
@@ -422,25 +542,20 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
MP_RETURN_IF_ERROR(ConvertToDetections(boxes.data(), detection_scores.data(),
|
MP_RETURN_IF_ERROR(ConvertToDetections(boxes.data(), detection_scores.data(),
|
||||||
detection_classes.data(),
|
detection_classes.data(),
|
||||||
output_detections));
|
output_detections));
|
||||||
|
|
||||||
#else
|
#else
|
||||||
LOG(ERROR) << "GPU input on non-Android not supported yet.";
|
LOG(ERROR) << "GPU input on non-Android not supported yet.";
|
||||||
#endif // defined(__ANDROID__)
|
#endif
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::Close(
|
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::Close(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
gpu_helper_.RunInGlContext([this] {
|
gpu_helper_.RunInGlContext([this] { gpu_data_.reset(); });
|
||||||
decode_program_.reset();
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
score_program_.reset();
|
gpu_data_.reset();
|
||||||
decoded_boxes_buffer_.reset();
|
#endif
|
||||||
raw_boxes_buffer_.reset();
|
|
||||||
raw_anchors_buffer_.reset();
|
|
||||||
scored_boxes_buffer_.reset();
|
|
||||||
raw_scores_buffer_.reset();
|
|
||||||
});
|
|
||||||
#endif // __ANDROID__
|
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
@@ -530,6 +645,7 @@ REGISTER_CALCULATOR(TfLiteTensorsToDetectionsCalculator);
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -586,12 +702,16 @@ Detection TfLiteTensorsToDetectionsCalculator::ConvertToDetection(
|
|||||||
return detection;
|
return detection;
|
||||||
}
|
}
|
||||||
|
|
||||||
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::GlSetup(
|
::mediapipe::Status TfLiteTensorsToDetectionsCalculator::GpuInit(
|
||||||
CalculatorContext* cc) {
|
CalculatorContext* cc) {
|
||||||
#if defined(__ANDROID__)
|
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
// A shader to decode detection boxes.
|
MP_RETURN_IF_ERROR(gpu_helper_.RunInGlContext([this]()
|
||||||
const std::string decode_src = absl::Substitute(
|
-> ::mediapipe::Status {
|
||||||
R"( #version 310 es
|
gpu_data_ = absl::make_unique<GPUData>();
|
||||||
|
|
||||||
|
// A shader to decode detection boxes.
|
||||||
|
const std::string decode_src = absl::Substitute(
|
||||||
|
R"( #version 310 es
|
||||||
|
|
||||||
layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
|
layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
|
||||||
|
|
||||||
@@ -665,7 +785,7 @@ void main() {
|
|||||||
if (num_keypoints > int(0)){
|
if (num_keypoints > int(0)){
|
||||||
for (int k = 0; k < num_keypoints; ++k) {
|
for (int k = 0; k < num_keypoints; ++k) {
|
||||||
int kp_offset =
|
int kp_offset =
|
||||||
int(g_idx * num_coords) + keypt_coord_offset + k * num_values_per_keypt;
|
int(g_idx * num_coords) + keypt_coord_offset + k * num_values_per_keypt;
|
||||||
float kp_y, kp_x;
|
float kp_y, kp_x;
|
||||||
if (reverse_output_order == int(0)) {
|
if (reverse_output_order == int(0)) {
|
||||||
kp_y = raw_boxes.data[kp_offset + int(0)];
|
kp_y = raw_boxes.data[kp_offset + int(0)];
|
||||||
@@ -679,55 +799,37 @@ void main() {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
})",
|
})",
|
||||||
options_.num_coords(), // box xywh
|
options_.num_coords(), // box xywh
|
||||||
options_.reverse_output_order() ? 1 : 0,
|
options_.reverse_output_order() ? 1 : 0,
|
||||||
options_.apply_exponential_on_box_size() ? 1 : 0,
|
options_.apply_exponential_on_box_size() ? 1 : 0,
|
||||||
options_.box_coord_offset(), options_.num_keypoints(),
|
options_.box_coord_offset(), options_.num_keypoints(),
|
||||||
options_.keypoint_coord_offset(), options_.num_values_per_keypoint());
|
options_.keypoint_coord_offset(), options_.num_values_per_keypoint());
|
||||||
|
|
||||||
// Shader program
|
// Shader program
|
||||||
GlShader decode_shader;
|
GlShader decode_shader;
|
||||||
auto status =
|
RET_CHECK_CALL(
|
||||||
GlShader::CompileShader(GL_COMPUTE_SHADER, decode_src, &decode_shader);
|
GlShader::CompileShader(GL_COMPUTE_SHADER, decode_src, &decode_shader));
|
||||||
if (!status.ok()) {
|
RET_CHECK_CALL(GpuProgram::CreateWithShader(decode_shader,
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
&gpu_data_->decode_program));
|
||||||
}
|
// Outputs
|
||||||
decode_program_ = absl::make_unique<GlProgram>();
|
size_t decoded_boxes_length = num_boxes_ * num_coords_;
|
||||||
status = GlProgram::CreateWithShader(decode_shader, decode_program_.get());
|
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
|
||||||
if (!status.ok()) {
|
decoded_boxes_length, &gpu_data_->decoded_boxes_buffer));
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
// Inputs
|
||||||
}
|
size_t raw_boxes_length = num_boxes_ * num_coords_;
|
||||||
// Outputs
|
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
|
||||||
size_t decoded_boxes_length = num_boxes_ * num_coords_;
|
raw_boxes_length, &gpu_data_->raw_boxes_buffer));
|
||||||
decoded_boxes_buffer_ = absl::make_unique<GlBuffer>();
|
size_t raw_anchors_length = num_boxes_ * kNumCoordsPerBox;
|
||||||
status = CreateReadWriteShaderStorageBuffer<float>(
|
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
|
||||||
decoded_boxes_length, decoded_boxes_buffer_.get());
|
raw_anchors_length, &gpu_data_->raw_anchors_buffer));
|
||||||
if (!status.ok()) {
|
// Parameters
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
glUseProgram(gpu_data_->decode_program.id());
|
||||||
}
|
glUniform4f(0, options_.x_scale(), options_.y_scale(), options_.w_scale(),
|
||||||
// Inputs
|
options_.h_scale());
|
||||||
size_t raw_boxes_length = num_boxes_ * num_coords_;
|
|
||||||
raw_boxes_buffer_ = absl::make_unique<GlBuffer>();
|
|
||||||
status = CreateReadWriteShaderStorageBuffer<float>(raw_boxes_length,
|
|
||||||
raw_boxes_buffer_.get());
|
|
||||||
if (!status.ok()) {
|
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
|
||||||
}
|
|
||||||
size_t raw_anchors_length = num_boxes_ * kNumCoordsPerBox;
|
|
||||||
raw_anchors_buffer_ = absl::make_unique<GlBuffer>();
|
|
||||||
status = CreateReadWriteShaderStorageBuffer<float>(raw_anchors_length,
|
|
||||||
raw_anchors_buffer_.get());
|
|
||||||
if (!status.ok()) {
|
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
|
||||||
}
|
|
||||||
// Parameters
|
|
||||||
glUseProgram(decode_program_->id());
|
|
||||||
glUniform4f(0, options_.x_scale(), options_.y_scale(), options_.w_scale(),
|
|
||||||
options_.h_scale());
|
|
||||||
|
|
||||||
// A shader to score detection boxes.
|
// A shader to score detection boxes.
|
||||||
const std::string score_src = absl::Substitute(
|
const std::string score_src = absl::Substitute(
|
||||||
R"( #version 310 es
|
R"( #version 310 es
|
||||||
|
|
||||||
layout(local_size_x = 1, local_size_y = $0, local_size_z = 1) in;
|
layout(local_size_x = 1, local_size_y = $0, local_size_z = 1) in;
|
||||||
|
|
||||||
@@ -781,6 +883,227 @@ void main() {
|
|||||||
scored_boxes.data[g_idx * uint(2) + uint(0)] = max_score;
|
scored_boxes.data[g_idx * uint(2) + uint(0)] = max_score;
|
||||||
scored_boxes.data[g_idx * uint(2) + uint(1)] = max_class;
|
scored_boxes.data[g_idx * uint(2) + uint(1)] = max_class;
|
||||||
}
|
}
|
||||||
|
})",
|
||||||
|
num_classes_, options_.sigmoid_score() ? 1 : 0,
|
||||||
|
options_.has_score_clipping_thresh() ? 1 : 0,
|
||||||
|
options_.has_score_clipping_thresh() ? options_.score_clipping_thresh()
|
||||||
|
: 0,
|
||||||
|
!ignore_classes_.empty() ? 1 : 0);
|
||||||
|
|
||||||
|
// # filter classes supported is hardware dependent.
|
||||||
|
int max_wg_size; // typically <= 1024
|
||||||
|
glGetIntegeri_v(GL_MAX_COMPUTE_WORK_GROUP_SIZE, 1,
|
||||||
|
&max_wg_size); // y-dim
|
||||||
|
CHECK_LT(num_classes_, max_wg_size)
|
||||||
|
<< "# classes must be < " << max_wg_size;
|
||||||
|
// TODO support better filtering.
|
||||||
|
CHECK_LE(ignore_classes_.size(), 1) << "Only ignore class 0 is allowed";
|
||||||
|
|
||||||
|
// Shader program
|
||||||
|
GlShader score_shader;
|
||||||
|
RET_CHECK_CALL(
|
||||||
|
GlShader::CompileShader(GL_COMPUTE_SHADER, score_src, &score_shader));
|
||||||
|
RET_CHECK_CALL(
|
||||||
|
GpuProgram::CreateWithShader(score_shader, &gpu_data_->score_program));
|
||||||
|
// Outputs
|
||||||
|
size_t scored_boxes_length = num_boxes_ * 2; // score, class
|
||||||
|
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
|
||||||
|
scored_boxes_length, &gpu_data_->scored_boxes_buffer));
|
||||||
|
// Inputs
|
||||||
|
size_t raw_scores_length = num_boxes_ * num_classes_;
|
||||||
|
RET_CHECK_CALL(CreateReadWriteShaderStorageBuffer<float>(
|
||||||
|
raw_scores_length, &gpu_data_->raw_scores_buffer));
|
||||||
|
|
||||||
|
return ::mediapipe::OkStatus();
|
||||||
|
}));
|
||||||
|
|
||||||
|
#elif defined(MEDIAPIPE_IOS)
|
||||||
|
|
||||||
|
gpu_data_ = absl::make_unique<GPUData>();
|
||||||
|
id<MTLDevice> device = gpu_helper_.mtlDevice;
|
||||||
|
|
||||||
|
// A shader to decode detection boxes.
|
||||||
|
std::string decode_src = absl::Substitute(
|
||||||
|
R"(
|
||||||
|
#include <metal_stdlib>
|
||||||
|
|
||||||
|
using namespace metal;
|
||||||
|
|
||||||
|
kernel void decodeKernel(
|
||||||
|
device float* boxes [[ buffer(0) ]],
|
||||||
|
device float* raw_boxes [[ buffer(1) ]],
|
||||||
|
device float* raw_anchors [[ buffer(2) ]],
|
||||||
|
uint2 gid [[ thread_position_in_grid ]]) {
|
||||||
|
|
||||||
|
uint num_coords = uint($0);
|
||||||
|
int reverse_output_order = int($1);
|
||||||
|
int apply_exponential = int($2);
|
||||||
|
int box_coord_offset = int($3);
|
||||||
|
int num_keypoints = int($4);
|
||||||
|
int keypt_coord_offset = int($5);
|
||||||
|
int num_values_per_keypt = int($6);
|
||||||
|
)",
|
||||||
|
options_.num_coords(), // box xywh
|
||||||
|
options_.reverse_output_order() ? 1 : 0,
|
||||||
|
options_.apply_exponential_on_box_size() ? 1 : 0,
|
||||||
|
options_.box_coord_offset(), options_.num_keypoints(),
|
||||||
|
options_.keypoint_coord_offset(), options_.num_values_per_keypoint());
|
||||||
|
decode_src += absl::Substitute(
|
||||||
|
R"(
|
||||||
|
float4 scale = float4(($0),($1),($2),($3));
|
||||||
|
)",
|
||||||
|
options_.x_scale(), options_.y_scale(), options_.w_scale(),
|
||||||
|
options_.h_scale());
|
||||||
|
decode_src += R"(
|
||||||
|
uint g_idx = gid.x;
|
||||||
|
uint box_offset = g_idx * num_coords + uint(box_coord_offset);
|
||||||
|
uint anchor_offset = g_idx * uint(4); // check kNumCoordsPerBox
|
||||||
|
|
||||||
|
float y_center, x_center, h, w;
|
||||||
|
|
||||||
|
if (reverse_output_order == int(0)) {
|
||||||
|
y_center = raw_boxes[box_offset + uint(0)];
|
||||||
|
x_center = raw_boxes[box_offset + uint(1)];
|
||||||
|
h = raw_boxes[box_offset + uint(2)];
|
||||||
|
w = raw_boxes[box_offset + uint(3)];
|
||||||
|
} else {
|
||||||
|
x_center = raw_boxes[box_offset + uint(0)];
|
||||||
|
y_center = raw_boxes[box_offset + uint(1)];
|
||||||
|
w = raw_boxes[box_offset + uint(2)];
|
||||||
|
h = raw_boxes[box_offset + uint(3)];
|
||||||
|
}
|
||||||
|
|
||||||
|
float anchor_yc = raw_anchors[anchor_offset + uint(0)];
|
||||||
|
float anchor_xc = raw_anchors[anchor_offset + uint(1)];
|
||||||
|
float anchor_h = raw_anchors[anchor_offset + uint(2)];
|
||||||
|
float anchor_w = raw_anchors[anchor_offset + uint(3)];
|
||||||
|
|
||||||
|
x_center = x_center / scale.x * anchor_w + anchor_xc;
|
||||||
|
y_center = y_center / scale.y * anchor_h + anchor_yc;
|
||||||
|
|
||||||
|
if (apply_exponential == int(1)) {
|
||||||
|
h = exp(h / scale.w) * anchor_h;
|
||||||
|
w = exp(w / scale.z) * anchor_w;
|
||||||
|
} else {
|
||||||
|
h = (h / scale.w) * anchor_h;
|
||||||
|
w = (w / scale.z) * anchor_w;
|
||||||
|
}
|
||||||
|
|
||||||
|
float ymin = y_center - h / 2.0;
|
||||||
|
float xmin = x_center - w / 2.0;
|
||||||
|
float ymax = y_center + h / 2.0;
|
||||||
|
float xmax = x_center + w / 2.0;
|
||||||
|
|
||||||
|
boxes[box_offset + uint(0)] = ymin;
|
||||||
|
boxes[box_offset + uint(1)] = xmin;
|
||||||
|
boxes[box_offset + uint(2)] = ymax;
|
||||||
|
boxes[box_offset + uint(3)] = xmax;
|
||||||
|
|
||||||
|
if (num_keypoints > int(0)){
|
||||||
|
for (int k = 0; k < num_keypoints; ++k) {
|
||||||
|
int kp_offset =
|
||||||
|
int(g_idx * num_coords) + keypt_coord_offset + k * num_values_per_keypt;
|
||||||
|
float kp_y, kp_x;
|
||||||
|
if (reverse_output_order == int(0)) {
|
||||||
|
kp_y = raw_boxes[kp_offset + int(0)];
|
||||||
|
kp_x = raw_boxes[kp_offset + int(1)];
|
||||||
|
} else {
|
||||||
|
kp_x = raw_boxes[kp_offset + int(0)];
|
||||||
|
kp_y = raw_boxes[kp_offset + int(1)];
|
||||||
|
}
|
||||||
|
boxes[kp_offset + int(0)] = kp_x / scale.x * anchor_w + anchor_xc;
|
||||||
|
boxes[kp_offset + int(1)] = kp_y / scale.y * anchor_h + anchor_yc;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})";
|
||||||
|
|
||||||
|
{
|
||||||
|
// Shader program
|
||||||
|
NSString* library_source =
|
||||||
|
[NSString stringWithUTF8String:decode_src.c_str()];
|
||||||
|
NSError* error = nil;
|
||||||
|
id<MTLLibrary> library = [device newLibraryWithSource:library_source
|
||||||
|
options:nullptr
|
||||||
|
error:&error];
|
||||||
|
RET_CHECK(library != nil) << "Couldn't create shader library "
|
||||||
|
<< [[error localizedDescription] UTF8String];
|
||||||
|
id<MTLFunction> kernel_func = nil;
|
||||||
|
kernel_func = [library newFunctionWithName:@"decodeKernel"];
|
||||||
|
RET_CHECK(kernel_func != nil) << "Couldn't create kernel function.";
|
||||||
|
gpu_data_->decode_program =
|
||||||
|
[device newComputePipelineStateWithFunction:kernel_func error:&error];
|
||||||
|
RET_CHECK(gpu_data_->decode_program != nil)
|
||||||
|
<< "Couldn't create pipeline state "
|
||||||
|
<< [[error localizedDescription] UTF8String];
|
||||||
|
// Outputs
|
||||||
|
size_t decoded_boxes_length = num_boxes_ * num_coords_ * sizeof(float);
|
||||||
|
gpu_data_->decoded_boxes_buffer =
|
||||||
|
[device newBufferWithLength:decoded_boxes_length
|
||||||
|
options:MTLResourceStorageModeShared];
|
||||||
|
// Inputs
|
||||||
|
size_t raw_boxes_length = num_boxes_ * num_coords_ * sizeof(float);
|
||||||
|
gpu_data_->raw_boxes_buffer =
|
||||||
|
[device newBufferWithLength:raw_boxes_length
|
||||||
|
options:MTLResourceStorageModeShared];
|
||||||
|
size_t raw_anchors_length = num_boxes_ * kNumCoordsPerBox * sizeof(float);
|
||||||
|
gpu_data_->raw_anchors_buffer =
|
||||||
|
[device newBufferWithLength:raw_anchors_length
|
||||||
|
options:MTLResourceStorageModeShared];
|
||||||
|
}
|
||||||
|
|
||||||
|
// A shader to score detection boxes.
|
||||||
|
const std::string score_src = absl::Substitute(
|
||||||
|
R"(
|
||||||
|
#include <metal_stdlib>
|
||||||
|
|
||||||
|
using namespace metal;
|
||||||
|
|
||||||
|
float optional_sigmoid(float x) {
|
||||||
|
int apply_sigmoid = int($1);
|
||||||
|
int apply_clipping_thresh = int($2);
|
||||||
|
float clipping_thresh = float($3);
|
||||||
|
if (apply_sigmoid == int(0)) return x;
|
||||||
|
if (apply_clipping_thresh == int(1)) {
|
||||||
|
x = clamp(x, -clipping_thresh, clipping_thresh);
|
||||||
|
}
|
||||||
|
x = 1.0 / (1.0 + exp(-x));
|
||||||
|
return x;
|
||||||
|
}
|
||||||
|
|
||||||
|
kernel void scoreKernel(
|
||||||
|
device float* scored_boxes [[ buffer(0) ]],
|
||||||
|
device float* raw_scores [[ buffer(1) ]],
|
||||||
|
uint2 tid [[ thread_position_in_threadgroup ]],
|
||||||
|
uint2 gid [[ thread_position_in_grid ]]) {
|
||||||
|
|
||||||
|
uint num_classes = uint($0);
|
||||||
|
int apply_sigmoid = int($1);
|
||||||
|
int apply_clipping_thresh = int($2);
|
||||||
|
float clipping_thresh = float($3);
|
||||||
|
int ignore_class_0 = int($4);
|
||||||
|
|
||||||
|
uint g_idx = gid.x; // box idx
|
||||||
|
uint s_idx = tid.y; // score/class idx
|
||||||
|
|
||||||
|
// load all scores into shared memory
|
||||||
|
threadgroup float local_scores[$0];
|
||||||
|
float score = raw_scores[g_idx * num_classes + s_idx];
|
||||||
|
local_scores[s_idx] = optional_sigmoid(score);
|
||||||
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||||
|
|
||||||
|
// find max score in shared memory
|
||||||
|
if (s_idx == uint(0)) {
|
||||||
|
float max_score = -FLT_MAX;
|
||||||
|
float max_class = -1.0;
|
||||||
|
for (int i=ignore_class_0; i<int(num_classes); ++i) {
|
||||||
|
if (local_scores[i] > max_score) {
|
||||||
|
max_score = local_scores[i];
|
||||||
|
max_class = float(i);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
scored_boxes[g_idx * uint(2) + uint(0)] = max_score;
|
||||||
|
scored_boxes[g_idx * uint(2) + uint(1)] = max_class;
|
||||||
|
}
|
||||||
})",
|
})",
|
||||||
num_classes_, options_.sigmoid_score() ? 1 : 0,
|
num_classes_, options_.sigmoid_score() ? 1 : 0,
|
||||||
options_.has_score_clipping_thresh() ? 1 : 0,
|
options_.has_score_clipping_thresh() ? 1 : 0,
|
||||||
@@ -788,42 +1111,44 @@ void main() {
|
|||||||
: 0,
|
: 0,
|
||||||
ignore_classes_.size() ? 1 : 0);
|
ignore_classes_.size() ? 1 : 0);
|
||||||
|
|
||||||
// # filter classes supported is hardware dependent.
|
|
||||||
int max_wg_size; // typically <= 1024
|
|
||||||
glGetIntegeri_v(GL_MAX_COMPUTE_WORK_GROUP_SIZE, 1, &max_wg_size); // y-dim
|
|
||||||
CHECK_LT(num_classes_, max_wg_size) << "# classes must be < " << max_wg_size;
|
|
||||||
// TODO support better filtering.
|
// TODO support better filtering.
|
||||||
CHECK_LE(ignore_classes_.size(), 1) << "Only ignore class 0 is allowed";
|
CHECK_LE(ignore_classes_.size(), 1) << "Only ignore class 0 is allowed";
|
||||||
|
|
||||||
// Shader program
|
{
|
||||||
GlShader score_shader;
|
// Shader program
|
||||||
status = GlShader::CompileShader(GL_COMPUTE_SHADER, score_src, &score_shader);
|
NSString* library_source =
|
||||||
if (!status.ok()) {
|
[NSString stringWithUTF8String:score_src.c_str()];
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
NSError* error = nil;
|
||||||
}
|
id<MTLLibrary> library = [device newLibraryWithSource:library_source
|
||||||
score_program_ = absl::make_unique<GlProgram>();
|
options:nullptr
|
||||||
status = GlProgram::CreateWithShader(score_shader, score_program_.get());
|
error:&error];
|
||||||
if (!status.ok()) {
|
RET_CHECK(library != nil) << "Couldn't create shader library "
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
<< [[error localizedDescription] UTF8String];
|
||||||
}
|
id<MTLFunction> kernel_func = nil;
|
||||||
// Outputs
|
kernel_func = [library newFunctionWithName:@"scoreKernel"];
|
||||||
size_t scored_boxes_length = num_boxes_ * 2; // score, class
|
RET_CHECK(kernel_func != nil) << "Couldn't create kernel function.";
|
||||||
scored_boxes_buffer_ = absl::make_unique<GlBuffer>();
|
gpu_data_->score_program =
|
||||||
status = CreateReadWriteShaderStorageBuffer<float>(
|
[device newComputePipelineStateWithFunction:kernel_func error:&error];
|
||||||
scored_boxes_length, scored_boxes_buffer_.get());
|
RET_CHECK(gpu_data_->score_program != nil)
|
||||||
if (!status.ok()) {
|
<< "Couldn't create pipeline state "
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
<< [[error localizedDescription] UTF8String];
|
||||||
}
|
// Outputs
|
||||||
// Inputs
|
size_t scored_boxes_length = num_boxes_ * 2 * sizeof(float); // score,class
|
||||||
size_t raw_scores_length = num_boxes_ * num_classes_;
|
gpu_data_->scored_boxes_buffer =
|
||||||
raw_scores_buffer_ = absl::make_unique<GlBuffer>();
|
[device newBufferWithLength:scored_boxes_length
|
||||||
status = CreateReadWriteShaderStorageBuffer<float>(raw_scores_length,
|
options:MTLResourceStorageModeShared];
|
||||||
raw_scores_buffer_.get());
|
// Inputs
|
||||||
if (!status.ok()) {
|
size_t raw_scores_length = num_boxes_ * num_classes_ * sizeof(float);
|
||||||
return ::mediapipe::InternalError(status.error_message());
|
gpu_data_->raw_scores_buffer =
|
||||||
|
[device newBufferWithLength:raw_scores_length
|
||||||
|
options:MTLResourceStorageModeShared];
|
||||||
|
// # filter classes supported is hardware dependent.
|
||||||
|
int max_wg_size = gpu_data_->score_program.maxTotalThreadsPerThreadgroup;
|
||||||
|
CHECK_LT(num_classes_, max_wg_size) << "# classes must be <" << max_wg_size;
|
||||||
}
|
}
|
||||||
|
|
||||||
#endif // defined(__ANDROID__)
|
#endif // !defined(MEDIAPIPE_DISABLE_GL_COMPUTE)
|
||||||
|
|
||||||
return ::mediapipe::OkStatus();
|
return ::mediapipe::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user