c_cpp Kinect v2坐标系统映射

Posted 2021-05-14

#include <iostream>
#include <sstream>

#include <Windows.h>
#include <Kinect.h>
#include <opencv2/opencv.hpp>

#include <atlbase.h>

// Error Check 
#define ERROR_CHECK( ret )                                        \
    if( (ret) != S_OK ){                                          \
        std::stringstream ss;                                     \
        ss << "failed " #ret " " << std::hex << ret << std::endl; \
        throw std::runtime_error( ss.str().c_str() );             \
    }

class Kinect
{
private:

    // Sensor
    CComPtr<IKinectSensor> kinect;
    CComPtr<ICoordinateMapper> mapper;

    // Color
    CComPtr<IColorFrameReader> colorFrameReader;
    std::vector<BYTE> colorBuffer;
    int colorWidth;
    int colorHeight;
    unsigned int colorBytesPerPixel;
    cv::Mat colorMat;

    // Depth
    CComPtr<IDepthFrameReader> depthFrameReader;
    std::vector<UINT16> depthBuffer;
    int depthWidth;
    int depthHeight;
    cv::Mat depthMat;

public:

    // Constructor
    Kinect()
    {
        // Initialize
        initialize();
    }

    ~Kinect()
    {
        // Finalize
        finalize();
    }

    // Run
    void run()
    {
        while ( 1 ) {
            // Update
            update();

            // Draw
            draw();

            // Show
            show();

            // Key Check
            const int key = cv::waitKey( 10 );
            if( key == VK_ESCAPE ){
                break;
            }
        }
    }

private:

    // Initialize
    void initialize()
    {
        cv::setUseOptimized( true );

        // Initialize Sensor
        initializeSensor();

        // Initialize Color
        initializeColor();

        // Initialize Depth
        initializeDepth();
    }

    // Initialize Sensor
    void initializeSensor()
    {
        // Open Sensor
        ERROR_CHECK( GetDefaultKinectSensor( &kinect ) );

        ERROR_CHECK( kinect->Open() );

        // Check Open
        BOOLEAN isOpen = FALSE;
        ERROR_CHECK( kinect->get_IsOpen( &isOpen ) );
        if( !isOpen ){
            throw std::runtime_error( "failed IKinectSensor::get_IsOpen( &isOpen )" );
        }

        // Retrieved Coordinate Mapper
        ERROR_CHECK( kinect->get_CoordinateMapper( &mapper ) );
    }

    // Initialize Color
    void initializeColor()
    {
        // Open Color Reader
        CComPtr<IColorFrameSource> colorFrameSource;
        ERROR_CHECK( kinect->get_ColorFrameSource( &colorFrameSource ) );
        ERROR_CHECK( colorFrameSource->OpenReader( &colorFrameReader ) );

        // Retrieved Color Description
        CComPtr<IFrameDescription> colorFrameDescription;
        ERROR_CHECK( colorFrameSource->CreateFrameDescription( ColorImageFormat::ColorImageFormat_Bgra, &colorFrameDescription ) );
        ERROR_CHECK( colorFrameDescription->get_Width( &colorWidth ) ); // 1920
        ERROR_CHECK( colorFrameDescription->get_Height( &colorHeight ) ); // 1080
        ERROR_CHECK( colorFrameDescription->get_BytesPerPixel( &colorBytesPerPixel ) ); // 4

        // Allocation Color Buffer
        colorBuffer.resize( colorWidth * colorHeight * colorBytesPerPixel );
    }

    // Initialize Depth
    void initializeDepth()
    {
        // Open Depth Reader
        CComPtr<IDepthFrameSource> depthFrameSource;
        ERROR_CHECK( kinect->get_DepthFrameSource( &depthFrameSource ) );
        ERROR_CHECK( depthFrameSource->OpenReader( &depthFrameReader ) );

        // Retrieved Depth Description
        CComPtr<IFrameDescription> depthFrameDescription;
        ERROR_CHECK( depthFrameSource->get_FrameDescription( &depthFrameDescription ) );
        ERROR_CHECK( depthFrameDescription->get_Width( &depthWidth ) ); // 512
        ERROR_CHECK( depthFrameDescription->get_Height( &depthHeight ) ); // 424

        // Allocation Depth Buffer
        depthBuffer.resize( depthWidth * depthHeight );
    }

    // Finalize
    void finalize()
    {
        // Close Sensor
        ERROR_CHECK( kinect->Close() );
    }

    // Update
    void update()
    {
        // Update Color
        updateColor();

        // Update Depth
        updateDepth();
    }

    // Update Color
    void updateColor()
    {
        // Retrieved Color Frame
        CComPtr<IColorFrame> colorFrame;
        const HRESULT ret = colorFrameReader->AcquireLatestFrame( &colorFrame );
        if( FAILED( ret ) ){
            return;
        }

        // Copy Data and Convert Format ( YUY2 -> BGRA )
        ERROR_CHECK( colorFrame->CopyConvertedFrameDataToArray( static_cast<UINT>( colorBuffer.size() ), &colorBuffer[0], ColorImageFormat::ColorImageFormat_Bgra ) );

    }

