物体検出やセマンティックセグメンテーションを行う際に、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 の全スクリプトです。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
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') |
最後の Issue は 2019/04 の日付でしたし、今後充実していくのかは怪しいところです。
レポジトリはこちら→ https://github.com/StacyYang/gluoncv-torch
MXNet 勉強しようかな・・・。
最後まで読んでいただきありがとうございました。