cudacodec - Add 10 bit YUV420 and YUV444 encoding functionality

Author: cudawarped

modules/cudacodec/include/opencv2/cudacodec.hpp

@@ -106,6 +106,8 @@ enum ColorFormat {
     NV_IYUV = 9, //!< Nvidia Buffer Format - Planar YUV [Y plane followed by U and V planes]. VideoWriter only.
     NV_YUV444 = 10, //!< Nvidia Buffer Format - Planar YUV [Y plane followed by U and V planes]. VideoWriter only.
     NV_AYUV = 11, //!< Nvidia Buffer Format - 8 bit Packed A8Y8U8V8. This is a word-ordered format where a pixel is represented by a 32-bit word with V in the lowest 8 bits, U in the next 8 bits, Y in the 8 bits after that and A in the highest 8 bits. VideoWriter only.
+    NV_YUV420_10BIT = 12, //!< Nvidia Buffer Format - 10 bit Semi-Planar YUV [Y plane followed by interleaved UV plane]. Each pixel of size 2 bytes. Most Significant 10 bits contain pixel data. VideoWriter only.
+    NV_YUV444_10BIT = 13, //!< Nvidia Buffer Format - 10 bit Planar YUV444 [Y plane followed by U and V planes]. Each pixel of size 2 bytes. Most Significant 10 bits contain pixel data. VideoWriter only.
 #ifndef CV_DOXYGEN
     PROP_NOT_SUPPORTED
 #endif

modules/cudacodec/src/ffmpeg_video_source.cpp

@@ -89,23 +89,6 @@ Codec FourccToCodec(int codec)
     CV_Error(Error::StsUnsupportedFormat, msg);
 }
 
-static
-void FourccToChromaFormat(const int pixelFormat, ChromaFormat &chromaFormat, int & nBitDepthMinus8)
-{
-    switch (pixelFormat)
-    {
-    case CV_FOURCC_MACRO('I', '4', '2', '0'):
-        chromaFormat = YUV420;
-        nBitDepthMinus8 = 0;
-        break;
-    default:
-        CV_LOG_WARNING(NULL, cv::format("ChromaFormat not recognized: 0x%08X (%s). Assuming I420", pixelFormat, fourccToString(pixelFormat).c_str()));
-        chromaFormat = YUV420;
-        nBitDepthMinus8 = 0;
-        break;
-    }
-}
-
 static
 int StartCodeLen(unsigned char* data, const int sz) {
     if (sz >= 3 && data[0] == 0 && data[1] == 0 && data[2] == 1)
@@ -145,14 +128,12 @@ cv::cudacodec::detail::FFmpegVideoSource::FFmpegVideoSource(const String& fname,
         extraData = tmpExtraData.clone();
 
     int codec = (int)cap.get(CAP_PROP_FOURCC);
-    int pixelFormat = (int)cap.get(CAP_PROP_CODEC_PIXEL_FORMAT);
     format_.codec = FourccToCodec(codec);
     format_.height = cap.get(CAP_PROP_FRAME_HEIGHT);
     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);
 }
 
 cv::cudacodec::detail::FFmpegVideoSource::~FFmpegVideoSource()

modules/cudacodec/src/video_decoder.cpp

@@ -45,7 +45,44 @@
 
 #ifdef HAVE_NVCUVID
 
