[OFA] 機械学習でImage Captioning, VQA [Python]

2022年3月24日木曜日

Artificial Intelligence

本記事では、OFAと呼ばれる機械学習手法を用いて、Image Captioning, Visual Question Answeringなどを行う方法をご紹介します。

アイキャッチ

OFA

概要

OFAは、Vision, LanguageなどのモダリティとImage Captioning, Visula Question Anseringなどのタスクを単純なシーケンスに統合する統合マルチモーダル事前トレーニングモデルです。

各タスクのトレーニングために固有のレイヤーを追加することなく、以下の様々なタスクをモデル一つで実現します。

  1. Visual Grounding
  2. Grounded Captioning
  3. Image-Text Matching
  4. Image Captioning
  5. Visual Question Answering
  6. Object Detection
  7. Image Infilling
  8. Text Infilling

ちなみに、上記アイキャッチもOFAにテキストを入力して生成してもらった画像です。

このように様々なタスクを実現するOFAですが、マルチモーダルでありながら、それぞれのタスク固有のユニモーダル手法と同等のパフォーマンスに達すると論文では述べられています。

OFAアーキテクチャ
出典: Unifying Architectures, Tasks, and Modalities Through a Simple Sequence-to-Sequence Learning Framework

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

本記事では、上記手法を用いて、任意の画像を用いてImage Captioning, VQAなどを実施する方法をご紹介します。

デモ(Colaboratory)

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

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

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

環境セットアップ

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

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

はじめに、論文提供元のGitHubからソースコードを取得します。

%cd /content

!git clone https://github.com/OFA-Sys/OFA.git

次に、学習済みモデルをダウンロードします。

%cd /content

!mkdir -p /content/OFA/checkpoints/
!wget https://ofa-silicon.oss-us-west-1.aliyuncs.com/checkpoints/ofa_large_clean.pt
!mv ofa_large_clean.pt OFA/checkpoints/ofa_large.pt

以降の動作で必要なライブラリをインストールします。
この時、事前にGoogle Colaboratoryにセットアップされているライブラリのバージョンによっては、 「RESTART RUNTIME」が表示される場合があります。その場合は「ランタイム」から「ランタイムを再起動」を押下してください。

# fairseq
# RESTART RUNTIMEが表示された場合「ランタイムを再起動」
%cd /content
!git clone https://github.com/pytorch/fairseq.git

%cd /content/fairseq
!pip install --use-feature=in-tree-build ./

requirements.txtを用いて他のライブラリもインストールします。

%cd /content/OFA

# 1行目は削除
!sed '1d' requirements.txt | xargs -I {} pip install {}

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

モデルのセットアップ

続いて、学習済みモデルのロードなどの準備をしていきます。

import torch
import numpy as np
from fairseq import checkpoint_utils, distributed_utils, options, tasks, utils
from fairseq.dataclass.utils import convert_namespace_to_omegaconf
from tasks.mm_tasks.refcoco import RefcocoTask

from models.ofa import OFAModel
from PIL import Image

import cv2
import numpy
from google.colab.patches import cv2_imshow

tasks.register_task('refcoco', RefcocoTask)

# turn on cuda if GPU is available
use_cuda = torch.cuda.is_available()
# use fp16 only when GPU is available
use_fp16 = False

# specify some options for evaluation
parser = options.get_generation_parser()
input_args = ["", "--task=refcoco", "--beam=10", "--path=checkpoints/ofa_large.pt", "--bpe-dir=utils/BPE"]
args = options.parse_args_and_arch(parser, input_args)
cfg = convert_namespace_to_omegaconf(args)

必要なライブラリをインポートした後に、モデルをロードします。

# configファイルと学習済みモデルのロード
task = tasks.setup_task(cfg.task)
models, cfg = checkpoint_utils.load_model_ensemble(
    utils.split_paths(cfg.common_eval.path),
    task=task
)

