diff --git a/soccer_control/soccer_pycontrol/src/soccer_pycontrol/model/sensors.py b/soccer_control/soccer_pycontrol/src/soccer_pycontrol/model/sensors.py index e7ef0734b..d9ec04eb6 100644 --- a/soccer_control/soccer_pycontrol/src/soccer_pycontrol/model/sensors.py +++ b/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 \ No newline at end of file diff --git a/soccer_control/soccer_pycontrol/test/test_perception.py b/soccer_control/soccer_pycontrol/test/test_perception.py new file mode 100644 index 000000000..5722a181f --- /dev/null +++ b/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() \ No newline at end of file