PnP works, kinda

find_cute_object.py

@@ -1,3 +1,4 @@
+import re
 import cv2
 from cv2 import aruco
 from datetime import datetime
@@ -33,21 +34,94 @@ def main():
     camera_matrix = cast(MatLike, ak.to_numpy(cal["camera_matrix"]))
     distortion_coefficients = cast(MatLike, ak.to_numpy(cal["distortion_coefficients"]))
     ops = ak.from_parquet(OBJECT_POINTS_PARQUET)
-    board = aruco.CharucoBoard(
-        size=(3, 3), squareLength=0.127, markerLength=0.097, dictionary=aruco_dict
+    detector = aruco.ArucoDetector(
+        dictionary=aruco_dict, detectorParams=aruco.DetectorParameters()
     )
-    detector = aruco.CharucoDetector(board)
+
+    total_ids = cast(NDArray, ak.to_numpy(ops["ids"])).flatten()
+    total_corners = cast(NDArray, ak.to_numpy(ops["corners"])).reshape(-1, 4, 3)
+    ops_map: dict[int, NDArray] = dict(zip(total_ids, total_corners))
+    logger.info("ops_map={}", ops_map)
 
     for frame in gen():
         grey = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
         # pylint: disable-next=unpacking-non-sequence
-        diamond_corners, diamond_ids, markers, marker_ids = detector.detectDiamonds(
-            grey
-        )
+        markers, ids, rejected = detector.detectMarkers(grey)
         # `markers` is [N, 1, 4, 2]
         # `ids` is [N, 1]
-        if diamond_ids is not None:
-            aruco.drawDetectedDiamonds(frame, diamond_corners, diamond_ids)
+        if ids is not None:
+            markers = np.reshape(markers, (-1, 4, 2))
+            ids = np.reshape(ids, (-1, 1))
+            # logger.info("markers={}, ids={}", np.array(markers).shape, np.array(ids).shape)
+            ips_map: dict[int, NDArray] = {}
+            for cs, id in zip(markers, ids):
+                id = int(id)
+                cs = cast(NDArray, cs)
+                ips_map[id] = cs
+                center = np.mean(cs, axis=0).astype(int)
+                GREY = (128, 128, 128)
+                # logger.info("id={}, center={}", id, center)
+                cv2.circle(frame, tuple(center), 5, GREY, -1)
+                cv2.putText(
+                    frame,
+                    str(id),
+                    tuple(center),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    1,
+                    GREY,
+                    2,
+                )
+                # BGR
+                RED = (0, 0, 255)
+                GREEN = (0, 255, 0)
+                BLUE = (255, 0, 0)
+                YELLOW = (0, 255, 255)
+                color_map = [RED, GREEN, BLUE, YELLOW]
+                for color, corners in zip(color_map, cs):
+                    corners = corners.astype(int)
+                    frame = cv2.circle(frame, corners, 5, color, -1)
+            # https://docs.opencv.org/4.x/d9/d0c/group__calib3d.html#ga50620f0e26e02caa2e9adc07b5fbf24e
+            ops: NDArray = np.empty((0, 3), dtype=np.float32)
+            ips: NDArray = np.empty((0, 2), dtype=np.float32)
+            for id, ip in ips_map.items():
+                try:
+                    op = ops_map[id]
+                    assert ip.shape == (4, 2), f"corners.shape={ip.shape}"
+                    assert op.shape == (4, 3), f"op.shape={op.shape}"
+                    ops = np.concatenate((ops, op), axis=0)
+                    ips = np.concatenate((ips, ip), axis=0)
+                except KeyError:
+                    logger.warning("No object points for id={}", id)
+                    continue
+            assert len(ops) == len(ips), f"len(ops)={len(ops)} != len(ips)={len(ips)}"
+            if len(ops) > 0:
+                # https://docs.opencv.org/4.x/d5/d1f/calib3d_solvePnP.html
+                # https://docs.opencv.org/4.x/d5/d1f/calib3d_solvePnP.html#calib3d_solvePnP_flags
+                ret, rvec, tvec= cv2.solvePnP(
+                    objectPoints=ops,
+                    imagePoints=ips,
+                    cameraMatrix=camera_matrix,
+                    distCoeffs=distortion_coefficients,
+                    flags=cv2.SOLVEPNP_SQPNP,
+                )
+                # ret, rvec, tvec, inliners = cv2.solvePnPRansac(
+                #     objectPoints=ops,
+                #     imagePoints=ips,
+                #     cameraMatrix=camera_matrix,
+                #     distCoeffs=distortion_coefficients,
+                #     flags=cv2.SOLVEPNP_SQPNP,
+                # )
+                if ret:
+                    cv2.drawFrameAxes(
+                        frame,
+                        camera_matrix,
+                        distortion_coefficients,
+                        rvec,
+                        tvec,
+                        MARKER_LENGTH,
+                    )
+                else:
+                    logger.warning("Failed to solvePnPRansac")
         cv2.imshow("frame", frame)
         if (k := cv2.waitKey(1)) == ord("q"):
             logger.info("Exiting")