はじめに

オートエンコーダーのSSIMの実装方法を整理します。

SSIMとは

SSIM（Structural Similarity Index Measure)とは、2004年に発表された画像の類似度を測定する指標です。
オートエンコーダーの課題の中で、オリジナル画像から作られる生成画像が鮮明ではない特徴があります。英語ではBlurred Generated Imageといいます。この場合、小さい異常の検知が難しい問題があります。
既存のオートエンコーダーにはmse(mean squared error)を損失関数として使い、どうしても鮮明ではない画像しか作られませんでした。ところで、SSIMを損失関数として使ったら、より鮮明な生成画像が得られ、結果的に異常検知のOK/NGの判別性能が向上する結果となりました。めでたしめでたし。

2.実装の手順

学習フェーズと推論フェーズに分けて説明します。

2.1.学習フェーズ

SSIMのカスタマイズ損失関数の定義

Tensorflowのtf.image.ssim関数を利用し、ssim_loss()のカスタム損失関数を定義します。
パラメータ設定には、float32とuint8によって異なりますので、それに関してはTensorflowのドキュメントを参照してください。

import tensorflow as tf

def ssim_loss(y_true, y_pred):
    return 1-tf.reduce_mean(tf.image.ssim(y_true, y_pred, 
                                          max_val = 1.0,filter_size=11,
                                        filter_sigma=1.5, k1=0.01, k2=0.03 ))

Compileにlossとして指定

autoencoderというmodelを用意します。
mseと使うcompile文と、ssim_lossを使うcompile文を併記しますので、参考にしてください。
SSIMは、このコードを使えば、問題なく学習ができます。

#損失関数　mse
autoencoder.compile(optimizer = 'adam', loss = 'mse')

#損失関数　SSIM
autoencoder.compile(optimizer = 'adam', loss = ssim_loss)

2.2.推論フェーズ

学習したモデル（h5ファイル）を読み込む時、「このモデルはSSIMというカスタム損失関数を使ってるよ」とプログラムに教える必要があります。
例えば、autoencoder-ssim.h5とうモデルを読み込む場合は、下記のようにコードを書きます。

load_modelの追記

from tensorflow.keras.models import load_model

loaded_model = load_model('autoencoder-ssim.h5', custom_objects={'ssim_loss':ssim_loss}

custom_objects部分を見てみましょう。

custom_objects={'ssim_loss':ssim_loss}

最初の'ssim_loss'は、autoencoder.compile(optimizer = 'adam', loss = ssim_loss)のloss='ssim_loss'のことです。
2番目のssim_lossはカスタム損失関数名になります。

SSIM関数の記述

上記の理由で、推論のコードにもカスタム損失関数を記述します。学習用のコードに書いた同じカスタム損失関数を書けばOKです。

def ssim_loss(y_true, y_pred):
    return 1-tf.reduce_mean(tf.image.ssim(y_true, y_pred, 
                                          max_val = 1.0,filter_size=11,
                                        filter_sigma=1.5, k1=0.01, k2=0.03 ))

これで問題なく、SSIMの学習モデルを読み込み、新しいデータでの推論が可能です。

まとめ

SSIMのカスタム関数を用いた、オートエンコーダーの学習・推論のコードの実装方法について解説しました。

参考資料

１．Improving Unsupervised Defect Segmentation by Applying Structural Similarity to Autoencoders
https://arxiv.org/abs/1807.02011

２．【論文読み】新・オートエンコーダで異常検知
https://qiita.com/shinmura0/items/ee074172ec3c818b614e