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

この記事で学ぶこと

この記事では、PyTorch Hubからインテルによって提供されているMiDaSというモデルを使った簡単なチュートリアルを行います。

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

これはカラー画像から深度（距離）の予測をするものです。ステレオ画像からでなく一枚の画像から推論することができるのがミソです。

さっそく、環境設定から始めていきましょう。

環境設定

依存するライブラリーをインストールします。

mkdir midas_test
cd midas_test

python3 -m venv venv
source venv/bin/activate
pip install --upgrade pip

pip install torch torch vision opencv-python matplotlib timm

timmはMiDaSが必要とするライブラリです。詳しくはこちら。

ちなみに、PyTorchのバージョンは3.10.7で、PyTorchのバージョンは1.13.0でした。

画像データのダウンロード

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

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があれば実行速度は速いですが、このサンプル画像１枚だけならCPUでも大丈夫です。

また、M1あるいはM2のMacだと、mpsというデバイスが使え、Metal プログラミング フレームワークを使用して、MacOS デバイスの GPU でCPUより高パフォーマンスのトレーニングを可能になります。

なお、近年のMacでもIntel CPU搭載型は、CPUしか選択できません。昔はMacBook ProにNVIDIAのGPUが載っていた時代もありましたが。。。

ちなみに、2022年版のMacBook Air（M2, メモリ16GB）でMiDaSを試してまったく問題ありませんでした。

import torch

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

# NVIDIAのGPUがあればcudaを指定。
# AppシリコンのMacならmpsを指定。
# さもなければ、CPUで推論を実行するのでcpuを指定。

if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")
print(device)

訓練済モデルのダウンロードと準備

次に、訓練済みのモデルをダウンロードします。今回は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])とでます。

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

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

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])とプリントされます。最初の１はバッチサイズです。

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

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

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

出力画像の表示

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

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

これでtorch.Size([1, 192, 256])から１を省いて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)

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

# モデルのダウンロード
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()

MiDaSからの出力結果

終わりに

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

以上、お疲れ様でした。