import gradio as gr import random import os import json import yaml import pandas as pd import numpy as np import plotly.graph_objects as go from collections import defaultdict import plotly.colors as pc import plotly.express as px # File paths for storing data DATA_DIR = "data" RATINGS_FILE = os.path.join(DATA_DIR, "ratings.json") PAIRWISE_FILE = os.path.join(DATA_DIR, "pairwise_results.json") ITEMS_CONFIG_FILE = "items_config.yaml" # Ensure data directory exists os.makedirs(DATA_DIR, exist_ok=True) def load_items_config(): with open(ITEMS_CONFIG_FILE, 'r') as f: return yaml.safe_load(f) # Load items configuration items_config = load_items_config() # Initialize Elo ratings and other statistics modalities = {modality: {} for modality in items_config.keys()} pairwise_results = {} def load_data(): global modalities, pairwise_results if os.path.exists(RATINGS_FILE): with open(RATINGS_FILE, 'r') as f: modalities = json.load(f) if os.path.exists(PAIRWISE_FILE): with open(PAIRWISE_FILE, 'r') as f: pairwise_results = json.load(f) # Initialize pairwise_results for modalities if not present for modality in modalities: if modality not in pairwise_results: pairwise_results[modality] = {} def save_data(): with open(RATINGS_FILE, 'w') as f: json.dump(modalities, f) with open(PAIRWISE_FILE, 'w') as f: json.dump(pairwise_results, f) # Load existing data load_data() def initialize_new_items(): for modality, items in items_config.items(): if modality not in modalities: modalities[modality] = {} if modality not in pairwise_results: pairwise_results[modality] = {} # Get the set of current item names from items_config current_item_names = set(item['name'] for item in items) # Check for items to remove items_to_remove = set(modalities[modality].keys()) - current_item_names for item_name in items_to_remove: # Remove from modalities del modalities[modality][item_name] # Remove from pairwise_results if item_name in pairwise_results[modality]: del pairwise_results[modality][item_name] # Remove this item from other items' pairwise results for other_item in pairwise_results[modality]: if item_name in pairwise_results[modality][other_item]: del pairwise_results[modality][other_item][item_name] for item in items: item_name = item['name'] if item_name not in modalities[modality]: modalities[modality][item_name] = { "rating": 1400, # Initial Elo rating "image": item['image'], "plays": 0 } # Initialize pairwise results for the new item if item_name not in pairwise_results[modality]: pairwise_results[modality][item_name] = {} for other_item in items: other_item_name = other_item['name'] if other_item_name != item_name: if other_item_name not in pairwise_results[modality][item_name]: pairwise_results[modality][item_name][other_item_name] = [0, 0] # Initialize new items initialize_new_items() # Save updated data save_data() current_modality = list(modalities.keys())[0] # Default to the first modality current_items = random.sample(list(modalities[current_modality].keys()), 2) def calculate_elo(rating1, rating2, k=32, win=1): """Calculate new Elo ratings""" expected1 = 1 / (1 + 10 ** ((rating2 - rating1) / 400)) expected2 = 1 / (1 + 10 ** ((rating1 - rating2) / 400)) new_rating1 = rating1 + k * (win - expected1) new_rating2 = rating2 + k * ((1 - win) - expected2) return new_rating1, new_rating2 def update_ratings(winner, loser): """Update ratings based on the winner""" items = modalities[current_modality] new_winner_rating, new_loser_rating = calculate_elo( items[winner]["rating"], items[loser]["rating"] ) items[winner]["rating"] = new_winner_rating items[loser]["rating"] = new_loser_rating items[winner]["plays"] += 1 items[loser]["plays"] += 1 pairwise_results[current_modality][winner][loser][0] += 1 pairwise_results[current_modality][winner][loser][1] += 1 pairwise_results[current_modality][loser][winner][1] += 1 save_data() def get_random_items(): """Get two random items from the current modality""" return random.sample(list(modalities[current_modality].keys()), 2) def choose_item(choice): """Handle item selection""" global current_items item1, item2 = current_items if choice == item1: update_ratings(item1, item2) else: update_ratings(item2, item1) # Get new random items current_items = get_random_items() return ( modalities[current_modality][current_items[0]]["image"], modalities[current_modality][current_items[1]]["image"], current_items[0], current_items[1], get_ratings_table(), get_pairwise_heatmap() ) def get_ratings_table(): """Generate a DataFrame of current ratings""" items = modalities[current_modality] data = [(name, info["rating"], info["plays"]) for name, info in items.items()] df = pd.DataFrame(data, columns=["Item", "Rating", "Plays"]) df = df.sort_values("Rating", ascending=False).reset_index(drop=True) df.index += 1 # Start index from 1 instead of 0 return df def get_pairwise_heatmap(): """Generate an interactive heatmap of pairwise win rates""" items = modalities[current_modality] # Sort items by rating sorted_items = sorted(items.items(), key=lambda x: x[1]['rating'], reverse=True) item_names = [item[0] for item in sorted_items] n = len(item_names) win_rates = np.zeros((n, n)) for i, item1 in enumerate(item_names): for j, item2 in enumerate(item_names): if i != j: wins, total = pairwise_results[current_modality][item1][item2] win_rates[i, j] = wins / total if total > 0 else 0.5 else: win_rates[i, j] = 0.5 # Set diagonal to 0.5 (neutral) # Create a custom colorscale from red to white to green colorscale = [ [0, "rgb(255,0,0)"], [0.5, "rgb(255,255,0)"], [1, "rgb(0,255,0)"] ] fig = go.Figure(data=go.Heatmap( z=win_rates, x=item_names, y=item_names, # Reverse the order for y-axis hoverongaps=False, colorscale=px.colors.diverging.RdYlGn, zmin=0, zmax=1, text=np.round(win_rates, 2), texttemplate="%{text}", textfont={"size": 10}, )) fig.update_layout( # title=f"Pairwise Win Rates - {current_modality}", xaxis_title="Opponent", yaxis_title="Item", width=800, height=800, ) # Add ratings to the y-axis labels (in reverse order) y_labels = [item for item in item_names[::-1]] fig.update_yaxes(ticktext=y_labels, tickvals=item_names[::-1]) fig.update_yaxes(autorange="reversed") fig.update_layout(xaxis = dict(side ="top")) fig.update_xaxes(tickangle=90) return fig def change_modality(new_modality): global current_modality, current_items current_modality = new_modality current_items = get_random_items() return ( modalities[current_modality][current_items[0]]["image"], modalities[current_modality][current_items[1]]["image"], current_items[0], current_items[1], get_ratings_table(), get_pairwise_heatmap() ) with gr.Blocks() as demo: gr.Markdown("# Elo Rating System") with gr.Row(): with gr.Column(scale=1): modality_dropdown = gr.Dropdown( choices=list(modalities.keys()), value=current_modality, label="Select Modality" ) with gr.Column(scale=4): with gr.Tabs(): with gr.TabItem("Compare Items"): gr.Markdown("Click on the button of the item you prefer. The Elo ratings will be updated accordingly.") with gr.Row(): img1 = gr.Image(value=modalities[current_modality][current_items[0]]["image"], label="Item 1", width=500, # Set a fixed width height=500) # Set a fixed height img2 = gr.Image(value=modalities[current_modality][current_items[1]]["image"], label="Item 2", width=500, # Set a fixed width height=500) # Set a fixed height with gr.Row(): btn1 = gr.Button(current_items[0]) btn2 = gr.Button(current_items[1]) with gr.TabItem("Ratings Table"): ratings_table = gr.Dataframe(value=get_ratings_table()) with gr.TabItem("Pairwise Results"): pairwise_heatmap = gr.Plot() modality_dropdown.change( change_modality, inputs=[modality_dropdown], outputs=[img1, img2, btn1, btn2, ratings_table, pairwise_heatmap] ) btn1.click( lambda: choose_item(current_items[0]), outputs=[img1, img2, btn1, btn2, ratings_table, pairwise_heatmap] ) btn2.click( lambda: choose_item(current_items[1]), outputs=[img1, img2, btn1, btn2, ratings_table, pairwise_heatmap] ) demo.launch()