From a88e0bd83c15f097df2a5fdd593992497b236589 Mon Sep 17 00:00:00 2001
From: Quanyi Li <quanyili0057@gmail.com>
Date: Tue, 24 Oct 2023 19:20:43 +0100
Subject: [PATCH] Update Sensor API in scripts (#39)

* add semantic cam

* update API

* format
---
 scenarionet/tests/script/generate_sensor.py | 41 ++++++++++++---------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git a/scenarionet/tests/script/generate_sensor.py b/scenarionet/tests/script/generate_sensor.py
index 2f008de..1a76598 100644
--- a/scenarionet/tests/script/generate_sensor.py
+++ b/scenarionet/tests/script/generate_sensor.py
@@ -1,4 +1,7 @@
 import pygame
+from metadrive.component.sensors.semantic_camera import SemanticCamera
+from metadrive.component.sensors.depth_camera import DepthCamera
+from metadrive.component.sensors.rgb_camera import RGBCamera
 from metadrive.engine.asset_loader import AssetLoader
 from metadrive.envs.scenario_env import ScenarioEnv
 from metadrive.policy.replay_policy import ReplayEgoCarPolicy
@@ -10,7 +13,8 @@ if __name__ == "__main__":
         {
             "use_render": True,
             "agent_policy": ReplayEgoCarPolicy,
-            "show_interface": False,
+            "show_interface": True,
+            "interface_panel": ["semantic_camera"],
             # "need_lane_localization": False,
             "show_logo": False,
             "no_traffic": False,
@@ -19,9 +23,9 @@ if __name__ == "__main__":
             "reactive_traffic": False,
             "show_fps": False,
             # "debug": True,
-            "render_pipeline": True,
+            "render_pipeline": False,
             "daytime": "08:10",
-            "window_size": (1600, 900),
+            "window_size": (800, 450),
             "camera_dist": 0.8,
             "camera_height": 1.5,
             "camera_pitch": None,
@@ -36,36 +40,37 @@ if __name__ == "__main__":
             # "show_coordinates": True,
             # "force_destroy": True,
             # "default_vehicle_in_traffic": True,
+            "default_vehicle_in_traffic": True,
+            "sensors": dict(semantic_camera=(SemanticCamera, 1600, 900)),
             "vehicle_config": dict(
-                # light=True,
-                # random_color=True,
                 show_navi_mark=False,
                 use_special_color=False,
-                image_source="depth_camera",
-                rgb_camera=(1600, 900),
-                depth_camera=(1600, 900, True),
-                # no_wheel_friction=True,
+                image_source="semantic_camera",
                 lidar=dict(num_lasers=120, distance=50),
                 lane_line_detector=dict(num_lasers=0, distance=50),
                 side_detector=dict(num_lasers=12, distance=50)
             ),
             "data_directory": AssetLoader.file_path("nuscenes", return_raw_style=False),
+
             # "image_observation": True,
         }
     )
 
     # 0,1,3,4,5,6
-    for seed in range(10):
-        env.reset(seed=seed)
+    # for seed in range(10):
+    while True:
+        env.reset(seed=6)
         for t in range(10000):
-            env.capture("rgb_deluxe_{}_{}.jpg".format(env.current_seed, t))
-            ret = env.render(
-                mode="topdown", screen_size=(1600, 900), film_size=(9000, 9000), target_vehicle_heading_up=True
-            )
-            pygame.image.save(ret, "top_down_{}_{}.png".format(env.current_seed, t))
+            # env.capture("rgb_deluxe_{}_{}.jpg".format(env.current_seed, t))
+            # ret = env.render(
+            #     mode="topdown", screen_size=(1600, 900), film_size=(9000, 9000), target_vehicle_heading_up=True
+            # )
+            # pygame.image.save(ret, "top_down_{}_{}.png".format(env.current_seed, t))
             # env.vehicle.get_camera("depth_camera").save_image(env.vehicle, "depth_{}.jpg".format(t))
             # env.vehicle.get_camera("rgb_camera").save_image(env.vehicle, "rgb_{}.jpg".format(t))
+            if t == 100:
+                # env.engine.get_sensor("semantic_camera").save_image(env.vehicle, "depth_{}.jpg".format(t))
+                env.engine.get_sensor("semantic_camera").save_image(env.vehicle, "semantic_{}.jpg".format(t))
+                break
             env.step([1, 0.88])
             # if d:
-            if env.episode_step >= env.engine.data_manager.current_scenario_length:
-                break