Tensorflow2.0之圖像說明文字生成

項目介紹

在此項目中，我們希望能在輸入一張圖像之後得到一句話來描述該圖像，比如輸入下面這張圖像後會輸出 “a man in an all excited to a podium” 之類的說明性文字。

注意：這裏我們使用 MS-COCO 數據集來練習，這個數據集比較大，大概有13G左右。

代碼實現

1、導入需要的庫

import tensorflow as tf

# You'll generate plots of attention in order to see which parts of an image
# our model focuses on during captioning
import matplotlib.pyplot as plt

# Scikit-learn includes many helpful utilities
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle

import re
import numpy as np
import os
import time
import json
from glob import glob
from PIL import Image
import pickle

2、下載數據集

# Download image files
image_folder = '/train2014/'
if not os.path.exists(os.path.abspath('.') + image_folder):
    image_zip = tf.keras.utils.get_file('train2014.zip',
                                      cache_subdir=os.path.abspath('.'),
                                      origin = 'http://images.cocodataset.org/zips/train2014.zip',
                                      extract = True)
    PATH = os.path.dirname(image_zip) + image_folder
    os.remove(image_zip)
else:
    PATH = os.path.abspath('.') + image_folder
# Download caption annotation files
annotation_folder = '/annotations/'
if not os.path.exists(os.path.abspath('.') + annotation_folder):
    annotation_zip = tf.keras.utils.get_file('captions.zip',
                                          cache_subdir=os.path.abspath('.'),
                                          origin = 'http://images.cocodataset.org/annotations/annotations_trainval2014.zip',
                                          extract = True)
    annotation_file = os.path.dirname(annotation_zip)+'/annotations/captions_train2014.json'
    os.remove(annotation_zip)

其中 image_zip 是下載的圖片壓縮包；annotation_zip 是下載的說明性文字壓縮包，其中 captions_train2014.json 文件中包含訓練集的說明性文字以及對應的圖片名稱。

3、讀取 json 文件

with open(annotation_file, 'r') as f:
    annotations = json.load(f)
annotations.keys()

dict_keys(['info', 'images', 'licenses', 'annotations'])

通過打印 annotations 類型可知它是一個字典，它的鍵值包括 ‘info’，‘images’，‘licenses’ 和 ‘annotations’，其中我們只需要用到 ‘annotations’。

annotations[‘annotations’] 返回一個列表，這個列表中的每個元素都是含有三個鍵值的字典，這三個鍵值分別是 ‘image_id’，‘id’ 和 ‘caption’，其中我們需要 ‘caption’ 來返回說明性文字和 ‘image_id’ 來返回這個說明性文字對應的圖片代號。

爲節約訓練時間，我們只取其中30000個樣本訓練。

all_captions = []
all_img_name_vector = []

for annot in annotations['annotations']:
    caption = '<start> ' + annot['caption'] + ' <end>'
    image_id = annot['image_id']
    full_coco_image_path = PATH + 'COCO_train2014_' + '%012d.jpg' % (image_id)

    all_img_name_vector.append(full_coco_image_path)
    all_captions.append(caption)

train_captions, img_name_vector = shuffle(all_captions,
                                          all_img_name_vector,
                                          random_state=1)

num_examples = 30000
train_captions = train_captions[:num_examples]
img_name_vector = img_name_vector[:num_examples]

關於 sklearn.utils.shuffle() 的用法請參考：用 sklearn.utils.shuffle 來打亂樣本順序。

4、載入圖片

def load_image(image_path):
    img = tf.io.read_file(image_path)
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, (299, 299))
    img = tf.keras.applications.inception_v3.preprocess_input(img)
    return img, image_path

其中 tf.io.read_file(image_path) 輸出的是圖片信息，我們在將其進行解碼（ tf.image.decode_jpeg(img, channels=3) ）後可以得到圖片像素。

因爲我們用 Inception_v3 網絡來提取圖片特徵，所以我們在這裏要將圖片像素範圍轉化爲 Inception_v3 網絡需要的範圍，所以要使用 tf.keras.applications.inception_v3.preprocess_input(img)。

5、載入模型

在這裏我們使用 Inception_v3 網絡來提取圖片特徵，輸出爲該網絡中最後一個卷積層的輸出。

image_model = tf.keras.applications.InceptionV3(include_top=False,
                                                weights='imagenet')
new_input = image_model.input
hidden_layer = image_model.layers[-1].output

image_features_extract_model = tf.keras.Model(new_input, hidden_layer)

6、獲取圖片特徵

6.1 刪除重複的圖片

因爲一張圖片可能對應不同的說明性文字，所以圖片數據集中存在重複的問題。

encode_train = sorted(set(img_name_vector))

此時，encode_train 列表中的圖片名稱是不重複的。

