Zenn
⭐

Docker × Gradio × FastSAM × Gemini 1.5 Pro で髪色レコメンド web アプリを作ってみた

2024/06/02に公開
Docker
Python
sam
Gradio
Gemini
tech

こんにちは、週4ゲーセン通いの1年目新卒エンジニアのみっちーです。
今回は新卒研修で生成AIを用いたアプリを作ってみたので、それの共有をさせていただこうかなと思っています。
お試しになりたい方は、以下のGitHubリポジトリからcloneしてください。
https://github.com/limich001/hair_color_recommender

想定読者

  • 業務でPythonを普段メインで用いており、生成AIの使い方の一例を知りたい人
  • GradioをDockerで使ってみたい人

前提知識

軽くDocker、Gradio、Gemini 1.5 Proについて説明します。

Docker とはアプリケーションとその依存関係（ライブラリや設定ファイルなど）を一緒にパッケージングする仮想化環境です。開発者ごとのPCの環境に依存せずアプリケーションを実行できるという点でメリットがあります。詳しくは公式ドキュメントをご参照ください。
https://www.docker.com/ja-jp/

Gradio とは機械学習アプリケーションを手軽にデモできるPythonライブラリです。いい感じにフロントエンドとバックエンドをやってくれるので、機械学習モデルの実験が簡単にできるという点がメリットです。詳しくは公式ドキュメントをご参照ください。
https://www.gradio.app/

FastSAM とは高速に物体セグメンテーションできる手法です。本アプリケーションでは髪型をセグメントするために用いています。詳しい手法に関しては以下のリンクをご参照ください。

GitHub
ArXiv

Gemini 1.5 Pro とはGoogleが開発した今話題の大規模生成モデルの一つです。今回作成したアプリケーションでは髪色レコメンド時に使用しています。性能や何が強みなのかについては以下の記事がわかりやすかたのでご参照ください。
https://weel.co.jp/media/innovator/gemini-1-5-pro/#index_id0

アプリケーションの概要図

^app_flow
アプリケーションの流れ

上図はアプリケーションの流れです。人物の顔画像を入力として、FastSAMによって髪の部分をセグメントします。また、なりたい髪色をGemini 1.5 Proに入力してDict形式のRGBの値の配列を得ます。セグメントされた部分をレコメンドされたRGBの値で塗りつぶすことによって髪色を変えるという流れになっています。
※用いた画像のリンク: https://www.pakutaso.com/20120211040post-1179.html

Gradio Components

Gradioのボタン、画像アップロード、テキスト入力の構成について以下に示します。

main.py
def hair_color_recommender(user_input) -> str:
    ### Hair Color Recommendation で説明 ###

def uploaded_image(image, state):
    new_state = {
        "user_input_text": state["user_input_text"],
        "user_input_image": image,
    }

    if state["user_input_image"] != None and state["user_input_text"] != "":
        return new_state
    else:
        return new_state

def input_text(text, state):
    new_state = {
        "user_input_text": text,
        "user_input_image": state["user_input_image"]
    }

    if state["user_input_image"] != None and state["user_input_text"] != "":
        return new_state
    else:
        return new_state

def userImge_to_haircolorImage(input_image, recommend_color_json):
    ### Hair Segmentation で説明 ###

def recommend_hair_color_image(state):
    recommend_color_json = hair_color_recommender(state["user_input_text"])
    output_image = userImge_to_haircolorImage(state["user_input_image"], recommend_color_json)

    return output_image # レコメンド後の画像を出力

with gr.Blocks() as demo:
    initial_state = {
        "user_input_text": "",
        "user_input_image":  None
    }
    state = gr.State(initial_state)

    user_input_image = gr.Image(type="pil", label="あなたの顔画像をアップロードしてください")
    user_input_text = gr.Textbox()
    hair_color_image_output_button = gr.Button(value="あなたのおすすめの髪色を表示", visible=True)
    hair_color_image_output = gr.Image(type="pil", label="あなたにおすすめの髪色です")

    user_input_image.upload(fn=uploaded_image, inputs=[user_input_image, state], outputs=[state])
    user_input_text.change(fn=input_text, inputs=[user_input_text, state], outputs=[state])
    hair_color_image_output_button.click(fn=recommend_hair_color_image, inputs=[state], outputs=[hair_color_image_output])

demo.launch(debug=True, server_name="0.0.0.0")

