Skip to content

Add PnP and keypoint geometry across poses/points2D/pinhole - #393

Open
oarriaga wants to merge 4 commits into
paz-jaxfrom
port/backend-keypoints-geometry
Open

Add PnP and keypoint geometry across poses/points2D/pinhole#393
oarriaga wants to merge 4 commits into
paz-jaxfrom
port/backend-keypoints-geometry

Conversation

@oarriaga

@oarriaga oarriaga commented Jul 8, 2026

Copy link
Copy Markdown
Owner

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 from
master.

Before this, backend geometry (solve_PnP, solve_PnP_RANSAC,
project_points3D, build_cube_corners, Pose6D, the solver constants) lived
inside paz/applications/pose_estimators.py, and the pix2pose example carried
its own solve_pose_from_nocs with an inline Rodrigues conversion. There was no
paz/backend/keypoints.py at all, and the keypoint-transform layer that
master's keypoint pipelines are built from was absent.

Changes

  • New paz/backend/keypoints.py (exposed as paz.keypoints):
    • PnP toolkit moved out of applications/pose_estimators.py: Pose6D, solver
      constants, solve_PnP, solve_PnP_RANSAC, project_points3D,
      build_cube_points3D.
    • Missing keypoint transforms ported from master as JAX (pure array math):
      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-pose
      kernel lifted from the pix2pose example (RANSAC PnP, then Rodrigues to a
      rotation matrix). rotation_vector_to_matrix helper.
  • applications/pose_estimators.py now imports the geometry from the
    backend; keeps only HeadPoseKeypointNet2D32, build_face_points3D, and
    draw_boxes3D.
  • pix2pose example (pipeline.py, demo.py, validate.py) imports from
    paz.keypoints; solve_pose_from_nocs now just builds correspondences and
    calls solve_pose_matrix_RANSAC.
  • Renamed the dimension-based cube builder to build_cube_points3D (master's
    name) to avoid colliding with the pre-existing
    pinhole.build_cube_corners, which takes min/max extents instead of
    width/height/depth.

Why

  • Puts geometry in the matching backend per the agent guide ("new functions in
    the matching backend"), fixing the layering the pix2pose PR introduced.
  • Delivers the keypoint-transform layer needed to port the keypoint_estimation
    training/augmentation pipeline next.
  • Removes a confusing name collision (build_cube_corners meaning two different
    things).

Verification

  • pytest paz/backend/keypoints_test.py — 10 pass (JAX transforms vs numpy
    reference, 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).
  • pyflakes clean on all changed files; paz.applications and the head-pose
    app import unchanged; pinhole.build_cube_corners untouched.

oarriaga added 2 commits July 8, 2026 12:28
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.
@oarriaga

oarriaga commented Jul 8, 2026

Copy link
Copy Markdown
Owner Author

Follow-up commit 098a3683: added the keypoint-aware augmentation training pipeline functions (master's RandomKeypointRotation / RandomKeypointTranslation ported as JAX) and adopted them in an 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 — transform an image and its keypoints together so labels track the content.
  • examples/probabilistic_keypoint_estimation/train.pyKeypointSequence now does rotate + translate + brightness through these backend helpers (was image-brightness only) and normalizes per batch.

Verified with hot-pixel tests (keypoint lands on the transformed pixel) and a 6-seed end-to-end check; pytest paz/backend/keypoints_test.py paz/backend/image_test.py → 22 pass.

oarriaga added 2 commits July 8, 2026 13:52
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.
@oarriaga oarriaga changed the title Add backend/keypoints for PnP and keypoint geometry Add PnP and keypoint geometry across poses/points2D/pinhole Jul 8, 2026
@oarriaga

oarriaga commented Jul 8, 2026

Copy link
Copy Markdown
Owner Author

Follow-up 4e13dcc2: reworked the layout after review. Instead of a new paz.keypoints backend, the geometry is distributed into the backends that already own each math domain (the earlier monolith duplicated points2D.transform, points2D.denormalize, and poses.project_to_image):

  • paz.posesPose6D, PnP solver flags, solve_PnP, solve_PnP_RANSAC, rotation_vector_to_matrix, solve_pose_matrix_RANSAC, project_points3D.
  • paz.points2Dnormalize/denormalize_keypoints2D, rotate_point2D, rotate_keypoints2D, flip_keypoints_left_right, translate_keypoints, uv_to_vu (transform_keypoint dropped — the existing batched points2D.transform covers it).
  • paz.pinholebuild_cube_points3D (beside build_cube_corners).

paz.keypoints is removed. Tests moved to poses_test, points2D_test, pinhole_test. pytest over the affected suites → 29 pass; the probabilistic-keypoints example still tracks keypoints through augmentation (6/6 seeds).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant