Force VideoReader to decode all YUV video formats to NV12

modules/cudacodec/include/opencv2/cudacodec.hpp

@@ -326,12 +326,13 @@ enum class VideoReaderProps {
 #endif
 };
 
-/** @brief ColorFormat for the frame returned by the decoder.
+/** @brief ColorFormat for the frame returned by nextFrame()/retrieve().
 */
 enum class ColorFormat {
     BGRA = 1,
     BGR = 2,
     GRAY = 3,
+    YUV = 4,
 #ifndef CV_DOXYGEN
     PROP_NOT_SUPPORTED
 #endif
@@ -394,7 +395,11 @@ 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 Set the desired ColorFormat for the frame returned by nextFrame()/retrieve().
+
+    @param colorFormat Value of the ColorFormat.
+     */
+    CV_WRAP virtual void set(const ColorFormat colorFormat) = 0;
 
     /** @brief Returns the specified VideoReader property
 

modules/cudacodec/src/video_decoder.cpp

@@ -45,14 +45,36 @@
 
 #ifdef HAVE_NVCUVID
 
+static const char* GetVideoChromaFormatString(cudaVideoChromaFormat eChromaFormat) {
+    static struct {
+        cudaVideoChromaFormat eChromaFormat;
+        const char* name;
+    } aChromaFormatName[] = {
+        { cudaVideoChromaFormat_Monochrome, "YUV 400 (Monochrome)" },
+        { cudaVideoChromaFormat_420,        "YUV 420"              },
+        { cudaVideoChromaFormat_422,        "YUV 422"              },
+        { cudaVideoChromaFormat_444,        "YUV 444"              },
+    };
+
+    if (eChromaFormat >= 0 && eChromaFormat < sizeof(aChromaFormatName) / sizeof(aChromaFormatName[0])) {
+        return aChromaFormatName[eChromaFormat].name;
+    }
+    return "Unknown";
+}
+
 void cv::cudacodec::detail::VideoDecoder::create(const FormatInfo& videoFormat)
 {
-    if (videoFormat.nBitDepthMinus8 > 0 || videoFormat.chromaFormat == YUV444)
-        CV_Error(Error::StsUnsupportedFormat, "NV12 output currently supported for 8 bit YUV420, YUV422 and Monochrome inputs.");
-
     videoFormat_ = videoFormat;
     const cudaVideoCodec _codec = static_cast<cudaVideoCodec>(videoFormat.codec);
     const cudaVideoChromaFormat _chromaFormat = static_cast<cudaVideoChromaFormat>(videoFormat.chromaFormat);
+    if (videoFormat.nBitDepthMinus8 > 0) {
+        std::ostringstream warning;
+        warning << "NV12 (8 bit luma, 4 bit chroma) is currently the only supported decoder output format. Video input is " << videoFormat.nBitDepthMinus8 + 8 << " bit " \
+            << std::string(GetVideoChromaFormatString(_chromaFormat)) << ".  Truncating luma to 8 bits";
+        if (videoFormat.chromaFormat != YUV420)
+            warning << " and chroma to 4 bits";
+        CV_LOG_WARNING(NULL, warning.str());
+    }
     const cudaVideoCreateFlags videoCreateFlags = (_codec == cudaVideoCodec_JPEG || _codec == cudaVideoCodec_MPEG2) ?
                                             cudaVideoCreate_PreferCUDA :
                                             cudaVideoCreate_PreferCUVID;
@@ -98,7 +120,7 @@ void cv::cudacodec::detail::VideoDecoder::create(const FormatInfo& videoFormat)
     cuSafeCall(cuCtxPushCurrent(ctx_));
     cuSafeCall(cuvidGetDecoderCaps(&decodeCaps));
     cuSafeCall(cuCtxPopCurrent(NULL));
-    if (!decodeCaps.bIsSupported)
+    if (!(decodeCaps.bIsSupported && (decodeCaps.nOutputFormatMask & (1 << cudaVideoSurfaceFormat_NV12))))
         CV_Error(Error::StsUnsupportedFormat, "Video source is not supported by hardware video decoder");
 
     CV_Assert(videoFormat.ulWidth >= decodeCaps.nMinWidth &&
@@ -115,6 +137,7 @@ void cv::cudacodec::detail::VideoDecoder::create(const FormatInfo& videoFormat)
     createInfo_.ulHeight            = videoFormat.ulHeight;
     createInfo_.ulNumDecodeSurfaces = videoFormat.ulNumDecodeSurfaces;
     createInfo_.ChromaFormat    = _chromaFormat;
+    createInfo_.bitDepthMinus8 = videoFormat.nBitDepthMinus8;
     createInfo_.OutputFormat    = cudaVideoSurfaceFormat_NV12;
     createInfo_.DeinterlaceMode = static_cast<cudaVideoDeinterlaceMode>(videoFormat.deinterlaceMode);
     createInfo_.ulTargetWidth       = videoFormat.width;

modules/cudacodec/src/video_reader.cpp

@@ -56,23 +56,26 @@ Ptr<VideoReader> cv::cudacodec::createVideoReader(const Ptr<RawVideoSource>&, co
 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)
+    Stream stream)
 {
     if (colorFormat == ColorFormat::BGRA) {
-        nv12ToBgra(decodedFrame, outFrame, width, height, stream);
+        nv12ToBgra(decodedFrame, outFrame, width, height, StreamAccessor::getStream(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;
+        nppStreamCtx.hStream = StreamAccessor::getStream(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);
+        cudaMemcpy2DAsync(outFrame.ptr(), outFrame.step, decodedFrame.ptr(), decodedFrame.step, width, height, cudaMemcpyDeviceToDevice, StreamAccessor::getStream(stream));
+    }
+    else if (colorFormat == ColorFormat::YUV) {
+        decodedFrame.copyTo(outFrame, stream);
     }
 }
 
@@ -96,7 +99,7 @@ namespace
 
         bool set(const VideoReaderProps propertyId, const double propertyVal) CV_OVERRIDE;
 
-        void VideoReaderImpl::set(const ColorFormat _colorFormat) CV_OVERRIDE;
+        void set(const ColorFormat _colorFormat) CV_OVERRIDE;
 
         bool get(const VideoReaderProps propertyId, double& propertyVal) const CV_OVERRIDE;
 
@@ -217,7 +220,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(), colorFormat, StreamAccessor::getStream(stream));
+            videoDecPostProcessFrame(decodedFrame, frame, videoDecoder_->targetWidth(), videoDecoder_->targetHeight(), colorFormat, stream);
 
             // unmap video frame
             // unmapFrame() synchronizes with the VideoDecode API (ensures the frame has finished decoding)

modules/cudacodec/test/test_video.cpp

@@ -187,6 +187,7 @@ CUDA_TEST_P(Video, Reader)
         {cudacodec::ColorFormat::GRAY,1},
         {cudacodec::ColorFormat::BGR,3},
         {cudacodec::ColorFormat::BGRA,4},
+        {cudacodec::ColorFormat::YUV,1}
     };
 
     std::string inputFile = std::string(cvtest::TS::ptr()->get_data_path()) + "../" + GET_PARAM(1);
@@ -201,7 +202,8 @@ CUDA_TEST_P(Video, Reader)
         ASSERT_TRUE(reader->nextFrame(frame));
         if(!fmt.valid)
             fmt = reader->format();
-        ASSERT_TRUE(frame.cols == fmt.width && frame.rows == fmt.height);
+        const int height = formatToChannels.first == cudacodec::ColorFormat::YUV ? 1.5 * fmt.height : fmt.height;
+        ASSERT_TRUE(frame.cols == fmt.width && frame.rows == height);
         ASSERT_FALSE(frame.empty());
         ASSERT_TRUE(frame.channels() == formatToChannels.second);
     }