飞桨PP-HumanSeg本地实时视频推理代码解读

文章同样发布在百度AIStudio,Fork后即可在线运行，请点击这里

本人希望基于PaddleSeg对视频实时进行图像分割，但在AiStudio中检索分割和实时两个关键词后并没有得到理想的结果，大部分项目是依赖于notebook对图片或对视频切片进行分割，而不是我想要的实时。之前在论坛提问过实时对视频进行图像分割的问题，无意间发现PP-HumanSeg可以实现视频换背景，于是对其本地部署代码进行挖掘，在此进行代码解读。了解代码后，可以自己编写项目实现对目标内容的实时分割。

1、 根目录文件内容概述

返回Aistudio目录，可以看到PaddleSeg包已经添加完毕，打开contrib可以看到PP-HumanSeg源文件和代码。

这里我用本地打开的文件示例。

configs中包含的是本项目的配置文件，包含数据集训练验证方法、训练轮次、优化器、模型等一系列数据。在训练和部署过程需要使用。

data文件夹包含了用于验证的图片、背景以及视频下载方式。

datasets是数据集文件。

xx_model中提供了部署好的模型和预训练模型的下载方式。

bg_replace.py是本地的执行文件。

predict,train,val文件请查看PaddleSeg文档。

2、bg_replace.py代码解读.

def parse_args():
    parser = argparse.ArgumentParser(
        description='PP-HumanSeg inference for video')
    parser.add_argument(
        "--config",
        dest="cfg",
        help="The config file.",
        default=None,
        type=str,
        required=True)
    parser.add_argument(
        "--input_shape",
        dest="input_shape",
        help="The image shape [h, w] for net inputs.",
        nargs=2,
        default=[192, 192],
        type=int)
    parser.add_argument(
        '--img_path',
        dest='img_path',
        help='Image including human',
        type=str,
        default=None)
    parser.add_argument(
        '--video_path',
        dest='video_path',
        help='Video path for inference',
        type=str,
        default=None)
    parser.add_argument(
        '--bg_img_path',
        dest='bg_img_path',
        help=
        'Background image path for replacing. If not specified, a white background is used',
        type=str,
        default=None)
    parser.add_argument(
        '--bg_video_path',
        dest='bg_video_path',
        help='Background video path for replacing',
        type=str,
        default=None)
    parser.add_argument(
        '--save_dir',
        dest='save_dir',
        help='The directory for saving the inference results',
        type=str,
        default='./output')

        ....
        
    return parser.parse_args()

第一部分是定义命令行参数，这也是可以用 !python 命令执行的原因。代码中的help对对应命令做了介绍。

我们只需要关注--config --img_path --video_path 以及 --bg_img_path --bg_video_path这几个参数

--config参数是我们需要的config文件参数，即.yml文件

其余几个path分别对应需要的图像，视频以及要更换的背景图像，背景视频文件路径。

通过更改目标参数的default参数，就可以实现直接在本地运行代码而不需要命令行调用。

---

之后是主文件，我们来逐步看看其如何实现不同输入的背景替换功能

首先选择是否在GPU上运行，以及创建对应的文件目录

def background_replace(args):
    env_info = get_sys_env()
    args.use_gpu = True if env_info['Paddle compiled with cuda'] and env_info[
        'GPUs used'] else False
    predictor = Predictor(args)

    if not osp.exists(args.save_dir):
        os.makedirs(args.save_dir)

我们来分块解读，首先第一部分，图像背景替换

其实代码十分易读，即获取图像和背景，预测后导出

重点在于comb=predictor.run(img,bg) 这一行代码

先搁置一下，稍后看它的运行流程

    # 图像背景替换
    if args.img_path is not None:
        if not osp.exists(args.img_path):
            raise Exception('The --img_path is not existed: {}'.format(
                args.img_path))
        img = cv2.imread(args.img_path)
        #读取待预测图片
        bg = get_bg_img(args.bg_img_path, img.shape)
        #读取待替换路径
        comb = predictor.run(img, bg)
        #预测
        save_name = osp.basename(args.img_path)
        save_path = osp.join(args.save_dir, save_name)
        cv2.imwrite(save_path, comb)
        #保存

下一段是视频预测，视频预测分为输入视频预测和摄像头输入预测，二者本质上是一套流程。

