#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Copyright 2018 The TensorFlow Authors.
End of explanation
"""


import numpy as np
import time

import PIL.Image as Image
import matplotlib.pylab as plt

import tensorflow as tf
import tensorflow_hub as hub

import datetime

%load_ext tensorboard

"""
Explanation: TensorFlow Hub による転移学習
<table class="tfo-notebook-buttons" align="left">
  <td>     <a target="_blank" href="https://www.tensorflow.org/tutorials/images/transfer_learning_with_hub"><img src="https://www.tensorflow.org/images/tf_logo_32px.png">View on TensorFlow.org</a>   </td>
  <td>     <a target="_blank" href="https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/tutorials/images/transfer_learning_with_hub.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png">Run in Google Colab</a>   </td>
  <td><a target="_blank" href="https://github.com/tensorflow/docs-l10n/blob/master/site/ja/tutorials/images/transfer_learning_with_hub.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png">GitHub で表示</a></td>
  <td>     <a href="https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ja/tutorials/images/transfer_learning_with_hub.ipynb" class="_active_edit_href"><img src="https://www.tensorflow.org/images/download_logo_32px.png">Download notebook</a>   </td>
  <td><a href="https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4" class="_active_edit_href"><img src="https://www.tensorflow.org/images/hub_logo_32px.png">TF Hub モデルを参照</a></td>
</table>

TensorFlow Hub は、トレーニング済みの TensorFlow モデルのリポジトリです。
このチュートリアルでは、以下の方法を実演します。

TensorFlow Hub からのモデルを tf.keras で利用する。
TensorFlow Hub からの画像分類モデルを使用する。
独自の画像クラスのモデルを微調整するためにシンプルな転移学習を行う。

セットアップ
End of explanation
"""


mobilenet_v2 ="https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification/4"
inception_v3 = "https://tfhub.dev/google/imagenet/inception_v3/classification/5"

classifier_model = mobilenet_v2 #@param ["mobilenet_v2", "inception_v3"] {type:"raw"}

IMAGE_SHAPE = (224, 224)

classifier = tf.keras.Sequential([
    hub.KerasLayer(classifier_model, input_shape=IMAGE_SHAPE+(3,))
])

"""
Explanation: ImageNet の分類器
ImageNet ベンチマークデータセットで事前トレーニングされた分類器モデルを使用するため、初期トレーニングは不要です！
分類器のダウンロード
TensorFlow Hub から事前トレーニング済みの <a href="https://arxiv.org/abs/1801.04381" class="external">MobileNetV2</a> モデルを選択し、Keras レイヤーとして hub.KerasLayer でラップします。ここでは、TensorFlow Hub からであれば、以下のドロップダウンに提供されている Example も含み、<a href="https://tfhub.dev/s?q=tf2&amp;module-type=image-classification/" class="external">互換性のあるどの画像分類器モデル</a>でも構いません。
End of explanation
"""


grace_hopper = tf.keras.utils.get_file('image.jpg','https://storage.googleapis.com/download.tensorflow.org/example_images/grace_hopper.jpg')
grace_hopper = Image.open(grace_hopper).resize(IMAGE_SHAPE)
grace_hopper

grace_hopper = np.array(grace_hopper)/255.0
grace_hopper.shape

"""
Explanation: 1 枚の画像で実行する
モデルを試すために、画像を 1 枚ダウンロードします。
End of explanation
"""


result = classifier.predict(grace_hopper[np.newaxis, ...])
result.shape

"""
Explanation: バッチの次元を追加し、画像をモデルに入力します。
End of explanation
"""


predicted_class = tf.math.argmax(result[0], axis=-1)
predicted_class

"""
Explanation: 結果は、1001 要素のベクトルのロジットで、画像の各クラスの確率を評価します。
そのため、トップのクラス ID は argmax を使うことでみつけることができます:
End of explanation
"""


labels_path = tf.keras.utils.get_file('ImageNetLabels.txt','https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
imagenet_labels = np.array(open(labels_path).read().splitlines())

plt.imshow(grace_hopper)
plt.axis('off')
predicted_class_name = imagenet_labels[predicted_class]
_ = plt.title("Prediction: " + predicted_class_name.title())

"""
Explanation: 推論結果のデコード
predicted_class ID（653 など）を取り、ImageNet データセットラベルをフェッチして予測をデコードします。
End of explanation
"""


data_root = tf.keras.utils.get_file(
  'flower_photos',
  'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',
   untar=True)

"""
Explanation: シンプルな転移学習
ただし、元の ImageNet データセット（事前トレーニング済みモデルがトレーニングされたデータセット）に含まれないクラスを持つ独自のデータセットを使用してカスタム分類器を作成する場合はどうでしょうか。
これは、以下のようにして行います。

