{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import cv2.aruco as aruco\n",
    "from typing import Sequence, cast\n",
    "import awkward as ak\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from typing import Final\n",
    "from matplotlib import pyplot as plt\n",
    "from cv2.typing import MatLike"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A_PATH = Path(\"output/af_03.parquet\")\n",
    "B_PATH = Path(\"output/ae_08.parquet\")\n",
    "\n",
    "a_params = ak.from_parquet(A_PATH)[0]\n",
    "b_params = ak.from_parquet(B_PATH)[0]\n",
    "display(a_params)\n",
    "display(b_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_new_aruco_marker_origin(marker_length: float):\n",
    "    \"\"\"\n",
    "    Create a new ArUco marker origin with the given length.\n",
    "\n",
    "    0 -> x\n",
    "    |\n",
    "    v\n",
    "    y\n",
    "\n",
    "    0---1\n",
    "    |   |\n",
    "    3---2\n",
    "\n",
    "    So that the center of the marker is the origin for this PnP problem.\n",
    "\n",
    "    Args:\n",
    "        marker_length: The length of the marker.\n",
    "    \"\"\"\n",
    "    return np.array(\n",
    "        [\n",
    "            [-marker_length / 2, marker_length / 2, 0],\n",
    "            [marker_length / 2, marker_length / 2, 0],\n",
    "            [marker_length / 2, -marker_length / 2, 0],\n",
    "            [-marker_length / 2, -marker_length / 2, 0],\n",
    "        ]\n",
    "    ).astype(np.float32)\n",
    "\n",
    "\n",
    "DICTIONARY: Final[int] = aruco.DICT_4X4_50\n",
    "# 400mm\n",
    "MARKER_LENGTH: Final[float] = 0.4\n",
    "aruco_dict = aruco.getPredefinedDictionary(DICTIONARY)\n",
    "detector = aruco.ArucoDetector(\n",
    "    dictionary=aruco_dict, detectorParams=aruco.DetectorParameters()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_img = cv2.imread(str(Path(\"dumped/marker/video-20241205-152716-board.png\")))\n",
    "a_mtx = ak.to_numpy(a_params[\"camera_matrix\"])\n",
    "a_dist = ak.to_numpy(a_params[\"distortion_coefficients\"])\n",
    "\n",
    "b_img = cv2.imread(str(Path(\"dumped/marker/video-20241205-152721-board.png\")))\n",
    "b_mtx =  ak.to_numpy(b_params[\"camera_matrix\"])\n",
    "b_dist = ak.to_numpy(b_params[\"distortion_coefficients\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_corners, a_ids, _a_rejected = detector.detectMarkers(a_img)\n",
    "b_corners, b_ids, _b_rejected = detector.detectMarkers(b_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_corners"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ok, a_rvec, a_tvec = cv2.solvePnP(create_new_aruco_marker_origin(MARKER_LENGTH), a_corners[0], a_mtx, a_dist)\n",
    "if not ok:\n",
    "    raise ValueError(\"Failed to solve PnP for A\")\n",
    "a_img_output =  cv2.drawFrameAxes(a_img, a_mtx, a_dist, a_rvec, a_tvec, MARKER_LENGTH)\n",
    "plt.imshow(cv2.cvtColor(a_img_output, cv2.COLOR_BGR2RGB))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ok, b_rvec, b_tvec = cv2.solvePnP(create_new_aruco_marker_origin(MARKER_LENGTH), b_corners[0], b_mtx, b_dist)\n",
    "if not ok:\n",
    "    raise ValueError(\"Failed to solve PnP for B\")\n",
    "b_img_output =  cv2.drawFrameAxes(b_img, b_mtx, b_dist, b_rvec, b_tvec, MARKER_LENGTH)\n",
    "plt.imshow(cv2.cvtColor(b_img_output, cv2.COLOR_BGR2RGB))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import TypeVar, Union\n",
    "\n",
    "\n",
    "T = TypeVar(\"T\")\n",
    "\n",
    "\n",
    "def create_transform_matrix(rvec: MatLike, tvec: MatLike, dtype: type = np.float32):\n",
    "    assert rvec.shape == (3, 1)\n",
    "    assert tvec.shape == (3, 1)\n",
    "    R, _ = cv2.Rodrigues(rvec)\n",
    "    transform = np.eye(4, dtype=dtype)\n",
    "    transform[:3, :3] = R\n",
    "    transform[:3, 3] = tvec.flatten()\n",
    "    return transform\n",
    "\n",
    "\n",
    "def extract_translation(transform: MatLike):\n",
    "    assert transform.shape == (4, 4)\n",
    "    return transform[:3, 3]\n",
    "\n",
    "\n",
    "def extract_rotation(transform: MatLike):\n",
    "    assert transform.shape == (4, 4)\n",
    "    return transform[:3, :3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_trans = create_transform_matrix(a_rvec, a_tvec)\n",
    "display(a_trans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.linalg.inv(a_trans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Converts a rotation matrix to a rotation vector or vice versa\n",
    "a_rmtx, _ = cv2.Rodrigues(a_rvec)\n",
    "b_rmtx, _ = cv2.Rodrigues(b_rvec)\n",
    "a_camera_coord = -(a_rmtx.T@ a_tvec)\n",
    "b_camera_coord = -(b_rmtx.T @ b_tvec)\n",
    "distance = np.linalg.norm(a_camera_coord - b_camera_coord)\n",
    "a_distance = np.linalg.norm(a_camera_coord)\n",
    "b_distance = np.linalg.norm(b_camera_coord)\n",
    "display(\"d_ab={:.4}m a={:.4}m b={:.4}m\".format(distance, a_distance, b_distance))\n",
    "display(\"a_coord={}\".format(a_camera_coord.T))\n",
    "display(\"b_coord={}\".format(b_camera_coord.T))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}