RSVP pour votre événement TensorFlow Everywhere local dès aujourd'hui!
Cette page a été traduite par l'API Cloud Translation.
Switch to English

Super résolution avec TensorFlow Lite

Voir sur TensorFlow.org Exécuter dans Google Colab Afficher la source sur GitHub Télécharger le carnet Voir le modèle TF Hub

Aperçu

La tâche consistant à récupérer une image haute résolution (HR) à partir de son homologue basse résolution est communément appelée Super résolution d'image unique (SISR).

Le modèle utilisé ici est ESRGAN ( ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks ). Et nous allons utiliser TensorFlow Lite pour exécuter l'inférence sur le modèle pré-entraîné.

Le modèle TFLite est converti à partir de cette implémentation hébergée sur TF Hub. Notez que le modèle que nous avons converti suréchantillonne une image basse résolution 50x50 en une image haute résolution 200x200 (facteur d'échelle = 4). Si vous souhaitez une taille d'entrée ou un facteur d'échelle différent, vous devez reconvertir ou réentraîner le modèle d'origine.

Installer

Installons d'abord les bibliothèques requises.

pip install -q matplotlib tensorflow tensorflow-hub

Importez les dépendances.

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

Téléchargez et convertissez le modèle ESRGAN

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'

Téléchargez une image test (tête d'insecte).

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
8192/6432 [======================================] - 0s 0us/step

Générer une image de super résolution à l'aide de 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)

Visualisez le résultat

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());
<matplotlib.image.AxesImage at 0x7f53e43f7d68>

png

png

Benchmarks de performance

Les numéros de référence de performance sont générés avec l'outil décrit ici .

Nom du modèle Taille du modèle Dispositif CPU GPU
super résolution (ESRGAN) 4,8 Mo Pixel 3 586,8 ms * 128,6 ms
Pixel 4 385,1 ms * 130,3 ms

* 4 fils utilisés