はじめに
この記事では、streamlitで作成した衣服の画像認識アプリの開発について説明します。
streamlitはデータ分析や学習モデルの動作確認などに使用され、学習コストが比較的低いことから簡易的なアプリ開発にされることが多いです。
そのため、より複雑な機能をもつアプリ開発に興味がある方は他の方の記事を参考にすることをおすすめします。
1. streamlitとは
Streamlitは、データサイエンティストやデベロッパーが簡単にウェブアプリケーションを構築するためのPythonライブラリおよびフレームワークです。Streamlitを使用することで、データ分析や機械学習モデルの結果を簡単にウェブベースのインタラクティブなアプリケーションに変換できます。
また Streamlit Cloud というサービスも提供されています。
Streamlit Cloudは、Streamlitアプリケーションをデプロイおよびホスティングするためのクラウドプラットフォームでです。
Githubと連携することで、コミットをすぐさまWebアプリに適用することが出来ます。
2. 開発環境
・モデル学習、評価
Docker
nvidia/cuda:12.2.2-devel-ubuntu22.04
torch==2.0.1+cu117
torchvision==0.15.2+cu117
numpy==1.26.0
GeForce RTX 3060ti
・実行環境
streamlit cloud
torch==2.0.1+cu117
torchvision==0.15.2+cu117
pillow=9.3.0
streamlit==1.28.1
3. 学習・評価
訓練データにはポピュラーなデータセットの一つであるFashion-Mnistを使用しました。また、学習モデルにはResNetの学習済みモデルをファインチューニングしました。
以下は実行したコードになります。
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.datasets import FashionMNIST
from torchvision.models import resnet50
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
import os
transforms = transforms.Compose([
transforms.CenterCrop(224), transforms.ToTensor()
])
root = "/root/"
train_dataset = FashionMNIST(
root=root, train=True, transform=transforms, download=True
)
val_dataset = FashionMNIST(
root=root, train=False, transform=transforms, download=True
)
class Resnet(nn.Module):
def __init__(self, num_classes):
super().__init__()
self._initModel(num_classes)
def _initModel(self, num_classes):
pre_model = resnet50(pretrained=True)
num_fc_in_features = pre_model.fc.in_features
pre_model.fc = nn.Linear(num_fc_in_features, num_classes)
pre_model.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
self.model = pre_model
def forward(self, images):
outputs = self.model(images)
return outputs
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Resnet(num_classes=10)
model = model.to(device)
class EarlyStopping:
def __init__(self, patience):
self.check_count = 0
self.patience = patience
def checkCount(self, _bool):
if _bool:
self.check_count = 0
else:
self.check_count += 1
if self.check_count == self.patience:
return 0
else:
return None
early_stopping = EarlyStopping(patience=3)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
epochs = 20
output_dir = "./output"
check = None
train_loss_list = []
train_correct_list =[]
val_loss_list = []
val_correct_list = []
for epoch in range(epochs):
for step in ["train", "val"]:
running_loss = 0.0
running_corrects = 0.0
min_loss = np.inf
if step == "train":
model.train()
dataloader = train_dataloader
else:
model.eval()
dataloader = val_dataloader
for batch, (images, labels) in enumerate(dataloader):
images = images.to(device)
labels = labels.to(device)
with torch.set_grad_enabled(step=="train"):
outputs = model(images)
pred = torch.argmax(outputs, dim=-1)
loss = (criterion(outputs, labels))
correct = (torch.sum(pred == labels)) / batch_size
if step == "train":
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
running_corrects += correct.item()
if step == "train":
print(f"STEP: Train Epoch: {epoch + 1}/{epochs} Iteration: {batch + 1}/{len(dataloader)} train_loss{running_loss / (batch + 1)}")
else:
print(f"STEP: VAL Epoch: {epoch + 1}/{epochs} Iteration: {batch + 1}/{len(dataloader)} val_loss{running_loss / (batch + 1)} val_correct: {running_corrects / (batch + 1)}")
if step == "train":
train_loss_list.append(running_loss / (batch + 1))
train_correct_list.append(running_corrects / (batch + 1))
else:
val_loss_list.append(running_loss / (batch + 1))
val_correct_list.append(running_corrects / (batch + 1))
if running_loss < min_loss:
early_stopping.checkCount(True)
print("Model Save !")
min_loss = running_loss
if not os.path.exists(output_dir):
os.mkdir(output_dir)
torch.save(model.state_dict(), os.path.join(output_dir, "model.pth"))
else:
check = early_stopping.checkCount(False)
if check == 0:
print("Early Stopping !")
break
学習結果は以下のようなグラフになりました。
評価データによる正解率（val_correct）は 85％ 程度に落ち着きました。初期は ResNet152を使用していたのですが、Githubでは100MBを超えるファイルは扱えないとのことで、パラメータが軽いResNet50を使用しました。
4. 実装
実装には streamlit cloud を使用しました。Githubと連携することでクラウド上にWebアプリを実装することが出来ます。
以下のフォームに実装したアプリのコードが記載されているリポジトリを入力することで、クラウド上でアプリを実行することが出来ます。
実行に必要なファイルは以下になります。
main_app.py: streamlitが記載されたコード
model.py: モデルの構成が記載されたコード
model.pth: モデルのパラメータが記載されたファイル
requirements.txt: 実行に必要なライブラリが記載されたファイル
それぞれのファイルのコードを以下に記載します。
1. main_app.py:
import streamlit as st
import torch
import torchvision
from model import Resnet
from PIL import Image
def predict(image, labels, model):
transform = torchvision.transforms.Compose([
torchvision.transforms.CenterCrop(224),
torchvision.transforms.Grayscale(num_output_channels=1),
torchvision.transforms.ToTensor()
])
image = transform(image)
image = image.unsqueeze(0)
model.eval()
outputs = model(image)
y_prob = torch.nn.functional.softmax(outputs.squeeze(0), dim=-1)
sorted_prob, sorted_indices = torch.sort(y_prob, descending=True)
results = []
for prob, idx in zip(sorted_prob, sorted_indices):
results.append((labels[idx.item()], prob.item()))
return results
def main():
model = Resnet(num_classes=10)
model.load_state_dict(torch.load("model.pth", map_location='cpu'))
with open("labels.txt", "r") as f:
labels = [line.strip() for line in f.readlines()]
st.sidebar.title("衣服の画像認識アプリ")
st.sidebar.write("画像認識モデルを使って衣服の種類を判定します。")
st.sidebar.write("判別が可能な種類は以下の通りです。")
for label in labels:
st.sidebar.write(f"{label}")
st.sidebar.write("")
img_source = st.sidebar.radio("画像のソースを選択してください",
("画像をアップロード", "画像を撮影"))
if img_source == "画像をアップロード":
img_file = st.sidebar.file_uploader("画像を選択してください。", type=["png", "jpg", "jpeg"])
else:
img_file = st.camera_input("カメラで撮影")
if img_file is not None:
with st.spinner("計算中・・・"):
img = Image.open(img_file)
st.image(img, caption="対象画像")
st.write("")
results = predict(img, labels, model)
st.subheader("判定結果")
num_top = 5
for result in results[:num_top]:
st.write(str(round(result[1] * 100, 2)) + "%の確率で" + result[0] + "です。")
st.sidebar.write("")
st.sidebar.write("")
st.sidebar.caption('"このアプリは[FashionMnist]を訓練データとして扱っています \n \
Copyright (c) 2017 Zalando SE \n \
Released under the MIT license \n \
https://github.com/zalandoresearch/fashion-mnist#license"')
if __name__ == "__main__":
main()
2. model.py
import torch.nn as nn
from torchvision.models import resnet50
class Resnet(nn.Module):
def __init__(self, num_classes):
super().__init__()
self._initModel(num_classes)
def _initModel(self, num_classes):
pre_model = resnet50(weights=None)
num_fc_in_features = pre_model.fc.in_features
pre_model.fc = nn.Linear(num_fc_in_features, num_classes)
pre_model.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
self.model = pre_model
def forward(self, images):
outputs = self.model(images)
return outputs
3. requirements.txt
torch==2.0.1+cu117
torchvision==0.15.2+cu117
pillow=9.3.0
streamlit==1.28.1
5. 実行
しっかりと分類結果を表示できていますね。
まとめ
ここまで、読んでくださいありがとうございます。
今回はstreamlitを使用した衣服の画像認識アプリの作成について説明しました。使用データセットはグレースケールであること、一方向のみの画像であることから実際の写真を分類すると判別精度が著しく下がる可能性があります。
今回は基本的な前処理しか行いませんでしたが、より学習モデルにロバスト性を持たせたいなら左右反転や回転などの前処理を行うことをおすすめします。
Streamlitは実装が簡単なことから、学生や機械学習の初学者にはおすすめのツールといえます。もし理論は理解しているが実装の経験が少ないという方は、この機会に実装してはいかがでしょうか？
次回は、テキストを使用したWebアプリを実装したいと考えています。