Update app.py

app.py

@@ -68,6 +68,8 @@ def load_leaderboard():
     result.drop('arch_name', axis=1, inplace=True)
     result.drop('crop_pct', axis=1, inplace=True)
     result.drop('interpolation', axis=1, inplace=True)
+
+    result['highlighted'] = False
     
     # Round numerical values
     result = result.round(2)
@@ -106,6 +108,7 @@ def filter_leaderboard(df, model_name, sort_by):
     
     return filtered_df
 
+
 def create_scatter_plot(df, x_axis, y_axis):
     fig = px.scatter(
         df,
@@ -116,10 +119,15 @@ def create_scatter_plot(df, x_axis, y_axis):
         hover_data=['model'],
         trendline='ols',
         trendline_options=dict(log_x=True, log_y=True),
+        color='highlighted',
+        color_discrete_map={True: 'red', False: 'blue'},
         title=f'{y_axis} vs {x_axis}'
     )
+    fig.update_layout(showlegend=False)
+    
     return fig
 
+
 # Load the leaderboard data
 full_df = load_leaderboard()
 
@@ -132,14 +140,30 @@ DEFAULT_SORT = "avg_top1"
 DEFAULT_X = "infer_samples_per_sec"
 DEFAULT_Y = "avg_top1"
 
-def update_leaderboard_and_plot(model_name=DEFAULT_SEARCH, sort_by=DEFAULT_SORT, x_axis=DEFAULT_X, y_axis=DEFAULT_Y):
-    filtered_df = filter_leaderboard(
-        full_df,  # in outer scope
-        model_name,
-        sort_by,
-    )
-    fig = create_scatter_plot(filtered_df, x_axis, y_axis)
-    return filtered_df, fig
+def update_leaderboard_and_plot(
+        model_name=DEFAULT_SEARCH,
+        highlight_name=None,
+        sort_by=DEFAULT_SORT,
+        x_axis=DEFAULT_X,
+        y_axis=DEFAULT_Y,
+):
+    filtered_df = filter_leaderboard(full_df, model_name, sort_by)    
+    
+    # Apply the highlight filter to the entire dataset so the output will be union (comparison) if the filters are disjoint
+    highlight_df = filter_leaderboard(full_df, highlight_name, sort_by) if highlight_name else None
+    
+    # Combine filtered_df and highlight_df, removing duplicates
+    if highlight_df is not None:
+        combined_df = pd.concat([filtered_df, highlight_df]).drop_duplicates().reset_index(drop=True)
+        combined_df = combined_df.sort_values(by=sort_by, ascending=False)
+        combined_df['highlighted'] = combined_df['model'].isin(highlight_df['model'])
+    else:
+        combined_df = filtered_df
+           
+    fig = create_scatter_plot(combined_df, x_axis, y_axis)
+    highlighted_df = combined_df.style.apply(lambda x: ['background-color: #ffcccc' if x['highlighted'] else '' for _ in x], axis=1)
+    
+    return highlighted_df, fig
 
 
 with gr.Blocks(title="The timm Leaderboard") as app:
@@ -148,8 +172,11 @@ with gr.Blocks(title="The timm Leaderboard") as app:
     gr.HTML("<p>Search tips:<br>- Use wildcards (* or ?) for pattern matching<br>- Use 're:' prefix for regex search<br>- Otherwise, fuzzy matching will be used</p>")
     
     with gr.Row():
-        search_bar = gr.Textbox(lines=1, label="Search Model", placeholder="e.g. resnet*, re:^vit, efficientnet", scale=3)
+        search_bar = gr.Textbox(lines=1, label="Model Filter", placeholder="e.g. resnet*, re:^vit, efficientnet", scale=3)
         sort_dropdown = gr.Dropdown(choices=sort_columns, label="Sort by", value=DEFAULT_SORT, scale=1)
+
+    with gr.Row():
+        highlight_bar = gr.Textbox(lines=1, label="Model Highlight/Compare Filter", placeholder="e.g. convnext*, re:^efficient")
     
     with gr.Row():
         x_axis = gr.Dropdown(choices=plot_columns, label="X-axis", value=DEFAULT_X)
@@ -164,13 +191,17 @@ with gr.Blocks(title="The timm Leaderboard") as app:
     
     search_bar.submit(
         update_leaderboard_and_plot,
-        inputs=[search_bar, sort_dropdown, x_axis, y_axis],
+        inputs=[search_bar, highlight_bar, sort_dropdown, x_axis, y_axis],
+        outputs=[leaderboard, plot]
+    )
+    highlight_bar.submit(
+        update_leaderboard_and_plot,
+        inputs=[search_bar, highlight_bar, sort_dropdown, x_axis, y_axis],
         outputs=[leaderboard, plot]
     )
-    
     update_btn.click(
         update_leaderboard_and_plot,
-        inputs=[search_bar, sort_dropdown, x_axis, y_axis],
+        inputs=[search_bar, highlight_bar, sort_dropdown, x_axis, y_axis],
         outputs=[leaderboard, plot]
     )