Add PnP and keypoint geometry across poses/points2D/pinhole

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.poses import solve_PnP_RANSAC, project_points3D
+from paz.pinhole import 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)

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.poses.solve_pose_matrix_RANSAC(*args)

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.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
@@ -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]

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,62 @@
 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.points2D.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.points2D.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)
+    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
+
+
 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__":

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.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):
@@ -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):

paz/backend/image.py

@@ -496,6 +496,13 @@ def random_rotation(
     return rotate(image, angle, order, mode, cval)
 
 
+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)
+    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)

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)

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])

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.