Merge pull request #3098 from cudawarped:rtsp_video_write

Add writing encoded video streams to disk when streaming from RTSP sources using cudacodec

* Add missing codecs to cudacodec which uses Nvidia Video Codec SDK including checks to ensure codec used in input video file is supported on the current device.

* Update cudacoded to
1) automatically write the raw encoded video stream to a video file and,
2) use less GPU memory by more closely mirroring the Nvidia samples.  Specifically querying the decoder for the number of decode surfaces (h265 commonly uses 4) instead of always using 20 and not using adaptive deinterlacing when the video sequence is progressive.  Additional updates to mirror the Nvidia sample include initializing the decoder so that HandleVideoSequence() gets called every time before the decoder is initialized, ensuring all the parameters for the decoder are provided by nvcudec.
Added facility to decode AV1, not tested as VideoCapture doesn't return a valid fourcc for this.
Add facility to decode MPEG4 video - requires modification to VideoCapture see pull request.

* Prevent adding parameter sets twice and add zero padding to output files to that they play in vlc.
Notes: VideoCapture - returns mpeg as the codec for mpeg4 files, so files written as .m4v from mpeg4 sources cannot currently be decoded.  This is also true for AV1 sources where cap.get(CAP_PROP_FOURCC) returns 0.
Added mpeg4 test file which can be decoded when VideoCapture adds the extra_data.

* Update to account for the extraData being passed from cap.retrieve instead of appended to the first packet.

* Update to be compatible with changes to VideoCapture

* Remove redundant test.

* Add check to ensure retrieve is successful.

* Remove writeToFile and allow VideoReader to return raw encoded data at the same time as decoded frames.

* Fix missing documentation.

Author: cudawarped

modules/cudacodec/include/opencv2/cudacodec.hpp

@@ -231,7 +231,7 @@ CV_EXPORTS_W Ptr<cudacodec::VideoWriter> createVideoWriter(const Ptr<EncoderCall
 want to work with raw video stream.
 @param frameSize Size of the input video frames.
 @param fps Framerate of the created video stream.
-@param params Encoder parameters. See cudacodec::EncoderParams .
+@param params Encoder parameters. See cudacodec::EncoderParams.
 @param format Surface format of input frames ( SF_UYVY , SF_YUY2 , SF_YV12 , SF_NV12 ,
 SF_IYUV , SF_BGR or SF_GRAY). BGR or gray frames will be converted to YV12 format before
 encoding, frames with other formats will be used as is.
@@ -265,7 +265,7 @@ enum Codec
     Uncompressed_UYVY   = (('U'<<24)|('Y'<<16)|('V'<<8)|('Y'))    //!< UYVY (4:2:2)
 };
 
-/** @brief Chroma formats supported by cudacodec::VideoReader .
+/** @brief Chroma formats supported by cudacodec::VideoReader.
  */
 enum ChromaFormat
 {
@@ -276,17 +276,47 @@ enum ChromaFormat
     NumFormats
 };
 
+/** @brief Deinterlacing mode used by decoder.
+* @param Weave Weave both fields (no deinterlacing). For progressive content and for content that doesn't need deinterlacing.
+* Bob Drop one field.
+* @param Adaptive Adaptive deinterlacing needs more video memory than other deinterlacing modes.
+* */
+enum DeinterlaceMode
+{
+    Weave = 0,
+    Bob = 1,
+    Adaptive = 2
+};
+
 /** @brief Struct providing information about video file format. :
  */
 struct FormatInfo
 {
     Codec codec;
     ChromaFormat chromaFormat;
     int nBitDepthMinus8 = -1;
-    int width = 0;//!< Width of the decoded frame returned by nextFrame(frame)
-    int height = 0;//!< Height of the decoded frame returned by nextFrame(frame)
+    int ulWidth = 0;//!< Coded sequence width in pixels.
+    int ulHeight = 0;//!< Coded sequence height in pixels.
+    int width = 0;//!< Width of the decoded frame returned by nextFrame(frame).
+    int height = 0;//!< Height of the decoded frame returned by nextFrame(frame).
+    int ulMaxWidth = 0;
+    int ulMaxHeight = 0;
     Rect displayArea;//!< ROI inside the decoded frame returned by nextFrame(frame), containing the useable video frame.
     bool valid = false;
+    double fps = 0;
+    int ulNumDecodeSurfaces = 0;//!< Maximum number of internal decode surfaces.
+    DeinterlaceMode deinterlaceMode;
+};
+
+/** @brief cv::cudacodec::VideoReader generic properties identifier.
+*/
+enum class VideoReaderProps {
+    PROP_DECODED_FRAME_IDX = 0, //!< Index for retrieving the decoded frame using retrieve().
+    PROP_EXTRA_DATA_INDEX = 1, //!< Index for retrieving the extra data associated with a video source using retrieve().
+    PROP_RAW_PACKAGES_BASE_INDEX = 2, //!< Base index for retrieving raw encoded data using retrieve().
+    PROP_NUMBER_OF_RAW_PACKAGES_SINCE_LAST_GRAB = 3, //!< Number of raw packages recieved since the last call to grab().
+    PROP_RAW_MODE = 4, //!< Status of raw mode.
+    PROP_LRF_HAS_KEY_FRAME = 5 //!< FFmpeg source only - Indicates whether the Last Raw Frame (LRF), output from VideoReader::retrieve() when VideoReader is initialized in raw mode, contains encoded data for a key frame.
 };
 
 /** @brief Video reader interface.
@@ -310,6 +340,48 @@ class CV_EXPORTS_W VideoReader
     /** @brief Returns information about video file format.
     */
     virtual FormatInfo format() const = 0;
