🤗

HuggingFaceのAIモデルで画像の特徴をテキスト化しよう

2024/05/24に公開

はじめに

こんにちは。よこやんです。
株式会社バニッシュ・スタンダードという会社でサーバーサイドエンジニアをやっています。

以前、GPTを使った新規機能実装で学んだプロンプトエンジニアリングというタイトルでブログを書かせていただきました。
AIの専門家でもないのにAIネタの記事を書いてしまい、少々こっ恥ずかしいですが、今回も懲りずにAIネタの記事を投稿させていただきます。

HuggingFaceとは?

まず、昨今のAI界隈の盛り上がりを受けて、AI系の技術記事も多数投稿されるようになっておりますが、その中で「Hugging Face」というワードを見たことある方も多いのではないでしょうか?
有名なところなので調べればわかることですが、こちらでも改めて。
Hugging FaceとはAIモデルやデータセットなどの共有、利用、管理をできるプラットフォームです。
AIモデルやデータセットなどのgithubのようなものだと思えばそこまでズレていないのではないかと思います。
LLM(大規模言語モデル)を使った開発などでは大量のデータセットが必要になるため、それを一から作る、というのはなかなか骨が折れることだと思います。
そんな時、HuggingFaceがあればこのような「車輪の再発明」を行わずに済む、というわけですね。
素晴らしい!!!

AIモデルをつかってみよう

ではこのツールの使い方を見ていきましょう。
今回はAIモデルを例にHuggingFaceを利用してみたいと思います。

早速HuggingFaceのサイト移動し、画面上部メニューにある「models」をクリックしてください。
すると2024年5月19日現在では以下のようなページに移動します。

画面左側にImage-Text-to-TextとかDepth Estimationなどが並んでいますが、「Tasks」、「Libraries」 など、AIモデルの各種属性がタグのような形で用意されており、これらを選択することで絞り込みを行うことができるようになっています。

Image-to-Text

この中から今回はImage-to-Textを選択します。
するとそのタスクにあったAIモデルが絞り込まれますので、これで目的にあったものを探しやすくなりました。
今回はサクッと使ってみたいのでサンプルコードがあるものを探してみます。
わりとダウンロード数が多いものの中から
nlpconnect/vit-gpt2-image-captioningのAIモデルを見てみます。
ライセンスもapache-2.0なので、ブログのネタに使わせていただいても多分大丈夫だと思う。。
今回解析に使う画像はこちらです。

こちらの画像データから、例えば服の特徴・カテゴリ「デニムスカート、カーディガン(茶)、ゆったり」などの情報を出力できないか、試してみましょう。

サンプルコードありましたがちょっと使いにくかったので以下のように書き換えました。

sample.py
import torch
from PIL import Image
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer

def load_model_and_tokenizer():
    model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
    feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
    tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    return model, feature_extractor, tokenizer, device

def predict_step(image_paths, model, feature_extractor, tokenizer, device):
    images = []
    for image_path in image_paths:
        i_image = Image.open(image_path)
        if i_image.mode != "RGB":
            i_image = i_image.convert(mode="RGB")
        images.append(i_image)

    pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values
    pixel_values = pixel_values.to(device)

    gen_kwargs = {"max_length": 16, "num_beams": 4}
    output_ids = model.generate(pixel_values, **gen_kwargs)

    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
    preds = [pred.strip() for pred in preds]
    return preds

def main():
    model, feature_extractor, tokenizer, device = load_model_and_tokenizer()
    image_paths = ['/root/opt/sample_coordinate.jpg']  # ここを実際の画像ファイルパスに置き換えてください
    captions = predict_step(image_paths, model, feature_extractor, tokenizer, device)
    for idx, caption in enumerate(captions):
        print(f"Caption for image {image_paths[idx]}: {caption}")

if __name__ == "__main__":
    main()

また上記コードを実行するとライブラリ不足でエラーが発生したので以下のコマンドで依存パッケージをインストールします

# pip install torch transformers pillow

再実行すると以下の通りになりました。

# python sample.py
Caption for image /root/opt/sample_coordinate.jpg: a woman in a blue dress holding a blue umbrella

「青い傘を持った青いドレスを着た女性」
とのことです。
青い傘は持ってないし、青いのはスカートだけですね。。。
もしアパレルに特化させた学習をさせた場合、もっと良い結果得られそうです。
今回は素の状態で使ってみて、どの程度できるかが確認できたのでよしとしましょう。

Image Classification

次に画像分類(Image Classification)のAIモデルも見てみましょう。
先ほどと同様にmodelsページからソートで良さそうなものを上位に上げます。
今回はgoogle/vit-base-patch16-224をみてみます。

Vision Transformer(ViT)は、大規模な画像コレクションであるImageNet-21k(解像度224x224ピクセル)で教師あり学習によって事前学習されたTransformerエンコーダーモデルで、このモデルはImageNet(ILSVRC2012とも呼ばれる)という100万枚の画像と1,000クラスを含むデータセットで微調整したもののようです。

こちらもサンプルコードを一部修正して使用してみました。

sample.py
import torch
from PIL import Image
from transformers import ViTImageProcessor, ViTForImageClassification

def load_model_and_processor():
    processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
    model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
    return processor, model

def predict(image_path, processor, model):
    image = Image.open(image_path)
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    return model.config.id2label[predicted_class_idx]

def main():
    image_path = '/root/opt/sample2.jpg'  # ここを実際の画像ファイルパスに置き換えてください
    processor, model = load_model_and_processor()
    predicted_class = predict(image_path, processor, model)
    print("Predicted class:", predicted_class)

if __name__ == "__main__":
    main()

実行すると以下のようになります。

# python sample.py
Predicted class: overskirt

うーん、、オーバースカートと出てますね。
違うと思うけど、、

まとめ

今回はHuggingFaceのAIモデルの紹介と実際の使い方について記事にしてみました。
今回画像からテキスト情報を取得するAIモデルをいじってみましたが、残念ながら思ったような結果は得られませんでした。もっと、こう、カテゴリ的なものを出力できると良かったのですが。。
しかし、今回はあくまで出来合いのものをいじってみただけで、望む結果が得られないのであれば他に良いAIモデルがないか探したり、さらには自身でも改良できないか試すなどアプローチもできそうです。

また、今回のサンプルコードは見てもらえたようにPythonで書かれています。
実は普段、Go言語とReactを用いたサービス開発を行っており、Pythonは業務では触っておりません。なので今回はこれらAIモデルを使い慣れたGo言語で動かせないか、というものを試してみたかったのですが、難しかった。
やはり適材適所というのは大事ですね。。。

最後に

皆さん、この記事を最後まで読んでいただき、ありがとうございます。私たちのプロジェクトや技術的な挑戦に興味を持っていただけたなら、さらに嬉しい限りです。そして、この機会に弊社「株式会社バニッシュ・スタンダード」では現在エンジニア、デザイナーを積極的に探していることをお伝えしたいと思います。

私たちは会社のミッションである「つまらない常識を革めるプロダクトを開発し、おもしろく生きる人で世界をいっぱいにする。」を実現するため、この理念に共感してくれる新たな仲間を求めています。

私たちと一緒に、技術の最前線で働き、新しい未来を創造していきませんか?あなたの才能と情熱を、私たちのチームで発揮してください。興味のある方は、ぜひお問い合わせください。私たちはあなたからの応募を心からお待ちしています。

https://v-standard.notion.site/Engineering-at-VANISH-STANDARD-929c17c0252c4d4a92fefc9a90b04d1a

Discussion