H264DecoderNvCodecHelper.h

#pragma once
 
#include <assert.h>
#include <stdint.h>
#include <mutex>
#include <vector>
#include <string>
#include <iostream>
#include <sstream>
#include <string.h>
#include "nvcuvid.h"
#include <boost/shared_ptr.hpp>
#include "CommonDefs.h"
#include "CudaCommon.h"
#include <queue>

class NVDECException : public std::exception
{
public:
    NVDECException(const std::string& errorStr, const CUresult errorCode)
        : m_errorString(errorStr), m_errorCode(errorCode) {}
 
    virtual ~NVDECException() throw() {}
    virtual const char* what() const throw() { return m_errorString.c_str(); }
    CUresult  getErrorCode() const { return m_errorCode; }
    const std::string& getErrorString() const { return m_errorString; }
    static NVDECException makeNVDECException(const std::string& errorStr, const CUresult errorCode,
        const std::string& functionName, const std::string& fileName, int lineNo);
private:
    std::string m_errorString;
    CUresult m_errorCode;
};
 
inline NVDECException NVDECException::makeNVDECException(const std::string& errorStr, const CUresult errorCode, const std::string& functionName,
    const std::string& fileName, int lineNo)
{
    std::ostringstream errorLog;
    errorLog << functionName << " : " << errorStr << " at " << fileName << ":" << lineNo << std::endl;
    NVDECException exception(errorLog.str(), errorCode);
    return exception;
}
 
#define NVDEC_THROW_ERROR( errorStr, errorCode )                                                         \
    do                                                                                                   \
    {                                                                                                    \
        throw NVDECException::makeNVDECException(errorStr, errorCode, __FUNCTION__, __FILE__, __LINE__); \
    } while (0)
 
 
#define NVDEC_API_CALL( cuvidAPI )                                                                                 \
    do                                                                                                             \
    {                                                                                                              \
        CUresult errorCode = cuvidAPI;                                                                             \
        if( errorCode != CUDA_SUCCESS)                                                                             \
        {                                                                                                          \
            std::ostringstream errorLog;                                                                           \
            errorLog << #cuvidAPI << " returned error " << errorCode;                                              \
            throw NVDECException::makeNVDECException(errorLog.str(), errorCode, __FUNCTION__, __FILE__, __LINE__); \
        }                                                                                                          \
    } while (0)
 
struct Rect {
    int l, t, r, b;
};
 
struct Dim {
    int w, h;
};

class NvDecoder {
 
public:
    NvDecoder(CUcontext cuContext, int nWidth, int nHeight, bool bUseDeviceFrame, cudaVideoCodec eCodec, std::mutex* pMutex = NULL,
        bool bLowLatency = false, bool bDeviceFramePitched = false, const Rect* pCropRect = NULL, const Dim* pResizeDim = NULL, int maxWidth = 0, int maxHeight = 0);
    NvDecoder() {}
    ~NvDecoder() ;

    CUcontext GetContext() { return m_cuContext; }

    int GetWidth() { assert(m_nWidth); return m_nWidth; }

    int GetHeight() { assert(m_nHeight); return m_nHeight; }

    int GetFrameSize() { assert(m_nWidth); return m_nWidth * m_nHeight * 3 / (m_nBitDepthMinus8 ? 1 : 2); }

    int GetDeviceFramePitch() { assert(m_nWidth); return m_nDeviceFramePitch ? (int)m_nDeviceFramePitch : m_nWidth * (m_nBitDepthMinus8 ? 2 : 1); }

    int GetBitDepth() { assert(m_nWidth); return m_nBitDepthMinus8 + 8; }

    CUVIDEOFORMAT GetVideoFormatInfo() { assert(m_nWidth); return m_videoFormat; }

    std::string GetVideoInfo() const { return m_videoInfo.str(); }

    bool Decode(const uint8_t* pData, int nSize, uint8_t*** pppFrame, int* pnFrameReturned, uint32_t flags = 0, int64_t** ppTimestamp = NULL, int64_t timestamp = 0, CUstream stream = 0);

