{ "cells": [ { "cell_type": "code", "execution_count": 151, "id": "b491ec50", "metadata": {}, "outputs": [], "source": [ "import json\n", "import numpy as np\n", "from pathlib import Path\n", "import awkward as ak\n", "from jaxtyping import Num, Array, Float, jaxtyped\n", "from typing import TypeAlias, TypedDict, Optional, Sequence, Generator, Any, TypeVar\n", "import jax.numpy as jnp\n", "from app.camera import Detection, Camera, CameraParams\n", "import jax\n", "from datetime import datetime, timedelta\n", "from beartype import beartype\n", "from scipy.spatial.transform import Rotation as R\n", "from more_itertools import partition\n", "NDArray: TypeAlias = np.ndarray\n", "DetectionGenerator: TypeAlias = Generator[Detection, None, None]" ] }, { "cell_type": "code", "execution_count": 152, "id": "ca140fe1", "metadata": {}, "outputs": [], "source": [ "class KeypointDataset(TypedDict):\n", " frame_index: int\n", " boxes: Num[NDArray, \"N 4\"]\n", " kps: Num[NDArray, \"N J 2\"]\n", " kps_scores: Num[NDArray, \"N J\"]\n", "\n", "T = TypeVar(\"T\")\n", "\n", "def unwrap(val: Optional[T]) -> T:\n", " if val is None:\n", " raise ValueError(\"None\")\n", " return val\n", "\n", "\n", "\n", "class Resolution(TypedDict):\n", " width: int\n", " height: int\n", "\n", "\n", "class Intrinsic(TypedDict):\n", " camera_matrix: Num[Array, \"3 3\"]\n", " \"\"\"\n", " K\n", " \"\"\"\n", " distortion_coefficients: Num[Array, \"N\"]\n", " \"\"\"\n", " distortion coefficients; usually 5\n", " \"\"\"\n", "\n", "\n", "class Extrinsic(TypedDict):\n", " rvec: Num[NDArray, \"3\"]\n", " tvec: Num[NDArray, \"3\"]\n", "\n", "class ExternalCameraParams(TypedDict):\n", " name: str\n", " port: int\n", " intrinsic: Intrinsic\n", " extrinsic: Extrinsic\n", " resolution: Resolution\n" ] }, { "cell_type": "code", "execution_count": 153, "id": "3df87e4e", "metadata": {}, "outputs": [], "source": [ "\"\"\"获得所有机位的相机内外参\"\"\"\n", "def get_camera_params(camera_path: Path) -> ak.Array:\n", " camera_dataset: ak.Array = ak.from_parquet(camera_path / \"camera_params.parquet\")\n", " return camera_dataset\n" ] }, { "cell_type": "code", "execution_count": 154, "id": "0c19ad25", "metadata": {}, "outputs": [], "source": [ "# 相机内外参路径\n", "CAMERA_PATH = Path(\n", " \"/home/admin/Documents/ActualTest_QuanCheng/camera_ex_params_1_2025_4_20/camera_params\"\n", ")\n", "# 所有机位的相机内外参\n", "AK_CAMERA_DATASET: ak.Array = get_camera_params(CAMERA_PATH)" ] }, { "cell_type": "code", "execution_count": 155, "id": "f6144d97", "metadata": {}, "outputs": [], "source": [ "@jaxtyped(typechecker=beartype)\n", "def to_transformation_matrix(\n", " rvec: Num[NDArray, \"3\"], tvec: Num[NDArray, \"3\"]\n", ") -> Num[NDArray, \"4 4\"]:\n", " res = np.eye(4)\n", " res[:3, :3] = R.from_rotvec(rvec).as_matrix()\n", " res[:3, 3] = tvec\n", " return res\n", "\n", "def from_camera_params(camera: ExternalCameraParams) -> Camera:\n", " rt = jnp.array(\n", " to_transformation_matrix(\n", " ak.to_numpy(camera[\"extrinsic\"][\"rvec\"]),\n", " ak.to_numpy(camera[\"extrinsic\"][\"tvec\"]),\n", " )\n", " )\n", " K = jnp.array(camera[\"intrinsic\"][\"camera_matrix\"]).reshape(3, 3)\n", " dist_coeffs = jnp.array(camera[\"intrinsic\"][\"distortion_coefficients\"])\n", " image_size = jnp.array(\n", " (camera[\"resolution\"][\"width\"], camera[\"resolution\"][\"height\"])\n", " )\n", " return Camera(\n", " id=camera[\"name\"],\n", " params=CameraParams(\n", " K=K,\n", " Rt=rt,\n", " dist_coeffs=dist_coeffs,\n", " image_size=image_size,\n", " ),\n", " )" ] }, { "cell_type": "code", "execution_count": 156, "id": "4dc00edf", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[{name: 'AF_01', port: 5601, intrinsic: {...}, extrinsic: {...}, ...},\n",
" {name: 'AF_02', port: 5602, intrinsic: {...}, extrinsic: {...}, ...},\n",
" {name: 'AF_03', port: 5603, intrinsic: {...}, extrinsic: {...}, ...},\n",
" {name: 'AF_04', port: 5604, intrinsic: {...}, extrinsic: {...}, ...},\n",
" {name: 'AF_05', port: 5605, intrinsic: {...}, extrinsic: {...}, ...},\n",
" {name: 'AF_06', port: 5606, intrinsic: {...}, extrinsic: {...}, ...},\n",
" {name: 'AE_01', port: 5607, intrinsic: {...}, extrinsic: {...}, ...},\n",
" {name: 'AE_1A', port: 5608, intrinsic: {...}, extrinsic: {...}, ...},\n",
" {name: 'AE_08', port: 5609, intrinsic: {...}, extrinsic: {...}, ...}]\n",
"------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\n",
"backend: cpu\n",
"nbytes: 2.3 kB\n",
"type: 9 * {\n",
" name: string,\n",
" port: int64,\n",
" intrinsic: {\n",
" camera_matrix: var * var * float64,\n",
" distortion_coefficients: var * float64\n",
" },\n",
" extrinsic: {\n",
" rvec: var * float64,\n",
" tvec: var * float64\n",
" },\n",
" resolution: {\n",
" width: int64,\n",
" height: int64\n",
" }\n",
"}"
],
"text/plain": [
"