ImageBindをgoogle colabで試してみた。

ImageBindとは

ImageBindは、MetaAIから出された画像、テキスト、オーディオ、Depth、Thremal Map、IMUデータの6つの異なるモダリティにわたる共通埋め込みを学習しています。クロスモーダル検索、算術的なモダリティの合成、クロスモーダル検出および生成など、新しい埋め込み表現アプリケーションの実現が可能となっています。

ImageBind: Holistic AI learning across six modalities

Visual Language modelやwhisper+Segment Anything+zerishot object detectionなど組み合わせることでクロスモーダルなAIアプリケーションを可能にしていた例はありましたが全て含んでいてしかも6つ、、凄すぎです。

リンク

Colab
Github

準備

Google Colabを開き、メニューから「ランタイム→ランタイムのタイプを変更」でランタイムを「GPU」に変更します。

環境構築

インストール手順は以下の通りです。

!git clone https://github.com/facebookresearch/ImageBind.git
!pip install git+https://github.com/facebookresearch/pytorchvideo.git@28fe037d212663c6a24f373b94cc5d478c8c1a1d timm==0.6.7 ftfy regex einops fvcore decord==0.6.0

推論

(1) データのチェック
ImageBindに用意しているデモサンプルの中身を確認しておきます。textの部分を「dog, car, bird」の中から選ぶとみれます。

%cd /content/ImageBind
import IPython
from PIL import Image

text = 'car'
display(Image.open(f".assets/{text}_image.jpg"))
IPython.display.Audio(f".assets/{text}_audio.wav")

こんな感じです。

(2)推論
推論のコードは以下の通りです。今回はサンプルにある通りImage, Audio, Textの3つで行ってみます。

import data
import torch
from models import imagebind_model
from models.imagebind_model import ModalityType

text_list=["A dog.", "A car", "A bird"]
image_paths=[".assets/dog_image.jpg", ".assets/car_image.jpg", ".assets/bird_image.jpg"]
audio_paths=[".assets/dog_audio.wav", ".assets/car_audio.wav", ".assets/bird_audio.wav"]

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

# Instantiate model
model = imagebind_model.imagebind_huge(pretrained=True)
model.eval()
model.to(device)

# Load data
inputs = {
    ModalityType.TEXT: data.load_and_transform_text(text_list, device),
    ModalityType.VISION: data.load_and_transform_vision_data(image_paths, device),
    ModalityType.AUDIO: data.load_and_transform_audio_data(audio_paths, device),
}

with torch.no_grad():
    embeddings = model(inputs)

print(
    "Vision x Text: ",
    torch.softmax(embeddings[ModalityType.VISION] @ embeddings[ModalityType.TEXT].T, dim=-1),
)
print(
    "Audio x Text: ",
    torch.softmax(embeddings[ModalityType.AUDIO] @ embeddings[ModalityType.TEXT].T, dim=-1),
)
print(
    "Vision x Audio: ",
    torch.softmax(embeddings[ModalityType.VISION] @ embeddings[ModalityType.AUDIO].T, dim=-1),
)

output

Vision x Text:  tensor([[9.9684e-01, 3.1311e-03, 2.5929e-05],
        [5.4494e-05, 9.9993e-01, 2.0353e-05],
        [4.4847e-05, 1.3246e-02, 9.8671e-01]], device='cuda:0')
Audio x Text:  tensor([[1., 0., 0.],
        [0., 1., 0.],
        [0., 0., 1.]], device='cuda:0')
Vision x Audio:  tensor([[0.8064, 0.1051, 0.0885],
        [0.1284, 0.7205, 0.1511],
        [0.0016, 0.0022, 0.9962]], device='cuda:0')

ちゃんと分類できちゃってますね。。すご。これが6GBのGPUで動く世界が来てしまっています。とち狂ってます。

Advanced Application

Depthに関するデモがなかったのでDPTというdepth estimatorを利用してやってみたいと思います。

app.py · nielsr/dpt-depth-estimation at main

(1) Depth Imageの準備
DPT largeを利用してdog, car, birdのdepth estimatorを利用してDepth Imageを取得します。

transformersをインストールします。

!pip install transformers

ライブラリのインポートです。

from transformers import DPTFeatureExtractor, DPTForDepthEstimation
import torch
import numpy as np
from PIL import Image

モデルのロードです。

feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")

depth mapを作成します。

text = "bird"  # dog, bird, car
image = Image.open(f"/content/ImageBind/.assets/{text}_image.jpg")

encoding = feature_extractor(image, return_tensors="pt")
    
# forward pass
with torch.no_grad():
  outputs = model(**encoding)
  predicted_depth = outputs.predicted_depth
    
