背景消除器

概述

此应用程序从输入视频源获取帧，在当前图像上运行算法，然后计算前景掩码。输出前景掩码将保存到视频文件中。

说明

命令行参数为

<backend> <input video>

其中

• backend: 定义将执行处理的后端。目前仅支持 cpu 和 cuda。
• input video: 输入视频文件名，它接受 .mp4、.avi 以及可能其他的格式，具体取决于 OpenCV 的支持。

这是一个示例

• C++

  ./vpi_sample_14_background_subtractor cpu ../assets/pedestrians.mp4

• Python

  python3 main.py cpu ../assets/pedestrians.mp4

该应用程序将处理 pedestrians.mp4 并创建 fgmask_cpu.mp4 和 bgimage_cpu.mp4。

结果

输入视频	前景视频	背景视频

源代码

为了方便起见，以下代码也安装在 samples 目录中。

语言 C++ Python

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

# ----------------------------
# Parse command line arguments

parser = ArgumentParser()
parser.add_argument('backend', choices=['cpu','cuda'],
 help='Backend to be used for processing')

parser.add_argument('input',
 help='Input video to be denoised')

args = parser.parse_args();

if args.backend == 'cuda'
 backend = vpi.Backend.CUDA
else
 assert args.backend == 'cpu'
 backend = vpi.Backend.CPU

# -----------------------------
# Open input and output videos

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)

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

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

#--------------------------------------------------------------
# Create the Background Subtractor object using the backend specified by the user
with backend
 bgsub = vpi.BackgroundSubtractor(inSize, vpi.Format.BGR8)

#--------------------------------------------------------------
# Main processing loop
idxFrame = 0
while True
 print("Processing frame {}".format(idxFrame))
 idxFrame+=1

 # Read one input frame
 ret, cvFrame = inVideo.read()
 if not ret
 break

 # Get the foreground mask and background image estimates
 fgmask, bgimage = bgsub(vpi.asimage(cvFrame, vpi.Format.BGR8), learnrate=0.01)

 # Mask needs to be converted to BGR8 for output
 fgmask = fgmask.convert(vpi.Format.BGR8, backend=vpi.Backend.CUDA);

 # Write images to output videos
 with fgmask.rlock_cpu(), bgimage.rlock_cpu()
 outVideoFGMask.write(fgmask.cpu())
 outVideoBGImage.write(bgimage.cpu())

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

#include <vpi/Array.h>
#include <vpi/Image.h>
#include <vpi/ImageFormat.h>
#include <vpi/Pyramid.h>
#include <vpi/Status.h>
#include <vpi/Stream.h>
#include <vpi/algo/BackgroundSubtractor.h>
#include <vpi/algo/ConvertImageFormat.h>

#include <iostream>
#include <sstream>

#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 image that will be wrapped by a VPIImage.
 // Define it here so that it's destroyed *after* wrapper is destroyed
 cv::Mat cvCurFrame;

 // VPI objects that will be used
 VPIStream stream = NULL;
 VPIImage imgCurFrame = NULL;
 VPIImage bgimage = NULL;
 VPIImage fgmask = NULL;
 VPIPayload payload = NULL;

 int retval = 0;

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

 // Parse input parameters
 std::string strBackend = argv[1];
 std::string strInputVideo = argv[2];

 VPIBackend backend;
 if (strBackend == "cpu")
 {
 backend = VPI_BACKEND_CPU;
 }
 else if (strBackend == "cuda")
 {
 backend = VPI_BACKEND_CUDA;
 }
 else
 {
 throw std::runtime_error("Backend '" + strBackend + "' not recognized.");
 }

 // Load the input video
 cv::VideoCapture invid;
 if (!invid.open(strInputVideo))
 {
 throw std::runtime_error("Can't open '" + strInputVideo + "'");
 }

 int32_t width = invid.get(cv::CAP_PROP_FRAME_WIDTH);
 int32_t height = invid.get(cv::CAP_PROP_FRAME_HEIGHT);

 // Create the stream where processing will happen. We'll use user-provided backend.
 CHECK_STATUS(vpiStreamCreate(backend, &stream));

 // Create background subtractor payload to be executed on the given backend
 // OpenCV delivers us BGR8 images, so the algorithm is configured to accept that.
 CHECK_STATUS(vpiCreateBackgroundSubtractor(backend, width, height, VPI_IMAGE_FORMAT_BGR8, &payload));

 // Create foreground image
 CHECK_STATUS(vpiImageCreate(width, height, VPI_IMAGE_FORMAT_U8, 0, &fgmask));

 // Create background image
 CHECK_STATUS(vpiImageCreate(width, height, VPI_IMAGE_FORMAT_BGR8, 0, &bgimage));

 int fourcc = cv::VideoWriter::fourcc('M', 'P', 'E', 'G');
 double fps = invid.get(cv::CAP_PROP_FPS);

 cv::VideoWriter outVideo("fgmask_" + strBackend + ".mp4", fourcc, fps, cv::Size(width, height), false);
 if (!outVideo.isOpened())
 {
 throw std::runtime_error("Can't create output video");
 }

 cv::VideoWriter bgimageVideo("bgimage_" + strBackend + ".mp4", fourcc, fps, cv::Size(width, height));
 if (!outVideo.isOpened())
 {
 throw std::runtime_error("Can't create output video");
 }

 // Fetch a new frame until video ends
 int idxFrame = 1;

 while (invid.read(cvCurFrame))
 {
 printf("Processing frame %d\n", idxFrame++);
 // Wrap frame into a VPIImage
 if (imgCurFrame == NULL)
 {
 CHECK_STATUS(vpiImageCreateWrapperOpenCVMat(cvCurFrame, 0, &imgCurFrame));
 }
 else
 {
 CHECK_STATUS(vpiImageSetWrappedOpenCVMat(imgCurFrame, cvCurFrame));
 }

 VPIBackgroundSubtractorParams params;
 CHECK_STATUS(vpiInitBackgroundSubtractorParams(&params));
 params.learningRate = 0.01;

 CHECK_STATUS(

 vpiSubmitBackgroundSubtractor(stream, backend, payload, imgCurFrame, fgmask, bgimage, &params));

 // Wait for processing to finish.
 CHECK_STATUS(vpiStreamSync(stream));

 {
 // Now add it to the output video stream
 VPIImageData imgdata;
 CHECK_STATUS(vpiImageLockData(fgmask, VPI_LOCK_READ, VPI_IMAGE_BUFFER_HOST_PITCH_LINEAR, &imgdata));

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

 outVideo << outFrame;

 CHECK_STATUS(vpiImageUnlock(fgmask));
 }

 {
 VPIImageData bgdata;
 CHECK_STATUS(vpiImageLockData(bgimage, VPI_LOCK_READ, VPI_IMAGE_BUFFER_HOST_PITCH_LINEAR, &bgdata));

 cv::Mat outFrame;
 CHECK_STATUS(vpiImageDataExportOpenCVMat(bgdata, &outFrame));

 bgimageVideo << outFrame;

 CHECK_STATUS(vpiImageUnlock(bgimage));
 }
 }
 }
 catch (std::exception &e)
 {
 std::cerr << e.what() << std::endl;
 retval = 1;
 }

 // Destroy all resources used
 vpiStreamDestroy(stream);
 vpiPayloadDestroy(payload);

 vpiImageDestroy(imgCurFrame);
 vpiImageDestroy(fgmask);
 vpiImageDestroy(bgimage);

 return retval;
}