camera_calibrate/calibrate_intri.py

import os.path as osp
import glob
import cv2 as cv
import numpy as np
import datetime
import argparse
from tqdm import tqdm
from calib_tools import write_json, read_json
from calib_tools import read_img_paths, create_output_folder
from calib_tools import DataPath


def format_intri_json_data(mtx, dist, image_shape, error):
    data = {
        "time": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "K": mtx.tolist(),
        "dist": dist.tolist(),
        "image_shape": image_shape,
        "error": error
    }
    return data

def calibrate_camera(camera, chessboardSize, squareSize, visualization):
    # 设置输出目录
    if visualization:
        outputFolder = create_output_folder(DataPath.intri_chessboard_vis, osp.basename(camera))

    # 图片路径
    imgPaths = read_img_paths(camera)
    if len(imgPaths) == 0:
        print("No images found!\n")
        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("Reprojection 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, mean_error


# calibrate_cameras函数中，照片按照相机编号进行分类
def calibrate_cameras(chessboardSize, squareSize, visualization):
    cameras_path = glob.glob(osp.join(DataPath.intri_chessboard_data, "cam[0-7]"))
    if len(cameras_path) == 0:
        print("No camera folders found!")
        return

    data = {}
    for camera_path in tqdm(cameras_path, desc="Processing Cameras", ncols=100):
        cameraId = osp.basename(camera_path)
        print("\nCalibrating camera {}... ".format(cameraId))
        mtx, dist, image_shape, error = calibrate_camera(camera_path, chessboardSize, squareSize, visualization)
        data[cameraId] = format_intri_json_data(mtx, dist, image_shape, error)
    write_json(data, osp.join(DataPath.intri_json_path, "intri.json"))
    print("Calibration data saved to: ", osp.join(DataPath.intri_json_path, "intri.json"))


# 去除图像畸变
def unistort_img(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


# 用于去除整个文件夹中的图像畸变
def unistort_imgs(mtx, dist):
    imgPaths = read_img_paths(DataPath.intri_undistort_data)
    if len(imgPaths) == 0:
        print("No images found!")
        return

    for imgPath in imgPaths:
        img = cv.imread(imgPath)
        dst = unistort_img(img, mtx, dist)
        cv.imwrite(osp.join(DataPath.intri_undistort_result, osp.basename(imgPath)), dst)

    print("Distortion corrected images saved to: ", DataPath.intri_undistort_result)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="相机内参标定和图像去畸变")
    parser.add_argument("--action", type=str, required=True, choices=["cameras", "undistort"],
                        help=" --action cameras: 标定多个相机"
                             " --action distortion: 去除图像畸变")
    parser.add_argument("--chessboardSize", type=str, default="11,8",
                        help="棋盘格角点数 (列数, 行数)，例如 '11,8'")
    parser.add_argument("--squareSize", type=float, default=60.0,
                        help="棋盘格方块的实际边长（单位与数据一致，例如 mm 或 m）")
    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 == "cameras":
        calibrate_cameras(chessboardSize, args.squareSize, args.vis)
    elif args.action == "undistort":
        print("Removing image distortion, require input folder")
        data = read_json(osp.join(DataPath.intri_json_path, "intri.json"))
        mtx = np.array(data["K"])
        dist = np.array(data["dist"])
        unistort_imgs(mtx, dist)
    else:
        print("Invalid action!")
        parser.print_help()