KeepTopNPeople in PoseExtractor

include/openpose/pose/poseExtractor.hpp

@@ -3,6 +3,7 @@
 
 #include <openpose/core/common.hpp>
 #include <openpose/core/enumClasses.hpp>
+#include <openpose/core/keepTopNPeople.hpp>
 #include <openpose/pose/poseParameters.hpp>
 #include <openpose/pose/poseExtractorNet.hpp>
 #include <openpose/experimental/tracking/personIdExtractor.hpp>
@@ -14,6 +15,7 @@ namespace op
     {
     public:
         PoseExtractor(const std::shared_ptr<PoseExtractorNet>& poseExtractorNet,
+                      const std::shared_ptr<KeepTopNPeople>& keepTopNPeople = nullptr,
                       const std::shared_ptr<PersonIdExtractor>& personIdExtractor = nullptr,
                       const std::vector<std::shared_ptr<PersonTracker>>& personTracker = {},
                       const int numberPeopleMax = -1, const int tracking = -1);
@@ -38,6 +40,9 @@ namespace op
 
         float getScaleNetToOutput() const;
 
+        // KeepTopNPeople functions
+        void keepTopPeople(Array<float>& poseKeypoints, const Array<float>& poseScores) const;
+
         // PersonIdExtractor functions
         // Not thread-safe
         Array<long long> extractIds(const Array<float>& poseKeypoints, const cv::Mat& cvMatInput,
@@ -60,6 +65,7 @@ namespace op
         const int mNumberPeopleMax;
         const int mTracking;
         const std::shared_ptr<PoseExtractorNet> spPoseExtractorNet;
+        const std::shared_ptr<KeepTopNPeople> spKeepTopNPeople;
         const std::shared_ptr<PersonIdExtractor> spPersonIdExtractor;
         const std::vector<std::shared_ptr<PersonTracker>> spPersonTrackers;
 

include/openpose/pose/wPoseExtractor.hpp

@@ -70,6 +70,8 @@ namespace op
                     tDatum.poseKeypoints = spPoseExtractor->getPoseKeypoints().clone();
                     tDatum.poseScores = spPoseExtractor->getPoseScores().clone();
                     tDatum.scaleNetToOutput = spPoseExtractor->getScaleNetToOutput();
+                    // Keep desired top N people
+                    spPoseExtractor->keepTopPeople(tDatum.poseKeypoints, tDatum.poseScores);
                     // ID extractor (experimental)
                     tDatum.poseIds = spPoseExtractor->extractIdsLockThread(tDatum.poseKeypoints, tDatum.cvInputData,
                                                                            i, tDatum.id);

include/openpose/wrapper/wrapper.hpp

@@ -606,6 +606,15 @@ namespace op
                     spWPoseExtractors.resize(poseExtractorNets.size());
                     const auto personIdExtractor = (wrapperStructPose.identification
                         ? std::make_shared<PersonIdExtractor>() : nullptr);
+                    // Keep top N people
+                    // Added right after PoseExtractorNet to avoid:
+                    // 1) Rendering people that are later deleted (wrong visualization).
+                    // 2) Processing faces and hands on people that will be deleted (speed up).
+                    // 3) Running tracking before deleting the people.
+                    // Add KeepTopNPeople for each PoseExtractorNet
+                    const auto keepTopNPeople = (wrapperStructPose.numberPeopleMax > 0 ?
+                        std::make_shared<KeepTopNPeople>(wrapperStructPose.numberPeopleMax)
+                        : nullptr);
                     // Person tracker
                     std::vector<std::shared_ptr<PersonTracker>> personTrackers;
                     if (wrapperStructPose.tracking > -1)
@@ -613,10 +622,10 @@ namespace op
                                               std::make_shared<PersonTracker>(wrapperStructPose.tracking == 0));
                     for (auto i = 0u; i < spWPoseExtractors.size(); i++)
                     {
-
-                        // OpenPose keypoint detector + ID extractor (experimental) + tracking (experimental)
+                        // OpenPose keypoint detector + keepTopNPeople
+                        //    + ID extractor (experimental) + tracking (experimental)
                         const auto poseExtractor = std::make_shared<PoseExtractor>(
-                            poseExtractorNets.at(i), personIdExtractor, personTrackers,
+                            poseExtractorNets.at(i), keepTopNPeople, personIdExtractor, personTrackers,
                             wrapperStructPose.numberPeopleMax, wrapperStructPose.tracking);
                         spWPoseExtractors.at(i) = {std::make_shared<WPoseExtractor<TDatumsPtr>>(poseExtractor)};
                         // // Just OpenPose keypoint detector
@@ -624,17 +633,18 @@ namespace op
                         //     poseExtractorNets.at(i))};
                     }
 
-                    // Added right after PoseExtractorNet to avoid:
-                    // 1) Rendering people that are later deleted (wrong visualization).
-                    // 2) Processing faces and hands on people that will be deleted (speed up).
-                    if (wrapperStructPose.numberPeopleMax > 0)
-                    {
-                        // Add KeepTopNPeople for each PoseExtractorNet
-                        const auto keepTopNPeople = std::make_shared<KeepTopNPeople>(
-                            wrapperStructPose.numberPeopleMax);
-                        for (auto& wPose : spWPoseExtractors)
-                            wPose.emplace_back(std::make_shared<WKeepTopNPeople<TDatumsPtr>>(keepTopNPeople));
-                    }
+                    // // (Before tracking / id extractor)
+                    // // Added right after PoseExtractorNet to avoid:
+                    // // 1) Rendering people that are later deleted (wrong visualization).
+                    // // 2) Processing faces and hands on people that will be deleted (speed up).
+                    // if (wrapperStructPose.numberPeopleMax > 0)
+                    // {
+                    //     // Add KeepTopNPeople for each PoseExtractorNet
+                    //     const auto keepTopNPeople = std::make_shared<KeepTopNPeople>(
+                    //         wrapperStructPose.numberPeopleMax);
+                    //     for (auto& wPose : spWPoseExtractors)
+                    //         wPose.emplace_back(std::make_shared<WKeepTopNPeople<TDatumsPtr>>(keepTopNPeople));
+                    // }
                 }
 
 

