190 lines
7.2 KiB
Python
190 lines
7.2 KiB
Python
import os
|
||
import os.path as osp
|
||
import glob
|
||
import cv2 as cv
|
||
import numpy as np
|
||
import json
|
||
import datetime
|
||
import argparse
|
||
|
||
|
||
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):
|
||
# 设置输出目录
|
||
if visualization:
|
||
outputFolder = create_output_folder(baseFolder, "calibrateOutputImages")
|
||
|
||
# 图片路径
|
||
imgPaths = read_img_paths(baseFolder)
|
||
if len(imgPaths) == 0:
|
||
print("No images found!")
|
||
return
|
||
|
||
# 存储世界坐标和像素坐标
|
||
# 计算出棋盘格中每个网格角点的坐标，之后当成世界坐标
|
||
board_w, board_h = chessboardSize
|
||
board_grid = np.zeros((board_w * board_h, 3), np.float32)
|
||
board_grid[:, :2] = np.mgrid[0:board_w, 0:board_h].T.reshape(-1, 2) * squareSize
|
||
|
||
pointsWorld = []
|
||
pointsPixel = []
|
||
|
||
# 遍历图片
|
||
for imgPath in imgPaths:
|
||
img = cv.imread(imgPath)
|
||
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
|
||
|
||
# 查找角点
|
||
ret, corners = cv.findChessboardCorners(gray, (board_w, board_h), None)
|
||
if ret:
|
||
cv.cornerSubPix(gray, corners, (11, 11), (-1, -1),
|
||
(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)
|
||
|
||
# 标定相机
|
||
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))
|
||
|
||
# 计算重投影误差
|
||
nimg = len(pointsWorld)
|
||
img_error = np.zeros(nimg)
|
||
for i in range(nimg):
|
||
imgpoints2, _ = cv.projectPoints(pointsWorld[i], rvecs[i], tvecs[i], mtx, dist)
|
||
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])
|
||
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, image_shape, None, None)
|
||
# print("Intrinsic matrix:\n", mtx.astype(np.float32))
|
||
# print("Distortion coefficients:\n", dist.astype(np.float32))
|
||
|
||
data = format_json_data(mtx, dist, image_shape)
|
||
# 在文件夹根目录下保存相机内参
|
||
outputJsonPath = osp.join(baseFolder, "intri_calib.json")
|
||
write_json(data, outputJsonPath)
|
||
return mtx, dist, image_shape
|
||
|
||
|
||
# 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)
|
||
x, y, w, h = roi
|
||
dst = dst[y:y + h, x:x + w]
|
||
return dst
|
||
|
||
# 用于去除整个文件夹中的图像畸变，保存到文件夹下的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 = create_output_folder(baseFolder, "distortion_corrected_images")
|
||
|
||
for imgPath in imgPaths:
|
||
img = cv.imread(imgPath)
|
||
dst = remove_image_distortion(img, mtx, dist)
|
||
cv.imwrite(osp.join(outputFolder, osp.basename(imgPath)), dst)
|
||
|
||
print("Distortion corrected images saved to: ", outputFolder)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
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(",")))
|
||
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()
|
||
|