cudacodec::VideoReader add colour format selection functionality

modules/cudacodec/include/opencv2/cudacodec.hpp

@@ -317,6 +317,18 @@ enum class VideoReaderProps {
     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.
+    PROP_COLOR_FORMAT = 6, //!< Set the ColorFormat of the decoded frame.  This can be changed before every call to nextFrame() and retrieve().
+#ifndef CV_DOXYGEN
+    PROP_NOT_SUPPORTED
+#endif
+};
+
+/** @brief ColorFormat for the frame returned by the decoder.
+*/
+enum class ColorFormat {
+    BGRA = 1,
+    BGR = 2,
+    GRAY = 3,
 #ifndef CV_DOXYGEN
     PROP_NOT_SUPPORTED
 #endif
@@ -379,6 +391,8 @@ class CV_EXPORTS_W VideoReader
      */
     CV_WRAP virtual bool set(const VideoReaderProps propertyId, const double propertyVal) = 0;
 
+    CV_WRAP virtual void set(const ColorFormat _colorFormat) = 0;
+
     /** @brief Returns the specified VideoReader property
 
     @param propertyId Property identifier from cv::cudacodec::VideoReaderProps (eg. cv::cudacodec::VideoReaderProps::PROP_DECODED_FRAME_IDX,

modules/cudacodec/src/cuda/nv12_to_rgb.cu

@@ -60,7 +60,7 @@
 using namespace cv;
 using namespace cv::cudev;
 
-void videoDecPostProcessFrame(const GpuMat& decodedFrame, GpuMat& _outFrame, int width, int height, cudaStream_t stream);
+void nv12ToBgra(const GpuMat& decodedFrame, GpuMat& outFrame, int width, int height, cudaStream_t stream);
 
 namespace
 {
@@ -112,7 +112,7 @@ namespace
     #define COLOR_COMPONENT_BIT_SIZE 10
     #define COLOR_COMPONENT_MASK     0x3FF
 
-    __global__ void NV12_to_RGB(const uchar* srcImage, size_t nSourcePitch,
+    __global__ void NV12_to_BGRA(const uchar* srcImage, size_t nSourcePitch,
                                   uint* dstImage, size_t nDestPitch,
                                   uint width, uint height)
     {
@@ -186,22 +186,16 @@ namespace
     }
 }
 
-void videoDecPostProcessFrame(const GpuMat& decodedFrame, GpuMat& outFrame, int width, int height, cudaStream_t stream)
+void nv12ToBgra(const GpuMat& decodedFrame, GpuMat& outFrame, int width, int height, cudaStream_t stream)
 {
-    // Final Stage: NV12toARGB color space conversion
-
     outFrame.create(height, width, CV_8UC4);
-
     dim3 block(32, 8);
     dim3 grid(divUp(width, 2 * block.x), divUp(height, block.y));
-
-    NV12_to_RGB<<<grid, block, 0, stream>>>(decodedFrame.ptr<uchar>(), decodedFrame.step,
-                                 outFrame.ptr<uint>(), outFrame.step,
-                                 width, height);
-
-    CV_CUDEV_SAFE_CALL( cudaGetLastError() );
+    NV12_to_BGRA<< <grid, block, 0, stream >> > (decodedFrame.ptr<uchar>(), decodedFrame.step,
+        outFrame.ptr<uint>(), outFrame.step, width, height);
+    CV_CUDEV_SAFE_CALL(cudaGetLastError());
     if (stream == 0)
-      CV_CUDEV_SAFE_CALL( cudaDeviceSynchronize() );
+        CV_CUDEV_SAFE_CALL(cudaDeviceSynchronize());
 }
 
 #endif

modules/cudacodec/src/video_reader.cpp

@@ -53,7 +53,28 @@ Ptr<VideoReader> cv::cudacodec::createVideoReader(const Ptr<RawVideoSource>&, co
 
 #else // HAVE_NVCUVID
 
-void videoDecPostProcessFrame(const GpuMat& decodedFrame, GpuMat& _outFrame, int width, int height, cudaStream_t stream);
+void nv12ToBgra(const GpuMat& decodedFrame, GpuMat& outFrame, int width, int height, cudaStream_t stream);
+
+void videoDecPostProcessFrame(const GpuMat& decodedFrame, GpuMat& outFrame, int width, int height, const ColorFormat colorFormat,
+    cudaStream_t stream)
+{
+    if (colorFormat == ColorFormat::BGRA) {
+        nv12ToBgra(decodedFrame, outFrame, width, height, stream);
+    }
+    else if (colorFormat == ColorFormat::BGR) {
+        outFrame.create(height, width, CV_8UC3);
+        Npp8u* pSrc[2] = { decodedFrame.data, &decodedFrame.data[decodedFrame.step * height] };
+        NppiSize oSizeROI = { width,height };
+        NppStreamContext nppStreamCtx;
+        nppSafeCall(nppGetStreamContext(&nppStreamCtx));
+        nppStreamCtx.hStream = stream;
+        nppSafeCall(nppiNV12ToBGR_8u_P2C3R_Ctx(pSrc, decodedFrame.step, outFrame.data, outFrame.step, oSizeROI, nppStreamCtx));
+    }
+    else if (colorFormat == ColorFormat::GRAY) {
+        outFrame.create(height, width, CV_8UC1);
+        cudaMemcpy2DAsync(outFrame.ptr(), outFrame.step, decodedFrame.ptr(), decodedFrame.step, width, height, cudaMemcpyDeviceToDevice, stream);
+    }
+}
 
 using namespace cv::cudacodec::detail;
 
@@ -75,6 +96,8 @@ namespace
 
         bool set(const VideoReaderProps propertyId, const double propertyVal) CV_OVERRIDE;
 
+        void VideoReaderImpl::set(const ColorFormat _colorFormat) CV_OVERRIDE;
+
         bool get(const VideoReaderProps propertyId, double& propertyVal) const CV_OVERRIDE;
 
         bool get(const int propertyId, double& propertyVal) const CV_OVERRIDE;
@@ -96,6 +119,7 @@ namespace
         static const int decodedFrameIdx = 0;
         static const int extraDataIdx = 1;
         static const int rawPacketsBaseIdx = 2;
+        ColorFormat colorFormat = ColorFormat::BGRA;
     };
 
     FormatInfo VideoReaderImpl::format() const
@@ -193,7 +217,7 @@ namespace
 
             // perform post processing on the CUDA surface (performs colors space conversion and post processing)
             // comment this out if we include the line of code seen above
-            videoDecPostProcessFrame(decodedFrame, frame, videoDecoder_->targetWidth(), videoDecoder_->targetHeight(), StreamAccessor::getStream(stream));
+            videoDecPostProcessFrame(decodedFrame, frame, videoDecoder_->targetWidth(), videoDecoder_->targetHeight(), colorFormat, StreamAccessor::getStream(stream));
 
             // unmap video frame
             // unmapFrame() synchronizes with the VideoDecode API (ensures the frame has finished decoding)
@@ -237,11 +261,14 @@ namespace
         switch (propertyId) {
         case VideoReaderProps::PROP_RAW_MODE :
             videoSource_->SetRawMode(static_cast<bool>(propertyVal));
-            break;
         }
         return true;
     }
 
+    void VideoReaderImpl::set(const ColorFormat _colorFormat) {
+        colorFormat = _colorFormat;
+    }
+
     bool VideoReaderImpl::get(const VideoReaderProps propertyId, double& propertyVal) const {
         switch (propertyId)
         {

modules/cudacodec/test/test_video.cpp

@@ -179,17 +179,27 @@ CUDA_TEST_P(Video, Reader)
     if (GET_PARAM(1) == "cv/video/768x576.avi" && !videoio_registry::hasBackend(CAP_FFMPEG))
         throw SkipTestException("FFmpeg backend not found");
 
+    const std::vector<std::pair< cudacodec::ColorFormat, int>> formatsToChannels = {
+        {cudacodec::ColorFormat::GRAY,1},
+        {cudacodec::ColorFormat::BGR,3},
+        {cudacodec::ColorFormat::BGRA,4},
+    };
+
     std::string inputFile = std::string(cvtest::TS::ptr()->get_data_path()) + "../" + GET_PARAM(1);
     cv::Ptr<cv::cudacodec::VideoReader> reader = cv::cudacodec::createVideoReader(inputFile);
     cv::cudacodec::FormatInfo fmt = reader->format();
     cv::cuda::GpuMat frame;
     for (int i = 0; i < 100; i++)
     {
+        // request a different colour format for each frame
+        const std::pair< cudacodec::ColorFormat, int>& formatToChannels = formatsToChannels[i % formatsToChannels.size()];
+        reader->set(formatToChannels.first);
         ASSERT_TRUE(reader->nextFrame(frame));
         if(!fmt.valid)
             fmt = reader->format();
         ASSERT_TRUE(frame.cols == fmt.width && frame.rows == fmt.height);
         ASSERT_FALSE(frame.empty());
+        ASSERT_TRUE(frame.channels() == formatToChannels.second);
     }
 }