Rendue de déformation géométrique

#!/usr/bin/python3
import numpy as np
import pickle as pkl
from matplotlib.pyplot import *
import seaborn as sns

sns.set_style("white")

size = (1280//2, 960//2, 50)

[x, y] = np.meshgrid(
    np.arange(-size[0],size[0],size[2]),
    np.arange(-size[1],size[1],size[2]),
)

def plotCorrection(dx, dy, f, name):
    figure(figsize=(7,4))
    scale = np.sqrt(dx**2+dy**2)*133
    title(name.split('-')[-1])
    #title('scale = ' + str(scale))
    quiver(
        x+size[0], y+size[1],
        dx, dy, scale,
        angles='xy',
        scale_units='inches',
        width=0.002,
        minshaft=0.5
    );
    tight_layout()
    
    import tikzplotlib
    tikzplotlib.save('lens/'+name+'.tex')
    
    savefig('lens/'+name+'.jpg')
    

def radialDistortion(k1, k2, k3):
    r2 = x**2 + y**2;
    k = k1*r2 + k2*r2**2 + k3*r2**3;
    dx = x*k
    dy = y*k
    return dx, dy
    

def tangentialDistortion (p1, p2):
    r2 = x**2 + y**2;
    dx = 2*p1*x*y + (r2 + 2*x**2)*p2
    dy = 2*p2*x*y + (r2 + 2*y**2)*p1
    return dx, dy
    

def prismDistortion (s1, s2, s3, s4):
    r2 = x**2 + y**2;
    dx = s1*r2 + s2*r2**2
    dy = s3*r2 + s4*r2**2
    return dx, dy
    

def lensCorrection(k1, k2, k3, p1, p2, s1, s2, s3, s4):
    r2 = x**2 + y**2;
    k = k1*r2 + k2*r2**2 + k3*r2**3;
    dx = x*k + 2*p1*x*y + p2*(r2 + 2*x**2) + s1*r2 + s2*r2**2
    dy = y*k + 2*p2*x*y + p1*(r2 + 2*y**2) + s3*r2 + s4*r2**2
    return dx, dy
    

#k1,k2,k3,p1,p2,k4,k5,k6,s1,s2,s3,s4 = np.load('lens.npy')[0]
X = pkl.load(open('config.pkl', 'rb'))
k1,k2,k3,p1,p2,k4,k5,k6,s1,s2,s3,s4 = X['len_dist_450'][0]

print('radial', k1, k2, k3)
print('prism', s1, s2, s3, s4)
print('tangential', p1, p2)
print('optional', k4,k5,k6)

rdx, rdy = radialDistortion(k1,k2,k3)
pdx, pdy = prismDistortion(s1,s2,s3,s4)
tdx, tdy = tangentialDistortion(p1,p2)
dx, dy = lensCorrection(k1,k2,k3,p1,p2,s1,s2,s3,s4)
#dx, dy = rdx+tdx+pdx, rdy+tdy+tdx
scale = 5

plotCorrection(rdx, rdy, scale, 'lens-correction-radial')
plotCorrection(tdx, tdy, scale, 'lens-correction-tangential')
plotCorrection(pdx, pdy, scale, 'lens-correction-prismatic')
plotCorrection( dx,  dy, scale, 'lens-correction-complet')