6.2 切片、分批

首先將列表 encode_train 轉化爲 dataset 類型的數據。

image_dataset = tf.data.Dataset.from_tensor_slices(encode_train)

然後我們把它映射到 load_image() 函數實現從圖片名稱到圖片的轉換。

image_dataset = image_dataset.map(load_image, num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(16)

6.3 將圖片輸入網絡以獲取特徵

for img, path in image_dataset:
    batch_features = image_features_extract_model(img)
    batch_features = tf.reshape(batch_features,
                                (batch_features.shape[0], -1, batch_features.shape[3]))

    for bf, p in zip(batch_features, path):
        path_of_feature = p.numpy().decode("utf-8")
        np.save(path_of_feature, bf.numpy())

這裏的 np.save(path_of_feature, bf.numpy()) 的目的是：在調用 np.load(path_of_feature+’.npy’) 的時候能輸出這個路徑下的圖片對應的特徵 bf.numpy()。

最終得到的圖片特徵的維度爲：（batch_size, 64, 2048）。

7、文本 $\rightarrow$ 數字向量

7.1 構建分詞器

爲了節省內存，我們把詞彙表大小限制在前5000個單詞，其他的單詞用 “<unk>” 代替。

top_k = 5000
tokenizer = tf.keras.preprocessing.text.Tokenizer(num_words=top_k,
                                                  oov_token="<unk>",
                                                  filters='!"#$%&()*+.,-/:;=?@[\]^_`{|}~ ')
tokenizer.fit_on_texts(train_captions)

7.2 構建數字向量

train_seqs = tokenizer.texts_to_sequences(train_captions)

7.3 將數字向量填充到同一長度

cap_vector = tf.keras.preprocessing.sequence.pad_sequences(train_seqs, padding='post')

8、劃分訓練集和驗證集

img_name_train, img_name_val, cap_train, cap_val = train_test_split(img_name_vector,
                                                                    cap_vector,
                                                                    test_size=0.2,
                                                                    random_state=0)

9、建立 tf.data 數據集

在 6.2 部分中建立的數據集是爲了將圖片輸入 Inception_v3 網絡得到特徵，而在這裏建立的數據集對應的樣本是圖片的（解碼前的）名稱，標籤是這張圖片對應的說明性文字的數字向量。

dataset = tf.data.Dataset.from_tensor_slices((img_name_train, cap_train))

接着，我們要將數據集中的圖片名稱轉換爲這張圖片對應的特徵，所以我們要使用 np.load() 函數。

def map_func(img_name, cap):
    img_tensor = np.load(img_name.decode('utf-8')+'.npy')
    return img_tensor, cap

dataset = dataset.map(lambda item1, item2: tf.numpy_function(
          map_func, [item1, item2], [tf.float32, tf.int32]),
          num_parallel_calls=tf.data.experimental.AUTOTUNE)

接着，對數據集進行打亂、分批：

dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

10、編碼器

因爲在之前我們已經用卷積神經網絡提取了特徵（batch_size, 64, 2048），所以在這個編碼器中我們只需要定義全連接層（其神經元個數是詞嵌入維度）即可。

class CNN_Encoder(tf.keras.Model):
    def __init__(self, embedding_dim):
        super(CNN_Encoder, self).__init__()
        self.fc = tf.keras.layers.Dense(embedding_dim)

    def call(self, x):
        x = self.fc(x)
        x = tf.nn.relu(x)
        return x

經過編碼器後，圖像特徵的形狀變爲（batch_size, 64, embedding_dim）。

11、Bahdanau 注意力

相關論文參考：BahdanauAttention。

class BahdanauAttention(tf.keras.Model):
    def __init__(self, units):
        super(BahdanauAttention, self).__init__()
        self.W1 = tf.keras.layers.Dense(units)
        self.W2 = tf.keras.layers.Dense(units)
        self.V = tf.keras.layers.Dense(1)

    def call(self, features, hidden):
        hidden_with_time_axis = tf.expand_dims(hidden, 1)

        score = tf.nn.tanh(self.W1(features) + self.W2(hidden_with_time_axis))

        attention_weights = tf.nn.softmax(self.V(score), axis=1)

        context_vector = attention_weights * features
        context_vector = tf.reduce_sum(context_vector, axis=1)

        return context_vector, attention_weights

這裏的 features 其實就是編碼器中輸出的結果，經過含 units 個神經元的 Dense 層之後，其形狀從 (batch_size, 64, embedding_dim) 變成了 (batch_size, 64, units)。

這裏的 hidden 其實就是解碼器中輸出的隱層向量，我們需要將其維度從 (batch_size, embedding_dim) 變成 (batch_size, 1, embedding_dim) 來執行之後的加法以計算分數，將增加維度後的向量經過含 units 個神經元的 Dense 層之後，其形狀從 (batch_size, 1, embedding_dim) 變成了 (batch_size, 1, units)。

將以上兩個輸出相加得到的形狀爲 (batch_size, 64, units)，經過含1個神經元的 Dense 層之後得到 score，其形狀變成 (batch_size, 64, 1)。

Softmax 默認被應用於最後一個軸，但是這裏我們想將它應用於第二個軸（即 axis=1），因爲分數 （score） 的形狀是 (batch_size, 64, 1)。我們想爲每個輸入的特徵 (batch_size, 64, embedding_dim) 分配一個權重，所以 softmax 應該用在 64 這個軸上。經過 Softmax 層之後，得到的注意力權重形狀和 score 的形狀相同，都是 (batch_size, 64, 1)。
【注】Softmax 的不同的軸的計算規則參考：tf.nn.softmax(x, axis)裏axis起什麼作用？

將注意力權重和 features 相乘，得到上下文向量，其形狀爲 (batch_size, 16, embedding_dim)。此向量也就是加了權重的編碼向量。將上下文向量基於第二個軸求和（原因與之前相同），得到最終的上下文向量，其形狀爲 (batch_size, embedding_dim)。

12、解碼器

class RNN_Decoder(tf.keras.Model):
    def __init__(self, embedding_dim, units, vocab_size):
        super(RNN_Decoder, self).__init__()
        self.units = units

        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
        self.gru = tf.keras.layers.GRU(self.units,
                                       return_sequences=True,
                                       return_state=True,
                                       recurrent_initializer='glorot_uniform')
        self.fc1 = tf.keras.layers.Dense(self.units)
        self.fc2 = tf.keras.layers.Dense(vocab_size)
        self.attention = BahdanauAttention(self.units)

    def call(self, x, features, hidden):
        context_vector, attention_weights = self.attention(features, hidden)
        x = self.embedding(x)
        x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)
        output, state = self.gru(x)
        x = self.fc1(output)
        x = tf.reshape(x, (-1, x.shape[2]))
        x = self.fc2(x)
        return x, state, attention_weights

    def reset_state(self, batch_size):
        return tf.zeros((batch_size, self.units))

