如何将照片或者视频中的背景图抠掉，机器学习开源项目使用 | 机器学习

目录

前言

开源项目介绍

项目结构

模型下载

测试数据下载

训练数据下载

依赖下载

测试图片的抠图效果

测试视频的抠图效果

总结

---

前言

本文主要介绍一个神奇有意思的开源项目: BackgroundMattingV2

该项目可以将照片或者视频中人的背景扣掉，变成带alpha通道的图片。怎么样？可以设想一下有什么其他用法，可以给个思路，如果可以拿到抠出背景后的图片，可以贴在任意背景上，不就可以任意替换照片人物的背景了吗？还挺有意思哦。

该项目的训练部分，我会简述。

开源项目介绍

我们进入项目GitHub页面，在README部分，我们可以看到作者给出的一些效果。还是很不错的。

 项目对照片与视频的效果不太一样。

1）照片是通过原始照片与原始照片的背景图，两张图，抠出带透明通道的png。

2）视频是通过原始视频与背景图，抠出纯绿色背景的人物视频。

项目给出的素材下载的地址很全。我下面按照项目验证效果的顺序，来把项目跑起来。

项目结构

Pycharm把Github地址项目clone下来。

颜色不一致的是我自己添加的代码，可忽略。

我们看到作者提供了现成的inference，可以直接去使用。

模型下载

模型下载地址项目README给出了：模型地址

选择pytorch里面pytorch_resnet50.pth，另一个也可以，看心情。

 

测试数据下载

我们下载一些测试数据：测试数据

选择一些图片、视频以及对应的背景图，下载下来做测试使用。

训练数据下载

这部分数据看你的需要，如果直接使用作者提供的模型可以忽略这部分。地址：dataset地址

数据还是

 

依赖下载

项目中的requirements.txt提供了项目所需要的依赖。如下：

kornia==0.4.1
tensorboard==2.3.0
torch==1.7.0
torchvision==0.8.1
tqdm==4.51.0
opencv-python==4.4.0.44
onnxruntime==1.6.0

 

需要注意一下，本机电脑的NVIDIA显卡驱动是否安装正确。具体可以参考我的另一篇文章，里面有介绍：Pycharm代码docker容器运行调试 | 机器学习系列_阿良的博客-CSDN博客介绍常规的本地化运行机器学习代码，安装Anaconda+cuda显卡驱动支持，许多文章都有介绍，不在此多做赘述了。本文主要是为了解决在工作环境中，本机电脑没有显卡，需要将程序运行在带显卡的远程服务器上。本文会介绍如何部署使用显卡的docker容器、如何使用pycharm连接docker容器运行机器学习代码。版本Pycharm: 2020.1.3docker:19.03.12python: 3.6.13demo算法: BackgroundMattingV2部署下面我会按照.https://huyi-aliang.blog.csdn.net/article/details/120556923

测试图片的抠图效果

先观察一下作者给的inferance_images.py，需要那些参数：

parser = argparse.ArgumentParser(description='Inference images')

parser.add_argument('--model-type', type=str, required=True, choices=['mattingbase', 'mattingrefine'])
parser.add_argument('--model-backbone', type=str, required=True, choices=['resnet101', 'resnet50', 'mobilenetv2'])
parser.add_argument('--model-backbone-scale', type=float, default=0.25)
parser.add_argument('--model-checkpoint', type=str, required=True)
parser.add_argument('--model-refine-mode', type=str, default='sampling', choices=['full', 'sampling', 'thresholding'])
parser.add_argument('--model-refine-sample-pixels', type=int, default=80_000)
parser.add_argument('--model-refine-threshold', type=float, default=0.7)
parser.add_argument('--model-refine-kernel-size', type=int, default=3)

parser.add_argument('--images-src', type=str, required=True)
parser.add_argument('--images-bgr', type=str, required=True)

parser.add_argument('--device', type=str, choices=['cpu', 'cuda'], default='cuda')
parser.add_argument('--num-workers', type=int, default=0, 
    help='number of worker threads used in DataLoader. Note that Windows need to use single thread (0).')
parser.add_argument('--preprocess-alignment', action='store_true')

parser.add_argument('--output-dir', type=str, required=True)
parser.add_argument('--output-types', type=str, required=True, nargs='+', choices=['com', 'pha', 'fgr', 'err', 'ref'])
parser.add_argument('-y', action='store_true')

args = parser.parse_args()

主要看一下required为True的参数，都属于必填参数。为了比较直观的看到参数的使用，我们不采用直接命令执行的方式。我们通过subprocess调用命令，代码如下：

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/11/13 22:43
# @Author  : 至尊宝
# @Site    : 
# @File    : main1.py
import subprocess

