From aa883a4eeac87a29a56d20f7db084ff7874e1fd4 Mon Sep 17 00:00:00 2001 From: Octavio Arriaga Date: Wed, 8 Jul 2026 12:28:08 +0200 Subject: [PATCH 1/4] Add backend/keypoints for PnP and keypoint geometry 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. --- examples/pix2pose/demo.py | 6 +- examples/pix2pose/pipeline.py | 12 +--- examples/pix2pose/validate.py | 6 +- paz/__init__.py | 1 + paz/applications/pose_estimators.py | 85 ++-------------------- paz/backend/keypoints.py | 107 ++++++++++++++++++++++++++++ paz/backend/keypoints_test.py | 103 ++++++++++++++++++++++++++ 7 files changed, 224 insertions(+), 96 deletions(-) create mode 100644 paz/backend/keypoints.py create mode 100644 paz/backend/keypoints_test.py diff --git a/examples/pix2pose/demo.py b/examples/pix2pose/demo.py index 8ead1444e..fc8cb56ef 100644 --- a/examples/pix2pose/demo.py +++ b/examples/pix2pose/demo.py @@ -9,8 +9,8 @@ import jax.numpy as jp import paz -from paz.applications.pose_estimators import solve_PnP_RANSAC -from paz.applications.pose_estimators import project_points3D, build_cube_corners +from paz.keypoints import solve_PnP_RANSAC +from paz.keypoints import project_points3D, build_cube_points3D import scenes Camera = namedtuple("Camera", ["intrinsics", "distortion"]) @@ -70,7 +70,7 @@ print(f"correspondences: {len(points3D)} | " f"mean reprojection error: {reprojection_error:.2f} px") -cube = paz.to_numpy(build_cube_corners(*extents)) +cube = paz.to_numpy(build_cube_points3D(*extents)) drawn = paz.applications.pose_estimators.draw_boxes3D( image.copy(), [pose6D], cube, camera, paz.draw.GREEN, 2, 3) paz.image.write(args.output, drawn) diff --git a/examples/pix2pose/pipeline.py b/examples/pix2pose/pipeline.py index 126b8427b..f69eb44c2 100644 --- a/examples/pix2pose/pipeline.py +++ b/examples/pix2pose/pipeline.py @@ -3,13 +3,11 @@ from functools import partial import numpy as np -import cv2 import jax import jax.numpy as jp import keras import paz -from paz.applications.pose_estimators import solve_PnP_RANSAC Camera = namedtuple("Camera", ["intrinsics", "distortion"]) @@ -81,11 +79,5 @@ def solve_pose_from_nocs(nocs, mask, extents, camera, max_points=1500, seed=0): return None points2D = np.stack([cols, rows], axis=1).astype("float64") points3D = extents * (nocs[rows, cols] - 0.5) - if len(points3D) > max_points: - choice = np.random.RandomState(seed).choice(len(points3D), max_points, False) # fmt: skip - points2D, points3D = points2D[choice], points3D[choice] - pose6D = solve_PnP_RANSAC(points2D, points3D, camera) - if pose6D is None: - return None - rotation = cv2.Rodrigues(pose6D.rotation_vector)[0] - return rotation, np.asarray(pose6D.translation).reshape(3) + args = (points2D, points3D, camera, max_points, seed) + return paz.keypoints.solve_pose_matrix_RANSAC(*args) diff --git a/examples/pix2pose/validate.py b/examples/pix2pose/validate.py index 792169164..b5d0ee306 100644 --- a/examples/pix2pose/validate.py +++ b/examples/pix2pose/validate.py @@ -9,8 +9,8 @@ import jax.numpy as jp import paz -from paz.applications.pose_estimators import solve_PnP_RANSAC -from paz.applications.pose_estimators import project_points3D, build_cube_corners +from paz.keypoints import solve_PnP_RANSAC +from paz.keypoints import project_points3D, build_cube_points3D from paz.applications.pose_estimators import draw_boxes3D import scenes import pipeline @@ -84,7 +84,7 @@ pose_errors.append(np.linalg.norm(projected - truth2D, axis=1).mean()) if arg < args.num_show and pose6D is not None: - cube = paz.to_numpy(build_cube_corners(*extents)) + cube = paz.to_numpy(build_cube_points3D(*extents)) overlay = draw_boxes3D(image.copy(), [pose6D], cube, camera, paz.draw.GREEN, 2, 3) # fmt: skip row = [image, (nocs_true * 255).astype("uint8"), (np.clip(nocs_pred, 0, 1) * 255).astype("uint8"), overlay] diff --git a/paz/__init__.py b/paz/__init__.py index 8bb9d68f4..d9e774649 100644 --- a/paz/__init__.py +++ b/paz/__init__.py @@ -27,6 +27,7 @@ from paz.backend import points2D from paz.backend import gaussian_mixture from paz.backend import poses +from paz.backend import keypoints from paz.backend import algebra from paz.backend import scene from paz.backend import plane diff --git a/paz/applications/pose_estimators.py b/paz/applications/pose_estimators.py index 7439525d0..0dd789394 100644 --- a/paz/applications/pose_estimators.py +++ b/paz/applications/pose_estimators.py @@ -1,62 +1,9 @@ -from collections import namedtuple -import cv2 import numpy as np -import jax.numpy as jp import paz - - -UPNP = cv2.SOLVEPNP_UPNP -LEVENBERG_MARQUARDT = cv2.SOLVEPNP_ITERATIVE -EPNP = cv2.SOLVEPNP_EPNP -MIN_REQUIRED_POINTS = 4 - -Pose6D = namedtuple("Pose6D", ["rotation_vector", "translation"]) - - -def build_cube_corners(width, height, depth): - """Build the 3D points of a cube in the openCV coordinate system: - 4--------1 - /| /| - / | / | - 3--------2 | - | 8_____|__5 - | / | / - |/ |/ - 7--------6 - - Z (depth) - / - /_____X (width) - | - | - Y (height) - - # Arguments - height: float, height of the 3D box. - width: float, width of the 3D box. - depth: float, width of the 3D box. - - # Returns - Numpy array of shape ``(8, 3)'' corresponding to 3D keypoints of a cube - """ - half_height, half_width, half_depth = height / 2.0, width / 2.0, depth / 2.0 - point1 = [+half_width, -half_height, +half_depth] - point2 = [+half_width, -half_height, -half_depth] - point3 = [-half_width, -half_height, -half_depth] - point4 = [-half_width, -half_height, +half_depth] - point5 = [+half_width, +half_height, +half_depth] - point6 = [+half_width, +half_height, -half_depth] - point7 = [-half_width, +half_height, -half_depth] - point8 = [-half_width, +half_height, +half_depth] - points = [point1, point2, point3, point4, point5, point6, point7, point8] - return jp.array(points) - - -def project_points3D(points3D, pose6D, camera): - args = (pose6D.translation, camera.intrinsics, camera.distortion) - points2D, _ = cv2.projectPoints(points3D, pose6D.rotation_vector, *args) - points2D = jp.squeeze(points2D, axis=1) # openCV shape (num_points, 1, 2) - return points2D +from paz.backend.keypoints import LEVENBERG_MARQUARDT +from paz.backend.keypoints import build_cube_points3D +from paz.backend.keypoints import project_points3D +from paz.backend.keypoints import solve_PnP def draw_boxes3D(image, poses, points3D, camera, color, thickness=5, radius=2): @@ -67,28 +14,6 @@ def draw_boxes3D(image, poses, points3D, camera, color, thickness=5, radius=2): return image -def solve_PnP(points2D, points3D, camera, solver=LEVENBERG_MARQUARDT): - points2D = np.array(points2D, np.float64).reshape((len(points3D), 1, 2)) - args = (camera.intrinsics, camera.distortion, None, None, False, solver) - (_, rotation_vector, translation) = cv2.solvePnP(points3D, points2D, *args) - return Pose6D(rotation_vector, translation) - - -def solve_PnP_RANSAC(points2D, points3D, camera, inlier_thresh=5.0, - iterations=100): - if len(points3D) < MIN_REQUIRED_POINTS: - return None - points2D = np.array(points2D, np.float64).reshape((len(points3D), 1, 2)) - points3D = np.array(points3D, np.float64) - args = (camera.intrinsics, camera.distortion, None, None, False, - iterations, inlier_thresh, 0.99, None, EPNP) - success, rotation, translation, inliers = cv2.solvePnPRansac( - points3D, points2D, *args) - if not success: - return None - return Pose6D(rotation, translation) - - def build_face_points3D(): points3D = np.array( [ @@ -122,7 +47,7 @@ def HeadPoseKeypointNet2D32(camera, box_scale=1.2, draw=None): solve_pose = paz.lock(solve_PnP, points3D, camera, LEVENBERG_MARQUARDT) if draw is None: - cube = paz.to_numpy(build_cube_corners(900, 1200, 800)) + cube = paz.to_numpy(build_cube_points3D(900, 1200, 800)) draw = paz.lock(draw_boxes3D, cube, camera, paz.draw.GREEN, 3, 5) def call(image): diff --git a/paz/backend/keypoints.py b/paz/backend/keypoints.py new file mode 100644 index 000000000..3b4671764 --- /dev/null +++ b/paz/backend/keypoints.py @@ -0,0 +1,107 @@ +from collections import namedtuple + +import cv2 +import numpy as np +import jax.numpy as jp + + +UPNP = cv2.SOLVEPNP_UPNP +LEVENBERG_MARQUARDT = cv2.SOLVEPNP_ITERATIVE +EPNP = cv2.SOLVEPNP_EPNP +MIN_REQUIRED_POINTS = 4 + +Pose6D = namedtuple("Pose6D", ["rotation_vector", "translation"]) + + +def build_cube_points3D(width, height, depth): + half_width, half_height, half_depth = width / 2, height / 2, depth / 2 + point1 = [+half_width, -half_height, +half_depth] + point2 = [+half_width, -half_height, -half_depth] + point3 = [-half_width, -half_height, -half_depth] + point4 = [-half_width, -half_height, +half_depth] + point5 = [+half_width, +half_height, +half_depth] + point6 = [+half_width, +half_height, -half_depth] + point7 = [-half_width, +half_height, -half_depth] + point8 = [-half_width, +half_height, +half_depth] + points = [point1, point2, point3, point4, point5, point6, point7, point8] + return jp.array(points) + + +def normalize_keypoints2D(points2D, height, width): + image_shape = jp.array([width, height]) + return (2.0 * points2D / image_shape) - 1.0 + + +def denormalize_keypoints2D(points2D, height, width): + image_shape = jp.array([width, height]) + return (points2D + 1.0) / 2.0 * image_shape + + +def rotate_point2D(point2D, rotation_angle): + angle = jp.pi * rotation_angle / 180.0 + sin_angle, cos_angle = jp.sin(angle), jp.cos(angle) + x = point2D[0] * cos_angle - point2D[1] * sin_angle + y = point2D[0] * sin_angle + point2D[1] * cos_angle + return jp.array([x, y]) + + +def transform_keypoint(keypoint, transform): + point = jp.array([keypoint[0], keypoint[1], 1.0]) + return transform @ point + + +def flip_keypoints_left_right(keypoints, width): + x, y = jp.split(keypoints, 2, axis=1) + return jp.concatenate([width - x, y], axis=1) + + +def translate_keypoints(keypoints, translation): + return keypoints + translation + + +def uv_to_vu(keypoints): + return keypoints[:, ::-1] + + +def project_points3D(points3D, pose6D, camera): + args = (pose6D.translation, camera.intrinsics, camera.distortion) + points2D, _ = cv2.projectPoints(points3D, pose6D.rotation_vector, *args) + return jp.squeeze(points2D, axis=1) # openCV shape (num_points, 1, 2) + + +def solve_PnP(points2D, points3D, camera, solver=LEVENBERG_MARQUARDT): + points2D = np.array(points2D, np.float64).reshape((len(points3D), 1, 2)) + args = (camera.intrinsics, camera.distortion, None, None, False, solver) + (_, rotation_vector, translation) = cv2.solvePnP(points3D, points2D, *args) + return Pose6D(rotation_vector, translation) + + +def solve_PnP_RANSAC(points2D, points3D, camera, inlier_thresh=5.0, + iterations=100): + if len(points3D) < MIN_REQUIRED_POINTS: + return None + points2D = np.array(points2D, np.float64).reshape((len(points3D), 1, 2)) + points3D = np.array(points3D, np.float64) + args = (camera.intrinsics, camera.distortion, None, None, False, + iterations, inlier_thresh, 0.99, None, EPNP) + success, rotation, translation, inliers = cv2.solvePnPRansac( + points3D, points2D, *args) + if not success: + return None + return Pose6D(rotation, translation) + + +def rotation_vector_to_matrix(rotation_vector): + return cv2.Rodrigues(rotation_vector)[0] + + +def solve_pose_matrix_RANSAC(points2D, points3D, camera, max_points=1500, + seed=0): + if len(points3D) > max_points: + choice = np.random.RandomState(seed).choice(len(points3D), max_points, False) # fmt: skip + points2D, points3D = points2D[choice], points3D[choice] + pose6D = solve_PnP_RANSAC(points2D, points3D, camera) + if pose6D is None: + return None + rotation = rotation_vector_to_matrix(pose6D.rotation_vector) + return rotation, np.asarray(pose6D.translation).reshape(3) diff --git a/paz/backend/keypoints_test.py b/paz/backend/keypoints_test.py new file mode 100644 index 000000000..49e303828 --- /dev/null +++ b/paz/backend/keypoints_test.py @@ -0,0 +1,103 @@ +from collections import namedtuple + +import numpy as np +import jax.numpy as jp +import cv2 + +from paz.backend import keypoints +from paz.backend.poses import project_to_image + +Camera = namedtuple("Camera", ["intrinsics", "distortion"]) + + +def build_camera(size=128, focal=150.0): + center = size / 2.0 + intrinsics = np.array([[focal, 0, center], + [0, focal, center], + [0, 0, 1.0]]) + return Camera(intrinsics, np.zeros((4, 1))) + + +def normalize_reference(points2D, height, width): + image_shape = np.array([width, height]) + return 2.0 * (points2D / image_shape) - 1.0 + + +def test_normalize_keypoints2D_matches_numpy_reference(): + points2D = np.array([[0.0, 0.0], [128.0, 64.0], [32.0, 96.0]]) + result = np.asarray(keypoints.normalize_keypoints2D(points2D, 128, 128)) + reference = normalize_reference(points2D, 128, 128) + assert np.allclose(result, reference) + + +def test_normalize_denormalize_round_trip(): + points2D = jp.array([[10.0, 20.0], [50.0, 5.0], [127.0, 63.0]]) + normalized = keypoints.normalize_keypoints2D(points2D, 128, 64) + recovered = keypoints.denormalize_keypoints2D(normalized, 128, 64) + assert np.allclose(np.asarray(recovered), np.asarray(points2D)) + + +def test_rotate_point2D_ninety_degrees(): + rotated = keypoints.rotate_point2D(jp.array([1.0, 0.0]), 90.0) + assert np.allclose(np.asarray(rotated), [0.0, 1.0], atol=1e-6) + + +def test_flip_keypoints_left_right(): + points = jp.array([[0.0, 5.0], [32.0, 10.0]]) + flipped = np.asarray(keypoints.flip_keypoints_left_right(points, 32.0)) + assert np.allclose(flipped, [[32.0, 5.0], [0.0, 10.0]]) + + +def test_transform_keypoint_translation(): + transform = jp.array([[1.0, 0.0, 3.0], [0.0, 1.0, -2.0], [0, 0, 1.0]]) + moved = keypoints.transform_keypoint(jp.array([4.0, 5.0]), transform) + assert np.allclose(np.asarray(moved)[:2], [7.0, 3.0]) + + +def test_uv_to_vu(): + flipped = keypoints.uv_to_vu(jp.array([[1.0, 2.0], [3.0, 4.0]])) + assert np.allclose(np.asarray(flipped), [[2.0, 1.0], [4.0, 3.0]]) + + +def test_build_cube_points3D_shape_and_center(): + cube = np.asarray(keypoints.build_cube_points3D(2.0, 4.0, 6.0)) + assert cube.shape == (8, 3) + assert np.allclose(cube.mean(axis=0), [0.0, 0.0, 0.0]) + assert np.allclose(np.abs(cube).max(axis=0), [1.0, 2.0, 3.0]) + + +def test_solve_PnP_recovers_known_pose(): + camera = build_camera() + points3D = keypoints.build_cube_points3D(0.1, 0.1, 0.1) + points3D = np.asarray(points3D, np.float64) + rotation = cv2.Rodrigues(np.array([0.2, -0.1, 0.3]))[0] + translation = np.array([0.02, -0.01, 0.5]) + points2D = project_to_image(rotation, translation, points3D, + camera.intrinsics) + pose6D = keypoints.solve_PnP(points2D, points3D, camera) + recovered = keypoints.rotation_vector_to_matrix(pose6D.rotation_vector) + assert np.allclose(recovered, rotation, atol=1e-4) + assert np.allclose(np.asarray(pose6D.translation).reshape(3), + translation, atol=1e-4) + + +def test_solve_pose_matrix_RANSAC_recovers_known_pose(): + camera = build_camera() + grid = np.linspace(-0.05, 0.05, 5) + points3D = np.array([[x, y, z] for x in grid for y in grid for z in grid]) + rotation = cv2.Rodrigues(np.array([0.1, 0.2, -0.15]))[0] + translation = np.array([0.0, 0.0, 0.6]) + points2D = project_to_image(rotation, translation, points3D, + camera.intrinsics) + result = keypoints.solve_pose_matrix_RANSAC(points2D, points3D, camera) + assert result is not None + recovered_rotation, recovered_translation = result + assert np.allclose(recovered_rotation, rotation, atol=1e-3) + assert np.allclose(recovered_translation, translation, atol=1e-3) + + +def test_solve_PnP_RANSAC_returns_none_below_minimum(): + camera = build_camera() + points2D = np.zeros((3, 2)) + points3D = np.zeros((3, 3)) + assert keypoints.solve_PnP_RANSAC(points2D, points3D, camera) is None From 098a368300e5504d0dd87a219e5f351d8846b63f Mon Sep 17 00:00:00 2001 From: Octavio Arriaga Date: Wed, 8 Jul 2026 13:39:45 +0200 Subject: [PATCH 2/4] Add keypoint-aware augmentation and adopt it in the keypoints example 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. --- .../train.py | 37 +++++++++++++----- paz/backend/image.py | 7 ++++ paz/backend/keypoints.py | 31 +++++++++++++++ paz/backend/keypoints_test.py | 39 +++++++++++++++++++ 4 files changed, 104 insertions(+), 10 deletions(-) diff --git a/examples/probabilistic_keypoint_estimation/train.py b/examples/probabilistic_keypoint_estimation/train.py index b8205866e..e14eb7075 100644 --- a/examples/probabilistic_keypoint_estimation/train.py +++ b/examples/probabilistic_keypoint_estimation/train.py @@ -3,6 +3,7 @@ os.environ["KERAS_BACKEND"] = "jax" import argparse +from functools import partial import numpy as np import jax @@ -13,29 +14,45 @@ import facial_keypoints +def augment_image_and_keypoints(key, image, keypoints, rotation_range, + delta_scale): + rotate_key, translate_key, bright_key = jax.random.split(key, 3) + image, keypoints = paz.keypoints.rotate_image_and_keypoints( + rotate_key, image, keypoints, rotation_range) + image, keypoints = paz.keypoints.translate_image_and_keypoints( + translate_key, image, keypoints, delta_scale) + image = paz.image.random_brightness(bright_key, image) + return image, keypoints + + class KeypointSequence(keras.utils.PyDataset): - def __init__(self, images, keypoints, batch_size, augment=False, seed=0): + def __init__(self, images, keypoints, batch_size, augment=False, + rotation_range=jp.pi / 12, delta_scale=(0.1, 0.1), seed=0): super().__init__() self.images = np.asarray(images, "uint8") - self.keypoints = paz.gaussian_mixture.normalize_points( - np.asarray(keypoints, "float32"), 96, 96) + self.keypoints = np.asarray(keypoints, "float32") self.batch_size = batch_size self.augment = augment self.key = jax.random.PRNGKey(seed) - self.augment_images = jax.jit(jax.vmap(paz.image.random_brightness)) + augment_one = partial(augment_image_and_keypoints, + rotation_range=rotation_range, + delta_scale=delta_scale) + self.augment_batch = jax.jit(jax.vmap(augment_one)) def __len__(self): return len(self.images) // self.batch_size def __getitem__(self, index): chunk = slice(index * self.batch_size, (index + 1) * self.batch_size) - images = self.images[chunk][..., None] + images = jp.asarray(self.images[chunk][..., None], "float32") + keypoints = jp.asarray(self.keypoints[chunk]) if self.augment: - key = jax.random.fold_in(self.key, index) - keys = jax.random.split(key, len(images)) - images = self.augment_images(keys, jp.asarray(images)) - images = np.asarray(images, "float32") / 255.0 - return images, np.asarray(self.keypoints[chunk]) + self.key, batch_key = jax.random.split(self.key) + keys = jax.random.split(batch_key, len(images)) + images, keypoints = self.augment_batch(keys, images, keypoints) + images = np.asarray(images) / 255.0 + keypoints = paz.gaussian_mixture.normalize_points(keypoints, 96, 96) + return images, np.asarray(keypoints) if __name__ == "__main__": diff --git a/paz/backend/image.py b/paz/backend/image.py index 3ac772ace..aaf410893 100644 --- a/paz/backend/image.py +++ b/paz/backend/image.py @@ -496,6 +496,13 @@ def random_rotation( return rotate(image, angle, order, mode, cval) +def translate_image(image, translation, order=1, mode="nearest", cval=0.0): + """Translates image content by a ``(x, y)`` pixel shift.""" + offset = jp.array([-translation[1], -translation[0], 0.0]) + matrix = paz.SE3.to_affine_matrix(jp.eye(3), offset) + return affine_transform(image, matrix, order=order, mode=mode, cval=cval) + + def random_flip_left_right(key, image): do_flip = jax.random.bernoulli(key) return jax.lax.cond(do_flip, flip_left_right, lambda x: x, image) diff --git a/paz/backend/keypoints.py b/paz/backend/keypoints.py index 3b4671764..6d694cd7b 100644 --- a/paz/backend/keypoints.py +++ b/paz/backend/keypoints.py @@ -2,8 +2,11 @@ import cv2 import numpy as np +import jax import jax.numpy as jp +import paz + UPNP = cv2.SOLVEPNP_UPNP LEVENBERG_MARQUARDT = cv2.SOLVEPNP_ITERATIVE @@ -63,6 +66,34 @@ def uv_to_vu(keypoints): return keypoints[:, ::-1] +def rotate_keypoints2D(keypoints, angle, center): + cos_angle, sin_angle = jp.cos(angle), jp.sin(angle) + rotation = jp.array([[cos_angle, -sin_angle], [sin_angle, cos_angle]]) + return (keypoints - center) @ rotation.T + center + + +def image_center2D(image): + height, width = image.shape[0], image.shape[1] + return jp.array([(width - 1) / 2.0, (height - 1) / 2.0]) + + +def rotate_image_and_keypoints(key, image, keypoints, rotation_range): + angle = jax.random.uniform(key, (), minval=-rotation_range, + maxval=rotation_range) + rotated_image = paz.image.rotate(image, angle) + rotated_keypoints = rotate_keypoints2D(keypoints, angle, image_center2D(image)) # fmt: skip + return rotated_image, rotated_keypoints + + +def translate_image_and_keypoints(key, image, keypoints, delta_scale): + height, width = image.shape[0], image.shape[1] + scale = jp.array([width * delta_scale[0], height * delta_scale[1]]) + translation = jax.random.uniform(key, (2,), minval=-scale, maxval=scale) + translated_image = paz.image.translate_image(image, translation) + translated_keypoints = translate_keypoints(keypoints, translation) + return translated_image, translated_keypoints + + def project_points3D(points3D, pose6D, camera): args = (pose6D.translation, camera.intrinsics, camera.distortion) points2D, _ = cv2.projectPoints(points3D, pose6D.rotation_vector, *args) diff --git a/paz/backend/keypoints_test.py b/paz/backend/keypoints_test.py index 49e303828..18dd7016c 100644 --- a/paz/backend/keypoints_test.py +++ b/paz/backend/keypoints_test.py @@ -1,9 +1,11 @@ from collections import namedtuple import numpy as np +import jax import jax.numpy as jp import cv2 +import paz from paz.backend import keypoints from paz.backend.poses import project_to_image @@ -101,3 +103,40 @@ def test_solve_PnP_RANSAC_returns_none_below_minimum(): points2D = np.zeros((3, 2)) points3D = np.zeros((3, 3)) assert keypoints.solve_PnP_RANSAC(points2D, points3D, camera) is None + + +def hot_pixel_location(image): + row, col = np.unravel_index(np.argmax(image[..., 0]), image.shape[:2]) + return col, row + + +def test_rotate_keypoints2D_tracks_image_rotation(): + image = np.zeros((31, 31, 3), "float32") + image[5, 20] = 1.0 + rotated = paz.image.rotate(jp.asarray(image), 0.5) + col, row = hot_pixel_location(np.asarray(rotated)) + center = keypoints.image_center2D(jp.asarray(image)) + keypoint = jp.array([[20.0, 5.0]]) # (x=col, y=row) + moved = np.asarray(keypoints.rotate_keypoints2D(keypoint, 0.5, center))[0] + assert abs(moved[0] - col) <= 1.5 and abs(moved[1] - row) <= 1.5 + + +def test_translate_image_and_keypoints_track(): + image = np.zeros((31, 31, 3), "float32") + image[10, 8] = 1.0 + translation = jp.array([4.0, -3.0]) # (x, y) shift + translated = paz.image.translate_image(jp.asarray(image), translation) + col, row = hot_pixel_location(np.asarray(translated)) + keypoint = jp.array([[8.0, 10.0]]) + moved = np.asarray(keypoints.translate_keypoints(keypoint, translation))[0] + assert (moved[0], moved[1]) == (col, row) + + +def test_rotate_image_and_keypoints_preserves_shapes(): + image = jp.zeros((96, 96, 1)) + points = jp.array([[10.0, 20.0], [50.0, 40.0]]) + key = jax.random.PRNGKey(0) + out_image, out_points = keypoints.rotate_image_and_keypoints( + key, image, points, jp.pi / 12) + assert out_image.shape == image.shape + assert out_points.shape == points.shape From 020c5f33f1eed3d9d22e3b9441a0a598930e0c37 Mon Sep 17 00:00:00 2001 From: Octavio Arriaga Date: Wed, 8 Jul 2026 13:52:04 +0200 Subject: [PATCH 3/4] Rename image.translate and move keypoint augmenters to the example 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. --- .../train.py | 25 ++++++++++++++++--- paz/backend/image.py | 2 +- paz/backend/keypoints.py | 25 ------------------- paz/backend/keypoints_test.py | 17 +++---------- 4 files changed, 25 insertions(+), 44 deletions(-) diff --git a/examples/probabilistic_keypoint_estimation/train.py b/examples/probabilistic_keypoint_estimation/train.py index e14eb7075..b4a2b5658 100644 --- a/examples/probabilistic_keypoint_estimation/train.py +++ b/examples/probabilistic_keypoint_estimation/train.py @@ -14,13 +14,30 @@ import facial_keypoints +def rotate_image_and_keypoints(key, image, keypoints, rotation_range): + angle = jax.random.uniform(key, (), minval=-rotation_range, + maxval=rotation_range) + height, width = image.shape[0], image.shape[1] + center = jp.array([(width - 1) / 2.0, (height - 1) / 2.0]) + image = paz.image.rotate(image, angle) + return image, paz.keypoints.rotate_keypoints2D(keypoints, angle, center) + + +def translate_image_and_keypoints(key, image, keypoints, delta_scale): + height, width = image.shape[0], image.shape[1] + scale = jp.array([width * delta_scale[0], height * delta_scale[1]]) + shift = jax.random.uniform(key, (2,), minval=-scale, maxval=scale) + image = paz.image.translate(image, shift) + return image, paz.keypoints.translate_keypoints(keypoints, shift) + + def augment_image_and_keypoints(key, image, keypoints, rotation_range, delta_scale): rotate_key, translate_key, bright_key = jax.random.split(key, 3) - image, keypoints = paz.keypoints.rotate_image_and_keypoints( - rotate_key, image, keypoints, rotation_range) - image, keypoints = paz.keypoints.translate_image_and_keypoints( - translate_key, image, keypoints, delta_scale) + args = (rotate_key, image, keypoints, rotation_range) + image, keypoints = rotate_image_and_keypoints(*args) + args = (translate_key, image, keypoints, delta_scale) + image, keypoints = translate_image_and_keypoints(*args) image = paz.image.random_brightness(bright_key, image) return image, keypoints diff --git a/paz/backend/image.py b/paz/backend/image.py index aaf410893..58fe59da6 100644 --- a/paz/backend/image.py +++ b/paz/backend/image.py @@ -496,7 +496,7 @@ def random_rotation( return rotate(image, angle, order, mode, cval) -def translate_image(image, translation, order=1, mode="nearest", cval=0.0): +def translate(image, translation, order=1, mode="nearest", cval=0.0): """Translates image content by a ``(x, y)`` pixel shift.""" offset = jp.array([-translation[1], -translation[0], 0.0]) matrix = paz.SE3.to_affine_matrix(jp.eye(3), offset) diff --git a/paz/backend/keypoints.py b/paz/backend/keypoints.py index 6d694cd7b..ca9cbef9f 100644 --- a/paz/backend/keypoints.py +++ b/paz/backend/keypoints.py @@ -2,11 +2,8 @@ import cv2 import numpy as np -import jax import jax.numpy as jp -import paz - UPNP = cv2.SOLVEPNP_UPNP LEVENBERG_MARQUARDT = cv2.SOLVEPNP_ITERATIVE @@ -72,28 +69,6 @@ def rotate_keypoints2D(keypoints, angle, center): return (keypoints - center) @ rotation.T + center -def image_center2D(image): - height, width = image.shape[0], image.shape[1] - return jp.array([(width - 1) / 2.0, (height - 1) / 2.0]) - - -def rotate_image_and_keypoints(key, image, keypoints, rotation_range): - angle = jax.random.uniform(key, (), minval=-rotation_range, - maxval=rotation_range) - rotated_image = paz.image.rotate(image, angle) - rotated_keypoints = rotate_keypoints2D(keypoints, angle, image_center2D(image)) # fmt: skip - return rotated_image, rotated_keypoints - - -def translate_image_and_keypoints(key, image, keypoints, delta_scale): - height, width = image.shape[0], image.shape[1] - scale = jp.array([width * delta_scale[0], height * delta_scale[1]]) - translation = jax.random.uniform(key, (2,), minval=-scale, maxval=scale) - translated_image = paz.image.translate_image(image, translation) - translated_keypoints = translate_keypoints(keypoints, translation) - return translated_image, translated_keypoints - - def project_points3D(points3D, pose6D, camera): args = (pose6D.translation, camera.intrinsics, camera.distortion) points2D, _ = cv2.projectPoints(points3D, pose6D.rotation_vector, *args) diff --git a/paz/backend/keypoints_test.py b/paz/backend/keypoints_test.py index 18dd7016c..c00b1d795 100644 --- a/paz/backend/keypoints_test.py +++ b/paz/backend/keypoints_test.py @@ -1,7 +1,6 @@ from collections import namedtuple import numpy as np -import jax import jax.numpy as jp import cv2 @@ -115,28 +114,18 @@ def test_rotate_keypoints2D_tracks_image_rotation(): image[5, 20] = 1.0 rotated = paz.image.rotate(jp.asarray(image), 0.5) col, row = hot_pixel_location(np.asarray(rotated)) - center = keypoints.image_center2D(jp.asarray(image)) + center = jp.array([(31 - 1) / 2.0, (31 - 1) / 2.0]) keypoint = jp.array([[20.0, 5.0]]) # (x=col, y=row) moved = np.asarray(keypoints.rotate_keypoints2D(keypoint, 0.5, center))[0] assert abs(moved[0] - col) <= 1.5 and abs(moved[1] - row) <= 1.5 -def test_translate_image_and_keypoints_track(): +def test_translate_keypoints_tracks_image_translation(): image = np.zeros((31, 31, 3), "float32") image[10, 8] = 1.0 translation = jp.array([4.0, -3.0]) # (x, y) shift - translated = paz.image.translate_image(jp.asarray(image), translation) + translated = paz.image.translate(jp.asarray(image), translation) col, row = hot_pixel_location(np.asarray(translated)) keypoint = jp.array([[8.0, 10.0]]) moved = np.asarray(keypoints.translate_keypoints(keypoint, translation))[0] assert (moved[0], moved[1]) == (col, row) - - -def test_rotate_image_and_keypoints_preserves_shapes(): - image = jp.zeros((96, 96, 1)) - points = jp.array([[10.0, 20.0], [50.0, 40.0]]) - key = jax.random.PRNGKey(0) - out_image, out_points = keypoints.rotate_image_and_keypoints( - key, image, points, jp.pi / 12) - assert out_image.shape == image.shape - assert out_points.shape == points.shape From 4e13dcc2e7f3bee5df9e60127000212055e4b94e Mon Sep 17 00:00:00 2001 From: Octavio Arriaga Date: Wed, 8 Jul 2026 16:06:20 +0200 Subject: [PATCH 4/4] Distribute keypoint geometry into existing backends; drop paz.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. --- examples/pix2pose/demo.py | 4 +- examples/pix2pose/pipeline.py | 2 +- examples/pix2pose/validate.py | 4 +- .../train.py | 4 +- paz/__init__.py | 1 - paz/applications/pose_estimators.py | 8 +- paz/backend/keypoints.py | 113 --------------- paz/backend/keypoints_test.py | 131 ------------------ paz/backend/pinhole.py | 14 ++ paz/backend/pinhole_test.py | 14 ++ paz/backend/points2D.py | 37 +++++ paz/backend/points2D_test.py | 70 ++++++++++ paz/backend/poses.py | 54 ++++++++ paz/backend/poses_test.py | 61 ++++++++ 14 files changed, 261 insertions(+), 256 deletions(-) delete mode 100644 paz/backend/keypoints.py delete mode 100644 paz/backend/keypoints_test.py create mode 100644 paz/backend/pinhole_test.py create mode 100644 paz/backend/points2D_test.py diff --git a/examples/pix2pose/demo.py b/examples/pix2pose/demo.py index fc8cb56ef..85157e99b 100644 --- a/examples/pix2pose/demo.py +++ b/examples/pix2pose/demo.py @@ -9,8 +9,8 @@ import jax.numpy as jp import paz -from paz.keypoints import solve_PnP_RANSAC -from paz.keypoints import project_points3D, build_cube_points3D +from paz.poses import solve_PnP_RANSAC, project_points3D +from paz.pinhole import build_cube_points3D import scenes Camera = namedtuple("Camera", ["intrinsics", "distortion"]) diff --git a/examples/pix2pose/pipeline.py b/examples/pix2pose/pipeline.py index f69eb44c2..8e066ed49 100644 --- a/examples/pix2pose/pipeline.py +++ b/examples/pix2pose/pipeline.py @@ -80,4 +80,4 @@ def solve_pose_from_nocs(nocs, mask, extents, camera, max_points=1500, seed=0): points2D = np.stack([cols, rows], axis=1).astype("float64") points3D = extents * (nocs[rows, cols] - 0.5) args = (points2D, points3D, camera, max_points, seed) - return paz.keypoints.solve_pose_matrix_RANSAC(*args) + return paz.poses.solve_pose_matrix_RANSAC(*args) diff --git a/examples/pix2pose/validate.py b/examples/pix2pose/validate.py index b5d0ee306..d24450e10 100644 --- a/examples/pix2pose/validate.py +++ b/examples/pix2pose/validate.py @@ -9,8 +9,8 @@ import jax.numpy as jp import paz -from paz.keypoints import solve_PnP_RANSAC -from paz.keypoints import project_points3D, build_cube_points3D +from paz.poses import solve_PnP_RANSAC, project_points3D +from paz.pinhole import build_cube_points3D from paz.applications.pose_estimators import draw_boxes3D import scenes import pipeline diff --git a/examples/probabilistic_keypoint_estimation/train.py b/examples/probabilistic_keypoint_estimation/train.py index b4a2b5658..2699b4e7f 100644 --- a/examples/probabilistic_keypoint_estimation/train.py +++ b/examples/probabilistic_keypoint_estimation/train.py @@ -20,7 +20,7 @@ def rotate_image_and_keypoints(key, image, keypoints, rotation_range): height, width = image.shape[0], image.shape[1] center = jp.array([(width - 1) / 2.0, (height - 1) / 2.0]) image = paz.image.rotate(image, angle) - return image, paz.keypoints.rotate_keypoints2D(keypoints, angle, center) + return image, paz.points2D.rotate_keypoints2D(keypoints, angle, center) def translate_image_and_keypoints(key, image, keypoints, delta_scale): @@ -28,7 +28,7 @@ def translate_image_and_keypoints(key, image, keypoints, delta_scale): scale = jp.array([width * delta_scale[0], height * delta_scale[1]]) shift = jax.random.uniform(key, (2,), minval=-scale, maxval=scale) image = paz.image.translate(image, shift) - return image, paz.keypoints.translate_keypoints(keypoints, shift) + return image, paz.points2D.translate_keypoints(keypoints, shift) def augment_image_and_keypoints(key, image, keypoints, rotation_range, diff --git a/paz/__init__.py b/paz/__init__.py index d9e774649..8bb9d68f4 100644 --- a/paz/__init__.py +++ b/paz/__init__.py @@ -27,7 +27,6 @@ from paz.backend import points2D from paz.backend import gaussian_mixture from paz.backend import poses -from paz.backend import keypoints from paz.backend import algebra from paz.backend import scene from paz.backend import plane diff --git a/paz/applications/pose_estimators.py b/paz/applications/pose_estimators.py index 0dd789394..08a9f0e93 100644 --- a/paz/applications/pose_estimators.py +++ b/paz/applications/pose_estimators.py @@ -1,9 +1,9 @@ import numpy as np import paz -from paz.backend.keypoints import LEVENBERG_MARQUARDT -from paz.backend.keypoints import build_cube_points3D -from paz.backend.keypoints import project_points3D -from paz.backend.keypoints import solve_PnP +from paz.backend.poses import LEVENBERG_MARQUARDT +from paz.backend.poses import project_points3D +from paz.backend.poses import solve_PnP +from paz.backend.pinhole import build_cube_points3D def draw_boxes3D(image, poses, points3D, camera, color, thickness=5, radius=2): diff --git a/paz/backend/keypoints.py b/paz/backend/keypoints.py deleted file mode 100644 index ca9cbef9f..000000000 --- a/paz/backend/keypoints.py +++ /dev/null @@ -1,113 +0,0 @@ -from collections import namedtuple - -import cv2 -import numpy as np -import jax.numpy as jp - - -UPNP = cv2.SOLVEPNP_UPNP -LEVENBERG_MARQUARDT = cv2.SOLVEPNP_ITERATIVE -EPNP = cv2.SOLVEPNP_EPNP -MIN_REQUIRED_POINTS = 4 - -Pose6D = namedtuple("Pose6D", ["rotation_vector", "translation"]) - - -def build_cube_points3D(width, height, depth): - half_width, half_height, half_depth = width / 2, height / 2, depth / 2 - point1 = [+half_width, -half_height, +half_depth] - point2 = [+half_width, -half_height, -half_depth] - point3 = [-half_width, -half_height, -half_depth] - point4 = [-half_width, -half_height, +half_depth] - point5 = [+half_width, +half_height, +half_depth] - point6 = [+half_width, +half_height, -half_depth] - point7 = [-half_width, +half_height, -half_depth] - point8 = [-half_width, +half_height, +half_depth] - points = [point1, point2, point3, point4, point5, point6, point7, point8] - return jp.array(points) - - -def normalize_keypoints2D(points2D, height, width): - image_shape = jp.array([width, height]) - return (2.0 * points2D / image_shape) - 1.0 - - -def denormalize_keypoints2D(points2D, height, width): - image_shape = jp.array([width, height]) - return (points2D + 1.0) / 2.0 * image_shape - - -def rotate_point2D(point2D, rotation_angle): - angle = jp.pi * rotation_angle / 180.0 - sin_angle, cos_angle = jp.sin(angle), jp.cos(angle) - x = point2D[0] * cos_angle - point2D[1] * sin_angle - y = point2D[0] * sin_angle + point2D[1] * cos_angle - return jp.array([x, y]) - - -def transform_keypoint(keypoint, transform): - point = jp.array([keypoint[0], keypoint[1], 1.0]) - return transform @ point - - -def flip_keypoints_left_right(keypoints, width): - x, y = jp.split(keypoints, 2, axis=1) - return jp.concatenate([width - x, y], axis=1) - - -def translate_keypoints(keypoints, translation): - return keypoints + translation - - -def uv_to_vu(keypoints): - return keypoints[:, ::-1] - - -def rotate_keypoints2D(keypoints, angle, center): - cos_angle, sin_angle = jp.cos(angle), jp.sin(angle) - rotation = jp.array([[cos_angle, -sin_angle], [sin_angle, cos_angle]]) - return (keypoints - center) @ rotation.T + center - - -def project_points3D(points3D, pose6D, camera): - args = (pose6D.translation, camera.intrinsics, camera.distortion) - points2D, _ = cv2.projectPoints(points3D, pose6D.rotation_vector, *args) - return jp.squeeze(points2D, axis=1) # openCV shape (num_points, 1, 2) - - -def solve_PnP(points2D, points3D, camera, solver=LEVENBERG_MARQUARDT): - points2D = np.array(points2D, np.float64).reshape((len(points3D), 1, 2)) - args = (camera.intrinsics, camera.distortion, None, None, False, solver) - (_, rotation_vector, translation) = cv2.solvePnP(points3D, points2D, *args) - return Pose6D(rotation_vector, translation) - - -def solve_PnP_RANSAC(points2D, points3D, camera, inlier_thresh=5.0, - iterations=100): - if len(points3D) < MIN_REQUIRED_POINTS: - return None - points2D = np.array(points2D, np.float64).reshape((len(points3D), 1, 2)) - points3D = np.array(points3D, np.float64) - args = (camera.intrinsics, camera.distortion, None, None, False, - iterations, inlier_thresh, 0.99, None, EPNP) - success, rotation, translation, inliers = cv2.solvePnPRansac( - points3D, points2D, *args) - if not success: - return None - return Pose6D(rotation, translation) - - -def rotation_vector_to_matrix(rotation_vector): - return cv2.Rodrigues(rotation_vector)[0] - - -def solve_pose_matrix_RANSAC(points2D, points3D, camera, max_points=1500, - seed=0): - if len(points3D) > max_points: - choice = np.random.RandomState(seed).choice(len(points3D), max_points, False) # fmt: skip - points2D, points3D = points2D[choice], points3D[choice] - pose6D = solve_PnP_RANSAC(points2D, points3D, camera) - if pose6D is None: - return None - rotation = rotation_vector_to_matrix(pose6D.rotation_vector) - return rotation, np.asarray(pose6D.translation).reshape(3) diff --git a/paz/backend/keypoints_test.py b/paz/backend/keypoints_test.py deleted file mode 100644 index c00b1d795..000000000 --- a/paz/backend/keypoints_test.py +++ /dev/null @@ -1,131 +0,0 @@ -from collections import namedtuple - -import numpy as np -import jax.numpy as jp -import cv2 - -import paz -from paz.backend import keypoints -from paz.backend.poses import project_to_image - -Camera = namedtuple("Camera", ["intrinsics", "distortion"]) - - -def build_camera(size=128, focal=150.0): - center = size / 2.0 - intrinsics = np.array([[focal, 0, center], - [0, focal, center], - [0, 0, 1.0]]) - return Camera(intrinsics, np.zeros((4, 1))) - - -def normalize_reference(points2D, height, width): - image_shape = np.array([width, height]) - return 2.0 * (points2D / image_shape) - 1.0 - - -def test_normalize_keypoints2D_matches_numpy_reference(): - points2D = np.array([[0.0, 0.0], [128.0, 64.0], [32.0, 96.0]]) - result = np.asarray(keypoints.normalize_keypoints2D(points2D, 128, 128)) - reference = normalize_reference(points2D, 128, 128) - assert np.allclose(result, reference) - - -def test_normalize_denormalize_round_trip(): - points2D = jp.array([[10.0, 20.0], [50.0, 5.0], [127.0, 63.0]]) - normalized = keypoints.normalize_keypoints2D(points2D, 128, 64) - recovered = keypoints.denormalize_keypoints2D(normalized, 128, 64) - assert np.allclose(np.asarray(recovered), np.asarray(points2D)) - - -def test_rotate_point2D_ninety_degrees(): - rotated = keypoints.rotate_point2D(jp.array([1.0, 0.0]), 90.0) - assert np.allclose(np.asarray(rotated), [0.0, 1.0], atol=1e-6) - - -def test_flip_keypoints_left_right(): - points = jp.array([[0.0, 5.0], [32.0, 10.0]]) - flipped = np.asarray(keypoints.flip_keypoints_left_right(points, 32.0)) - assert np.allclose(flipped, [[32.0, 5.0], [0.0, 10.0]]) - - -def test_transform_keypoint_translation(): - transform = jp.array([[1.0, 0.0, 3.0], [0.0, 1.0, -2.0], [0, 0, 1.0]]) - moved = keypoints.transform_keypoint(jp.array([4.0, 5.0]), transform) - assert np.allclose(np.asarray(moved)[:2], [7.0, 3.0]) - - -def test_uv_to_vu(): - flipped = keypoints.uv_to_vu(jp.array([[1.0, 2.0], [3.0, 4.0]])) - assert np.allclose(np.asarray(flipped), [[2.0, 1.0], [4.0, 3.0]]) - - -def test_build_cube_points3D_shape_and_center(): - cube = np.asarray(keypoints.build_cube_points3D(2.0, 4.0, 6.0)) - assert cube.shape == (8, 3) - assert np.allclose(cube.mean(axis=0), [0.0, 0.0, 0.0]) - assert np.allclose(np.abs(cube).max(axis=0), [1.0, 2.0, 3.0]) - - -def test_solve_PnP_recovers_known_pose(): - camera = build_camera() - points3D = keypoints.build_cube_points3D(0.1, 0.1, 0.1) - points3D = np.asarray(points3D, np.float64) - rotation = cv2.Rodrigues(np.array([0.2, -0.1, 0.3]))[0] - translation = np.array([0.02, -0.01, 0.5]) - points2D = project_to_image(rotation, translation, points3D, - camera.intrinsics) - pose6D = keypoints.solve_PnP(points2D, points3D, camera) - recovered = keypoints.rotation_vector_to_matrix(pose6D.rotation_vector) - assert np.allclose(recovered, rotation, atol=1e-4) - assert np.allclose(np.asarray(pose6D.translation).reshape(3), - translation, atol=1e-4) - - -def test_solve_pose_matrix_RANSAC_recovers_known_pose(): - camera = build_camera() - grid = np.linspace(-0.05, 0.05, 5) - points3D = np.array([[x, y, z] for x in grid for y in grid for z in grid]) - rotation = cv2.Rodrigues(np.array([0.1, 0.2, -0.15]))[0] - translation = np.array([0.0, 0.0, 0.6]) - points2D = project_to_image(rotation, translation, points3D, - camera.intrinsics) - result = keypoints.solve_pose_matrix_RANSAC(points2D, points3D, camera) - assert result is not None - recovered_rotation, recovered_translation = result - assert np.allclose(recovered_rotation, rotation, atol=1e-3) - assert np.allclose(recovered_translation, translation, atol=1e-3) - - -def test_solve_PnP_RANSAC_returns_none_below_minimum(): - camera = build_camera() - points2D = np.zeros((3, 2)) - points3D = np.zeros((3, 3)) - assert keypoints.solve_PnP_RANSAC(points2D, points3D, camera) is None - - -def hot_pixel_location(image): - row, col = np.unravel_index(np.argmax(image[..., 0]), image.shape[:2]) - return col, row - - -def test_rotate_keypoints2D_tracks_image_rotation(): - image = np.zeros((31, 31, 3), "float32") - image[5, 20] = 1.0 - rotated = paz.image.rotate(jp.asarray(image), 0.5) - col, row = hot_pixel_location(np.asarray(rotated)) - center = jp.array([(31 - 1) / 2.0, (31 - 1) / 2.0]) - keypoint = jp.array([[20.0, 5.0]]) # (x=col, y=row) - moved = np.asarray(keypoints.rotate_keypoints2D(keypoint, 0.5, center))[0] - assert abs(moved[0] - col) <= 1.5 and abs(moved[1] - row) <= 1.5 - - -def test_translate_keypoints_tracks_image_translation(): - image = np.zeros((31, 31, 3), "float32") - image[10, 8] = 1.0 - translation = jp.array([4.0, -3.0]) # (x, y) shift - translated = paz.image.translate(jp.asarray(image), translation) - col, row = hot_pixel_location(np.asarray(translated)) - keypoint = jp.array([[8.0, 10.0]]) - moved = np.asarray(keypoints.translate_keypoints(keypoint, translation))[0] - assert (moved[0], moved[1]) == (col, row) diff --git a/paz/backend/pinhole.py b/paz/backend/pinhole.py index 353c2536f..3034a9c59 100644 --- a/paz/backend/pinhole.py +++ b/paz/backend/pinhole.py @@ -162,6 +162,20 @@ def build_cube_corners(min_extents, max_extents): ) +def build_cube_points3D(width, height, depth): + half_width, half_height, half_depth = width / 2, height / 2, depth / 2 + point1 = [+half_width, -half_height, +half_depth] + point2 = [+half_width, -half_height, -half_depth] + point3 = [-half_width, -half_height, -half_depth] + point4 = [-half_width, -half_height, +half_depth] + point5 = [+half_width, +half_height, +half_depth] + point6 = [+half_width, +half_height, -half_depth] + point7 = [-half_width, +half_height, -half_depth] + point8 = [-half_width, +half_height, +half_depth] + points = [point1, point2, point3, point4, point5, point6, point7, point8] + return jp.array(points) + + def compute_AABB(vertices): min_extents = jp.min(vertices, axis=0) max_extents = jp.max(vertices, axis=0) diff --git a/paz/backend/pinhole_test.py b/paz/backend/pinhole_test.py new file mode 100644 index 000000000..c2ca189bc --- /dev/null +++ b/paz/backend/pinhole_test.py @@ -0,0 +1,14 @@ +import os + +os.environ.setdefault("KERAS_BACKEND", "jax") + +import numpy as np + +from paz.backend import pinhole + + +def test_build_cube_points3D_shape_and_center(): + cube = np.asarray(pinhole.build_cube_points3D(2.0, 4.0, 6.0)) + assert cube.shape == (8, 3) + assert np.allclose(cube.mean(axis=0), [0.0, 0.0, 0.0]) + assert np.allclose(np.abs(cube).max(axis=0), [1.0, 2.0, 3.0]) diff --git a/paz/backend/points2D.py b/paz/backend/points2D.py index bb526323c..46f20e3b5 100644 --- a/paz/backend/points2D.py +++ b/paz/backend/points2D.py @@ -36,6 +36,43 @@ def shift_to_box_origin(points, box): return points + jp.array([x_min, y_min]) +def normalize_keypoints2D(points2D, height, width): + image_shape = jp.array([width, height]) + return (2.0 * points2D / image_shape) - 1.0 + + +def denormalize_keypoints2D(points2D, height, width): + image_shape = jp.array([width, height]) + return (points2D + 1.0) / 2.0 * image_shape + + +def rotate_point2D(point2D, rotation_angle): + angle = jp.pi * rotation_angle / 180.0 + sin_angle, cos_angle = jp.sin(angle), jp.cos(angle) + x = point2D[0] * cos_angle - point2D[1] * sin_angle + y = point2D[0] * sin_angle + point2D[1] * cos_angle + return jp.array([x, y]) + + +def rotate_keypoints2D(keypoints, angle, center): + cos_angle, sin_angle = jp.cos(angle), jp.sin(angle) + rotation = jp.array([[cos_angle, -sin_angle], [sin_angle, cos_angle]]) + return (keypoints - center) @ rotation.T + center + + +def flip_keypoints_left_right(keypoints, width): + x, y = jp.split(keypoints, 2, axis=1) + return jp.concatenate([width - x, y], axis=1) + + +def translate_keypoints(keypoints, translation): + return keypoints + translation + + +def uv_to_vu(keypoints): + return keypoints[:, ::-1] + + def denormalize(keypoints, H, W): """Transform nomralized points2D to image UV coordinates i.e. [-1, 1] -> [U, V]. UV have maximum values of [W, H] respectively. diff --git a/paz/backend/points2D_test.py b/paz/backend/points2D_test.py new file mode 100644 index 000000000..1de2d1747 --- /dev/null +++ b/paz/backend/points2D_test.py @@ -0,0 +1,70 @@ +import os + +os.environ.setdefault("KERAS_BACKEND", "jax") + +import numpy as np +import jax.numpy as jp + +import paz +from paz.backend import points2D + + +def normalize_reference(points, height, width): + image_shape = np.array([width, height]) + return 2.0 * (points / image_shape) - 1.0 + + +def hot_pixel_location(image): + row, col = np.unravel_index(np.argmax(image[..., 0]), image.shape[:2]) + return col, row + + +def test_normalize_keypoints2D_matches_numpy_reference(): + points = np.array([[0.0, 0.0], [128.0, 64.0], [32.0, 96.0]]) + result = np.asarray(points2D.normalize_keypoints2D(points, 128, 128)) + assert np.allclose(result, normalize_reference(points, 128, 128)) + + +def test_normalize_denormalize_round_trip(): + points = jp.array([[10.0, 20.0], [50.0, 5.0], [127.0, 63.0]]) + normalized = points2D.normalize_keypoints2D(points, 128, 64) + recovered = points2D.denormalize_keypoints2D(normalized, 128, 64) + assert np.allclose(np.asarray(recovered), np.asarray(points)) + + +def test_rotate_point2D_ninety_degrees(): + rotated = points2D.rotate_point2D(jp.array([1.0, 0.0]), 90.0) + assert np.allclose(np.asarray(rotated), [0.0, 1.0], atol=1e-6) + + +def test_flip_keypoints_left_right(): + points = jp.array([[0.0, 5.0], [32.0, 10.0]]) + flipped = np.asarray(points2D.flip_keypoints_left_right(points, 32.0)) + assert np.allclose(flipped, [[32.0, 5.0], [0.0, 10.0]]) + + +def test_uv_to_vu(): + flipped = points2D.uv_to_vu(jp.array([[1.0, 2.0], [3.0, 4.0]])) + assert np.allclose(np.asarray(flipped), [[2.0, 1.0], [4.0, 3.0]]) + + +def test_rotate_keypoints2D_tracks_image_rotation(): + image = np.zeros((31, 31, 3), "float32") + image[5, 20] = 1.0 + rotated = paz.image.rotate(jp.asarray(image), 0.5) + col, row = hot_pixel_location(np.asarray(rotated)) + center = jp.array([(31 - 1) / 2.0, (31 - 1) / 2.0]) + keypoint = jp.array([[20.0, 5.0]]) # (x=col, y=row) + moved = np.asarray(points2D.rotate_keypoints2D(keypoint, 0.5, center))[0] + assert abs(moved[0] - col) <= 1.5 and abs(moved[1] - row) <= 1.5 + + +def test_translate_keypoints_tracks_image_translation(): + image = np.zeros((31, 31, 3), "float32") + image[10, 8] = 1.0 + translation = jp.array([4.0, -3.0]) # (x, y) shift + translated = paz.image.translate(jp.asarray(image), translation) + col, row = hot_pixel_location(np.asarray(translated)) + keypoint = jp.array([[8.0, 10.0]]) + moved = np.asarray(points2D.translate_keypoints(keypoint, translation))[0] + assert (moved[0], moved[1]) == (col, row) diff --git a/paz/backend/poses.py b/paz/backend/poses.py index 99914d214..c2d6b40ab 100644 --- a/paz/backend/poses.py +++ b/paz/backend/poses.py @@ -1,9 +1,19 @@ +from collections import namedtuple + import cv2 import numpy as np from paz.datasets import human36m +UPNP = cv2.SOLVEPNP_UPNP +LEVENBERG_MARQUARDT = cv2.SOLVEPNP_ITERATIVE +EPNP = cv2.SOLVEPNP_EPNP +MIN_REQUIRED_POINTS = 4 + +Pose6D = namedtuple("Pose6D", ["rotation_vector", "translation"]) + + def match_poses(boxes, poses, prior_boxes, iou_threshold=0.5): """Assigns ground-truth poses to prior boxes by IoU, appending a positive flag column. Returns an array of shape `(num_priors, poses_dim + 1)`.""" @@ -67,6 +77,50 @@ def project_to_image(rotation, translation, points3D, camera_intrinsics): return np.concatenate([fx * (x / z) + cx, fy * (y / z) + cy], axis=1) +def project_points3D(points3D, pose6D, camera): + args = (pose6D.translation, camera.intrinsics, camera.distortion) + points2D, _ = cv2.projectPoints(points3D, pose6D.rotation_vector, *args) + return np.squeeze(points2D, axis=1) # openCV shape (num_points, 1, 2) + + +def solve_PnP(points2D, points3D, camera, solver=LEVENBERG_MARQUARDT): + points2D = np.array(points2D, np.float64).reshape((len(points3D), 1, 2)) + args = (camera.intrinsics, camera.distortion, None, None, False, solver) + (_, rotation_vector, translation) = cv2.solvePnP(points3D, points2D, *args) + return Pose6D(rotation_vector, translation) + + +def solve_PnP_RANSAC(points2D, points3D, camera, inlier_thresh=5.0, + iterations=100): + if len(points3D) < MIN_REQUIRED_POINTS: + return None + points2D = np.array(points2D, np.float64).reshape((len(points3D), 1, 2)) + points3D = np.array(points3D, np.float64) + args = (camera.intrinsics, camera.distortion, None, None, False, + iterations, inlier_thresh, 0.99, None, EPNP) + success, rotation, translation, inliers = cv2.solvePnPRansac( + points3D, points2D, *args) + if not success: + return None + return Pose6D(rotation, translation) + + +def rotation_vector_to_matrix(rotation_vector): + return cv2.Rodrigues(rotation_vector)[0] + + +def solve_pose_matrix_RANSAC(points2D, points3D, camera, max_points=1500, + seed=0): + if len(points3D) > max_points: + choice = np.random.RandomState(seed).choice(len(points3D), max_points, False) # fmt: skip + points2D, points3D = points2D[choice], points3D[choice] + pose6D = solve_PnP_RANSAC(points2D, points3D, camera) + if pose6D is None: + return None + rotation = rotation_vector_to_matrix(pose6D.rotation_vector) + return rotation, np.asarray(pose6D.translation).reshape(3) + + def filter_keypoints3D(keypoints3D, args_to_joints3D): keypoints3D = np.reshape(keypoints3D, [len(keypoints3D), 32, 3]) return keypoints3D[:, args_to_joints3D, :] diff --git a/paz/backend/poses_test.py b/paz/backend/poses_test.py index aaecaf3b3..5f31c08fb 100644 --- a/paz/backend/poses_test.py +++ b/paz/backend/poses_test.py @@ -2,10 +2,23 @@ os.environ.setdefault("KERAS_BACKEND", "jax") +from collections import namedtuple + import numpy as np +import cv2 from paz.backend import poses +Camera = namedtuple("Camera", ["intrinsics", "distortion"]) + + +def build_camera(size=128, focal=150.0): + center = size / 2.0 + intrinsics = np.array([[focal, 0, center], + [0, focal, center], + [0, 0, 1.0]]) + return Camera(intrinsics, np.zeros((4, 1))) + def test_rotation_matrix_to_axis_angle_identity(): identity = np.eye(3).reshape(1, 9) @@ -31,3 +44,51 @@ def test_match_poses_shapes_and_flag(): assert matched.shape == (2, 10) assert matched[0, -1] == 1.0 assert matched[1, -1] == 0.0 + + +def test_solve_PnP_recovers_known_pose(): + camera = build_camera() + grid = np.linspace(-0.05, 0.05, 3) + points3D = np.array([[x, y, z] for x in grid for y in grid for z in grid]) + rotation = cv2.Rodrigues(np.array([0.2, -0.1, 0.3]))[0] + translation = np.array([0.02, -0.01, 0.5]) + points2D = poses.project_to_image(rotation, translation, points3D, + camera.intrinsics) + pose6D = poses.solve_PnP(points2D, points3D, camera) + recovered = poses.rotation_vector_to_matrix(pose6D.rotation_vector) + assert np.allclose(recovered, rotation, atol=1e-4) + assert np.allclose(np.asarray(pose6D.translation).reshape(3), + translation, atol=1e-4) + + +def test_solve_pose_matrix_RANSAC_recovers_known_pose(): + camera = build_camera() + grid = np.linspace(-0.05, 0.05, 5) + points3D = np.array([[x, y, z] for x in grid for y in grid for z in grid]) + rotation = cv2.Rodrigues(np.array([0.1, 0.2, -0.15]))[0] + translation = np.array([0.0, 0.0, 0.6]) + points2D = poses.project_to_image(rotation, translation, points3D, + camera.intrinsics) + result = poses.solve_pose_matrix_RANSAC(points2D, points3D, camera) + assert result is not None + recovered_rotation, recovered_translation = result + assert np.allclose(recovered_rotation, rotation, atol=1e-3) + assert np.allclose(recovered_translation, translation, atol=1e-3) + + +def test_solve_PnP_RANSAC_returns_none_below_minimum(): + camera = build_camera() + points2D, points3D = np.zeros((3, 2)), np.zeros((3, 3)) + assert poses.solve_PnP_RANSAC(points2D, points3D, camera) is None + + +def test_project_points3D_matches_project_to_image(): + camera = build_camera() + points3D = np.array([[0.01, 0.0, 0.5], [-0.02, 0.03, 0.6]]) + rotation = cv2.Rodrigues(np.array([0.1, 0.0, 0.0]))[0] + translation = np.array([0.0, 0.0, 0.5]) + pose6D = poses.Pose6D(cv2.Rodrigues(rotation)[0], translation) + cv2_points = np.asarray(poses.project_points3D(points3D, pose6D, camera)) + analytic = poses.project_to_image(rotation, translation, points3D, + camera.intrinsics) + assert np.allclose(cv2_points, analytic, atol=1e-3)