画像から背景を分割（人体を抽出）

この記事では主に rembg を使用して写真から人体を抽出し、対応するマスクと抽出された人体画像を取得します。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
このコードには主にmain関数、画像読み込み関数、マスク取得関数、人体抽出関数が含まれています。

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

分割の例を次に示します。

質問

プログラムの実行に時間がかかる場合、またはプログラムが直接フリーズする場合は、通常、ソケットの接続に問題があります。事前トレーニングされたモデルを手動でダウンロードして、ホームフォルダーの下の .u2net フォルダーに置くことができます。ダウンロードリンク

チャット

CSDNの記事掲載基準についていくつか提案があるのですが、「記事の質が低いとレコメンデーションに影響する」ということを思い出すたびにドキドキしながら、結局自分なりに苦労して書き上げました。このプロンプトを受け取ったとき、なぜそのプロンプトを受け取ったのか分かりませんでした。その後、仕方なくカスタマーサービスに相談したところ、この章の文字数が少なすぎるため、単語を追加するよう勧められました。ただし、私の主な目的はコードを共有することです。必要なテキスト説明にコードコメントを追加しましたが、変更方法がわからないので、Small talkを後ろに置いています。CSDN がこの注意喚起を行う際には引き続き慎重であることを願っています。作成者はどの記事にも力を入れています。これは本当に不快です。

参照

Xiu, Y.、Yang, J.、Tzionas, D.、Black, MJ (2022)。アイコン: 法線から取得した暗黙の服を着た人間。コンピュータービジョンとパターン認識に関する IEEE/CVF 会議議事録 (pp. 13296-13306)。