Add PnP and keypoint geometry across poses/points2D/pinhole - #393
Add PnP and keypoint geometry across poses/points2D/pinhole#393oarriaga wants to merge 4 commits into
Conversation
Consolidate the PnP and keypoint geometry that was scattered across the applications layer and the pix2pose example into a single backend module, paz.keypoints, and port the missing 2D keypoint-transform helpers. - New paz/backend/keypoints.py holds the PnP toolkit (Pose6D, solver constants, solve_PnP, solve_PnP_RANSAC, project_points3D, build_cube_points3D) moved out of applications/pose_estimators.py, plus the JAX keypoint transforms ported from master (normalize/denormalize keypoints2D, rotate_point2D, transform_keypoint, flip_keypoints_left_right, translate_keypoints, uv_to_vu). - Lift the pix2pose dense-correspondence pose solve into a reusable solve_pose_matrix_RANSAC (RANSAC PnP then Rodrigues to a rotation matrix); the example's solve_pose_from_nocs now calls it. - Rename the dimension-based cube builder to build_cube_points3D to match master and avoid colliding with pinhole.build_cube_corners, which takes min/max extents. - Expose the module as paz.keypoints; pose_estimators and the pix2pose example import from it. Co-located keypoints_test.py checks the JAX transforms against the numpy reference, a normalize/denormalize round trip, and PnP pose recovery.
Port master's keypoint augmentation (RandomKeypointRotation / RandomKeypointTranslation) as JAX functions built on the new keypoints backend, and use them in the probabilistic keypoint example. - paz.image.translate_image: pure-image (x, y) pixel shift via affine_transform, next to rotate. - paz.keypoints.rotate_keypoints2D, rotate_image_and_keypoints, translate_image_and_keypoints, image_center2D: rotate/translate an image and its keypoints together so labels track the image content. - probabilistic_keypoint_estimation/train.py: KeypointSequence now applies rotate + translate + brightness through these backend helpers and normalizes per batch, instead of augmenting only image brightness. Tests draw a hot pixel, transform it, and check the keypoint lands on the moved pixel, validating that image and keypoints stay consistent.
|
Follow-up commit
Verified with hot-pixel tests (keypoint lands on the transformed pixel) and a 6-seed end-to-end check; |
Address naming and layering feedback on the augmentation helpers. - Rename paz.image.translate_image to paz.image.translate, matching the module's own convention (rotate, crop, resize, pad do not repeat "image"). - Move the two-type augmentation recipes (rotate/translate an image together with its keypoints) out of paz.keypoints into the keypoints example, where they read as user code composing the backend primitives; drop image_center2D. The backend keeps only single-type primitives: paz.image.rotate/translate and paz.keypoints.rotate_keypoints2D/translate_keypoints.
The JAX tree already organizes point geometry by math domain (points2D, pointcloud, pinhole, poses); a dedicated keypoints backend cut across all of them and duplicated points2D.transform, points2D.denormalize and poses.project_to_image. Dissolve it into the backends that already fit. - paz.poses: Pose6D, PnP solver constants, solve_PnP, solve_PnP_RANSAC, rotation_vector_to_matrix, solve_pose_matrix_RANSAC, project_points3D (next to project_to_image and the existing pose helpers). - paz.points2D: normalize/denormalize_keypoints2D, rotate_point2D, rotate_keypoints2D, flip_keypoints_left_right, translate_keypoints, uv_to_vu. transform_keypoint is dropped in favor of the existing points2D.transform. - paz.pinhole: build_cube_points3D (beside build_cube_corners). - Callers (pose_estimators, pix2pose example, probabilistic keypoints example) import from these backends; paz.keypoints and its module are removed. Tests move alongside: PnP recovery and project consistency in poses_test, 2D transforms in points2D_test, cube corners in pinhole_test.
|
Follow-up
|
What
Consolidates the PnP and keypoint geometry that was scattered across the
applications layer and the pix2pose example into a single backend module
paz.keypoints, and ports the missing 2D keypoint-transform helpers frommaster.
Before this, backend geometry (
solve_PnP,solve_PnP_RANSAC,project_points3D,build_cube_corners,Pose6D, the solver constants) livedinside
paz/applications/pose_estimators.py, and the pix2pose example carriedits own
solve_pose_from_nocswith an inline Rodrigues conversion. There was nopaz/backend/keypoints.pyat all, and the keypoint-transform layer thatmaster's keypoint pipelines are built from was absent.
Changes
paz/backend/keypoints.py(exposed aspaz.keypoints):applications/pose_estimators.py:Pose6D, solverconstants,
solve_PnP,solve_PnP_RANSAC,project_points3D,build_cube_points3D.normalize_keypoints2D,denormalize_keypoints2D,rotate_point2D,transform_keypoint,flip_keypoints_left_right,translate_keypoints,uv_to_vu. cv2 is reserved for the actual PnP solve only.solve_pose_matrix_RANSAC: the reusable dense-correspondence → 6D-posekernel lifted from the pix2pose example (RANSAC PnP, then Rodrigues to a
rotation matrix).
rotation_vector_to_matrixhelper.applications/pose_estimators.pynow imports the geometry from thebackend; keeps only
HeadPoseKeypointNet2D32,build_face_points3D, anddraw_boxes3D.pipeline.py,demo.py,validate.py) imports frompaz.keypoints;solve_pose_from_nocsnow just builds correspondences andcalls
solve_pose_matrix_RANSAC.build_cube_points3D(master'sname) to avoid colliding with the pre-existing
pinhole.build_cube_corners, which takes min/max extents instead ofwidth/height/depth.
Why
the matching backend"), fixing the layering the pix2pose PR introduced.
keypoint_estimationtraining/augmentation pipeline next.
build_cube_cornersmeaning two differentthings).
Verification
pytest paz/backend/keypoints_test.py— 10 pass (JAX transforms vs numpyreference, normalize/denormalize round trip, PnP + RANSAC pose recovery to
1e-4, below-minimum returns None).
pytest paz/losses/pix2pose_test.py paz/callbacks/evaluate_pose_test.py paz/backend/evaluation_test.py— 15 pass (no regression in the pose stack).pyflakesclean on all changed files;paz.applicationsand the head-poseapp import unchanged;
pinhole.build_cube_cornersuntouched.