cudacodec - Add 10 bit YUV420 and YUV444 encoding functionality

opencv · Dec 30, 2024 · bfea856 · bfea856
1 parent 89529d7
commit bfea856
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 4 deletions.
diff --git a/modules/cudacodec/include/opencv2/cudacodec.hpp b/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

diff --git a/modules/cudacodec/src/video_writer.cpp b/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;
     }
 }

diff --git a/modules/cudacodec/test/test_video.cpp b/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)