+
+    /** @brief Grabs the next frame from the video source.
+
+    @return `true` (non-zero) in the case of success.
+
+    The method/function grabs the next frame from video file or camera and returns true (non-zero) in
+    the case of success.
+
+    The primary use of the function is for reading both the encoded and decoded video data when rawMode is enabled.  With rawMode enabled
+    retrieve() can be called following grab() to retrieve all the data associated with the current video source since the last call to grab() or the creation of the VideoReader.
+     */
+    CV_WRAP virtual bool grab(Stream& stream = Stream::Null()) = 0;
+
+    /** @brief Returns previously grabbed video data.
+
+    @param [out] frame The returned data which depends on the provided idx.  If there is no new data since the last call to grab() the image will be empty.
+    @param idx Determins the returned data inside image. The returned data can be the:
+    Decoded frame, idx = get(PROP_DECODED_FRAME_IDX).
+    Extra data if available, idx = get(PROP_EXTRA_DATA_INDEX).
+    Raw encoded data package.  To retrieve package i,  idx = get(PROP_RAW_PACKAGES_BASE_INDEX) + i with i < get(PROP_NUMBER_OF_RAW_PACKAGES_SINCE_LAST_GRAB)
+    @return `false` if no frames has been grabbed
+
+    The method returns data associated with the current video source since the last call to grab() or the creation of the VideoReader. If no data is present
+    the method returns false and the function returns an empty image.
+     */
+    CV_WRAP virtual bool retrieve(CV_OUT OutputArray frame, const size_t idx = static_cast<size_t>(VideoReaderProps::PROP_DECODED_FRAME_IDX)) const = 0;
+
+    /** @brief Sets a property in the VideoReader.
+
+    @param property Property identifier from cv::cudacodec::VideoReaderProps (eg. cv::cudacodec::PROP_DECODED_FRAME_IDX, cv::cudacodec::PROP_EXTRA_DATA_INDEX, ...)
+    @param propertyVal Value of the property.
+    @return `true` if the property has been set.
+     */
+    CV_WRAP virtual bool set(const VideoReaderProps property, const double propertyVal) = 0;
+
+    /** @brief Returns the specified VideoReader property
+
+    @param property Property identifier from cv::cudacodec::VideoReaderProps (eg. cv::cudacodec::PROP_DECODED_FRAME_IDX, cv::cudacodec::PROP_EXTRA_DATA_INDEX, ...)
+    @param propertyVal Optional value for the property.
+    @return Value for the specified property. Value -1 is returned when querying a property that is not supported.
+    */
+    CV_WRAP virtual int get(const VideoReaderProps property, const int propertyVal = -1) const = 0;
 };
 
 /** @brief Interface for video demultiplexing. :
@@ -328,26 +400,39 @@ class CV_EXPORTS_W RawVideoSource
      */
     virtual bool getNextPacket(unsigned char** data, size_t* size) = 0;
 
+    /** @brief Returns true if the last packet contained a key frame.
+     */
+    virtual bool lastPacketContainsKeyFrame() const { return false; }
+
     /** @brief Returns information about video file format.
     */
     virtual FormatInfo format() const = 0;
 
     /** @brief Updates the coded width and height inside format.
     */
