时域噪声消除

概述

时域噪声消除 示例应用程序接受一个带有噪声的输入视频，降低其噪声并将结果写入磁盘。您可以定义将用于处理的后端。

说明

命令行参数为

<后端> <输入视频>

其中

• backend：cuda 或 vic（仅在 Jetson 设备上）。它定义了将执行处理的后端。
• input video：要降低噪声的视频文件；它接受 .mp4、.avi 以及可能取决于 OpenCV 支持的其他格式。

VPI 示例安装程序包含一个人工添加噪声的示例视频，位于 /opt/nvidia/vpi3/samples/assets/noisy.mp4。

这是一个调用示例

• C++

  ./vpi_sample_09_tnr cuda ../assets/noisy.mp4

• Python

  python3 main.py cuda ../assets/noisy.mp4

应用程序将处理 noisy.mp4 并创建 denoised_cuda.mp4，其中包含输入视频的降噪版本。

结果

输入视频	降噪视频

源代码

为了方便起见，这里提供了也安装在示例目录中的代码。

语言 C++ Python

import sys
import vpi
import numpy as np
from argparse import ArgumentParser
import cv2

# ----------------------------
# 解析命令行参数

parser = ArgumentParser()
parser.add_argument('backend', choices=['cuda','vic'],
 help='用于处理的后端')

parser.add_argument('input',
 help='要降噪的输入视频')

args = parser.parse_args();

if args.backend == 'cuda'
 backend = vpi.Backend.CUDA
else
 assert args.backend == 'vic'
 backend = vpi.Backend.VIC

# -----------------------------
# 打开输入和输出视频

inVideo = cv2.VideoCapture(args.input)

fourcc = cv2.VideoWriter_fourcc(*'MPEG')
inSize = (int(inVideo.get(cv2.CAP_PROP_FRAME_WIDTH)), int(inVideo.get(cv2.CAP_PROP_FRAME_HEIGHT)))
fps = inVideo.get(cv2.CAP_PROP_FPS)

outVideo = cv2.VideoWriter('denoised_python'+str(sys.version_info[0])+'_'+args.backend+'.mp4',
 fourcc, fps, inSize)

#--------------------------------------------------------------
# 使用用户指定的后端创建 TNR 对象
with backend
 tnr = vpi.TemporalNoiseReduction(inSize, vpi.Format.NV12_ER)

#--------------------------------------------------------------
# 主处理循环
curFrame = 0
while True
 curFrame+=1
 print("帧: {}".format(curFrame))

 # 读取一个输入帧
 ret, cvFrame = inVideo.read()
 if not ret
 break

 # 将其转换为 VPI 使用的 NV12_ER 格式
 with vpi.Backend.CUDA
 frame = vpi.asimage(cvFrame).convert(vpi.Format.NV12_ER)

 # 检索相应的降噪帧
 denoised = tnr(frame, preset=vpi.TNRPreset.INDOOR_MEDIUM_LIGHT, strength=1)

 # 使用 CUDA 后端将其转换为 RGB8 以进行输出
 with vpi.Backend.CUDA
 denoised = denoised.convert(vpi.Format.RGB8)

 # 将降噪帧写入输出视频
 with denoised.rlock_cpu() as data
 outVideo.write(data)

#include <opencv2/core/version.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/videoio.hpp>
#include <vpi/OpenCVInterop.hpp>

#include <vpi/Image.h>
#include <vpi/Status.h>
#include <vpi/Stream.h>
#include <vpi/algo/ConvertImageFormat.h>
#include <vpi/algo/TemporalNoiseReduction.h>

#include <algorithm>
#include <cstring> // for memset
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>

#define CHECK_STATUS(STMT) \
 do \
 { \
 VPIStatus status = (STMT); \
 if (status != VPI_SUCCESS) \
 { \
 char buffer[VPI_MAX_STATUS_MESSAGE_LENGTH]; \
 vpiGetLastStatusMessage(buffer, sizeof(buffer)); \
 std::ostringstream ss; \
 ss << vpiStatusGetName(status) << ": " << buffer; \
 throw std::runtime_error(ss.str()); \
 } \
 } while (0);