# GPUに載せる
for model in models:
    model.eval()
    if use_fp16:
        model.half()
    if use_cuda and not cfg.distributed_training.pipeline_model_parallel:
        model.cuda()
    model.prepare_for_inference_(cfg)

# generatorの初期化
generator = task.build_generator(models, cfg.generation)

次に、transform用関数を定義しておきます。

# Image transform
from torchvision import transforms
mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]

patch_resize_transform = transforms.Compose([
    lambda image: image.convert("RGB"),
    transforms.Resize((task.cfg.patch_image_size, task.cfg.patch_image_size), interpolation=Image.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize(mean=mean, std=std),
])

# Text preprocess
bos_item = torch.LongTensor([task.src_dict.bos()])
eos_item = torch.LongTensor([task.src_dict.eos()])
pad_idx = task.src_dict.pad()


def get_symbols_to_strip_from_output(generator):
    if hasattr(generator, "symbols_to_strip_from_output"):
        return generator.symbols_to_strip_from_output
    else:
        return {generator.bos, generator.eos}


def decode_fn(x, tgt_dict, bpe, generator, tokenizer=None):
    x = tgt_dict.string(x.int().cpu(), extra_symbols_to_ignore=get_symbols_to_strip_from_output(generator))
    token_result = []
    bin_result = []
    img_result = []
    for token in x.strip().split():
      if token.startswith('<bin_'):
        bin_result.append(token)
      elif token.startswith('<code_'):
        img_result.append(token)
      else:
        if bpe is not None:
          token = bpe.decode('{}'.format(token))
        if tokenizer is not None:
          token = tokenizer.decode(token)
        if token.startswith(' ') or len(token_result) == 0:
          token_result.append(token.strip())
        else:
          token_result[-1] += token

    return ' '.join(token_result), ' '.join(bin_result), ' '.join(img_result)


def coord2bin(coords, w_resize_ratio, h_resize_ratio):
    coord_list = [float(coord) for coord in coords.strip().split()]
    bin_list = []
    bin_list += ["<bin_{}>".format(int((coord_list[0] * w_resize_ratio / task.cfg.max_image_size * (task.cfg.num_bins - 1))))]
    bin_list += ["<bin_{}>".format(int((coord_list[1] * h_resize_ratio / task.cfg.max_image_size * (task.cfg.num_bins - 1))))]
    bin_list += ["<bin_{}>".format(int((coord_list[2] * w_resize_ratio / task.cfg.max_image_size * (task.cfg.num_bins - 1))))]
    bin_list += ["<bin_{}>".format(int((coord_list[3] * h_resize_ratio / task.cfg.max_image_size * (task.cfg.num_bins - 1))))]
    return ' '.join(bin_list)


def bin2coord(bins, w_resize_ratio, h_resize_ratio):
    bin_list = [int(bin[5:-1]) for bin in bins.strip().split()]
    coord_list = []
    coord_list += [bin_list[0] / (task.cfg.num_bins - 1) * task.cfg.max_image_size / w_resize_ratio]
    coord_list += [bin_list[1] / (task.cfg.num_bins - 1) * task.cfg.max_image_size / h_resize_ratio]
    coord_list += [bin_list[2] / (task.cfg.num_bins - 1) * task.cfg.max_image_size / w_resize_ratio]
    coord_list += [bin_list[3] / (task.cfg.num_bins - 1) * task.cfg.max_image_size / h_resize_ratio]
    return coord_list


def encode_text(text, length=None, append_bos=False, append_eos=False):
    line = [
      task.bpe.encode(' {}'.format(word.strip())) 
      if not word.startswith('<code_') and not word.startswith('<bin_') else word
      for word in text.strip().split()
    ]
    line = ' '.join(line)
    s = task.tgt_dict.encode_line(
        line=line,
        add_if_not_exist=False,
        append_eos=False
    ).long()
    if length is not None:
        s = s[:length]
    if append_bos:
        s = torch.cat([bos_item, s])
    if append_eos:
        s = torch.cat([s, eos_item])
    return s

def construct_sample(image: Image, instruction: str):
    patch_image = patch_resize_transform(image).unsqueeze(0)
    patch_mask = torch.tensor([True])

    instruction = encode_text(' {}'.format(instruction.lower().strip()), append_bos=True, append_eos=True).unsqueeze(0)
    instruction_length = torch.LongTensor([s.ne(pad_idx).long().sum() for s in instruction])
    sample = {
        "id":np.array(['42']),
        "net_input": {
            "src_tokens": instruction,
            "src_lengths": instruction_length,
            "patch_images": patch_image,
            "patch_masks": patch_mask,
        }
    }
    return sample
  
# Function to turn FP32 to FP16
def apply_half(t):
    if t.dtype is torch.float32:
        return t.to(dtype=torch.half)
    return t

Image Captioning

それではまず、Image Captioningタスクから試してみます。 このタスクは画像を入力すると、Captionを出力します。

入力画像は、こちらの画像を使用させて頂きます。

%cd /content/OFA
!mkdir test_imgs
%cd test_imgs

# https://www.pakutaso.com/20180239038post-15116.html
!wget https://www.pakutaso.com/shared/img/thumb/smIMGL4174_TP_V4.jpg

%cd /content/OFA
image = Image.open('/content/OFA/test_imgs/smIMGL4174_TP_V4.jpg')

Image Captioningを実行します。

instruction = "what does the image describe?"

# Construct input sample & preprocess for GPU if cuda available
sample = construct_sample(image, instruction)
sample = utils.move_to_cuda(sample) if use_cuda else sample
sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample

# Generate result
with torch.no_grad():
    hypos = task.inference_step(generator, models, sample)
    tokens1, bins1, imgs1 = decode_fn(hypos[0][0]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens2, bins2, imgs2 = decode_fn(hypos[0][1]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens3, bins3, imgs3 = decode_fn(hypos[0][2]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens4, bins4, imgs4 = decode_fn(hypos[0][3]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens5, bins5, imgs5 = decode_fn(hypos[0][4]["tokens"], task.tgt_dict, task.bpe, generator)

# display result
display(image)
print('Instruction: {}'.format(instruction))
print('OFA\'s Output1: {}, Probs: {}'.format(tokens1, hypos[0][0]["score"].exp().item()))
print('OFA\'s Output2: {}, Probs: {}'.format(tokens2, hypos[0][1]["score"].exp().item()))
print('OFA\'s Output3: {}, Probs: {}'.format(tokens3, hypos[0][2]["score"].exp().item()))
print('OFA\'s Output4: {}, Probs: {}'.format(tokens4, hypos[0][3]["score"].exp().item()))
print('OFA\'s Output5: {}, Probs: {}'.format(tokens5, hypos[0][4]["score"].exp().item()))

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

テスト画像
Instruction: what does the image describe?
OFA's Output1: how to make a smoothie with fruits and vegetables, Probs: 0.4354555010795593
OFA's Output2: young girl holding a lemon in front of a juicer, Probs: 0.3779199421405792
OFA's Output3: young girl holding a lemon in front of a food processor, Probs: 0.3760034739971161
OFA's Output4: young girl eating a fruit smoothie, Probs: 0.34694352746009827
OFA's Output5: young woman eating a fruit smoothie, Probs: 0.3448621332645416

Visual Question Answering: VQA

VQAタスクを試してみます。 このタスクは画像と質問となるテキストを入力すると、答えとなるテキストを出力します。

先ほどと同様に画像をセットアップします。

%cd /content/OFA
!mkdir test_imgs
%cd test_imgs

# https://www.pakutaso.com/20180239038post-15116.html
!wget https://www.pakutaso.com/shared/img/thumb/smIMGL4174_TP_V4.jpg

%cd /content/OFA
image = Image.open('/content/OFA/test_imgs/smIMGL4174_TP_V4.jpg')

次に、質問を定義します。
今回は左手に何を持っているか質問してみました。

instruction = "What does a woman have in her left hand?" #@param {type:"string"}

VQAを実行します。

# Construct input sample & preprocess for GPU if cuda available
sample = construct_sample(image, instruction)
sample = utils.move_to_cuda(sample) if use_cuda else sample
sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample

# Generate result
with torch.no_grad():
    hypos = task.inference_step(generator, models, sample)
    tokens1, bins1, imgs1 = decode_fn(hypos[0][0]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens2, bins2, imgs2 = decode_fn(hypos[0][1]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens3, bins3, imgs3 = decode_fn(hypos[0][2]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens4, bins4, imgs4 = decode_fn(hypos[0][3]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens5, bins5, imgs5 = decode_fn(hypos[0][4]["tokens"], task.tgt_dict, task.bpe, generator)

# display result
display(image)
print('Instruction: {}'.format(instruction))
print('OFA\'s Output1: {}, Probs: {}'.format(tokens1, hypos[0][0]["score"].exp().item()))
print('OFA\'s Output2: {}, Probs: {}'.format(tokens2, hypos[0][1]["score"].exp().item()))
print('OFA\'s Output3: {}, Probs: {}'.format(tokens3, hypos[0][2]["score"].exp().item()))
print('OFA\'s Output4: {}, Probs: {}'.format(tokens4, hypos[0][3]["score"].exp().item()))
print('OFA\'s Output5: {}, Probs: {}'.format(tokens5, hypos[0][4]["score"].exp().item()))

実行結果は以下の通りです。

テスト画像
Instruction: What does a woman have in her left hand?
OFA's Output1: an apple, Probs: 0.8495906591415405
OFA's Output2: a jar, Probs: 0.26569849252700806
OFA's Output3: an orange, Probs: 0.22921980917453766
OFA's Output4: apple, Probs: 0.17912006378173828
OFA's Output5: fruit, Probs: 0.11936525255441666

Grounded QA

Grounded QAタスクを試してみます。 このタスクは画像内の領域を座標で指定し、その領域に関する質問を入力すると、答えとなるテキストを出力します。

先ほどと同様に画像をセットアップします。

%cd /content/OFA
!mkdir test_imgs
%cd test_imgs

# https://www.pakutaso.com/20180239038post-15116.html
!wget https://www.pakutaso.com/shared/img/thumb/smIMGL4174_TP_V4.jpg

%cd /content/OFA
image = Image.open('/content/OFA/test_imgs/smIMGL4174_TP_V4.jpg')

座標と、質問を定義します。

coords = "522.0 396.0 595.0 469.0" #@param {type:"string"}
w, h = image.size
w_resize_ratio = task.cfg.patch_image_size / w
h_resize_ratio = task.cfg.patch_image_size / h
bins = coord2bin(coords, w_resize_ratio, h_resize_ratio)
question = "What's in the region?" #@param {type:"string"}
instruction = "\"" + question + " region: \" + bins"

Grounded QAを実行します。

# Construct input sample & preprocess for GPU if cuda available
sample = construct_sample(image, instruction)
sample = utils.move_to_cuda(sample) if use_cuda else sample
sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample

# Generate result
with torch.no_grad():
    hypos = task.inference_step(generator, models, sample)
    tokens1, bins1, imgs1 = decode_fn(hypos[0][0]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens2, bins2, imgs2 = decode_fn(hypos[0][1]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens3, bins3, imgs3 = decode_fn(hypos[0][2]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens4, bins4, imgs4 = decode_fn(hypos[0][3]["tokens"], task.tgt_dict, task.bpe, generator)
    tokens5, bins5, imgs5 = decode_fn(hypos[0][4]["tokens"], task.tgt_dict, task.bpe, generator)

# display result
img = cv2.cvtColor(numpy.asarray(image), cv2.COLOR_RGB2BGR)
coord_list = bin2coord(bins, w_resize_ratio, h_resize_ratio)
cv2.rectangle(
    img,
    (int(coord_list[0]), int(coord_list[1])),
    (int(coord_list[2]), int(coord_list[3])),
    (0, 255, 0),
    3
)
cv2_imshow(img)

print('Instruction: {}'.format(instruction))
print('OFA\'s Output1: {}, Probs: {}'.format(tokens1, hypos[0][0]["score"].exp().item()))
print('OFA\'s Output2: {}, Probs: {}'.format(tokens2, hypos[0][1]["score"].exp().item()))
print('OFA\'s Output3: {}, Probs: {}'.format(tokens3, hypos[0][2]["score"].exp().item()))
print('OFA\'s Output4: {}, Probs: {}'.format(tokens4, hypos[0][3]["score"].exp().item()))
print('OFA\'s Output5: {}, Probs: {}'.format(tokens5, hypos[0][4]["score"].exp().item()))

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

Grounded
Instruction: "What's in the region? region: " + bins
OFA's Output1: fruits, Probs: 0.416251003742218
OFA's Output2: apples, Probs: 0.35842400789260864
OFA's Output3: a blender, Probs: 0.28545939922332764
OFA's Output4: fruit, Probs: 0.23230677843093872
OFA's Output5: apple, Probs: 0.216787189245224

Visual Grounding

Visual Groundingタスクを試してみます。 このタスクは画像と質問を入力し、質問に該当する領域を示す座標を出力します。

先ほどと同様に画像をセットアップします。

%cd /content/OFA
!mkdir test_imgs
%cd test_imgs

# https://www.pakutaso.com/20180239038post-15116.html
!wget https://www.pakutaso.com/shared/img/thumb/smIMGL4174_TP_V4.jpg

%cd /content/OFA
image = Image.open('/content/OFA/test_imgs/smIMGL4174_TP_V4.jpg')

次に質問を定義します。
ここでは、女性が左手に持っているフルーツを示すように質問しています。

question = "Fruit that a woman has in her right hand"  #@param {type:"string"}
instruction = 'which region does the text \" ' + question + '\" describe?'

Visual Groundingを実行します。

# Construct input sample & preprocess for GPU if cuda available
sample = construct_sample(image, instruction)
sample = utils.move_to_cuda(sample) if use_cuda else sample
sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample

# Generate result
with torch.no_grad():
    hypos = task.inference_step(generator, models, sample)
    tokens, bins, imgs = decode_fn(hypos[0][0]["tokens"], task.tgt_dict, task.bpe, generator)

# display result
w_resize_ratio = task.cfg.patch_image_size / w
h_resize_ratio = task.cfg.patch_image_size / h
img = cv2.cvtColor(numpy.asarray(image), cv2.COLOR_RGB2BGR)
coord_list = bin2coord(bins, w_resize_ratio, h_resize_ratio)
cv2.rectangle(
    img,
    (int(coord_list[0]), int(coord_list[1])),
    (int(coord_list[2]), int(coord_list[3])),
    (0, 255, 0),
    3
)
cv2_imshow(img)

実行結果は以下の通りです。

Visual Grounding結果

左手と右手を認識しながら正確にフルーツの領域を認識しています。

まとめ

本記事では、OFAを用いたImage Captioning, VQA, Grounded QA, Visual Groundingを行いました。

様々なタスクをモデル一つで実現している本手法を見ると、こうして汎用型人工知能は出来上がっていくのかと考えさせられます。
これを機に機械学習に興味を持つ方が一人でもいらっしゃいましたら幸いです。

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


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

参考文献

1.  論文 - Unifying Architectures, Tasks, and Modalities Through a Simple Sequence-to-Sequence Learning Framework

2. GitHub - OFA-Sys/OFA

AIで副業ならココから!

まずは無料会員登録

プロフィール

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

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


Twitter

カテゴリ

このブログを検索

ブログ アーカイブ

TeDokology