src/openpose/pose/poseExtractor.cpp

@@ -6,12 +6,14 @@ namespace op
                                      " `--number_people_max 1` in order to run the person tracker (`--tracking`).";
 
     PoseExtractor::PoseExtractor(const std::shared_ptr<PoseExtractorNet>& poseExtractorNet,
+                                 const std::shared_ptr<KeepTopNPeople>& keepTopNPeople,
                                  const std::shared_ptr<PersonIdExtractor>& personIdExtractor,
                                  const std::vector<std::shared_ptr<PersonTracker>>& personTrackers,
                                  const int numberPeopleMax, const int tracking) :
         mNumberPeopleMax{numberPeopleMax},
         mTracking{tracking},
         spPoseExtractorNet{poseExtractorNet},
+        spKeepTopNPeople{keepTopNPeople},
         spPersonIdExtractor{personIdExtractor},
         spPersonTrackers{personTrackers}
     {
@@ -116,6 +118,20 @@ namespace op
         }
     }
 
+    void PoseExtractor::keepTopPeople(Array<float>& poseKeypoints, const Array<float>& poseScores) const
+    {
+        try
+        {
+            // Keep only top N people
+            if (spKeepTopNPeople)
+                poseKeypoints = spKeepTopNPeople->keepTopPeople(poseKeypoints, poseScores);
+        }
+        catch (const std::exception& e)
+        {
+            error(e.what(), __LINE__, __FUNCTION__, __FILE__);
+        }
+    }
+
     Array<long long> PoseExtractor::extractIds(const Array<float>& poseKeypoints, const cv::Mat& cvMatInput,
                                                const unsigned long long imageViewIndex)
     {