diff --git a/src/poseforge/pose/camera.py b/src/poseforge/pose/camera.py index b7718ed3..89d83744 100644 --- a/src/poseforge/pose/camera.py +++ b/src/poseforge/pose/camera.py @@ -95,8 +95,17 @@ def __init__( def __call__(self, xy, depth): """Map 2D camera coordinates + depth to 3D world coordinates. + IMPORTANT: ``xy`` MUST be expressed in the SAME pixel space the mapper + was built with, i.e. in ``rendering_size`` pixels. The focal length and + principal point scale linearly with ``rendering_size``, so feeding xy + that live in a different pixel resolution injects an anisotropic + distortion (in-plane x, y scaled by the resolution ratio, depth left + exact). If your predictions are at resolution S, build the mapper with + ``rendering_size=(S, S)``. + Args: - xy: (..., 2) array of 2D camera coordinates (in pixels). + xy: (..., 2) array of 2D camera coordinates, in ``rendering_size`` + pixels. depth: (...,) array of depth values (in mm). Returns: diff --git a/src/poseforge/pose/keypoints3d/scripts/run_keypoints3d_inference.py b/src/poseforge/pose/keypoints3d/scripts/run_keypoints3d_inference.py index f59d16af..d55755d7 100644 --- a/src/poseforge/pose/keypoints3d/scripts/run_keypoints3d_inference.py +++ b/src/poseforge/pose/keypoints3d/scripts/run_keypoints3d_inference.py @@ -97,9 +97,32 @@ def run_keypoints3d_inference( print("========== Model Summary ==========") summary(pipeline.model, input_size=(3, *inference_image_size), device="cuda") - # Set up camera mapper + # Set up camera mapper. + # + # IMPORTANT: the mapper's intrinsics (focal length and principal point) + # scale linearly with the image/sensor size it is built with, so it MUST be + # built at the pixel space the predictions actually live in. The model + # consumes images resized to `inference_image_size` and emits `pred_xy` in + # that same pixel space. Therefore the mapper is built at + # `inference_image_size`, NOT at `camera_rendering_size`. + # + # `camera_rendering_size` is the resolution at which the *training* data was + # rendered in simulation; it is NOT the space the predictions are in. Using + # it here would scale the in-plane (x, y) coordinates by + # inference_image_size / camera_rendering_size while leaving depth exact, + # injecting an anisotropic distortion that corrupts downstream IK joint + # angles. See tests/test_camera_unprojection.py. + if tuple(camera_rendering_size) != tuple(inference_image_size): + print( + "Note: camera_rendering_size " + f"{tuple(camera_rendering_size)} differs from inference_image_size " + f"{tuple(inference_image_size)}. The camera mapper is built at " + "inference_image_size because pred_xy is in inference_image_size " + "pixels; camera_rendering_size (the simulation render size) is " + "intentionally not used for unprojection." + ) cam_mapper = CameraToWorldMapper( - camera_pos, camera_fov_deg, camera_rendering_size, camera_rotation_euler + camera_pos, camera_fov_deg, inference_image_size, camera_rotation_euler ) # Run inference @@ -287,7 +310,7 @@ def start(): output_basedir=output_basedir, batch_size=batch_size, n_workers=n_workers, - c=inference_image_size, + inference_image_size=inference_image_size, camera_pos=camera_pos, camera_fov_deg=camera_fov_deg, camera_rendering_size=camera_rendering_size, diff --git a/src/poseforge/production/spotlight/keypoints3d.py b/src/poseforge/production/spotlight/keypoints3d.py index 544a006e..dc90187e 100644 --- a/src/poseforge/production/spotlight/keypoints3d.py +++ b/src/poseforge/production/spotlight/keypoints3d.py @@ -62,14 +62,37 @@ def predict_keypoints3d( logger.info("Creating 3D keypoints inference pipeline") pipeline = keypoints3d.Pose2p5DPipeline(model, device=device, use_float16=True) - # Set up camera mapper + # Set up camera mapper. + # + # IMPORTANT: the mapper's intrinsics (focal length and principal point) + # scale linearly with the image/sensor size it is built with, so it MUST be + # built at the pixel space the predictions actually live in. The model + # consumes images resized to `working_size` and emits `pred_xy` in that same + # pixel space, so the mapper is built at `(working_size, working_size)`, NOT + # at `camera_rendering_size`. `camera_rendering_size` is the simulation + # render resolution of the training data and is NOT the prediction space; + # using it would scale in-plane (x, y) by working_size / camera_rendering_size + # while leaving depth exact, an anisotropic distortion that corrupts + # downstream IK joint angles. See tests/test_camera_unprojection.py. + working_size = keypoints3d_model_config["working_size"] + if tuple(camera_rendering_size) != (working_size, working_size): + logger.info( + "camera_rendering_size {} differs from working_size {}; building the " + "camera mapper at working_size because pred_xy is in working_size " + "pixels (camera_rendering_size is the simulation render size and is " + "intentionally not used for unprojection).", + tuple(camera_rendering_size), + (working_size, working_size), + ) cam_mapper = CameraToWorldMapper( - camera_pos, camera_fov_deg, camera_rendering_size, camera_rotation_euler + camera_pos, + camera_fov_deg, + (working_size, working_size), + camera_rotation_euler, ) # Create video loader logger.info("Creating video loader for 3D keypoints prediction") - working_size = keypoints3d_model_config["working_size"] video_loader = SimpleVideoCollectionLoader( [aligned_behavior_video_path], transform=transforms.Resize((working_size, working_size)), diff --git a/tests/test_camera_unprojection.py b/tests/test_camera_unprojection.py new file mode 100644 index 00000000..2293525b --- /dev/null +++ b/tests/test_camera_unprojection.py @@ -0,0 +1,199 @@ +"""Tests for ``poseforge.pose.camera.CameraToWorldMapper``. + +These tests pin down two facts that motivate the bug fix in +``run_keypoints3d_inference.py`` and ``production/spotlight/keypoints3d.py``: + +1. The mapper is a correct inverse of its own forward projection (round-trip + recovery to ~1e-6). +2. The mapper's intrinsics scale with the pixel resolution it is built with: + the in-plane (x, y) focal scaling is proportional to the image size while + depth is left untouched. Consequently, unprojecting predictions that live in + ``inference_image_size`` (=256) pixels with a mapper built at 256 is correct, + whereas building it at the simulation render size (=464) injects an + anisotropic distortion (in-plane factor 256/464 ~= 0.55, depth factor 1.0). + +``camera.py`` depends only on numpy and scipy, so these tests run standalone. +The module is loaded directly from the source file so the tests always exercise +the code under review regardless of where an editable install resolves +``poseforge`` to. +""" + +import importlib.util +from pathlib import Path + +import numpy as np +import pytest + +# Load camera.py directly from this repo's source tree (sibling of tests/). +_CAMERA_PY = ( + Path(__file__).resolve().parent.parent + / "src" + / "poseforge" + / "pose" + / "camera.py" +) +_spec = importlib.util.spec_from_file_location("_poseforge_camera_under_test", _CAMERA_PY) +_camera = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_camera) +CameraToWorldMapper = _camera.CameraToWorldMapper + + +# Camera parameters matching the inference/production defaults. +CAMERA_POS = (0.0, 0.0, -100.0) +CAMERA_FOV_DEG = 5.0 +CAMERA_ROTATION_EULER = (0.0, np.pi, -np.pi / 2) + +INFERENCE_SIZE = 256 # pixel space the model's pred_xy actually lives in +RENDERING_SIZE = 464 # simulation render size -- NOT the prediction space + + +def _make_mapper(size: int) -> "CameraToWorldMapper": + return CameraToWorldMapper( + CAMERA_POS, CAMERA_FOV_DEG, (size, size), CAMERA_ROTATION_EULER + ) + + +def _forward_project(mapper, world_xyz: np.ndarray): + """Project world points to (pixel_xy, depth) with the mapper's own + ``camera_matrix``. This is the exact inverse operation of ``__call__``. + + Args: + world_xyz: (N, 3) world coordinates in mm. + + Returns: + (pixel_xy (N, 2), depth (N,)). + """ + hom = np.hstack([world_xyz, np.ones((world_xyz.shape[0], 1))]) # (N, 4) + proj = (mapper.camera_matrix @ hom.T).T # (N, 3): [x*z, y*z, z] + depth = proj[:, 2] + pixel_xy = proj[:, :2] / proj[:, 2:3] + return pixel_xy, depth + + +def _to_camera_frame(mapper, world_xyz: np.ndarray) -> np.ndarray: + """Map world coords into the mapper's camera frame. The camera-frame z is + the depth; x, y are the in-plane coordinates.""" + hom = np.hstack([world_xyz, np.ones((world_xyz.shape[0], 1))]) + return (mapper.camera_extrinsic_mat @ hom.T).T[:, :3] + + +@pytest.mark.parametrize("size", [128, 256, 464]) +def test_unprojection_inverts_projection(size): + """Unprojecting the mapper's own forward projection recovers the original + world points to numerical precision.""" + mapper = _make_mapper(size) + + # Several known world points spread around the camera target region (mm). + world_xyz = np.array( + [ + [0.5, -1.0, 0.2], + [1.0, -2.0, 0.5], + [-0.3, 0.4, -0.1], + [2.0, 1.5, 1.0], + [-1.2, -0.7, 0.9], + [0.0, 0.0, 0.0], + ] + ) + + pixel_xy, depth = _forward_project(mapper, world_xyz) + recovered = mapper(pixel_xy, depth) + + assert recovered.shape == world_xyz.shape + np.testing.assert_allclose(recovered, world_xyz, atol=1e-6, rtol=0) + + +def test_focal_scales_with_size_depth_does_not(): + """The in-plane focal scaling is linear in image size; depth is independent + of image size. + + This is the root cause of the bug: the mapper's geometry depends on the + pixel space it is built with, so it must be built at the resolution the + predictions live in. + """ + m256 = _make_mapper(INFERENCE_SIZE) + m464 = _make_mapper(RENDERING_SIZE) + + focal_256 = -m256.focal_transform_mat[0, 0] + focal_464 = -m464.focal_transform_mat[0, 0] + + # In-plane focal scaling is proportional to image size. + np.testing.assert_allclose( + focal_256 / focal_464, INFERENCE_SIZE / RENDERING_SIZE, rtol=1e-12 + ) + # ~0.55, NOT 1.0 -- so resolution matters for in-plane coords. + assert abs(focal_256 / focal_464 - 0.55) < 0.01 + + # The depth axis of the focal transform is the identity regardless of size. + assert m256.focal_transform_mat[2, 2] == 1.0 + assert m464.focal_transform_mat[2, 2] == 1.0 + + +def test_same_input_different_mapper_size_anisotropic_distortion(): + """Unprojecting the SAME (256-space) xy + depth with a 256-mapper vs a + 464-mapper differs, with in-plane scaling ~= 256/464 and depth scaling + exactly 1.0. + + This confirms that 256 is the correct resolution for 256-space predictions: + using a 464-mapper shrinks the recovered in-plane geometry by ~0.55 while + leaving depth untouched -- precisely the anisotropy that corrupts the + downstream inverse-kinematics joint angles. + """ + m256 = _make_mapper(INFERENCE_SIZE) # correct for 256-space pred_xy + m464 = _make_mapper(RENDERING_SIZE) # buggy: 256-space xy fed to a 464 mapper + + # pred_xy lives in 256-px space; depths are mm. + xy = np.array( + [ + [100.0, 120.0], + [200.0, 50.0], + [10.0, 240.0], + [128.0, 128.0], + [60.0, 200.0], + ] + ) + depth = np.array([95.0, 102.0, 100.0, 99.0, 98.0]) + + world_256 = m256(xy, depth) + world_464 = m464(xy, depth) + + # The two unprojections genuinely disagree (this is the bug's effect). + assert not np.allclose(world_256, world_464) + + cam_256 = _to_camera_frame(m256, world_256) + cam_464 = _to_camera_frame(m464, world_464) + + # Depth (camera-frame z) is exactly the input depth in BOTH cases: depth is + # unaffected by the mapper resolution -> depth ratio is exactly 1.0. + np.testing.assert_allclose(cam_256[:, 2], depth, atol=1e-9) + np.testing.assert_allclose(cam_464[:, 2], depth, atol=1e-9) + np.testing.assert_allclose(cam_256[:, 2] / cam_464[:, 2], 1.0, atol=1e-9) + + # In-plane geometry, however, is scaled. Hold the pixel displacement away + # from each sensor's own principal point fixed (this isolates 1/focal): the + # same +(dx, dy) pixel step maps to a camera-frame in-plane offset whose + # magnitude scales as focal_464 / focal_256 = 464/256. Equivalently, the + # 256-mapper's in-plane extent is 256/464 ~= 0.55x that of the 464-mapper + # for the same predicted pixels. + fixed_depth = np.array([100.0, 100.0]) + step = np.array([40.0, 25.0]) # arbitrary fixed pixel displacement + + def inplane_step_magnitude(mapper): + pp = (mapper.rendering_size[0] - 1) / 2.0 + base_px = np.array([pp, pp]) + query_xy = np.stack([base_px, base_px + step]) + cam = _to_camera_frame(mapper, mapper(query_xy, fixed_depth)) + return np.linalg.norm(cam[1, :2] - cam[0, :2]) + + inplane_256 = inplane_step_magnitude(m256) + inplane_464 = inplane_step_magnitude(m464) + + # In-plane displacement ratio is the inverse of the focal ratio. + np.testing.assert_allclose( + inplane_256 / inplane_464, RENDERING_SIZE / INFERENCE_SIZE, rtol=1e-9 + ) + # Equivalently: the 464-mapper's in-plane response is 256/464 ~= 0.55x the + # 256-mapper's, while depth response is identical (1.0x). Anisotropy. + np.testing.assert_allclose( + inplane_464 / inplane_256, INFERENCE_SIZE / RENDERING_SIZE, rtol=1e-9 + ) + assert abs(inplane_464 / inplane_256 - 0.55) < 0.01