NiTE2(Beta) 入門(2) :スケルトンを追跡する

続いてスケルトンの追跡です。

スケルトンもKinect SDKっぽくなっています。ユーザーの認識から追跡を手動で開始させるのは相変わらずですが。

#include <iostream>

#include<opencv2\opencv.hpp>

#include <NiTE.h>

cv::Scalar colors[] = {
    cv::Scalar( 1, 0, 0 ),
    cv::Scalar( 0, 1, 0 ),
    cv::Scalar( 0, 0, 1 ),
    cv::Scalar( 1, 1, 0 ),
    cv::Scalar( 1, 0, 1 ),
    cv::Scalar( 0, 1, 1 ),
};

cv::Mat drawUser( nite::UserTrackerFrameRef& userFrame )
{
    cv::Mat depthImage;

    openni::VideoFrameRef depthFrame = userFrame.getDepthFrame();
    if ( depthFrame.isValid() ) {
        openni::VideoMode videoMode = depthFrame.getVideoMode();
        depthImage = cv::Mat( videoMode.getResolutionY(),
                                videoMode.getResolutionX(),
                                CV_8UC4 );

        openni::DepthPixel* depth = (openni::DepthPixel*)depthFrame.getData();
        const nite::UserId* pLabels = userFrame.getUserMap().getPixels();
        for (int i = 0; i < (depthFrame.getDataSize()/sizeof(openni::DepthPixel)); ++i) {
            // 画像インデックスを生成
            int index = i * 4;

            // 距離データを画像化する
            UCHAR* data = &depthImage.data[index];
            if ( pLabels[i] != 0 ) {
                data[0] *= colors[pLabels[i]][0];
                data[1] *= colors[pLabels[i]][1];
                data[2] *= colors[pLabels[i]][2];
            }
            else {
                // 0-255のグレーデータを作成する
                // distance : 10000 = gray : 255
                int gray = ~((depth[i] * 255) / 10000);
                data[0] = gray;
                data[1] = gray;
                data[2] = gray;
            }
        }
    }

    return depthImage;
}

void drawSkeleton( cv::Mat& depthImage, nite::UserTracker& userTracker,
                   nite::UserTrackerFrameRef& userFrame )
{
    const nite::Array<nite::UserData>& users = userFrame.getUsers();
    for ( int i = 0; i < users.getSize(); ++i ) {
        const nite::UserData& user = users[i];
        if ( user.isNew() ) {
            userTracker.startSkeletonTracking( user.getId() );
        }
        else if ( !user.isLost() ) {
            const nite::Skeleton& skeelton = user.getSkeleton();
            if ( skeelton.getState() == nite::SkeletonState::SKELETON_TRACKED ) {
                for ( int j = 0; j <= 14; ++j ) {
                    const nite::SkeletonJoint& joint = skeelton.getJoint((nite::JointType)j);
                    if ( joint.getPositionConfidence() >= 0.7f ) {
                        const nite::Point3f& position = joint.getPosition();
                        float x = 0, y = 0;
                        userTracker.convertJointCoordinatesToDepth(
                            position.x, position.y, position.z, &x, &y );

                        cv::circle( depthImage, cvPoint( (int)x, (int)y ),
                            5, cv::Scalar( 0, 0, 255 ), 5 );
                    }
                }
            }
        }
    }
}

void main()
{
    try {
        nite::Status niteRet = nite::NiTE::initialize();

        nite::UserTracker userTracker;
        niteRet = userTracker.create();
        if ( niteRet != nite::STATUS_OK ) {
            throw std::runtime_error( "userTracker.create" );
        }

        cv::Mat depthImage;

        while ( 1 ) {
            nite::UserTrackerFrameRef userFrame;
            userTracker.readFrame( &userFrame );

            depthImage = drawUser( userFrame );
            drawSkeleton( depthImage, userTracker, userFrame );

            cv::imshow( "Skeleton", depthImage );

            int key = cv::waitKey( 10 );
            if ( key == 'q' ) {
                break;
            }
        }
    }
    catch ( std::exception& ) {
        std::cout << openni::OpenNI::getExtendedError() << std::endl;
    }
}

前回のコードに drawSkeleton() を追加しました。

追跡しているユーザーの情報(nite::UserDataの配列)がnite::UserTrackerFrameRef::getUsers() で取得できるので、各ユーザーの状態によって処理を行います。

新しく追跡を開始したユーザーは nite::UserData::isNew() が true を返すので、nite::UserTracker::startSkeletonTracking() を呼び出しスケルトンの追跡を開始します。

以降のフレームからユーザーを見失うまで nite::UserData::isLost() はfalseを返し続けるので、その間はスケルトンの描画を行います。スケルトンはnite::UserData::getSkeleton() でnite::Skeleton型のスケルトン情報が取得できます。nite::Skeleton::getState() で状態が取得できるので、追跡している(nite::SkeletonState::SKELETON_TRACKEDである)場合はジョイント(関節)の描画を行います。

ジョイントは15点認識します。各ジョイントの情報はnite::Skeleton::getJoint() でnite::SkeletonJoint型として取得できます。nite::SkeletonJoint::getPositionConfidence() でジョイント座標の信頼性が取得できるので、ある一定値以上でフィルタし描画します。ジョイントの座標はnite::SkeletonJoint::getPosition()で取得できますが、これは3次元座標のため、nite::UserTracker::convertJointCoordinatesToDepth() でDepthの2次元座標に変換して表示します。

このように、スケルトンの扱いはKinect SDKと似ていますので扱いやすいです。スケルトンもデバイスごとの認識人数を調べる必要があります。