apple/DepthProをローカルで動かす

記事の概要

リンゴ屋さんがdepth系のモデルをリリースしていたので、ローカルで動かしてみた。（VRAM8G環境。おそらくVRAM6Gあれば十分）
手順を残しておく。

ちなみにデモをするだけなら以下のページで試せる。

正直、depthmap以外の余計な情報を含んだファイルが出力されるのは使い勝手悪いと感じる。
そのため、depthmapのみを出力するプログラムも作成している。

環境

OS:Windows 11
GPU:GeForce RTX 3060 laptop
CPU:i7-10750H
memory:16G
python:3.10.11
pytorch:2.4.1
CUDA:11.8
cuDNN:9.0.1

手順

作業するフォルダでコマンドプロンプトを起動し、以下のコマンドを実行する。

# 仮想環境作成
python -m venv depthpro_env

# 仮想環境アクティベイト
depthpro_env\Scripts\activate

# フォルダ移動
cd depthpro_env

# gitクローン
git clone https://github.com/apple/ml-depth-pro.git

# フォルダ移動
cd ml-depth-pro

# ライブラリインストール
pip install -e .

# モデルを配置するフォルダ
mkdir checkpoints

以下から「depth_pro.pt」をローカルにダウンロードする。（「LFS」という文字列の右のアイコンをクリック）

上記のコマンドで「checkpoints」フォルダにダウンロードしたファイルを配置する。

これまでの手順でインストールしたpytorchはCPUのため、pytorchをアンインストールする。

pip uninstall torch torchvision torchaudio

上記のコマンドを実行すると、アンインストールするか確認されますが、アンインストールを続行する。（３回「Y」を選択する。）

アンインストールが完了したら、以下のコマンドでpytorchのCUDA版をインストールする。（CUDA:11.8の場合は以下のコマンドなので、別のバージョンの場合はコマンドの修正が必要）

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

「ml-depth-pro」フォルダで以下のコマンドを実行する。

depth-pro-run -i ./data/example.jpg

以下の結果が得られる想定。

depthmapのみが欲しい場合があると思うので、「ml-depth-pro」フォルダで以下のプログラムを実行する。

import torch
import numpy as np
from PIL import Image
import depth_pro

def generate_depth_map(input_path, output_path):
    # GPUが利用可能な場合はGPUを使用
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    
    # モデルとトランスフォームの読み込み
    model, transform = depth_pro.create_model_and_transforms()
    model = model.to(device)
    model.eval()
    
    # 画像の読み込みと前処理
    image, _, f_px = depth_pro.load_rgb(input_path)
    image = transform(image).unsqueeze(0).to(device)
    
    # 推論の実行
    with torch.no_grad():
        prediction = model.infer(image, f_px=f_px)
    
    # デプスマップの取得
    depth = prediction["depth"].squeeze().cpu().numpy()
    
    # デプスマップの正規化 (0-255の範囲に)
    depth_normalized = ((depth - depth.min()) / (depth.max() - depth.min()) * 255).astype(np.uint8)
    
    # デプスマップをグレースケール画像として保存
    depth_image = Image.fromarray(depth_normalized)
    depth_image.save(output_path)
    
    print(f"Depth map saved to {output_path}")

if __name__ == "__main__":
    input_image_path = "./data/example.jpg"  # 入力画像のパスを指定
    output_image_path = "./depth_map_output.png"  # 出力画像のパスを指定
    generate_depth_map(input_image_path, output_image_path)

以下の結果が出力される想定。