-
Notifications
You must be signed in to change notification settings - Fork 1
/
calibration_functions.py
115 lines (86 loc) · 3.85 KB
/
calibration_functions.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import glob
import os.path as path
import pickle
import cv2
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import numpy as np
from globals import m_chess, n_chess
def calibrateCamera_decorater(func):
"""
Decorater for the calibrateCamera_SLOW() to avoid the calculation of the calibration parameters everytime we run the program.
"""
calibration_data = 'camera_cal/calibration_data.pickle'
def timesaver(*args, **kwargs):
if path.exists(calibration_data):
print('Loading calculated camera calibration...', end=' ')
with open(calibration_data, 'rb') as dump_file:
calibration_param = pickle.load(dump_file)
else:
print('Computing camera calibration_param...', end=' ')
calibration_param = func(*args, **kwargs)
with open(calibration_data, 'wb') as dump_file:
pickle.dump(calibration_param, dump_file)
print('Done.')
return calibration_param
return timesaver
@calibrateCamera_decorater
def calibrateCamera_SLOW(calib_images_directory, verbose=False):
assert path.exists(calib_images_directory), '"{}" The folder must exist and contain calibration images.'.format(
calib_images_directory)
# prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
objp = np.zeros((n_chess*m_chess, 3), np.float32)
objp[:, :2] = np.mgrid[0:m_chess, 0:n_chess].T.reshape(-1, 2)
# Arrays to store object points and image points from all the images.
objpoints = [] # 3d points in real world space
imgpoints = [] # 2d points in image plane.
# Make a list of calibration images
images = glob.glob(path.join(calib_images_directory, 'calibration*.jpg'))
# Step through the list and search for chessboard corners
for idx, fname in enumerate(images):
img = cv2.imread(fname)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# Find the chessboard corners
ret, corners = cv2.findChessboardCorners(gray, (m_chess, n_chess), None)
# If found, add object points, image points
if ret == True:
objpoints.append(objp)
imgpoints.append(corners)
if verbose:
# Draw and display the corners
cv2.drawChessboardCorners(img, (m_chess, n_chess), corners, ret)
#write_name = 'corners_found'+str(idx)+'.jpg'
#cv2.imwrite(write_name, img)
cv2.imshow('img', img)
cv2.waitKey(500)
if verbose:
cv2.destroyAllWindows()
# Do camera calibration given object points and image points
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(
objpoints, imgpoints, img.shape[1::-1], None, None)
return ret, mtx, dist, rvecs, tvecs
def undistort(img, mtx, dist, verbose=False):
"""
Undistort the image givine the camera's distortion and matrix coefficients
:param image: input image
:param mtx: camera matrix
:param dist: distortion coefficients
:param verbose: if True, show frame before/after distortion correction
:return: undistorted image
"""
image_undistorted = cv2.undistort(
img, mtx, dist, None, newCameraMatrix=mtx)
if verbose:
fig, ax = plt.subplots(nrows=1, ncols=2)
ax[0].imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
ax[1].imshow(cv2.cvtColor(image_undistorted, cv2.COLOR_BGR2RGB))
plt.show()
return image_undistorted
if __name__ == '__main__':
ret, mtx, dist, rvecs, tvecs = calibrateCamera_SLOW(
calib_images_directory='camera_cal', verbose=False)
img = cv2.imread('test_images/test2.jpg')
img_undistorted = undistort(img, mtx, dist, verbose=True)
cv2.imwrite('calibration_output/test2_before_calibration.jpg', img)
cv2.imwrite('calibration_output/test2_after_calibration.jpg',
img_undistorted)