app.py

import gradio as gr
import random
import json
import os

initial_data = [
    ["dalle_desc_25", 1000, ""],
    ["dalle_desc_50", 1000, ""],
    ["dalle_desc_100", 1000, ""],
    ["dalle_desc_150", 1000, ""],
    ["dalle_desc_250", 1000, ""],
    ["desc_25_threshold_250", 1000, ""],
    ["desc_25_threshold_500", 1000, ""],
    ["desc_25_threshold_1000", 1000, ""],
    ["desc_250_threshold_250", 1000, ""],
    ["desc_250_threshold_500", 1000, ""],
    ["desc_250_threshold_1000", 1000, ""],
    ["jpeg_scale_2", 1000, ""],
    ["jpeg_scale_4", 1000, ""],
    ["jpeg_scale_8", 1000, ""],
    ["jpeg_scale_16", 1000, ""],
    ["jpeg_scale_32", 1000, ""],
    ["sa30_desc_50", 1000, ""],
    ["sa30_desc_100", 1000, ""],
    ["sa30_desc_150", 1000, ""],
    ["sa30_desc_250", 1000, ""],
    ["sd30_desc_25", 1000, ""],
    ["sd35_desc_25", 1000, ""],
    ["sd35_desc_50", 1000, ""],
    ["sd35_desc_100", 1000, ""],
    ["sd35_desc_150", 1000, ""],
    ["sd35_desc_250", 1000, ""],
]

# File to store leaderboard data
LEADERBOARD_FILE = "leaderboard_data.json"

def load_leaderboard_data():
    if os.path.exists(LEADERBOARD_FILE):
        with open(LEADERBOARD_FILE, 'r') as f:
            return json.load(f)
    else:
        # Initial leaderboard data if file doesn't exist
        save_leaderboard_data(initial_data)
        return initial_data

def save_leaderboard_data(data):
    with open(LEADERBOARD_FILE, 'w') as f:
        json.dump(data, f)

# Initialize global variables
objs = ['axe', 'barrel', 'bed', 'bottle', 'canon', 'car', 'chair', 'chair2', 'chair3', 'chair4']
leaderboard_data = load_leaderboard_data()
images = [f"3d/{e[0]}/OBJ.png" for e in initial_data]
current_images = [0, 0]

def get_new_images():
    global current_images
    random.seed()
    idx1, idx2 = random.sample(range(len(images)), 2)
    current_images = [idx1, idx2]
    obj = random.choice(objs)
    new_images = [i.replace('OBJ', obj) for i in images]
    original = f"3d/original/{obj}.png"
    return {
        "original": original,
        "image1": new_images[idx1],
        "image2": new_images[idx2],
        "label1": "Left",
        "label2": "Right"
    }

def update_elo(winner, loser, k=32):
    global leaderboard_data
    winner_elo = leaderboard_data[winner][1]
    loser_elo = leaderboard_data[loser][1]
    expected_winner = 1 / (1 + 10 ** ((loser_elo - winner_elo) / 400))
    expected_loser = 1 / (1 + 10 ** ((winner_elo - loser_elo) / 400))
    leaderboard_data[winner][1] += round(k * (1 - expected_winner))
    leaderboard_data[loser][1] += round(k * (0 - expected_loser))
    save_leaderboard_data(leaderboard_data)
    return leaderboard_data

def vote_and_randomize(choice):
    global current_images
    
    if choice == "left":
        updated_data = update_elo(current_images[0], current_images[1])
        winner_name = leaderboard_data[current_images[0]][0]
    else:
        updated_data = update_elo(current_images[1], current_images[0])
        winner_name = leaderboard_data[current_images[1]][0]
    
    new_state = get_new_images()
    
    return (
        f"Thanks for voting for {winner_name}!", 
        new_state["original"],
        new_state["image1"], 
        new_state["image2"],
        new_state["label1"],
        new_state["label2"],
        updated_data
    )

def start_voting():
    # Get initial random images
    initial_state = get_new_images()
    
    return (
        gr.update(visible=False),  # Hide start button
        gr.update(visible=True),   # Show voting container
        initial_state["original"],
        initial_state["image1"],
        initial_state["image2"],
        initial_state["label1"],
        initial_state["label2"]
    )

with gr.Blocks(css="""
    #main-image {
        margin: auto; /* Center the image */
        display: block;
    }
    """) as demo:
    with gr.Tabs():
        # Tab 1: Voting
        with gr.Tab("Voting"):
            gr.Markdown("### Vote for your favorite option!")
            
            # Start button (centered)
            with gr.Column(elem_id="start-container"):
                start_btn = gr.Button("Start!", scale=0.5)
            
            # Voting interface (initially hidden)
            with gr.Column(visible=False) as voting_container:
                # Image Comparison Grid
                # justify in the center
                with gr.Row(equal_height=True):
                    main_image = gr.Image(value=None, label="Original", interactive=False, show_download_button=True, elem_id="main-image", scale=0.25)

                with gr.Row():
                    left_image = gr.Image(value=None, label="Left Option", interactive=False, show_download_button=False)
                    right_image = gr.Image(value=None, label="Right Option", interactive=False, show_download_button=False)

                with gr.Row():
                    vote_1 = gr.Button(value="")
                    vote_2 = gr.Button(value="")

                output = gr.Textbox(label="Vote Result", interactive=False)
        
        # Tab 2: Leaderboard
        with gr.Tab("Leaderboard"):
            gr.Markdown("### Leaderboard")
            leaderboard_table = gr.DataFrame(
                headers=["Name", "Elo", "Description"],
                value=load_leaderboard_data(),
                interactive=False
            )
    
    # Handle start button click
    start_btn.click(
        fn=start_voting,
        outputs=[
            start_btn,
            voting_container,
            main_image,
            left_image,
            right_image,
            vote_1,
            vote_2
        ]
    )
    
    # Handle voting buttons
    vote_1.click(
        fn=lambda: vote_and_randomize("left"),
        outputs=[output, main_image, left_image, right_image, vote_1, vote_2, leaderboard_table]
    )
    vote_2.click(
        fn=lambda: vote_and_randomize("right"),
        outputs=[output, main_image, left_image, right_image, vote_1, vote_2, leaderboard_table]
    )

if __name__ == "__main__":
    demo.launch()