Update the script for generating multi-sensors images (#61)

* fix broken script

* format code

* introduce offscreen rendering

* try debug

* fix

* fix

* up

* up

* remove fix

* fix

* WIP

scenarionet/tests/script/generate_sensor.py

@@ -1,47 +1,31 @@
+import pathlib
+
 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.component.sensors.semantic_camera import SemanticCamera
 from metadrive.engine.asset_loader import AssetLoader
 from metadrive.envs.scenario_env import ScenarioEnv
 from metadrive.policy.replay_policy import ReplayEgoCarPolicy
 
-NuScenesEnv = ScenarioEnv
-
 if __name__ == "__main__":
-    env = NuScenesEnv(
+    env = ScenarioEnv(
         {
+            # To enable onscreen rendering, set this config to True.
             "use_render": True,
+
+            # !!!!! To enable offscreen rendering, set this config to True !!!!!
+            "image_observation": False,
+
+            # ===== The scenario and MetaDrive config =====
             "agent_policy": ReplayEgoCarPolicy,
-            "show_interface": True,
-            "interface_panel": ["semantic_camera"],
-            # "need_lane_localization": False,
-            "show_logo": False,
             "no_traffic": False,
-            "drivable_region_extension": 15,
             "sequential_seed": True,
             "reactive_traffic": False,
-            "show_fps": False,
-            # "debug": True,
-            "render_pipeline": False,
-            "daytime": "08:10",
-            "window_size": (800, 450),
-            "camera_dist": 0.8,
-            "camera_height": 1.5,
-            "camera_pitch": None,
-            "camera_fov": 60,
             "start_scenario_index": 0,
             "num_scenarios": 10,
-            # "force_reuse_object_name": True,
-            # "data_directory": "/home/shady/Downloads/test_processed",
             "horizon": 1000,
             "no_static_vehicles": False,
-            # "show_policy_mark": True,
-            # "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(
                 show_navi_mark=False,
                 use_special_color=False,
@@ -52,25 +36,56 @@ if __name__ == "__main__":
             ),
             "data_directory": AssetLoader.file_path("nuscenes", unix_style=False),
 
-            # "image_observation": True,
+            # ===== Set some sensor and visualization configs =====
+            "daytime": "08:10",
+            "window_size": (800, 450),
+            "camera_dist": 0.8,
+            "camera_height": 1.5,
+            "camera_pitch": None,
+            "camera_fov": 60,
+            "interface_panel": ["semantic_camera"],
+            "sensors": dict(
+                semantic_camera=(SemanticCamera, 1600, 900),
+                depth_camera=(DepthCamera, 800, 600),
+                rgb_camera=(RGBCamera, 800, 600),
+            ),
+
+            # ===== Remove useless items in the images =====
+            "show_logo": False,
+            "show_fps": False,
+            "show_interface": True,
         }
     )
 
-    # 0,1,3,4,5,6
+    file_dir = pathlib.Path("images")
-    # for seed in range(10):
+    file_dir.mkdir(exist_ok=True)
-    while True:
+
+    for ep in range(1):
         env.reset(seed=6)
+
+        # Run it once to initialize the TopDownRenderer
+        env.render(
+            mode="topdown",
+            screen_size=(1600, 900),
+            film_size=(9000, 9000),
+            target_vehicle_heading_up=True,
+            semantic_map=True,
+        )
+
         for t in range(10000):
-            # env.capture("rgb_deluxe_{}_{}.jpg".format(env.current_seed, t))
+            env.capture("rgb_deluxe_{}_{}.jpg".format(env.current_seed, t))
-            # ret = env.render(
+            ret = env.render(
-            #     mode="topdown", screen_size=(1600, 900), film_size=(9000, 9000), target_vehicle_heading_up=True
+                mode="topdown",
-            # )
+                screen_size=(1600, 900),
-            # pygame.image.save(ret, "top_down_{}_{}.png".format(env.current_seed, t))
+                film_size=(9000, 9000),
-            # env.vehicle.get_camera("depth_camera").save_image(env.vehicle, "depth_{}.jpg".format(t))
+                target_vehicle_heading_up=True,
-            # env.vehicle.get_camera("rgb_camera").save_image(env.vehicle, "rgb_{}.jpg".format(t))
+                semantic_map=True,
-            if t == 100:
+                to_image=False
-                # 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))
+            pygame.image.save(ret, str(file_dir / "bev_{}.png".format(t)))
+            env.engine.get_sensor("depth_camera").save_image(env.agent, str(file_dir / "depth_{}.jpg".format(t)))
+            env.engine.get_sensor("rgb_camera").save_image(env.agent, str(file_dir / "rgb_{}.jpg".format(t)))
+            env.engine.get_sensor("semantic_camera").save_image(env.agent, str(file_dir / "semantic_{}.jpg".format(t)))
+            if t == 30:
                 break
             env.step([1, 0.88])
-            # if d:

scenarionet/tests/script/generate_sensor_offscreen.py

@@ -0,0 +1,95 @@
+import pathlib
+
+import pygame
+from metadrive.component.sensors.depth_camera import DepthCamera
+from metadrive.component.sensors.rgb_camera import RGBCamera
+from metadrive.component.sensors.semantic_camera import SemanticCamera
+from metadrive.engine.asset_loader import AssetLoader
+from metadrive.envs.scenario_env import ScenarioEnv
+from metadrive.policy.replay_policy import ReplayEgoCarPolicy
+
+if __name__ == "__main__":
+    env = ScenarioEnv(
+        {
+            # To enable onscreen rendering, set this config to True.
+            "use_render": False,
+
+            # !!!!! To enable offscreen rendering, set this config to True !!!!!
+            "image_observation": True,
+
+            # ===== The scenario and MetaDrive config =====
+            "agent_policy": ReplayEgoCarPolicy,
+            "no_traffic": False,
+            "sequential_seed": True,
+            "reactive_traffic": False,
+            "start_scenario_index": 0,
+            "num_scenarios": 10,
+            "horizon": 1000,
+            "no_static_vehicles": False,
+            "vehicle_config": dict(
+                show_navi_mark=False,
+                use_special_color=False,
+                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", unix_style=False),
+
+            # ===== Set some sensor and visualization configs =====
+            "daytime": "08:10",
+            "window_size": (800, 450),
+            "camera_dist": 0.8,
+            "camera_height": 1.5,
+            "camera_pitch": None,
+            "camera_fov": 60,
+            "interface_panel": ["semantic_camera"],
+            "sensors": dict(
+                semantic_camera=(SemanticCamera, 1600, 900),
+                depth_camera=(DepthCamera, 800, 600),
+                rgb_camera=(RGBCamera, 800, 600),
+            ),
+
+            # ===== Remove useless items in the images =====
+            "show_logo": False,
+            "show_fps": False,
+            "show_interface": True,
+        }
+    )
+
+    file_dir = pathlib.Path("images_offscreen")
+    file_dir.mkdir(exist_ok=True)
+
+    for ep in range(1):
+        env.reset(seed=6)
+
+        # Run it once to initialize the TopDownRenderer
+        env.render(
+            mode="topdown",
+            screen_size=(1600, 900),
+            film_size=(9000, 9000),
+            target_vehicle_heading_up=True,
+            semantic_map=True,
+        )
+
+        for t in range(10000):
+
+            # We don't have interface in offscreen. So comment out:
+            # 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,
+                semantic_map=True,
+                to_image=False
+            )
+            pygame.image.save(ret, str(file_dir / "bev_{}.png".format(t)))
+            env.engine.get_sensor("depth_camera").save_image(env.agent, str(file_dir / "depth_{}.jpg".format(t)))
+            env.engine.get_sensor("rgb_camera").save_image(env.agent, str(file_dir / "rgb_{}.jpg".format(t)))
+            env.engine.get_sensor("semantic_camera").save_image(env.agent, str(file_dir / "semantic_{}.jpg".format(t)))
+            print("Image at step {} is saved at: {}".format(t, file_dir))
+            if t == 30:
+                break
+            env.step([1, 0.88])