    bool DecodeLockFrame(const uint8_t* pData, int nSize, uint8_t*** pppFrame, int* pnFrameReturned, uint32_t flags = 0, int64_t** ppTimestamp = NULL, int64_t timestamp = 0, CUstream stream = 0);

    void UnlockFrame(uint8_t** ppFrame, int nFrame);

    int setReconfigParams(const Rect* pCropRect, const Dim* pResizeDim);
 
private:
    static int CUDAAPI HandleVideoSequenceProc(void* pUserData, CUVIDEOFORMAT* pVideoFormat) { return ((NvDecoder*)pUserData)->HandleVideoSequence(pVideoFormat); }

    static int CUDAAPI HandlePictureDecodeProc(void* pUserData, CUVIDPICPARAMS* pPicParams) { return ((NvDecoder*)pUserData)->HandlePictureDecode(pPicParams); }

    static int CUDAAPI HandlePictureDisplayProc(void* pUserData, CUVIDPARSERDISPINFO* pDispInfo) { return ((NvDecoder*)pUserData)->HandlePictureDisplay(pDispInfo); }

    int HandleVideoSequence(CUVIDEOFORMAT* pVideoFormat);

    int HandlePictureDecode(CUVIDPICPARAMS* pPicParams);

    int HandlePictureDisplay(CUVIDPARSERDISPINFO* pDispInfo);

    int ReconfigureDecoder(CUVIDEOFORMAT* pVideoFormat);
 
private:
    CUcontext m_cuContext = NULL;
    CUvideoctxlock m_ctxLock;
    std::mutex* m_pMutex;
    CUvideoparser m_hParser = NULL;
    CUvideodecoder m_hDecoder = NULL;
    bool m_bUseDeviceFrame;
    // dimension of the output
    unsigned int m_nWidth = 0, m_nHeight = 0;
    // height of the mapped surface 
    int m_nSurfaceHeight = 0;
    int m_nSurfaceWidth = 0;
    cudaVideoCodec m_eCodec = cudaVideoCodec_NumCodecs;
    cudaVideoChromaFormat m_eChromaFormat;
    int m_nBitDepthMinus8 = 0;
    CUVIDEOFORMAT m_videoFormat = {};
    Rect m_displayRect = {};
    // stock of frames
    std::vector<uint8_t*> m_vpFrame;
    // decoded frames for return
    std::vector<uint8_t*> m_vpFrameRet;
    // timestamps of decoded frames
    std::vector<int64_t> m_vTimestamp;
    int m_nDecodedFrame = 0, m_nDecodedFrameReturned = 0;
    int m_nDecodePicCnt = 0, m_nPicNumInDecodeOrder[32];
    bool m_bEndDecodeDone = false;
    std::mutex m_mtxVPFrame;
    int m_nFrameAlloc = 0;
    CUstream m_cuvidStream = 0;
    bool m_bDeviceFramePitched = false;
    size_t m_nDeviceFramePitch = 0;
    Rect m_cropRect = {};
    Dim m_resizeDim = {};
 
    std::ostringstream m_videoInfo;
    unsigned int m_nMaxWidth = 0, m_nMaxHeight = 0;
    bool m_bReconfigExternal = false;
    bool m_bReconfigExtPPChange = false;
};
 
class H264DecoderNvCodecHelper : public NvDecoder
{
public:
    H264DecoderNvCodecHelper(int mWidth, int mHieght);
    H264DecoderNvCodecHelper() {}
    ~H264DecoderNvCodecHelper() {}
 
    bool init(std::function<void(frame_sp&)> send, std::function<frame_sp()> makeFrame);
    void ConvertToPlanar(uint8_t* pHostFrame, int nWidth, int nHeight, int nBitDepth);
    void process(void* inputFrameBuffer, size_t inputFrameSize, uint64_t inputFrameTS);
    std::function<void( frame_sp&)> send;
    std::function<frame_sp()> makeFrame;
private:
    boost::shared_ptr<NvDecoder> helper;
    std::queue<uint64_t> framesTimestampEntry;
};