GradioにはBlockという概念があり、入力フォームや画像アップロードなどのUIコンポーネントとコンポーネントに対してどのようなイベントを定義するかのイベントリスナーがあります。以下にそれぞれのブロックが何をやっているかを説明します。

State

main.py
initial_state = {
        "user_input_text": "",
        "user_input_image":  None
    }
state = gr.State(initial_state)

現在のユーザの初期ステータスを定義しています。初期状態では、ユーザはテキスト入力、画像アップロードを行っていないので、それぞれ""、Noneを指定しています。

Upload Image

main.py
def uploaded_image(image, state):
    new_state = {
        "user_input_text": state["user_input_text"],
        "user_input_image": image,
    }

    if state["user_input_image"] != None and state["user_input_text"] != "":
        return new_state
    else:
        return new_state
...
user_input_image = gr.Image(type="pil", label="あなたの顔画像をアップロードしてください")
...
user_input_image.upload(fn=uploaded_image, inputs=[user_input_image, state], outputs=[state])

user_input_imageというコンポーネントに対して、uploaded_imageというイベントリスナーを定義しています。ユーザが画像アップロードを行った際、stateのuser_input_imageに入力された画像が格納されます。

Input Text

main.py
def input_text(text, state):
    new_state = {
        "user_input_text": text,
        "user_input_image": state["user_input_image"]
    }

    if state["user_input_image"] != None and state["user_input_text"] != "":
        return new_state
    else:
        return new_state
...
user_input_text = gr.Textbox()
...
user_input_text.change(fn=input_text, inputs=[user_input_text, state], outputs=[state])

user_input_textというコンポーネントに対して、input_textというイベントリスナーを定義しています。ユーザがどのような髪色になりたいかをテキストボックスに入力した際、stateのuser_input_textに入力されたテキストが格納されます。

main.py
def hair_color_recommender(user_input) -> str:
    ### Hair Color Recommendation で説明 ###
...
def userImge_to_haircolorImage(input_image, recommend_color_json):
    ### Hair Segmentation で説明 ###
...
def recommend_hair_color_image(state):
    recommend_color_json = hair_color_recommender(state["user_input_text"])
    output_image = userImge_to_haircolorImage(state["user_input_image"], recommend_color_json)
...
hair_color_image_output_button = gr.Button(value="あなたのおすすめの髪色を表示", visible=True)
    hair_color_image_output = gr.Image(type="pil", label="あなたにおすすめの髪色です")
...
hair_color_image_output_button.click(fn=recommend_hair_color_image, inputs=[state], outputs=[hair_color_image_output])

hair_color_image_output_buttonというボタンコンポーネントに対して、recommend_hair_color_imageというイベントリスナーを定義しています。hair_color_image_output_buttonがクリックされた際、hair_color_recommender(Hair Color Recommendationで詳しく説明)で入力されたテキストに対しておすすめの色がレコメンドされ、その値を用いてuserImge_to_haircolorImage(Hair Segmentationで詳しく説明)で髪型セグメントをしつつ、髪色を変化させるという処理を行っています。そして、出力としてhair_color_image_outputというコンポーネントで画像を出力するという流れになっています。

以下は出来上がった画面です。

Hair Color Recommendation

Gradio Componentsで少し触れたhair_color_recommender関数について説明します。

main.py
def hair_color_recommender(user_input) -> str:
  # Google Generative AI（Gemini API）のAPIキー設定
  genai.configure(api_key=os.environ['GEMINI_API_KEY'])

  # Geminiモデルの設定
  model = genai.GenerativeModel(
    model_name='gemini-1.5-pro-latest',
    system_instruction=[
      "あなたはヘアスタイリストです。",
      "これからユーザがおすすめの色を提案してくるので、それに基づいた回答をRGBの値で、次のJSON形式でお願いします。",
      "{\"color\": [\"RED\", \"GREEN\", \"BLUE\"]}"
    ]
  )

  # Gemini APIを使って応答を生成
  response = model.generate_content(user_input)

  # 応答をテキストとして取得
  assistant_response = response.text
  print(assistant_response)

  return json.loads(assistant_response)