TensorFlow Hub から事前トレーニング済みモデルを選択します。
カスタムデータセットのクラスを認識できるよう、最上位（最後）のレイヤーを保持します。

データセット
この例では、TensorFlow flowers データセットを使用します。
End of explanation
"""


batch_size = 32
img_height = 224
img_width = 224

train_ds = tf.keras.utils.image_dataset_from_directory(
  str(data_root),
  validation_split=0.2,
  subset="training",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size
)

val_ds = tf.keras.utils.image_dataset_from_directory(
  str(data_root),
  validation_split=0.2,
  subset="validation",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size
)

"""
Explanation: まず、tf.keras.utils.image_dataset_from_directory を使用して、このデータをディスクの画像データを使ったモデルに読み込みます。これにより、tf.data.Dataset が生成されます。
End of explanation
"""


class_names = np.array(train_ds.class_names)
print(class_names)

"""
Explanation: flowers データセットには 5 つのクラスがあります。
End of explanation
"""


normalization_layer = tf.keras.layers.Rescaling(1./255)
train_ds = train_ds.map(lambda x, y: (normalization_layer(x), y)) # Where x—images, y—labels.
val_ds = val_ds.map(lambda x, y: (normalization_layer(x), y)) # Where x—images, y—labels.

"""
Explanation: 次に、画像モデルに使用される TensorFlow Hub の規則では[0, 1] 範囲の浮動小数点数の入力が期待されるため、tf.keras.layers.Rescaling 前処理レイヤーを使用してこれを達成します。
注意: モデルには、tf.keras.layers.Rescaling レイヤーも含めることができます。トレードオフに関する議論について、前処理レイヤーの操作ガイドをご覧ください。
End of explanation
"""


AUTOTUNE = tf.data.AUTOTUNE
train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

for image_batch, labels_batch in train_ds:
  print(image_batch.shape)
  print(labels_batch.shape)
  break

"""
Explanation: 3 番目に、Dataset.prefetch を使って、バッファリングされたプリフェッチで入力パイプラインを終了します。これで、I/O ブロッキングの問題が生じずにディスクからデータを生成することができます。
これらが、データを読み込む際に使用することが推奨される、いくつかの最も重要な tf.data メソッドです。さらに関心がある場合は、これらのメソッド、ディスクへのデータのキャッシュ方法、およびその他の手法について、tf.data API によるパフォーマンスの改善ガイドをご覧ください。
End of explanation
"""


result_batch = classifier.predict(train_ds)

predicted_class_names = imagenet_labels[tf.math.argmax(result_batch, axis=-1)]
predicted_class_names

"""
Explanation: 分類器で画像をバッチ処理する
分類器で画像をバッチ処理していきます。
End of explanation
"""


plt.figure(figsize=(10,9))
plt.subplots_adjust(hspace=0.5)
for n in range(30):
  plt.subplot(6,5,n+1)
  plt.imshow(image_batch[n])
  plt.title(predicted_class_names[n])
  plt.axis('off')
_ = plt.suptitle("ImageNet predictions")

"""
Explanation: これらの予測が画像とどれくらい整合しているかを確認します。
End of explanation
"""


mobilenet_v2 = "https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4"
inception_v3 = "https://tfhub.dev/google/tf2-preview/inception_v3/feature_vector/4"

feature_extractor_model = mobilenet_v2 #@param ["mobilenet_v2", "inception_v3"] {type:"raw"}

"""
Explanation: 注意: すべての画像は CC-BY のライセンス下にあります。作成者のリストは LICENSE.txt ファイルをご覧ください。
結果は完全とは決して言えませんが、これらはモデルがトレーニングされたクラスではないこと（「daisy」を除く）考慮すれば、合理的です。
ヘッドレスモデルのダウンロード
TensorFlow Hub は最上位の分類層を含まないモデルも配布しています。これらは転移学習に簡単に利用することができます。
<a href="https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4" class="external">TensorFlow Hub から</a>事前トレーニング済みの <a href="https://arxiv.org/abs/1801.04381" class="external">MobileNetV2</a> モデルを選択します。ここでは、TensorFlow Hub からであれば、以下のドロップダウンに提供されている Example も含み、<a href="https://tfhub.dev/s?module-type=image-feature-vector&amp;q=tf2" class="external">互換性のあるどの画像分類器モデル</a>でも構いません。
End of explanation
"""


feature_extractor_layer = hub.KerasLayer(
    feature_extractor_model,
    input_shape=(224, 224, 3),
    trainable=False)

"""
Explanation: hub.KerasLayer を使用して、事前トレーニング済みモデルを Keras レイヤーとしてラップし、特徴量抽出器を作成します。trainable=False 引数を使用して変数を凍結し、トレーニングのみが新しい分類器レイヤーを変更できるようにします。
End of explanation
"""


feature_batch = feature_extractor_layer(image_batch)
print(feature_batch.shape)

"""
Explanation: 特徴量抽出器は、画像ごとに 1280 長のベクトルを返します（この例では、画像バッチサイズは 32 のママになります）。
End of explanation
"""


num_classes = len(class_names)

model = tf.keras.Sequential([
  feature_extractor_layer,
  tf.keras.layers.Dense(num_classes)
])

model.summary()

predictions = model(image_batch)

predictions.shape

"""
Explanation: 上位の分類レイヤーを接合する
モデルを完成するために、特徴量抽出器レイヤーをtf.keras.Sequential モデルにラップし、分類用に全結合レイヤーを追加します。
End of explanation
"""


model.compile(
  optimizer=tf.keras.optimizers.Adam(),
  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
  metrics=['acc'])

log_dir = "logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(
    log_dir=log_dir,
    histogram_freq=1) # Enable histogram computation for every epoch.

"""
Explanation: モデルのトレーニング
Model.compile を使用してトレーニングプロセスを構成し、tf.keras.callbacks.TensorBoard コールバックを追加してログの作成と保存を行います。
End of explanation
"""


NUM_EPOCHS = 10

history = model.fit(train_ds,
                    validation_data=val_ds,
                    epochs=NUM_EPOCHS,
                    callbacks=tensorboard_callback)

"""
Explanation: 次に、Model.fit メソッドを使用して、モデルをトレーニングします。
この例を短くするために、10 エポックだけトレーニングします。後で TensorBoard にトレーニングプロセスを可視化できるよう、TensorBoard コールバックでログを作成して保存します。
End of explanation
"""


%tensorboard --logdir logs/fit

"""
Explanation: エポックごとに指標がどのように変化しているかを表示し、他のスカラー値を追跡するために、TensorBoard を起動します。
End of explanation
"""


predicted_batch = model.predict(image_batch)
predicted_id = tf.math.argmax(predicted_batch, axis=-1)
predicted_label_batch = class_names[predicted_id]
print(predicted_label_batch)

"""
Explanation: <!-- <img class="tfo-display-only-on-site" src="https://github.com/tensorflow/docs/blob/master/site/en/tutorials/images/images/tensorboard_transfer_learning_with_hub.png?raw=1"/> -->

