Merge branch 'master' into 891-vis_walking_metrics

soccer_control/soccer_pycontrol/src/soccer_pycontrol/model/sensors.py

@@ -81,3 +81,36 @@ def get_foot_pressure_sensors(self, floor: pb.loadURDF) -> List[bool]:
         #     index = np.signbit(np.matmul(left_tr, point[5] - left_center))[0:2]
         #     locations[index[1] + (index[0] * 2) + 4] = True
         return locations
+
+    def get_camera_image(self):
+        """
+        Captures the image from the camera mounted on the robot
+        """
+        robot_position, robot_orientation = pb.getBasePositionAndOrientation(self.body)
+
+        # Add more offsets later
+        camera_position = robot_position
+        camera_target = [robot_position[0] + 1, robot_position[1], robot_position[2]]
+
+        camera_up = [0, 0, 1]
+
+        view_matrix = pb.computeViewMatrix(camera_position, camera_target, camera_up)
+
+        width, height = 1280, 720
+        fov = 60
+        aspect = width / height
+        near = 0.1
+        far = 50
+
+        projection_matrix = pb.computeProjectionMatrixFOV(fov, aspect, near, far)
+
+        images = pb.getCameraImage(width, height, view_matrix, projection_matrix, shadow=False,
+                                   renderer=pb.ER_BULLET_HARDWARE_OPENGL)
+
+        # NOTE: the ordering of height and width change based on the conversion
+        rgb_opengl = np.reshape(images[2], (height, width, 4))[:, :, :3] * 1. / 255.
+        # depth_buffer_opengl = np.reshape(images[3], [width, height])
+        # depth_opengl = far * near / (far - (far - near) * depth_buffer_opengl)
+        # seg_opengl = np.reshape(images[4], [width, height]) * 1. / 255.
+
+        return rgb_opengl

soccer_control/soccer_pycontrol/test/test_perception.py

@@ -0,0 +1,46 @@
+import unittest
+import matplotlib.pyplot as plt
+from soccer_pycontrol.model.bez import Bez  # Import the Bez class
+from soccer_pycontrol.pybullet_usage.pybullet_world import PybulletWorld  # Import the PybulletWorld class
+from soccer_common import Transformation
+
+
+class TestSensors(unittest.TestCase):
+
+    def setUp(self):
+        """
+        Set up the PyBullet environment and the robot with sensors.
+        """
+        self.world = PybulletWorld(
+            camera_yaw=90,
+            real_time=True,  # Change to False if you don't want real time simulation
+            rate=200,
+        )
+
+        self.robot_model = "bez2"
+        self.bez = Bez(robot_model=self.robot_model, pose=Transformation())
+
+    def tearDown(self):
+        """
+        Close the PyBullet environment.
+        """
+        self.world.close()
+        del self.bez
+        del self.world
+
+    def test_camera_image_capture(self):
+        """
+        Test to capture an image from the robot's camera using the Sensors class.
+        """
+        image = self.bez.sensors.get_camera_image()
+
+        # self.assertEqual(image.shape, (480, 640, 3), "Captured image does not have the correct dimensions.")
+
+        # Display the captured image
+        plt.imshow(image)
+        plt.title('Captured Camera Image')
+        plt.axis('off')
+        plt.show()
+
+if __name__ == "__main__":
+    unittest.main()