app.py

import streamlit as st
import torch
from PIL import Image
from utils import (
    load_models,
    load_data,
    search_by_text,
    search_by_image,
    generate_caption,
)

st.set_page_config(page_title="🍱 Food Search App", layout="wide")
st.title("🍽️ Food Image & Text Search App")

device = "cuda" if torch.cuda.is_available() else "cpu"

with st.spinner("🔄 Loading models and data..."):
    clip_model, clip_processor, blip_model, blip_processor = load_models(device)
    df, image_embeddings = load_data()

tab1, tab2, tab3 = st.tabs(["🔤 Text Search", "🖼️ Image Search", "📝 Describe Image"])

with tab1:
    st.subheader("Search by Text")
    query = st.text_input("Type a food description (e.g. 'spicy noodles'):")
    if st.button("Search", key="text_search") and query.strip():
        results = search_by_text(query, clip_processor, clip_model, image_embeddings, df, device=device)
        cols = st.columns(5)
        for col, item in zip(cols, results):
            col.image(item["image"], caption=item["label"], use_column_width=True)

with tab2:
    st.subheader("Search by Image")
    uploaded_img = st.file_uploader("Upload a food image", type=["jpg", "jpeg", "png"], key="img_search")
    if uploaded_img:
        image = Image.open(uploaded_img)
        st.image(image, caption="Uploaded image", use_column_width=True)
        if st.button("Find Similar Foods", key="search_image_button"):
            results = search_by_image(image, clip_processor, clip_model, image_embeddings, df, device=device)
            cols = st.columns(5)
            for col, item in zip(cols, results):
                col.image(item["image"], caption=item["label"], use_column_width=True)

with tab3:
    st.subheader("Describe an Image (Auto Caption)")
    uploaded_caption_img = st.file_uploader("Upload a food image", type=["jpg", "jpeg", "png"], key="caption_img")
    if uploaded_caption_img:
        image = Image.open(uploaded_caption_img)
        st.image(image, caption="Uploaded image", use_column_width=True)
        if st.button("Generate Description", key="caption_button"):
            caption = generate_caption(image, blip_processor, blip_model, device=device)
            st.success(f"**Generated Caption:** {caption}")