Update app.py

app.py

@@ -4,7 +4,7 @@ from tempfile import TemporaryDirectory
 import numpy as np
 import pandas as pd
 import torch
-from rlgym_tools.rocket_league.misc.serialize import serialize_replay_frame, RF_SCOREBOARD, RF_EPISODE_SECONDS_REMAINING, SB_GAME_TIMER_SECONDS, SB_BLUE_SCORE, SB_ORANGE_SCORE
+from rlgym_tools.rocket_league.misc.serialize import serialize_game_state, serialize_scoreboard, RF_SCOREBOARD, RF_EPISODE_SECONDS_REMAINING, SB_GAME_TIMER_SECONDS, SB_BLUE_SCORE, SB_ORANGE_SCORE
 from rlgym_tools.rocket_league.replays.convert import replay_to_rlgym
 from rlgym_tools.rocket_league.replays.parsed_replay import ParsedReplay
 from tqdm import trange, tqdm
@@ -26,7 +26,6 @@ MODEL.eval()
 @torch.inference_mode()
 def infer(model, replay_file,
           nullify_goal_difference=False,
-          predict_with_all_quadrants=True,
           ignore_ties=False):
     num_outputs = 123
     swap_team_idx = torch.arange(num_outputs)
@@ -37,20 +36,24 @@ def infer(model, replay_file,
     replay = ParsedReplay.load(replay_file)
     it = tqdm(replay_to_rlgym(replay), desc="Loading replay", total=len(replay.game_df))
     replay_frames = []
-    serialized = []
+    serialized_states = []
+    serialized_scoreboards = []
+    seconds_remaining = []
     for replay_frame in it:
         replay_frames.append(replay_frame)
-        s = serialize_replay_frame(replay_frame)
-        serialized.append(s)
-    serialized = np.stack(serialized)
-    serialized = torch.from_numpy(serialized)
+        sstate = serialize_game_state(replay_frame.state)
+        sscoreboard = serialize_scoreboard(replay_frame.scoreboard)
+        serialized_states.append(sstate)
+        serialized_scoreboards.append(sscoreboard)
+        seconds_remaining.append(replay_frame.episode_seconds_remaining)
+    serialized_states = torch.from_numpy(np.stack(serialized_states))
+    serialized_scoreboards = torch.from_numpy(np.stack(serialized_scoreboards))
+    seconds_remaining = torch.tensor(seconds_remaining)
     it.close()
 
-    scoreboard = serialized[:, RF_SCOREBOARD]
-
-    timer = scoreboard[:, SB_GAME_TIMER_SECONDS]
+    timer = serialized_scoreboards[:, SB_GAME_TIMER_SECONDS]
     is_ot = timer > 450
-    ot_time_remaining = serialized[is_ot, RF_EPISODE_SECONDS_REMAINING]
+    ot_time_remaining = seconds_remaining[is_ot]
     if len(ot_time_remaining) > 0:
         ot_timer = ot_time_remaining[0] - ot_time_remaining
         timer[is_ot] = -ot_timer  # Negate to indicate overtime
@@ -63,9 +66,10 @@ def infer(model, replay_file,
 
     bs = 512
     predictions = []
-    it = trange(len(serialized), desc="Running model")
-    for i in range(0, len(serialized), bs):
-        batch = serialized[i:i + bs].clone().to(DEVICE)
+    it = trange(len(serialized_states), desc="Running model")
+    for i in range(0, len(serialized_states), bs):
+        batch = (serialized_states[i:i + bs].clone().to(DEVICE), 
+                 scoreboard[i:i + bs].clone().to(DEVICE))
         if nullify_goal_difference:
             batch[:, SB_BLUE_SCORE] = 0
             batch[:, SB_ORANGE_SCORE] = 0