using pretrained model

pages/1_Image_Classifier.py

@@ -1,54 +1,63 @@
 import streamlit as st
 from PIL import Image
-import requests
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+import torch
 
-st.set_page_config(
-    page_title="Food Classifier",
-    page_icon="🍕",
-    layout="wide"
-)
+# Streamlit page config
+st.set_page_config(page_title="Food Classifier", layout="wide")
 
 st.title("Food Classification")
-st.markdown("AI-powered food recognition")
-
-API_URL = "https://api-inference.huggingface.co/models/nateraw/vit-base-beans"
-
-def classify_image(image_bytes):
-    try:
-        response = requests.post(API_URL, data=image_bytes, timeout=30)
-        return response.json()
-    except:
-        return None
-
+st.markdown("AI-powered food recognition using a pretrained Food101 model")
+
+# Load processor and model
+processor = AutoImageProcessor.from_pretrained("nateraw/food")
+model = AutoModelForImageClassification.from_pretrained("nateraw/food")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+model.eval()  # set model to evaluation mode
+
+# Function to preprocess and classify image
+def classify_image(image: Image.Image):
+    inputs = processor(images=image.convert("RGB"), return_tensors="pt")
+    # Move to device
+    for k, v in inputs.items():
+        inputs[k] = v.to(device)
+    # Forward pass
+    with torch.no_grad():
+        outputs = model(**inputs)
+    # Softmax probabilities
+    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    # Get top 5 predictions
+    top_prob, top_class = torch.topk(probs, k=5)
+    id2label = model.config.id2label
+    results = []
+    for prob, idx in zip(top_prob[0], top_class[0]):
+        results.append({"label": id2label[idx.item()], "score": prob.item()})
+    return results
+
+# Streamlit layout
 col1, col2 = st.columns(2)
 
 with col1:
     st.subheader("Upload Image")
-    uploaded_file = st.file_uploader("Choose a food image", type=['jpg', 'jpeg', 'png'])
-    
+    uploaded_file = st.file_uploader("Choose a food image", type=["jpg", "jpeg", "png"])
     if uploaded_file:
         image = Image.open(uploaded_file)
-        st.image(image, use_column_width=True)
+        st.image(image, use_container_width=True)
 
 with col2:
     st.subheader("Results")
-    
     if uploaded_file:
         with st.spinner("Analyzing..."):
-            results = classify_image(uploaded_file.getvalue())
-            
-            if results and isinstance(results, list):
-                for i, result in enumerate(results[:5], 1):
-                    label = result.get('label', 'Unknown')
-                    score = result.get('score', 0)
-                    
-                    st.write(f"**{i}. {label}**")
-                    st.progress(score)
-                    st.caption(f"{score*100:.1f}%")
+            results = classify_image(image)
+            if results:
+                for i, r in enumerate(results, 1):
+                    st.write(f"{i}. {r['label']} — {r['score']*100:.2f}%")
+                    st.progress(r['score'])
             else:
-                st.info("Model loading. Wait 20 seconds and retry.")
+                st.info("Could not classify the image. Try again.")
     else:
-        st.info("Upload an image to classify")
+        st.info("Upload an image to classify.")
 
 st.markdown("---")
-st.caption("Week 1 Project - Image Classification Pipeline")
+st.caption("Week 1 Project - Image Classification Pipeline")