test_renderer.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

import unittest

import numpy as np
import pymomentum.camera as pym_camera
import pymomentum.geometry as pym_geometry
import pymomentum.renderer as pym_renderer
import torch
from numpy.typing import NDArray


class TestRendering(unittest.TestCase):
    def test_lighting(self) -> None:
        """Test that directional lighting works."""

        image_width, image_height = 30, 30
        camera = pym_camera.Camera(
            pym_camera.PinholeIntrinsicsModel(
                image_width=image_width,
                image_height=image_height,
            )
        )

        def rasterize_spheres_with_lights(
            lights: list[pym_renderer.Light] | None,
            camera_param: pym_camera.Camera = camera,
        ) -> torch.Tensor:
            z_buffer = pym_renderer.create_z_buffer(camera_param)
            rgb_buffer = pym_renderer.create_rgb_buffer(camera_param)

            pym_renderer.rasterize_spheres(
                center=torch.as_tensor([[0, 0, 20]], dtype=torch.float32),
                lights=lights,
                radius=torch.as_tensor([5], dtype=torch.float32),
                camera=camera_param,
                rgb_buffer=rgb_buffer,
                z_buffer=z_buffer,
            )
            return rgb_buffer

        def mean_absolute_difference(t1: torch.Tensor, t2: torch.Tensor) -> float:
            return torch.mean(torch.abs(t1 - t2)).item()

        default_lighting = rasterize_spheres_with_lights(None)
        lights_at_camera = rasterize_spheres_with_lights(
            [pym_renderer.Light.create_point_light(camera.center_of_projection)]
        )
        # pyre-fixme[6]: For 2nd argument expected `SupportsDunderLT[_T]` but got
        #  `float`.
        self.assertGreater(torch.mean(default_lighting), 0.15)
        front_directional_lighting = rasterize_spheres_with_lights(
            [
                pym_renderer.Light.create_directional_light(
                    camera.world_space_principle_axis
                )
            ]
        )

        # One light at the camera should be identical to the default lighting.
        self.assertLess(
            mean_absolute_difference(default_lighting, lights_at_camera), 1e-4
        )
        # Directional lighting is a bit different from a point light at the camera but
        # should be pretty close.
        self.assertLess(
            mean_absolute_difference(default_lighting, front_directional_lighting), 0.05
        )
        self.assertGreater(
            torch.mean(front_directional_lighting),
            # pyre-fixme[6]: For 2nd argument expected `SupportsDunderLT[_T]` but
            #  got `Tensor`.
            torch.mean(default_lighting),
        )

        back_lighting = rasterize_spheres_with_lights(
            [
                pym_renderer.Light.create_directional_light(
                    -camera.world_space_principle_axis
                )
            ]
        )

        # We should see very little light coming from the back:
        self.assertLess(
            mean_absolute_difference(back_lighting, torch.zeros_like(default_lighting)),
            1e-4,
        )

        # All lighting should be from the left side:
        left_lighting = rasterize_spheres_with_lights(
            [
                pym_renderer.Light.create_directional_light(
                    np.asarray([1, 0, 0], dtype=np.float32)
                )
            ]
        )
        # pyre-fixme[6]: For 2nd argument expected `SupportsDunderLT[_T]` but got
        #  `float`.
        self.assertGreater(torch.mean(left_lighting[:, : image_width // 2]), 0.1)
        # pyre-fixme[6]: For 2nd argument expected `SupportsDunderGT[_T]` but got
        #  `float`.
        self.assertLess(torch.mean(left_lighting[:, image_width // 2 :]), 1e-5)

        # All lighting should be from the right side:
        right_lighting = rasterize_spheres_with_lights(
            [
                pym_renderer.Light.create_directional_light(
                    np.asarray([-1, 0, 0], dtype=np.float32)
                )
            ]
        )
        # pyre-fixme[6]: For 2nd argument expected `SupportsDunderGT[_T]` but got
        #  `float`.
        self.assertLess(torch.mean(right_lighting[:, : image_width // 2]), 1e-5)
        # pyre-fixme[6]: For 2nd argument expected `SupportsDunderLT[_T]` but got
        #  `float`.
        self.assertGreater(torch.mean(right_lighting[:, image_width // 2 :]), 0.1)

    def test_subdivide_uneven(self) -> None:
        """Check that we can do uneven subdivision where only some edges are split."""

        def _compute_area(
            mesh_positions: NDArray[np.float32], triangles: NDArray[np.int32]
        ) -> float:
            v1 = mesh_positions[triangles[:, 1]] - mesh_positions[triangles[:, 0]]
            v2 = mesh_positions[triangles[:, 2]] - mesh_positions[triangles[:, 0]]
            cross_prod = np.cross(v1, v2)
            return 0.5 * np.sum(np.linalg.norm(cross_prod, axis=-1))

        def _compute_triangle_normals(
            mesh_positions: NDArray[np.float32], triangles: NDArray[np.int32]
        ) -> NDArray[np.float32]:
            v1 = mesh_positions[triangles[:, 1]] - mesh_positions[triangles[:, 0]]
            v2 = mesh_positions[triangles[:, 2]] - mesh_positions[triangles[:, 0]]
            cross_prod = np.cross(v1, v2)
            return cross_prod / np.linalg.norm(cross_prod, axis=-1, keepdims=True)

        mesh_positions_orig = np.asarray([[0, 0, 0], [0.25, 0, 0], [0, 1, 0]])

        mesh_normals_orig = np.asarray([[0, 0, 1]] * 3)
        mesh_triangles_orig = np.asarray([[0, 1, 2]], dtype=np.int32)

        self.assertAlmostEqual(
            _compute_area(mesh_positions_orig, mesh_triangles_orig), 0.125
        )

        # Check basic invariants:
        #   total area is unchanged
        #   all triangles have the correct winding
        def _check_invariants(
            max_edge_length: float,
            expected_triangles: int,
            positions: NDArray[np.float32] = mesh_positions_orig,
            normals: NDArray[np.float32] = mesh_normals_orig,
            triangles: NDArray[np.int32] = mesh_triangles_orig,
        ) -> None:
            (
                mesh_positions_subd,
                mesh_normals_subd,
                mesh_triangles_subd,
                texture_coords_subd,
                texture_triangles_subd,
            ) = pym_renderer.subdivide_mesh(
                positions,
                normals,
                triangles,
                levels=1,
                max_edge_length=max_edge_length,
            )

            self.assertEqual(mesh_triangles_subd.shape[0], expected_triangles)

            # Since we didn't provide texture coords, should get just all-zeros texture coords:
            self.assertTrue(np.array_equal(mesh_triangles_subd, texture_triangles_subd))
            self.assertFalse(np.any(texture_coords_subd))

            area_new = _compute_area(mesh_positions_subd, mesh_triangles_subd)
            self.assertAlmostEqual(area_new, 0.125)

            normals_new = _compute_triangle_normals(
                mesh_positions_subd, mesh_triangles_subd
            )
            self.assertEqual(normals_new.shape[0], expected_triangles)
            for i in range(normals_new.shape[0]):
                self.assertLess(np.linalg.norm(normals_new[i] - normals[0]), 1e-3)

        _check_invariants(10, 1)
        _check_invariants(1.01, 2)
        _check_invariants(0.99, 3)
        _check_invariants(0.001, 4)

    def test_shadow_projection_matrix_point_light(self) -> None:
        """Check that the projection matrix actually projects points."""

        points = np.asarray(
            [[0, 0, 0], [0, 10, 0], [0, 0, 10], [0, 10, 10], [10, 20, 20]],
            dtype=np.float32,
        )

        pos_light = pym_renderer.Light.create_point_light(
            np.asarray([0, 50, 0], dtype=np.float32)
        )
        pos_proj_mat = pym_renderer.create_shadow_projection_matrix(pos_light)
        plane_normal = np.asarray([0, 1, 0], dtype=np.float32)

        for i in range(points.shape[0]):
            p = np.concatenate([points[i, :], np.ones([1], dtype=np.float32)])
            p_proj_unnormalized = np.matmul(pos_proj_mat, p)
            p_proj_normalized = p_proj_unnormalized[0:3] / p_proj_unnormalized[3]

            # vector pointing from p toward light:
            light_vec = pos_light.position - p[0:3]
            light_vec = light_vec / np.linalg.norm(light_vec)

            # intersect with plane:
            # (p + t * light_vec) . n = 0
            # p . n = -(light_vec . n) * t
            # Solve for t:
            t = -np.dot(p[0:3], plane_normal) / np.dot(light_vec, plane_normal)
            p_proj = p[0:3] + t * light_vec
            self.assertTrue(np.allclose(p_proj, p_proj_normalized))

    def test_shadow_projection_matrix_dir_light(self) -> None:
        points = np.asarray(
            [[0, 0, 0], [0, 10, 0], [0, 0, 10], [0, 10, 10], [10, 20, 20]],
            dtype=np.float32,
        )

        dir_light = pym_renderer.Light.create_directional_light(
            np.asarray([-1, -1, -2], dtype=np.float32)
        )
        plane_normal = np.asarray([0, 1, 0], dtype=np.float32)
        plane_origin = np.asarray([0, -1, 0], dtype=np.float32)
        plane_offset = np.dot(plane_normal, plane_origin)
        dir_proj_mat = pym_renderer.create_shadow_projection_matrix(
            dir_light, plane_normal=plane_normal, plane_origin=plane_origin
        )

        for i in range(points.shape[0]):
            p = np.concatenate([points[i, :], np.ones([1], dtype=np.float32)])
            p_proj_unnormalized = np.matmul(dir_proj_mat, p)
            p_proj_normalized = p_proj_unnormalized[0:3] / p_proj_unnormalized[3]

            # vector pointing from p toward light:
            light_vec = dir_light.position
            light_vec = light_vec / np.linalg.norm(light_vec)

            # intersect with plane:
            # (p + t * light_vec) . n = plane_offset
            # t * (light_vec . n) = plane_offset - (p . n)
            # Solve for t:
            t = (plane_offset - np.dot(p[0:3], plane_normal)) / np.dot(
                light_vec, plane_normal
            )
            p_proj = p[0:3] + t * light_vec
            self.assertTrue(np.allclose(p_proj, p_proj_normalized))
            self.assertAlmostEqual(p_proj[1], -1)

    def _create_test_character_for_camera(self) -> pym_geometry.Character:
        joints = [
            pym_geometry.Joint(
                name="b_root",
                parent=-1,
                pre_rotation=np.asarray([0, 0, 0, 1], dtype=np.float32),
                translation_offset=np.asarray([0, 0, 1], dtype=np.float32),
            ),
            pym_geometry.Joint(
                name="b_spine1",
                parent=0,
                pre_rotation=np.asarray([0, 0, 0, 1], dtype=np.float32),
                translation_offset=np.asarray([0, 0, 1], dtype=np.float32),
            ),
            pym_geometry.Joint(
                name="b_spine3",
                parent=0,
                pre_rotation=np.asarray([0, 0, 0, 1], dtype=np.float32),
                translation_offset=np.asarray([0, 0, 1], dtype=np.float32),
            ),
            pym_geometry.Joint(
                name="b_left_arm",
                parent=1,
                pre_rotation=np.asarray([0, 0, 0, 1], dtype=np.float32),
                translation_offset=np.asarray([0, 0, 1], dtype=np.float32),
            ),
        ]

        skeleton = pym_geometry.Skeleton(joints)

        # Create an empty parameter transform (no model parameters)
        # Transform matrix shape: (7 * n_joints) x n_params = (7 * 3) x 0 = 21 x 0
        n_joints = len(joints)
        empty_transform = np.zeros((7 * n_joints, 0), dtype=np.float32)

        return pym_geometry.Character(
            name="test_camera",
            skeleton=skeleton,
            parameter_transform=pym_geometry.ParameterTransform(
                names=[], skeleton=skeleton, transform=empty_transform
            ),
        )

    def test_create_camera_for_body_single_frame(self) -> None:
        """Test that create_camera_for_body matches build_cameras_for_body for unbatched input."""
        character = self._create_test_character_for_camera()

        # Create random joint parameters
        torch.manual_seed(42)
        joint_parameters = torch.randn(character.skeleton.size * 7).reshape(
            character.skeleton.size, 7
        )
        skeleton_states = torch.from_numpy(
            pym_geometry.joint_parameters_to_skeleton_state(
                character, joint_parameters.unsqueeze(0).numpy()
            )
        ).squeeze(0)

        # Create cameras using old function
        cameras_old = pym_renderer.build_cameras_for_body(
            character, joint_parameters, 512, 512
        )

        # Create cameras using new function
        camera_new = pym_renderer.create_camera_for_body(
            character, skeleton_states.numpy(), 512, 512
        )

        # Verify we got the same number of cameras
        self.assertEqual(len(cameras_old), 1)

        # Verify camera properties match
        self._compare_cameras(cameras_old[0], camera_new)

    def test_create_camera_for_body_multiple_frames(self) -> None:
        """Test that create_camera_for_body matches build_cameras_for_body for batched input."""
        character = self._create_test_character_for_camera()

        # Create random joint parameters with batch size 3
        torch.manual_seed(42)
        n_frames = 3
        joint_parameters = torch.randn(n_frames, character.skeleton.size * 7).reshape(
            n_frames, character.skeleton.size, 7
        )
        skeleton_states = torch.from_numpy(
            pym_geometry.joint_parameters_to_skeleton_state(
                character, joint_parameters.numpy()
            )
        )

        # Create cameras using old function
        cameras_old = pym_renderer.build_cameras_for_body(
            character, joint_parameters.unsqueeze(0), 512, 512
        )

        # Create cameras using new function
        camera_new = pym_renderer.create_camera_for_body(
            character, skeleton_states.numpy(), 512, 512
        )

        # Verify we got the same number of cameras
        self.assertEqual(len(cameras_old), 1)

        # Verify camera properties match for each batch
        self._compare_cameras(cameras_old[0], camera_new)

    def _compare_cameras(
        self, camera_old: pym_camera.Camera, camera_new: pym_camera.Camera
    ) -> None:
        """Compare two cameras to ensure they have matching properties."""
        # Compare basic properties
        self.assertEqual(camera_old.image_width, camera_new.image_width)
        self.assertEqual(camera_old.image_height, camera_new.image_height)
        self.assertAlmostEqual(camera_old.fx, camera_new.fx, places=4)
        self.assertAlmostEqual(camera_old.fy, camera_new.fy, places=4)

        # Compare camera transforms
        T_old = camera_old.T_eye_from_world
        T_new = camera_new.T_eye_from_world
        np.testing.assert_allclose(T_old, T_new, rtol=1e-4, atol=1e-4)

        # Compare center of projection
        cop_old = camera_old.center_of_projection
        cop_new = camera_new.center_of_projection
        np.testing.assert_allclose(cop_old, cop_new, rtol=1e-4, atol=1e-4)