diff --git a/demo/csrc/cpp/pose_tracker.cxx b/demo/csrc/cpp/pose_tracker.cxx index c50bdfbe14..484fe0ebd7 100644 --- 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 index 59b0df415a..1a73121676 100644 --- 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 index e9d8493f58..df93cfea8f 100644 --- 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 used(pts_size); std::vector is_end_point(pts_size); for (size_t i = 0; i < skel.links.size(); ++i) {