int main(int argc, char *argv[])
{
 // OpenCV 图像，将被 VPIImage 包裹。
 // 在此处定义它，以便在包装器销毁*后*销毁它
 cv::Mat cvFrame;

 // 在此处声明我们将需要的所有 VPI 对象，以便我们
 // 可以在最后销毁它们。
 VPIStream stream = NULL;
 VPIImage imgPrevious = NULL, imgCurrent = NULL, imgOutput = NULL;
 VPIImage frameBGR = NULL;
 VPIPayload tnr = NULL;

 // 主要返回值
 int retval = 0;

 try
 {
 // =============================
 // 解析命令行参数

 if (argc != 3)
 {
 throw std::runtime_error(std::string("用法: ") + argv[0] + " <vic|cuda> <input_video>");
 }

 std::string strBackend = argv[1];
 std::string strInputVideo = argv[2];

 // 现在解析后端
 VPIBackend backend;

 if (strBackend == "cuda")
 {
 backend = VPI_BACKEND_CUDA;
 }
 else if (strBackend == "vic")
 {
 backend = VPI_BACKEND_VIC;
 }
 else
 {
 throw std::runtime_error("后端 '" + strBackend + "' 无法识别，它必须是 cuda 或 vic。");
 }

 // ===============================
 // 准备输入和输出视频

 // 加载输入视频
 cv::VideoCapture invid;
 if (!invid.open(strInputVideo))
 {
 throw std::runtime_error("无法打开 '" + strInputVideo + "'");
 }

 // 使用输入的特性打开输出视频以进行写入
 int w = invid.get(cv::CAP_PROP_FRAME_WIDTH);
 int h = invid.get(cv::CAP_PROP_FRAME_HEIGHT);
 int fourcc = cv::VideoWriter::fourcc('M', 'P', 'E', 'G');
 double fps = invid.get(cv::CAP_PROP_FPS);

 // 创建输出视频
 cv::VideoWriter outVideo("denoised_" + strBackend + ".mp4", fourcc, fps, cv::Size(w, h));
 if (!outVideo.isOpened())
 {
 throw std::runtime_error("无法创建输出视频");
 }

 // =================================
 // 分配所有需要的 VPI 资源

 // 我们将使用传递的后端运行重映射算法和 CUDA 来进行图像格式
 // 转换，因此我们必须强制启用 CUDA 后端，以及
 // 所需的后端。
 CHECK_STATUS(vpiStreamCreate(VPI_BACKEND_CUDA | backend, &stream));

 CHECK_STATUS(vpiImageCreate(w, h, VPI_IMAGE_FORMAT_NV12_ER, 0, &imgPrevious));
 CHECK_STATUS(vpiImageCreate(w, h, VPI_IMAGE_FORMAT_NV12_ER, 0, &imgCurrent));
 CHECK_STATUS(vpiImageCreate(w, h, VPI_IMAGE_FORMAT_NV12_ER, 0, &imgOutput));

 // 创建配置为在室内中等光线下处理 NV12_ER
 // 帧的时域噪声消除负载
 CHECK_STATUS(vpiCreateTemporalNoiseReduction(backend, w, h, VPI_IMAGE_FORMAT_NV12_ER, VPI_TNR_DEFAULT, &tnr));

 // ====================
 // 主处理循环

 int curFrame = 0;
 while (invid.read(cvFrame))
 {
 printf("帧: %d\n", ++curFrame);

 // frameBGR 尚未分配？
 if (frameBGR == NULL)
 {
 // 创建一个包裹帧的 VPIImage
 CHECK_STATUS(vpiImageCreateWrapperOpenCVMat(cvFrame, 0, &frameBGR));
 }
 else
 {
 // 重用现有的 VPIImage 包装器来包裹新帧。
 CHECK_STATUS(vpiImageSetWrappedOpenCVMat(frameBGR, cvFrame));
 }

 // 首先将其转换为 NV12_ER
 CHECK_STATUS(vpiSubmitConvertImageFormat(stream, VPI_BACKEND_CUDA, frameBGR, imgCurrent, NULL));

 // 应用时域噪声消除
 // 对于第一帧，我们必须传递 NULL 作为上一帧，
 // 这将重置内部状态。
 VPITNRParams params;
 CHECK_STATUS(vpiInitTemporalNoiseReductionParams(&params));

 params.preset = VPI_TNR_PRESET_INDOOR_MEDIUM_LIGHT;
 params.strength = 1.0f;

 CHECK_STATUS(vpiSubmitTemporalNoiseReduction(stream, 0, tnr, curFrame == 1 ? NULL : imgPrevious, imgCurrent,
 imgOutput, &params));

 // 将输出转换回 BGR
 CHECK_STATUS(vpiSubmitConvertImageFormat(stream, VPI_BACKEND_CUDA, imgOutput, frameBGR, NULL));
 CHECK_STATUS(vpiStreamSync(stream));

 // 现在将其添加到输出视频流
 VPIImageData imgdata;
 CHECK_STATUS(vpiImageLockData(frameBGR, VPI_LOCK_READ, VPI_IMAGE_BUFFER_HOST_PITCH_LINEAR, &imgdata));

 cv::Mat outFrame;
 CHECK_STATUS(vpiImageDataExportOpenCVMat(imgdata, &outFrame));
 outVideo << outFrame;

 CHECK_STATUS(vpiImageUnlock(frameBGR));

 // 此迭代的输出将是下一个迭代的上一帧。将被丢弃的上一帧将被重用
 // 以存储下一帧。
 std::swap(imgPrevious, imgOutput);
 };
 }
 catch (std::exception &e)
 {
 std::cerr << e.what() << std::endl;
 retval = 1;
 }

 // =========================
 // 销毁所有 VPI 资源
 vpiStreamDestroy(stream);
 vpiPayloadDestroy(tnr);
 vpiImageDestroy(imgPrevious);
 vpiImageDestroy(imgCurrent);
 vpiImageDestroy(imgOutput);
 vpiImageDestroy(frameBGR);

 return retval;
}