Fix Instance_Camera not updating/color recycling

metadrive/component/sensors/instance_camera.py

@@ -1,3 +1,4 @@
+from numpy import ndarray
 from panda3d.core import RenderState, LightAttrib, ColorAttrib, ShaderAttrib, TextureAttrib, FrameBufferProperties
 
 from metadrive.component.sensors.base_camera import BaseCamera
@@ -6,6 +7,13 @@
 from metadrive.engine.engine_utils import get_engine
 
 
+from typing import Union
+
+import cv2
+import numpy as np
+from panda3d.core import NodePath
+
+
 class InstanceCamera(BaseCamera):
     CAM_MASK = CamMask.SemanticCam
 
@@ -20,10 +28,18 @@ def track(self, new_parent_node, position, hpr):
         """
         self._setup_effect()
         super().track(new_parent_node, position, hpr)
+
+
+
+    def perceive(
+        self, to_float=True, new_parent_node: Union[NodePath, None] = None, position=None, hpr=None
+    ) -> np.ndarray:
+        self._setup_effect()
+        return super().perceive(to_float, new_parent_node, position, hpr)
 
     def _setup_effect(self):
         """
-        Use tag to apply color to different object class
+        Use tag to apply color to different objects(instances)
         Returns: None
 
         """
@@ -34,15 +50,23 @@ def _setup_effect(self):
         else:
             mapping = get_engine().id_c
             spawned_objects = get_engine().get_objects()
+
+            ##Ensure consistency between color mapping and the objects actually active in the engine.
+            mapping_set, object_set = set(list(mapping.keys())),set(list(spawned_objects.keys()))
+            assert(len(mapping_set.difference(object_set))==0)
+
+
             for id, obj in spawned_objects.items():
                 obj.origin.setTag(CameraTagStateKey.ID, id)
             cam = self.get_cam().node()
             cam.setTagStateKey(CameraTagStateKey.ID)
             cam.setInitialState(
                 RenderState.make(
                     ShaderAttrib.makeOff(), LightAttrib.makeAllOff(), TextureAttrib.makeOff(),
-                    ColorAttrib.makeFlat((0, 0, 0, 1)), 1
+                    ColorAttrib.makeFlat((1, 1, 1, 1)), 1
                 )
             )
-            for id, c in mapping.items():
+
+            for id in spawned_objects.keys():
+                c = mapping[id]
                 cam.setTagState(id, RenderState.make(ColorAttrib.makeFlat((c[0], c[1], c[2], 1)), 1))

metadrive/engine/base_engine.py

@@ -187,11 +187,14 @@ def _clean_color(self, id):
         """
         if id in self.id_c.keys():
             my_color = self.id_c[id]
-            BaseEngine.COLORS_OCCUPIED.remove(my_color)
+            if my_color in BaseEngine.COLORS_OCCUPIED:
+                BaseEngine.COLORS_OCCUPIED.remove(my_color)
             BaseEngine.COLORS_FREE.add(my_color)
             # print("After cleaning:,", len(BaseEngine.COLORS_OCCUPIED), len(BaseEngine.COLORS_FREE))
-            self.id_c.pop(id)
-            self.c_id.pop(my_color)
+            if id in self.id_c.keys():
+                self.id_c.pop(id)
+            if my_color in self.c_id.keys():
+                self.c_id.pop(my_color)
 
     def id_to_color(self, id):
         if id in self.id_c.keys():
@@ -261,14 +264,15 @@ def clear_objects(self, filter: Optional[Union[Callable, List]], force_destroy=F
         else:
             raise ValueError("filter should be a list or a function")
         for id, obj in exclude_objects.items():
+            self._clean_color(id)
             self._spawned_objects.pop(id)
             if id in self._object_tasks:
                 self._object_tasks.pop(id)
             if id in self._object_policies:
                 policy = self._object_policies.pop(id)
                 policy.destroy()
             if force_destroy_this_obj:
-                self._clean_color(obj.id)
+                #self._clean_color(obj.id)
                 obj.destroy()
             else:
                 obj.detach_from_world(self.physics_world)
@@ -283,7 +287,7 @@ def clear_objects(self, filter: Optional[Union[Callable, List]], force_destroy=F
                 if len(self._dying_objects[obj.class_name]) < self.global_config["num_buffering_objects"]:
                     self._dying_objects[obj.class_name].append(obj)
                 else:
-                    self._clean_color(obj.id)
+                    #self._clean_color(obj.id)
                     obj.destroy()
             if self.global_config["record_episode"] and not self.replay_episode and record:
                 self.record_manager.add_clear_info(obj)