Update app.py

app.py

@@ -67,9 +67,11 @@ def infer(model, replay_file,
     for i in range(0, len(serialized_states), bs):
         batch = (serialized_states[i:i + bs].clone().to(DEVICE), 
                  serialized_scoreboards[i:i + bs].clone().to(DEVICE))
-        if nullify_goal_difference:
-            batch[:, SB_BLUE_SCORE] = 0
-            batch[:, SB_ORANGE_SCORE] = 0
+        if nullify_goal_difference or ignore_ties:
+            batch[1][:, SB_BLUE_SCORE] = 0
+            batch[1][:, SB_ORANGE_SCORE] = 0
+        if ignore_ties:
+            batch[1][:, SB_GAME_TIMER_SECONDS] = float("inf")
         out = model(*batch)
         it.update(len(batch[0]))
 
@@ -112,13 +114,17 @@ def infer(model, replay_file,
     preds = preds[main_cols + [c for c in preds.columns if c not in main_cols]]
     # Set index name
     preds.index.name = "Frame"
+    if ignore_ties:
+        tie_probs = preds["Tie"]
+        preds = preds.drop("Tie")
+        preds[[c for c in preds.columns if c.startswith("Blue") or c.startswith("Orange")] /= (1 - tie_probs)
 
     return preds
 
 
 def plot_plotly(preds: pd.DataFrame):
     import plotly.graph_objects as go
-    preds_df = preds * 100
+    preds_df = preds.drop(["Touch", "Timer", "Goal"], axis=1) * 100
     timer = preds["Timer"]
     fig = go.Figure()
 
@@ -138,10 +144,11 @@ def plot_plotly(preds: pd.DataFrame):
         go.Scatter(x=preds_df.index, y=preds_df["Orange"],
                    mode='lines', name='Orange', line=dict(color='orange'),
                    customdata=timer_text, hovertemplate=hovertemplate))
-    fig.add_trace(
-        go.Scatter(x=preds_df.index, y=preds_df["Tie"],
-                   mode='lines', name='Tie', line=dict(color='gray'),
-                   customdata=timer_text, hovertemplate=hovertemplate))
+    if "Tie" in preds.columns:
+        fig.add_trace(
+            go.Scatter(x=preds_df.index, y=preds_df["Tie"],
+                       mode='lines', name='Tie', line=dict(color='gray'),
+                       customdata=timer_text, hovertemplate=hovertemplate))
 
     # Add the horizontal line at y=50%
     fig.add_hline(y=50, line_dash="dash", line_color="black", name="50% Probability")
@@ -242,20 +249,24 @@ def gradio_app():
             # Use gr.Column to stack components vertically
             with gr.Column():
                 file_input = gr.File(label="Upload Replay File", type="filepath", file_types=[".replay"])
+                checkboxes = gr.CheckboxGroup(["Nullify goal difference", "Ignore ties"])
                 submit_button = gr.Button("Generate Predictions")
                 plot_output = gr.Plot(label="Predictions")
                 download_button = gr.DownloadButton("Download Predictions", visible=False)
 
             # Make plot on button click
-            def make_plot(replay_file, progress=gr.Progress(track_tqdm=True)):
+            def make_plot(replay_file, checkbox_options, progress=gr.Progress(track_tqdm=True)):
+                nullify_goal_difference, ignore_ties = checkbox_options
                 print(f"Processing file: {replay_file}")
                 replay_stem = os.path.splitext(os.path.basename(replay_file))[0]
-                preds_file = os.path.join(temp_dir, f"predictions_{replay_stem}.csv")
+                preds_file = os.path.join(temp_dir, f"predictions_{replay_stem}_{options}.csv")
                 if os.path.exists(preds_file):
                     print(f"Predictions file already exists: {preds_file}")
                     preds = pd.read_csv(preds_file)
                 else:
-                    preds = infer(MODEL, replay_file)
+                    preds = infer(MODEL, replay_file, 
+                                 nullify_goal_difference=nullify_goal_difference,
+                                 ignore_ties=ignore_ties)
                 plt = plot_plotly(preds)
                 print(f"Plot generated for file: {replay_file}")
                 preds.to_csv(preds_file)
@@ -267,7 +278,7 @@ def gradio_app():
 
             submit_button.click(
                 fn=make_plot,
-                inputs=file_input,
+                inputs=[file_input, checkboxes],
                 outputs=[plot_output, download_button],
                 show_progress="full",
             )