YOLOを使ったWEBアプリを作った話

はじめに

こんにちは。株式会社アイデミーのエンジニアの和泉です。普段は、WEBアプリケーションの開発業務に携わっています。

今回は、StreamlitとYOLOを使った物体を検出するアプリを作ったので、ご紹介したいと思います。

Stremlitとは

機械学習やデータサイエンスのためのWebアプリを簡単に作成・共有できるオープンソースのPythonライブラリです。
公式サイトはこちらになります。

YOLOとは

物体検出や画像セグメンテーションを行うアルゴリズムです。
名前の由来は、「You Only Look Once（一度だけ見る）」です。
現在の最新バージョンは、8.0.228で、
公式サイトはこちらになります。

実際のソースコード

import cv2
import os
import streamlit as st
from ultralytics import YOLO

TMP_DIR_PATH = "/opt/app/tmp"
MODEL_DIR_PATH = "/opt/app/src/models"
if not os.path.exists(TMP_DIR_PATH):
    os.makedirs(TMP_DIR_PATH)

st.title("YOLO WEB App")

# ファイルのアップロード
img_file = st.file_uploader("Upload an image",
                            type=["png", "jpg", "jpeg"])

if img_file is not None:
    # アップロードした画像の保存
    file_path = os.path.join(TMP_DIR_PATH, img_file.name)
    with open(file_path, "wb") as f:
        f.write(img_file.getvalue())

    # YOLOの実行
    model = YOLO(os.path.join(MODEL_DIR_PATH, "yolov8n.pt"))
    results = model.predict(file_path, save=True)

    # 予測結果の描画
    img = cv2.imread(file_path)
    for point in results[0].boxes.xyxy:
        cv2.rectangle(img,
                      (int(point[0]), int(point[1])),
                      (int(point[2]), int(point[3])),
                      (0, 0, 255),
                      thickness=5)

    # 解析画像の保存
    analysis_img_path = os.path.join(TMP_DIR_PATH,
                                     f"analysis_{img_file.name}")
    cv2.imwrite(analysis_img_path, img)

    # 画像の表示
    st.image(analysis_img_path,
             use_column_width=True

重要な点としては、以下の部分です。
Stremlitの場合、アップロードしたファイルは、以下のように保存する必要があります。

# ファイルのアップロード
img_file = st.file_uploader("Upload an image",
                            type=["png", "jpg", "jpeg"])

if img_file is not None:
    # アップロードした画像の保存
    file_path = os.path.join(TMP_DIR_PATH, img_file.name)
    with open(file_path, "wb") as f:
        f.write(img_file.getvalue()

また、結果を元画像に出力する際には、以下のように、バウンディングボックスの結果がresults[0].boxes.xyxyに格納されておりますので、1つずつ取り出して描画するようにしてください。

for point in results[0].boxes.xyxy:
        cv2.rectangle(img,
                      (int(point[0]), int(point[1])),
                      (int(point[2]), int(point[3])),
                      (0, 0, 255),
                      thickness=5)

まとめ

今回、YOLOv8とStremlitを用いて、WEBアプリケーションを作りました。
他にも、学習する方法とかも、比較的簡単に実装できますので、試しに作ってみるなどは良いかもしれません。
もし、誰かの役に立てればと幸いです。