Add counting_vector_size_calculator

mediapipe/calculators/util/BUILD

@@ -311,6 +311,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"],

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<std::vector<::mediapipe::NormalizedLandmarkList>> 
+    CountingNormalizedLandmarkListVectorSizeCalculator;
+
+REGISTER_CALCULATOR(CountingNormalizedLandmarkListVectorSizeCalculator);
+}  // namespace mediapipe
+
+

mediapipe/calculators/util/counting_vector_size_calculator.h

@@ -0,0 +1,69 @@
+// 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 count input landmarksList size.
+//
+// Count input landmark(std::vector<NormalizedLandmarkList>) and return this 
+// value to ouput_stream. Input IMAGE has no effect on calculation, but is used to
+// ensure that the calculator works even when the landmark is empty. And if the 
+// input landmark is empty, the number of faces found is zero.
+//
+// Example config:
+// node {
+//   calculator: "CountingVectorSizeCalculator"
+//   input_stream: "IMAGE:input_image"
+//   input_stream: "LANDMARKS:multi_face_landmarks"
+//   output_stream: "COUNT:face_count"
+// }
+
+template <typename VectorT>
+class CountingVectorSizeCalculator : public CalculatorBase { 
+  public:
+    static ::mediapipe::Status GetContract(CalculatorContract* cc) {
+        // Check tag.
+        RET_CHECK(cc->Inputs().HasTag("CLOCK"));
+        cc->Inputs().Tag("CLOCK").SetAny();
+        RET_CHECK(cc->Inputs().HasTag("VECTOR"));
+        cc->Inputs().Tag("VECTOR").Set<VectorT>();
+        RET_CHECK(cc->Outputs().HasTag("COUNT"));
+        cc->Outputs().Tag("COUNT").Set<int>();
+
+        return ::mediapipe::OkStatus();
+    }
+
+    ::mediapipe::Status Process(CalculatorContext* cc) {
+        std::unique_ptr<int> face_count;
+        if (!cc->Inputs().Tag("VECTOR").IsEmpty()) {
+            const auto& landmarks = cc->Inputs().Tag("VECTOR").Get<VectorT>();
+            face_count = absl::make_unique<int>(landmarks.size());
+        } else {
+            face_count = absl::make_unique<int>(0);
+        }
+        cc->Outputs().Tag("COUNT").Add(face_count.release(), cc->InputTimestamp());
+
+        return ::mediapipe::OkStatus();
+    };    
+};
+
+}
+
+#endif  // MEDIAPIPE_CALCULATORS_UTIL_COUNTING_VECTOR_SIZE_CALCULATOR_H

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",

mediapipe/examples/desktop/demo_run_graph_main.cc

@@ -18,6 +18,7 @@
 #include "mediapipe/framework/calculator_framework.h"
 #include "mediapipe/framework/formats/image_frame.h"
 #include "mediapipe/framework/formats/image_frame_opencv.h"
+#include "mediapipe/framework/formats/landmark.pb.h"
 #include "mediapipe/framework/port/commandlineflags.h"
 #include "mediapipe/framework/port/file_helpers.h"
 #include "mediapipe/framework/port/opencv_highgui_inc.h"
@@ -26,8 +27,11 @@
 #include "mediapipe/framework/port/parse_text_proto.h"
 #include "mediapipe/framework/port/status.h"
 
+
 constexpr char kInputStream[] = "input_video";
-constexpr char kOutputStream[] = "output_video";
+constexpr char kOutputImageStream[] = "output_video";
+constexpr char kOutputFaceCountStream[] = "face_count";
+constexpr char kOutputLandmarksStream[] = "multi_face_landmarks";
 constexpr char kWindowName[] = "MediaPipe";
 
 DEFINE_string(
@@ -76,12 +80,17 @@ DEFINE_string(output_video_path, "",
   }
 
   LOG(INFO) << "Start running the calculator graph.";
-  ASSIGN_OR_RETURN(mediapipe::OutputStreamPoller poller,
-                   graph.AddOutputStreamPoller(kOutputStream));
+  ASSIGN_OR_RETURN(mediapipe::OutputStreamPoller image_poller,
+                   graph.AddOutputStreamPoller(kOutputImageStream));
+  ASSIGN_OR_RETURN(mediapipe::OutputStreamPoller face_count_poller,
+                   graph.AddOutputStreamPoller(kOutputFaceCountStream));
+  ASSIGN_OR_RETURN(mediapipe::OutputStreamPoller landmarks_poller,
+                   graph.AddOutputStreamPoller(kOutputLandmarksStream));
   MP_RETURN_IF_ERROR(graph.StartRun({}));
 
   LOG(INFO) << "Start grabbing and processing frames.";
   bool grab_frames = true;
+  int prev_face_count = -1;
   while (grab_frames) {
     // Capture opencv camera or video frame.
     cv::Mat camera_frame_raw;
@@ -108,9 +117,29 @@ DEFINE_string(output_video_path, "",
                           .At(mediapipe::Timestamp(frame_timestamp_us))));
 
     // Get the graph result packet, or stop if that fails.
-    mediapipe::Packet packet;
-    if (!poller.Next(&packet)) break;
-    auto& output_frame = packet.Get<mediapipe::ImageFrame>();
+    mediapipe::Packet image_packet;
+    if (!image_poller.Next(&image_packet)) break;
+    auto& output_frame = image_packet.Get<mediapipe::ImageFrame>();
+
+    mediapipe::Packet face_count_packet;
+    if (!face_count_poller.Next(&face_count_packet)) break;
+    auto& face_count = face_count_packet.Get<int>();
+
+    if (face_count != prev_face_count) {
+      LOG(INFO) << "Found face count : " << face_count;
+      prev_face_count = face_count;
+
+      if (face_count != 0) {
+        mediapipe::Packet landmarks_packet;
+        if (!landmarks_poller.Next(&landmarks_packet)) break;
+        auto& landmarks = landmarks_packet.Get<std::vector<::mediapipe::NormalizedLandmarkList>>();
+
+        LOG(INFO) << "Landmarks size : " << landmarks[0].landmark_size();
+
+        // Do something with landmarks
+      }
+    }
+    
 
     // Convert back to opencv for display or saving.
     cv::Mat output_frame_mat = mediapipe::formats::MatView(&output_frame);

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<NormalizedLandmarkList>)
 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.

mediapipe/modules/face_landmark/BUILD

@@ -75,6 +75,7 @@ mediapipe_simple_subgraph(
         "//mediapipe/calculators/util:association_norm_rect_calculator",
         "//mediapipe/calculators/util:collection_has_min_size_calculator",
         "//mediapipe/modules/face_detection:face_detection_front_cpu",
+        "//mediapipe/calculators/util:counting_vector_size_calculator",
     ],
 )
 

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<NormalizedRect>)
 output_stream: "ROIS_FROM_DETECTIONS:face_rects_from_detections"
+# (int)
+output_stream: "FACE_COUNT_FROM_LANDMARKS:face_count"
 
 # Determines if an input vector of NormalizedRect has a size greater than or
 # equal to the provided num_faces.
@@ -189,6 +191,14 @@ node {
   output_stream: "ITERABLE:multi_face_landmarks"
 }
 
+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.