Skip to content

Commit f5f4216

Browse files
asmorkalovasmorkalov
authored
Merge pull request #3784 from cudawarped:cudacodec_videowriter_fix_ts_in_container
`cudacodec`: update `VideoWriter` to correctly set the PTS inside the encapsulated video stream
2 parents 9bfa192 + 234297d commit f5f4216

File tree

5 files changed

+69
-19
lines changed

5 files changed

+69
-19
lines changed

‎modules/cudacodec/include/opencv2/cudacodec.hpp‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -208,8 +208,15 @@ class CV_EXPORTS_W EncoderCallback {
208208
/** @brief Callback function to signal that the encoded bitstream for one or more frames is ready.
209209
210210
@param vPacket The raw bitstream for one or more frames.
211+
@param pts Presentation timestamps for each frame in vPacket using the FPS time base. e.g. fps = 25, pts = 3, presentation time = 3/25 seconds.
211212
*/
212-
virtual void onEncoded(const std::vector<std::vector<uint8_t>>& vPacket) = 0;
213+
virtual void onEncoded(const std::vector<std::vector<uint8_t>>& vPacket, const std::vector<uint64_t>& pts) = 0;
214+
215+
/** @brief Set the GOP pattern used by the encoder.
216+
217+
@param frameIntervalP Specify the GOP pattern as follows : \p frameIntervalP = 0: I, 1 : IPP, 2 : IBP, 3 : IBBP.
218+
*/
219+
virtual bool setFrameIntervalP(const int frameIntervalP) = 0;
213220

214221
/** @brief Callback function to that the encoding has finished.
215222
* */

‎modules/cudacodec/src/NvEncoder.cpp‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -305,9 +305,10 @@ void NvEncoder::MapResources(uint32_t bfrIdx)
305305
m_vMappedInputBuffers[bfrIdx] = mapInputResource.mappedResource;
306306
}
307307

308-
void NvEncoder::EncodeFrame(std::vector<std::vector<uint8_t>>& vPacket, NV_ENC_PIC_PARAMS* pPicParams)
308+
void NvEncoder::EncodeFrame(std::vector<std::vector<uint8_t>>& vPacket, std::vector<uint64_t>& outputTimeStamps, NV_ENC_PIC_PARAMS* pPicParams)
309309
{
310310
vPacket.clear();
311+
outputTimeStamps.clear();
311312
if (!IsHWEncoderInitialized())
312313
{
313314
NVENC_THROW_ERROR("Encoder device not found", NV_ENC_ERR_NO_ENCODE_DEVICE);
@@ -322,7 +323,7 @@ void NvEncoder::EncodeFrame(std::vector<std::vector<uint8_t>>& vPacket, NV_ENC_P
322323
if (nvStatus == NV_ENC_SUCCESS || nvStatus == NV_ENC_ERR_NEED_MORE_INPUT)
323324
{
324325
m_iToSend++;
325-
GetEncodedPacket(m_vBitstreamOutputBuffer, vPacket, true);
326+
GetEncodedPacket(m_vBitstreamOutputBuffer, vPacket, outputTimeStamps, true);
326327
}
327328
else
328329
{
@@ -353,6 +354,7 @@ NVENCSTATUS NvEncoder::DoEncode(NV_ENC_INPUT_PTR inputBuffer, NV_ENC_OUTPUT_PTR
353354
{
354355
picParams = *pPicParams;
355356
}
357+
picParams.inputTimeStamp = m_iInputFrame++;
356358
picParams.version = NV_ENC_PIC_PARAMS_VER;
357359
picParams.pictureStruct = NV_ENC_PIC_STRUCT_FRAME;
358360
picParams.inputBuffer = inputBuffer;
@@ -376,7 +378,7 @@ void NvEncoder::SendEOS()
376378
NVENC_API_CALL(m_nvenc.nvEncEncodePicture(m_hEncoder, &picParams));
377379
}
378380

379-
void NvEncoder::EndEncode(std::vector<std::vector<uint8_t>>& vPacket)
381+
void NvEncoder::EndEncode(std::vector<std::vector<uint8_t>>& vPacket, std::vector<uint64_t>& outputTimeStamps)
380382
{
381383
vPacket.clear();
382384
if (!IsHWEncoderInitialized())
@@ -386,10 +388,10 @@ void NvEncoder::EndEncode(std::vector<std::vector<uint8_t>>& vPacket)
386388

387389
SendEOS();
388390

389-
GetEncodedPacket(m_vBitstreamOutputBuffer, vPacket, false);
391+
GetEncodedPacket(m_vBitstreamOutputBuffer, vPacket, outputTimeStamps, false);
390392
}
391393

392-
void NvEncoder::GetEncodedPacket(std::vector<NV_ENC_OUTPUT_PTR>& vOutputBuffer, std::vector<std::vector<uint8_t>>& vPacket, bool bOutputDelay)
394+
void NvEncoder::GetEncodedPacket(std::vector<NV_ENC_OUTPUT_PTR>& vOutputBuffer, std::vector<std::vector<uint8_t>>& vPacket, std::vector<uint64_t>& outputTimeStamps, bool bOutputDelay)
393395
{
394396
unsigned i = 0;
395397
int iEnd = bOutputDelay ? m_iToSend - m_nOutputDelay : m_iToSend;
@@ -402,6 +404,7 @@ void NvEncoder::GetEncodedPacket(std::vector<NV_ENC_OUTPUT_PTR>& vOutputBuffer,
402404
lockBitstreamData.doNotWait = false;
403405
NVENC_API_CALL(m_nvenc.nvEncLockBitstream(m_hEncoder, &lockBitstreamData));
404406

407+
outputTimeStamps.push_back(lockBitstreamData.outputTimeStamp);
405408
uint8_t* pData = (uint8_t*)lockBitstreamData.bitstreamBufferPtr;
406409
if (vPacket.size() < i + 1)
407410
{
@@ -499,7 +502,8 @@ void NvEncoder::FlushEncoder()
499502
try
500503
{
501504
std::vector<std::vector<uint8_t>> vPacket;
502-
EndEncode(vPacket);
505+
std::vector<uint64_t> outputTimeStamps;
506+
EndEncode(vPacket, outputTimeStamps);
503507
}
504508
catch (...)
505509
{

‎modules/cudacodec/src/NvEncoder.h‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ class NvEncoder
100100
* data, which has been copied to an input buffer obtained from the
101101
* GetNextInputFrame() function.
102102
*/
103-
virtual void EncodeFrame(std::vector<std::vector<uint8_t>>& vPacket, NV_ENC_PIC_PARAMS* pPicParams = nullptr);
103+
virtual void EncodeFrame(std::vector<std::vector<uint8_t>>& vPacket, std::vector<uint64_t>& outputTimeStamps, NV_ENC_PIC_PARAMS* pPicParams = nullptr);
104104

105105
/**
106106
* @brief This function to flush the encoder queue.
@@ -109,7 +109,7 @@ class NvEncoder
109109
* from the encoder. The application must call this function before destroying
110110
* an encoder session.
111111
*/
112-
virtual void EndEncode(std::vector<std::vector<uint8_t>>& vPacket);
112+
virtual void EndEncode(std::vector<std::vector<uint8_t>>& vPacket, std::vector<uint64_t>& outputTimeStamps);
113113

114114
/**
115115
* @brief This function is used to query hardware encoder capabilities.
@@ -317,7 +317,7 @@ class NvEncoder
317317
* This is called by DoEncode() function. If there is buffering enabled,
318318
* this may return without any output data.
319319
*/
320-
void GetEncodedPacket(std::vector<NV_ENC_OUTPUT_PTR>& vOutputBuffer, std::vector<std::vector<uint8_t>>& vPacket, bool bOutputDelay);
320+
void GetEncodedPacket(std::vector<NV_ENC_OUTPUT_PTR>& vOutputBuffer, std::vector<std::vector<uint8_t>>& vPacket, std::vector<uint64_t>& outputTimeStamps, bool bOutputDelay);
321321

322322
/**
323323
* @brief This is a private function which is used to initialize the bitstream buffers.
@@ -369,6 +369,7 @@ class NvEncoder
369369
int32_t m_iGot = 0;
370370
int32_t m_nEncoderBuffer = 0;
371371
int32_t m_nOutputDelay = 0;
372+
int32_t m_iInputFrame = 0;
372373

373374
private:
374375
void* m_pDevice;

‎modules/cudacodec/src/video_writer.cpp‎

Lines changed: 32 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,10 @@ Ptr<cudacodec::VideoWriter> createVideoWriter(const String&, const Size, const C
5353

5454
#else // !defined HAVE_NVCUVENC
5555

56+
#if defined(WIN32) // remove when FFmpeg wrapper includes PR25874
57+
#define WIN32_WAIT_FOR_FFMPEG_WRAPPER_UPDATE
58+
#endif
59+
5660
NV_ENC_BUFFER_FORMAT EncBufferFormat(const ColorFormat colorFormat);
5761
int NChannels(const ColorFormat colorFormat);
5862
GUID CodecGuid(const Codec codec);
@@ -72,8 +76,9 @@ class FFmpegVideoWriter : public EncoderCallback
7276
public:
7377
FFmpegVideoWriter(const String& fileName, const Codec codec, const int fps, const Size sz, const int idrPeriod);
7478
~FFmpegVideoWriter();
75-
void onEncoded(const std::vector<std::vector<uint8_t>>& vPacket);
79+
void onEncoded(const std::vector<std::vector<uint8_t>>& vPacket, const std::vector<uint64_t>& pts);
7680
void onEncodingFinished();
81+
bool setFrameIntervalP(const int frameIntervalP);
7782
private:
7883
cv::VideoWriter writer;
7984
};
@@ -95,21 +100,32 @@ FFmpegVideoWriter::~FFmpegVideoWriter() {
95100
onEncodingFinished();
96101
}
97102

98-
void FFmpegVideoWriter::onEncoded(const std::vector<std::vector<uint8_t>>& vPacket) {
99-
for (auto& packet : vPacket) {
103+
void FFmpegVideoWriter::onEncoded(const std::vector<std::vector<uint8_t>>& vPacket, const std::vector<uint64_t>& pts) {
104+
CV_Assert(vPacket.size() == pts.size());
105+
for (int i = 0; i < vPacket.size(); i++){
106+
std::vector<uint8_t> packet = vPacket.at(i);
100107
Mat wrappedPacket(1, packet.size(), CV_8UC1, (void*)packet.data());
108+
const double ptsDouble = static_cast<double>(pts.at(i));
109+
CV_Assert(static_cast<uint64_t>(ptsDouble) == pts.at(i));
110+
#if !defined(WIN32_WAIT_FOR_FFMPEG_WRAPPER_UPDATE)
111+
CV_Assert(writer.set(VIDEOWRITER_PROP_PTS, ptsDouble));
112+
#endif
101113
writer.write(wrappedPacket);
102114
}
103115
}
104116

117+
bool FFmpegVideoWriter::setFrameIntervalP(const int frameIntervalP) {
118+
return writer.set(VIDEOWRITER_PROP_DTS_DELAY, static_cast<double>(frameIntervalP - 1));
119+
}
105120

106121
class RawVideoWriter : public EncoderCallback
107122
{
108123
public:
109124
RawVideoWriter(const String fileName);
110125
~RawVideoWriter();
111-
void onEncoded(const std::vector<std::vector<uint8_t>>& vPacket);
126+
void onEncoded(const std::vector<std::vector<uint8_t>>& vPacket, const std::vector<uint64_t>& pts);
112127
void onEncodingFinished();
128+
bool setFrameIntervalP(const int) { return false;}
113129
private:
114130
std::ofstream fpOut;
115131
};
@@ -128,7 +144,7 @@ RawVideoWriter::~RawVideoWriter() {
128144
onEncodingFinished();
129145
}
130146

131-
void RawVideoWriter::onEncoded(const std::vector<std::vector<uint8_t>>& vPacket) {
147+
void RawVideoWriter::onEncoded(const std::vector<std::vector<uint8_t>>& vPacket, const std::vector<uint64_t>&) {
132148
for (auto& packet : vPacket)
133149
fpOut.write(reinterpret_cast<const char*>(packet.data()), packet.size());
134150
}
@@ -208,8 +224,9 @@ VideoWriterImpl::VideoWriterImpl(const Ptr<EncoderCallback>& encoderCallBack_, c
208224
}
209225

210226
void VideoWriterImpl::release() {
211-
pEnc->EndEncode(vPacket);
212-
encoderCallback->onEncoded(vPacket);
227+
std::vector<uint64_t> pts;
228+
pEnc->EndEncode(vPacket, pts);
229+
encoderCallback->onEncoded(vPacket, pts);
213230
encoderCallback->onEncodingFinished();
214231
}
215232

@@ -316,6 +333,11 @@ void VideoWriterImpl::InitializeEncoder(const GUID codec, const double fps)
316333
initializeParams.encodeConfig->rcParams.maxBitRate = encoderParams.maxBitRate;
317334
initializeParams.encodeConfig->rcParams.targetQuality = encoderParams.targetQuality;
318335
initializeParams.encodeConfig->gopLength = encoderParams.gopLength;
336+
#if !defined(WIN32_WAIT_FOR_FFMPEG_WRAPPER_UPDATE)
337+
if (initializeParams.encodeConfig->frameIntervalP > 1) {
338+
CV_Assert(encoderCallback->setFrameIntervalP(initializeParams.encodeConfig->frameIntervalP));
339+
}
340+
#endif
319341
if (codec == NV_ENC_CODEC_H264_GUID)
320342
initializeParams.encodeConfig->encodeCodecConfig.h264Config.idrPeriod = encoderParams.idrPeriod;
321343
else if (codec == NV_ENC_CODEC_HEVC_GUID)
@@ -383,8 +405,9 @@ void VideoWriterImpl::CopyToNvSurface(const InputArray src)
383405
void VideoWriterImpl::write(const InputArray frame) {
384406
CV_Assert(frame.channels() == nSrcChannels);
385407
CopyToNvSurface(frame);
386-
pEnc->EncodeFrame(vPacket);
387-
encoderCallback->onEncoded(vPacket);
408+
std::vector<uint64_t> pts;
409+
pEnc->EncodeFrame(vPacket, pts);
410+
encoderCallback->onEncoded(vPacket, pts);
388411
};
389412

390413
EncoderParams VideoWriterImpl::getEncoderParams() const {

‎modules/cudacodec/test/test_video.cpp‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -650,6 +650,9 @@ struct TransCode : testing::TestWithParam<cv::cuda::DeviceInfo>
650650
}
651651
};
652652

653+
#if defined(WIN32) // remove when FFmpeg wrapper includes PR25874
654+
#define WIN32_WAIT_FOR_FFMPEG_WRAPPER_UPDATE
655+
#endif
653656

654657
CUDA_TEST_P(TransCode, H264ToH265)
655658
{
@@ -691,6 +694,10 @@ CUDA_TEST_P(TransCode, H264ToH265)
691694
for (int i = 0; i < nFrames; ++i) {
692695
cap >> frame;
693696
ASSERT_FALSE(frame.empty());
697+
#if !defined(WIN32_WAIT_FOR_FFMPEG_WRAPPER_UPDATE)
698+
const int pts = static_cast<int>(cap.get(CAP_PROP_PTS));
699+
ASSERT_EQ(i, pts > 0 ? pts : 0); // FFmpeg back end returns dts if pts is zero.
700+
#endif
694701
}
695702
}
696703
ASSERT_EQ(0, remove(outputFile.c_str()));
@@ -773,6 +780,10 @@ CUDA_TEST_P(Write, Writer)
773780
for (int i = 0; i < nFrames; ++i) {
774781
cap >> frame;
775782
ASSERT_FALSE(frame.empty());
783+
#if !defined(WIN32_WAIT_FOR_FFMPEG_WRAPPER_UPDATE)
784+
const int pts = static_cast<int>(cap.get(CAP_PROP_PTS));
785+
ASSERT_EQ(i, pts > 0 ? pts : 0); // FFmpeg back end returns dts if pts is zero.
786+
#endif
776787
}
777788
}
778789
ASSERT_EQ(0, remove(outputFile.c_str()));
@@ -867,6 +878,10 @@ CUDA_TEST_P(EncoderParams, Writer)
867878
const bool keyFrameActual = capRaw.get(CAP_PROP_LRF_HAS_KEY_FRAME) == 1.0;
868879
const bool keyFrameReference = i % idrPeriod == 0;
869880
ASSERT_EQ(keyFrameActual, keyFrameReference);
881+
#if !defined(WIN32_WAIT_FOR_FFMPEG_WRAPPER_UPDATE)
882+
const int pts = static_cast<int>(cap.get(CAP_PROP_PTS));
883+
ASSERT_EQ(i, pts > 0 ? pts : 0); // FFmpeg back end returns dts if pts is zero.
884+
#endif
870885
}
871886
}
872887
}

0 commit comments

Comments
 (0)