in preprocess.py [0:0]
def lens_distortion_calibration(args):
import cv2
# alpha = 1 # 0 crops a very tight image to eliminate as much invalid/black area as possible, 1 keeps all pixels from the input, which leads to a large invalid/black area. values in between 0 and 1 are possible.
input_folder = os.path.join(args.input, "images")
output_folder = args.output
images = sorted(os.listdir(input_folder))
images = [filename for filename in images if filename[-4:] in [".png", ".jpg"]]
if args.visualize_detections:
detected_folder = os.path.join(output_folder, "detected/")
create_folder(detected_folder)
# based on https://docs.opencv.org/3.4/dc/dbb/tutorial_py_calibration.html
# termination criteria
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
# prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
objp = np.zeros((args.checkerboard_width * args.checkerboard_height, 3), np.float32)
objp[:, :2] = np.mgrid[
0 : args.checkerboard_height, 0 : args.checkerboard_width
].T.reshape(-1, 2)
# Arrays to store object points and image points from all the images.
objpoints = [] # 3d point in real world space
imgpoints = [] # 2d points in image plane.
for i, filename in enumerate(images):
print(str(i) + " / " + str(len(images)) + " " + filename, flush=True)
img = cv2.imread(os.path.join(input_folder, filename))
height, width = img.shape[:2]
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# Find the chess board corners
corners = np.array([])
ret, corners = cv2.findChessboardCorners(
gray, (args.checkerboard_height, args.checkerboard_width), corners, flags=0
)
# If found, add object points, image points (after refining them)
if ret == True:
objpoints.append(objp)
corners2 = cv2.cornerSubPix(
gray, corners, (5, 5), (-1, -1), criteria
) # corners2 == corners
imgpoints.append(corners)
if args.visualize_detections:
# Draw and display the corners
cv2.drawChessboardCorners(
img,
(args.checkerboard_height, args.checkerboard_width),
corners2,
ret,
)
cv2.imwrite(os.path.join(detected_folder, filename), img)
# cv2.imshow('img', img)
# cv2.waitKey(500)
if args.visualize_detections:
# cv2.destroyAllWindows() # uncomment if using imshow()
pass
print("computing calibration...", flush=True)
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(
objpoints, imgpoints, gray.shape[::-1], None, None
)
print("RMSE (in pixel units): " + str(ret), flush=True)
newcameramtx = mtx
roi = (0, 0, width, height)
# newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (width, height), alpha, (width, height))
calibration_paramters = {
"newcameramtx": newcameramtx.tolist(),
"mtx": mtx.tolist(),
"dist": dist.tolist(),
"roi": roi,
}
print("storing calibration...", flush=True)
import json
with open(os.path.join(output_folder, "lens_distortion.json"), "w") as json_file:
json.dump(calibration_paramters, json_file, indent=4)
if args.undistort_calibration_images:
undistorted_folder = os.path.join(output_folder, "undistorted/")
create_folder(undistorted_folder)
for i, filename in enumerate(images):
distorted_image = cv2.imread(os.path.join(input_folder, filename))
newcameramtx = np.array(calibration_paramters["newcameramtx"])
mtx = np.array(calibration_paramters["mtx"])
dist = np.array(calibration_paramters["dist"])
roi = np.array(calibration_paramters["roi"])
# https://docs.opencv.org/3.4/dc/dbb/tutorial_py_calibration.html
# undistort
undistorted_image = cv2.undistort(
distorted_image, mtx, dist, None, newcameramtx
)
# crop the image
x, y, w, h = roi
undistorted_image = undistorted_image[y : y + h, x : x + w]
output_file = os.path.join(undistorted_folder, filename)
cv2.imwrite(output_file, undistorted_image)
# rgb mask
if i == 0:
Image.fromarray(
(255.0 * (np.sum(undistorted_image == 0.0, axis=-1) != 3)).astype(
"uint8"
),
mode="L",
).save(os.path.join(undistorted_folder, "mask.png"))