Label / app.py
bgamazay's picture
Update app.py
66838c3 verified
raw
history blame
4.9 kB
import streamlit as st
import pandas as pd
from PIL import Image, ImageDraw, ImageFont
import io
def main():
# Sidebar for dropdown, buttons, and instructions
st.sidebar.image("logo.png", use_container_width=True) # Display the logo at the top
st.sidebar.title("Label Generator")
st.sidebar.write("### Instructions:")
st.sidebar.write("1. Select a model from the dropdown.")
st.sidebar.write("2. Download the label.")
st.sidebar.write("3. Share your label in technical reports, announcements, etc.")
st.sidebar.markdown("[AI Energy Score Leaderboard](https://huggingface.co/spaces/AIEnergyScore/Leaderboard)")
# Read Data from CSV
try:
data_df = pd.read_csv("data.csv")
except FileNotFoundError:
st.sidebar.error("Could not find 'data.csv'! Please make sure it's present.")
return
# Ensure the CSV has 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
# Dropdown for selecting a model
model_options = data_df["model"].unique().tolist()
selected_model = st.sidebar.selectbox("Select a Model:", model_options)
# Filter the data for the selected model
model_data = data_df[data_df["model"] == selected_model].iloc[0]
# Dynamically select the background image based on the score
try:
score = int(model_data["score"]) # Convert to int
background_path = f"{score}.png" # E.g., "1.png", "2.png"
background = Image.open(background_path).convert("RGBA")
# Proportional scaling to fit within the target size
target_size = (800, 600) # Maximum width and height
background.thumbnail(target_size, Image.Resampling.LANCZOS)
except FileNotFoundError:
st.sidebar.error(f"Could not find background image '{score}.png'. Using default background.")
background = Image.open("default_background.png").convert("RGBA")
background.thumbnail(target_size, Image.Resampling.LANCZOS) # Resize default image proportionally
except ValueError:
st.sidebar.error(f"Invalid score '{model_data['score']}'. Score must be an integer.")
return
# Generate the label with text
generated_label = create_label(background, model_data)
# Display the label
st.image(generated_label, caption="Generated Label Preview")
# Download button for the label
img_buffer = io.BytesIO()
generated_label.save(img_buffer, format="PNG")
img_buffer.seek(0)
st.sidebar.download_button(
label="Download Label as PNG",
data=img_buffer,
file_name="AIEnergyScore.png",
mime="image/png"
)
def create_label(background_image, model_data):
"""
Create the label image by adding text from model_data to the background image.
"""
label_img = background_image.convert("RGBA")
draw = ImageDraw.Draw(label_img)
# Load the Inter variable font (no LAYOUT_RAQM)
try:
inter_font_path = "Inter-VariableFont_opsz,wght.ttf"
title_font = ImageFont.truetype(inter_font_path, 24) # Bold for title
details_font = ImageFont.truetype(inter_font_path, 20) # Medium for details
energy_font = ImageFont.truetype(inter_font_path, 22) # Medium for energy
# Set bold weight for title font
title_font = title_font.font_variant(weight=700) # Set font weight to bold
except Exception as e:
st.error(f"Font loading failed: {e}")
return label_img
# Define positions for each text group
title_x, title_y = 28, 124
details_x, details_y = 375, 208
energy_x, energy_y = 350, 388
# Group 1: Title (Left-Justified, no prefixes)
draw.text((title_x, title_y), str(model_data['model']), font=title_font, fill="black")
draw.text((title_x, title_y + 30), str(model_data['provider']), font=title_font, fill="black")
# Group 2: Details (Right-Justified, no prefixes)
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]
draw.text((details_x - text_width, details_y + i * 40), line, font=details_font, fill="black")
# Group 3: Energy (Bottom-Center, no prefixes)
energy_text = str(model_data['energy']) # Ensure this is a string
bbox = draw.textbbox((0, 0), energy_text, font=energy_font)
energy_text_width = bbox[2] - bbox[0]
draw.text((energy_x - energy_text_width // 2, energy_y), energy_text, font=energy_font, fill="black")
return label_img
if __name__ == "__main__":
main()