-    virtual void updateFormat(const int codedWidth, const int codedHeight) = 0;
+    virtual void updateFormat(const FormatInfo& videoFormat) = 0;
+
+    /** @brief Returns any extra data associated with the video source.
+
+    @param extraData 1D cv::Mat containing the extra data if it exists.
+     */
+    virtual void getExtraData(cv::Mat& extraData) const = 0;
 };
 
 /** @brief Creates video reader.
 
 @param filename Name of the input video file.
+@param rawMode Allow the raw encoded data which has been read up until the last call to grab() to be retrieved by calling retrieve(rawData,RAW_DATA_IDX).
 
 FFMPEG is used to read videos. User can implement own demultiplexing with cudacodec::RawVideoSource
  */
-CV_EXPORTS_W Ptr<VideoReader> createVideoReader(const String& filename);
+CV_EXPORTS_W Ptr<VideoReader> createVideoReader(const String& filename, const bool rawMode = false);
+
 /** @overload
 @param source RAW video source implemented by user.
+@param rawMode Allow the raw encoded data which has been read up until the last call to grab() to be retrieved by calling retrieve(rawData,RAW_DATA_IDX).
 */
-CV_EXPORTS_W Ptr<VideoReader> createVideoReader(const Ptr<RawVideoSource>& source);
+CV_EXPORTS_W Ptr<VideoReader> createVideoReader(const Ptr<RawVideoSource>& source, const bool rawMode = false);
 
 //! @}
 

modules/cudacodec/src/cuvid_video_source.cpp

@@ -76,6 +76,8 @@ cv::cudacodec::detail::CuvidVideoSource::CuvidVideoSource(const String& fname)
     format_.height = vidfmt.coded_height;
     format_.displayArea = Rect(Point(vidfmt.display_area.left, vidfmt.display_area.top), Point(vidfmt.display_area.right, vidfmt.display_area.bottom));
     format_.valid = true;
+    if (vidfmt.frame_rate.numerator != 0 && vidfmt.frame_rate.denominator != 0)
+        format_.fps = vidfmt.frame_rate.numerator / (double)vidfmt.frame_rate.denominator;
 }
 
 cv::cudacodec::detail::CuvidVideoSource::~CuvidVideoSource()
@@ -88,10 +90,9 @@ FormatInfo cv::cudacodec::detail::CuvidVideoSource::format() const
     return format_;
 }
 
-void cv::cudacodec::detail::CuvidVideoSource::updateFormat(const int codedWidth, const int codedHeight)
+void cv::cudacodec::detail::CuvidVideoSource::updateFormat(const FormatInfo& videoFormat)
 {
-    format_.width = codedWidth;
-    format_.height = codedHeight;
+    format_ = videoFormat;
     format_.valid = true;
 }
 
@@ -119,7 +120,7 @@ int CUDAAPI cv::cudacodec::detail::CuvidVideoSource::HandleVideoData(void* userD
 {
     CuvidVideoSource* thiz = static_cast<CuvidVideoSource*>(userData);
 
-    return thiz->parseVideoData(packet->payload, packet->payload_size, (packet->flags & CUVID_PKT_ENDOFSTREAM) != 0);
+    return thiz->parseVideoData(packet->payload, packet->payload_size, thiz->RawModeEnabled(), false, (packet->flags & CUVID_PKT_ENDOFSTREAM) != 0);
 }
 
 #endif // HAVE_NVCUVID

modules/cudacodec/src/cuvid_video_source.hpp

@@ -55,7 +55,7 @@ class CuvidVideoSource : public VideoSource
     ~CuvidVideoSource();
 
     FormatInfo format() const CV_OVERRIDE;
-    void updateFormat(const int codedWidth, const int codedHeight);
+    void updateFormat(const FormatInfo& videoFormat) CV_OVERRIDE;
     void start() CV_OVERRIDE;
     void stop() CV_OVERRIDE;
     bool isStarted() const CV_OVERRIDE;

modules/cudacodec/src/ffmpeg_video_source.cpp

@@ -75,6 +75,7 @@ Codec FourccToCodec(int codec)
     case CV_FOURCC_MACRO('M', 'P', 'G', '1'): return MPEG1;
     case CV_FOURCC_MACRO('M', 'P', 'G', '2'): return MPEG2;
     case CV_FOURCC_MACRO('X', 'V', 'I', 'D'): // fallthru
+    case CV_FOURCC_MACRO('m', 'p', '4', 'v'): // fallthru
     case CV_FOURCC_MACRO('D', 'I', 'V', 'X'): return MPEG4;
     case CV_FOURCC_MACRO('W', 'V', 'C', '1'): return VC1;
     case CV_FOURCC_MACRO('H', '2', '6', '4'): // fallthru
@@ -112,6 +113,22 @@ void FourccToChromaFormat(const int pixelFormat, ChromaFormat &chromaFormat, int
     }
 }
 
+int StartCodeLen(unsigned char* data, const int sz) {
+    if (sz >= 3 && data[0] == 0 && data[1] == 0 && data[2] == 1)
+        return 3;
+    else if (sz >= 4 && data[0] == 0 && data[1] == 0 && data[2] == 0 && data[3] == 1)
+        return 4;
+    else
+        return 0;
+}
+
+bool ParamSetsExist(unsigned char* parameterSets, const int szParameterSets, unsigned char* data, const int szData) {
+    const int paramSetStartCodeLen = StartCodeLen(parameterSets, szParameterSets);
+    const int packetStartCodeLen = StartCodeLen(data, szData);
+    // weak test to see if the parameter set has already been included in the RTP stream
+    return paramSetStartCodeLen != 0 && packetStartCodeLen != 0 && parameterSets[paramSetStartCodeLen] == data[packetStartCodeLen];
+}
+
 cv::cudacodec::detail::FFmpegVideoSource::FFmpegVideoSource(const String& fname)
 {
     if (!videoio_registry::hasBackend(CAP_FFMPEG))
@@ -125,6 +142,11 @@ cv::cudacodec::detail::FFmpegVideoSource::FFmpegVideoSource(const String& fname)
         CV_Error(Error::StsUnsupportedFormat, "Fetching of RAW video streams is not supported");
     CV_Assert(cap.get(CAP_PROP_FORMAT) == -1);
 
+    const int codecExtradataIndex = static_cast<int>(cap.get(CAP_PROP_CODEC_EXTRADATA_INDEX));
+    Mat tmpExtraData;
+    if (cap.retrieve(tmpExtraData, codecExtradataIndex) && tmpExtraData.total())
+        extraData = tmpExtraData.clone();
+
     int codec = (int)cap.get(CAP_PROP_FOURCC);
     int pixelFormat = (int)cap.get(CAP_PROP_CODEC_PIXEL_FORMAT);
 
@@ -133,6 +155,7 @@ cv::cudacodec::detail::FFmpegVideoSource::FFmpegVideoSource(const String& fname)
     format_.width = cap.get(CAP_PROP_FRAME_WIDTH);
     format_.displayArea = Rect(0, 0, format_.width, format_.height);
     format_.valid = false;
+    format_.fps = cap.get(CAP_PROP_FPS);
     FourccToChromaFormat(pixelFormat, format_.chromaFormat, format_.nBitDepthMinus8);
 }
 
@@ -147,10 +170,9 @@ FormatInfo cv::cudacodec::detail::FFmpegVideoSource::format() const
     return format_;
 }
 
-void cv::cudacodec::detail::FFmpegVideoSource::updateFormat(const int codedWidth, const int codedHeight)
+void cv::cudacodec::detail::FFmpegVideoSource::updateFormat(const FormatInfo& videoFormat)
 {
-    format_.width = codedWidth;
-    format_.height = codedHeight;
+    format_ = videoFormat;
     format_.valid = true;
 }
 
@@ -159,7 +181,26 @@ bool cv::cudacodec::detail::FFmpegVideoSource::getNextPacket(unsigned char** dat
     cap >> rawFrame;
     *data = rawFrame.data;
     *size = rawFrame.total();
+    if (iFrame++ == 0 && extraData.total()) {
+        if (format_.codec == Codec::MPEG4 ||
+            ((format_.codec == Codec::H264 || format_.codec == Codec::HEVC) && !ParamSetsExist(extraData.data, extraData.total(), *data, *size)))
+        {
+            const size_t nBytesToTrimFromData = format_.codec == Codec::MPEG4 ? 3 : 0;
+            const size_t newSz = extraData.total() + *size - nBytesToTrimFromData;
+            dataWithHeader = Mat(1, newSz, CV_8UC1);
+            memcpy(dataWithHeader.data, extraData.data, extraData.total());
+            memcpy(dataWithHeader.data + extraData.total(), (*data) + nBytesToTrimFromData, *size - nBytesToTrimFromData);
+            *data = dataWithHeader.data;
+            *size = newSz;
+        }
+    }
+
     return *size != 0;
 }
 
+bool cv::cudacodec::detail::FFmpegVideoSource::lastPacketContainsKeyFrame() const
+{
+    return cap.get(CAP_PROP_LRF_HAS_KEY_FRAME);
+}
+
 #endif // HAVE_CUDA

modules/cudacodec/src/ffmpeg_video_source.hpp

@@ -56,15 +56,19 @@ class FFmpegVideoSource : public RawVideoSource
 
     bool getNextPacket(unsigned char** data, size_t* size) CV_OVERRIDE;
 
