samples/dnn/openpose.cpp

//
// 这个示例演示了如何使用带有 OpenCV 的 dnn 模块的预训练 OpenPose 网络。
//
// 它可以用于人体姿势检测，使用 COCO 模型（18 个部位）
// http://posefs1.perception.cs.cmu.edu/OpenPose/models/pose/coco/pose_iter_440000.caffemodel
// https://raw.githubusercontent.com/opencv/opencv_extra/4.x/testdata/dnn/openpose_pose_coco.prototxt
//
// 或者 MPI 模型（16 个部位）
// http://posefs1.perception.cs.cmu.edu/OpenPose/models/pose/mpi/pose_iter_160000.caffemodel
// https://raw.githubusercontent.com/opencv/opencv_extra/4.x/testdata/dnn/openpose_pose_mpi_faster_4_stages.prototxt
//
//（为了简化这个示例，身体模型被限制为单个人。）
//
//
// 你也可以尝试手部姿势模型
// http://posefs1.perception.cs.cmu.edu/OpenPose/models/hand/pose_iter_102000.caffemodel
// https://raw.githubusercontent.com/CMU-Perceptual-Computing-Lab/openpose/master/models/hand/pose_deploy.prototxt
//
 
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
using namespace cv;
using namespace cv::dnn;
 
#include <iostream>
using namespace std;
 
 
// 连接表，格式为 [模型 ID][对 ID][从/到]
// 请查看底部的详细解释
// https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/output.md
//
const int POSE_PAIRS[3][20][2] = {
{ // COCO 身体
    {1,2}, {1,5}, {2,3},
    {3,4}, {5,6}, {6,7},
    {1,8}, {8,9}, {9,10},
    {1,11}, {11,12}, {12,13},
    {1,0}, {0,14},
    {14,16}, {0,15}, {15,17}
},
{ // MPI 身体
    {0,1}, {1,2}, {2,3},
    {3,4}, {1,5}, {5,6},
    {6,7}, {1,14}, {14,8}, {8,9},
    {9,10}, {14,11}, {11,12}, {12,13}
},
{ // 手
{0,1}, {1,2}, {2,3}, {3,4}, // 大拇指
{0,5}, {5,6}, {6,7}, {7,8}, // 小拇指
{0,9}, {9,10}, {10,11}, {11,12}, // 中指
{0,13}, {13,14}, {14,15}, {15,16}, // 无名指
{0,17}, {17,18}, {18,19}, {19,20} // 小指
}};
 
int main(int argc, char **argv)
{
    CommandLineParser parser(argc, argv,
        "{ h help | false | 打印此帮助消息 }"
        "{ p proto | | （必需）模型配置，例如 hand/pose.prototxt }"
        "{ m model | | （必需）模型权重，例如 hand/pose_iter_102000.caffemodel }"
        "{ i image | | （必需）图像文件的路径（包含单个人或手）}"
        "{ d dataset | | 指定训练的模型类型。它可以是 (COCO, MPI, HAND)，具体取决于数据集。}"
        "{ width | 368 | 通过调整到特定宽度来预处理输入图像。}"
        "{ height | 368 | 通过调整到特定高度来预处理输入图像。}"
        "{ t threshold | 0.1 | 热图的阈值或置信度值 }"
        "{ s scale | 0.003922 | blob 的比例 }"
    );
 
    String modelTxt = samples::findFile(parser.get<string>("proto"));
    String modelBin = samples::findFile(parser.get<string>("model"));
    String imageFile = samples::findFile(parser.get<String>("image"));
    String dataset = parser.get<String>("dataset");
    int W_in = parser.get<int>("width");
    int H_in = parser.get<int>("height");
    float thresh = parser.get<float>("threshold");
    float scale = parser.get<float>("scale");
 
    if (parser.get<bool>("help") || modelTxt.empty() || modelBin.empty() || imageFile.empty())
    {
cout << "一个演示使用预训练的 OpenPose dnn 检测人体或手部姿势的示例应用程序。" << endl;
parser.printMessage();
        return 0;
    }
 
    int midx, npairs, nparts;
         if (!dataset.compare("COCO")) { midx = 0; npairs = 17; nparts = 18; }
    else if (!dataset.compare("MPI")) { midx = 1; npairs = 14; nparts = 16; }
    else if (!dataset.compare("HAND")) { midx = 2; npairs = 20; nparts = 22; }
    else
    {
std::cerr << "无法解释数据集参数：" << dataset << std::endl;
exit(-1);
    }
 
    // 读取网络模型
    Net net = readNet(modelBin, modelTxt);
    // 以及图像
    Mat img = imread(imageFile);
    if (img.empty())
    {
std::cerr << "无法从文件读取图像：" << imageFile << std::endl;
exit(-1);
    }
 
    // 将其发送到网络
    Mat inputBlob = blobFromImage(img, scale, Size(W_in, H_in), Scalar(0, 0, 0), false, false);
net.setInput(inputBlob);
    Mat result = net.forward();
    // 结果是“热图”数组，即身体部位位于 x,y 位置的概率
 
    int H = result.size[2];
    int W = result.size[3];
 
    // 找到身体部位的位置
vector<Point> points(22);
    for (int n=0; n<nparts; n++)
    {
        // 对相应身体部位的热图进行切片。
        Mat heatMap(H, W, CV_32F, result.ptr(0,n));
        // 每个热图 1 个最大值
        Point p(-1,-1),pm;
        double conf;
        minMaxLoc(heatMap, 0, &conf, 0, &pm);
        if (conf > thresh)
p = pm;
points[n] = p;
    }
 
    // 连接身体部位并绘制它！
    float SX = float(img.cols) / W;
    float SY = float(img.rows) / H;
    for (int n=0; n<npairs; n++)
    {
        // 查找 2 个连接的身体/手部部位
        Point2f a = points[POSE_PAIRS[midx][n][0]];
        Point2f b = points[POSE_PAIRS[midx][n][1]];
 
        // 我们之前没有找到足够的置信度
        if (a.x<=0 || a.y<=0 || b.x<=0 || b.y<=0)
            continue;
 
        // 缩放到图像大小
a.x*=SX; a.y*=SY;
b.x*=SX; b.y*=SY;
 
        line(img, a, b, Scalar(0,200,0), 2);
        circle(img, a, 3, Scalar(0,0,200), -1);
        circle(img, b, 3, Scalar(0,0,200), -1);
    }
 
    imshow("OpenPose", img);
    waitKey();
 
    return 0;
}