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

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)化為一個(gè) TFLite 模型。注意在這里我們使用了動(dòng)態(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)支持動(dòng)態(tài)大小的輸入，所以你也可以在模型轉(zhuǎn)化的時(shí)候不指定輸入圖片的大小，而在運(yùn)行的時(shí)候動(dòng)態(tài)指定。如果你想使用動(dòng)態(tài)輸入大小，請(qǐng)參考這個(gè)例子。

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

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

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ù)集中的一張蝴蝶圖片中切割出來(lái). ESRGAN (x4): ESRGAN 模型生成的超分圖片，單邊分辨率提升4倍. Bicubic: 雙三次插值生成圖片. 在這里大家可以很容易看出來(lái)，雙三次插值生成的圖片要比 ESRGAN 模型生成的超分圖片模糊很多

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

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

我們?cè)陬A(yù)先編譯好的 AAR 文件中包含了 TFLite C API需要的頭文件和庫(kù) （包括核心庫(kù)和 GPU 庫(kù)）。為了正確的設(shè)置好 Android 項(xiàng)目，我們首先需要下載兩個(gè) JAR 文件并將相應(yīng)的頭文件和庫(kù)抽取出來(lái)。我們可以在一個(gè) download.gradle 文件中定義這些任務(wù)，然后將這些任務(wù)導(dǎo)入 build.gradle。下面我們先定義下載 TFLite JAR 文件的兩個(gè)任務(wù)：

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

然后我們定義另一個(gè)任務(wù)來(lái)講頭文件和庫(kù)解壓然后放到正確的位置：

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" }

因?yàn)槲覀兪怯冒沧?NDK 來(lái)編譯這個(gè) app，我們需要讓 Android Studio 知道如何處理對(duì)應(yīng)的原生文件。我們?cè)?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})

我們?cè)?app 里包含了3個(gè)示例圖片，這樣用戶可能會(huì)運(yùn)行同一個(gè)模型多次，這意味著為了提高運(yùn)行效率，我們需要將 TFLite 解釋執(zhí)行器進(jìn)行緩存。這一點(diǎn)我們可以在解釋執(zhí)行器成功建立后通過(guò)將其指針從 C++ 傳回到 Java 來(lái)實(shí)現(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í)行器建立之后，運(yùn)行模型實(shí)際上就非常簡(jiǎn)單了，我們只需要按照 TFLite C API 來(lái)就好。不過(guò)我們需要注意的是如何從每個(gè)像素中抽取 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

運(yùn)行完模型后我們需要再將 RGB 值再打包進(jìn)每個(gè)像素：

// 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); }

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

對(duì)應(yīng)的代碼庫(kù)中
https://github.com/tensorflow/examples/tree/master/lite/examples/super_resolution

致謝

作者十分感謝 @captain__pool 將他實(shí)現(xiàn)的 ESRGAN 模型上傳到 TFHub, 以及 TFLite 團(tuán)隊(duì)的 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 代碼實(shí)現(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.

責(zé)任編輯：xj

原文標(biāo)題：學(xué)習(xí)教程 | 使用 TensorFlow Lite 在 Android App 中生成超分圖片

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