+    bool lastPacketContainsKeyFrame() const;
+
     FormatInfo format() const CV_OVERRIDE;
 
-    void updateFormat(const int codedWidth, const int codedHeight);
+    void updateFormat(const FormatInfo& videoFormat) CV_OVERRIDE;
 
+    void getExtraData(cv::Mat& _extraData) const CV_OVERRIDE { _extraData = extraData; }
 
 private:
     FormatInfo format_;
     VideoCapture cap;
-    Mat rawFrame;
+    Mat rawFrame, extraData, dataWithHeader;
+    int iFrame = 0;
 };
 
 }}}

modules/cudacodec/src/frame_queue.cpp

@@ -45,13 +45,22 @@
 
 #ifdef HAVE_NVCUVID
 
-cv::cudacodec::detail::FrameQueue::FrameQueue() :
-    endOfDecode_(0),
-    framesInQueue_(0),
-    readPosition_(0)
-{
-    std::memset(displayQueue_, 0, sizeof(displayQueue_));
-    std::memset((void*) isFrameInUse_, 0, sizeof(isFrameInUse_));
+RawPacket::RawPacket(const unsigned char* _data, const size_t _size, const bool _containsKeyFrame) : size(_size), containsKeyFrame(_containsKeyFrame) {
+    data = cv::makePtr<unsigned char*>(new unsigned char[size]);
+    memcpy(*data, _data, size);
+};
+
+cv::cudacodec::detail::FrameQueue::~FrameQueue() {
+    if (isFrameInUse_)
+        delete[] isFrameInUse_;
+}
+
+void cv::cudacodec::detail::FrameQueue::init(const int _maxSz) {
+    AutoLock autoLock(mtx_);
+    maxSz = _maxSz;
+    displayQueue_ = std::vector<CUVIDPARSERDISPINFO>(maxSz, CUVIDPARSERDISPINFO());
+    isFrameInUse_ = new volatile int[maxSz];
+    std::memset((void*)isFrameInUse_, 0, sizeof(*isFrameInUse_) * maxSz);
 }
 
 bool cv::cudacodec::detail::FrameQueue::waitUntilFrameAvailable(int pictureIndex)
@@ -68,7 +77,7 @@ bool cv::cudacodec::detail::FrameQueue::waitUntilFrameAvailable(int pictureIndex
     return true;
 }
 
-void cv::cudacodec::detail::FrameQueue::enqueue(const CUVIDPARSERDISPINFO* picParams)
+void cv::cudacodec::detail::FrameQueue::enqueue(const CUVIDPARSERDISPINFO* picParams, const std::vector<RawPacket> rawPackets)
 {
     // Mark the frame as 'in-use' so we don't re-use it for decoding until it is no longer needed
     // for display
@@ -82,10 +91,12 @@ void cv::cudacodec::detail::FrameQueue::enqueue(const CUVIDPARSERDISPINFO* picPa
         {
             AutoLock autoLock(mtx_);
 
-            if (framesInQueue_ < MaximumSize)
+            if (framesInQueue_ < maxSz)
             {
-                int writePosition = (readPosition_ + framesInQueue_) % MaximumSize;
-                displayQueue_[writePosition] = *picParams;
+                const int writePosition = (readPosition_ + framesInQueue_) % maxSz;
+                displayQueue_.at(writePosition) = *picParams;
+                for (const auto& rawPacket : rawPackets)
+                    rawPacketQueue.push(rawPacket);
                 framesInQueue_++;
                 isFramePlaced = true;
             }
@@ -99,15 +110,19 @@ void cv::cudacodec::detail::FrameQueue::enqueue(const CUVIDPARSERDISPINFO* picPa
     } while (!isEndOfDecode());
 }
 
-bool cv::cudacodec::detail::FrameQueue::dequeue(CUVIDPARSERDISPINFO& displayInfo)
+bool cv::cudacodec::detail::FrameQueue::dequeue(CUVIDPARSERDISPINFO& displayInfo, std::vector<RawPacket>& rawPackets)
 {
     AutoLock autoLock(mtx_);
 
     if (framesInQueue_ > 0)
     {
         int entry = readPosition_;
-        displayInfo = displayQueue_[entry];
-        readPosition_ = (entry + 1) % MaximumSize;
+        displayInfo = displayQueue_.at(entry);
+        while (!rawPacketQueue.empty()) {
+            rawPackets.push_back(rawPacketQueue.front());
+            rawPacketQueue.pop();
+        }
+        readPosition_ = (entry + 1) % maxSz;
         framesInQueue_--;
         return true;
     }