【ATU Book-i.MX8系列 - TFLite 進階】風格轉換應用

一. 概述

在邊緣運算的重點技術之中，模組輕量化網路架構 是不可或缺的一環，如何高效的利用硬體資源來達到最佳目標，特別是在效能與準確度的衡量上，是個非常有趣的議題。此章節再來探討深度學習熱門的研究項目之一 風格轉換(Style Transform) ，主要用途是讓機器繪畫出不同風格的圖片，比如說油畫、水彩畫等等。其中具代表性的神經網路架構為 Fast Style Transfer、Neural Style Transfer、Adaptive Style Transfer、a Arbitrary Style Transfer 等等。而本範例採用所謂的 任意風格轉換(Arbitrary Style Transfer) 結合 輕量化網路架構 MobileNet 作應用，後續將介紹演算法的基本概念。

若新讀者欲理解更多人工智慧、機器學習以及深度學習的資訊，可點選查閱下方博文
大大通精彩博文 【ATU Book-i.MX8系列】博文索引

TensorFlow Lite 進階系列博文-文章架構示意圖

二. 算法介紹

MobileNet 神經網路架構探討 :

核心概念是利用拆分的概念，將原本的卷積層拆成深度卷積(Depthwise Convolution) 與逐點卷積(Pointwise Convolution) 兩個部分，稱作深層可分離卷積(Depthwise Separable Convolution) 。以此方式進行運算，能夠大幅度減少參數量，以達到加快運算速度。(用途擷取特徵)

MobileNet 輕量化概念示意圖
圖文來源 - 參考 LaptrihnX 網站

Arbitrary Style Transfer :

任意風格轉換(Arbitrary Style Transfer) 引入 編碼器與解碼器(Encoder-Decoder) 概念,讓機器可以學習任意一張圖片的風格進行轉換。如下圖所示，具體一點就是拿任意一張圖片(image) 給機器(已訓練的模組, 或稱編碼器 Encoder) 學習風格 之後，再拿任意一張圖片(content_image) 透過機器(已訓練的模組, 或稱解碼器 decoder) 即可 繼承與轉換風格 來產生新的圖片(stylized_image)。

如下圖所示，任意風格轉換概念其實就是由 編碼器(Encoder) 與 解碼器(Decoder) 組成。

編碼器(Encoder) : 此稱作 Style Prediction Network。主要由 MobileNet 架構來提取特徵，並將學到的風格資訊進行編碼。
解碼器(decoder) : 此稱作 Style Transform Network。主要用途為產生影像，並於原神經網路架構的特徵層，將風格特徵的均值和標準轉換至該影像上，來產生新的影像。

任意風格轉換(Arbitrary Style Transfer) 簡易概念示意圖
圖文來源 - 參考 TensorFlow Blog 網站

任意風格轉換(Arbitrary Style Transfer) 的另一個重點概念就是 自適應實例標準(Adaptive Instance Normalization, AdaIN) 的轉換概念。
如下圖所示，編碼器與解碼器之間存在 AdaIN 轉接器。其中 AdaIN 公式的意思可想像成先去風格化 (減去自身均值u(x) 再除以自身標準差σ(x) )，再進行添加風格化 (乘 style image 的標準差 σ(y) 再 style image 的加均值u(y) )。

任意風格轉換(Arbitrary Style Transfer) 之 AdaIN 概念示意圖
圖文來源 - 參考 TensorFlow Blog 網站

故編碼器、解碼器、AdaIN 三大概念就是任意風格轉換的基本概念，有興趣者可參考此連結。

三. 算法實現

Google 官方有提供效果極佳的 magenta_arbitrary-image-stylization.tflite 模組，可查看第六步驟直接下載使用。
故這裡利用遷移學習方法與 TF-Slim實現 任意風格轉換(Arbitrary Style Transfer)。

實現步驟如下:

第一步 : 開啟 Colab 設定環境

%tensorflow_version 1.x
!python -c 'import matplotlib as tf; print(tf.__version__)' # Check the version of the tensorflow

