【TouchDesigner】PyTorchでMiDaSによる深度推定

TouchDesignerでPyTorchを用いてMiDasによる深度推定を行います。

https://github.com/isl-org/MiDaS
https://pytorch.org/hub/intelisl_midas_v2/

この記事の内容は以下の環境で検証しています。

• OS: Windows 11 Home
• TouchDesigner: 2023.11340 (Python: 3.11.1)

はじめにPowershellで仮想環境を作成して必要なパッケージをインストールします。PowerShell以外を使用している場合は利用しているツールで同様の操作をしてください。また、ここではPyTorchをCUDA 12.1を用いる設定でインストールしていますが、PyTorch公式サイトのGet Startedを参考にして各々の環境に合わせてインストールしてください。

# TouchDesignerのファイルがあるフォルダへ移動
> cd \directory\to\touchdesigner\file
# pyenvでTouchDesignerのPythonバージョンに合わせる
> pyenv local 3.11.1
# 仮想環境の作成と有効化
> python -m venv .venv
> .\.venv\Scripts\activate
# 必要なパッケージのインストール
> pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
> pip install timm

インストールしたパッケージを用いて、Script TOPで入力画像をMiDaSで深度推定します。

# 仮想環境にインストールしたパッケージの読み込み
# ref: http://satoruhiga.com/post/extending-touchdesigner/
import sys, os
path = os.path.join(os.path.abspath("."), ".venv\Lib\site-packages")
if not path in sys.path:
	sys.path.append(path)
	print('Append new module search path:', path)

import numpy as np
import torch

# MiDaSモデルの選択
# model_type = "DPT_Large"
model_type = "DPT_Hybrid"
# model_type = "MiDaS_small"

# MiDaSモデルとトランスフォームの読み込み
midas = torch.hub.load("intel-isl/MiDaS", model_type)
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
midas.to(device)
midas.eval()
midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
	transform = midas_transforms.dpt_transform
else:
	transform = midas_transforms.small_transform

def onCook(scriptOp):
	# 一つ目のTOP入力の読み込み
	frame = scriptOp.inputs[0].numpyArray(delayed=True)
	if frame is None:
		return
	# アルファチャンネルの削除、numpyArrayで取得できる配列は原点が左下なので上下反転
	frame = np.flipud(frame[:, :, :3]).copy()

	# MiDaSで深度推定
	input_batch = transform(frame).to(device)
	with torch.no_grad():
		prediction = midas(input_batch)
		prediction = torch.nn.functional.interpolate(
			prediction.unsqueeze(1),
			size=frame.shape[:2],
			mode="bicubic",
			align_corners=False
		).squeeze()

	depth_frame = prediction.cpu().numpy()
	# 深度画像を上下反転(元に戻す)
	depth_frame = np.flipud(depth_frame).copy()
	# 画像として分かりやすいように正規化
	depth_frame /= np.max(depth_frame)
	# shapeを(height, width)から(height, width, 1)に変換
	depth_frame = np.expand_dims(depth_frame, axis=2)
	# 深度画像を出力
	scriptOp.copyNumpyArray(depth_frame)

Video Device In TOPでWebカメラの映像から深度推定すると以下のようになります。わかりやすいようにLookup TOPとRamp TOPを用いてグラデーションをつけています。