Optimal Transport for Image Processing

Image processing: Transforming MNIST digits using optimal transport

The goal of this post is to create the following transformation which maps a number 2 to number 4 using Optimal Transport, where the two numbers are hand written digit images, see MNIST digit dataset. The Python code is available on my github.

Here, we briefly recap the Optimal Transport, which aims to move (mass of) points from “source” to the (mass of) points in the “target”.

Mathematically, it is a linear programming problem:
$$\begin{aligned} \min_{P \in \mathbb R^{n \times n}} \hspace{0.5cm} & \sum_{i,j} C_{ij}P_{ij}\\ s.t. \hspace{0.5cm} & P_{ij}\geq 0, \forall i, j =1,..., n\\ & \sum_j P_{ij} = a_i, \forall i=1,...,n\\ & \sum_i P_{ij} = b_j, \forall j=1,...,n \end{aligned}\tag{OT}$$

In this formulation, $a$ represents the vector of resources, $b$ is the vector of targets, and $C_{ij}$ denotes the cost associated with moving a unit from resource $a_i$ to target $b_j$. The matrix $P$ is the solution, representing the plan, with $P_{ij}$ indicating the quantity of product being transported from $a_i$ to $b_j$. In practical applications, the cost matrix $C$ is often constructed using squared distance:

$$C_{ij} = ||x^{(1)}_i-x^{(2)}_j||_2^2$$

Here $x^{(1)}$ and $x^{(2)}$ (in $\mathbb R^{n\times 2}$) denote the locations of resources $a \in \mathbb R^n$ and targets $b\in \mathbb R^n$, respectively.

For simplicity, we require that $\sum_{i=1}^n a_i=\sum_{i=1}^n b_i$.

Returning to the images labeled 2 and 4, we extract $n$ pixels from each image, assigning unit mass to each pixel. Subsequently, we apply Optimal Transport to determine the corresponding pairs (assuming that such pairing exists) and employ linear interpolation between these points. This process results in the creation of the GIF image showcased at the outset of this blog post.