【Python】BiRefNetで画像の背景を除去する方法

pip install transformers pillow

from pathlib import Path
import time

import torch
from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageSegmentation


# ============================================================
# 設定
# ============================================================
# 背景除去したい画像のファイルパスを指定する！！！
INPUT_PATH = Path("target/dog/00015_Toy_Poodle.png")

# 背景透過PNG
OUTPUT_PATH = Path("output_transparent.png")

# 白黒マスク画像
MASK_PATH = Path("output_mask.png")

MODEL_NAME = "ZhengPeng7/BiRefNet"

# BiRefNetへ入力する解像度
MODEL_INPUT_SIZE = (512, 512)


# ============================================================
# モデル読み込み
# ============================================================

def load_model(
    device: torch.device,
) -> AutoModelForImageSegmentation:
    """BiRefNetを読み込む。"""

    print(f"Loading model: {MODEL_NAME}")

    model = AutoModelForImageSegmentation.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True,
    )

    model.to(device)
    model.eval()

    return model


# ============================================================
# 前処理
# ============================================================

def create_transform() -> transforms.Compose:
    """BiRefNet用の前処理を作成する。"""

    return transforms.Compose(
        [
            transforms.Resize(MODEL_INPUT_SIZE),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225],
            ),
        ]
    )


# ============================================================
# 背景除去
# ============================================================

def remove_background(
    image: Image.Image,
    model: AutoModelForImageSegmentation,
    device: torch.device,
) -> tuple[Image.Image, Image.Image]:
    """
    BiRefNetで背景を除去する。

    戻り値:
        transparent_image:
            背景を透明化したRGBA画像

        mask:
            前景を白、背景を黒としたマスク画像
    """

    rgb_image = image.convert("RGB")
    original_size = rgb_image.size

    transform = create_transform()

    input_tensor = transform(rgb_image).unsqueeze(0)

    # モデルと入力Tensorのdtypeを揃える
    model_dtype = next(model.parameters()).dtype

    input_tensor = input_tensor.to(
        device=device,
        dtype=model_dtype,
    )

    print(f"Model dtype: {model_dtype}")
    print(f"Input dtype: {input_tensor.dtype}")

    with torch.inference_mode():
        outputs = model(input_tensor)

        prediction = outputs[-1].sigmoid().float().cpu()

    prediction = prediction[0].squeeze()

    mask = transforms.ToPILImage()(prediction)

    mask = mask.resize(
        original_size,
        resample=Image.Resampling.BILINEAR,
    )

    transparent_image = rgb_image.convert("RGBA")
    transparent_image.putalpha(mask)

    return transparent_image, mask


# ============================================================
# 実行
# ============================================================

def main() -> None:
    if not INPUT_PATH.exists():
        raise FileNotFoundError(
            f"入力画像が見つかりません: {INPUT_PATH.resolve()}"
        )

    device = torch.device(
        "cuda" if torch.cuda.is_available() else "cpu"
    )

    print(f"Device: {device}")

    if device.type == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}")

    total_start = time.perf_counter()

    image = Image.open(INPUT_PATH)

    model_start = time.perf_counter()
    model = load_model(device)
    model_elapsed = time.perf_counter() - model_start

    inference_start = time.perf_counter()

    transparent_image, mask = remove_background(
        image=image,
        model=model,
        device=device,
    )

    inference_elapsed = time.perf_counter() - inference_start

    OUTPUT_PATH.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    MASK_PATH.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    transparent_image.save(OUTPUT_PATH)
    mask.save(MASK_PATH)

    total_elapsed = time.perf_counter() - total_start

    print()
    print(f"Transparent image: {OUTPUT_PATH.resolve()}")
    print(f"Mask image       : {MASK_PATH.resolve()}")
    print(f"Model load time  : {model_elapsed:.2f} sec")
    print(f"Inference time   : {inference_elapsed:.2f} sec")
    print(f"Total time       : {total_elapsed:.2f} sec")


if __name__ == "__main__":
    main()

INPUT_PATH = Path("target/dog/00015_Toy_Poodle.png")
OUTPUT_PATH = Path("output_transparent.png")
MASK_PATH = Path("output_mask.png")
MODEL_NAME = "ZhengPeng7/BiRefNet"

model = AutoModelForImageSegmentation.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
)

model.to(device)
model.eval()

transforms.Resize(MODEL_INPUT_SIZE),
transforms.ToTensor(),
transforms.Normalize(
           mean=[0.485, 0.456, 0.406],
           std=[0.229, 0.224, 0.225],
           )

input_tensor = transform(rgb_image).unsqueeze(0)

outputs = model(input_tensor)

prediction = outputs[-1].sigmoid().float().cpu()

mask = mask.resize(
        original_size,
        resample=Image.Resampling.BILINEAR,
    )

transparent_image = rgb_image.convert("RGBA")
transparent_image.putalpha(mask)