当前图像/视频增强算法的官方代码往往不包含在视频文件上的推理,我们试图建立一个简单的框架,可以快速验证各种图像增强算法(包括超分、去噪、纹理色彩增强、超帧、sdr2hdr等)在视频上的效果。
ffmpeg具有强大的功能,但编写ffmpeg filter比较耗时,而ffmpeg-python使得我们可以经过简单改写算法代码,就可以完成视频推理。
ffmpeg with x264/x265 (or other encoder you need) support
Required Python packages: numpy ffmpeg-python
pip install ffmpeg-python
继承video_infer创建新类,并重写单帧推理函数 def forward(self, x), 注意x为ffmpeg的解除的yuv帧,尺寸为[3, h, w]
将彩色视频转为黑白视频
from ffmpeg_infer import video_infer
class gray_video_infer(video_infer):
def __init__(self, in_path, out_path, encode_params, model=None, scale=1, in_pix_fmt="yuv444p", out_pix_fmt="yuv444p"):
super(gray_video_infer, self).__init__(in_path, out_path, encode_params, model, scale, in_pix_fmt, out_pix_fmt)
def forward(self, x):
y = x.copy()
y[1, :, :] = 128
y[2, :, :] = 128
return y
in_path = "input.mp4"
out_path = "gray.mp4"
encode_params = ("libx264", "x264opts", "qp=12:bframes=3") # 输出视频编码参数
gray = gray_video_infer(in_path, out_path, encode_params, out_pix_fmt="yuv444p")
gray.infer()
del gray