#!/usr/bin/env python
# coding: utf-8

# ![](pics/header-2.png)
# 
# # Undistortion
# 
# Kevin Walcko
# 
# ---
# 
# `StereoCamera` has a method that will give you a `UndistortStereo` object which allows you to undistort stereo images easily.

# In[1]:


# reload library
get_ipython().run_line_magic('load_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')


# In[27]:


import numpy as np
import cv2
from matplotlib import pyplot as plt
from opencv_camera import StereoCamera
from opencv_camera import visualizeDistortion
from opencv_camera import UndistortStereo, UnDistort

np.set_printoptions(precision=3)
np.set_printoptions(suppress=True)


# In[3]:


img = cv2.imread("aruco-imgs-2/0.png",0)
h,w = img.shape[:2]
imgl = img[:,:w//2]
imgr = img[:,w//2:]


# In[5]:


plt.imshow(np.hstack((imgl, imgr)), cmap="gray")
plt.axis("off");


# In[25]:


sc = StereoCamera.from_yaml("camera.yml")
print(sc)


# In[26]:


h,w = imgl.shape[:2]
uds = sc.getUndistortion(h,w)

a, b = uds.undistort(imgl, imgr)

plt.figure(figsize=(10,5))
plt.imshow(np.hstack((imgl, imgr)), cmap="gray");
plt.axis("off")
plt.title("raw");

plt.figure(figsize=(10,5))
plt.imshow(np.hstack((a, b)), cmap="gray");
plt.axis("off")
plt.title("undistorted");


# # Appendix: Source Code

# In[30]:


from jtb import getCodeUrl, getCodeFile, getCodeImport

getCodeImport(UnDistort)


# In[29]:


getCodeImport(UndistortStereo)


# In[ ]: