加入了图像矫正函数

calibrate_intri.py

@@ -6,6 +6,7 @@ import numpy as np
 import json
 import argparse
 
+
 def calibrate_camera(imgFolder, chessboardSize, squareSize):
     # 设置输出目录
     outputFolder = osp.join(imgFolder, "output")
@@ -37,7 +38,8 @@ def calibrate_camera(imgFolder, chessboardSize, squareSize):
         # 查找角点
         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))
+            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)
 
@@ -67,7 +69,7 @@ def calibrate_camera(imgFolder, chessboardSize, squareSize):
     elif len(good_img) == nimg:
         print("All images have error < 0.5")
         pass
-    else :
+    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)
@@ -77,6 +79,47 @@ def calibrate_camera(imgFolder, chessboardSize, squareSize):
     return mtx, dist
 
 
+def undistort_image(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 undistort_images(imgFolder, mtx, dist):
+    imgPaths = []
+    for extension in ["jpg", "png", "jpeg"]:
+        imgPaths += glob.glob(osp.join(imgFolder, "*.{}".format(extension)))
+    if len(imgPaths) == 0:
+        print("No images found!")
+        return
+
+    outputFolder = osp.join(imgFolder, "undistorted_images")
+    if not osp.exists(outputFolder):
+        os.makedirs(outputFolder)
+
+    for imgPath in imgPaths:
+        img = cv.imread(imgPath)
+        dst = undistort_image(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(),
@@ -85,7 +128,14 @@ def write_json_data(mtx, dist, outputFolder):
     with open(osp.join(outputFolder, "calibration.json"), "w") as f:
         json.dump(data, f, indent=4)
 
-if __name__ == "__main__":
-    mtx, dist = calibrate_camera("./test/intri_imgs", (11, 8), 60)
-    write_json_data(mtx, dist, "./test/intri_imgs/output")
+
+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)
+    args = parser.parse_args()
+
+    chessboardSize = tuple(map(int, args.chessboardSize.split(",")))
+    squareSize = float(args.squareSize)