add coco whole-body skeleton (#1736)

open-mmlab · Feb 9, 2023 · 31b099a · 31b099a
1 parent a3311b0
commit 31b099a
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 3 deletions.
diff --git a/demo/csrc/cpp/pose_tracker.cxx b/demo/csrc/cpp/pose_tracker.cxx
@@ -17,7 +17,8 @@ DEFINE_int32(output_size, 0, "Long-edge of output frames");
 DEFINE_int32(flip, 0, "Set to 1 for flipping the input horizontally");
 DEFINE_int32(show, 1, "Delay passed to `cv::waitKey` when using `cv::imshow`; -1: disable");
 
-DEFINE_string(skeleton, "coco", R"(Path to skeleton data or name of predefined skeletons: "coco")");
+DEFINE_string(skeleton, "coco",
+              R"(Path to skeleton data or name of predefined skeletons: "coco", "coco-wholebody")");
 DEFINE_string(background, "default",
               R"(Output background, "default": original image, "black": black background)");
 

diff --git a/demo/csrc/cpp/utils/skeleton.h b/demo/csrc/cpp/utils/skeleton.h
@@ -19,7 +19,7 @@ struct Skeleton {
   static Skeleton get(const std::string& path);
 };
 
-const Skeleton& gCocoSkeleton() {
+const Skeleton& gSkeletonCoco() {
   static const Skeleton inst{
       {
           {15, 13}, {13, 11}, {16, 14}, {14, 12}, {11, 12}, {5, 11}, {6, 12},
@@ -38,6 +38,41 @@ const Skeleton& gCocoSkeleton() {
   return inst;
 }
 
+const Skeleton& gSkeletonCocoWholeBody() {
+  static const Skeleton inst{
+      {
+          {15, 13},   {13, 11},   {16, 14},   {14, 12},   {11, 12},   {5, 11},    {6, 12},
+          {5, 6},     {5, 7},     {6, 8},     {7, 9},     {8, 10},    {1, 2},     {0, 1},
+          {0, 2},     {1, 3},     {2, 4},     {3, 5},     {4, 6},     {15, 17},   {15, 18},
+          {15, 19},   {16, 20},   {16, 21},   {16, 22},   {91, 92},   {92, 93},   {93, 94},
+          {94, 95},   {91, 96},   {96, 97},   {97, 98},   {98, 99},   {91, 100},  {100, 101},
+          {101, 102}, {102, 103}, {91, 104},  {104, 105}, {105, 106}, {106, 107}, {91, 108},
+          {108, 109}, {109, 110}, {110, 111}, {112, 113}, {113, 114}, {114, 115}, {115, 116},
+          {112, 117}, {117, 118}, {118, 119}, {119, 120}, {112, 121}, {121, 122}, {122, 123},
+          {123, 124}, {112, 125}, {125, 126}, {126, 127}, {127, 128}, {112, 129}, {129, 130},
+          {130, 131}, {131, 132},
+      },
+      {
+          {51, 153, 255},
+          {0, 255, 0},
+          {255, 128, 0},
+          {255, 255, 255},
+          {255, 153, 255},
+          {102, 178, 255},
+          {255, 51, 51},
+      },
+      {1, 1, 2, 2, 0, 0, 0, 0, 1, 2, 1, 2, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1,
+       2, 2, 2, 2, 2, 2, 2, 4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6, 1, 1, 1,
+       1, 2, 2, 2, 2, 4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6, 1, 1, 1, 1},
+      {0, 0, 0, 0, 0, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3,
+       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
+       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3,
+       3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6,
+       1, 1, 1, 1, 3, 2, 2, 2, 2, 4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6, 1, 1, 1, 1},
+  };
+  return inst;
+}
+
 // n_links
 // u0, v0, u1, v1, ..., un-1, vn-1
 // n_palette
@@ -48,7 +83,9 @@ const Skeleton& gCocoSkeleton() {
 // j0, j1, ..., jn-1
 inline Skeleton Skeleton::get(const std::string& path) {
   if (path == "coco") {
-    return gCocoSkeleton();
+    return gSkeletonCoco();
+  } else if (path == "coco-wholebody") {
+    return gSkeletonCocoWholeBody();
   }
   std::ifstream ifs(path);
   if (!ifs.is_open()) {

diff --git a/demo/csrc/cpp/utils/visualize.h b/demo/csrc/cpp/utils/visualize.h
@@ -198,6 +198,11 @@ class Visualize {
 
     void add_pose(const mmdeploy_point_t* pts, const float* scores, int32_t pts_size, double thr) {
       auto& skel = v_.skeleton_;
+      if (skel.point_colors.size() != pts_size) {
+        std::cout << "error: mismatched number of keypoints: " << skel.point_colors.size() << " vs "
+                  << pts_size << ", skip pose visualization.\n";
+        return;
+      }
       std::vector<int> used(pts_size);
       std::vector<int> is_end_point(pts_size);
       for (size_t i = 0; i < skel.links.size(); ++i) {