# interpolate to original size
prediction = torch.nn.functional.interpolate(
                        predicted_depth.unsqueeze(1),
                        size=image.size[::-1],
                        mode="bicubic",
                        align_corners=False,
    ).squeeze()
output = prediction.cpu().numpy()
formatted = (output * 255 / np.max(output)).astype('uint8')
img = Image.fromarray(formatted)

display(img)  # 結果の表示
img.save(f"/content/ImageBind/.assets/{text}_depth.jpg")

depthのイメージ結果です。

(2) Depthを用いてMulti Modality
depth dataをロードするための関数を作成します。おそらくこの段階でgoogle colabを再起動したほうがいいと思われます。（メモリの関係上）

%cd /content/ImageBind
from torchvision import transforms
from PIL import Image
def load_and_transform_depth_data(depth_paths, device):
    if depth_paths is None:
        return None

    depth_ouputs = []
    for depth_path in depth_paths:
        data_transform = transforms.Compose(
            [
                transforms.Resize(
                    224, interpolation=transforms.InterpolationMode.BICUBIC
                ),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                # transforms.Normalize((0.5, ), (0.5, ))
            ]
        )
        with open(depth_path, "rb") as fopen:
            image = Image.open(fopen).convert("L")

        image = data_transform(image).to(device)
        depth_ouputs.append(image)
    return torch.stack(depth_ouputs, dim=0)

推論のコードです。

import data
import torch
from models import imagebind_model
from models.imagebind_model import ModalityType

text_list=["A dog.", "A car", "A bird"]
image_paths=[".assets/dog_image.jpg", ".assets/car_image.jpg", ".assets/bird_image.jpg"]
audio_paths=[".assets/dog_audio.wav", ".assets/car_audio.wav", ".assets/bird_audio.wav"]
depth_paths = [".assets/dog_depth.jpg", ".assets/car_depth.jpg", ".assets/bird_depth.jpg"]

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

# Instantiate model
model = imagebind_model.imagebind_huge(pretrained=True)
model.eval()
model.to(device)

# Load data
inputs = {
    ModalityType.TEXT: data.load_and_transform_text(text_list, device),
    ModalityType.VISION: data.load_and_transform_vision_data(image_paths, device),
    ModalityType.AUDIO: data.load_and_transform_audio_data(audio_paths, device),
    ModalityType.DEPTH: load_and_transform_depth_data(depth_paths, device),
}

with torch.no_grad():
    embeddings = model(inputs)

print(
    "Vision x Depth: ",
    torch.softmax(embeddings[ModalityType.VISION] @ embeddings[ModalityType.DEPTH].T, dim=-1),
)
print(
    "Text x Depth: ",
    torch.softmax(embeddings[ModalityType.TEXT] @ embeddings[ModalityType.DEPTH].T, dim=-1),
)
print(
    "Depth x Audio: ",
    torch.softmax(embeddings[ModalityType.DEPTH] @ embeddings[ModalityType.AUDIO].T, dim=-1),
)

output

Vision x Depth:  tensor([[0.3444, 0.3040, 0.3516],
        [0.3451, 0.2363, 0.4186],
        [0.3517, 0.3634, 0.2849]], device='cuda:0')
Text x Depth:  tensor([[9.5571e-01, 4.4270e-02, 1.5210e-05],
        [5.6266e-01, 4.3734e-01, 9.7014e-10],
        [4.6230e-06, 1.0000e+00, 7.2704e-15]], device='cuda:0')
Depth x Audio:  tensor([[1.9618e-01, 1.4769e-02, 7.8905e-01],
        [1.5248e-02, 4.6171e-03, 9.8014e-01],
        [1.5896e-04, 1.8075e-02, 9.8177e-01]], device='cuda:0')

え、なんか全然ダメなんですが。。なんか前処理とか違ってるのかな、、
ちょっとわからないのでIssueを投げておきました。

解決してくれーー

最後に

今回はMetaAIから出されたMultiModalityのAI toolsであるImageBindについて試してみました。
公式で紹介されているImage, Audio, Textの例はかなりいい精度で分類できているようですごい限りですね。
紹介されていなかったDepthに関してもやってみましたがこちらの方はあまりうまくできませんでした。（ひとえに自分の解決力のなさの問題な気もする）こちらが解決すれば色々な場面で利用可能なアルゴリズムですね。個人的にLLMで解決しようとしていないところがGood Pointです。
また状況のアップデート等ありましたら新しいアプリケーションについても実装してみたいですね。

今後ともLLM, Diffusion model, Image Analysis, 3Dに関連する試した記事を投稿していく予定なのでよろしくお願いします。