app.py

import streamlit as st
import pandas as pd
from PIL import Image, ImageDraw, ImageFont
import io

def main():
    # Sidebar logo and title
    with st.sidebar:
        col1, col2 = st.columns([1, 5])  # Shrink the logo column and expand the text column
    
        with col1:
            logo = Image.open("logo.png")
            resized_logo = logo.resize((40, 40))  # Resize the logo
            st.image(resized_logo)
    
        with col2:
            st.markdown(
                """
                <div style="
                    display: flex; 
                    align-items: center; 
                    gap: 10px; 
                    margin: 0; 
                    padding: 0; 
                    font-family: 'Inter', sans-serif; 
                    font-size: 26px; 
                    font-weight: bold;">
                    AI Energy Score
                </div>
                """,
                unsafe_allow_html=True,
            )

    # Sidebar instructions and link
    st.sidebar.markdown(
    """
    <h1 style="text-align: center; font-size: 24px; font-weight: bold;">
        Generate a Label to Display your 
        <a href="https://huggingface.co/spaces/AIEnergyScore/Leaderboard" target="_blank" style="text-decoration: none; color: inherit;">
            AI Energy Score
        </a>
    </h1>
    """,
    unsafe_allow_html=True,
    )

    st.sidebar.markdown("<hr style='border: 1px solid gray; margin: 15px 0;'>", unsafe_allow_html=True)

    st.sidebar.write("### Instructions:")
    st.sidebar.write("#### 1. Select task(s)")
    
    # Define the ordered list of tasks.
    task_order = [
        "Text Generation", 
        "Image Generation", 
        "Text Classification", 
        "Image Classification", 
        "Image Captioning", 
        "Summarization", 
        "ASR", 
        "Object Detection", 
        "Question Answering", 
        "Sentence Similarity"
    ]
    
    # Multi-select dropdown for tasks.
    selected_tasks = st.sidebar.multiselect("Select Task(s)", options=task_order, default=task_order)
    if not selected_tasks:
        st.sidebar.error("Please select at least one task.")
        st.stop()
    
    st.sidebar.write("#### 2. Select a model below")
    
    # Mapping from task to CSV file name.
    task_to_file = {
        "Text Generation": "text_gen_energyscore.csv",
        "Image Generation": "image_generation_energyscore.csv",
        "Text Classification": "text_classification_energyscore.csv",
        "Image Classification": "image_classification_energyscore.csv",
        "Image Captioning": "image_caption_energyscore.csv",
        "Summarization": "summarization_energyscore.csv",
        "ASR": "asr_energyscore.csv",
        "Object Detection": "object_detection_energyscore.csv",
        "Question Answering": "question_answering_energyscore.csv",
        "Sentence Similarity": "sentence_similarity_energyscore.csv"
    }
    
    dfs = []
    # Load and process each CSV corresponding to the selected tasks.
    for task in selected_tasks:
        file_name = task_to_file[task]
        try:
            df = pd.read_csv(file_name)
        except FileNotFoundError:
            st.sidebar.error(f"Could not find '{file_name}' for task {task}!")
            continue
        except Exception as e:
            st.sidebar.error(f"Error reading '{file_name}' for task {task}: {e}")
            continue
        
        # Split the "Model" column into 'provider' (before the "/") and 'model' (after the "/")
        df[['provider', 'model']] = df['Model'].str.split('/', 1, expand=True)
        # Round total_gpu_energy to 3 decimal places and assign to 'energy'
        df['energy'] = df['total_gpu_energy'].round(3)
        # Use the energy_score column as 'score'
        df['score'] = df['energy_score'].astype(int)
        # Hardcode date and hardware
        df['date'] = "February 2025"
        df['hardware'] = "NVIDIA H100-80GB"
        # Set the task from the file name mapping
        df['task'] = task
        
        dfs.append(df)
    
    if not dfs:
        st.sidebar.error("No data available for the selected task(s).")
        return
    
    data_df = pd.concat(dfs, ignore_index=True)
    
    # Check required columns
    required_columns = ["model", "provider", "date", "task", "hardware", "energy", "score"]
    for col in required_columns:
        if col not in data_df.columns:
            st.sidebar.error(f"The CSV file must contain a column named '{col}'.")
            return

    model_options = data_df["model"].unique().tolist()
    selected_model = st.sidebar.selectbox(
        "Scored Models",
        model_options,
        help="Start typing to search for a model"
    )
    
    model_data = data_df[data_df["model"] == selected_model].iloc[0]

    # Select background by score
    try:
        score = int(model_data["score"])
        background_path = f"{score}.png"
        background = Image.open(background_path).convert("RGBA")
    except FileNotFoundError:
        st.sidebar.error(f"Could not find background image '{score}.png'. Using default background.")
        background = Image.open("default_background.png").convert("RGBA")
    except ValueError:
        st.sidebar.error(f"Invalid score '{model_data['score']}'. Score must be an integer.")
        return

    # Keep the final label size at 520×728
    final_size = (520, 728)
    generated_label = create_label_single_pass(background, model_data, final_size)

    st.image(generated_label, caption="Generated Label Preview", width=520)

    img_buffer = io.BytesIO()
    generated_label.save(img_buffer, format="PNG")
    img_buffer.seek(0)

    st.sidebar.download_button(
        label="Download",
        data=img_buffer,
        file_name="AIEnergyScore.png",
        mime="image/png"
    )

    st.sidebar.write("#### 3. Share your label! [Guidelines](https://huggingface.github.io/AIEnergyScore/#labelusage)")
    st.sidebar.markdown("<hr style='border: 1px solid gray; margin: 15px 0;'>", unsafe_allow_html=True)
    st.sidebar.write("### Key Links")
    st.sidebar.write("- [Leaderboard](https://huggingface.co/spaces/AIEnergyScore/Leaderboard)")
    st.sidebar.write("- [Submission Portal](https://huggingface.co/spaces/AIEnergyScore/submission_portal)")    
    st.sidebar.write("- [FAQ](https://huggingface.github.io/AIEnergyScore/#faq)")
    st.sidebar.write("- [Documentation](https://huggingface.github.io/AIEnergyScore/#documentation)")
    