-#if (CUDART_VERSION < 9000)
+#if (CUDART_VERSION >= 9000)
+static const char* GetVideoCodecString(cudaVideoCodec eCodec) {
+    static struct {
+        cudaVideoCodec eCodec;
+        const char* name;
+    } aCodecName[] = {
+        { cudaVideoCodec_MPEG1,     "MPEG-1"       },
+        { cudaVideoCodec_MPEG2,     "MPEG-2"       },
+        { cudaVideoCodec_MPEG4,     "MPEG-4 (ASP)" },
+        { cudaVideoCodec_VC1,       "VC-1/WMV"     },
+        { cudaVideoCodec_H264,      "AVC/H.264"    },
+        { cudaVideoCodec_JPEG,      "M-JPEG"       },
+        { cudaVideoCodec_H264_SVC,  "H.264/SVC"    },
+        { cudaVideoCodec_H264_MVC,  "H.264/MVC"    },
+        { cudaVideoCodec_HEVC,      "H.265/HEVC"   },
+        { cudaVideoCodec_VP8,       "VP8"          },
+        { cudaVideoCodec_VP9,       "VP9"          },
+        { cudaVideoCodec_AV1,       "AV1"          },
+        { cudaVideoCodec_NumCodecs, "Invalid"      },
+        { cudaVideoCodec_YUV420,    "YUV  4:2:0"   },
+        { cudaVideoCodec_YV12,      "YV12 4:2:0"   },
+        { cudaVideoCodec_NV12,      "NV12 4:2:0"   },
+        { cudaVideoCodec_YUYV,      "YUYV 4:2:2"   },
+        { cudaVideoCodec_UYVY,      "UYVY 4:2:2"   },
+    };
+
+    if (eCodec >= 0 && eCodec <= cudaVideoCodec_NumCodecs) {
+        return aCodecName[eCodec].name;
+    }
+    for (int i = cudaVideoCodec_NumCodecs + 1; i < sizeof(aCodecName) / sizeof(aCodecName[0]); i++) {
+        if (eCodec == aCodecName[i].eCodec) {
+            return aCodecName[eCodec].name;
+        }
+    }
+    return "Unknown";
+}
+#endif
+
 static const char* GetVideoChromaFormatString(cudaVideoChromaFormat eChromaFormat) {
     static struct {
         cudaVideoChromaFormat eChromaFormat;
@@ -62,7 +99,6 @@ static const char* GetVideoChromaFormatString(cudaVideoChromaFormat eChromaForma
     }
     return "Unknown";
 }
-#endif
 
 void cv::cudacodec::detail::VideoDecoder::create(const FormatInfo& videoFormat)
 {
@@ -141,7 +177,7 @@ void cv::cudacodec::detail::VideoDecoder::create(const FormatInfo& videoFormat)
     cuSafeCall(cuCtxPopCurrent(NULL));
 
     if (!decodeCaps.bIsSupported) {
-        CV_Error(Error::StsUnsupportedFormat, "Video codec is not supported by this GPU hardware video decoder refer to Nvidia's GPU Support Matrix to confirm your GPU supports hardware decoding of the video source's codec.");
+        CV_Error(Error::StsUnsupportedFormat, std::to_string(decodeCaps.nBitDepthMinus8 + 8) + " bit " + GetVideoCodecString(_codec) + " with " + GetVideoChromaFormatString(_chromaFormat) + " chroma format is not supported by this GPU hardware video decoder.  Please refer to Nvidia's GPU Support Matrix to confirm your GPU supports hardware decoding of this video source.");
     }
 
     if (!(decodeCaps.nOutputFormatMask & (1 << surfaceFormat)))

modules/cudacodec/src/video_writer.cpp

@@ -188,22 +188,26 @@ NV_ENC_BUFFER_FORMAT EncBufferFormat(const ColorFormat colorFormat) {
     case ColorFormat::NV_IYUV: return NV_ENC_BUFFER_FORMAT_IYUV;
     case ColorFormat::NV_YUV444: return NV_ENC_BUFFER_FORMAT_YUV444;
     case ColorFormat::NV_AYUV: return NV_ENC_BUFFER_FORMAT_AYUV;
+    case ColorFormat::NV_YUV420_10BIT: return NV_ENC_BUFFER_FORMAT_YUV420_10BIT;
+    case ColorFormat::NV_YUV444_10BIT: return NV_ENC_BUFFER_FORMAT_YUV444_10BIT;
     default: return NV_ENC_BUFFER_FORMAT_UNDEFINED;
     }
 }
 
 int NChannels(const ColorFormat colorFormat) {
     switch (colorFormat) {
     case ColorFormat::BGR:
-    case ColorFormat::RGB:
-    case ColorFormat::NV_IYUV:
-    case ColorFormat::NV_YUV444: return 3;
+    case ColorFormat::RGB: return 3;
     case ColorFormat::RGBA:
     case ColorFormat::BGRA:
     case ColorFormat::NV_AYUV: return 4;
     case ColorFormat::GRAY:
     case ColorFormat::NV_NV12:
-    case ColorFormat::NV_YV12: return 1;
+    case ColorFormat::NV_IYUV:
+    case ColorFormat::NV_YV12:
+    case ColorFormat::NV_YUV420_10BIT:
+    case ColorFormat::NV_YUV444:
+    case ColorFormat::NV_YUV444_10BIT: return 1;
     default: return 0;
     }
 }

modules/cudacodec/test/test_video.cpp

@@ -1065,6 +1065,67 @@ CUDA_TEST_P(H264ToH265, Transcode)
 
 INSTANTIATE_TEST_CASE_P(CUDA_Codec, H264ToH265, ALL_DEVICES);
 
+CV_ENUM(YuvColorFormats, cudacodec::ColorFormat::NV_YUV444, cudacodec::ColorFormat::NV_YUV420_10BIT, cudacodec::ColorFormat::NV_YUV444_10BIT)
+PARAM_TEST_CASE(YUVFormats, cv::cuda::DeviceInfo, YuvColorFormats)
+{
+};
+
+CUDA_TEST_P(YUVFormats, Transcode)
+{
+    cv::cuda::setDevice(GET_PARAM(0).deviceID());
+    const std::string inputFile = std::string(cvtest::TS::ptr()->get_data_path()) + "../highgui/video/big_buck_bunny.h265";
+    const cv::cudacodec::ColorFormat writerColorFormat = static_cast<cudacodec::ColorFormat>(static_cast<int>(GET_PARAM(1)));
+    constexpr double fps = 25;
+    const cudacodec::Codec codec = cudacodec::Codec::HEVC;
+    const std::string ext = ".mp4";
+    const std::string outputFile = cv::tempfile(ext.c_str());
+    constexpr int nFrames = 5;
+    vector<Mat> bgrGs;
+    {
+        VideoCapture cap(inputFile);
+        cv::Ptr<cv::cudacodec::VideoWriter> writer;
+        Mat frame, yuv, bgr;
+        cv::cudacodec::EncoderParams params;
+        params.tuningInfo = cv::cudacodec::EncodeTuningInfo::ENC_TUNING_INFO_LOSSLESS;
+        params.rateControlMode = cv::cudacodec::EncodeParamsRcMode::ENC_PARAMS_RC_CONSTQP;
+        for (int i = 0; i < nFrames; ++i) {
+            ASSERT_TRUE(cap.read(frame));
+            ASSERT_FALSE(frame.empty());
+            cudacodec::SurfaceFormat yuvFormat = cudacodec::SurfaceFormat::SF_YUV444;
+            cudacodec::BitDepth bitDepth = cudacodec::BitDepth::EIGHT;
+            if (writerColorFormat == cudacodec::ColorFormat::NV_YUV444_10BIT) {
+                yuvFormat = cudacodec::SurfaceFormat::SF_YUV444_16Bit;
+                bitDepth = cudacodec::BitDepth::SIXTEEN;
+            }
+            else if (writerColorFormat == cudacodec::ColorFormat::NV_YUV420_10BIT){
+                yuvFormat = cudacodec::SurfaceFormat::SF_P016;
+                bitDepth = cudacodec::BitDepth::SIXTEEN;
+            }
+            generateTestImages(frame, yuv, bgr, yuvFormat, cudacodec::ColorFormat::BGR, bitDepth, false);
+            bgrGs.push_back(bgr.clone());
+            if (writer.empty())
+                writer = cv::cudacodec::createVideoWriter(outputFile, frame.size(), codec, fps, writerColorFormat, params);
+            writer->write(yuv);
+        }
+    }
+
+    {
+        cv::Ptr<cv::cudacodec::VideoReader> reader = cv::cudacodec::createVideoReader(outputFile);
+        reader->set(cudacodec::ColorFormat::BGR);
+        cv::cuda::GpuMat frame, frameGs;
+        Mat frameHost, frameGsHost;
+        for (int i = 0; i < nFrames; ++i) {
+            ASSERT_TRUE(reader->nextFrame(frame));
+            frame.download(frameHost);
+            frameGsHost = bgrGs[i];
+            const int diff = writerColorFormat == cudacodec::ColorFormat::NV_YUV420_10BIT || writerColorFormat == cudacodec::ColorFormat::NV_YUV444_10BIT ? 512 : 1;
+            EXPECT_MAT_NEAR(frameHost, frameGsHost, diff);
+        }
+    }
+    ASSERT_EQ(0, remove(outputFile.c_str()));
+}
+
+INSTANTIATE_TEST_CASE_P(CUDA_Codec, YUVFormats, testing::Combine(ALL_DEVICES, YuvColorFormats::all()));
 #endif
 
 #if defined(HAVE_NVCUVENC)