在解碼器中，我們首先通過 Bahdanau 注意力得到上下文向量和注意力權重，他們的形狀分別爲 (batch_size, embedding_dim) 和 (batch_size, 64, 1)。

這裏的輸入 x 是一個詞對應的數字（第一個輸入模型的數字一定是 “<start>” 對應的數字），其形狀爲（batch_size, 1），經過詞嵌入層之後，其形狀變爲（batch_size, embedding_dim），將其與增加了第二維度的上下文向量合併後得到形狀爲 (batch_size, 64, 2*embedding_dim)。

將其輸入 GRU，得到輸出爲 (batch_size, 1, units)，隱藏狀態爲 (batch_size, units)。

然後，經過全連接層後，得到 (batch_size, vocab_size)。

13、設置超參數建立模型

BATCH_SIZE = 64
BUFFER_SIZE = 1000
embedding_dim = 256
units = 512
vocab_size = top_k
num_steps = len(img_name_train) // BATCH_SIZE
features_shape = 2048
attention_features_shape = 64

encoder = CNN_Encoder(embedding_dim)
decoder = RNN_Decoder(embedding_dim, units, vocab_size)

14、初始化優化器

optimizer = tf.keras.optimizers.Adam()

15、損失函數

當輸入的向量中出現0元素，說明這個元素所在的文本已經結束了，這個文本不再參與損失的計算，所以在計算損失的時候，要使用掩膜處理，將已結束文本的損失置零。

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True, reduction='none')

def loss_function(real, pred):
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    loss_ = loss_object(real, pred)

    mask = tf.cast(mask, dtype=loss_.dtype)
    loss_ *= mask

    return tf.reduce_mean(loss_)

16、配置檢查點

checkpoint_path = "./checkpoints/train"
ckpt = tf.train.Checkpoint(encoder=encoder,
                           decoder=decoder,
                           optimizer = optimizer)
ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)

# 如果檢查點存在，則恢復最新的檢查點。
start_epoch = 0
if ckpt_manager.latest_checkpoint:
    start_epoch = int(ckpt_manager.latest_checkpoint.split('-')[-1])
    # restoring the latest checkpoint in checkpoint_path
    ckpt.restore(ckpt_manager.latest_checkpoint)

17、梯度下降

