import os import json import time import random import logging import threading from pathlib import Path from datetime import datetime, timedelta import numpy as np import gradio as gr from dotenv import load_dotenv from datasets import load_dataset from huggingface_hub import CommitScheduler from db import ( compute_elo_scores, get_all_votes, add_vote, is_running_in_space, fill_database_once ) # Load environment variables load_dotenv() huggingface_token = os.getenv("HUGGINGFACE_HUB_TOKEN") # Configure logging logging.basicConfig(level=logging.INFO) # Load datasets and initialize database dataset = load_dataset("bgsys/background-removal-arena-green", split='train') fill_database_once() # Directory setup for JSON dataset JSON_DATASET_DIR = Path("data/json_dataset") JSON_DATASET_DIR.mkdir(parents=True, exist_ok=True) # Initialize CommitScheduler if running in space commit_scheduler = CommitScheduler( repo_id="bgsys/votes_datasets_test2", repo_type="dataset", folder_path=JSON_DATASET_DIR, path_in_repo="data", token=huggingface_token ) if is_running_in_space() else None def fetch_elo_scores(): """Fetch and log Elo scores.""" try: elo_scores = compute_elo_scores() logging.info("Elo scores successfully computed.") return elo_scores except Exception as e: logging.error("Error computing Elo scores: %s", str(e)) return None def update_rankings_table(): """Update and return the rankings table based on Elo scores.""" elo_scores = fetch_elo_scores() or {} default_score = 1000 rankings = [ ["Photoroom", int(elo_scores.get("Photoroom", default_score))], ["RemoveBG", int(elo_scores.get("RemoveBG", default_score))], ["BRIA RMBG 2.0", int(elo_scores.get("BRIA RMBG 2.0", default_score))], ] rankings.sort(key=lambda x: x[1], reverse=True) return rankings def select_new_image(): """Select a new image and its segmented versions.""" max_attempts = 10 last_image_index = None for _ in range(max_attempts): available_indices = [i for i in range(len(dataset)) if i != last_image_index] if not available_indices: logging.error("No available images to select from.") return None random_index = random.choice(available_indices) sample = dataset[random_index] input_image = sample['original_image'] segmented_images = [sample.get(key) for key in ['clipdrop_image', 'bria_image', 'photoroom_image', 'removebg_image']] segmented_sources = ['Clipdrop', 'BRIA RMBG 2.0', 'Photoroom', 'RemoveBG'] if segmented_images.count(None) > 2: logging.error("Not enough segmented images found for: %s. Resampling another image.", sample['original_filename']) last_image_index = random_index continue try: selected_indices = random.sample([i for i, img in enumerate(segmented_images) if img is not None], 2) model_a_index, model_b_index = selected_indices return ( sample['original_filename'], input_image, segmented_images[model_a_index], segmented_images[model_b_index], segmented_sources[model_a_index], segmented_sources[model_b_index] ) except Exception as e: logging.error("Error processing images: %s. Resampling another image.", str(e)) last_image_index = random_index logging.error("Failed to select a new image after %d attempts.", max_attempts) return None def get_notice_markdown(): """Generate the notice markdown with dynamic vote count.""" total_votes = len(get_all_votes()) return f""" # ⚔️ Background Removal Arena: Compare & Test the Best Background Removal Models ## 📜 How It Works - **Blind Test**: You will see two images with their background removed from two anonymous background removal models (Clipdrop, RemoveBG, Photoroom, BRIA RMBG 2.0). - **Vote for the Best**: Choose the best result, if none stand out choose "Tie". ## 📊 Stats - **Total #votes**: {total_votes} ## 👇 Test now! """ def compute_mask_difference(segmented_a, segmented_b): """Compute the absolute difference between two image masks, ignoring green background.""" mask_a = np.asarray(segmented_a) mask_b = np.asarray(segmented_b) # Define the green background color green_background = (0, 255, 0, 255) # Create a binary mask where non-green and non-transparent pixels are marked as 1 mask_a_1d = np.where( (mask_a[..., :3] != green_background[:3]).any(axis=-1) & (mask_a[..., 3] != 0), 1, 0 ) mask_b_1d = np.where( (mask_b[..., :3] != green_background[:3]).any(axis=-1) & (mask_b[..., 3] != 0), 1, 0 ) # Compute the absolute difference between the masks return np.abs(mask_a_1d - mask_b_1d) js = r""" function load_zoom() { setTimeout(function() { // Select all images from the three displayed image containers. const images = document.querySelectorAll('.image-container img'); // Set transform origin so scaling and translating feels "natural". images.forEach(img => { img.style.transformOrigin = 'top left'; img.style.transition = 'transform 0.1s ease-out'; img.style.cursor = 'zoom-in'; }); // Choose a scale factor const scale = 2; function handleMouseMove(e) { const rect = e.currentTarget.getBoundingClientRect(); const xPercent = (e.clientX - rect.left) / rect.width; const yPercent = (e.clientY - rect.top) / rect.height; const offsetX = xPercent * (scale - 1) * 100; const offsetY = yPercent * (scale - 1) * 100; images.forEach(img => { img.style.transform = `translate(-${offsetX}%, -${offsetY}%) scale(${scale})`; }); } function handleMouseEnter(e) { e.currentTarget.addEventListener('mousemove', handleMouseMove); } function handleMouseLeave(e) { e.currentTarget.removeEventListener('mousemove', handleMouseMove); images.forEach(img => { img.style.transform = 'translate(0,0) scale(1)'; }); } const containers = document.querySelectorAll('.image-container'); containers.forEach(container => { container.addEventListener('mouseenter', handleMouseEnter); container.addEventListener('mouseleave', handleMouseLeave); }); }, 1000); // 1 second timeout } """ def gradio_interface(): """Create and return the Gradio interface.""" with gr.Blocks(js=js) as demo: gr.Markdown("# Background Removal Arena") button_name = "Difference between masks" with gr.Tabs() as tabs: with gr.Tab("⚔️ Arena (battle)", id=0): image_width = 800 image_height = 800 notice_markdown = gr.Markdown(get_notice_markdown(), elem_id="notice_markdown") with gr.Row(equal_height=True): def on_enter_contest(username): feedback_message = f"Thank you, {username or 'anonymous'}! You can see how you rank in the Hall of Fame." logging.info(feedback_message) return feedback_message with gr.Column(scale=2): username_input = gr.Textbox( label="Enter your username (optional)", placeholder="✨ Enter your username (optional)", show_label=False, submit_btn="Enter", interactive=True ) with gr.Column(scale=3): feedback_output = gr.Textbox( label="Feedback", interactive=False, show_label=False ) username_input.submit( fn=on_enter_contest, inputs=username_input, outputs=feedback_output ) def refresh_states(): filename, input_image, segmented_a, segmented_b, model_a_name, model_b_name = select_new_image() mask_difference = compute_mask_difference(segmented_a, segmented_b) image_a = gr.Image( value=segmented_a, label="Image", width=image_width, height=image_height ) input_image_display = gr.AnnotatedImage( value=(input_image, [(mask_difference > 0, button_name)]), label="Input Image", width=image_width, height=image_height ) image_b = gr.Image( value=segmented_b, label="Image", width=image_width, height=image_height ) zoomed_state_a = gr.State(False) zoomed_state_b = gr.State(False) state_model_a_name = gr.State(model_a_name) state_model_b_name = gr.State(model_b_name) state_filename = gr.State(filename) state_segmented_a = gr.Image(value=segmented_a, visible=False) state_segmented_b = gr.Image(value=segmented_b, visible=False) outputs = [ state_filename, image_a, image_b, state_model_a_name, state_model_b_name, input_image_display, zoomed_state_a, zoomed_state_b, state_segmented_a, state_segmented_b ] return outputs with gr.Row(): state_filename, image_a, image_b, state_model_a_name, state_model_b_name, input_image_display, zoomed_state_a, zoomed_state_b, state_segmented_a, state_segmented_b = refresh_states() with gr.Row(): vote_a_button = gr.Button("👈 A is better") vote_tie_button = gr.Button("🤝 Tie") vote_b_button = gr.Button("👉 B is better") def vote_for_model(choice, original_filename, model_a_name, model_b_name, user_username): """Submit a vote for a model and return updated images and model names.""" logging.info("Voting for model: %s", choice) vote_data = { "image_id": original_filename.value, "model_a": model_a_name.value, "model_b": model_b_name.value, "winner": choice, "user_id": user_username or "anonymous" } try: logging.debug("Adding vote data to the database: %s", vote_data) result = add_vote(vote_data) logging.info("Vote successfully recorded in the database with ID: %s", result["id"]) except Exception as e: logging.error("Error recording vote: %s", str(e)) outputs = refresh_states() new_notice_markdown = get_notice_markdown() return outputs + [new_notice_markdown] vote_a_button.click( fn=lambda username: vote_for_model("model_a", state_filename, state_model_a_name, state_model_b_name, username), inputs=[username_input], outputs=[ state_filename, image_a, image_b, state_model_a_name, state_model_b_name, input_image_display, zoomed_state_a, zoomed_state_b, state_segmented_a, state_segmented_b, notice_markdown ] ) vote_b_button.click( fn=lambda username: vote_for_model("model_b", state_filename, state_model_a_name, state_model_b_name, username), inputs=[username_input], outputs=[ state_filename, image_a, image_b, state_model_a_name, state_model_b_name, input_image_display, zoomed_state_a, zoomed_state_b, state_segmented_a, state_segmented_b, notice_markdown ] ) vote_tie_button.click( fn=lambda username: vote_for_model("tie", state_filename, state_model_a_name, state_model_b_name, username), inputs=[username_input], outputs=[ state_filename, image_a, image_b, state_model_a_name, state_model_b_name, input_image_display, zoomed_state_a, zoomed_state_b, state_segmented_a, state_segmented_b, notice_markdown ] ) def handle_zoom(current_image, zoomed_state_input, original_image, other_image, event: gr.SelectData): """Toggle between zoomed and original image based on click events.""" # Determine the current zoom state zoomed_state = zoomed_state_input.value if isinstance(zoomed_state_input, gr.State) else zoomed_state_input if zoomed_state: return ( gr.Image(value=original_image, label="Image", width=image_width, height=image_height), False, gr.Image(value=other_image, label="Image", width=image_width, height=image_height), False ) start_row, start_col = event.index[1], event.index[0] zoom_size = max(10, min(current_image.shape[:2]) // 10) row_start, row_end = max(start_row - zoom_size, 0), min(start_row + zoom_size, current_image.shape[0]) col_start, col_end = max(start_col - zoom_size, 0), min(start_col + zoom_size, current_image.shape[1]) grey_image = np.mean(current_image, axis=-1, keepdims=True).astype(current_image.dtype) grey_image = np.repeat(grey_image, current_image.shape[-1], axis=-1) output_image = grey_image.copy() zoomed_area = current_image[row_start:row_end, col_start:col_end] upscale_factor = 6 zoomed_area_upscaled = np.kron(zoomed_area, np.ones((upscale_factor, upscale_factor, 1))) center_row, center_col = start_row, start_col row_start_upscaled = max(center_row - zoomed_area_upscaled.shape[0] // 2, 0) row_end_upscaled = min(center_row + zoomed_area_upscaled.shape[0] // 2, output_image.shape[0]) col_start_upscaled = max(center_col - zoomed_area_upscaled.shape[1] // 2, 0) col_end_upscaled = min(center_col + zoomed_area_upscaled.shape[1] // 2, output_image.shape[1]) row_start_zoomed = max(0, -row_start_upscaled) row_end_zoomed = row_start_zoomed + (row_end_upscaled - row_start_upscaled) col_start_zoomed = max(0, -col_start_upscaled) col_end_zoomed = col_start_zoomed + (col_end_upscaled - col_start_upscaled) row_end_zoomed = min(row_end_zoomed, zoomed_area_upscaled.shape[0]) col_end_zoomed = min(col_end_zoomed, zoomed_area_upscaled.shape[1]) output_image[row_start_upscaled:row_end_upscaled, col_start_upscaled:col_end_upscaled] = \ zoomed_area_upscaled[row_start_zoomed:row_end_zoomed, col_start_zoomed:col_end_zoomed] # Apply the same zoom to the other image other_output_image = grey_image.copy() other_zoomed_area = other_image[row_start:row_end, col_start:col_end] other_zoomed_area_upscaled = np.kron(other_zoomed_area, np.ones((upscale_factor, upscale_factor, 1))) other_output_image[row_start_upscaled:row_end_upscaled, col_start_upscaled:col_end_upscaled] = \ other_zoomed_area_upscaled[row_start_zoomed:row_end_zoomed, col_start_zoomed:col_end_zoomed] return output_image, True, other_output_image, True image_a.select(handle_zoom, [image_a, zoomed_state_a, state_segmented_a, state_segmented_b], [image_a, zoomed_state_a, image_b, zoomed_state_b]) image_b.select(handle_zoom, [image_b, zoomed_state_b, state_segmented_b, state_segmented_a], [image_b, zoomed_state_b, image_a, zoomed_state_a]) with gr.Tab("🏆 Leaderboard", id=1) as leaderboard_tab: rankings_table = gr.Dataframe( headers=["Model", "Ranking"], value=update_rankings_table(), label="Current Model Rankings", column_widths=[180, 60], row_count=4 ) leaderboard_tab.select( fn=lambda: update_rankings_table(), outputs=rankings_table ) with gr.Tab("📊 Vote Data", id=2) as vote_data_tab: def update_vote_data(): votes = get_all_votes() return [[vote.id, vote.image_id, vote.model_a, vote.model_b, vote.winner, vote.user_id, vote.timestamp] for vote in votes] vote_table = gr.Dataframe( headers=["ID", "Image ID", "Model A", "Model B", "Winner", "user_id", "Timestamp"], value=update_vote_data(), label="Vote Data", column_widths=[20, 150, 100, 100, 100, 100, 150], row_count=0 ) vote_data_tab.select( fn=lambda: update_vote_data(), outputs=vote_table ) with gr.Tab("👥 Hall of Fame", id=3) as user_leaderboard_tab: current_time = datetime.now() start_of_week = current_time - timedelta(days=current_time.weekday()) def get_weekly_user_leaderboard(): """Get the leaderboard of users with the most votes in the current week, excluding anonymous votes.""" votes = get_all_votes() weekly_votes = [ vote for vote in votes if vote.timestamp >= start_of_week and vote.user_id and vote.user_id != "anonymous" ] user_vote_count = {} for vote in weekly_votes: user_vote_count[vote.user_id] = user_vote_count.get(vote.user_id, 0) + 1 sorted_users = sorted(user_vote_count.items(), key=lambda x: x[1], reverse=True) # Add medals for the top 3 users medals = ["🥇", "🥈", "🥉"] leaderboard = [] for index, (user, count) in enumerate(sorted_users): medal = medals[index] if index < len(medals) else "" leaderboard.append([f"{medal} {user}", count]) return leaderboard user_leaderboard_table = gr.Dataframe( headers=["User", "Votes"], value=get_weekly_user_leaderboard(), label="User Vote Leaderboard (This Week)", column_widths=[150, 100], row_count=0 ) leaderboard_info = gr.Markdown( value=f""" This leaderboard shows the ranking of users based on the number of votes they have cast in the current week. The current ranking is based on votes cast from {start_of_week.strftime('%Y-%m-%d')} to {current_time.strftime('%Y-%m-%d')}. It will be updated each week. """ ) user_leaderboard_tab.select( fn=lambda: get_weekly_user_leaderboard(), outputs=user_leaderboard_table ) return demo def dump_database_to_json(): """Dump the database to a JSON file and upload it to Hugging Face.""" if not is_running_in_space(): logging.info("Not running in Hugging Face Spaces. Skipping database dump.") return votes = get_all_votes() json_data = [ { "id": vote.id, "image_id": vote.image_id, "model_a": vote.model_a, "model_b": vote.model_b, "winner": vote.winner, "user_id": vote.user_id, "timestamp": vote.timestamp.isoformat() } for vote in votes ] json_file_path = JSON_DATASET_DIR / "votes.json" # Upload to Hugging Face with commit_scheduler.lock: with json_file_path.open("w") as f: json.dump(json_data, f, indent=4) logging.info("Database dumped to JSON") def schedule_dump_database(interval=60): """Schedule the database dump to JSON every specified interval in seconds.""" def run(): while True: logging.info("Starting database dump to JSON.") dump_database_to_json() logging.info("Database dump completed. Sleeping for %d seconds.", interval) time.sleep(interval) if is_running_in_space(): logging.info("Initializing database dump scheduler with interval: %d seconds.", interval) thread = threading.Thread(target=run, daemon=True) thread.start() logging.info("Database dump scheduler started.") else: logging.info("Not running in Hugging Face Spaces. Database dump scheduler not started.") if __name__ == "__main__": schedule_dump_database() # Start the periodic database dump demo = gradio_interface() demo.launch()