[CapDec] AIで画像のキャプションを生成する

2022年11月12日土曜日

Artificial Intelligence

本記事では、CapDecと呼ばれる機械学習手法を用いて、任意の画像のキャプションを生成する方法をご紹介します。

アイキャッチ
出典: DavidHuji/CapDec

CapDec

概要

CapDecは、画像を入力し、画像の説明文となるキャプションを生成するImage captioning技術です。

CapDecは、トレーニング時にCLIPモデルと追加のテキストデータのみを使用し、キャプション付き画像を必要としません。

このアプローチは、CLIPが視覚的なembeddeingとテキストのembeddingを類似させるようにトレーニングされていることに基づき、CapDecはテキストembeddingをテキストに変換する方法をトレーニングさせることによりImage captioningを実現しています。

論文中では、このアプローチが4つのベンチマークで最先端の精度を達成していると示されています。

Architecture
出典: Text-Only Training for Image Captioning using Noise-Injected CLIP

詳細はこちらの論文をご参照ください。

本記事では上記手法を用いて、任意の画像からキャプションを生成していきます。

デモ(Colaboratory)

それでは、実際に動かしながらImage captioningタスクを試していきます。
ソースコードは本記事にも記載していますが、下記のGitHubでも取得可能です。
GitHub - Colaboratory demo

また、下記から直接Google Colaboratoryで開くこともできます。
Open In Colab

なお、このデモはPythonで実装しています。
Pythonの実装に不安がある方、Pythonを使った機械学習について詳しく勉強したい方は、以下の書籍やオンライン講座などがおすすめです。

環境セットアップ

それではセットアップしていきます。 Colaboratoryを開いたら下記を設定しGPUを使用するようにしてください。

「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更

初めににライブラリをインストールします。

%cd /content

!pip install git+https://github.com/openai/CLIP.git@d50d76daa670286dd6cacf3bcd80b5e4823fc8e1
!pip install diffusers==0.7.2 transformers scipy ftfy
!pip install --upgrade gdown

続いて、CapDecにキャプションを生成させる画像をStable Diffusionで生成するためアクセストークンを設定します。
Stable Diffusionのアクセストークン取得方法は以下をご覧ください。

access_tokens="ここにアクセストークンを設定" # @param {type:"string"}

最後にライブラリをインポートします。

%cd /content

import os

# For Stable Diffusion
import torch
from torch import autocast
device = 'cuda' if torch.cuda.is_available() else "cpu"
print("using device is", device)
from diffusers import StableDiffusionPipeline
import matplotlib.pyplot as plt

# For CapDec
from typing import Tuple, List, Union, Optional
import numpy as np

from transformers import GPT2LMHeadModel, GPT2Tokenizer
from torch import nn
import torch.nn.functional as nnf

import clip

from PIL import Image

以上で環境セットアップは完了です。

学習済みモデルのセットアップ

ここでは、Image captioningの事前学習済みをダウンロードします。

# 015.pt
model_path = './015.pt'

!gdown --id 1zfadJ41bo8HrCKTdoVsKGHWO4UBXCy9W \
        -O {model_path}

Text to Image

ここではまず、CapDecに入力する画像をStable Diffusionで生成していきます。

テキストA man and a woman posing for a picture next to tokyo tower.を入力し、画像を生成します。

# load pretrain model
stable_diffusion_model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=access_tokens)
stable_diffusion_model.to("cuda")

prompt = "A man and a woman posing for a picture next to tokyo tower." #@param {type:"string"}

!mkdir outputs

create_img = "outputs/test_01.png"

# 乱数固定
generator = torch.Generator(device).manual_seed(0)

# モデルにpromptを入力し画像生成
image = stable_diffusion_model(
    prompt,
    generator=generator,
    width=720, height=512)[0][0]
# 保存
image.save(create_img)

plt.imshow(plt.imread(create_img))
plt.axis('off')
plt.show()

Stable Diffusionが生成した画像は以下の通りです。
東京タワーらしき建物が2棟立っています。

StableDiffusion生成画像

Image captioning

それでは、本題のCapDecによるImage captioningを動かしていきます。
まず、モデルをロードします。

# モデルの定義部分を省略しています。
# Google Colab(https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/CapDec_demo.ipynb)を参照してください。

# build model
cap_model = ClipCaptionModel()

# load weight
cap_model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))) 
cap_model = cap_model.eval() 
cap_model = cap_model.to(device)

# load clip
clip_model, preprocess = clip.load("RN50x4", device=device, jit=False)
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

ロードしたCapDecを用いて、キャプションを生成していきます。

pil_image = Image.open(create_img)
image = preprocess(pil_image).unsqueeze(0).to(device)

with torch.no_grad():
  prefix = clip_model.encode_image(image).to(device, dtype=torch.float32)
  prefix_embed = cap_model.clip_project(prefix).reshape(1, 40, -1)

generated_text_prefix = generate_beam(cap_model, tokenizer, embed=prefix_embed)[0]

print("\nCaption:", generated_text_prefix)

出力結果は以下の通りです。
直訳すると高層ビルの隣の桟橋の上に立っている 2 人の人々が出力されました。
東京タワーとはキャプションされませんでしたが、概ね意味の通るキャプションが生成されています。

Caption: A couple of people standing on top of a pier next to a tall building.

Image captioning to Image

最後におまけとして、キャプションを入力に、再度Stable Diffusionで画像を生成してみます。

create_img = "outputs/test_02.png"

# 乱数固定

# モデルにpromptを入力し画像生成
image = stable_diffusion_model(
    generated_text_prefix,
    generator=generator,
    width=720, height=512)[0][0]
# 保存
image.save(create_img)

plt.imshow(plt.imread(create_img))
plt.axis('off')
plt.show()

出力結果は以下の通りです。

captioning画像

キャプションを反映して桟橋ギリギリに立つ2人の画像が生成されています。

まとめ

本記事では、CapDecを用いて任意の画像からキャプションを生成する方法をご紹介しました。
精度向上と共に、データセットの自動生成などにも活用できそうです。

また本記事では、機械学習を動かすことにフォーカスしてご紹介しました。
もう少し学術的に体系立てて学びたいという方には以下の書籍などがお勧めです。ぜひご一読下さい。


また動かせるだけから理解して応用できるエンジニアの足掛かりに下記のUdemyなどもお勧めです。

参考文献

1.  論文 - Text-Only Training for Image Captioning using Noise-Injected CLIP

2. GitHub - DavidHuji/CapDec

AIで副業ならココから!

まずは無料会員登録

プロフィール

メーカーで研究開発を行う現役エンジニア
組み込み機器開発や機会学習モデル開発に従事しています

本ブログでは最新AI技術を中心にソースコード付きでご紹介します


Twitter

カテゴリ

このブログを検索

ブログ アーカイブ

TeDokology