进入视频预测流程，会判断是否有视频背景输入，因为图片替换不设计视频背景。

之后会通过是否有待预测视频来判断是否打开摄像头。 请读者自行打开bg_replace.py文件阅读

下面的代码定义了opencv打开视频和写入视频文件的变量。

cap_video = cv2.VideoCapture(args.video_path)
#打开摄像头
fps = cap_video.get(cv2.CAP_PROP_FPS)
width = int(cap_video.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
#取得输入参数 h，w，fps
save_name = osp.basename(args.video_path)
save_name = save_name.split('.')[0]
save_path = osp.join(args.save_dir, save_name + '.avi')
#设定保存位置等参数
cap_out = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps,(width, height))

之后是视频的逐帧预测和保存流程，这一部分二者十分相似。如图

 

唯一的区别是，摄像头读出的视频直接逐帧展示出来，而视频的导出数据会存入指定路径。

与图片预测类似的是， 使用了comb=predictor.run(img,bg) 就可以得到预测结果，那么这一步是怎么定义的？

我们马上进入infer.py查看它的定义和运行流程

在此之前，还有最后一部分代码，它定义了背景替换图片即bg这个参数的由来

很好理解，如果没有输入背景图像就用纯白色图像替换，当然你也可以通过吧ones→zeros转为黑色图片

如果输入了背景图就把其作为bg返回

 

def get_bg_img(bg_img_path, img_shape):
    if bg_img_path is None:
        bg = 255 * np.ones(img_shape)
#如果没有输入背景图像就用纯白色图像替换
    elif not osp.exists(bg_img_path):
        raise Exception(
            'The --bg_img_path is not existed: {}'.format(bg_img_path))
    else:
        bg = cv2.imread(bg_img_path)
#如果输入了背景图就把其作为bg返回
    return bg

3、infer.py代码解读

同样的，import部分不再赘述

DeployConfig的作用是吧.yml即--config中我们输入的配置文件导入进来，获得其变换以及模型的一些信息

注意，这一部分是我们部署时export.py导出的文件信息，包含了模型结构和参数等

class DeployConfig:
    def __init__(self, path):
        with codecs.open(path, 'r', 'utf-8') as file:
            self.dic = yaml.load(file, Loader=yaml.FullLoader)

        self._transforms = self._load_transforms(
            self.dic['Deploy']['transforms'])
        self._dir = os.path.dirname(path)

    @property
    def transforms(self):
        return self._transforms
#获得变换信息，如裁剪等
    @property
    def model(self):
        return os.path.join(self._dir, self.dic['Deploy']['model'])
#获得模型信息
    @property
    def params(self):
        return os.path.join(self._dir, self.dic['Deploy']['params'])
#获得参数信息
    def _load_transforms(self, t_list):
        com = manager.TRANSFORMS
        transforms = []
        for t in t_list:
            ctype = t.pop('type')
            transforms.append(com[ctype](**t))

        return transforms

Predictor是重要的一部分，还记得bg_replace.py import了infer.py中的Predictor吗，这就是预测模型的关键

它分为预处理、预测和后处理部分

我们先来看定义部分

class Predictor:
    def __init__(self, args):
        self.cfg = DeployConfig(args.cfg)
        #导入config文件
        self.args = args
        self.compose = T.Compose(self.cfg.transforms)
        #导入转换信息
        resize_h, resize_w = args.input_shape

        self.disflow = cv2.DISOpticalFlow_create(
            cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST)
        #这一部分是光流算法，可以减少闪烁
        self.prev_gray = np.zeros((resize_h, resize_w), np.uint8)
        self.prev_cfd = np.zeros((resize_h, resize_w), np.float32)
        self.is_init = True

        pred_cfg = PredictConfig(self.cfg.model, self.cfg.params)
        #配置模型和参数
        pred_cfg.disable_glog_info()
        if self.args.use_gpu:
            pred_cfg.enable_use_gpu(100, 0)
        #判断是否是GPU环境，如果是，分配初始显存并指定序号
        self.predictor = create_predictor(pred_cfg)
        #这一句生成了预测器
        if self.args.test_speed:
            self.cost_averager = TimeAverager()
        #如果开启了测速，会计算平均预测时间