推論結果の確認
モデルの予測からクラス名の番号付きリストを取得します。
End of explanation
"""


plt.figure(figsize=(10,9))
plt.subplots_adjust(hspace=0.5)

for n in range(30):
  plt.subplot(6,5,n+1)
  plt.imshow(image_batch[n])
  plt.title(predicted_label_batch[n].title())
  plt.axis('off')
_ = plt.suptitle("Model predictions")

"""
Explanation: 結果をプロットします
End of explanation
"""


t = time.time()

export_path = "/tmp/saved_models/{}".format(int(t))
model.save(export_path)

export_path

"""
Explanation: モデルのエクスポート
モデルのトレーニングが完了したので、後で再利用するために、SavedModel としてエクスポートします。
End of explanation
"""


reloaded = tf.keras.models.load_model(export_path)

result_batch = model.predict(image_batch)
reloaded_result_batch = reloaded.predict(image_batch)

abs(reloaded_result_batch - result_batch).max()

reloaded_predicted_id = tf.math.argmax(reloaded_result_batch, axis=-1)
reloaded_predicted_label_batch = class_names[reloaded_predicted_id]
print(reloaded_predicted_label_batch)

plt.figure(figsize=(10,9))
plt.subplots_adjust(hspace=0.5)
for n in range(30):
  plt.subplot(6,5,n+1)
  plt.imshow(image_batch[n])
  plt.title(reloaded_predicted_label_batch[n].title())
  plt.axis('off')
_ = plt.suptitle("Model predictions")

"""
Explanation: SavedModel を再読み込みできることと、モデルが同じ結果を出力できることを確認します。
End of explanation
"""


