src/streamlit_app.py

import streamlit as st
import pandas as pd
import requests
from PIL import Image, ImageStat
from io import BytesIO
import numpy as np
import time
import base64

st.set_page_config(page_title="Amazon Image Optimizer", layout="wide")

st.title("Amazon Image Optimizer")
st.markdown("""
This app helps reorganize your product images to comply with Amazon's listing requirements.
Upload your CSV file containing SKU and image URLs, and the app will:
1. Detect which image has a white background
2. Ensure the white background image is placed in the image1 column
3. Allow you to download the reorganized CSV
""")

def is_white_background(image_url, threshold=240):
    """Determine if an image has a predominantly white background"""
    try:
        response = requests.get(image_url, timeout=5)
        img = Image.open(BytesIO(response.content)).convert('RGB')
        
        # Get the edges of the image (10% from each border)
        width, height = img.size
        border_width = int(width * 0.1)
        border_height = int(height * 0.1)
        
        # Create masks for the edges
        left_edge = img.crop((0, 0, border_width, height))
        right_edge = img.crop((width - border_width, 0, width, height))
        top_edge = img.crop((0, 0, width, border_height))
        bottom_edge = img.crop((0, height - border_height, width, height))
        
        # Calculate average RGB values for edges
        edges = [left_edge, right_edge, top_edge, bottom_edge]
        edge_stats = [ImageStat.Stat(edge) for edge in edges]
        edge_means = [stat.mean for stat in edge_stats]
        
        # Check if edges are predominantly white
        is_white = all(all(channel > threshold for channel in mean) for mean in edge_means)
        
        # Also check overall brightness
        overall_stat = ImageStat.Stat(img)
        overall_brightness = sum(overall_stat.mean) / 3
        
        return is_white and overall_brightness > threshold
    except Exception as e:
        st.error(f"Error processing image: {e}")
        return False

def process_csv(df):
    """Process the dataframe to reorder images putting white background first"""
    progress_bar = st.progress(0)
    status_text = st.empty()
    
    # Create a copy of the original dataframe to preserve all columns
    result_df = df.copy()
    
    # Get the image column names
    image_columns = [col for col in df.columns if col.startswith('image') and col != 'image1']
    
    for i, row in df.iterrows():
        sku = row['sku']
        status_text.text(f"Processing SKU: {sku} ({i+1}/{len(df)})")
        
        white_bg_found = False
        img_urls = []
        white_bg_url = None
        
        # Check if image1 already has white background
        if pd.notna(row['image1']) and is_white_background(row['image1']):
            white_bg_found = True
            white_bg_url = row['image1']
        
        # Collect all other image URLs
        for col in image_columns:
            if pd.notna(row[col]):
                if not white_bg_found and is_white_background(row[col]):
                    white_bg_found = True
                    white_bg_url = row[col]
                else:
                    img_urls.append(row[col])
        
        # Reorganize images if white background found
        if white_bg_found:
            # Store the white background image in image1
            result_df.at[i, 'image1'] = white_bg_url
            
            # Fill in remaining image slots
            remaining_img_count = 0
            for idx, url in enumerate(img_urls, start=0):
                col_name = f'image{idx+2}'
                if col_name in result_df.columns:
                    result_df.at[i, col_name] = url
                    remaining_img_count += 1
        
        progress_bar.progress((i + 1) / len(df))
    
    status_text.text("Processing complete!")
    return result_df

def get_download_link(df, filename):
    """Generate a download link for the dataframe"""
    # Prevent SKUs from being converted to scientific notation
    csv = df.to_csv(index=False, float_format='%.0f')
    b64 = base64.b64encode(csv.encode()).decode()
    href = f'data:file/csv;base64,{b64}'
    return href

# File uploader
uploaded_file = st.file_uploader("Upload your product CSV file", type=["csv"])

if uploaded_file is not None:
    # Load and display the uploaded file
    # Treat SKUs as strings to prevent scientific notation issues
    df = pd.read_csv(uploaded_file, dtype={'sku': str})
    
    st.subheader("Original Data Preview")
    st.dataframe(df.head())
    
    # Check if required columns exist
    required_columns = ['sku', 'image1']
    missing_columns = [col for col in required_columns if col not in df.columns]
    
    if missing_columns:
        st.error(f"Missing required columns: {', '.join(missing_columns)}")
    else:
        if st.button("Process Images"):
            with st.spinner("Processing images..."):
                result_df = process_csv(df)
                
                st.subheader("Results")
                st.dataframe(result_df.head())
                
                # Create download link
                dl_link = get_download_link(result_df, "amazon_optimized_images.csv")
                st.markdown(f'<a href="{dl_link}" download="amazon_optimized_images.csv">Download Processed CSV</a>', unsafe_allow_html=True)
                
                # Summary
                total_rows = len(df)
                white_bg_found = sum(1 for _, row in result_df.iterrows() if pd.notna(row['image1']))
                
                st.subheader("Summary")
                st.write(f"Total products processed: {total_rows}")
                st.write(f"Products with white background images: {white_bg_found}")
                st.write(f"Products missing white background images: {total_rows - white_bg_found}")

# Add instructions and tips
st.sidebar.header("Instructions")
st.sidebar.markdown("""
### CSV Format Requirements
Your CSV file should include:
- A 'sku' column with product identifiers
- Image columns named 'image1', 'image2', etc. containing image URLs

### How It Works
The app analyzes the edges of each image to detect white backgrounds.
It then reorganizes the URLs to ensure white background images are in the image1 position.

### Tips
- Ensure your image URLs are publicly accessible
- The process may take some time for large datasets
- For best results, make sure product images have clear contrasts between product and background
""")