From 867fc0885ba636046be6ba1a91307fff02ed28a3 Mon Sep 17 00:00:00 2001
From: crosstyan <crosstyan@outlook.com>
Date: Tue, 3 Dec 2024 19:46:49 +0800
Subject: [PATCH] parquet

---
 .gitattributes             |  1 +
 .gitignore                 |  1 +
 output/calibration.parquet |  3 +++
 test.py                    | 48 ++++++++++++++++++++++++++++++++++----
 4 files changed, 48 insertions(+), 5 deletions(-)
 create mode 100644 .gitattributes
 create mode 100644 output/calibration.parquet

diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 0000000..b57f642
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1 @@
+*.parquet filter=lfs diff=lfs merge=lfs -text
diff --git a/.gitignore b/.gitignore
index 8198a97..8ea6538 100644
--- a/.gitignore
+++ b/.gitignore
@@ -165,3 +165,4 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+.DS_Store
diff --git a/output/calibration.parquet b/output/calibration.parquet
new file mode 100644
index 0000000..1e7497a
--- /dev/null
+++ b/output/calibration.parquet
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8518bdc6922d9145872404899807ae3af1bb7c2dc27e8c152fc25d41f6df6fa
+size 21421
diff --git a/test.py b/test.py
index 1bd9e3b..df3726e 100644
--- a/test.py
+++ b/test.py
@@ -6,6 +6,8 @@ from enum import Enum
 from pathlib import Path
 from loguru import logger
 from itertools import chain
+from typing import Optional, TypedDict, cast
+import awkward as ak
 
 from matplotlib.pyplot import stem
 from numpy import ndarray
@@ -38,6 +40,14 @@ class ArucoDictionary(Enum):
 IMAGE_FOLDER = Path("xx")
 OUTPUT_FOLDER = Path("output")
 DICTIONARY = ArucoDictionary.Dict_4X4_50
+CALIBRATION_PARQUET: Optional[Path] = OUTPUT_FOLDER / "calibration.parquet"
+
+
+class CameraParams(TypedDict):
+    camera_matrix: MatLike
+    distortion_coefficients: MatLike
+    rotation_vectors: MatLike
+    translation_vectors: MatLike
 
 
 def main():
@@ -57,6 +67,18 @@ def main():
     all_image_points: list[MatLike] = []
     all_object_points: list[MatLike] = []
     last_shape = np.array((0, 0))
+    calibration: Optional[ak.Record] = None
+
+    def has_cal():
+        return CALIBRATION_PARQUET is not None and CALIBRATION_PARQUET.exists()
+
+    try:
+        if has_cal():
+            calibration_arr = ak.from_parquet(CALIBRATION_PARQUET)
+            calibration = calibration_arr[0]
+            logger.info(f"Loaded calibration parameters: {calibration}")
+    except Exception as e:
+        logger.error(f"Failed to load calibration parameters: {e}")
     for img_path in images:
         img = cv2.imread(str(img_path))
         last_shape = img.shape
@@ -82,6 +104,14 @@ def main():
             op, ip = board.matchImagePoints(ch_corners, ch_ids)  # type: ignore
             all_object_points.append(op)
             all_image_points.append(ip)
+            if calibration is not None:
+                mtx = cast(MatLike, ak.to_numpy(calibration["camera_matrix"]))
+                dist = cast(MatLike,ak.to_numpy(calibration["distortion_coefficients"]))
+                ret, rvec, tvec = cv2.solvePnP(op, ip, mtx, dist)
+                if ret:
+                    img = cv2.drawFrameAxes(img, mtx, dist, rvec, tvec, 0.1)
+                else:
+                    logger.warning(f"Failed to draw frame axes in {img_path}")
         else:
             logger.warning(f"Failed to detect Charuco board in {img_path}")
             continue
@@ -90,17 +120,25 @@ 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:
+
+    # compute calibration
+    if calibration is None and 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
+            all_object_points, all_image_points, (last_shape[0], last_shape[1]), None, None  # type: ignore
+        )  # 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)
+        parameters = {
+            "camera_matrix": mtx,
+            "distortion_coefficients": dist,
+            "rotation_vectors": rvecs,
+            "translation_vectors": tvecs,
+        }
+        ak.to_parquet([parameters], OUTPUT_FOLDER / "calibration.parquet")
     else:
-        logger.warning("No Charuco board detected in any image")
+        logger.warning("no calibration data calculated; either no images or already calibrated")
 
 
 if __name__ == "__main__":