Refactor Jupyter notebooks for marker processing: deleted boom.ipynb and compute_3d_maybe.ipynb, added calculate_box_coord_naive.ipynb, calculate_box_face_coord_naive.ipynb, estimate_extrinstic.ipynb, find_aruco_points_with_image.ipynb, find_aruco_points.py, and find_extrinsic_object.py. Updated .gitignore to include new output files.
This commit is contained in:
150
find_extrinsic_object.py
Normal file
150
find_extrinsic_object.py
Normal file
@ -0,0 +1,150 @@
|
||||
import cv2
|
||||
from cv2 import aruco
|
||||
from datetime import datetime
|
||||
from loguru import logger
|
||||
from pathlib import Path
|
||||
from typing import Optional, cast, Final
|
||||
import awkward as ak
|
||||
from cv2.typing import MatLike
|
||||
import numpy as np
|
||||
|
||||
NDArray = np.ndarray
|
||||
CALIBRATION_PARQUET = Path("output") / "usbcam_cal.parquet"
|
||||
OBJECT_POINTS_PARQUET = Path("output") / "object_points.parquet"
|
||||
DICTIONARY: Final[int] = aruco.DICT_4X4_50
|
||||
# 400mm
|
||||
MARKER_LENGTH: Final[float] = 0.4
|
||||
|
||||
|
||||
def gen():
|
||||
API = cv2.CAP_AVFOUNDATION
|
||||
cap = cv2.VideoCapture(0, API)
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
logger.warning("Failed to grab frame")
|
||||
break
|
||||
yield frame
|
||||
|
||||
|
||||
def main():
|
||||
aruco_dict = aruco.getPredefinedDictionary(DICTIONARY)
|
||||
cal = ak.from_parquet(CALIBRATION_PARQUET)[0]
|
||||
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)
|
||||
detector = aruco.ArucoDetector(
|
||||
dictionary=aruco_dict, detectorParams=aruco.DetectorParameters()
|
||||
)
|
||||
|
||||
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)
|
||||
writer: Optional[cv2.VideoWriter] = None
|
||||
|
||||
for frame in gen():
|
||||
grey = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
# pylint: disable-next=unpacking-non-sequence
|
||||
markers, ids, rejected = detector.detectMarkers(grey)
|
||||
# `markers` is [N, 1, 4, 2]
|
||||
# `ids` is [N, 1]
|
||||
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 writer is not None:
|
||||
writer.write(frame)
|
||||
if (k := cv2.waitKey(1)) == ord("q"):
|
||||
logger.info("Exiting")
|
||||
break
|
||||
elif k == ord("s"):
|
||||
now = datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
file_name = f"aruco_{now}.png"
|
||||
logger.info("Saving to {}", file_name)
|
||||
cv2.imwrite(file_name, frame)
|
||||
elif k == ord("r"):
|
||||
if writer is not None:
|
||||
writer.release()
|
||||
writer = None
|
||||
logger.info("Recording stopped")
|
||||
else:
|
||||
now = datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
file_name = f"aruco_{now}.mp4"
|
||||
logger.info("Recording to {}", file_name)
|
||||
fourcc = cv2.VideoWriter.fourcc(*"mp4v")
|
||||
writer = cv2.VideoWriter(file_name, fourcc, 20.0, frame.shape[:2][::-1])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user