fix camera render shenanigans

jjshoots · Jul 25, 2023 · 42cb219 · 42cb219
1 parent 5f6ba72
commit 42cb219
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 16 deletions.
diff --git a/PyFlyt/gym_envs/fixedwing_envs/fixedwing_base_env.py b/PyFlyt/gym_envs/fixedwing_envs/fixedwing_base_env.py
@@ -271,12 +271,10 @@ def render(self):
         ), "Please set `render_mode='human'` or `render_mode='rgb_array'` to use this function."
 
         _, _, rgbaImg, _, _ = self.env.getCameraImage(
-            width=self.render_resolution[0],
-            height=self.render_resolution[1],
+            width=self.render_resolution[1],
+            height=self.render_resolution[0],
             viewMatrix=self.env.drones[0].camera.view_mat,
             projectionMatrix=self.env.drones[0].camera.proj_mat,
         )
 
-        return np.array(rgbaImg, dtype=np.uint8).reshape(
-            self.render_resolution[1], self.render_resolution[0], -1
-        )
+        return rgbaImg
diff --git a/PyFlyt/gym_envs/quadx_envs/quadx_base_env.py b/PyFlyt/gym_envs/quadx_envs/quadx_base_env.py
@@ -265,12 +265,10 @@ def render(self):
         ), "Please set `render_mode='human'` or `render_mode='rgb_array'` to use this function."
 
         _, _, rgbaImg, _, _ = self.env.getCameraImage(
-            width=self.render_resolution[0],
-            height=self.render_resolution[1],
+            width=self.render_resolution[1],
+            height=self.render_resolution[0],
             viewMatrix=self.camera_parameters[2],
             projectionMatrix=self.camera_parameters[3],
         )
 
-        return np.array(rgbaImg, dtype=np.uint8).reshape(
-            self.render_resolution[1], self.render_resolution[0], -1
-        )
+        return rgbaImg
diff --git a/PyFlyt/gym_envs/rocket_envs/rocket_base_env.py b/PyFlyt/gym_envs/rocket_envs/rocket_base_env.py
@@ -327,12 +327,9 @@ def render(self):
         ), "Please set `render_mode='human'` or `render_mode='rgb_array'` to use this function."
 
         _, _, rgbaImg, _, _ = self.env.getCameraImage(
-            width=self.render_resolution[0],
-            height=self.render_resolution[1],
+            width=self.render_resolution[1],
+            height=self.render_resolution[0],
             viewMatrix=self.env.drones[0].camera.view_mat,
             projectionMatrix=self.env.drones[0].camera.proj_mat,
         )
-
-        return np.array(rgbaImg, dtype=np.uint8).reshape(
-            self.render_resolution[1], self.render_resolution[0], -1
-        )
+        return rgbaImg
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -114,6 +114,14 @@ def test_camera():
     for i in range(100):
         env.step()
 
+        # check the camera image
+        assert isinstance(env.drones[0].rgbaImg, np.ndarray), f"Expected camera image to be of type `np.ndarray`, got {type(env.drones[0].rgbaImg)}."
+        assert isinstance(env.drones[0].depthImg, np.ndarray), f"Expected depth image to be of type `np.ndarray`, got {type(env.drones[0].depthImg)}."
+        assert isinstance(env.drones[0].segImg, np.ndarray), f"Expected segmented image to be of type `np.ndarray`, got {type(env.drones[0].segImg)}."
+        assert env.drones[0].rgbaImg.shape[-1] == 4, f"Expected 4 channels in the rendered image, got {env.drones[0].rgbaImg.shape[-1]}."
+        assert len(env.drones[0].depthImg.shape) == 2, f"Expected depth image to only have 2 dimensions, got {len(env.drones[0].depthImg.shape)}."
+        assert len(env.drones[0].segImg.shape) == 2, f"Expected segmented image to have only 2 dimensions, got {len(env.drones[0].segImg.shape)}"
+
     env.disconnect()
 
 

diff --git a/tests/test_gym_envs.py b/tests/test_gym_envs.py
@@ -5,6 +5,7 @@
 import pytest
 from gymnasium.error import Error
 from gymnasium.utils.env_checker import check_env, data_equivalence
+import numpy as np
 
 import PyFlyt.gym_envs
 
@@ -161,4 +162,8 @@ def test_render(env_config):
         frames.append(env.render())
         env.step(env.action_space.sample())
 
+    for frame in frames:
+        assert isinstance(frame, np.ndarray), f"Expected render frames to be of type `np.ndarray`, got {type(frame)}."
+        assert frame.shape[-1] == 4, f"Expected 4 channels in the rendered image, got {frame.shape[-1]}."
+
     env.close()