render to video using sumo-gui (#784)

* render to video using sumo-gui

* flake

* allow usage of sumo_gui in render

* cleanup

flow/envs/base.py

@@ -7,6 +7,8 @@
 import traceback
 import numpy as np
 import random
+import shutil
+import subprocess
 from flow.renderer.pyglet_renderer import PygletRenderer as Renderer
 from flow.utils.flow_warnings import deprecated_attribute
 
@@ -214,7 +216,10 @@ def __init__(self,
             # render a frame
             self.render(reset=True)
         elif self.sim_params.render in [True, False]:
-            pass  # default to sumo-gui (if True) or sumo (if False)
+            # default to sumo-gui (if True) or sumo (if False)
+            if (self.sim_params.render is True) and self.sim_params.save_render:
+                self.path = os.path.expanduser('~')+'/flow_rendering/' + self.network.name
+                os.makedirs(self.path, exist_ok=True)
         else:
             raise FatalFlowError(
                 'Mode %s is not supported!' % self.sim_params.render)
@@ -670,6 +675,15 @@ def terminate(self):
             # close pyglet renderer
             if self.sim_params.render in ['gray', 'dgray', 'rgb', 'drgb']:
                 self.renderer.close()
+            # generate video
+            elif (self.sim_params.render is True) and self.sim_params.save_render:
+                images_dir = self.path.split('/')[-1]
+                speedup = 10  # multiplier: renders video so that `speedup` seconds is rendered in 1 real second
+                fps = speedup//self.sim_step
+                p = subprocess.Popen(["ffmpeg", "-y", "-r", str(fps), "-i", self.path+"/frame_%06d.png",
+                                      "-pix_fmt", "yuv420p", "%s/../%s.mp4" % (self.path, images_dir)])
+                p.wait()
+                shutil.rmtree(self.path)
         except FileNotFoundError:
             # Skip automatic termination. Connection is probably already closed
             print(traceback.format_exc())
@@ -699,6 +713,9 @@ def render(self, reset=False, buffer_length=5):
                 if len(self.frame_buffer) > buffer_length:
                     self.frame_buffer.pop(0)
                     self.sights_buffer.pop(0)
+        elif (self.sim_params.render is True) and self.sim_params.save_render:
+            # sumo-gui render
+            self.k.kernel_api.gui.screenshot("View #0", self.path+"/frame_%06d.png" % self.time_counter)
 
     def pyglet_render(self):
         """Render a frame using pyglet."""

flow/visualize/visualizer_rllib.py

@@ -13,7 +13,6 @@
 """
 
 import argparse
-from datetime import datetime
 import gym
 import numpy as np
 import os
@@ -114,8 +113,9 @@ def visualizer_rllib(args):
     elif args.render_mode == 'no_render':
         sim_params.render = False
     if args.save_render:
-        sim_params.render = 'drgb'
-        sim_params.pxpm = 4
+        if args.render_mode != 'sumo_gui':
+            sim_params.render = 'drgb'
+            sim_params.pxpm = 4
         sim_params.save_render = True
 
     # Create and register a gym+rllib env
@@ -314,23 +314,6 @@ def visualizer_rllib(args):
         # delete the .xml version of the emission file
         os.remove(emission_path)
 
-    # if we wanted to save the render, here we create the movie
-    if args.save_render:
-        dirs = os.listdir(os.path.expanduser('~')+'/flow_rendering')
-        # Ignore hidden files
-        dirs = [d for d in dirs if d[0] != '.']
-        dirs.sort(key=lambda date: datetime.strptime(date, "%Y-%m-%d-%H%M%S"))
-        recent_dir = dirs[-1]
-        # create the movie
-        movie_dir = os.path.expanduser('~') + '/flow_rendering/' + recent_dir
-        save_dir = os.path.expanduser('~') + '/flow_movies'
-        if not os.path.exists(save_dir):
-            os.mkdir(save_dir)
-        os_cmd = "cd " + movie_dir + " && ffmpeg -i frame_%06d.png"
-        os_cmd += " -pix_fmt yuv420p " + dirs[-1] + ".mp4"
-        os_cmd += "&& cp " + dirs[-1] + ".mp4 " + save_dir + "/"
-        os.system(os_cmd)
-
 
 def create_parser():
     """Create the parser to capture CLI arguments."""