ゆるいDeep Learning

PyTorchのTorchVisionで用意されている学習済みモデルを調べてみた（識別、物体検出、セグメンテーション、ポーズ推定）

Deep Learning

モデルの学習に時間がとられている・・・ちょっとだけ試したいだけなのに

学習済みモデルを使用して、すぐに画像に関連するタスクを解決するんだ

モデルの学習をせずにかんたんにモデルを試したい

深層学習ではモデルの学習をしないとモデルの動作を確認できません。ですがPyTorchでは事前学習済みモデルが用意されており、そのバリエーションが多数あります。

下記のタスクのモデルが用意されています。

これらのモデルを使用すると簡単に各タスクを確認できます。 # 動画による各タスクの確認 では動画で具体的にこれらのモデルが使用しているところを確認してみます。

多様なモデルが用意されている！

識別

物体検出、セグメンテーション: Mask RCNNの物体検知のColab

ポーズ推定 > ポーズ検出を試せるColab

具体的なコードで確認してみよう

物体検出とセグメンテーションが可能なMask-RCNNを試してみる

Mask-RCNNはセグメンテーションと物体検出が可能なモデルです。

ライブラリを導入します。


import os
from os.path import exists, join, basename, splitext

import random
import PIL
import torchvision
import cv2
import numpy as np
import torch
torch.set_grad_enabled(False)

import time
import matplotlib
import matplotlib.pylab as plt
plt.rcParams["axes.grid"] = False

Mask-RCNNの学習済みモデルを導入して、GPUで動作可能にします。


model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
model = model.eval().cuda()

予測用の画像を取得します。

IMAGE_URL = 'https://raw.githubusercontent.com/tugstugi/dl-colab-notebooks/master/resources/dog.jpg'

image_file = basename(IMAGE_URL)
!wget -q -O {image_file} {IMAGE_URL}
plt.imshow(matplotlib.image.imread(image_file))

下記のような画像が表示されます。

画像データを前処理してGPUで処理できる形にします。モデルに入力して結果を推論します。


t = time.time()
image = PIL.Image.open(image_file)
image_tensor = torchvision.transforms.functional.to_tensor(image).cuda()
output = model([image_tensor])[0]
print('executed in %.3fs' % (time.time() - t))

取得した情報をもとに物体検出を行います。

  • ラベル情報を設定
  • ラベル情報からカラーを設定
  • 推論結果に物体検出のボックス、検出したクラス、検出した確信度を設定する
    • 確信とが0.5以上のみ対象
    • ボックスを記述するための設定
    • ラベルと確信度を記述するためのテキストの設定
    • ボックスとテキストを描画

coco_names = ['unlabeled', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'street sign', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'hat', 'backpack', 'umbrella', 'shoe', 'eye glasses', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'plate', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'mirror', 'dining table', 'window', 'desk', 'toilet', 'door', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'blender', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
colors = [[random.randint(0, 255) for _ in range(3)] for _ in coco_names]

result_image = np.array(image.copy())
for box, label, score in zip(output['boxes'], output['labels'], output['scores']):
  if score > 0.5:
    color = random.choice(colors)

    # draw box
    tl = round(0.002 * max(result_image.shape[0:2])) + 1  # line thickness
    c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
    cv2.rectangle(result_image, c1, c2, color, thickness=tl)
    # draw text
    display_txt = "%s: %.1f%%" % (coco_names[label], 100*score)
    tf = max(tl - 1, 1)  # font thickness
    t_size = cv2.getTextSize(display_txt, 0, fontScale=tl / 3, thickness=tf)[0]
    c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
    cv2.rectangle(result_image, c1, c2, color, -1)  # filled
    cv2.putText(result_image, display_txt, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

plt.figure(figsize=(20, 15))
plt.imshow(result_image)

セグメンテーションマスクの処理 – セグメンテーションマスクと確信度の情報を取得

  • 0.5以上の確信度のみマスクを設定
    • 全体のマスクから各マスクを比較して大きい値がある部分はそのマスク値に設定
    • マスク情報を表示
    • squeeze処理は余計な次元を減らす処理になります。
      • 具体例：(A×1×B×C×1×D) -> (A×B×C×D)

masks = None
for score, mask in zip(output['scores'], output['masks']):
  if score > 0.5:
    if masks is None:
      masks = mask
    else:
      masks = torch.max(masks, mask)

plt.imshow(masks.squeeze(0).cpu().numpy())
masks.cpu().shape, masks.squeeze(0).cpu().shape

これでモデルの学習をせずに簡単に試せるはず。Colabを使用すれば環境構築も不要だぞ！！

これでDeep Learningに奪われた時間を取り戻せる！！

画像は学習済みモデルで大丈夫そうだな。自然言語モデルの学習よろしく！！

ええええーーーーーー

その他の学習済みモデルとColab

学習済みモデルとColabがいくつかのタスクで用意されているので簡単に試すことができます。

Colabですぐに試してみましょう。

Close Bitnami banner
Bitnami