Fix semantic camera bug

metadrive/component/sensors/base_camera.py

@@ -30,10 +30,16 @@ class BaseCamera(ImageBuffer, BaseSensor):
     display_region_size = [1 / 3, 2 / 3, 0.8, 1.0]
     attached_object = None
 
-    def __init__(self, engine, setup_pbr=False, need_cuda=False):
+    def __init__(self, engine, setup_pbr=False, need_cuda=False, frame_buffer_property=None):
         self._enable_cuda = need_cuda
         super(BaseCamera, self).__init__(
-            self.BUFFER_W, self.BUFFER_H, Vec3(0., 0.8, 1.5), self.BKG_COLOR, setup_pbr=setup_pbr, engine=engine
+            self.BUFFER_W,
+            self.BUFFER_H,
+            Vec3(0., 0.8, 1.5),
+            self.BKG_COLOR,
+            setup_pbr=setup_pbr,
+            engine=engine,
+            frame_buffer_property=frame_buffer_property
         )
 
         width = self.BUFFER_W

metadrive/component/sensors/depth_camera.py

@@ -6,6 +6,7 @@
 from metadrive.constants import CamMask
 from metadrive.constants import RENDER_MODE_NONE
 from metadrive.engine.asset_loader import AssetLoader
+from panda3d.core import FrameBufferProperties
 
 
 class DepthCamera(BaseCamera):
@@ -20,6 +21,9 @@ class DepthCamera(BaseCamera):
     def __init__(self, width, height, engine, *, cuda=False):
         self.BUFFER_W, self.BUFFER_H = width, height
         self.VIEW_GROUND = True  # default true
+        frame_buffer_property = FrameBufferProperties()
+        frame_buffer_property.set_rgba_bits(8, 8, 8, 0)  # disable alpha for RGB camera
+        # TODO It can be made more efficient by only using one channel
         super(DepthCamera, self).__init__(engine, False, cuda)
         cam = self.get_cam()
         lens = self.get_lens()

metadrive/component/sensors/mini_map.py

@@ -2,6 +2,7 @@
 
 from metadrive.component.sensors.base_camera import BaseCamera
 from metadrive.constants import CamMask
+from panda3d.core import FrameBufferProperties
 
 
 class MiniMap(BaseCamera):
@@ -10,7 +11,9 @@ class MiniMap(BaseCamera):
 
     def __init__(self, width, height, z_pos, engine, *, cuda=False):
         self.BUFFER_W, self.BUFFER_H, height = width, height, z_pos
-        super(MiniMap, self).__init__(engine=engine, need_cuda=cuda)
+        frame_buffer_property = FrameBufferProperties()
+        frame_buffer_property.set_rgba_bits(8, 8, 8, 0)  # disable alpha for RGB camera
+        super(MiniMap, self).__init__(engine=engine, need_cuda=cuda, frame_buffer_property=frame_buffer_property)
 
         cam = self.get_cam()
         lens = self.get_lens()

metadrive/component/sensors/rgb_camera.py

@@ -2,6 +2,7 @@
 from metadrive.constants import CamMask
 from metadrive.engine.engine_utils import engine_initialized, get_global_config
 from direct.filter.CommonFilters import CommonFilters
+from panda3d.core import FrameBufferProperties
 
 
 class RGBCamera(BaseCamera):
@@ -13,7 +14,9 @@ class RGBCamera(BaseCamera):
 
     def __init__(self, width, height, engine, *, cuda=False):
         self.BUFFER_W, self.BUFFER_H = width, height
-        super(RGBCamera, self).__init__(engine, True, cuda)
+        frame_buffer_property = FrameBufferProperties()
+        frame_buffer_property.set_rgba_bits(8, 8, 8, 0)  # disable alpha for RGB camera
+        super(RGBCamera, self).__init__(engine, True, cuda, frame_buffer_property=frame_buffer_property)
         cam = self.get_cam()
         lens = self.get_lens()
         # cam.lookAt(0, 2.4, 1.3)

metadrive/component/sensors/semantic_camera.py

@@ -22,6 +22,7 @@ class SemanticCamera(BaseCamera):
     def __init__(self, width, height, engine, *, cuda=False):
         self.BUFFER_W, self.BUFFER_H = width, height
         self.VIEW_GROUND = True  # default true
+        # The framebuffer can not be 3 channel like RGB Camera...
         super(SemanticCamera, self).__init__(engine, False, cuda)
         cam = self.get_cam()
         lens = self.get_lens()