第二步 : Magenta Model 與套件下載與安裝

Magenta 是由 Google Brain 團隊研究人員所創建，供開發者研究使用。

# 安裝必要套件
%cd root
!pip install magenta pytest-pylint
!pip install tensorflow-probability==0.8.0
!pip install gast==0.2.2
!apt-get install build-essential libasound2-dev libjack-dev
# 下載與安裝 Magenta Model
!git clone https://github.com/tensorflow/magenta.git
%cd /root/magenta
!sudo apt-get install build-essential libasound2-dev libjack-dev portaudio19-dev
!pip install -e .
!python setup.py test
!python setup.py bdist_wheel --universal

第三步 : TensorFlow Slim 下載與安裝

import os
os.environ['PYTHONPATH'] += ':/root/models/research/:/root/models/research/slim/:/root/models/research/object_detection/utils/:/root/models/research/object_detection'
!pip install tf_slim
!python object_detection/builders/model_builder_test.py  # tf-slim 模組建立測試

第四步 : 下載資料庫

%cd /root/
!mkdir dataset
%cd /root/dataset
!wget https://www.robots.ox.ac.uk/~vgg/data/dtd/download/dtd-r1.0.1.tar.gz
!tar -xvzf dtd-r1.0.1.tar.gz
!rm dtd-r1.0.1.tar.gz

第五步 : 製作 TensorFlow Record

