Windowsで OpenMMLab の MMSegmentation v1.1.1 を使う

環境

WIndows 11
CUDA 11.7
Python 3.10

Python環境構築

pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 --index-url https://download.pytorch.org/whl/cu117
pip install mmcv==2.0.1 -f https://download.openmmlab.com/mmcv/dist/cu117/torch2.0.0/index.html
pip install openmim==0.3.9
pip install mmsegmentation==1.1.1

実行

人物のみをSegmentationしてみます。

モデルのダウンロードは自動で行われます。

import numpy as np
from PIL import Image
from mmseg.apis import MMSegInferencer

checkpoint_name = 'segformer_mit-b5_8xb2-160k_ade20k-640x640'

inferencer = MMSegInferencer(model=checkpoint_name)

result = inferencer('girl.jpg')['predictions']

mask = np.where(result==12, 255, 0).astype('uint8')
Image.fromarray(mask).save("mask.png")

結果

左が元画像、右がSegmentationの結果です。

元画像はこちらで作成したものになります。
touch-sp.hatenablog.com