lpips-tensorflow

Tensorflow port for the PyTorch implementation of the Learned Perceptual Image Patch Similarity (LPIPS) metric. This is done by exporting the model from PyTorch to ONNX and then to TensorFlow.

Getting started

Installation

Clone this repo.

git clone https://github.com/alexlee-gk/lpips-tensorflow.git
cd lpips-tensorflow

Install TensorFlow and dependencies from http://tensorflow.org/
Install other dependencies.

pip install -r requirements.txt

Using the LPIPS metric

The lpips TensorFlow function works with individual images or batches of images. It also works with images of any spatial dimensions (but the dimensions should be at least the size of the network's receptive field). This example computes the LPIPS distance between batches of images.

import numpy as np
import tensorflow as tf
import lpips_tf

batch_size = 32
image_shape = (batch_size, 64, 64, 3)
image0 = np.random.random(image_shape)
image1 = np.random.random(image_shape)
image0_ph = tf.placeholder(tf.float32)
image1_ph = tf.placeholder(tf.float32)

distance_t = lpips_tf.lpips(image0_ph, image1_ph, model='net-lin', net='alex')

with tf.Session() as session:
    distance = session.run(distance_t, feed_dict={image0_ph: image0, image1_ph: image1})

Exporting additional models

Export PyTorch model to TensorFlow through ONNX

Clone the PerceptualSimilarity submodule and add it to the PYTHONPATH.

git submodule update --init --recursive
export PYTHONPATH=PerceptualSimilarity:$PYTHONPATH

Install more dependencies.

pip install -r requirements-dev.txt

Export the model to ONNX *.onnx and TensorFlow *.pb files in the models directory.

python export_to_tensorflow.py --model net-lin --net alex

Known issues

The SqueezeNet model cannot be exported since ONNX cannot export one of the operators.