%cd /root/magenta/magenta/models/image_stylization
STYLE_IMAGES_PATHS="/root/dataset/dtd/images/cobwebbed/cobwebbed_0030.jpg"
RECORDIO_PATH="/root/dataset/dtd_cobwebbed.tfrecord"
!python image_stylization_create_dataset.py \
 --style_files=/root/dataset/dtd/images/cobwebbed/*.jpg \
 --output_file=$RECORDIO_PATH --compute_gram_matrices=False --logtostderr

*** 執行後，將於 /root/dataset 資料夾內產出 dtd_cobwebbed.tfrecord 檔案

第六步 : 下載訓練過的模組"

此步驟利用之前訓練過的模組資源重新訓練，即 遷移學習(Transfer Learning) 的技術。

%cd /root
!gdown https://storage.googleapis.com/download.magenta.tensorflow.org/models/arbitrary_style_transfer.tar.gz
!tar -xvzf arbitrary_style_transfer.tar.gz
!rm arbitrary_style_transfer.tar.gz

*** 因 ImageNet 資料庫不容易取得，故此小節先以 VGG-16 的 Arbitrary Style Transfer 作範例。
請下載官方現有資源(已訓練完成之模組)，並可直接進行實作步驟 :

學習風格之 Predict 模型 (編碼器) : https://tfhub.dev/google/lite-model/magenta/arbitrary-image-stylization-v1-256/int8/prediction/1?lite-format=tflite
轉換風格之 Transfer 模型 (解碼器) : https://tfhub.dev/google/lite-model/magenta/arbitrary-image-stylization-v1-256/int8/transfer/1?lite-format=tflite

第七步 : 進行訓練

%cd /root/magenta/magenta/models/arbitrary_image_stylization
!arbitrary_image_stylization_train \
      --batch_size=8 \
      --imagenet_data_dir=/root/dataset/train2017 \
      --style_dataset_file=/root/dataset/dtd_cobwebbed.tfrecord \
      --train_dir=/root/trained \
      --content_weights={\"vgg_16/conv3\":2.0} \
      --random_style_image_size=False \
      --augment_style_images=False \
      --center_crop=True \
      --logtostderr

第八步 : 產生Frozen Graph

此步驟可以調整模組輸出大小，比如說將原本輸入大小 224x224 改成 96x96 。

import tensorflow as tf

import tf_slim as slim
from magenta.models.arbitrary_image_stylization import arbitrary_image_stylization_build_model as build_model
from magenta.models.image_stylization import image_utils
with tf.Graph().as_default(), tf.Session() as sess:
style_img_ph = tf.placeholder(tf.float32, shape=[None, None, 3])
style_img_preprocessed = image_utils.resize_image(style_img_ph, 256)
content_img_ph = tf.placeholder(tf.float32, shape=[None, None, 3])
content_img_preprocessed = image_utils.resize_image(content_img_ph, 256)
stylized_images, _, _, bottleneck_feat = build_model.build_model(content_img_preprocessed, \
style_img_preprocessed, trainable=False, is_training=False, \
inception_end_point='Mixed_6e', style_prediction_bottleneck=100, adds_losses=False)
checkpoint = tf.train.latest_checkpoint("/root/arbitrary_style_transfer/model.ckpt")
init_fn = slim.assign_from_checkpoint_fn(checkpoint, slim.get_variables_to_restore())
sess.run([tf.local_variables_initializer()])
tf.train.write_graph(sess.graph_def, '/root/arbitrary_style_transfer/', 'arbitrary_style_transfer.pbtxt', as_text=True)
tf.train.write_graph(sess.graph_def, '/root/arbitrary_style_transfer/', 'arbitrary_style_transfer.pb', as_text=False)

第九步 : TensorFlow Lite 轉換

with tf.Graph().as_default(), tf.Session() as sess:
  style_image_tensor = tf.placeholder(tf.float32, shape=[1, 256, 256, 3], name='style_image')
  content_image_tensor = tf.placeholder(tf.float32, shape=[1, 256, 256, 3], name='content_image')
  stylized_images, _, _, bottleneck_feat = build_model.build_model(content_image_tensor, \
                                     style_image_tensor, trainable=False, is_training=False, \
                                     inception_end_point='Mixed_6e', style_prediction_bottleneck=100, adds_losses=False)
  model_saver = tf.train.Saver(tf.global_variables())
  model_saver.restore(sess, "/root/arbitrary_style_transfer/model.ckpt")
  #(1) predict 學習風格模組
  tf.saved_model.simple_save(sess,"/root/arbitrary_style_transfer/predict",
        inputs={style_image_tensor.name: style_image_tensor},
        outputs={'style_bottleneck': bottleneck_feat})
  #(2) transform 轉換風格模組
  tf.saved_model.simple_save(sess,"/root/arbitrary_style_transfer/transform",
        inputs={content_image_tensor.name: content_image_tensor,'style_bottleneck': bottleneck_feat},
        outputs={'stylized_image': stylized_images})

第十步 : TensorFlow Lite 轉換 - Predict Model

import tensorflow as tf
import numpy as np
def representative_dataset_gen():
    for _ in range(50):
        yield [np.random.uniform(0.0, 1.0, size=(1, 256, 256, 3)).astype(np.float32)] 
converter = tf.compat.v1.lite.TocoConverter.from_saved_model("/root/arbitrary_style_transfer/predict")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
converter.representative_dataset = representative_dataset_gen
tflite_quant_model = converter.convert()
with open('/root/arbitrary_style_transfer_predict.tflite', 'wb') as w:
    w.write(tflite_quant_model)

第十步 : TensorFlow Lite 轉換 - Transform Model

因 ImageNet 資料庫不容易取得，故此小節先以 VGG-16 的 Arbitrary Style Transfer 作範例，若欲想自行產出 MobielNet 的 Arbitrary Style Transfer 作範例，僅須依照此教程即可!!

import tensorflow as tf
import numpy as np
def representative_dataset_gen():
    for _ in range(50):
        yield [np.random.uniform(0.0, 1.0, size=(1, 256, 256, 3)).astype(np.float32)] 
converter = tf.compat.v1.lite.TocoConverter.from_saved_model("/root/arbitrary_style_transfer/ transform ")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8
converter.representative_dataset = representative_dataset_gen
tflite_quant_model = converter.convert()
with open('/root/arbitrary_style_transfer_transform.tflite', 'wb') as w:
    w.write(tflite_quant_model)

第十一步 : Arbitrary Style Transform 範例實現 (於 i.MX8M Plus 撰寫運行)

import io
import numpy as np
import cv2
from tflite_runtime.interpreter import Interpreter

# 解析 tensorflow lite 檔案 (Prediction)
interpreter_style = Interpreter("magenta_arbitrary-image-stylization-v1-256_int8_prediction.tflite") # 將模組移動至 i.MX8 平台
interpreter_style.allocate_tensors()
style_input_details  = interpreter_style.get_input_details()
style_output_details = interpreter_style.get_output_details()
_, height_predict, width_predict, _ = interpreter_style.get_input_details()[0]['shape']

# 解析 tensorflow lite 檔案 (Transfer)
interpreter_transfer = Interpreter("magenta_arbitrary-image-stylization-v1-256_int8_transfer.tflite") # 將模組移動至 i.MX8 平台
interpreter_transfer.allocate_tensors()
transfer_input_details  = interpreter_transfer.get_input_details()
transfer_output_details = interpreter_transfer.get_output_details()
_, height, width, _ = interpreter_transfer.get_input_details()[0]['shape']

# 學習風格
style         = cv2.imread("VanGogh_Star.jpg")
style_rgb     = cv2.cvtColor(style, cv2.COLOR_BGR2RGB) 
style_resized = cv2.resize(style_rgb, (width_predict, height_predict))
style_resized = np.expand_dims(style_resized, axis=0)
style_resized = np.array(style_resized, dtype=np.float32)
style_resized = style_resized/255
interpreter_style.set_tensor(style_input_details[0]["index"], style_resized)
interpreter_style.invoke()
style_bottleneck = interpreter_style.get_tensor(style_output_details[0]['index'])

# 載入新圖片，開始轉換風格
frame         = cv2.imread("belfry.jpg")
frame_rgb     = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame_resized = cv2.resize(frame_rgb, (width, height))
frame_resized = np.expand_dims(frame_resized, axis=0)
frame_resized = np.array(frame_resized, dtype=np.float32)
frame_resized = frame_resized/255
transfer_input_details = interpreter_transfer.get_input_details()
interpreter_transfer.set_tensor(transfer_input_details[0]["index"], frame_resized)
interpreter_transfer.set_tensor(transfer_input_details[1]["index"], style_bottleneck)
interpreter_transfer.invoke() # 進行轉換
stylized_image = interpreter_transfer.get_tensor(transfer_output_details[0]['index']) # 取得轉換之結果

# 顯示結果
v2.imshow(”StyleTransorm”,frame_rgb)
cv2.waitKey(0)
cv2.destroyAllWindows()

Arbitrary Style Transform 實現結果呈現

如下圖所示，成功將圖片轉換為梵谷星空的風格。
在 i.MX8M Plus 的 NPU 處理器，推理時間(Inference Time) 約 1.7 s。

四. 結語

風格轉換(Style Transform) 是個不錯的應用，可以學習任何一種圖片的畫風進行轉換。比較可惜的是，目前仍未找合適的作法將現有的浮點數模組轉換成整數模組，因此運行在 i.MX8MP 的 Vivante VIP8000 NPU，其推理時間可達每幀 1.7 s 的處理速度。下一章節將會介紹機器學習的新穎技術 “遷移學習(Transfer Learning)” ，敬請期待 !!。

五. 參考文件

[1] MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications
[2] Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization
[3] Optimizing style transfer to run on mobile with TFLite
[4] 从Style的角度理解Instance Normalization

如有任何相關 TensorFlow Lite 進階技術問題，歡迎至博文底下留言提問 !!
接下來還會分享更多 TensorFlow Lite 進階文章 !!敬請期待 【ATU Book-i.MX8系列 – TFLite 進階】 !!

★博文內容均由個人提供，與平台無關，如有違法或侵權，請與網站管理員聯繫。

★文明上網，請理性發言。內容一周內被舉報5次，發文人進小黑屋喔~

【ATU Book-i.MX8系列 - TFLite 進階】 風格轉換應用

評論

【ATU Book-i.MX8系列 - TFLite 進階】風格轉換應用