如何使用TensorFlow Hub的ESRGAN模型来在安卓app中生成超分图片

从一张低分辨率的图片生成一张对应的高分辨率图片的任务通常被称为单图超分(single image super resolution - sisr)。尽管可以使用传统的插值方法(如双线性插值和双三次插值)来完成这个任务,但是产生的图片质量却经常差强人意。深度学习,尤其是对抗生成网络 gan,已经被成功应用在超分任务上,比如 srgan 和 esrgan 都可以生成比较真实的超分图片。那么在本文里,我们将介绍一下如何使用tensorflow hub上的一个预训练的 esrgan 模型来在一个安卓 app 中生成超分图片。最终的 app 效果如下图,我们也已经将完整代码开源给大家参考。
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 模型,然后很容易的将其转化为一个 tflite 模型。注意在这里我们使用了动态范围量化(dynamic range quantization),并将输入图片的尺寸固定在50x50像素(我们已经将转化后的模型上传到 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(转化后模型)
https://hub.tensorflow.google.cn/captain-pool/lite-model/esrgan-tf2/1
现在 tflite 已经支持动态大小的输入,所以你也可以在模型转化的时候不指定输入图片的大小,而在运行的时候动态指定。如果你想使用动态输入大小,请参考这个例子。
例子
https://github.com/tensorflow/tensorflow/blob/c58c88b23122576fc99ecde988aab6041593809b/tensorflow/lite/python/lite_test.py#l529-l560
模型转化完之后,我们可以很快验证 esrgan 生成的超分图片质量确实比双三次插值要好很多。如果你想更多的了解 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 数据集中的一张蝴蝶图片中切割出来. esrgan (x4): esrgan 模型生成的超分图片,单边分辨率提升4倍. bicubic: 双三次插值生成图片. 在这里大家可以很容易看出来,双三次插值生成的图片要比 esrgan 模型生成的超分图片模糊很多
你可能已经知道,tensorflow lite 是 tensorflow 用于在端侧运行的官方框架,目前全球已有超过40亿台设备在运行 tflite,它可以运行在安卓,ios,基于 linux 的 iot 设备以及微处理器上。你可以使用 java, c/c++ 或其他编程语言来运行 tflite。在这篇文章中,我们将使用 tflite c api,因为有许多的开发者表示希望我们能提供这样一个范例。
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 解释执行器进行缓存。这一点我们可以在解释执行器成功建立后通过将其指针从 c++ 传回到 java 来实现:
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; } }
解释执行器建立之后,运行模型实际上就非常简单了,我们只需要按照 tflite c api 来就好。不过我们需要注意的是如何从每个像素中抽取 rgb 值:
// extract rgb values from each pixel float input_buffer[knumberofinputpixels * kimagechannels]; for (int i = 0, j = 0; 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 的关键步骤,我们可以用这个 app 来生成超分图片。您可以在对应的代码库中看到更多信息。我们希望这个范例能作为一个好的参考来帮助刚刚起步的开发者更快的掌握如何使用 tflite c/c++ api 来搭建自己的机器学习 app。
对应的代码库中
https://github.com/tensorflow/examples/tree/master/lite/examples/super_resolution
致谢
作者十分感谢 @captain__pool 将他实现的 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 代码实现
https://github.com/captain-pool/gsoc
[5] eirikur agustsson, radu timofte. 2017. ntire 2017 challenge on single image super-resolution: dataset and study.

原文标题:学习教程 | 使用 tensorflow lite 在 android app 中生成超分图片
文章出处:【微信公众号:tensorflow】欢迎添加关注!文章转载请注明出处。

模拟乘法器电路图
基于单片机的LED灯光立方设计
微信投票系统开发者生态爆发,公正微信投票平台快开发上线
阿根廷一银行和Bitex合作推出比特币跨境支付平台
新型激光扫描仪可在极短时间内检测到皮肤癌
如何使用TensorFlow Hub的ESRGAN模型来在安卓app中生成超分图片
Python证件照制作小程序源代码
探讨新能源汽车双积分政策背后的隐深意
5G设备和边缘计算是推动全球通信运营商资本支出增加的关键因素
二位五通电磁阀接线
邬贺铨院士:工业互联网ICT的新挑战
云计算与5G等新兴技术融合,推动传统企业数字化转型
关于GD32F350R8的单导联心电采集系统的介绍和分析
C语言实现面向对象的方式 C++中的class的运行原理
柔性供料器CCD视觉上料 柔性振动盘anyfeeder
深入了解开关电源BUCK电路各个元器件
惠普推Cortana智能扬声器 是其Windows 10 PC的配套设备
诺基亚财报让微软心寒 双方尚未找到互补契合点
正点原子开拓者FPGA:数码管动态显示实验
出售Agilent86108B电模块