generating video

src/contest/capture.py

@@ -11,7 +11,6 @@
 # Student side autograding was added by Brad Miller, Nick Hay, and
 # Pieter Abbeel (pabbeel@cs.berkeley.edu).
 
-
 # capture.py
 # ----------
 # Licensing Information: Please do not distribute or publish solutions to this
@@ -61,6 +60,7 @@
 from contest.game import Actions
 from contest.game import GameStateData, Game, Grid, Configuration
 from contest.util import nearestPoint, manhattanDistance
+from contest.captureGraphicsDisplay import *
 
 #DIR_SCRIPT = sys.path[0] + "/src/contest/"
 import contest
@@ -866,6 +866,7 @@ def read_command(argv):
         recorded['red_team_name'] = parsed_options.red
         recorded['blue_team_name'] = parsed_options.blue
         recorded['wait_end'] = False
+        recorded['replay_filename'] = parsed_options.replay
         replay_game(**recorded)
         sys.exit(0)
 
@@ -879,7 +880,7 @@ def read_command(argv):
         recorded['red_team_name'] = parsed_options.red
         recorded['blue_team_name'] = parsed_options.blue
         recorded['wait_end'] = False
-
+        recorded['replay_filename'] = parsed_options.replay
         replay_game(**recorded)
         sys.exit(0)
 
@@ -1018,7 +1019,9 @@ def load_agents(is_red, agent_file, cmd_line_args):
     return create_team_func(indices[0], indices[1], is_red, **args)
 
 
-def replay_game(layout, agents, actions, display, length, red_team_name, blue_team_name, wait_end=True, delay=1):
+def replay_game(layout, agents, actions, display, length, red_team_name, blue_team_name, replay_filename,wait_end=True, delay=1):
+    clear_directory(PS_DIR)
+    clear_directory(PNG_DIR)
     rules = CaptureRules()
     game = rules.new_game(layout, agents, display, length, False, False)
     state = game.state
@@ -1033,6 +1036,9 @@ def replay_game(layout, agents, actions, display, length, red_team_name, blue_te
         display.update(state.data)
         # Allow for game specific conditions (winning, losing, etc.)
         rules.process(state, game)
+        base_name = os.path.basename(replay_filename).replace(".replay", "")
+        saveFrame(base_name)
+
         time.sleep(delay)
 
     game.game_over = True
@@ -1066,9 +1072,18 @@ def replay_game(layout, agents, actions, display, length, red_team_name, blue_te
             input("PRESS ENTER TO CONTINUE")
         except:
             print("END")
+
+    if SAVE_POSTSCRIPT:
+        # Convert all PostScript files to PNG
+        convert_all_ps_to_png(base_name)
+        # Create videos
+        create_video_from_pngs(base_name)
 
-    display.finish()
+        #clear files
+        clear_directory(PS_DIR)
+        clear_directory(PNG_DIR)
 
+    display.finish()
 
 def run_games(layouts, agents, display, length, num_games, record, num_training, red_team_name, blue_team_name,
               contest_name="default", mute_agents=False, catch_exceptions=False, delay_step=0, match_id=0):

src/contest/captureGraphicsDisplay.py

@@ -802,17 +802,80 @@ def add(x, y):
 # convert -delay 7 -loop 1 -compress lzw -layers optimize frame* out.gif
 # convert is part of imagemagick (freeware)
 
-SAVE_POSTSCRIPT = False
-POSTSCRIPT_OUTPUT_DIR = 'frames'
+SAVE_POSTSCRIPT = True
+PS_DIR = './frames'
+PNG_DIR = './frames_png'
 FRAME_NUMBER = 0
 import os
+import subprocess
+import shutil
 
+def convert_ps_to_png(ps_file, png_file):
+    try:
+        cmd = ["convert", ps_file, png_file]
+        subprocess.run(cmd, check=True)
+    except subprocess.CalledProcessError as e:
+        print(f"Error converting {ps_file} to {png_file}: {e}")
 
-def saveFrame():
-    "Saves the current graphical output as a postscript file"
-    global SAVE_POSTSCRIPT, FRAME_NUMBER, POSTSCRIPT_OUTPUT_DIR
-    if not SAVE_POSTSCRIPT: return
-    if not os.path.exists(POSTSCRIPT_OUTPUT_DIR): os.mkdir(POSTSCRIPT_OUTPUT_DIR)
-    name = os.path.join(POSTSCRIPT_OUTPUT_DIR, 'frame_%08d.ps' % FRAME_NUMBER)
+def convert_all_ps_to_png(base_name):
+    for i in range(FRAME_NUMBER):
+        ps_name = os.path.join(PS_DIR, f'{base_name}_frame_%08d.ps' % i)
+        png_name = os.path.join(PNG_DIR, f'{base_name}_frame_%08d.png' % i)
+        convert_ps_to_png(ps_name, png_name)
+
+
+def convert_pngs_to_mp4(base_name):
+    """
+    Convert PNG images in PNG_DIR with the given base_name to an MP4 video.
+    """
+    output_file = os.path.join('./contest_video', f'{base_name}.mp4')
+    os.makedirs('./contest_video', exist_ok=True)
+
+    try:
+        cmd = [
+            'ffmpeg', 
+            '-framerate', '30',
+            '-i', os.path.join(PNG_DIR, f'{base_name}_frame_%08d.png'),
+            '-vf', 'pad=ceil(iw/2)*2:ceil(ih/2)*2',
+            '-c:v', 'libx264',
+            '-pix_fmt', 'yuv420p',
+            output_file
+        ]
+        subprocess.run(cmd, check=True)
+    except subprocess.CalledProcessError as e:
+        print(f"Error creating MP4 for {base_name}: {e}")
+
+
+
+def create_video_from_pngs(base_name):
+    convert_pngs_to_mp4(base_name)
+
+def clear_directory(directory):
+    """
+    Clear all files and subdirectories in the given directory.
+    """
+    for filename in os.listdir(directory):
+        file_path = os.path.join(directory, filename)
+        try:
+            if os.path.isfile(file_path) or os.path.islink(file_path):
+                os.unlink(file_path)
+            elif os.path.isdir(file_path):
+                shutil.rmtree(file_path)
+        except Exception as e:
+            print(f'Failed to delete {file_path}. Reason: {e}')
+
+
+def saveFrame(base_name):
+    global SAVE_POSTSCRIPT, FRAME_NUMBER, PS_DIR
+
+    if not SAVE_POSTSCRIPT: 
+        return FRAME_NUMBER
+
+    if not os.path.exists(PS_DIR): 
+        os.mkdir(PS_DIR)
+
+    ps_name = os.path.join(PS_DIR, f'{base_name}_frame_%08d.ps' % FRAME_NUMBER)
     FRAME_NUMBER += 1
-    writePostscript(name)  # writes the current canvas
+    writePostscript(ps_name) 
+
+    return FRAME_NUMBER