【GluonCV-Torch】torchvisionのような使い勝手でより様々なモデルが使えるライブラリを求めて

物体検出やセマンティックセグメンテーションを行う際に、torchvision,models.~~のようにAPIで複雑なモデルを呼び起こせるのがとても便利なのですが、torchvisionでは使えるモデルが少ないのが悲しいです。

似たような感覚で、もっと使えるモデルの選択肢を増やすことができないか、というところでライブラリ探しの旅に出ています。

今回は、そんな中でみつけた GluonCV-Torch についてメモしていきます。

GluonCV とは

PyTorch ではなく、MXNet というフレームワークでのコンピュータビジョン系の様々なモデルが簡単に使えるライブラリのようです。

本家で使えるモデルのバリエーションがとても充実しており、2020/06現在でドキュメントにのっていたものをざっくり表にしたのが以下です。

他にも、Image Classification や Pose Estimation のモデルも充実しています。これはアツい

詳細はこちらの Model Zoo を。https://gluon-cv.mxnet.io/contents.html

これをPyTorchで使えるように移植しているのが、GluonCV-Torch です。

ただし、現在では Image Classification と Semantic Segmenation の一部のモデルしか使えません。

セグメンテーションについては、ResNet101バックボーンの FCN, PSPNet, DeeplabV3 が使えます。

それぞれ、Pascal VOC, ADE20K での学習済みモデルが用意されています。

Tutorial

レポジトリに学習済みモデルを用いた推論のかんたんなチュートリアルが乗っています。

インストール

pip install gluoncv-torch

モデルの構築

import gluoncvth as gcv

model = gcv.models.get_deeplab_resnet101_ade(pretrained=True)

get_モデル名、で簡単にモデル構築ができます。

推論

output = model.evaluate(img)

カラーパレットも用意されており、 gcv.utils.get_mask_pallete() でマスク画像を色つけて可視化することもできます。

こちら↓がレポジトリの Quick Demo の全スクリプトです。

import torch
import gluoncvth

# Get the model
model = gluoncvth.models.get_deeplab_resnet101_ade(pretrained=True)
model.eval()

# Prepare the image
url = 'https://github.com/zhanghang1989/image-data/blob/master/encoding/' + \
    'segmentation/ade20k/ADE_val_00001142.jpg?raw=true'
filename = 'example.jpg'
img = gluoncvth.utils.load_image(
    gluoncvth.utils.download(url, filename)).unsqueeze(0)

# Make prediction
output = model.evaluate(img)
predict = torch.max(output, 1)[1].cpu().numpy() + 1

# Get color pallete for visualization
mask = gluoncvth.utils.get_mask_pallete(predict, 'ade20k')
mask.save('output.png')