Keras 以 ResNet-50 預訓練模型建立狗與貓辨識程式

這裡示範在 Keras 架構下以 ResNet-50 預訓練模型為基礎，建立可用來辨識狗與貓的 AI 程式。

在 Keras 的部落格中示範了使用 VGG16 模型建立狗與貓的辨識程式，準確率大約為 94%，而這裡則是改用 ResNet50 模型為基礎，並將輸入影像尺寸提高為 224×224，加上大量的 data augmentation，結果可讓辨識的準確率達到 99%。

準備資料

從 Kaggle 的網站上下載 train.zip 壓縮檔，然後解壓縮：

# 解壓縮 train.zip
unzip train.zip

建立好目錄結構：

# 建立目錄結構
mkdir -p sample/train/cats
mkdir -p sample/train/dogs
mkdir -p sample/valid/cats
mkdir -p sample/valid/dogs

選擇貓與狗各前 1000 張照片作為訓練資料集，另各取 400 張作為測試資料集：

# 複製圖片
cd train
cp cat.?.jpg cat.??.jpg cat.???.jpg ../sample/train/cats/
cp dog.?.jpg dog.??.jpg dog.???.jpg ../sample/train/dogs/
cp cat.1[0-3]??.jpg ../sample/valid/cats/
cp dog.1[0-3]??.jpg ../sample/valid/dogs/

建立與訓練模型

以下是以 ResNet-50 預訓練模型為基礎，建立與訓練狗與貓辨識模型的程式碼：

from tensorflow.python.keras import backend as K
from tensorflow.python.keras.models import Model
from tensorflow.python.keras.layers import Flatten, Dense, Dropout
from tensorflow.python.keras.applications.resnet50 import ResNet50
from tensorflow.python.keras.optimizers import Adam
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator

# 資料路徑
DATASET_PATH  = 'sample'

# 影像大小
IMAGE_SIZE = (224, 224)

# 影像類別數
NUM_CLASSES = 2

# 若 GPU 記憶體不足，可調降 batch size 或凍結更多層網路
BATCH_SIZE = 8

# 凍結網路層數
FREEZE_LAYERS = 2

# Epoch 數
NUM_EPOCHS = 20

# 模型輸出儲存的檔案
WEIGHTS_FINAL = 'model-resnet50-final.h5'

# 透過 data augmentation 產生訓練與驗證用的影像資料
train_datagen = ImageDataGenerator(rotation_range=40,
                                   width_shift_range=0.2,
                                   height_shift_range=0.2,
                                   shear_range=0.2,
                                   zoom_range=0.2,
                                   channel_shift_range=10,
                                   horizontal_flip=True,
                                   fill_mode='nearest')
train_batches = train_datagen.flow_from_directory(DATASET_PATH + '/train',
                                                  target_size=IMAGE_SIZE,
                                                  interpolation='bicubic',
                                                  class_mode='categorical',
                                                  shuffle=True,
                                                  batch_size=BATCH_SIZE)

valid_datagen = ImageDataGenerator()
valid_batches = valid_datagen.flow_from_directory(DATASET_PATH + '/valid',
                                                  target_size=IMAGE_SIZE,
                                                  interpolation='bicubic',
                                                  class_mode='categorical',
                                                  shuffle=False,
                                                  batch_size=BATCH_SIZE)

# 輸出各類別的索引值
for cls, idx in train_batches.class_indices.items():
    print('Class #{} = {}'.format(idx, cls))

# 以訓練好的 ResNet50 為基礎來建立模型，
# 捨棄 ResNet50 頂層的 fully connected layers
net = ResNet50(include_top=False, weights='imagenet', input_tensor=None,
               input_shape=(IMAGE_SIZE[0],IMAGE_SIZE[1],3))
x = net.output
x = Flatten()(x)

# 增加 DropOut layer
x = Dropout(0.5)(x)

# 增加 Dense layer，以 softmax 產生個類別的機率值
output_layer = Dense(NUM_CLASSES, activation='softmax', name='softmax')(x)

# 設定凍結與要進行訓練的網路層
net_final = Model(inputs=net.input, outputs=output_layer)
for layer in net_final.layers[:FREEZE_LAYERS]:
    layer.trainable = False
for layer in net_final.layers[FREEZE_LAYERS:]:
    layer.trainable = True

# 使用 Adam optimizer，以較低的 learning rate 進行 fine-tuning
net_final.compile(optimizer=Adam(lr=1e-5),
                  loss='categorical_crossentropy', metrics=['accuracy'])

# 輸出整個網路結構
print(net_final.summary())

# 訓練模型
net_final.fit_generator(train_batches,
                        steps_per_epoch = train_batches.samples // BATCH_SIZE,
                        validation_data = valid_batches,
                        validation_steps = valid_batches.samples // BATCH_SIZE,
                        epochs = NUM_EPOCHS)

# 儲存訓練好的模型
net_final.save(WEIGHTS_FINAL)

將這段 Python 程式碼儲存為 train_resnet50.py，然後執行它即可進行模型的訓練：

python3 train_resnet50.py

模型訓練完成之後，就會儲存於 model-resnet50-final.h5 這個檔案中。

辨識狗與貓

將模型訓練好之後，就可以使用以下的 Python 指令稿打造自己的狗與貓辨識程式了：

from tensorflow.python.keras import backend as K
from tensorflow.python.keras.models import load_model
from tensorflow.python.keras.preprocessing import image
import sys
import numpy as np

# 從參數讀取圖檔路徑
files = sys.argv[1:]

# 載入訓練好的模型
net = load_model('model-resnet50-final.h5')

cls_list = ['cats', 'dogs']

# 辨識每一張圖
for f in files:
    img = image.load_img(f, target_size=(224, 224))
    if img is None:
        continue
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis = 0)
    pred = net.predict(x)[0]
    top_inds = pred.argsort()[::-1][:5]
    print(f)
    for i in top_inds:
        print('    {:.3f}  {}'.format(pred[i], cls_list[i]))

將這段 Python 程式碼儲存為 predict_resnet50.py，然後就可以拿來辨識圖片了。我們可以直接使用 Kaggle 所提供的圖片進行測試：

# 辨識狗與貓
python3 predict_resnet50.py train/dog.1000[0-2].jpg train/cat.1000[0-2].jpg

train/dog.10000.jpg
    1.000  dogs
    0.000  cats
train/dog.10001.jpg
    1.000  dogs
    0.000  cats
train/dog.10002.jpg
    1.000  dogs
    0.000  cats
train/cat.10000.jpg
    0.998  cats
    0.002  dogs
train/cat.10001.jpg
    1.000  cats
    0.000  dogs
train/cat.10002.jpg
    1.000  cats
    0.000  dogs

也可以拿自己的圖片放進去進行辨識：

# 辨識自己的圖片
python3 predict_resnet50.py pug-20180918-01.jpg

pug-20180918-01.jpg
    0.999  dogs
    0.001  cats

參考資料：JK Jung’s blog