@@ -66,7 +67,7 @@ def __init__(self, width, height, engine, *, cuda=False):
             self.GROUND.generate()
 
     def track(self, base_object):
-        if self.VIEW_GROUND:
+        if self.VIEW_GROUND and base_object is not None:
             pos = base_object.origin.getPos()
             self.GROUND_MODEL.setPos(pos[0], pos[1], self.GROUND_HEIGHT)
             self.GROUND_MODEL.setH(base_object.origin.getH())

metadrive/component/vehicle/vehicle_type.py

@@ -314,6 +314,7 @@ def get_vehicle_type(length, np_random=None, need_default_vehicle=False):
         elif length <= 5.5:
             type_count[1] += 1
             vs = [LVehicle, MVehicle, SVehicle]
+            # vs = [SVehicle, LVehicle, MVehicle]
             if need_default_vehicle:
                 vs.append(TrafficDefaultVehicle)
             return vs[type_count[1] % len(vs)]

metadrive/engine/core/image_buffer.py

@@ -5,7 +5,7 @@
 from typing import Union, List
 
 import numpy as np
-from panda3d.core import NodePath, Vec3, Vec4, Camera, PNMImage, Shader, RenderState, ShaderAttrib, FrameBufferProperties
+from panda3d.core import NodePath, Vec3, Vec4, Camera, PNMImage, Shader, RenderState, ShaderAttrib
 
 from metadrive.constants import RENDER_MODE_ONSCREEN, BKG_COLOR, RENDER_MODE_NONE
 
@@ -52,9 +52,6 @@ def __init__(
             self.lens = self.cam.node().getLens()
             return
 
-        frame_buffer_property = FrameBufferProperties()
-        frame_buffer_property.set_rgba_bits(8, 8, 8, 0)  # disable alpha for RGB camera
-
         # self.texture = Texture()
         if frame_buffer_property is None:
             self.buffer = self.engine.win.makeTextureBuffer("camera", width, height)
@@ -105,9 +102,12 @@ def __init__(
     def get_rgb_array_cpu(self):
         origin_img = self.buffer.getDisplayRegion(1).getScreenshot()
         img = np.frombuffer(origin_img.getRamImage().getData(), dtype=np.uint8)
-        img = img.reshape((origin_img.getYSize(), origin_img.getXSize(), 3))
+        img = img.reshape((origin_img.getYSize(), origin_img.getXSize(), -1))
         # img = np.swapaxes(img, 1, 0)
         img = img[::-1]
+        if img.shape[-1] == 4:
+            # To 3 channel
+            img = img[..., :-1]
         return img
 
     @staticmethod

metadrive/manager/agent_manager.py

@@ -239,13 +239,20 @@ def try_actuate_agent(self, step_infos, stage="before_step"):
         assert stage == "before_step" or stage == "after_step"
         for agent_id in self.active_agents.keys():
             policy = self.get_policy(self._agent_to_object[agent_id])
-            cond_1 = stage == "before_step" and not isinstance(policy, ReplayTrafficParticipantPolicy)
-            cond_2 = stage == "after_step" and isinstance(policy, ReplayTrafficParticipantPolicy)
-            if cond_2 or cond_1:
-                assert policy is not None, "No policy is set for agent {}".format(agent_id)
-                action = policy.act(agent_id)
-                step_infos[agent_id] = policy.get_action_info()
-                step_infos[agent_id].update(self.get_agent(agent_id).before_step(action))
+            is_replay = isinstance(policy, ReplayTrafficParticipantPolicy)
+            assert policy is not None, "No policy is set for agent {}".format(agent_id)
+            if is_replay:
+                if stage == "after_step":
+                    policy.act(agent_id)
+                    step_infos[agent_id] = policy.get_action_info()
+                else:
+                    step_infos[agent_id] = self.get_agent(agent_id).before_step([0, 0])
+            else:
+                if stage == "before_step":
+                    action = policy.act(agent_id)
+                    step_infos[agent_id] = policy.get_action_info()
+                    step_infos[agent_id].update(self.get_agent(agent_id).before_step(action))
+
         return step_infos
 
     def before_step(self):

metadrive/tests/vis_functionality/vis_semantic_cam.py

@@ -20,6 +20,7 @@ def get_image(env):
             "use_render": True,
             "image_observation": True,
             "rgb_clip": True,
+            "show_interface": False,
             "agent_policy": ReplayEgoCarPolicy,
             "interface_panel": ["semantic_camera"],
             "sensors": dict(semantic_camera=(SemanticCamera, 800, 600)),