cudacodec::VideoReader Add python binding for videoFullRangeFlag to VideoReaderInitParams #3982
Description
Recently, I have been exploring video processing on the GPU and wish to implement my own gamma correction function. However, the results look awful and I believe it's related to the python binding for cudacodec::VideoReader not supporting videoFullRangeFlag
Test Scenarios
- Input video: h264, y420p(pc), bt.709 colorspace
- Output video: same, but with an mp4 container and custom gamma correction
- All cases are using FFmpeg as the backend
gamma_lutis just a reimplementation of FFmpeg'screate_lutand returns acv2.cuda.createLookupTable()
The following will produce a poorly-corrected video:
def pure_cv2_cuda(local_filename: str, output_filename: str = None, size=(1640, 1232)):
read_params = cv2.cudacodec.VideoReaderInitParams()
read_params.minNumDecodeSurfaces = 4
read_params.targetSz = size
reader = cv2.cudacodec.createVideoReader(
local_filename,
sourceParams=[],
params=read_params,
)
reader.set(colorFormat=cv2.cudacodec.BGR, bitDepth=cv2.CV_8UC3)
# Initialize VideoWriter
params = cv2.cudacodec.EncoderParams()
params.targetQuality = 10
params.videoFullRangeFlag = True
writer = cv2.cudacodec.createVideoWriter(
f"data/{output_filename}",
size,
cv2.cudacodec.H264,
25,
colorFormat=cv2.cudacodec.BGR,
params=params,
)
input_frame = cv2.cuda.GpuMat(size, cv2.CV_8UC3)
gamma_frame = cv2.cuda.GpuMat(size, cv2.CV_8UC3)
gamma_table = gamma_lut(gamma=2.5, mode="cv2")
while reader.nextFrame(input_frame)[0]:
gamma_table.transform(input_frame, gamma_frame)
writer.write(gamma_frame)
writer.release()However, the following scenario appears to work as intended:
def ffmpeg_to_cv2_cuda_using_gamma_table(local_filename: str, output_filename: str = None, size=(1640, 1232)):
command = [
"ffmpeg",
"-hwaccel",
"cuda",
"-i",
local_filename,
"-pix_fmt",
"bgr24",
"-vf",
"colorspace=all=bt709:iall=bt709:irange=pc:fast=1",
# "eq=gamma=2.5", # Using this instead of specifying colorspace and the gamma_lut also works
"-f",
"rawvideo",
"pipe:",
]
pipe = subprocess.Popen(
command,
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
)
params = cv2.cudacodec.EncoderParams()
params.targetQuality = 10
writer = cv2.cudacodec.createVideoWriter(
f"data/{output_filename}",
size,
cv2.cudacodec.H264,
25,
colorFormat=cv2.cudacodec.BGR,
params=params,
)
gamma_table = gamma_lut(gamma=2.5)
frame = np.zeros((size[1], size[0], 3), dtype=np.uint8)
gpu_frame = cv2.cuda.GpuMat(size, cv2.CV_8UC3)
output_frame = cv2.cuda.GpuMat(size, cv2.CV_8UC3)
while pipe.stdout.readinto(frame): # returns bytes_read, breaks on 0 bytes
gpu_frame.upload(frame)
gamma_table.transform(gpu_frame, output_frame)
writer.write(output_frame)
writer.release()
pipe.terminate()