完成标定和去畸变功能,加入argparse

This commit is contained in:
fuly 2024-11-06 20:55:17 +08:00
parent 3e63b4760a
commit 87b48807e9

View File

@ -4,19 +4,49 @@ import glob
import cv2 as cv
import numpy as np
import json
import datetime
import argparse
def calibrate_camera(imgFolder, chessboardSize, squareSize):
def format_json_data(mtx, dist, image_shape):
data = {
"time": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
"K": mtx.tolist(),
"dist": dist.tolist(),
"image_shape": image_shape
}
return data
def write_json(data, output_path):
with open(output_path, "w") as f:
json.dump(data, f, indent=4)
def read_json(input):
with open(input, "r") as f:
data = json.load(f)
return data
def read_img_paths(imgFolder):
imgPaths = []
for extension in ["jpg", "png", "jpeg", "bmp"]:
imgPaths += glob.glob(osp.join(imgFolder, "*.{}".format(extension)))
return imgPaths
def create_output_folder(baseFolder, outputFolder):
folder = osp.join(baseFolder, outputFolder)
if not osp.exists(folder):
os.makedirs(folder)
return folder
def calibrate_camera(baseFolder, chessboardSize, squareSize, visualization):
# 设置输出目录
outputFolder = osp.join(imgFolder, "output")
if not osp.exists(outputFolder):
os.makedirs(outputFolder)
if visualization:
outputFolder = create_output_folder(baseFolder, "calibrateOutputImages")
# 图片路径
imgPaths = []
for extension in ["jpg", "png", "jpeg"]:
imgPaths += glob.glob(osp.join(imgFolder, "*.{}".format(extension)))
imgPaths = read_img_paths(baseFolder)
if len(imgPaths) == 0:
print("No images found!")
return
@ -42,12 +72,13 @@ def calibrate_camera(imgFolder, chessboardSize, squareSize):
(cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001))
pointsWorld.append(board_grid)
pointsPixel.append(corners)
if visualization:
cv.drawChessboardCorners(img, (board_w, board_h), corners, ret)
cv.imwrite(osp.join(outputFolder, osp.basename(imgPath)), img)
# 标定相机
ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(pointsWorld, pointsPixel, gray.shape[::-1], None, None)
image_shape = gray.shape[::-1]
ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(pointsWorld, pointsPixel, image_shape, None, None)
print("Intrinsic matrix:\n", mtx.astype(np.float32))
print("Distortion coefficients:\n", dist.astype(np.float32))
@ -59,27 +90,49 @@ def calibrate_camera(imgFolder, chessboardSize, squareSize):
error = cv.norm(pointsPixel[i], imgpoints2, cv.NORM_L2) / len(imgpoints2)
img_error[i] = error
good_img = np.where(img_error < 0.5)[0]
mean_error = np.mean(img_error[good_img])
print("Reprojection error: ", mean_error)
# good_img = np.where(img_error < 0.5)[0]
# mean_error = np.mean(img_error[good_img])
mean_error = np.mean(img_error)
print("\nReprojection error: ", mean_error)
# 挑选出重投影误差小于1.0的图片,重新标定相机
if len(good_img) == 0:
print("No images with error < 0.5")
elif len(good_img) == nimg:
print("All images have error < 0.5")
pass
else:
pointsWorld2 = [pointsWorld[i] for i in good_img]
pointsPixel2 = [pointsPixel[i] for i in good_img]
ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(pointsWorld2, pointsPixel2, gray.shape[::-1], None, None)
print("Intrinsic matrix:\n", mtx.astype(np.float32))
print("Distortion coefficients:\n", dist.astype(np.float32))
# if len(good_img) == 0:
# print("No images with error < 0.5")
# elif len(good_img) == nimg:
# print("All images have error < 0.5")
# pass
# else:
# pointsWorld2 = [pointsWorld[i] for i in good_img]
# pointsPixel2 = [pointsPixel[i] for i in good_img]
# ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(pointsWorld2, pointsPixel2, image_shape, None, None)
# print("Intrinsic matrix:\n", mtx.astype(np.float32))
# print("Distortion coefficients:\n", dist.astype(np.float32))
return mtx, dist
data = format_json_data(mtx, dist, image_shape)
# 在文件夹根目录下保存相机内参
outputJsonPath = osp.join(baseFolder, "intri_calib.json")
write_json(data, outputJsonPath)
return mtx, dist, image_shape
def undistort_image(img, mtx, dist):
# calibrate_cameras函数中照片按照相机编号进行分类
def calibrate_cameras(baseFolder, chessboardSize, squareSize, visualization):
cameras = glob.glob(osp.join(baseFolder, '[0-9]'))
if len(cameras) == 0:
print("No camera folders found!")
return
outputJsonPath = osp.join(baseFolder, "intri_calib.json")
data = {}
for camera in cameras:
cameraId = osp.basename(camera)
print("\nCalibrating camera{}...".format(cameraId))
mtx, dist, image_shape = calibrate_camera(camera, chessboardSize, squareSize, visualization)
data[cameraId] = format_json_data(mtx, dist, image_shape)
write_json(data, outputJsonPath)
print("\nCalibration data saved to: ", outputJsonPath)
# 去除图像畸变
def remove_image_distortion(img, mtx, dist):
h, w = img.shape[:2]
newcameramtx, roi = cv.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))
dst = cv.undistort(img, mtx, dist, None, newcameramtx)
@ -87,55 +140,50 @@ def undistort_image(img, mtx, dist):
dst = dst[y:y + h, x:x + w]
return dst
def undistort_images(imgFolder, mtx, dist):
imgPaths = []
for extension in ["jpg", "png", "jpeg"]:
imgPaths += glob.glob(osp.join(imgFolder, "*.{}".format(extension)))
# 用于去除整个文件夹中的图像畸变保存到文件夹下的distortion_corrected_images文件夹中
def remove_images_distortion(baseFolder, mtx, dist):
imgPaths = read_img_paths(baseFolder)
if len(imgPaths) == 0:
print("No images found!")
return
outputFolder = osp.join(imgFolder, "undistorted_images")
if not osp.exists(outputFolder):
os.makedirs(outputFolder)
outputFolder = create_output_folder(baseFolder, "distortion_corrected_images")
for imgPath in imgPaths:
img = cv.imread(imgPath)
dst = undistort_image(img, mtx, dist)
dst = remove_image_distortion(img, mtx, dist)
cv.imwrite(osp.join(outputFolder, osp.basename(imgPath)), dst)
print("Undistorted images saved to: ", outputFolder)
def calibrate_cameras(imgFolder, chessboardSize, squareSize):
mtxs = []
dists = []
for folder in glob.glob(osp.join(imgFolder, "*")):
if not osp.isdir(folder):
continue
mtx, dist = calibrate_camera(folder, chessboardSize, squareSize)
mtxs.append(mtx)
dists.append(dist)
return mtxs, dists
def write_json_data(mtx, dist, outputFolder):
data = {
"intrinsic_matrix": mtx.tolist(),
"distortion_coefficients": dist.tolist()
}
with open(osp.join(outputFolder, "calibration.json"), "w") as f:
json.dump(data, f, indent=4)
print("Distortion corrected images saved to: ", outputFolder)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("imgFolder", help="Folder containing images for calibration")
parser.add_argument("--chessboardSize", help="Size of chessboard (rows, cols)", default="11,8")
parser.add_argument("--squareSize", help="Size of square in chessboard", default=60)
parser = argparse.ArgumentParser(description="相机内参标定和图像去畸变")
parser.add_argument("--action", type=str, required=True, choices=["camera", "cameras", "distortion"],
help=" --action camera: 标定单个相机"
" --action cameras: 标定多个相机"
" --action distortion: 去除图像畸变")
parser.add_argument("--folder", type=str, required=True, help="包含相机文件夹的基础文件夹")
parser.add_argument("--chessboardSize", type=str, default="11,8", help="棋盘格尺寸,格式为'w,h'")
parser.add_argument("--squareSize", type=float, default=60.0, help="棋盘格方块尺寸")
parser.add_argument("--no-vis", dest="vis", action="store_false", help="禁用标定结果的可视化输出")
args = parser.parse_args()
chessboardSize = tuple(map(int, args.chessboardSize.split(",")))
squareSize = float(args.squareSize)
if args.action == "camera":
print("Calibrating camera...")
mtx, dist, image_shape = calibrate_camera(args.folder, chessboardSize, args.squareSize, args.vis)
outputJsonPath = osp.join(args.folder, "intri_calib.json")
print("Calibration data saved to: ", outputJsonPath)
elif args.action == "cameras":
calibrate_cameras(args.folder, chessboardSize, args.squareSize, args.vis)
elif args.action == "distortion":
print("Removing image distortion...")
data = read_json(osp.join(args.folder, "intri_calib.json"))
mtx = np.array(data["K"])
dist = np.array(data["dist"])
remove_images_distortion(args.folder, mtx, dist)
else:
print("Invalid action!")
parser.print_help()