関数の中身は上記のようになっています。
まず、環境変数で設定しておいたGEMINI_API_KEYを用いてGemini API Keyを設定します。Gemini API Keyは自分で取得したキーを用いてください。
次に、Geminiモデルに入力するプロンプトを設定します。今回はモデルにはgemini-1.5-pro-latestモデルを設定し、システムプロンプトとしておすすめの色を出力してくれるような文章と出力形式を設定しました。
ユーザプロンプトには、ユーザが入力したテキストを設定し、モデルの応答をテキスト化し、内容をdict形式で返すという処理を行います。99.9%(体感)の確率で{\"color\": [\"RED\", \"GREEN\", \"BLUE\"]}の形式でレスポンスが返ってくるので、エラーハンドリングは今回していません（本当はちゃんとエラーハンドリングする必要がありますが、Gradio側でエラーが返ってきた時の処理をうまいことやってくれるのでそれに甘えてます笑）。

Hair Segmentation

こちらもGradio Componentsで少し触れたuserImge_to_haircolorImage関数について説明します。

main.py
def userImge_to_haircolorImage(input_image, recommend_color_json):
  model = FastSAM('FastSAM-x.pt')
  DEVICE = "cpu"

  everything_results = model(input_image, device=DEVICE, retina_masks=True, imgsz=1024, conf=0.4, iou=0.9,)
  prompt_process = FastSAMPrompt(input_image, everything_results, device=DEVICE)

  ann = prompt_process.text_prompt(text='hair')
  output_image = prompt_process.plot_to_result(annotations=ann, withContours=False, mask_random_color=False, color_dict=recommend_color_json)

  return Image.fromarray(output_image, 'RGB')

関数の中身は上記のようになっています。
まず、FastSAMのモデルを起動しています。今回はFastSAM-x.ptを指定しています。GPU環境を整えていないため、deviceはcpuを指定しています。
次に入力された画像に対してモデルを適用して物体セグメントを行います。セグメント結果に対して、hairというプロンプトを指定して髪型をセグメントするようにモデルに指示を与えています（ここら辺はCLIPを用いてセグメント結果に対して意味づけを行っているようです。）。
髪の部分のセグメントを行い、その部分のマスクに対してGeminiからレコメンドされた色で塗りつぶすという処理を行います。この処理はFastSAMのコードでは提供されておらず、FastSAMのコードに変更を加えて任意の色を指定して塗りつぶせるようにしました。詳しくはGitHubの実装をご覧ください。

Dockerfile の設定

Dockerfileの設定について以下に示します。ENV GEMINI_API_KEY={YOUR_GEMINI_API_KEY}の部分で自分のGemini API Keyを設定してください。

FROM python:3.11-slim

WORKDIR /workspace

# 必要なビルドツール、git、OpenGLライブラリ、GLibライブラリをインストール
RUN apt-get update && \
    apt-get install -y gcc python3-dev git libgl1-mesa-glx libglib2.0-0 && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

# FastSAM ディレクトリとします
ADD FastSAM /workspace/FastSAM

# main.py は個別に追加します
ADD main.py /workspace/

# requirements.txt をインストールします
RUN pip install -r /workspace/FastSAM/requirements.txt
RUN pip install gradio==3.39.0
RUN pip install google-generativeai

# OpenAIのCLIPをgitからインストールします
RUN pip install git+https://github.com/openai/CLIP.git

ENV GEMINI_API_KEY={YOUR_GEMINI_API_KEY}

CMD ["python", "main.py"]

Demo

こちらはデモ動画になります。

まとめと今後の課題

今回は、Docker、Gradio、FastSAM、Gemini 1.5 Proを使って髪色レコメンドwebアプリを作成しました。
概ねやりたいことはできたのですが、今後アプリを発展させていく上で課題があるかな〜と思いました。

  • 現状ローカルでしか動かないアプリなのでGCPを利用してデプロイしたい
  • 髪の部分を指定した色でベタ塗りするだけだと自然な感じの髪にならないので、自然な感じに馴染ませるような手法を採用して定性的な出力結果を改善する
  • GEMINI_API_KEY の管理方法が雑で今のままだとdocker imageに直接KEYがのる感じになるので、なんとかしてセキュアに管理したい

余談：この記事を書く中で、Geminiモデルを調査している中でLLMの性能を評価してるベンチマークの羅列みたいなのあって、内部的にどんな感じで評価してるのか気になったので次回はそれについてまとめて書こかな〜と思います。それか今回のアプリをデプロイまで持っていきたいです。

Discussion