Label / app.py
bgamazay's picture
Update app.py
a25259c verified
raw
history blame
3.73 kB
import streamlit as st
import pandas as pd
from PIL import Image, ImageDraw, ImageFont
import io
def main():
# 1. Page Title
st.title("AI Energy Score Label Generator")
# 2. Read Data from CSV
# Using pandas to read the 'data.csv' file
# Make sure 'data.csv' is in the same folder as 'app.py'
try:
data_df = pd.read_csv("data.csv")
except FileNotFoundError:
st.error("Could not find 'data.csv'! Please make sure it's present.")
return
# 3. Ensure the CSV has a "Model" column
if "Model" not in data_df.columns:
st.error("The CSV file must contain a column named 'Model'.")
return
# 4. Create a dropdown list based on unique values in the 'Model' column
model_options = data_df["Model"].unique().tolist()
selected_model = st.selectbox("Select a Model:", model_options)
# 5. Filter the data for the selected model
model_data = data_df[data_df["Model"] == selected_model].iloc[0]
# 6. Load the background image for the label
# Make sure 'background.png' is in the same folder as 'app.py'
try:
background = Image.open("background.png")
except FileNotFoundError:
st.error("Could not find 'background.png'! Please make sure it's present.")
return
# 7. Overlay the data on the image
# We'll create a function to do this cleanly.
generated_label = create_label(background, model_data)
# 8. Display the generated label in the Streamlit app
st.image(generated_label, caption="Generated Label Preview")
# 9. Provide a download button
# We'll create an in-memory file to let user download the image.
img_buffer = io.BytesIO()
generated_label.save(img_buffer, format="PNG")
img_buffer.seek(0)
st.download_button(
label="Download Label as PNG",
data=img_buffer,
file_name="AIEnergyScore.png",
mime="image/png"
)
def create_label(background_image, model_data):
"""
This function takes a background image and a row (model_data) from the CSV,
then draws text on the image. Finally, returns the modified image object.
"""
# Convert background to a format that can be edited (RGBA mode).
label_img = background_image.convert("RGBA")
# Create a Drawing context
draw = ImageDraw.Draw(label_img)
# Choose a font and size. Change path/size as needed.
# If you don't have a custom font file, you can use a PIL built-in font.
try:
font = ImageFont.truetype("Roboto-SemiBold.ttf", 30)
print("Font loaded successfully!")
except Exception as e:
font = ImageFont.load_default()
print(f"Font loading failed: {e}")
# You can customize the positions, colors, etc. For instance:
# We'll just place the text in a simple stacked format at a fixed position.
# Position variables (x, y) - Adjust as needed.
x_position = 75
y_position = 350
line_spacing = 50
# Extracting data from the row. Customize these lines based on your CSV columns.
model_name = f"Model: {model_data['Model']}"
if 'Task' in model_data:
task = f"Task: {model_data['Task']}"
else:
task = "Task: N/A"
if 'date' in model_data:
date = f"Date: {model_data['Date']}"
else:
date = "Date: N/A"
if 'Energy' in model_data:
energy = f"Energy: {model_data['Energy']}"
else:
energy = "Energy: N/A"
text_lines = [model_name, task, date, energy]
# Draw each line on the image
for line in text_lines:
draw.text((x_position, y_position), line, fill="black", font=font)
y_position += line_spacing
return label_img
if __name__ == "__main__":
main()