From cdd3835d13667c78d129a6fbe2083e7780560e14 Mon Sep 17 00:00:00 2001
From: Maksym Walczak <maksym@nimagna.com>
Date: Tue, 9 Aug 2022 14:27:36 +0200
Subject: [PATCH] Add the possibility to get landmark visibility

---
 mediapipe/pose_tracking_dll/README.md         | 15 +++++++++++++++
 mediapipe/pose_tracking_dll/pose_tracking.cpp |  9 ++++++++-
 mediapipe/pose_tracking_dll/pose_tracking.h   |  1 +
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/mediapipe/pose_tracking_dll/README.md b/mediapipe/pose_tracking_dll/README.md
index 6d886edcf..249c0ae15 100644
--- a/mediapipe/pose_tracking_dll/README.md
+++ b/mediapipe/pose_tracking_dll/README.md
@@ -77,6 +77,21 @@ bazel build -c opt --define MEDIAPIPE_DISABLE_GPU=1 --action_env PYTHON_BIN_PATH
 
 The results will be stored in the `bazel-bin\mediapipe\pose_tracking_dll` folder.
 
+In the case the build stalls, pressing Ctrl+C might not be sufficient to stop the task. In that case, if you try to (resume the) build again,
+the following message will be displayed:
+
+```
+Another command (pid=5300) is running. Waiting for it to complete on the server (server_pid=3684)
+```
+
+Unfortunately this process is hidden for some reason and can't be found in taskmgr. Fortunately, you can use the `taskkill` command to kill the process:
+
+```
+taskkill /F /PID 3684
+```
+
+After that, you should be able to run the build command again.
+
 ### Build debug symbols
 - `dbg` can be used in place of `opt` to build the library with debug symbols in Visual Studio pdb format.
 
diff --git a/mediapipe/pose_tracking_dll/pose_tracking.cpp b/mediapipe/pose_tracking_dll/pose_tracking.cpp
index 70a598fee..4267a0fc9 100644
--- a/mediapipe/pose_tracking_dll/pose_tracking.cpp
+++ b/mediapipe/pose_tracking_dll/pose_tracking.cpp
@@ -116,6 +116,7 @@ class PoseTrackingImpl {
       poseLandmarks[j].x = landmark.x();
       poseLandmarks[j].y = landmark.y();
       poseLandmarks[j].z = landmark.z();
+	  visibility[j] = landmark.visibility();
     }
 
     return absl::OkStatus();
@@ -124,6 +125,7 @@ class PoseTrackingImpl {
   nimagna::cv_wrapper::Point3f* lastDetectedLandmarks() { return poseLandmarks; }
 
   cv::Mat lastSegmentedFrame() { return segmentedMask; }
+  float* landmarksVisibility() { return visibility; }
 
   static constexpr size_t kLandmarksCount = 33u;
 
@@ -131,12 +133,13 @@ class PoseTrackingImpl {
   mediapipe::Packet poseLandmarksPacket;
   cv::Mat segmentedMask;
   nimagna::cv_wrapper::Point3f poseLandmarks[kLandmarksCount];
+  float visibility[kLandmarksCount] = {0};
   std::unique_ptr<mediapipe::OutputStreamPoller> maskPollerPtr;
   std::unique_ptr<mediapipe::OutputStreamPoller> landmarksPollerPtr;
   mediapipe::CalculatorGraph graph;
   const char* kInputStream = "input_video";
   const char* kOutputSegmentationStream = "segmentation_mask";
-  const char* kOutpuLandmarksStream = "pose_landmarks";
+  const char* kOutpuLandmarksStream = "pose_world_landmarks";
 };
 
 namespace nimagna {
@@ -153,6 +156,10 @@ cv_wrapper::Point3f* PoseTracking::lastDetectedLandmarks() {
   return mImplementation->lastDetectedLandmarks();
 }
 
+float* PoseTracking::lastLandmarksVisibility() {
+  return mImplementation->landmarksVisibility();
+}
+
 cv_wrapper::Mat PoseTracking::lastSegmentedFrame() {
   const cv::Mat result = mImplementation->lastSegmentedFrame();
 
diff --git a/mediapipe/pose_tracking_dll/pose_tracking.h b/mediapipe/pose_tracking_dll/pose_tracking.h
index 536547e06..78b93b0b5 100644
--- a/mediapipe/pose_tracking_dll/pose_tracking.h
+++ b/mediapipe/pose_tracking_dll/pose_tracking.h
@@ -108,6 +108,7 @@ class DLLEXPORT PoseTracking {
   bool processFrame(const cv_wrapper::Mat& inputRGB8Bit);
   cv_wrapper::Mat lastSegmentedFrame();
   cv_wrapper::Point3f* lastDetectedLandmarks();
+  float* lastLandmarksVisibility();
 
  private:
   PoseTrackingImpl* mImplementation;