Perception (#547)

* New branch created for instance camera/merging of Point-Clout LiDAR

* Finished instance camera and related APIs

* Formatted

* Fixing for tests

* Fixing for format

* Accomadate according to PR Review by Quanyi

* Formatting

metadrive/component/sensors/instance_camera.py

@@ -0,0 +1,48 @@
+from metadrive.component.sensors.semantic_camera import SemanticCamera
+import cv2
+from panda3d.core import GeoMipTerrain, PNMImage
+from panda3d.core import RenderState, LightAttrib, ColorAttrib, ShaderAttrib, TextureAttrib, LVecBase4, MaterialAttrib
+from metadrive.constants import Semantics
+from metadrive.component.sensors.base_camera import BaseCamera
+from metadrive.constants import CamMask
+from metadrive.constants import RENDER_MODE_NONE
+from metadrive.engine.asset_loader import AssetLoader
+from metadrive.engine.engine_utils import get_engine
+import random
+
+
+class InstanceCamera(SemanticCamera):
+    CAM_MASK = CamMask.SemanticCam
+
+    def __init__(self, width, height, engine, *, cuda=False):
+        super().__init__(width, height, engine, cuda=cuda)
+
+    def track(self, base_object):
+        self._setup_effect()
+        super().track(base_object)
+
+    def _setup_effect(self):
+        """
+        Use tag to apply color to different object class
+        Returns: None
+
+        """
+        # setup camera
+
+        if get_engine() is None:
+            super()._setup_effect()
+        else:
+            mapping = get_engine().id_c
+            spawned_objects = get_engine().get_objects()
+            for id, obj in spawned_objects.items():
+                obj.origin.setTag("id", id)
+            cam = self.get_cam().node()
+            cam.setTagStateKey("id")
+            cam.setInitialState(
+                RenderState.make(
+                    ShaderAttrib.makeOff(), LightAttrib.makeAllOff(), TextureAttrib.makeOff(),
+                    ColorAttrib.makeFlat((0, 0, 0, 1)), 1
+                )
+            )
+            for id, c in mapping.items():
+                cam.setTagState(id, RenderState.make(ColorAttrib.makeFlat((c[0], c[1], c[2], 1)), 1))

metadrive/engine/base_engine.py

@@ -9,6 +9,7 @@
 from collections import OrderedDict
 from typing import Callable, Optional, Union, List, Dict, AnyStr
 
+from collections import deque
 import numpy as np
 from panda3d.core import NodePath, Vec3
 
@@ -22,6 +23,21 @@
 logger = logging.getLogger(__name__)
 
 
+def generate_distinct_rgb_values():
+    distinct_rgb_values = []
+    step = 256 // 32  # 8 intervals for each RGB component (0-31, 32-63, ..., 224-255)
+
+    for r in range(step, 256, step):
+        for g in range(0, 256, step):
+            for b in range(0, 256, step):
+                distinct_rgb_values.append((round(r / 255, 5), round(g / 255, 5), round(b / 255, 5)))
+
+    return distinct_rgb_values[:4096]  # Return the first 4096 values
+
+
+COLOR_SPACE = generate_distinct_rgb_values()
+
+
 class BaseEngine(EngineCore, Randomizable):
     """
     Due to the feature of Panda3D, BaseEngine should only be created once(Singleton Pattern)
@@ -31,7 +47,13 @@ class BaseEngine(EngineCore, Randomizable):
     singleton = None
     global_random_seed = None
 
+    MAX_COLOR = len(COLOR_SPACE)
+    COLORS_OCCUPIED = set()
+    COLORS_FREE = set(COLOR_SPACE)
+
     def __init__(self, global_config):
+        self.c_id = dict()
+        self.id_c = dict()
         self.try_pull_asset()
         EngineCore.__init__(self, global_config)
         Randomizable.__init__(self, self.global_random_seed)
@@ -154,9 +176,64 @@ def spawn_object(
         if self.global_config["record_episode"] and not self.replay_episode and record:
             self.record_manager.add_spawn_info(obj, object_class, kwargs)
         self._spawned_objects[obj.id] = obj
+        color = self.pick_color(obj.id)
+        if color == (-1, -1, -1):
+            print("FK!~")
+            exit()
+
         obj.attach_to_world(self.pbr_worldNP if pbr_model else self.worldNP, self.physics_world)
         return obj
 
+    def pick_color(self, id):
+        """
+        Return a color multiplier representing a unique color for an object if some colors are available.
+        Return -1,-1,-1 if no color available
+
+        SideEffect: COLOR_PTR will no longer point to the available color
+        SideEffect: COLORS_OCCUPIED[COLOR_PTR] will not be avilable
+        """
+        if len(BaseEngine.COLORS_OCCUPIED) == BaseEngine.MAX_COLOR:
+            return (-1, -1, -1)
+        assert (len(BaseEngine.COLORS_FREE) > 0)
+        my_color = BaseEngine.COLORS_FREE.pop()
+        BaseEngine.COLORS_OCCUPIED.add(my_color)
+        #print("After picking:", len(BaseEngine.COLORS_OCCUPIED), len(BaseEngine.COLORS_FREE))
+        self.id_c[id] = my_color
+        self.c_id[my_color] = id
+        return my_color
+
+    def clean_color(self, id):
+        """
+        Relinquish a color once the object is focibly destroyed
+        SideEffect:
+        BaseEngins.COLORS_OCCUPIED += 1
+        BaseEngine.COLOR_PTR now points to the idx just released
+        BaseEngine.COLORS_RECORED
+        Mapping Destroyed
+
+        """
+        if id in self.id_c.keys():
+            my_color = self.id_c[id]
+            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)
+
+    def id_to_color(self, id):
+        if id in self.id_c.keys():
+            return self.id_c[id]
+        else:
+            print("Invalid ID: ", id)
+            return -1, -1, -1
+
+    def color_to_id(self, color):
+        if color in self.c_id.keys():
+            return self.c_id[color]
+        else:
+            print("Invalid color:", color)
+            return "NA"
+
     def get_objects(self, filter: Optional[Union[Callable, List]] = None):
         """
         Return objects spawned, default all objects. Filter_func will be applied on all objects.
@@ -194,6 +271,10 @@ def clear_objects(self, filter: Optional[Union[Callable, List]], force_destroy=F
         If force_destroy=True, we will destroy this element instead of storing them for next time using
 
         filter: A list of object ids or a function returning a list of object id
+        """
+        """
+        In addition, we need to remove a color mapping whenever an object is destructed.
+        
         """
         force_destroy_this_obj = True if force_destroy or self.global_config["force_destroy"] else False
 
@@ -214,6 +295,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)
                 obj.destroy()
             else:
                 obj.detach_from_world(self.physics_world)
@@ -228,6 +310,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)
                     obj.destroy()
             if self.global_config["record_episode"] and not self.replay_episode and record:
                 self.record_manager.add_clear_info(obj)
@@ -243,6 +326,7 @@ def clear_object_if_possible(self, obj, force_destroy):
                 obj in self._dying_objects[obj.class_name]:
             self._dying_objects[obj.class_name].remove(obj)
             if hasattr(obj, "destroy"):
+                self.clean_color(obj.id)
                 obj.destroy()
         del obj
 
@@ -343,6 +427,24 @@ def process_memory():
         for _ in range(5):
             self.graphicsEngine.renderFrame()
 
+        #reset colors
+        BaseEngine.COLORS_FREE = set(COLOR_SPACE)
+        BaseEngine.COLORS_OCCUPIED = set()
+        new_i2c = {}
+        new_c2i = {}
+        #print("rest objects", len(self.get_objects()))
+        for object in self.get_objects().values():
+            if object.id in self.id_c.keys():
+                id = object.id
+                color = self.id_c[object.id]
+                BaseEngine.COLORS_OCCUPIED.add(color)
+                BaseEngine.COLORS_FREE.remove(color)
+                new_i2c[id] = color
+                new_c2i[color] = id
+        #print(len(BaseEngine.COLORS_FREE), len(BaseEngine.COLORS_OCCUPIED))
+        self.c_id = new_c2i
+        self.id_c = new_i2c
+
     def before_step(self, external_actions: Dict[AnyStr, np.array]):
         """
         Entities make decision here, and prepare for step
@@ -449,9 +551,11 @@ def close(self):
                 self._object_policies.pop(id).destroy()
             if id in self._object_tasks:
                 self._object_tasks.pop(id).destroy()
+            self.clean_color(obj.id)
             obj.destroy()
         for cls, pending_obj in self._dying_objects.items():
             for obj in pending_obj:
+                self.clean_color(obj.id)
                 obj.destroy()
         if self.main_camera is not None:
             self.main_camera.destroy()

metadrive/tests/test_sensors/test_instance_cam.py

@@ -0,0 +1,86 @@
+import pytest
+
+from metadrive.component.sensors.instance_camera import InstanceCamera
+from metadrive.envs.metadrive_env import MetaDriveEnv
+import numpy as np
+blackbox_test_configs = dict(
+    # standard=dict(stack_size=3, width=256, height=128, rgb_clip=True),
+    small=dict(stack_size=1, width=64, height=32, rgb_clip=True),
+)
+
+
+@pytest.mark.parametrize("config", list(blackbox_test_configs.values()), ids=list(blackbox_test_configs.keys()))
+def test_instance_cam(config, render=False):
+    """
+    Test the output shape of Instance camera. This can NOT make sure the correctness of rendered image but only for
+    checking the shape of image output and image retrieve pipeline
+    Args:
+        config: test parameter
+        render: render with cv2
+
+    Returns: None
+
+    """
+    env = MetaDriveEnv(
+        {
+            "num_scenarios": 1,
+            "traffic_density": 0.1,
+            "map": "S",
+            "show_terrain": False,
+            "start_seed": 4,
+            "stack_size": config["stack_size"],
+            "vehicle_config": dict(image_source="camera"),
+            "sensors": {
+                "camera": (InstanceCamera, config["width"], config["height"])
+            },
+            "interface_panel": ["dashboard", "camera"],
+            "image_observation": True,  # it is a switch telling metadrive to use rgb as observation
+            "rgb_clip": config["rgb_clip"],  # clip rgb to range(0,1) instead of (0, 255)
+        }
+    )
+    env.reset()
+    base_free = len(env.engine.COLORS_FREE)
+    base_occupied = len(env.engine.COLORS_OCCUPIED)
+    assert base_free + base_occupied == 4096
+    try:
+        import cv2
+        import time
+        start = time.time()
+        for i in range(1, 10):
+            o, r, tm, tc, info = env.step([0, 1])
+            assert env.observation_space.contains(o)
+            # Reverse
+            assert o["image"].shape == (
+                config["height"], config["width"], InstanceCamera.num_channels, config["stack_size"]
+            )
+            image = o["image"][..., -1]
+            image = image.reshape(-1, 3)
+            unique_colors = np.unique(image, axis=0)
+            #Making sure every color observed correspond to an object
+            for unique_color in unique_colors:
+                if (unique_color != np.array((0, 0, 0))).all():  #Ignore the black background.
+                    color = unique_color.tolist()
+                    color = (
+                        round(color[2], 5), round(color[1], 5), round(color[0], 5)
+                    )  #In engine, we use 5-diigt float for keys
+                    assert color in env.engine.COLORS_OCCUPIED
+                    assert color not in env.engine.COLORS_FREE
+                    assert color in env.engine.c_id.keys()
+                    assert env.engine.id_c[env.engine.c_id[color]] == color  #Making sure the color-id is a bijection
+                    assert len(env.engine.c_id.keys()) == len(env.engine.COLORS_OCCUPIED)
+                    assert len(env.engine.id_c.keys()) == len(env.engine.COLORS_OCCUPIED)
+                    assert len(env.engine.COLORS_FREE) + len(env.engine.COLORS_OCCUPIED) == 4096
+            #Making sure every object in the engine(not necessarily observable) have corresponding color
+            for id, object in env.engine.get_objects().items():
+                assert id in env.engine.id_c.keys()
+            if render:
+                cv2.imshow('img', o["image"][..., -1])
+                cv2.waitKey(1)
+        print("FPS:", 10 / (time.time() - start))
+    finally:
+        env.close()
+
+
+if __name__ == '__main__':
+    test_instance_cam(config=blackbox_test_configs["small"], render=True)
+    my_dict = {(0, 0, 0): "Hello, World!"}