def create_label_single_pass(background_image, model_data, final_size=(520, 728)):
    """
    Resizes the background to 520×728, then draws text onto it.
    """
    # 1. Resize background to final_size
    bg_resized = background_image.resize(final_size, Image.Resampling.LANCZOS)
    draw = ImageDraw.Draw(bg_resized)

    # 2. Load fonts at sizes appropriate for a 520×728 label
    try:
        title_font = ImageFont.truetype("Inter_24pt-Bold.ttf", size=27)
        details_font = ImageFont.truetype("Inter_18pt-Regular.ttf", size=23)
        energy_font = ImageFont.truetype("Inter_18pt-Medium.ttf", size=24)
    except Exception as e:
        st.error(f"Font loading failed: {e}")
        return bg_resized

    # 3. Place your text. 
    #    You may need to experiment with x/y coordinates or font sizes
    #    to make it look right in 520×728.
    title_x, title_y = 33, 150
    details_x, details_y = 480, 256
    energy_x, energy_y = 480, 472

    # Text 1 (title) - model and provider separated on different lines
    draw.text((title_x, title_y), str(model_data['model']), font=title_font, fill="black")
    draw.text((title_x, title_y + 38), str(model_data['provider']), font=title_font, fill="black")

    # Text 2 (details)
    details_lines = [
        str(model_data['date']),
        str(model_data['task']),
        str(model_data['hardware'])
    ]
    for i, line in enumerate(details_lines):
        bbox = draw.textbbox((0, 0), line, font=details_font)
        text_width = bbox[2] - bbox[0]
        # Right-justify the details text at details_x
        draw.text((details_x - text_width, details_y + i*47), line, font=details_font, fill="black")

    # Text 3 (energy)
    energy_text = str(model_data['energy'])
    bbox = draw.textbbox((0, 0), energy_text, font=energy_font)
    energy_text_width = bbox[2] - bbox[0]
    # Right-align the energy text at energy_x
    draw.text((energy_x - energy_text_width, energy_y), energy_text, font=energy_font, fill="black")

    return bg_resized


if __name__ == "__main__":
    main()