diff --git a/mediapipe/calculators/util/BUILD b/mediapipe/calculators/util/BUILD index df6d5c6d6..eb9d9c8ea 100644 --- a/mediapipe/calculators/util/BUILD +++ b/mediapipe/calculators/util/BUILD @@ -236,6 +236,20 @@ cc_library( alwayslink = 1, ) +cc_library( + name = "counting_vector_size_calculator", + srcs = ["counting_vector_size_calculator.cc"], + hdrs = ["counting_vector_size_calculator.h"], + visibility = [ + "//visibility:public", + ], + deps = [ + "//mediapipe/framework:calculator_framework", + "//mediapipe/framework/formats:landmark_cc_proto", + ], + alwayslink = 1, +) + cc_library( name = "annotation_overlay_calculator", srcs = ["annotation_overlay_calculator.cc"], diff --git a/mediapipe/calculators/util/counting_vector_size_calculator.cc b/mediapipe/calculators/util/counting_vector_size_calculator.cc new file mode 100644 index 000000000..d483240eb --- /dev/null +++ b/mediapipe/calculators/util/counting_vector_size_calculator.cc @@ -0,0 +1,27 @@ +// Copyright 2020 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "mediapipe/calculators/util/counting_vector_size_calculator.h" + +#include "mediapipe/framework/formats/landmark.pb.h" + +namespace mediapipe { + +typedef CountingVectorSizeCalculator> + CountingNormalizedLandmarkListVectorSizeCalculator; + +REGISTER_CALCULATOR(CountingNormalizedLandmarkListVectorSizeCalculator); +} // namespace mediapipe + + diff --git a/mediapipe/calculators/util/counting_vector_size_calculator.h b/mediapipe/calculators/util/counting_vector_size_calculator.h new file mode 100644 index 000000000..12bcbe548 --- /dev/null +++ b/mediapipe/calculators/util/counting_vector_size_calculator.h @@ -0,0 +1,79 @@ +// 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. + +#ifndef MEDIAPIPE_CALCULATORS_UTIL_COUNTING_VECTOR_SIZE_CALCULATOR_H +#define MEDIAPIPE_CALCULATORS_UTIL_COUNTING_VECTOR_SIZE_CALCULATOR_H + +#include "mediapipe/framework/calculator_framework.h" +#include "mediapipe/framework/formats/landmark.pb.h" + +namespace mediapipe { + +// A calculator that counts the size of the input vector. It was created to +// aid in polling packets in the output stream synchronously. If there is +// a clock stream, it will output a value of 0 even if the input vector stream +// is empty. If not, it will output some value only if there is an input vector. +// The clock stream must have the same time stamp as the vector stream, and +// it must be the stream where packets are transmitted while the graph is running. +// (e.g. Any input stream of graph) +// +// It is designed to be used like: +// +// Example config: +// node { +// calculator: "CountingWithVectorSizeCalculator" +// input_stream: "CLOCK:triger_signal" +// input_stream: "VECTOR:input_vector" +// output_stream: "COUNT:vector_count" +// } +// +// node { +// calculator: "CountingWithVectorSizeCalculator" +// input_stream: "VECTOR:input_vector" +// output_stream: "COUNT:vector_count" +// } + +template +class CountingVectorSizeCalculator : public CalculatorBase { + public: + static ::mediapipe::Status GetContract(CalculatorContract* cc) { + if (cc->Inputs().HasTag("CLOCK")) { + cc->Inputs().Tag("CLOCK").SetAny(); + } + + RET_CHECK(cc->Inputs().HasTag("VECTOR")); + cc->Inputs().Tag("VECTOR").Set(); + RET_CHECK(cc->Outputs().HasTag("COUNT")); + cc->Outputs().Tag("COUNT").Set(); + + return ::mediapipe::OkStatus(); + } + + ::mediapipe::Status Process(CalculatorContext* cc) { + std::unique_ptr face_count; + if (!cc->Inputs().Tag("VECTOR").IsEmpty()) { + const auto& landmarks = cc->Inputs().Tag("VECTOR").Get(); + face_count = absl::make_unique(landmarks.size()); + } else { + face_count = absl::make_unique(0); + } + cc->Outputs().Tag("COUNT").Add(face_count.release(), cc->InputTimestamp()); + + return ::mediapipe::OkStatus(); + }; +}; + +} + +#endif // MEDIAPIPE_CALCULATORS_UTIL_COUNTING_VECTOR_SIZE_CALCULATOR_H \ No newline at end of file diff --git a/mediapipe/examples/desktop/BUILD b/mediapipe/examples/desktop/BUILD index 7772e21da..ddcabdd37 100644 --- a/mediapipe/examples/desktop/BUILD +++ b/mediapipe/examples/desktop/BUILD @@ -41,6 +41,7 @@ cc_library( "//mediapipe/framework:calculator_framework", "//mediapipe/framework/formats:image_frame", "//mediapipe/framework/formats:image_frame_opencv", + "//mediapipe/framework/formats:landmark_cc_proto", "//mediapipe/framework/port:commandlineflags", "//mediapipe/framework/port:file_helpers", "//mediapipe/framework/port:opencv_highgui", diff --git a/mediapipe/graphs/face_mesh/face_mesh_desktop_live.pbtxt b/mediapipe/graphs/face_mesh/face_mesh_desktop_live.pbtxt index 57654436a..f4e11023c 100644 --- a/mediapipe/graphs/face_mesh/face_mesh_desktop_live.pbtxt +++ b/mediapipe/graphs/face_mesh/face_mesh_desktop_live.pbtxt @@ -8,6 +8,7 @@ output_stream: "output_video" # Collection of detected/processed faces, each represented as a list of # landmarks. (std::vector) output_stream: "multi_face_landmarks" +output_stream: "face_count" # Throttles the images flowing downstream for flow control. It passes through # the very first incoming image unaltered, and waits for downstream nodes @@ -50,6 +51,7 @@ node { output_stream: "ROIS_FROM_LANDMARKS:face_rects_from_landmarks" output_stream: "DETECTIONS:face_detections" output_stream: "ROIS_FROM_DETECTIONS:face_rects_from_detections" + output_stream: "FACE_COUNT_FROM_LANDMARKS:face_count" } # Subgraph that renders face-landmark annotation onto the input image. diff --git a/mediapipe/modules/face_landmark/BUILD b/mediapipe/modules/face_landmark/BUILD index 7aa10fe54..e0f30a0fc 100644 --- a/mediapipe/modules/face_landmark/BUILD +++ b/mediapipe/modules/face_landmark/BUILD @@ -70,6 +70,7 @@ mediapipe_simple_subgraph( "//mediapipe/calculators/image:image_properties_calculator", "//mediapipe/calculators/util:association_norm_rect_calculator", "//mediapipe/calculators/util:collection_has_min_size_calculator", + "//mediapipe/calculators/util:counting_vector_size_calculator", "//mediapipe/modules/face_detection:face_detection_front_cpu", ], ) @@ -91,6 +92,7 @@ mediapipe_simple_subgraph( "//mediapipe/calculators/image:image_properties_calculator", "//mediapipe/calculators/util:association_norm_rect_calculator", "//mediapipe/calculators/util:collection_has_min_size_calculator", + "//mediapipe/calculators/util:counting_vector_size_calculator", "//mediapipe/modules/face_detection:face_detection_front_gpu", ], ) diff --git a/mediapipe/modules/face_landmark/face_landmark_front_cpu.pbtxt b/mediapipe/modules/face_landmark/face_landmark_front_cpu.pbtxt index 6086a7cd2..2d8a6e771 100644 --- a/mediapipe/modules/face_landmark/face_landmark_front_cpu.pbtxt +++ b/mediapipe/modules/face_landmark/face_landmark_front_cpu.pbtxt @@ -43,6 +43,8 @@ output_stream: "ROIS_FROM_LANDMARKS:face_rects_from_landmarks" # Regions of interest calculated based on face detections. # (std::vector) output_stream: "ROIS_FROM_DETECTIONS:face_rects_from_detections" +# (int) +output_stream: "FACE_COUNT_FROM_LANDMARKS:face_count" # Defines whether landmarks on the previous image should be used to help # localize landmarks on the current image. @@ -208,6 +210,15 @@ node { output_stream: "ITERABLE:multi_face_landmarks" } +# Counting a multi_faceLandmarks vector size. The image stream is only used to +# make the calculator work even when there is no input vector. +node { + calculator: "CountingNormalizedLandmarkListVectorSizeCalculator" + input_stream: "CLOCK:image" + input_stream: "VECTOR:multi_face_landmarks" + output_stream: "COUNT:face_count" +} + # Collects a NormalizedRect for each face into a vector. Upon receiving the # BATCH_END timestamp, outputs the vector of NormalizedRect at the BATCH_END # timestamp. diff --git a/mediapipe/modules/face_landmark/face_landmark_front_gpu.pbtxt b/mediapipe/modules/face_landmark/face_landmark_front_gpu.pbtxt index 1caa6baf6..b4b312f42 100644 --- a/mediapipe/modules/face_landmark/face_landmark_front_gpu.pbtxt +++ b/mediapipe/modules/face_landmark/face_landmark_front_gpu.pbtxt @@ -43,6 +43,8 @@ output_stream: "ROIS_FROM_LANDMARKS:face_rects_from_landmarks" # Regions of interest calculated based on face detections. # (std::vector) output_stream: "ROIS_FROM_DETECTIONS:face_rects_from_detections" +# (int) +output_stream: "FACE_COUNT_FROM_LANDMARKS:face_count" # Defines whether landmarks on the previous image should be used to help # localize landmarks on the current image. @@ -208,6 +210,15 @@ node { output_stream: "ITERABLE:multi_face_landmarks" } +# Counting a multi_faceLandmarks vector size. The image stream is only used to +# make the calculator work even when there is no input vector. +node { + calculator: "CountingNormalizedLandmarkListVectorSizeCalculator" + input_stream: "CLOCK:image" + input_stream: "VECTOR:multi_face_landmarks" + output_stream: "COUNT:face_count" +} + # Collects a NormalizedRect for each face into a vector. Upon receiving the # BATCH_END timestamp, outputs the vector of NormalizedRect at the BATCH_END # timestamp.