This commit is contained in:
2024-11-21 19:01:28 +08:00
parent 4a8bd957d2
commit 4ad22811d9

49
test.py
View File

@ -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,19 +72,16 @@ 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)
)
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
@ -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()