その他

【初心者向け】新入社員がAIでAIを勉強!! | ゴミ判別プログラム その3

その他
この記事は約7分で読めます。

はじめに

みなさん、こんにちは。インプルの新入社員 てる です!
4月から始まったrambleの記事投稿も、今回で3回目を迎えました。1回目からご覧いただいている皆様、ありがとうございます。まだご覧になられていない方は、ぜひ下記のリンクからご覧ください!

webアプリ化してみよう!

前回までの復習と、今回の内容

前回までは、Pytorchを使ってゴミ判別プログラムを作っていました。過学習でうまく判定してくれないこともありましたが、フレームワークを変えたことで正しくゴミを判別することができるようになりました。

今までは、ターミナルやコマンドプロンプトでコマンドを入力して操作するCUI(Character User Interface)という方法で判別を実施してきましたが、これではコマンドを知っている人しか操作できません。だれでも操作できるようにするためには、普段パソコンを使うように、マウスポインターなどで操作をするGUI(Graphical User Interface)という操作方法が一般的です。

そこで今回は、機械学習させてできたゴミ判別プログラムを、実際のwebサイトのような形で使えるように、webアプリ化していこうと思います!

今回つかうもの

今回はStreamlitという、フレームワークを使用していこうと思います。

Streamlit • A faster way to build and share data apps
Streamlit is an open-source Python framework for data scientists and AI/ML engineers to deliver interactive data apps – in only a few lines of code.

Streamlitは、PythonでWebアプリを公開したり開発したりできるフレームワークです。簡単なコードでさまざまなUIを再現できるので、私のような初心者向けのフレームワークかもしれません。早速作っていきましょう!

作っていこう!

Streamlit をインストール

まずはStreamlitをインストールします。
仮想環境を起動して、ターミナルに以下のコマンドを入力しましょう。実行すると、Streamlitのインストールが開始されます。

pip install streamlit pillow torch torchvision

プログラムを作っていく

続いて、エディターでプログラムを作っていきます。

まずは、前回学習させたモデルを読み込ませていきます。黄色のマーカー部分は、皆さんが学習させたモデルのパス(モデルの保存されている場所)を入力しましょう。

その後、クラス名を入力していきます。ここでは、判定されたゴミを何ゴミとして表示するのかクラス名を記入していきます。Pytorchを使ってゴミの画像で機械学習をさせた際に、学習させた画像のファイルは「burnable」「plastic」「pet」「cardboard」の順にしていたので、これに対応するように名前をつけてあげましょう。
もし今後、住んでる地域によってゴミの判定を変えたい!というときは、ここを変えると良さそうですね。今後、これもチャレンジしてみましょう!

そして、前処理をするプログラムを入れ、webアプリの最初の部分ができました。

# モデルのロード
model = torch.load("<span class="marker-under">gomi_model.pth</span>", weights_only=False)
model.eval()

# クラス名
class_names = ["燃えるごみ", "プラスチック", "ペットボトル", "段ボール"]

# 画像前処理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

続いて、ブラウザで見える部分を作っていきます。

Streamlitは、先述の通り、簡単なコードでさまざまなUIを再現できます。例えば、タイトルであれば、st.title()で太字のタイトルに当たる文字を表示させてくれますし、st.write()で説明文のような細字を表示させることもできます。ゴミの画像をアップする部分も、下記のようにパパッと作れてしまいます!

# タイトル
st.title("ごみ分別AI")
st.write("画像をアップロードして、種類を分類してみましょう!")

# 画像アップロード
uploaded_file = st.file_uploader("画像をアップロードしてください", type=["jpg", "png", "jpeg"])

続いて、判定部分を作っていきます。

この部分では、まず、アップロードされた画像をRGB変換して表示させます。ゴミの判定結果が出てくるだけでもすごいですが、せっかくなら画像と一緒に表示させたいですよね!

その後、前処理をして推論(何ゴミか判定させる)を行います。その結果を、一番下のこれは **{label}** です!”で、表示させます。ちなみに{label}の部分が、冒頭に紹介したクラス名にリンクして表示されます。また、st.successは、Streamlitでの成功メッセージを表すもので、緑色の装飾をして表示してくれます。UIも、基本的な部分は抑えているのでいいですね!

if uploaded_file is not None:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="アップロード画像", use_container_width=True)

    # 前処理して推論
    input_tensor = transform(image).unsqueeze(0)  # バッチ次元を追加

    with torch.no_grad():
        output = model(input_tensor)
        _, predicted = torch.max(output, 1)
        label = class_names[predicted.item()]

    st.success(f"これは **{label}** です!")

実行してみよう

実行してみる

それでは書いたプログラムを実行させてみましょう。
今まではPythonのファイルを実行させるときは、Python3 example.py のような形で入力していたと思いますが、Streamlitのときは以下のように入力をして実行させます。毎回のことですが、黄色マーカーの部分はみなさんのファイルの名前にして、実行してみてください。

streamlit run <span class="marker-under">example.py</span>

すると、次のようなウェルカムメッセージが表示されました。これは、Streamlitを始めて起動させた時に出てくるメッセージで、使い方のガイドや最新情報をメールで送るかどうか聞いているメッセージです。特に必要ない人は、そのままEnterキーを押しましょう。

     👋 Welcome to Streamlit!

      If you’d like to receive helpful onboarding emails, news, offers, promotions,
      and the occasional swag, please enter your email address below. Otherwise,
      leave this field blank.

      Email:  

その後にでてくる以下のようなメッセージが、先ほどPythonで制作したサイトのURLです。このメッセージが出てくると、ブラウザが勝手に立ち上がります。

  You can now view your Streamlit app in your browser.

  Local URL: http://localhost:8501

開いてみると

開いてみると、下のような画像が出てきました。どうやらうまくいっているようです。いままで作ってきた裏側の部分が、こうして見えるようになると、ちょっとワクワクしますね!

さて、アップロードフォルダから画像を選んでアップしてみましょう。今回はレシートの写真を選びました。なので、可燃ごみになると成功です。早速読み込ませてみると、画像が読み込まれて、画像の下に「これは可燃ごみです!」と判別されました。素晴らしいですね!

まとめ

今回は前回学習させたモデルを使って、Streamlitというフレームワークを使用しwebアプリを作ってみました。画像をアップするとゴミを判別できるようになったので、今まで機械学習をさせてきたモデルが、かなり実用的な形に仕上がったのではないかと思います。

実際のスマホの使用シーンを考えた時に、画像をアップするのではなく、写真を取ったらそのまま判定してくれるようになったら、もっと便利になりそうですよね。また、記事の途中でも少し触れましたが、これに位置情報を対応できたら、一番最初の記事で課題として挙げた、地域によるゴミ分別の違いも対応できそうですね。

ご覧いただき、ありがとうございました。