From 4ad22811d955ab52cd41bf19ef8d5fcbe8a75fe6 Mon Sep 17 00:00:00 2001
From: crosstyan <crosstyan@outlook.com>
Date: Thu, 21 Nov 2024 19:01:28 +0800
Subject: [PATCH] many

---
 test.py | 51 ++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 38 insertions(+), 13 deletions(-)

diff --git a/test.py b/test.py
index 415a352..1bd9e3b 100644
--- a/test.py
+++ b/test.py
@@ -1,4 +1,5 @@
 import cv2
+import numpy as np
 from cv2 import aruco
 from cv2.typing import MatLike
 from enum import Enum
@@ -7,6 +8,7 @@ from loguru import logger
 from itertools import chain
 
 from matplotlib.pyplot import stem
+from numpy import ndarray
 
 
 class ArucoDictionary(Enum):
@@ -37,15 +39,29 @@ IMAGE_FOLDER = Path("xx")
 OUTPUT_FOLDER = Path("output")
 DICTIONARY = ArucoDictionary.Dict_4X4_50
 
+
 def main():
     OUTPUT_FOLDER.mkdir(exist_ok=True)
-    images = chain(IMAGE_FOLDER.glob("*.jpeg"), IMAGE_FOLDER.glob("*.png"), IMAGE_FOLDER.glob("*.jpg"))
+    images = chain(
+        IMAGE_FOLDER.glob("*.jpeg"),
+        IMAGE_FOLDER.glob("*.png"),
+        IMAGE_FOLDER.glob("*.jpg"),
+    )
+    border_num_x = 10
+    border_num_y = 7
+    dictionary = aruco.getPredefinedDictionary(DICTIONARY.value)
+    board = aruco.CharucoBoard((border_num_x, border_num_y), 0.115, 0.09, dictionary)
+    detector = aruco.CharucoDetector(board)
+    all_ch_corners: list[MatLike] = []
+    all_ch_ids: list[MatLike] = []
+    all_image_points: list[MatLike] = []
+    all_object_points: list[MatLike] = []
+    last_shape = np.array((0, 0))
     for img_path in images:
         img = cv2.imread(str(img_path))
+        last_shape = img.shape
         # 10x7
         # minus 1 when dealing with normal chessboard
-        border_num_x = 10
-        border_num_y = 7
         # 115mm square
         # 90mm marker
         # https://docs.opencv.org/3.4/df/d4a/tutorial_charuco_detection.html
@@ -56,20 +72,17 @@ def main():
         # https://docs.opencv.org/4.x/d9/d0c/group__calib3d.html#ga93efa9b0aa890de240ca32b11253dd4a
         # https://github.com/opencv/opencv/issues/22083
         # OpenCV 4.10.x
-        dictionary = aruco.getPredefinedDictionary(DICTIONARY.value)
-        board = aruco.CharucoBoard(
-            (border_num_x, border_num_y), 0.115, 0.09, dictionary
-        )
-        detector = aruco.CharucoDetector(board)
-
         ch_corners, ch_ids, markers_corners, marker_ids = detector.detectBoard(img)
         # https://docs.opencv.org/4.10.0/d9/df5/classcv_1_1aruco_1_1CharucoDetector.html
         if ch_corners is not None:
             # https://docs.opencv.org/4.x/d4/db2/classcv_1_1aruco_1_1Board.html
-            aruco.drawDetectedCornersCharuco(
-                img, ch_corners, ch_ids, (0, 255, 0)
-            )
-        else: 
+            aruco.drawDetectedCornersCharuco(img, ch_corners, ch_ids, (0, 255, 0))
+            all_ch_corners.append(ch_corners)
+            all_ch_ids.append(ch_ids)
+            op, ip = board.matchImagePoints(ch_corners, ch_ids)  # type: ignore
+            all_object_points.append(op)
+            all_image_points.append(ip)
+        else:
             logger.warning(f"Failed to detect Charuco board in {img_path}")
             continue
         if markers_corners is not None:
@@ -77,6 +90,18 @@ def main():
         output_path = OUTPUT_FOLDER / (f"{img_path.stem}_output.jpg")
         logger.info(f"Saving to {output_path}")
         cv2.imwrite(str(output_path), img)
+    if len(all_image_points) > 0:
+        ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(
+            all_object_points, all_image_points, last_shape[::-1], None, None
+        ) # type: ignore
+        logger.info(f"Camera matrix: {mtx}")
+        logger.info(f"Distortion coefficients: {dist}")
+        logger.info(f"Rotation vectors: {rvecs}")
+        logger.info(f"Translation vectors: {tvecs}")
+        np.save("camera_matrix.npy", mtx)
+    else:
+        logger.warning("No Charuco board detected in any image")
+
 
 if __name__ == "__main__":
     main()