cmd = 'python inference_images.py \
        --model-type mattingrefine \
        --model-backbone resnet50 \
        --model-backbone-scale 0.25 \
        --model-refine-mode sampling \
        --model-refine-sample-pixels 80000 \
        --model-checkpoint "content/pytorch_resnet50.pth" \
        --images-src "content/img/" \
        --images-bgr "content/bgr/" \
        --output-dir "content/output/" \
        --output-type com'
print(cmd)
child = subprocess.Popen(cmd, shell=True)
child.wait()

命令说明

1、我们下载的模型是pytorch_resnet50.pth，所以model-backbone选择resnet50。

2、注意图片抠图的images-src与images-bgr为图片目录，其中根据文件名产生对应关系。

 3、output-types的类型可以支持'com', 'pha', 'fgr', 'err', 'ref'，我们就取com，最终输出的带alpha通道的图片。

执行看看效果

 

 

 

 效果很不错，头发也比较清楚的抠出来了。

测试视频的抠图效果

观察一下作者给出的inferance_video.py给出的参数。

parser = argparse.ArgumentParser(description='Inference video')

parser.add_argument('--model-type', type=str, required=True, choices=['mattingbase', 'mattingrefine'])
parser.add_argument('--model-backbone', type=str, required=True, choices=['resnet101', 'resnet50', 'mobilenetv2'])
parser.add_argument('--model-backbone-scale', type=float, default=0.25)
parser.add_argument('--model-checkpoint', type=str, required=True)
parser.add_argument('--model-refine-mode', type=str, default='sampling', choices=['full', 'sampling', 'thresholding'])
parser.add_argument('--model-refine-sample-pixels', type=int, default=80_000)
parser.add_argument('--model-refine-threshold', type=float, default=0.7)
parser.add_argument('--model-refine-kernel-size', type=int, default=3)

parser.add_argument('--video-src', type=str, required=True)
parser.add_argument('--video-bgr', type=str, required=True)
parser.add_argument('--video-target-bgr', type=str, default=None, help="Path to video onto which to composite the output (default to flat green)")
parser.add_argument('--video-resize', type=int, default=None, nargs=2)

parser.add_argument('--device', type=str, choices=['cpu', 'cuda'], default='cuda')
parser.add_argument('--preprocess-alignment', action='store_true')

parser.add_argument('--output-dir', type=str, required=True)
parser.add_argument('--output-types', type=str, required=True, nargs='+', choices=['com', 'pha', 'fgr', 'err', 'ref'])
parser.add_argument('--output-format', type=str, default='video', choices=['video', 'image_sequences'])

args = parser.parse_args()

 

我们通过subprocess调用命令，代码如下：

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/11/14 0:28
# @Author  : 至尊宝
# @Site    : 
# @File    : main2.py
import subprocess

cmd = 'python inference_video.py \
        --model-type mattingrefine \
        --model-backbone resnet50 \
        --model-backbone-scale 0.25 \
        --model-refine-mode sampling \
        --model-refine-sample-pixels 80000 \
        --model-checkpoint "content/pytorch_resnet50.pth" \
        --video-src "content/video/b8.mp4" \
        --video-bgr "content/video/b8.png" \
        --output-dir "content/output1/" \
        --output-type com'
print(cmd)
child = subprocess.Popen(cmd, shell=True)
child.wait()

参数说明

1、基本和图片处理的参数差不多，只是换了两个参数video-src与video-bgr。

2、准备好视频和背景图。

验证一下效果

 

 

效果还不错，头发还是可以看出来的。

总结

这个开源项目还是很优越的，之后我会把这个项目改造改造，变成可以直接用的工具。

说说项目缺点，都是小缺点。

1、图片的处理需要使用目录所谓参数，而不是一张图，如果需要对单一张图处理的话，仍然需要使用文件夹存储图片，有点麻烦。

2、每次都会询问是否需要创建output-dir，看一下代码中的逻辑：

# Create output directory
if os.path.exists(args.output_dir):
    if args.y or input(f'Directory {args.output_dir} already exists. Override? [Y/N]: ').lower() == 'y':
        shutil.rmtree(args.output_dir)
    else:
        exit()

 总的来说，该模型的抠图效果还是很不错的，只是需要我们二次创作。

分享：

        白马非马不算境界，白马是鹿才是境界。——《雪中悍刀行》

如果本文对你有用的话，请不要吝啬你的赞，谢谢！