第一部分是预处理过程，这一部分主要对图片进行归一化等处理，方便后续操作

输入的图像尺寸越小，这一部分的运行时间就会越快，实测对1280x720的图片，需要20ms，而对一半尺寸的图片只需要8ms

ori_shape可以在最后将图片恢复成原本大小以展示

def preprocess(self, img):
        ori_shapes = []
        processed_imgs = []
        processed_img = self.compose(img)[0]
        #主要是这一步对图像进行了归一化等操作
        processed_imgs.append(processed_img)
        ori_shapes.append(img.shape)
        return processed_imgs, ori_shapes

run函数self.predictor.get_input_names() 等函数获取输入和输出的句柄和名称

这一部分在PaddleSeg API文档中有说明，也可以参阅其他文档

有关test_speed的参数用于测速，可以先不管

运行流程是

• 获得名称和句柄
• 前处理、转换数组
• 预测
• 获得输出名称和句柄
• 后处理
• 返回

   def run(self, img, bg):
        input_names = self.predictor.get_input_names()
        input_handle = self.predictor.get_input_handle(input_names[0])
        #获取输入名称和句柄
        processed_imgs, ori_shapes = self.preprocess(img)
        #调用了预处理，即上一个函数
        data = np.array(processed_imgs)
        #转化为numpy数组
        input_handle.reshape(data.shape)
        input_handle.copy_from_cpu(data)
        #获得待预测参数
        if self.args.test_speed:
            start = time.time()

        self.predictor.run()
        #预测
        if self.args.test_speed:
            self.cost_averager.record(time.time() - start)
        output_names = self.predictor.get_output_names()
        output_handle = self.predictor.get_output_handle(output_names[0])
        output = output_handle.copy_to_cpu()
        #获取输出名称和句柄
        return self.postprocess(output, img, ori_shapes[0], bg)
        #后处理后返回参数

如果不开启光流处理， 会进入如下所示的图块，这是因为光流处理改变了其结构

第一步 pred是后处理传回的参数，一般为1x2x_x_，即已经分成了主体和背景。这里pred[:, 1, :, :] 是取出其主体信息，即待检测目标

通过np.newaxis将其扩充回4维tensor，之后送入下面程序恢复原本的形状

之后就可以进行背景替换，如果背景维度不足就加上一层维度

comb=... 代码利用alpha所含的概率信息，将属于原本图像的部分与高概率区域相乘，低概率区域取反后与背景相乘，这样就得到了背景替换的结果

    def postprocess(self, pred, img, ori_shape, bg):
        if not os.path.exists(self.args.save_dir):
            os.makedirs(self.args.save_dir)
        resize_w = pred.shape[-1]
        resize_h = pred.shape[-2]
        if self.args.soft_predict:
            if self.args.use_optic_flow:

            ....  

            
            else:
                score_map = pred[:, 1, :, :]
                score_map = score_map[np.newaxis, ...]
                score_map = reverse_transform(
                    paddle.to_tensor(score_map),
                    ori_shape,
                    self.cfg.transforms,
                    mode='bilinear')
                alpha = np.transpose(score_map.numpy().squeeze(0), [1, 2, 0])


            .... 
             
            
        # background replace
        h, w, _ = img.shape
        bg = cv2.resize(bg, (w, h))
        if bg.ndim == 2:
            bg = bg[..., np.newaxis]

        comb = (alpha * img + (1 - alpha) * bg).astype(np.uint8)
        return comb

总结

读者可以通过Debug在本地运行一遍此流程，了解各个位置参数的形态

重新梳理一下，bg_replace定义了命令行转化参数，当然我们可以通过修改其默认值后直接运行文件

在判断了是对图像或者视频预测后，将输入的图片或某一帧视频或摄像头信息送入infer.py，当然还包含了需要更换的背景信息

在infer.py中 经过前处理 预测 后处理阶段返回bg_replace.py

最后bg_repalce.py函数对返回的图片进行展示、保存等

笔者在本地1050环境下测试了模型，对摄像头输入并输出，fps可以稳定在30以上

笔者使用此流程预测了其他图像分割模型，通过修改config文件，使用PPlite算法在本地运行，针对640x360的图片或视频，实际运行帧率约10fps，当对前后处理流程优化之后，实际可以达到30fps

有任何疑问欢迎留言