Merge pull request #3868 from cudawarped:cudacodec_videowriter_add_full_range_flag

`cudacodec::VideoWriter` add flag for writing full range luma chroma video

Author: asmorkalov

modules/cudacodec/include/opencv2/cudacodec.hpp

@@ -186,7 +186,7 @@ struct CV_EXPORTS_W_SIMPLE EncoderParams
 public:
     CV_WRAP EncoderParams() : nvPreset(ENC_PRESET_P3), tuningInfo(ENC_TUNING_INFO_HIGH_QUALITY), encodingProfile(ENC_CODEC_PROFILE_AUTOSELECT),
         rateControlMode(ENC_PARAMS_RC_VBR), multiPassEncoding(ENC_MULTI_PASS_DISABLED), constQp({ 0,0,0 }), averageBitRate(0), maxBitRate(0),
-        targetQuality(30), gopLength(250), idrPeriod(250) {};
+        targetQuality(30), gopLength(250), idrPeriod(250), videoFullRangeFlag(false){};
     CV_PROP_RW EncodePreset nvPreset;
     CV_PROP_RW EncodeTuningInfo tuningInfo;
     CV_PROP_RW EncodeProfile encodingProfile;
@@ -198,6 +198,7 @@ struct CV_EXPORTS_W_SIMPLE EncoderParams
     CV_PROP_RW uint8_t targetQuality; //!< value 0 - 51 where video quality decreases as targetQuality increases, used with \ref ENC_PARAMS_RC_VBR.
     CV_PROP_RW int gopLength; //!< the number of pictures in one GOP, ensuring \ref idrPeriod >= \ref gopLength.
     CV_PROP_RW int idrPeriod; //!< IDR interval, ensuring \ref idrPeriod >= \ref gopLength.
+    CV_PROP_RW bool videoFullRangeFlag;//!< Indicates if the black level, luma and chroma of the source are represented using the full or limited range (AKA TV or "analogue" range) of values as defined in Annex E of the ITU-T Specification.
 };
 CV_EXPORTS bool operator==(const EncoderParams& lhs, const EncoderParams& rhs);
 

modules/cudacodec/src/video_writer.cpp

@@ -315,6 +315,13 @@ GUID EncodingPresetGuid(const EncodePreset nvPreset) {
     CV_Error(Error::StsUnsupportedFormat, msg);
 }
 
+std::string GetVideoCodecString(const GUID codec) {
+    if (codec == NV_ENC_CODEC_H264_GUID) return "AVC/H.264";
+    else if (codec == NV_ENC_CODEC_HEVC_GUID) return "H.265/HEVC";
+    else if (codec == NV_ENC_CODEC_AV1_GUID) return "AV1";
+    else return "Unknown";
+}
+
 void VideoWriterImpl::InitializeEncoder(const GUID codec, const double fps)
 {
     NV_ENC_INITIALIZE_PARAMS initializeParams = {};
@@ -334,10 +341,24 @@ void VideoWriterImpl::InitializeEncoder(const GUID codec, const double fps)
     if (initializeParams.encodeConfig->frameIntervalP > 1) {
         CV_Assert(encoderCallback->setFrameIntervalP(initializeParams.encodeConfig->frameIntervalP));
     }
-    if (codec == NV_ENC_CODEC_H264_GUID)
+    if (codec == NV_ENC_CODEC_H264_GUID) {
         initializeParams.encodeConfig->encodeCodecConfig.h264Config.idrPeriod = encoderParams.idrPeriod;
-    else if (codec == NV_ENC_CODEC_HEVC_GUID)
+        if (encoderParams.videoFullRangeFlag) {
+            initializeParams.encodeConfig->encodeCodecConfig.h264Config.h264VUIParameters.videoFullRangeFlag = 1;
+            initializeParams.encodeConfig->encodeCodecConfig.h264Config.h264VUIParameters.videoSignalTypePresentFlag = 1;
+        }
+    }
+    else if (codec == NV_ENC_CODEC_HEVC_GUID) {
         initializeParams.encodeConfig->encodeCodecConfig.hevcConfig.idrPeriod = encoderParams.idrPeriod;
+        if (encoderParams.videoFullRangeFlag) {
+            initializeParams.encodeConfig->encodeCodecConfig.hevcConfig.hevcVUIParameters.videoFullRangeFlag = 1;
+            initializeParams.encodeConfig->encodeCodecConfig.hevcConfig.hevcVUIParameters.videoSignalTypePresentFlag = 1;
+        }
+    }
+    else {
+        std::string msg = "videoFullRangeFlag is not supported by codec: " + GetVideoCodecString(codec);
+        CV_LOG_WARNING(NULL, msg);
+    }
     pEnc->CreateEncoder(&initializeParams);
 }
 

modules/cudacodec/test/test_video.cpp

@@ -1060,7 +1060,7 @@ 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)
+PARAM_TEST_CASE(YUVFormats, cv::cuda::DeviceInfo, YuvColorFormats, bool)
 {
 };
 
@@ -1069,6 +1069,7 @@ 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)));
+    const bool fullRange = GET_PARAM(2);
     constexpr double fps = 25;
     const cudacodec::Codec codec = cudacodec::Codec::HEVC;
     const std::string ext = ".mp4";
@@ -1082,6 +1083,7 @@ CUDA_TEST_P(YUVFormats, Transcode)
         cv::cudacodec::EncoderParams params;
         params.tuningInfo = cv::cudacodec::EncodeTuningInfo::ENC_TUNING_INFO_LOSSLESS;
         params.rateControlMode = cv::cudacodec::EncodeParamsRcMode::ENC_PARAMS_RC_CONSTQP;
+        params.videoFullRangeFlag = fullRange;
         for (int i = 0; i < nFrames; ++i) {
             ASSERT_TRUE(cap.read(frame));
             ASSERT_FALSE(frame.empty());
@@ -1095,7 +1097,7 @@ CUDA_TEST_P(YUVFormats, Transcode)
                 yuvFormat = cudacodec::SurfaceFormat::SF_P016;
                 bitDepth = cudacodec::BitDepth::SIXTEEN;
             }
-            generateTestImages(frame, yuv, bgr, yuvFormat, cudacodec::ColorFormat::BGR, bitDepth, false);
+            generateTestImages(frame, yuv, bgr, yuvFormat, cudacodec::ColorFormat::BGR, bitDepth, false, fullRange);
             bgrGs.push_back(bgr.clone());
             if (writer.empty())
                 writer = cv::cudacodec::createVideoWriter(outputFile, frame.size(), codec, fps, writerColorFormat, params);
@@ -1119,7 +1121,7 @@ CUDA_TEST_P(YUVFormats, Transcode)
     ASSERT_EQ(0, remove(outputFile.c_str()));
 }
 
-INSTANTIATE_TEST_CASE_P(CUDA_Codec, YUVFormats, testing::Combine(ALL_DEVICES, YuvColorFormats::all()));
+INSTANTIATE_TEST_CASE_P(CUDA_Codec, YUVFormats, testing::Combine(ALL_DEVICES, YuvColorFormats::all(), testing::Bool()));
 #endif
 
 #if defined(HAVE_NVCUVENC)