sayakpaul / mirnet-tflite-trt

TensorFlow Lite models for MIRNet for low-light image enhancement.

MIRNet-TFLite

This repository shows the TensorFlow Lite and TensorRT model conversion and inference processes for the MIRNet model as proposed by Learning Enriched Features for Real Image Restoration and Enhancement. This model is capable of enhancing low-light images upto a great extent.

Source

Model training code and pre-trained weights are provided by Soumik through this repository.

Comparison between the TensorFlow Lite and original models

TensorFlow Lite model (dynamic-range quantized)

Original model

About the notebooks

• MIRNet_TFLite.ipynb: Shows the model conversion and inference processes. Models converted in this notebook support dynamic shaped inputs.
• MIRNet_TFLite_Fixed_Shape.ipynb: Shows the model conversion and inference processes. Models converted in this notebook only support fixed shaped inputs.
• MIRNet_TRT.ipynb: Shows the model conversion process with TensorRT as well as the inference. Recommended if you would run inference with an NVIDIA GPU-enabled environment.
• Add_Metadata.ipynb: Adds metadata to TensorFlow Lite models. Metadata makes it easier for mobile developers to integrate the TensorFlow Lite models in their applications.

TensorFlow Lite models

• Dynamic shape (contains dynamic-range and fp16 quantized models)
• Fixed shape metadata-populated models on TensorFlow Hub

Benchmarking

Pixel 4 was used in order to run the benchmarking tests. Also, fixed-shape TensorFlow Lite models (accepting 400x400x3 images) were only benchmarked.

Notes

If you would run inference with an NVIDIA GPU-enabled environment then please follow along with this notebook - MIRNet_TRT.ipynb. If you use the TensorRT optimized model (as shown in that notebook) with an NVIDIA GPU-enabled environment the inference latency greatly improves (~0.6 seconds on a Tesla T4). Here's a demo of running the TensorRT optimized model on a low-light video.

mirnet-tflite-trt's Issues

Dynamic-shape integer quantized model errors out during inference

@khanhlvg,

I worked on converting the MIRNet model to TensorFlow Lite this weekend. This is model is capable of enhancing low-light images such as below -

Pretty exciting right?

While dynamic-range and fp16 quantizations went seamlessly I am encountering problems with the integer quantization. Even though the conversion process does not error out but I get a shape mismatch error during inference. Please note that the TFLite model can accept dynamic shapes so I am suspecting the representative dataset might need to be constructed accordingly. You can refer to this notebook to see all of this in action.

Could you shed some light on what could be done in order to mitigate this problem?

Metadata-populated model test and more

@khanhlvg,

I have populated the fixed-shape models with metadata (notebook). Would you be able to help verify if those work in Android Studio? If they do then I will go ahead and publish with TensorFlow Hub.

Here's where you can find the metadata populated models.

Note: I was not able to complete the inference with integer quantized fixed-shape models (notebook).

Trained weights are not accessible now

I am trying to run

inferer.download_weights('1sUlRD5MTRKKGxtqyYDpTv7T3jOW6aVAL').
Error
Cannot retrieve the public link of the file. You may need to change
the permission to 'Anyone with the link', or have had many accesses.
You may still be able to access the file from the browser:
https://drive.google.com/uc?id=1sUlRD5MTRKKGxtqyYDpTv7T3jOW6aVAL

But the file does not exist.

GPU delegate issue

@khanhlvg, while benchmarking the fixed-shape models I couldn't run them using the GPU delegate. Here are the logs -

• Dynamic-range
• Float16

Could you shed some light on this?