    // Update Depth
    void updateDepth()
    {
        // Retrieved Depth Frame
        CComPtr<IDepthFrame> depthFrame;
        const HRESULT ret = depthFrameReader->AcquireLatestFrame( &depthFrame );
        if( FAILED( ret ) ){
            return;
        }

        // Copy Data
        ERROR_CHECK( depthFrame->CopyFrameDataToArray( depthBuffer.size(), &depthBuffer[0] ) );
    }

    // Draw
    void draw()
    {
        // Draw Color
        drawColor();

        // Draw Depth
        drawDepth();
    }

    // Draw Color
    void drawColor()
    {
        // Retrieve Mapped Coordinates
        std::vector<ColorSpacePoint> colorSpace( depthWidth * depthHeight );
        ERROR_CHECK( mapper->MapDepthFrameToColorSpace( depthBuffer.size(), &depthBuffer[0], colorSpace.size(), &colorSpace[0] ) );

        // Mapping Depth to Color Resolution
        std::vector<BYTE> buffer( depthWidth * depthHeight * colorBytesPerPixel );

        for( int depthY = 0; depthY < depthHeight; depthY++ ){
            for( int depthX = 0; depthX < depthWidth; depthX++ ){
                unsigned int depthIndex = depthY * depthWidth + depthX;
                int colorX = static_cast<int>( colorSpace[depthIndex].X + 0.5f );
                int colorY = static_cast<int>( colorSpace[depthIndex].Y + 0.5f );
                if( ( 0 <= colorX ) && ( colorX < colorWidth ) && ( 0 <= colorY ) && ( colorY < colorHeight ) ){
                    unsigned int colorIndex = colorY * colorWidth + colorX;
                    buffer[depthIndex * colorBytesPerPixel + 0] = colorBuffer[colorIndex * colorBytesPerPixel + 0];
                    buffer[depthIndex * colorBytesPerPixel + 1] = colorBuffer[colorIndex * colorBytesPerPixel + 1];
                    buffer[depthIndex * colorBytesPerPixel + 2] = colorBuffer[colorIndex * colorBytesPerPixel + 2];
                    buffer[depthIndex * colorBytesPerPixel + 3] = colorBuffer[colorIndex * colorBytesPerPixel + 3];
                }
            }
        }

        // Set Coordinate Buffer to cv::Mat
        colorMat = cv::Mat( depthHeight, depthWidth, CV_8UC4, &buffer[0] ).clone();
    }

    // Draw Depth
    void drawDepth()
    {
        // Retrieve Mapped Coordinates
        std::vector<DepthSpacePoint> depthSpace( colorWidth * colorHeight );
        ERROR_CHECK( mapper->MapColorFrameToDepthSpace( depthBuffer.size(), &depthBuffer[0], depthSpace.size(), &depthSpace[0] ) );

        // Mapping Depth to Color Resolution
        std::vector<UINT16> buffer( colorWidth * colorHeight );

        for ( int colorY = 0; colorY < colorHeight; colorY++ ){
            for ( int colorX = 0; colorX < colorWidth; colorX++ ){
                unsigned int colorIndex = colorY * colorWidth + colorX;
                int depthX = static_cast<int>( depthSpace[colorIndex].X + 0.5f );
                int depthY = static_cast<int>( depthSpace[colorIndex].Y + 0.5f );
                if ( ( 0 <= depthX ) && ( depthX < depthWidth ) && ( 0 <= depthY ) && ( depthY < depthHeight ) ){
                    unsigned int depthIndex = depthY * depthWidth + depthX;
                    buffer[colorIndex] = depthBuffer[depthIndex];
                }
            }
        }

        // Set Coordinate Buffer to cv::Mat
        cv::Mat rawMat = cv::Mat( colorHeight, colorWidth, CV_16UC1, &buffer[0] ).clone();

        // Convert Raw(16bit) to 8bit
        rawMat.convertTo( depthMat, CV_8U, -255.0f / 8000.0f, 255.0f ); // 255(white) - 0(black)
        //rawMat.convertTo( depthMat, CV_8U, 255.0f / 8000.0f, 0.0f ); // 0(black) - 255(white)
    }

    // Show
    void show()
    {
        // Show Color
        showColor();

        // Show Depth
        showDepth();
    }

    // Show Color
    void showColor()
    {
        // Show Color Data
        cv::imshow( "Color", colorMat );
    }

    // Show Depth
    void showDepth()
    {
        // Resize Image
        const double scale = 0.5;
        cv::Mat resizeMat;
        cv::resize( depthMat, resizeMat, cv::Size(), scale, scale );

        // Show Depth Data
        cv::imshow( "Depth", resizeMat );
    }
};

void main()
{
    try {
        Kinect kinect;
        kinect.run();
    }
    catch ( std::exception& ex ){
        std::cout << ex.what() << std::endl;
    }
}