從一張低分辨率的圖片生成一張對應的高分辨率圖片的任務通常被稱為單圖超分(Single Image Super Resolution - SISR)。盡管可以使用傳統(tǒng)的插值方法（如雙線性插值和雙三次插值）來完成這個任務，但是產(chǎn)生的圖片質(zhì)量卻經(jīng)常差強人意。深度學習，尤其是對抗生成網(wǎng)絡 GAN，已經(jīng)被成功應用在超分任務上，比如 SRGAN 和 ESRGAN 都可以生成比較真實的超分圖片。那么在本文里，我們將介紹一下如何使用TensorFlow Hub上的一個預訓練的 ESRGAN 模型來在一個安卓 app 中生成超分圖片。最終的 app 效果如下圖，我們也已經(jīng)將完整代碼開源給大家參考。

SRGAN
https://arxiv.org/abs/1609.04802

ESRGAN
https://arxiv.org/abs/1809.00219

完整代碼
https://github.com/tensorflow/examples/tree/master/lite/examples/super_resolution

首先，我們可以很方便的從 TFHub 上加載 ESRGAN 模型，然后很容易的將其轉(zhuǎn)化為一個 TFLite 模型。注意在這里我們使用了動態(tài)范圍量化(dynamic range quantization)，并將輸入圖片的尺寸固定在50x50像素（我們已經(jīng)將轉(zhuǎn)化后的模型上傳到 TFHub 上了）：

model = hub.load("https://tfhub.dev/captain-pool/esrgan-tf2/1") concrete_func = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY] concrete_func.inputs[0].set_shape([1, 50, 50, 3]) converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func]) converter.optimizations = [tf.lite.Optimize.DEFAULT] tflite_model = converter.convert() # Save the TF Lite model. with tf.io.gfile.GFile('ESRGAN.tflite', 'wb') as f: f.write(tflite_model) esrgan_model_path = './ESRGAN.tflite'

TFHub
https://hub.tensorflow.google.cn/

TFHub（轉(zhuǎn)化后模型）
https://hub.tensorflow.google.cn/captain-pool/lite-model/esrgan-tf2/1

現(xiàn)在 TFLite 已經(jīng)支持動態(tài)大小的輸入，所以你也可以在模型轉(zhuǎn)化的時候不指定輸入圖片的大小，而在運行的時候動態(tài)指定。如果你想使用動態(tài)輸入大小，請參考這個例子。

例子
https://github.com/tensorflow/tensorflow/blob/c58c88b23122576fc99ecde988aab6041593809b/tensorflow/lite/python/lite_test.py#L529-L560

模型轉(zhuǎn)化完之后，我們可以很快驗證 ESRGAN 生成的超分圖片質(zhì)量確實比雙三次插值要好很多。如果你想更多的了解 ESRGAN 模型，我們還有另外一個教程可供參考：

lr = cv2.imread(test_img_path) lr = cv2.cvtColor(lr, cv2.COLOR_BGR2RGB) lr = tf.expand_dims(lr, axis=0) lr = tf.cast(lr, tf.float32) # Load TFLite model and allocate tensors. interpreter = tf.lite.Interpreter(model_path=esrgan_model_path) interpreter.allocate_tensors() # Get input and output tensors. input_details = interpreter.get_input_details() output_details = interpreter.get_output_details() # Run the model interpreter.set_tensor(input_details[0]['index'], lr) interpreter.invoke() # Extract the output and postprocess it output_data = interpreter.get_tensor(output_details[0]['index']) sr = tf.squeeze(output_data, axis=0) sr = tf.clip_by_value(sr, 0, 255) sr = tf.round(sr) sr = tf.cast(sr, tf.uint8)

教程
https://tensorflow.google.cn/hub/tutorials/image_enhancing

LR: 輸入的低分辨率圖片,該圖從 DIV2K 數(shù)據(jù)集中的一張蝴蝶圖片中切割出來. ESRGAN (x4): ESRGAN 模型生成的超分圖片，單邊分辨率提升4倍. Bicubic: 雙三次插值生成圖片. 在這里大家可以很容易看出來，雙三次插值生成的圖片要比 ESRGAN 模型生成的超分圖片模糊很多

