ช่วยปกป้อง Great Barrier Reef กับ TensorFlow บน Kaggle เข้าร่วมท้าทาย

ความละเอียดสูงสุดด้วย TensorFlow Lite

ดูบน TensorFlow.org ทำงานใน Google Colab ดูแหล่งที่มาบน GitHub ดาวน์โหลดโน๊ตบุ๊ค ดูรุ่น TF Hub

ภาพรวม

งานในการกู้คืนภาพที่มีความละเอียดสูง (HR) จากภาพที่มีความละเอียดต่ำมักเรียกว่า Single Image Super Resolution (SISR)

รูปแบบที่ใช้ที่นี่เป็น ESRGAN ( ESRGAN: ปรับปรุง Super-Resolution กำเนิดเครือข่ายขัดแย้ง ) และเราจะใช้ TensorFlow Lite เพื่อทำการอนุมานบนโมเดลที่ฝึกไว้ล่วงหน้า

รุ่น TFLite ถูกแปลงจากนี้ การดำเนินการ โฮสต์บน TF Hub โปรดทราบว่าโมเดลที่เราแปลงจะอัพตัวอย่างรูปภาพความละเอียดต่ำ 50x50 เป็นรูปภาพความละเอียดสูง 200x200 (ตัวคูณมาตราส่วน=4) หากคุณต้องการขนาดอินพุตหรือตัวคูณมาตราส่วนที่แตกต่างกัน คุณต้องแปลงหรือฝึกโมเดลดั้งเดิมอีกครั้ง

ติดตั้ง

มาติดตั้งไลบรารีที่จำเป็นก่อน

pip install matplotlib tensorflow tensorflow-hub

นำเข้าการอ้างอิง

import tensorflow as tf
import tensorflow_hub as hub
import matplotlib.pyplot as plt
print(tf.__version__)
2.7.0

ดาวน์โหลดและแปลงโมเดล ESRGAN

model = hub.load("https://tfhub.dev/captain-pool/esrgan-tf2/1")
concrete_func = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

@tf.function(input_signature=[tf.TensorSpec(shape=[1, 50, 50, 3], dtype=tf.float32)])
def f(input):
  return concrete_func(input);

converter = tf.lite.TFLiteConverter.from_concrete_functions([f.get_concrete_function()], model)
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'
WARNING:absl:Found untraced functions such as restored_function_body, restored_function_body, restored_function_body, restored_function_body, restored_function_body while saving (showing 5 of 335). These functions will not be directly callable after loading.
INFO:tensorflow:Assets written to: /tmp/tmpinlbbz0t/assets
INFO:tensorflow:Assets written to: /tmp/tmpinlbbz0t/assets
2021-11-16 12:15:19.621471: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:363] Ignored output_format.
2021-11-16 12:15:19.621517: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:366] Ignored drop_control_dependency.
WARNING:absl:Buffer deduplication procedure will be skipped when flatbuffer library is not properly loaded

ดาวน์โหลดภาพทดสอบ (หัวแมลง)

test_img_path = tf.keras.utils.get_file('lr.jpg', 'https://raw.githubusercontent.com/tensorflow/examples/master/lite/examples/super_resolution/android/app/src/main/assets/lr-1.jpg')
Downloading data from https://raw.githubusercontent.com/tensorflow/examples/master/lite/examples/super_resolution/android/app/src/main/assets/lr-1.jpg
16384/6432 [============================================================================] - 0s 0us/step

สร้างภาพความละเอียดสูงโดยใช้ TensorFlow Lite

lr = tf.io.read_file(test_img_path)
lr = tf.image.decode_jpeg(lr)
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)

เห็นภาพผลลัพธ์

lr = tf.cast(tf.squeeze(lr, axis=0), tf.uint8)
plt.figure(figsize = (1, 1))
plt.title('LR')
plt.imshow(lr.numpy());

plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)        
plt.title(f'ESRGAN (x4)')
plt.imshow(sr.numpy());

bicubic = tf.image.resize(lr, [200, 200], tf.image.ResizeMethod.BICUBIC)
bicubic = tf.cast(bicubic, tf.uint8)
plt.subplot(1, 2, 2)   
plt.title('Bicubic')
plt.imshow(bicubic.numpy());

png

png

เกณฑ์มาตรฐานประสิทธิภาพ

หมายเลขมาตรฐานประสิทธิภาพได้รับการสร้างขึ้นด้วยเครื่องมือ อธิบายไว้ที่นี่

ชื่อรุ่น ขนาดรุ่น อุปกรณ์ ซีพียู GPU
ความละเอียดสูงสุด (ESRGAN) 4.8 Mb Pixel 3 586.8ms* 128.6ms
Pixel 4 385.1ms* 130.3ms

* 4 หัวข้อใช้