# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import numpy as np
import cv2
import glob
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
# %matplotlib inline
import pickle

# Read calibration images
images = glob.glob('./calibration*.jpg')

# +
# Create arrays to store object and image points
objpoints = [] # 3D points in real world space
imgpoints = [] # 2D points in image space

# Prepare 3D object points like (0,0,0), (1,0,0), (2,0,0), ..., (nx-1,ny-1,0)
nx = 9
ny = 6
objp = np.zeros((nx*ny,3), np.float32)
objp[:,:2] = np.mgrid[0:nx,0:ny].T.reshape(-1,2)
# -

for fname in images:
    # Read each image in for loop
    img = mpimg.imread(fname)
    
    # Convert image to grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # Find the chessboard corners
    ret, corners = cv2.findChessboardCorners(gray, (nx,ny), None)
    
    # Fill object and image points array with corners
    if ret == True:
        imgpoints.append(corners)
        objpoints.append(objp)
        
        # Draw corners and write output files
        img = cv2.drawChessboardCorners(img, (nx,ny), corners, ret)
        file_name = 'corner_' + fname.split('/')[-1]
        cv2.imwrite(file_name, img)
    

# +
# Load example image for reference
img_ref = cv2.imread('./calibration4.jpg')
img_size = (img_ref.shape[1], img_ref.shape[0])

# Calibrate camera with determined object and image points
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, img_size,None,None)

# Undistort example image
dst = cv2.undistort(img_ref, mtx, dist, None, mtx)
# Save camera calibration results in a pickle file
cal_pickle = {}
cal_pickle["mtx"] = mtx
cal_pickle["dist"] = dist
pickle.dump(cal_pickle, open('./cal_pickle.p', "wb"))
# -

# Visualize undistortion
f, (ax1, ax2) = plt.subplots(1, 2, figsize=(20,10))
f.subplots_adjust(hspace = .2, wspace=.05)
ax1.imshow(img_ref)
ax1.set_title('Original Image', fontsize=30)
ax2.imshow(dst)
ax2.set_title('Undistorted Image', fontsize=30)

StarcoderdataJupyterScriptsDedupFiltered