你可能已經(jīng)知道，TensorFlow Lite 是 TensorFlow 用于在端側(cè)運行的官方框架，目前全球已有超過40億臺設備在運行 TFLite，它可以運行在安卓，iOS，基于 Linux 的 IoT 設備以及微處理器上。你可以使用 Java, C/C++ 或其他編程語言來運行 TFLite。在這篇文章中，我們將使用 TFLite C API，因為有許多的開發(fā)者表示希望我們能提供這樣一個范例。

DIV2K
https://data.vision.ee.ethz.ch/cvl/DIV2K/

Java, C/C++
https://tensorflow.google.cn/lite/guide/android

TFLite C API
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/c/c_api.h

我們在預先編譯好的 AAR 文件中包含了 TFLite C API需要的頭文件和庫 （包括核心庫和 GPU 庫）。為了正確的設置好 Android 項目，我們首先需要下載兩個 JAR 文件并將相應的頭文件和庫抽取出來。我們可以在一個 download.gradle 文件中定義這些任務，然后將這些任務導入 build.gradle。下面我們先定義下載 TFLite JAR 文件的兩個任務：

task downloadTFLiteAARFile() { download { src "https://bintray.com/google/tensorflow/download_file?file_path=org%2Ftensorflow%2Ftensorflow-lite%2F2.3.0%2Ftensorflow-lite-2.3.0.aar" dest "${project.rootDir}/libraries/tensorflow-lite-2.3.0.aar" overwrite false retries 5 } } task downloadTFLiteGPUDelegateAARFile() { download { src "https://bintray.com/google/tensorflow/download_file?file_path=org%2Ftensorflow%2Ftensorflow-lite-gpu%2F2.3.0%2Ftensorflow-lite-gpu-2.3.0.aar" dest "${project.rootDir}/libraries/tensorflow-lite-gpu-2.3.0.aar" overwrite false retries 5 } }

AAR 文件
https://tensorflow.google.cn/lite/guide/android#use_tflite_c_api

然后我們定義另一個任務來講頭文件和庫解壓然后放到正確的位置：

task fetchTFLiteLibs() { copy { from zipTree("${project.rootDir}/libraries/tensorflow-lite-2.3.0.aar") into "${project.rootDir}/libraries/tensorflowlite/" include "headers/tensorflow/lite/c/*h" include "headers/tensorflow/lite/*h" include "jni/**/libtensorflowlite_jni.so" } copy { from zipTree("${project.rootDir}/libraries/tensorflow-lite-gpu-2.3.0.aar") into "${project.rootDir}/libraries/tensorflowlite-gpu/" include "headers/tensorflow/lite/delegates/gpu/*h" include "jni/**/libtensorflowlite_gpu_jni.so" }

因為我們是用安卓 NDK 來編譯這個 app，我們需要讓 Android Studio 知道如何處理對應的原生文件。我們在 CMakeList.txt 文件中這樣寫：

set(TFLITE_LIBPATH "${CMAKE_CURRENT_SOURCE_DIR}/../../../../libraries/tensorflowlite/jni") set(TFLITE_INCLUDE "${CMAKE_CURRENT_SOURCE_DIR}/../../../../libraries/tensorflowlite/headers") set(TFLITE_GPU_LIBPATH "${CMAKE_CURRENT_SOURCE_DIR}/../../../../libraries/tensorflowlite-gpu/jni") set(TFLITE_GPU_INCLUDE "${CMAKE_CURRENT_SOURCE_DIR}/../../../../libraries/tensorflowlite-gpu/headers") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=gnu++14") set(CMAKE_CXX_STANDARD 14) add_library(SuperResolution SHARED SuperResolution_jni.cpp SuperResolution.cpp) add_library(lib_tensorflowlite SHARED IMPORTED) set_target_properties(lib_tensorflowlite PROPERTIES IMPORTED_LOCATION ${TFLITE_LIBPATH}/${ANDROID_ABI}/libtensorflowlite_jni.so) add_library(lib_tensorflowlite_gpu SHARED IMPORTED) set_target_properties(lib_tensorflowlite_gpu PROPERTIES IMPORTED_LOCATION ${TFLITE_GPU_LIBPATH}/${ANDROID_ABI}/libtensorflowlite_gpu_jni.so) find_library(log-lib log) include_directories(${TFLITE_INCLUDE}) target_include_directories(SuperResolution PRIVATE ${TFLITE_INCLUDE}) include_directories(${TFLITE_GPU_INCLUDE}) target_include_directories(SuperResolution PRIVATE ${TFLITE_GPU_INCLUDE}) target_link_libraries(SuperResolution android lib_tensorflowlite lib_tensorflowlite_gpu ${log-lib})

我們在 app 里包含了3個示例圖片，這樣用戶可能會運行同一個模型多次，這意味著為了提高運行效率，我們需要將 TFLite 解釋執(zhí)行器進行緩存。這一點我們可以在解釋執(zhí)行器成功建立后通過將其指針從 C++ 傳回到 Java 來實現(xiàn)：

extern "C" JNIEXPORT jlong JNICALL Java_org_tensorflow_lite_examples_superresolution_MainActivity_initWithByteBufferFromJNI(JNIEnv *env, jobject thiz, jobject model_buffer, jboolean use_gpu) { const void *model_data = static_cast(env->GetDirectBufferAddress(model_buffer)); jlong model_size_bytes = env->GetDirectBufferCapacity(model_buffer); SuperResolution *super_resolution = new SuperResolution(model_data, static_cast(model_size_bytes), use_gpu); if (super_resolution->IsInterpreterCreated()) { LOGI("Interpreter is created successfully"); return reinterpret_cast(super_resolution); } else { delete super_resolution; return 0; } }

解釋執(zhí)行器建立之后，運行模型實際上就非常簡單了，我們只需要按照 TFLite C API 來就好。不過我們需要注意的是如何從每個像素中抽取 RGB 值：

// Extract RGB values from each pixel float input_buffer[kNumberOfInputPixels * kImageChannels]; for (int i = 0, j = 0; i < kNumberOfInputPixels; i++) { // Alpha is ignored input_buffer[j++] = static_cast((lr_img_rgb[i] >> 16) & 0xff); input_buffer[j++] = static_cast((lr_img_rgb[i] >> 8) & 0xff); input_buffer[j++] = static_cast((lr_img_rgb[i]) & 0xff); }

TFLite C API
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/c/c_api.h

運行完模型后我們需要再將 RGB 值再打包進每個像素：

// Postprocess the output from TFLite int clipped_output[kImageChannels]; auto rgb_colors = std::make_unique(kNumberOfOutputPixels); for (int i = 0; i < kNumberOfOutputPixels; i++) { for (int j = 0; j < kImageChannels; j++) { clipped_output[j] = std::max(0, std::min(255, output_buffer[i * kImageChannels + j])); } // When we have RGB values, we pack them into a single pixel. // Alpha is set to 255. rgb_colors[i] = (255u & 0xff) << 24 | (clipped_output[0] & 0xff) << 16 | (clipped_output[1] & 0xff) << 8 | (clipped_output[2] & 0xff); }

那么到這里我們就完成了這個 app 的關(guān)鍵步驟，我們可以用這個 app 來生成超分圖片。您可以在對應的代碼庫中看到更多信息。我們希望這個范例能作為一個好的參考來幫助剛剛起步的開發(fā)者更快的掌握如何使用 TFLite C/C++ API 來搭建自己的機器學習 app。

對應的代碼庫中
https://github.com/tensorflow/examples/tree/master/lite/examples/super_resolution

致謝

作者十分感謝 @captain__pool 將他實現(xiàn)的 ESRGAN 模型上傳到 TFHub, 以及 TFLite 團隊的 Tian Lin 和 Jared Duke 提供十分有幫助的反饋。

— 參考 —

[1] Christian Ledig, Lucas Theis, Ferenc Huszar, Jose Caballero, Andrew Cunningham, Alejandro Acosta, Andrew Aitken, Alykhan Tejani, Johannes Totz, Zehan Wang, Wenzhe Shi. 2016. Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network.

[2] Xintao Wang, Ke Yu, Shixiang Wu, Jinjin Gu, Yihao Liu, Chao Dong, Chen Change Loy, Yu Qiao, Xiaoou Tang. 2018. ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks.

[3] Tensorflow 2.x based implementation of EDSR, WDSR and SRGAN for single image super-resolution

https://github.com/krasserm/super-resolution

[4] @captain__pool 的 ESGRAN 代碼實現(xiàn)

https://github.com/captain-pool/GSOC

[5] Eirikur Agustsson, Radu Timofte. 2017. NTIRE 2017 Challenge on Single Image Super-Resolution: Dataset and Study.

責任編輯：xj

原文標題：學習教程 | 使用 TensorFlow Lite 在 Android App 中生成超分圖片

文章出處：【微信公眾號：TensorFlow】歡迎添加關(guān)注！文章轉(zhuǎn)載請注明出處。