はじめに
前回までは streamlit 上で ind_knn_add の使用方法について説明しました。
streamlit 上では学習に使用できる画像が限られているので、それぞれのメソッド ( SPADE, Padim, PatchCore ) の性能に関して、正確に検証できませんでした。
今回は、streamlit を使用せずカスタムデータセットで学習させる方法について説明します。
前提条件
前提条件は以下の通りです。
- python3.9
- streamlit == 0.86.0
- torch == 1.12.1+cu113
ind_knn_ad の github はこちらです。
データセットについて
前回と同様に、こちらのデータセットを使用します。Class1, Class1_def を使用します。
datasets フォルダ内に custom という名前のフォルダを作成し、その中に train, test フォルダを作成します。良品データを good 、不良品データを defect としてフォルダを作成します。
cd datasets
mkdir custom
cd custom
mkdir train
mkdir test
cd train
mkdir good
cd ../test
mkdir good
mkdir defectまた、検証用のデータとして ind_knn_ad フォルダ内に good.png と defect.png を用意します。
学習の準備
学習用のプログラムを作成します。
indad フォルダ内に demo.py を作成します。
from models import SPADE, PaDiM, PatchCore
from data import MVTecDataset
import cv2
import torch
import numpy as np
from torchvision import transforms
from torch import tensor
IMAGENET_MEAN = tensor([.485, .456, .406])
IMAGENET_STD = tensor([.229, .224, .225])
SIZE = 224
model = SPADE(k=11, backbone_name="wide_resnet50_2")
# model = PaDiM(d_reduced=150, backbone_name="wide_resnet50_2")
# model = PatchCore(f_coreset=.10, backbone_name="wide_resnet50_2")
train_ds, test_ds = MVTecDataset("custom", SIZE).get_dataloaders()
# feed healthy dataset
model.fit(train_ds)
transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize(SIZE, interpolation=transforms.InterpolationMode.BICUBIC),
transforms.CenterCrop(SIZE),
transforms.ToTensor(),
transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD)
])
# get predictions
good_frame = cv2.imread("./good.png")
good_frame = cv2.cvtColor(good_frame,cv2.COLOR_BGR2RGB)
good_x = transform(good_frame)
good_x = good_x.unsqueeze(0)
defect_frame = cv2.imread("./defect.png")
defect_frame = cv2.cvtColor(defect_frame,cv2.COLOR_BGR2RGB)
defect_x = transform(defect_frame)
defect_x = defect_x.unsqueeze(0)
print(good_x.shape, defect_x.shape)
model.eval()
with torch.no_grad():
img_lvl_anom_score, pxl_lvl_anom_score = model.predict(good_x)
print("good frame score is: ", img_lvl_anom_score)
img_lvl_anom_score, pxl_lvl_anom_score = model.predict(defect_x)
print("defect frame score is: ", img_lvl_anom_score)SPADE の結果
上記を実行すると、以下のような出力が得られます。
100%|██████████| 1000/1000 [03:18<00:00, 5.03it/s]
torch.Size([1, 3, 224, 224]) torch.Size([1, 3, 224, 224])
good frame score is: tensor(6.3686)
defect frame score is: tensor(6.5281)良品画像は 6.3686、不良品画像は 6.5281 となりました。
0.16 ポイントしか差がありません。
PaDiM の場合
PaDiM の場合は、以下のような出力になります。
100%|██████████| 1000/1000 [03:07<00:00, 5.35it/s]
PaDiM: (randomly) reducing 1792 dimensions to 150.
torch.Size([1, 3, 224, 224]) torch.Size([1, 3, 224, 224])
good frame score is: tensor(21.0718)
defect frame score is: tensor(46.7788)良品画像は 21.0718, 不良品画像は 46.7788 となりました。
PaDiM の方がノイズに敏感です。
その分、メモリの消費量は大きいです。
PatchCore の場合
PatchCore の場合は、以下のようになります。
100%|██████████| 1000/1000 [02:36<00:00, 6.37it/s]
Fitting random projections. Start dim = torch.Size([784000, 1536]).
DONE. Transformed dim = torch.Size([784000, 335]).
100%|██████████| 78399/78399 [14:50<00:00, 87.99it/s]
torch.Size([1, 3, 224, 224]) torch.Size([1, 3, 224, 224])
good frame score is: tensor(22.1969)
defect frame score is: tensor(25.6561)良品画像は 22.1969, 不良品画像は 25.6561 となりました。
3.5 ポイントと、明確な差が算出されました。
おわりに
今回は streamlit を使用せず ind_knn_ad を使用する方法について説明しました。
これによって、独自のシステムに組み込むことも可能になります。
次回はコードの説明と、3種のモデルのハイパーパラメータを調整したときの結果について説明しようと思います。
コメント