从图片中分割背景（提取人体）

本文主要是一个利用rembg从图片中将人体提取出来，得到相应的mask以及抠出来的人体图像！rembg是一款不错的背景移除工具，其采用UNet网络进行训练得到分割模型。

环境配置

python==3.8
#torch根据自己的[CUDA版本安装](https://pytorch.org/get-started/previous-versions/)
pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
pip install
			argparse
			opencv-python
			rembg
			pillow

代码

Infer_mask.py
此代码主要包括主函数、图像加载函数、mask获取函数以及人体提取函数。

import cv2
import numpy as np
import os 
import io
import argparse
import torch
from torchvision import transforms
from PIL import Image
from rembg.bg import remove

def load_img(img_file):

    img = cv2.imread(img_file, cv2.IMREAD_UNCHANGED)
    if len(img.shape) == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    if not img_file.endswith("png"):
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    else:
        img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)

    return img

def get_image_mask(img_file):

    mask_to_origin_tensor = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.0, ), (1.0, ))
    ])
    img_ori = load_img(img_file)
    with torch.no_grad():
        buf = io.BytesIO()
        Image.fromarray(img_ori).save(buf, format='png')
        img_pil = Image.open(
            io.BytesIO(remove(buf.getvalue()))).convert("RGBA")
    img_mask = torch.tensor(1.0) - (mask_to_origin_tensor(img_pil.split()[-1]) <
                                    torch.tensor(0.5)).float()

    return img_mask

def apply_mask(image, mask):
    """Apply the given mask to the image.
    """
    for c in range(3):
        image[:, :, c] = np.where(mask == 0,
                                  255,
                                  image[:, :, c])
    return image

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--img_path", type=str, default="./examples/5a6a25963db2f667441d5076972c207c.png")
    parser.add_argument("--save_path", type=str, default="./results/")
    parser.add_argument("--save_seg", action="store_true")
    args = parser.parse_args()

    img = cv2.imread(args.img_path, cv2.IMREAD_UNCHANGED)
    base=os.path.basename(args.img_path)
    name=os.path.splitext(base)[0]
    mask=get_image_mask(args.img_path)

    if args.save_seg:
        cv2.imwrite(args.save_path+name+'_image.png',apply_mask(img,np.squeeze(mask.numpy(),0)))
    cv2.imwrite(args.save_path+name+'_mask.png',np.squeeze(mask.numpy(),0)*255)
    print(args.save_path+name+'_mask.png')```

运行

save_seg加入则会保存分割后的人体图像。

##################run single################
#To get a mask, result will be save at './results'
python Infer_mask.py --img_path ./00025_0052.png --save_path ./ --save_seg

############Run batch mask#############
image_path='/data1/supersmpl/dy/ReCurrence/A_Tools/Mask' #要处理图片的文件夹
save_path='./'  #结尾记得加反斜杠
for file_a in ${image_path}/*g; 
    do           #所有图像
        echo 'Process' `basename $file_a`
        python Infer_mask.py --img_path $file_a --save_path $save_path #--save_seg
    done

下面是一个分割后的示例：

问题

如果你的程序运行很久才运行出来，或者直接卡死，一般就是就是sock连接的问题！你可以手动下载预训练的模型，放在home文件夹下的.u2net文件夹中！下载连接

闲谈

对CSDN发文规范这个事情有点建议，每次提醒“文章质量较低，推荐受影响”心里都会咯噔一下，毕竟是自己幸苦写下来的。提示的时候根本不知道自己是因为什么原因收到这个提示。后来没有办法就去咨询客服，客服说是因为本章字数太少，建议我加一些字，但是我主要是分享代码，必要的文字解释我都加入代码注释了，不知道怎么改了，所以有了后面的闲谈。希望CSDN在给这个提示的时候还是慎重，创作者的每一篇文章都是付出心血的，这样真的让人很难受！

Reference

Xiu, Y., Yang, J., Tzionas, D., & Black, M. J. (2022). Icon: Implicit clothed humans obtained from normals. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 13296-13306).