キカベン
機械学習でより便利な世の中へ
G検定対策
お問い合わせ
   

PyTorch Hub:ダウンロードしたモデル(MiDaS)で画像から深度を推測

thumb image

今日は、ちょっと簡単にPyTorchを使ったチュートリアルをやります。こういった需要がどれほどあるのかわかりませんが、試しに書きます。

では、PyTorch Hubからインテルによって提供されているMiDaSというモデルを使ってみます。

https://pytorch.org/hub/intelisl_midas_v2/

これはカラー画像から深度(距離)の予測をするものです。

ステレオ画像からでなく一枚の画像から推論することができるのがミソですね。

ちなみに、PyTorchのバージョンは3.8を使っています。

まずは、サンプルプログラムに従って、画像データ(犬の写真)をダウンロードします。

import urllib.request

# 画像がある場所のURL
url = "https://github.com/pytorch/hub/raw/master/images/dog.jpg"

# 画像を保存する先のファイル名
filename = "dog.jpg"

# 画像をダウンロード
urllib.request.urlretrieve(url, filename)

つぎに、モデルをダウンロードするのですが、まず、モデルで推論を実行する時に使うデバイスを決めます。

これは、GPU(cuda) かCPUかになります。GPUはNVIDIAのGPUです。

GPUがあれば実行速度は速いですが、このサンプル画像1枚だけならCPUでも大丈夫です。

ちなみに、2018年版のMac mini(8Gメモリ)で試してまったく問題ありませんでした。

import torch

# まず、推論を実行するデバイスを決める。

# NVIDIAのGPUがあればcudaを指定。
# さもなければ、CPUで推論を実行するのでcpuを指定。

if torch.cuda.is_available():
    device = torch.device("cuda")  
else:
    device = torch.device("cpu")

次に、訓練済みのモデルをダウンロードします。今回はMiDaS_smallという小さい方のモデルを使います。

一度ダウンロードされたモデルはローカルのフォルダ(~/.cache/torch/hub/intel-isl_MiDaS_master)に格納されるので次回からはすぐに読み込めるようになります。

また、先ほど選んでおいた実行デバイスを指定します。

# 訓練済みモデルのダウンロード
midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small")

# モデルを実行するデバイスを指定する
midas.to(device)

さらに、モードを推論モードにします。

このeval()とはevaluationの意味です。これを呼ぶことでトレーニングだけで必要なこと(dropoutなど)が実行されなくなります。

ちなみに、トレーニングする前には、train()を呼び出します。

# (トレーニングではなく)推論モードにする。
midas.eval()

画像を扱うモデルでは、通常、画像に何らかの前処理を加えます。

よって、提供されている前処理用のトランスフォームをダウンロードします。

注意:smallモデルのためのトランスフォームを使います。

# 画像の前処理用のトランスフォームをダウンロード
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")

transform = midas_transforms.small_transform

先ほどダウンロードした犬の画像を読み込んで、BGRからRGBフォーマットに変換します。

これは、モデルを訓練する時にRGBが使われていたからです。BGRのままでも推論の実行はできますが結果の精度は悪くなるでしょう(ためしに、やってみてください)。

import cv2

img = cv2.imread(filename)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

この画像のデータに前処理を施せば、推論を実行する準備ができます。

ここでは、バッチ(input_batch)と呼ばれていますが、中身は一つの画像だけです。

推論を行う時には入力データは画像が一つだけでもバッチの形式にする必要があります。

# 画像に前処理を施してバッチの準備
input_batch = transform(img).to(device)

バッチのシェイプをプリントして確認してみてください。

print(input_batch.shape)

torch.Size([1, 3, 192, 256])とでます。

最初の1がバッチサイズ。画像が一つだけなので1です。

次がチャンネルの数。RGBの3チャンネルです。先ほども述べたように、BGRでも3チャンネルなので実行はできてしまいますが。

192と256は画像の高さと幅のピクセル数です。

実は、本来の画像のサイズは1213×1546なのですが、前処理によってサイズが調節されています。

ちなみに、犬はこんな感じです。

では、推論を実行しましょう。

PyTorchのモデルを実行する際には、no_grad()でグラディエントの計算は不要であると宣言します。

グラディエントはトレーニングに必要ですが、推論では通常は使わないので、with torch.no_grad()と指定したなかで推論を実行します。これで不要なグラディエントの計算が行われなくなります。

推論するために前処理された入力バッチをモデルに渡すと、推論結果が返されます。

# 推論の実行
with torch.no_grad():
    prediction = midas(input_batch)

ここで、predictionのshapeを見てみます。

print(prediction.shape)

predictionのshapeがtorch.Size([1, 192, 256])とプリントされます。最初の1はバッチサイズです。

192と256は画像の高さと幅です。

つまり、もともと3チャンネルあったのがなくなっています。

これは、深度の予測値を収納するのに複数チャンネルは必要ないからです。

最後に、推論結果をmatplotlibで画像表示するためコードを解説します。

まず、squeeze()を呼んで無駄な次元を取り除きます。

これでtorch.Size([1, 192, 256])から1を省いてtorch.Size([192, 256])になります。

output = prediction.squeeze()

この時点でoutputはまだGPUにあるかもしれないので、cpuにうつします。もともとcpuにある場合はなにも影響がありません。

output = output.cpu()

また、outputはtorch.Tensorの型なのでnumpyに変換することでmatplotlibで取り扱えるようになります。

output = output.numpy()

最後にmatplotlibで画像を表示します。

この際、配色のパターン(cmap)にplasmaを指定しました。

これで近くにあるものがより黄色く、遠くになるものがより青色に表示されます。

いろんなパターンがあるので試してください。

https://beiznotes.org/matplot-cmap-list/

import matplotlib.pyplot as plt

fig, (ax1, ax2) = plt.subplots(1, 2)
ax1.imshow(img)
ax1.axis('off')
ax2.imshow(output, cmap='plasma')
ax2.axis('off')
plt.show()

全部まとめるとこうなります。

import cv2
import matplotlib.pyplot as plt
import torch
import urllib.request

# 画像データのダウンロード
url = "https://github.com/pytorch/hub/raw/master/images/dog.jpg"
filename = "dog.jpg"

urllib.request.urlretrieve(url, filename)

# デバイスを決める
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# モデルのダウンロード
midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small")
midas.to(device)
midas.eval()

# 前処理用のトランスフォームをダウンロード
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
transform = midas_transforms.small_transform

# 画像を読み込む
img = cv2.imread(filename)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# 画像に前処理を施してバッチの準備
input_batch = transform(img).to(device)
print(input_batch.shape)

# 推論の実行
with torch.no_grad():
    prediction = midas(input_batch)

# matplotlibで扱えるように推論結果を変換
print(prediction.shape)
output = prediction.squeeze()
print(output.shape)

output = output.cpu().numpy()

# 入力画像と結果を表示
fig, (ax1, ax2) = plt.subplots(1, 2)
ax1.imshow(img)
ax1.axis('off')
ax2.imshow(output, cmap='plasma')
ax2.axis('off')
plt.show()

応用としてはOpenCVでビデオ入力を扱いリアルタイムで深度の推測をするなどが考えられます。その際には速度を早くするためにGPUが必要になるかもしれません。

以上、お疲れ様でした。



コメントを残す

メールアドレスは公開されません。