QBGPT

.streamlit/config.toml

@@ -0,0 +1,2 @@
+[theme]
+base="dark"

pages.py

@@ -2,7 +2,7 @@ import streamlit as st
 import plotly.express as px
 import pandas as pd
 import numpy as np
-from tools import generator
+from tools import generator, get_plot
 from PIL import Image
 
 def set_app_title_and_logo():
@@ -63,25 +63,10 @@ def qb_gpt_page(ref_df, ref, tokenizer, model):
         step1_true = QB_gen.prepare_for_plot(decoded_true)
         plot_true = pd.DataFrame(step1_true)
         
-        fig_gen = px.line(plot, x="input_ids_x", y="input_ids_y", animation_frame="pos_ids", color="OffDef", symbol="ids",
-                        text="position_ids", title="Generated players' trajectories Over Time", line_shape="linear",
-                        range_x=[0, 140], range_y=[0, 60], # Set X and Y axis ranges
-                        render_mode="svg")  # Render mode for smoother lines
-
-        # Customize the appearance of the plot
-        fig_gen.update_traces(marker=dict(size=10), selector=dict(mode='lines'))
-        fig_gen.update_layout(width=800, height=600) 
+        fig_gen = get_plot(plot, frames, "Generated")
         st.plotly_chart(fig_gen)
         
-        fig_true = px.line(plot_true, x="input_ids_x", y="input_ids_y", animation_frame="pos_ids", color="OffDef", symbol="ids",
-                    text="position_ids", title="True players' trajectories Over Time",
-                    range_x=[0, 140], range_y=[0, 60], # Set X and Y axis ranges
-                    line_shape="linear",  # Draw lines connecting points
-                    render_mode="svg")  # Render mode for smoother lines
-
-        # Customize the appearance of the plot
-        fig_true.update_traces(marker=dict(size=10), selector=dict(mode='lines'))
-        fig_true.update_layout(width=800, height=600) 
+        fig_true = get_plot(plot_true, frames, "True")
         st.plotly_chart(fig_true)
     
         

tools.py

@@ -2,6 +2,8 @@ import polars as pl
 import numpy as np
 import tensorflow as tf
 import pandas as pd
+import plotly.graph_objects as go
+
 
 class tokenizer:
     def __init__(self, 
@@ -373,3 +375,79 @@ class generator:
         cutted["ids"] = [[i for e in range(len(cutted["input_ids"][i]))] for i in range(len(cutted["input_ids"]))]
         merged = self.merge_cuts(cutted)
         return merged
+
+def get_plot(df, n_frames, name):
+    fig = go.Figure(
+    layout=go.Layout(
+        updatemenus=[dict(type="buttons", direction="right", x=0.9, y=1.16), ],
+        xaxis=dict(range=[0, 120],
+                   autorange=False, tickwidth=2,
+                   title_text="X"),
+        yaxis=dict(range=[0, 60],
+                   autorange=False,
+                   title_text="Y")
+    ))
+
+    # Add traces
+    i = 1
+    frames = {i: [] for i in df["pos_ids"].unique() if i !=0}
+
+    for id in df["ids"].unique():
+        spec = df[df["ids"] == id].reset_index(drop = True)
+        fig.add_trace(
+            go.Scatter(x=spec.input_ids_x[:i],
+                    y=spec.input_ids_y[:i],
+                    name= spec.position_ids.unique()[0],
+                    text= spec.position_ids.unique()[0],
+                    visible=True,
+                    line=dict(color="#f47738", dash="solid")))
+        
+        for k in range(i, spec.shape[0]):
+            current_frame = spec["pos_ids"][k]
+            frames[current_frame].append(go.Scatter(x=spec.input_ids_x[:k], y=spec.input_ids_y[:k]))
+            
+    frames = list(frames.values())
+    frames = [go.Frame(data = v) for v in frames]
+        
+
+    # Animation
+    fig.update(frames=frames)
+
+    fig.update_xaxes(ticks="outside", tickwidth=2, tickcolor='white', ticklen=10)
+    fig.update_yaxes(ticks="outside", tickwidth=2, tickcolor='white', ticklen=1)
+    fig.update_layout(yaxis_tickformat=',')
+    fig.update_layout(legend=dict(x=0, y=1.1), legend_orientation="h")
+
+    # Buttons
+    fig.update_layout(title=f"{name} play",
+                    xaxis_title="X",
+                    yaxis_title="Y",
+                    legend_title="Legend Title",
+                    showlegend=False,
+                    font=dict(
+                        family="Arial",
+                        size=14
+                    ),
+                    hovermode="x",
+                    updatemenus=[
+                            dict(
+                                buttons=list(
+                                    [
+                                    dict(label="Play",
+                                        method="animate",
+                                        args=[None, {"frame": {"duration": n_frames}}])
+                                    ]
+                                             ),
+                                type = "buttons",
+                                direction="right",
+                                pad={"r": 50, "t": 50},
+                                showactive=False,
+                                x=0.5,
+                                yanchor="top")
+                            ])
+    
+    fig.update_layout(template='plotly_dark'
+                 )
+
+    fig.update_layout(width=1200, height=600)
+    return fig