streamlit_app.py

import streamlit as st
import cv2
import mediapipe as mp
import numpy as np
import os
from PIL import Image

# -------------------------------
# MediaPipe Classifier Setup
# -------------------------------
BaseOptions = mp.tasks.BaseOptions
ImageClassifier = mp.tasks.vision.ImageClassifier
ImageClassifierOptions = mp.tasks.vision.ImageClassifierOptions

model_path = "classifier.tflite"
options = ImageClassifierOptions(
    base_options=BaseOptions(model_asset_path=model_path),
    max_results=5
)
classifier = ImageClassifier.create_from_options(options)

# -------------------------------
# Streamlit UI Setup
# -------------------------------
st.set_page_config(page_title="Image Classifier", layout="wide", page_icon="🛒")

# Compact layout fix — title fully visible
st.markdown(
    """
    <style>
    div.block-container {
        padding-top: 1.5rem;
        padding-bottom: 0.5rem;
        max-width: 100%;
    }
    h1, h2, h3, h4, h5 {font-size: 1rem;}
    img {max-width: 100%; height: auto;}
    .stSlider {margin-top: 0.2rem;}
    </style>
    """,
    unsafe_allow_html=True
)

st.title("E-Commerce Image Classifier")
st.write(
    "Try uploading an image or a folder to see automatic classification results. "
    "You can navigate between images using the arrow buttons below. "
    "This project is open source — check it out on [GitHub](https://github.com/travelmateen/image-classification-ecommerce). 🚀"
)
st.markdown("<style> div[data-testid='stStatusWidget']{display:none}</style>", unsafe_allow_html=True)

# ✅ Sidebar uploader and controls
with st.sidebar:
    st.title("User Configuration")

    num_classes = st.number_input(
        "Number of classes to display",
        min_value=1,
        max_value=5,
        value=3,
        help="Choose how many classification results to show (1-5)"
    )

    # Selection mode (Images or Directory)
    selection_mode = st.radio(
        "Choose upload type:",
        ["Directory", "Select Images"],
        index=0,
        horizontal=True,
    )

    st.header("Upload Your Files")

    if selection_mode == "Directory":
        uploaded_files = st.file_uploader(
            "Upload images from directory",
            accept_multiple_files="directory",
            type=["jpg", "jpeg", "png"],
        )
    else:
        uploaded_files = st.file_uploader(
            "Select individual images",
            type=["jpg", "jpeg", "png"],
            accept_multiple_files=True
        )

    with st.sidebar.expander("⚠️ Limitations & Tips"):
        st.write("""
        **Known Limitations:**
        - Pre-trained MediaPipe general classifier
        - 1000 ImageNet categories only
        - Not customized for specific domains
        - Max 10MB per image

        **For Best Results:**
        - Clear, single-subject images
        - Common objects and scenes
        - Good lighting and focus
        - Avoid ambiguous or complex scenes
        """)

# -------------------------------
# Default folder handling
# -------------------------------
if not uploaded_files:
    default_folder = "images"
    if os.path.exists(default_folder):
        image_files = [
            os.path.join(default_folder, f)
            for f in os.listdir(default_folder)
            if f.lower().endswith((".jpg", ".jpeg", ".png"))
        ]
        if image_files:
            uploaded_files = [open(img, "rb") for img in image_files]

# -------------------------------
# Classification Logic
# -------------------------------
if uploaded_files:
    total_images = len(uploaded_files)

    if 'foo' not in st.session_state:
        st.session_state['foo'] = 0

    current_index = st.session_state['foo']
    # Prevent out-of-range errors
    if current_index >= len(uploaded_files):
        current_index = len(uploaded_files) - 1
        st.session_state['foo'] = current_index
    elif current_index < 0:
        current_index = 0
        st.session_state['foo'] = 0

    current_image = uploaded_files[current_index]

    # --- Read image ---
    file_bytes = np.asarray(bytearray(current_image.read()), dtype=np.uint8)
    frame = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
    if frame is None:
        st.error("⚠️ Unable to read image.")
        st.stop()

    # --- Scale image to 50% ---
    frame = cv2.resize(frame, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA)
    
    # --- Convert to RGB ---
    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # --- Classify image ---
    mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb)
    result = classifier.classify(mp_image)

    # --- Layout: image + classification ---
    col1, col2 = st.columns([1, 1])

    with col1:
        st.subheader("Original Image")
        st.image(rgb, use_container_width=True)

        nav_col1, nav_col2, nav_col3 = st.columns([3, 4, 1], gap="small")
        with nav_col1:
            st.markdown("<div style='text-align:left; margin-top:2px;'>", unsafe_allow_html=True)
            if st.button("⬅️", key="prev") and current_index > 0:
                st.session_state['foo'] = current_index - 1
                st.rerun()
            st.markdown("</div>", unsafe_allow_html=True)
        with nav_col2:
            st.caption(f"🖼️ Image {current_index + 1} of {total_images}")
        with nav_col3:
            st.markdown("<div style='text-align:right; margin-top:2px;'>", unsafe_allow_html=True)
            if st.button("➡️", key="next") and current_index < total_images - 1:
                st.session_state['foo'] = current_index + 1
                st.rerun()
            st.markdown("</div>", unsafe_allow_html=True)

    with col2:
        st.subheader("Classification Results")
        if result.classifications:
            categories = result.classifications[0].categories
            for cat in categories[:num_classes]:
                st.write(f"**{cat.category_name}** ({cat.score:.2f})")
                st.progress(float(cat.score))
        else:
            st.write("No classification detected.")
else:
    st.info("📂 Please upload images using the sidebar to begin classification, or place images in the 'images' folder.")

# -------------------------------
# Footer
# -------------------------------
st.markdown("""
<hr style="border:0;border-top:1px solid #e6eef8;margin:8px 0 4px 0;">
<div style='text-align:center;color:#111F68;margin:0;padding:0;'>
  <p style="margin:0;">Made by <a href='https://techtics.ai' target='_blank' style='color:#042AFF;text-decoration:none;font-weight:700;'>Techtics.ai</a></p>
</div>
""", unsafe_allow_html=True)