Skip to content

Commit

Permalink
Fix the example for generating videos (#555)
Browse files Browse the repository at this point in the history 
* clean code

* format

* up
  • Loading branch information
@pengzhenghao
pengzhenghao committed Nov 16, 2023
1 parent 0df2f97 commit 55ab1a7
Showing 1 changed file with 22 additions and 85 deletions.
107 changes: 22 additions & 85 deletions metadrive/examples/generate_video_for_bev_and_interface.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,51 +4,29 @@
example_video_RANDOM-STRING/0_bev.mp4
example_video_RANDOM-STRING/0_interface.mp4
This file demonstrates how to use API to generate high-resolution video for demo.
This file demonstrates how to use API to generate high-resolution & temporal aligned videos for visualization.
Dependencies:
pip install mediapy
conda install ffmpeg
https://itsfoss.com/install-h-264-decoder-ubuntu/ (h264 decoder in Linux, or you can install VLC)
"""
import os
import uuid
from datetime import datetime

import cv2
import matplotlib.pyplot as plt
import mediapy
import pygame

from metadrive.envs import MetaDriveEnv


def plot(d):
# Can use this function to show images, if you wish to set breakpoints in this file.
plt.imshow(d)
plt.show()


class VideoRecorder:
def __init__(self, video_name, code, height, width, fps=40):
"""
code=mp4v, avc1, x264, h264 etc.
"""
assert video_name.endswith(".mp4")
video = cv2.VideoWriter(video_name, cv2.VideoWriter_fourcc(*code), fps, (width, height))
self.video_writer = video

def write(self, image):
self.video_writer.write(image[..., ::-1])

def finish(self):
self.video_writer.release()
try:
cv2.destroyAllWindows()
except Exception as e:
print("Error happen: ", e)


if __name__ == '__main__':
num_scenarios = 1 # Will run each seed for 1 episode.
start_seed = 100
generate_video = True

if generate_video:
folder_name = "example_video_{}".format(str(uuid.uuid4())[:6])
folder_name = "example_video_{}".format(datetime.now().strftime("%Y-$m-%d_%H-%M-%S"))
video_bev = []
video_interface = []

env = MetaDriveEnv(
dict(
Expand All @@ -59,6 +37,7 @@ def finish(self):
window_size=(1600, 1200),
crash_vehicle_done=False,
manual_control=True, # For using expert policy. You don't need to control it.
horizon=10,
)
)

Expand All @@ -70,23 +49,12 @@ def finish(self):
env.vehicle.expert_takeover = True
env.engine.force_fps.disable()

if generate_video:
os.makedirs(folder_name, exist_ok=True)

video_base_name = "{}/{}".format(folder_name, ep_count)
video_name_bev = video_base_name + "_bev.mp4"
video_bev = VideoRecorder(video_name=video_name_bev, height=3000, width=3000, code="avc1")
# video_list_bev = []
video_name_interface = video_base_name + "_interface.mp4"
# video_list_interface = []
video_interface = VideoRecorder(video_name=video_name_interface, height=1200, width=1600, code="avc1")

while True:
o, r, tm, tc, info = env.step(env.action_space.sample())

img_interface = env.render("rgb_array")
img_interface = env.render(mode="rgb_array")
img_bev = env.render(
"topdown",
mode="topdown",
target_vehicle_heading_up=False,
draw_target_vehicle_trajectory=True,
film_size=(3000, 3000),
Expand All @@ -95,20 +63,10 @@ def finish(self):
)

if generate_video:
# Option 1: Directly feed image to video stream. Can save IO time and memory.
# Has the risk to generate broken video.
img_bev = pygame.surfarray.array3d(img_bev)
img_bev = img_bev.swapaxes(0, 1)
video_bev.write(img_bev)
video_interface.write(img_interface)

# Option 2: Save image to local temp folder and generate video later. Waste IO time but save memory.
# The most safe method.
# pygame.image.save(img_bev, "{}/{}.png".format(tmp_folder_bev, frame_count))
# env.capture(file_name="{}/{}.png".format(tmp_folder_interface, frame_count))

# Note that the interface image can also be used like this:
# pygame.image.save(img_interface, "{}/{}.png".format(tmp_folder_interface, frame_count))
video_bev.append(img_bev)
video_interface.append(img_interface)

frame_count += 1
step_count += 1
Expand All @@ -119,41 +77,20 @@ def finish(self):

env.engine.force_fps.disable()
if generate_video:
video_bev.finish()
video_interface.finish()
# image_list_to_video(video_name_bev, video_list_bev, code="avc1")
print("BEV Video is saved at: ", video_name_bev)
# image_list_to_video(video_name_interface, video_list_interface, code="avc1")
print("Interface Video is saved at: ", video_name_interface)
os.makedirs(folder_name, exist_ok=True)

video_base_name = "{}/{}".format(folder_name, ep_count)
video_name_bev = video_base_name + "_bev.mp4"
# video_list_bev = []
print("BEV video should be saved at: ", video_name_bev)
mediapy.write_video(video_name_bev, video_bev, fps=60)

video_name_interface = video_base_name + "_interface.mp4"
# video_list_interface = []
print("Interface video should be saved at: ", video_name_interface)
mediapy.write_video(video_name_interface, video_interface, fps=60)

if ep_count >= num_scenarios:
break

if generate_video:
video_bev = VideoRecorder(video_name=video_name_bev, height=3000, width=3000, code="avc1")
video_interface = VideoRecorder(video_name=video_name_interface, height=1200, width=1600, code="avc1")

o, _ = env.reset(seed=ep_count + start_seed)
env.vehicle.expert_takeover = True
env.engine.force_fps.disable()

# if generate_video:
# from metadrive.utils.image_to_video import image_list_to_video
# import shutil
# For Option 2: You need to send the temp folder to video recorder.
# It will read all image files and generate video.
# Then delete the temp folder.

# image_list_to_video(video_name_interface, tmp_folder_interface, code="avc1")
# print("Interface Video is saved at: ", video_name_interface)
# shutil.rmtree(tmp_folder_interface)
#
# image_list_to_video(video_name_bev, tmp_folder_bev, code="avc1")
# print("BEV Video is saved at: ", video_name_bev)
# shutil.rmtree(tmp_folder_bev)

0 comments on commit 55ab1a7

Please sign in to comment.