@tf.function
def train_step(img_tensor, target):
    loss = 0

    # initializing the hidden state for each batch
    # because the captions are not related from image to image
    hidden = decoder.reset_state(batch_size=target.shape[0])

    dec_input = tf.expand_dims([tokenizer.word_index['<start>']] * target.shape[0], 1)

    with tf.GradientTape() as tape:
        features = encoder(img_tensor)

        for i in range(1, target.shape[1]):
            # passing the features through the decoder
            predictions, hidden, _ = decoder(dec_input, features, hidden)

            loss += loss_function(target[:, i], predictions)

            # using teacher forcing
            dec_input = tf.expand_dims(target[:, i], 1)

    total_loss = (loss / int(target.shape[1]))

    trainable_variables = encoder.trainable_variables + decoder.trainable_variables

    gradients = tape.gradient(loss, trainable_variables)

    optimizer.apply_gradients(zip(gradients, trainable_variables))

    return loss, total_loss

18、訓練

loss_plot = []

EPOCHS = 20

# 訓練從 start_epoch 訓練到 EPOCHS
for epoch in range(start_epoch, EPOCHS):
    start = time.time()
    total_loss = 0

    for (batch, (img_tensor, target)) in enumerate(dataset):
        batch_loss, t_loss = train_step(img_tensor, target)
        total_loss += t_loss

        if batch % 100 == 0:
            print ('Epoch {} Batch {} Loss {:.4f}'.format(
              epoch + 1, batch, batch_loss.numpy() / int(target.shape[1])))
    # storing the epoch end loss value to plot later
    loss_plot.append(total_loss / num_steps)

    if epoch % 5 == 0:
        ckpt_manager.save()

    print ('Epoch {} Loss {:.6f}'.format(epoch + 1,
                                         total_loss/num_steps))
    print ('Time taken for 1 epoch {} sec\n'.format(time.time() - start))

19、驗證

19.1 驗證函數

• 1、初始化解碼器的隱藏狀態；
• 2、爲圖像添加維度；
• 3、提取圖像特徵並將所得形狀轉換成編碼器需要的形狀；
• 4、將圖像特徵輸入編碼器；
• 5、初始化解碼器輸入爲 ‘<start>’；
• 6、經解碼器得到的數據形狀爲 (batch_size, vocab_size)，要用 tf.random.categorical() 函數找出每個樣本的概率最大的下一個單詞。
• 7、將預測的單詞放到一個列表中。

def evaluate(image):
    attention_plot = np.zeros((max_length, attention_features_shape))

    hidden = decoder.reset_state(batch_size=1)

    temp_input = tf.expand_dims(load_image(image)[0], 0)
    img_tensor_val = image_features_extract_model(temp_input)
    img_tensor_val = tf.reshape(img_tensor_val, (img_tensor_val.shape[0], -1, img_tensor_val.shape[3]))

    features = encoder(img_tensor_val)

    dec_input = tf.expand_dims([tokenizer.word_index['<start>']], 0)
    result = []

    for i in range(max_length):
        predictions, hidden, attention_weights = decoder(dec_input, features, hidden)

        attention_plot[i] = tf.reshape(attention_weights, (-1, )).numpy()

        predicted_id = tf.random.categorical(predictions, 1)[0][0].numpy()
        result.append(tokenizer.index_word[predicted_id])

        if tokenizer.index_word[predicted_id] == '<end>':
            return result, attention_plot

        dec_input = tf.expand_dims([predicted_id], 0)

    attention_plot = attention_plot[:len(result), :]
    return result, attention_plot

19.2 畫注意力圖

def plot_attention(image, result, attention_plot):
    temp_image = np.array(Image.open(image))

    fig = plt.figure(figsize=(10, 10))

    len_result = len(result)
    for l in range(len_result):
        temp_att = np.resize(attention_plot[l], (8, 8))
        ax = fig.add_subplot(len_result//2, len_result//2, l+1)
        ax.set_title(result[l])
        img = ax.imshow(temp_image)
        ax.imshow(temp_att, cmap='gray', alpha=0.6, extent=img.get_extent())

    plt.tight_layout()
    plt.show()

19.3 隨機測試驗證集圖片

# captions on the validation set
rid = np.random.randint(0, len(img_name_val))
image = img_name_val[rid]
real_caption = ' '.join([tokenizer.index_word[i] for i in cap_val[rid] if i not in [0]])
result, attention_plot = evaluate(image)

print ('Real Caption:', real_caption)
print ('Prediction Caption:', ' '.join(result))
plot_attention(image, result, attention_plot)

19.4 測試自己的圖片

image_url = 'https://tensorflow.org/images/surf.jpg'
image_extension = image_url[-4:]
image_path = tf.keras.utils.get_file('image'+image_extension,
                                     origin=image_url)

result, attention_plot = evaluate(image_path)
print ('Prediction Caption:', ' '.join(result))
plot_attention(image_path, result, attention_plot)
# opening the image
Image.open(image_path)