Pythonを使って画像の背景を一瞬で消す方法がチート級に簡単すぎた【rembg】

2024年10月16日Python基礎

タイトルの通り、画像背景を一瞬で取り除けるPythonライブラリを紹介します。

1. 概要
2. 使い方
3. 試してみた
4. まとめ

概要

使用するのは、rembgというライブラリです。
以下のように、非常に精度高く背景が取り除けることが分かります。

使い方

事前準備

対応するPythonバージョンは以下の通り。

python: >3.7, <3.13

Python 3.13には未対応のようなので注意です。

ライブラリのインストール

CPU動作の場合は以下をインストール

pip install rembg
pip install "rembg[cli]"

普段、pip installするときに見かけないダブルクオーテーションがありますが、気にせずそのまま実行して大丈夫です。
“[]"が含まれているためダブルクオーテーションが必要になっているだけです。

GPU動作の場合は以下をインストール
※onnxruntime-gpuサポート必須です。

pip install "rembg[GPU]"
pip install "rembg[gpu,cli]"

CPU版と同様に、ダブルクオーテーションごと実行して問題ありません。
バッティング回避のため、CPU版とGPU版は両方入れずに片方のみのインストールに留めてください。

サンプルコード

バイト形式で入出力

from rembg import remove

input_path = 'input.png'
output_path = 'output.png'

with open(input_path, 'rb') as i:
    with open(output_path, 'wb') as o:
        input = i.read()
        output = remove(input)
        o.write(output)

PIL（Pillow）で入出力

from rembg import remove
from PIL import Image

input_path = 'input.png'
output_path = 'output.png'

input = Image.open(input_path)
output = remove(input)
output.save(output_path)

Numpy配列（OpenCV、CV2）で入出力

from rembg import remove
import cv2

input_path = 'input.png'
output_path = 'output.png'

input = cv2.imread(input_path)
output = remove(input)
cv2.imwrite(output_path, output)

とても簡単ですね。
機械学習で背景除去をしているためか、PILやOpenCVにも対応しているのがとても好印象です。

試してみた

実際に動かしてみました。

インプットしたのは適当にフリー画像として拾ってきた以下の画像です。

初回実行時は、機械学習モデルがダウンロードされます。

Downloading data from 'https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx' to file 'C:\Users\XXX\.u2net\u2net.onnx'.
100%|###############################################| 176M/176M [00:00<?, ?B/s]

実行結果はこんな感じになりました。
ちなみにGPU版（PILで処理）を使用したためか、処理速度はかなり速かったです。