Clean up nanobind typing and source layout

This commit is contained in:
2026-03-27 11:19:53 +08:00
parent 9d63177de0
commit 69e83d8430
19 changed files with 792 additions and 47 deletions
+18 -7
View File
@@ -34,7 +34,15 @@ if TYPE_CHECKING:
import numpy as np
import numpy.typing as npt
from ._helpers import CameraLike, CameraModelLike, DepthImageLike, Matrix3x3Like, PoseViewLike, VectorLike
from ._helpers import (
CameraLike,
CameraModelLike,
DepthImageLike,
Matrix3x3Like,
PoseViewLike,
TranslationVectorLike,
VectorLike,
)
PoseArray2D = npt.NDArray[np.float32]
PoseArray3D = npt.NDArray[np.float32]
@@ -64,22 +72,25 @@ def make_camera(
K: "Matrix3x3Like",
DC: "VectorLike",
R: "Matrix3x3Like",
T: "Sequence[Sequence[float]]",
T: "TranslationVectorLike",
width: int,
height: int,
model: "CameraModel | CameraModelLike",
) -> Camera:
"""Create an immutable camera and precompute its cached projection fields."""
"""Create an immutable camera and precompute its cached projection fields.
from ._helpers import _coerce_camera_model, _coerce_distortion
`T` may be a flat `[x, y, z]` vector or a nested translation matrix with shape `[1, 3]` or `[3, 1]`.
"""
from ._helpers import _coerce_camera_model, _coerce_distortion, _coerce_matrix3x3, _coerce_translation
camera_model = _coerce_camera_model(model)
return _make_camera(
name,
K,
_coerce_matrix3x3(K, "K").tolist(),
_coerce_distortion(DC, camera_model),
R,
T,
_coerce_matrix3x3(R, "R").tolist(),
_coerce_translation(T).tolist(),
width,
height,
camera_model,
+61 -20
View File
@@ -5,25 +5,34 @@ import numpy as np
import numpy.typing as npt
from ._core import (
AssociationReport,
AssociationStatus,
Camera,
CameraModel,
CoreProposalDebug,
FinalPoseAssociationDebug,
FullProposalDebug,
GroupingDebug,
MergeDebug,
PairCandidate,
PreviousPoseFilterDebug,
PreviousPoseMatch,
ProposalGroupDebug,
TriangulationConfig,
TriangulationOptions,
TriangulationResult,
TriangulationTrace,
AssociationReport as AssociationReport,
AssociationStatus as AssociationStatus,
Camera as Camera,
CameraModel as CameraModel,
CoreProposalDebug as CoreProposalDebug,
FinalPoseAssociationDebug as FinalPoseAssociationDebug,
FullProposalDebug as FullProposalDebug,
GroupingDebug as GroupingDebug,
MergeDebug as MergeDebug,
PairCandidate as PairCandidate,
PreviousPoseFilterDebug as PreviousPoseFilterDebug,
PreviousPoseMatch as PreviousPoseMatch,
ProposalGroupDebug as ProposalGroupDebug,
TriangulationConfig as TriangulationConfig,
TriangulationOptions as TriangulationOptions,
TriangulationResult as TriangulationResult,
TriangulationTrace as TriangulationTrace,
)
from ._helpers import (
CameraLike,
CameraModelLike,
DepthImageLike,
Matrix3x3Like,
PoseViewLike,
RoomParamsLike,
TranslationVectorLike,
VectorLike,
)
from ._helpers import CameraLike, CameraModelLike, DepthImageLike, Matrix3x3Like, PoseViewLike, RoomParamsLike, VectorLike
PoseArray2D: TypeAlias = npt.NDArray[np.float32]
PoseArray3D: TypeAlias = npt.NDArray[np.float32]
@@ -40,7 +49,7 @@ def make_camera(
K: Matrix3x3Like,
DC: VectorLike,
R: Matrix3x3Like,
T: Sequence[Sequence[float]],
T: TranslationVectorLike,
width: int,
height: int,
model: CameraModel | CameraModelLike,
@@ -155,4 +164,36 @@ def triangulate_with_report(
) -> TriangulationResult: ...
__all__: list[str]
__all__ = [
"Camera",
"CameraModel",
"AssociationReport",
"AssociationStatus",
"apply_depth_offsets",
"FinalPoseAssociationDebug",
"TriangulationConfig",
"TriangulationOptions",
"TriangulationResult",
"CoreProposalDebug",
"FullProposalDebug",
"GroupingDebug",
"MergeDebug",
"PairCandidate",
"PreviousPoseFilterDebug",
"PreviousPoseMatch",
"ProposalGroupDebug",
"TriangulationTrace",
"build_pair_candidates",
"convert_cameras",
"filter_pairs_with_previous_poses",
"lift_depth_poses_to_world",
"make_camera",
"make_triangulation_config",
"merge_rgbd_views",
"pack_poses_2d",
"reconstruct_rgbd",
"sample_depth_for_poses",
"triangulate_debug",
"triangulate_poses",
"triangulate_with_report",
]
+530
View File
@@ -0,0 +1,530 @@
from collections.abc import Sequence
import enum
from typing import Annotated, overload
import numpy
from numpy.typing import NDArray
class CameraModel(enum.Enum):
PINHOLE = 0
FISHEYE = 1
class Camera:
"""Immutable camera calibration with precomputed projection cache fields."""
@property
def name(self) -> str: ...
@property
def K(self) -> list[list[float]]: ...
@property
def DC(self) -> list[float]: ...
@property
def R(self) -> list[list[float]]: ...
@property
def T(self) -> list[list[float]]: ...
@property
def width(self) -> int: ...
@property
def height(self) -> int: ...
@property
def model(self) -> CameraModel: ...
@property
def invR(self) -> list[list[float]]: ...
@property
def center(self) -> list[float]: ...
@property
def newK(self) -> list[list[float]]: ...
@property
def invK(self) -> list[list[float]]: ...
def __repr__(self) -> str: ...
class TriangulationOptions:
"""Score and grouping thresholds used by triangulation."""
def __init__(self) -> None: ...
@property
def min_match_score(self) -> float: ...
@min_match_score.setter
def min_match_score(self, arg: float, /) -> None: ...
@property
def min_group_size(self) -> int: ...
@min_group_size.setter
def min_group_size(self, arg: int, /) -> None: ...
class TriangulationConfig:
"""Stable scene configuration used for triangulation."""
def __init__(self) -> None: ...
@property
def cameras(self) -> list[Camera]: ...
@cameras.setter
def cameras(self, arg: Sequence[Camera], /) -> None: ...
@property
def roomparams(self) -> list[list[float]]: ...
@roomparams.setter
def roomparams(self, arg: Sequence[Sequence[float]], /) -> None: ...
@property
def joint_names(self) -> list[str]: ...
@joint_names.setter
def joint_names(self, arg: Sequence[str], /) -> None: ...
@property
def options(self) -> TriangulationOptions: ...
@options.setter
def options(self, arg: TriangulationOptions, /) -> None: ...
class PairCandidate:
def __init__(self) -> None: ...
@property
def view1(self) -> int: ...
@view1.setter
def view1(self, arg: int, /) -> None: ...
@property
def view2(self) -> int: ...
@view2.setter
def view2(self, arg: int, /) -> None: ...
@property
def person1(self) -> int: ...
@person1.setter
def person1(self, arg: int, /) -> None: ...
@property
def person2(self) -> int: ...
@person2.setter
def person2(self, arg: int, /) -> None: ...
@property
def global_person1(self) -> int: ...
@global_person1.setter
def global_person1(self, arg: int, /) -> None: ...
@property
def global_person2(self) -> int: ...
@global_person2.setter
def global_person2(self, arg: int, /) -> None: ...
class PreviousPoseMatch:
def __init__(self) -> None: ...
@property
def previous_pose_index(self) -> int: ...
@previous_pose_index.setter
def previous_pose_index(self, arg: int, /) -> None: ...
@property
def previous_track_id(self) -> int: ...
@previous_track_id.setter
def previous_track_id(self, arg: int, /) -> None: ...
@property
def score_view1(self) -> float: ...
@score_view1.setter
def score_view1(self, arg: float, /) -> None: ...
@property
def score_view2(self) -> float: ...
@score_view2.setter
def score_view2(self, arg: float, /) -> None: ...
@property
def matched_view1(self) -> bool: ...
@matched_view1.setter
def matched_view1(self, arg: bool, /) -> None: ...
@property
def matched_view2(self) -> bool: ...
@matched_view2.setter
def matched_view2(self, arg: bool, /) -> None: ...
@property
def kept(self) -> bool: ...
@kept.setter
def kept(self, arg: bool, /) -> None: ...
@property
def decision(self) -> str: ...
@decision.setter
def decision(self, arg: str, /) -> None: ...
class PreviousPoseFilterDebug:
def __init__(self) -> None: ...
@property
def used_previous_poses(self) -> bool: ...
@used_previous_poses.setter
def used_previous_poses(self, arg: bool, /) -> None: ...
@property
def matches(self) -> list[PreviousPoseMatch]: ...
@matches.setter
def matches(self, arg: Sequence[PreviousPoseMatch], /) -> None: ...
@property
def kept_pair_indices(self) -> list[int]: ...
@kept_pair_indices.setter
def kept_pair_indices(self, arg: Sequence[int], /) -> None: ...
@property
def kept_pairs(self) -> list[PairCandidate]: ...
@kept_pairs.setter
def kept_pairs(self, arg: Sequence[PairCandidate], /) -> None: ...
class CoreProposalDebug:
def __init__(self) -> None: ...
@property
def pair_index(self) -> int: ...
@pair_index.setter
def pair_index(self, arg: int, /) -> None: ...
@property
def pair(self) -> PairCandidate: ...
@pair.setter
def pair(self, arg: PairCandidate, /) -> None: ...
@property
def score(self) -> float: ...
@score.setter
def score(self, arg: float, /) -> None: ...
@property
def kept(self) -> bool: ...
@kept.setter
def kept(self, arg: bool, /) -> None: ...
@property
def drop_reason(self) -> str: ...
@drop_reason.setter
def drop_reason(self, arg: str, /) -> None: ...
@property
def pose_3d(self) -> Annotated[NDArray[numpy.float32], dict(shape=(None, 4), order='C')]: ...
class ProposalGroupDebug:
def __init__(self) -> None: ...
@property
def center(self) -> list[float]: ...
@center.setter
def center(self, arg: Sequence[float], /) -> None: ...
@property
def proposal_indices(self) -> list[int]: ...
@proposal_indices.setter
def proposal_indices(self, arg: Sequence[int], /) -> None: ...
@property
def pose_3d(self) -> Annotated[NDArray[numpy.float32], dict(shape=(None, 4), order='C')]: ...
class GroupingDebug:
def __init__(self) -> None: ...
@property
def initial_groups(self) -> list[ProposalGroupDebug]: ...
@initial_groups.setter
def initial_groups(self, arg: Sequence[ProposalGroupDebug], /) -> None: ...
@property
def duplicate_pair_drops(self) -> list[int]: ...
@duplicate_pair_drops.setter
def duplicate_pair_drops(self, arg: Sequence[int], /) -> None: ...
@property
def groups(self) -> list[ProposalGroupDebug]: ...
@groups.setter
def groups(self, arg: Sequence[ProposalGroupDebug], /) -> None: ...
class FullProposalDebug:
def __init__(self) -> None: ...
@property
def source_core_proposal_index(self) -> int: ...
@source_core_proposal_index.setter
def source_core_proposal_index(self, arg: int, /) -> None: ...
@property
def pair(self) -> PairCandidate: ...
@pair.setter
def pair(self, arg: PairCandidate, /) -> None: ...
@property
def pose_3d(self) -> Annotated[NDArray[numpy.float32], dict(shape=(None, 4), order='C')]: ...
class MergeDebug:
def __init__(self) -> None: ...
@property
def group_proposal_indices(self) -> list[list[int]]: ...
@group_proposal_indices.setter
def group_proposal_indices(self, arg: Sequence[Sequence[int]], /) -> None: ...
@property
def merged_poses(self) -> Annotated[NDArray[numpy.float32], dict(shape=(None, None, 4), order='C')]: ...
class AssociationStatus(enum.Enum):
MATCHED = 0
NEW = 1
AMBIGUOUS = 2
class AssociationReport:
"""Track-association summary for a tracked triangulation call."""
def __init__(self) -> None: ...
@property
def pose_previous_indices(self) -> list[int]: ...
@pose_previous_indices.setter
def pose_previous_indices(self, arg: Sequence[int], /) -> None: ...
@property
def pose_previous_track_ids(self) -> list[int]: ...
@pose_previous_track_ids.setter
def pose_previous_track_ids(self, arg: Sequence[int], /) -> None: ...
@property
def pose_status(self) -> list[AssociationStatus]: ...
@pose_status.setter
def pose_status(self, arg: Sequence[AssociationStatus], /) -> None: ...
@property
def pose_candidate_previous_indices(self) -> list[list[int]]: ...
@pose_candidate_previous_indices.setter
def pose_candidate_previous_indices(self, arg: Sequence[Sequence[int]], /) -> None: ...
@property
def pose_candidate_previous_track_ids(self) -> list[list[int]]: ...
@pose_candidate_previous_track_ids.setter
def pose_candidate_previous_track_ids(self, arg: Sequence[Sequence[int]], /) -> None: ...
@property
def unmatched_previous_indices(self) -> list[int]: ...
@unmatched_previous_indices.setter
def unmatched_previous_indices(self, arg: Sequence[int], /) -> None: ...
@property
def unmatched_previous_track_ids(self) -> list[int]: ...
@unmatched_previous_track_ids.setter
def unmatched_previous_track_ids(self, arg: Sequence[int], /) -> None: ...
@property
def new_pose_indices(self) -> list[int]: ...
@new_pose_indices.setter
def new_pose_indices(self, arg: Sequence[int], /) -> None: ...
@property
def ambiguous_pose_indices(self) -> list[int]: ...
@ambiguous_pose_indices.setter
def ambiguous_pose_indices(self, arg: Sequence[int], /) -> None: ...
class FinalPoseAssociationDebug:
def __init__(self) -> None: ...
@property
def final_pose_index(self) -> int: ...
@final_pose_index.setter
def final_pose_index(self, arg: int, /) -> None: ...
@property
def source_core_proposal_indices(self) -> list[int]: ...
@source_core_proposal_indices.setter
def source_core_proposal_indices(self, arg: Sequence[int], /) -> None: ...
@property
def source_pair_indices(self) -> list[int]: ...
@source_pair_indices.setter
def source_pair_indices(self, arg: Sequence[int], /) -> None: ...
@property
def candidate_previous_indices(self) -> list[int]: ...
@candidate_previous_indices.setter
def candidate_previous_indices(self, arg: Sequence[int], /) -> None: ...
@property
def candidate_previous_track_ids(self) -> list[int]: ...
@candidate_previous_track_ids.setter
def candidate_previous_track_ids(self, arg: Sequence[int], /) -> None: ...
@property
def resolved_previous_index(self) -> int: ...
@resolved_previous_index.setter
def resolved_previous_index(self, arg: int, /) -> None: ...
@property
def resolved_previous_track_id(self) -> int: ...
@resolved_previous_track_id.setter
def resolved_previous_track_id(self, arg: int, /) -> None: ...
@property
def status(self) -> AssociationStatus: ...
@status.setter
def status(self, arg: AssociationStatus, /) -> None: ...
class TriangulationTrace:
"""
Full debug trace for triangulation, including pair, grouping, and association stages.
"""
def __init__(self) -> None: ...
@property
def pairs(self) -> list[PairCandidate]: ...
@pairs.setter
def pairs(self, arg: Sequence[PairCandidate], /) -> None: ...
@property
def previous_filter(self) -> PreviousPoseFilterDebug: ...
@previous_filter.setter
def previous_filter(self, arg: PreviousPoseFilterDebug, /) -> None: ...
@property
def core_proposals(self) -> list[CoreProposalDebug]: ...
@core_proposals.setter
def core_proposals(self, arg: Sequence[CoreProposalDebug], /) -> None: ...
@property
def grouping(self) -> GroupingDebug: ...
@grouping.setter
def grouping(self, arg: GroupingDebug, /) -> None: ...
@property
def full_proposals(self) -> list[FullProposalDebug]: ...
@full_proposals.setter
def full_proposals(self, arg: Sequence[FullProposalDebug], /) -> None: ...
@property
def merge(self) -> MergeDebug: ...
@merge.setter
def merge(self, arg: MergeDebug, /) -> None: ...
@property
def association(self) -> AssociationReport: ...
@association.setter
def association(self, arg: AssociationReport, /) -> None: ...
@property
def final_pose_associations(self) -> list[FinalPoseAssociationDebug]: ...
@final_pose_associations.setter
def final_pose_associations(self, arg: Sequence[FinalPoseAssociationDebug], /) -> None: ...
@property
def final_poses(self) -> Annotated[NDArray[numpy.float32], dict(shape=(None, None, 4), order='C')]: ...
class TriangulationResult:
"""
Tracked triangulation output containing poses and association metadata.
"""
def __init__(self) -> None: ...
@property
def association(self) -> AssociationReport: ...
@association.setter
def association(self, arg: AssociationReport, /) -> None: ...
@property
def poses_3d(self) -> Annotated[NDArray[numpy.float32], dict(shape=(None, None, 4), order='C')]: ...
def make_camera(name: str, K: Sequence[Sequence[float]], DC: Sequence[float], R: Sequence[Sequence[float]], T: Sequence[Sequence[float]], width: int, height: int, model: CameraModel) -> Camera: ...
def build_pair_candidates(poses_2d: Annotated[NDArray[numpy.float32], dict(shape=(None, None, None, 3), order='C', writable=False)], person_counts: Annotated[NDArray[numpy.uint32], dict(shape=(None,), order='C', writable=False)]) -> list[PairCandidate]: ...
def filter_pairs_with_previous_poses(poses_2d: Annotated[NDArray[numpy.float32], dict(shape=(None, None, None, 3), order='C', writable=False)], person_counts: Annotated[NDArray[numpy.uint32], dict(shape=(None,), order='C', writable=False)], config: TriangulationConfig, previous_poses_3d: Annotated[NDArray[numpy.float32], dict(shape=(None, None, 4), order='C', writable=False)], previous_track_ids: Annotated[NDArray[numpy.int64], dict(shape=(None,), order='C', writable=False)]) -> PreviousPoseFilterDebug: ...
@overload
def triangulate_debug(poses_2d: Annotated[NDArray[numpy.float32], dict(shape=(None, None, None, 3), order='C', writable=False)], person_counts: Annotated[NDArray[numpy.uint32], dict(shape=(None,), order='C', writable=False)], config: TriangulationConfig) -> TriangulationTrace: ...
@overload
def triangulate_debug(poses_2d: Annotated[NDArray[numpy.float32], dict(shape=(None, None, None, 3), order='C', writable=False)], person_counts: Annotated[NDArray[numpy.uint32], dict(shape=(None,), order='C', writable=False)], config: TriangulationConfig, previous_poses_3d: Annotated[NDArray[numpy.float32], dict(shape=(None, None, 4), order='C', writable=False)], previous_track_ids: Annotated[NDArray[numpy.int64], dict(shape=(None,), order='C', writable=False)]) -> TriangulationTrace: ...
def triangulate_poses(poses_2d: Annotated[NDArray[numpy.float32], dict(shape=(None, None, None, 3), order='C', writable=False)], person_counts: Annotated[NDArray[numpy.uint32], dict(shape=(None,), order='C', writable=False)], config: TriangulationConfig) -> Annotated[NDArray[numpy.float32], dict(shape=(None, None, 4), order='C')]: ...
def merge_rgbd_views(poses_3d: Annotated[NDArray[numpy.float32], dict(shape=(None, None, None, 4), order='C', writable=False)], person_counts: Annotated[NDArray[numpy.uint32], dict(shape=(None,), order='C', writable=False)], config: TriangulationConfig, max_distance: float = 0.5) -> Annotated[NDArray[numpy.float32], dict(shape=(None, None, 4), order='C')]: ...
def triangulate_with_report(poses_2d: Annotated[NDArray[numpy.float32], dict(shape=(None, None, None, 3), order='C', writable=False)], person_counts: Annotated[NDArray[numpy.uint32], dict(shape=(None,), order='C', writable=False)], config: TriangulationConfig, previous_poses_3d: Annotated[NDArray[numpy.float32], dict(shape=(None, None, 4), order='C', writable=False)], previous_track_ids: Annotated[NDArray[numpy.int64], dict(shape=(None,), order='C', writable=False)]) -> TriangulationResult: ...
+40 -12
View File
@@ -1,34 +1,44 @@
from __future__ import annotations
from collections.abc import Sequence
from typing import Literal, TypeAlias, TypedDict
from jaxtyping import Float
import numpy as np
import numpy.typing as npt
from ._core import Camera, CameraModel, TriangulationConfig, TriangulationOptions, make_camera
from ._core import Camera, CameraModel, TriangulationConfig, TriangulationOptions, make_camera as _make_camera
Matrix3x3Like: TypeAlias = Sequence[Sequence[float]]
VectorLike: TypeAlias = Sequence[float]
Matrix3x3: TypeAlias = Float[np.ndarray, "3 3"]
DistortionVector: TypeAlias = Float[np.ndarray, "coeffs"]
TranslationVector: TypeAlias = Float[np.ndarray, "3"]
TranslationColumn: TypeAlias = Float[np.ndarray, "3 1"]
TranslationRow: TypeAlias = Float[np.ndarray, "1 3"]
Matrix3x3Like: TypeAlias = Matrix3x3 | Sequence[Sequence[float]]
VectorLike: TypeAlias = DistortionVector | Sequence[float]
TranslationVectorLike: TypeAlias = (
TranslationVector | TranslationColumn | TranslationRow | Sequence[float] | Sequence[Sequence[float]]
)
RoomParamsLike: TypeAlias = npt.NDArray[np.generic] | Sequence[Sequence[float]]
PoseViewLike: TypeAlias = npt.NDArray[np.generic] | Sequence[Sequence[Sequence[float]]] | Sequence[Sequence[float]]
DepthImageLike: TypeAlias = npt.NDArray[np.generic] | Sequence[Sequence[float]]
class CameraDict(TypedDict, total=False):
class _CameraDictRequired(TypedDict):
name: str
K: Matrix3x3Like
DC: VectorLike
R: Matrix3x3Like
T: Sequence[Sequence[float]]
T: TranslationVectorLike
width: int
height: int
class CameraDict(_CameraDictRequired, total=False):
type: Literal["pinhole", "fisheye"]
model: Literal["pinhole", "fisheye"] | CameraModel
CameraModelLike: TypeAlias = CameraModel | Literal["pinhole", "fisheye"]
CameraLike = Camera | CameraDict
CameraLike: TypeAlias = Camera | CameraDict
DEFAULT_DEPTH_OFFSETS_METERS: dict[str, float] = {
"nose": 0.005,
@@ -79,6 +89,24 @@ def _coerce_distortion(distortion: VectorLike, camera_model: CameraModel) -> tup
return values
def _coerce_matrix3x3(matrix: object, field_name: str) -> Matrix3x3:
array = np.asarray(matrix, dtype=np.float32)
if array.shape != (3, 3):
raise ValueError(f"{field_name} must have shape [3, 3].")
return np.ascontiguousarray(array, dtype=np.float32)
def _coerce_translation(translation: object) -> TranslationColumn:
array = np.asarray(translation, dtype=np.float32)
if array.shape == (3,):
array = array[:, np.newaxis]
elif array.shape == (1, 3):
array = array.T
if array.shape != (3, 1):
raise ValueError("T must have shape [3], [1, 3], or [3, 1].")
return np.ascontiguousarray(array, dtype=np.float32)
def _coerce_depth_image(depth_image: DepthImageLike) -> npt.NDArray[np.float32]:
array = np.asarray(depth_image, dtype=np.float32)
if array.ndim == 3 and array.shape[-1] == 1:
@@ -99,12 +127,12 @@ def convert_cameras(cameras: Sequence[CameraLike]) -> list[Camera]:
camera_model = _coerce_camera_model(cam.get("model", cam.get("type", "pinhole")))
converted.append(
make_camera(
_make_camera(
str(cam["name"]),
cam["K"],
_coerce_matrix3x3(cam["K"], "K").tolist(),
_coerce_distortion(cam["DC"], camera_model),
cam["R"],
cam["T"],
_coerce_matrix3x3(cam["R"], "R").tolist(),
_coerce_translation(cam["T"]).tolist(),
int(cam["width"]),
int(cam["height"]),
camera_model,