Upload 3 files

app.py

@@ -0,0 +1,52 @@
+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}")

requirements.txt

@@ -1,3 +1,8 @@
-altair
-pandas
-streamlit
+streamlit
+transformers
+torch
+datasets
+scikit-learn
+torchvision
+pyarrow
+Pillow

utils.py

@@ -0,0 +1,61 @@
+import torch
+import numpy as np
+import pyarrow.parquet as pq
+from PIL import Image
+from io import BytesIO
+from sklearn.neighbors import NearestNeighbors
+from torchvision import transforms
+from transformers import (
+    CLIPProcessor,
+    CLIPModel,
+    BlipProcessor,
+    BlipForConditionalGeneration
+)
+
+MEAN = [0.48145466, 0.4578275, 0.40821073]
+STD = [0.26862954, 0.26130258, 0.27577711]
+
+def load_models(device):
+    clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+    clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+    blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
+    blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+    return clip_model, clip_processor, blip_model, blip_processor
+
+def load_data(parquet_path="food101_embeddings.parquet"):
+    table = pq.read_table(parquet_path)
+    df = table.to_pandas()
+    embeddings = np.vstack(df["embedding"].to_numpy())
+    return df, embeddings
+
+def bytes_to_pil(byte_data):
+    return Image.open(BytesIO(byte_data)).convert("RGB")
+
+def preprocess_image(image):
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=MEAN, std=STD),
+    ])
+    return transform(image).unsqueeze(0)
+
+def search_by_text(text, processor, model, embeddings, df, top_k=5, device="cpu"):
+    inputs = processor(text=[text], return_tensors="pt").to(device)
+    with torch.no_grad():
+        text_feat = model.get_text_features(**inputs).cpu().numpy()
+    nn = NearestNeighbors(n_neighbors=top_k, metric="cosine").fit(embeddings)
+    return [{"label": df.iloc[i]["label_name"], "image": bytes_to_pil(df.iloc[i]["image_bytes"])} for i in nn.kneighbors(text_feat, return_distance=False)[0]]
+
+def search_by_image(uploaded_image, processor, model, embeddings, df, top_k=5, device="cpu"):
+    image_tensor = preprocess_image(uploaded_image).to(device)
+    with torch.no_grad():
+        img_feat = model.get_image_features(image_tensor).cpu().numpy()
+    nn = NearestNeighbors(n_neighbors=top_k, metric="cosine").fit(embeddings)
+    return [{"label": df.iloc[i]["label_name"], "image": bytes_to_pil(df.iloc[i]["image_bytes"])} for i in nn.kneighbors(img_feat, return_distance=False)[0]]
+
+def generate_caption(uploaded_image, processor, model, device="cpu"):
+    image = uploaded_image.convert("RGB")
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        output = model.generate(**inputs)
+    return processor.decode(output[0], skip_special_tokens=True)