diff --git a/metadrive/component/sensors/instance_camera.py b/metadrive/component/sensors/instance_camera.py index 9adb6a780..8a2816ed0 100644 --- a/metadrive/component/sensors/instance_camera.py +++ b/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,6 +50,12 @@ 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() @@ -41,8 +63,10 @@ def _setup_effect(self): 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)) diff --git a/metadrive/engine/base_engine.py b/metadrive/engine/base_engine.py index 24b2b9f63..5e7cb415b 100644 --- a/metadrive/engine/base_engine.py +++ b/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,6 +264,7 @@ 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) @@ -268,7 +272,7 @@ def clear_objects(self, filter: Optional[Union[Callable, List]], force_destroy=F 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)