diff --git a/src/contest/capture.py b/src/contest/capture.py index 7869cd5..49090fa 100644 --- a/src/contest/capture.py +++ b/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): diff --git a/src/contest/captureGraphicsDisplay.py b/src/contest/captureGraphicsDisplay.py index 8f1bbfe..465f6b8 100644 --- a/src/contest/captureGraphicsDisplay.py +++ b/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