polynomial-inverse-warp.py

import numpy as np
from skimage.transform import warp, PolynomialTransform
from skimage import data
import matplotlib.pyplot as plt

def swirl(xy, x0, y0, R):
    r = np.sqrt((xy[:,1]-x0)**2 + (xy[:,0]-y0)**2)
    a = np.pi*r / R
    xy = np.copy(xy)
    xy[:, 1], xy[:, 0] = (
        (xy[:, 1]-x0)*np.cos(a) + (xy[:, 0]-y0)*np.sin(a) + x0,
        -(xy[:, 1]-x0)*np.sin(a) + (xy[:, 0]-y0)*np.cos(a) + y0,
    )
    return xy

# Step 1: get image
im = data.astronaut() / 255

# Step 2: warp image
im_warped = warp(im, swirl, map_args={'x0': 256, 'y0': 256, 'R': 500})

# Step 3: approximate warp with polynomial transform and unwarp
t = PolynomialTransform()
y, x = np.mgrid[:im.shape[0], :im.shape[1]]
dst_indices = np.hstack((x.reshape(-1, 1), y.reshape(-1,1)))
dst_indices = dst_indices[np.random.choice(dst_indices.shape[0], 200),:]
src_indices = swirl(dst_indices, 256, 256, 500).astype(int)
fit_success = t.estimate(src_indices, dst_indices, order=5)
print(fit_success)
im_unwarped = warp(im_warped, t, order=1, mode='constant', cval=0)

# Step 4: display results
fig, (ax0, ax1, ax2) = plt.subplots(1, 3)
ax0.imshow(im)
ax0.set_axis_off()
ax0.set_title('Input image')

ax1.imshow(im_warped)
ax1.set_axis_off()
ax1.set_title('Warped image')

ax2.imshow(im_unwarped)
ax2.set_axis_off()
ax2.set_title('Unwarped')

plt.show()