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(
"""
AI Energy Score
""",
unsafe_allow_html=True,
)
# Sidebar instructions and link
st.sidebar.markdown(
"""
Generate a Label to Display your
AI Energy Score
""",
unsafe_allow_html=True,
)
st.sidebar.markdown("
", 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("
", 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()