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; 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));
        // 每张热图只有一个最大值
        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++)
    {
        // 查找两个连接的身体/手部部件
        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; 0;
}