Load model directly from Hub (no Inference API)

pages/1_Image_Classifier.py

@@ -1,8 +1,7 @@
 import streamlit as st
 from PIL import Image
-import requests
-import json
-import time
+import torch
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 
 st.set_page_config(
     page_title="Food Classifier",
@@ -11,49 +10,31 @@ st.set_page_config(
 )
 
 st.title("Food-101 Image Classifier")
-st.markdown("### AI-powered food recognition using ResNet-50")
+st.markdown("### ResNet-50 trained on 75K food images")
 
-# Model API
-API_URL = "https://api-inference.huggingface.co/models/nateraw/food"
+# Load model (cached)
+@st.cache_resource
+def load_model():
+    processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
+    model = AutoModelForImageClassification.from_pretrained("suchithnj12/food101-resnet50")
+    model.eval()
+    return processor, model
 
-# Try to get token, ignore if not found
 try:
-    headers = {"Authorization": f"Bearer {st.secrets['HF_TOKEN']}"}
-except:
-    headers = {}
+    with st.spinner("Loading model (first time takes 30 seconds)..."):
+        processor, model = load_model()
+    st.success("✅ Model loaded!")
+except Exception as e:
+    st.error(f"Error loading model: {str(e)}")
+    st.stop()
 
-def query(image_bytes, max_retries=3):
-    """Query HuggingFace Inference API with retry logic"""
-    for attempt in range(max_retries):
-        try:
-            response = requests.post(API_URL, headers=headers, data=image_bytes, timeout=30)
-            
-            if response.status_code == 503:
-                return {"error": "loading", "message": "Model is loading. Please wait."}
-            
-            if response.status_code == 200:
-                return response.json()
-            
-            return {"error": "api_error", "status": response.status_code, "message": response.text}
-            
-        except requests.exceptions.Timeout:
-            if attempt < max_retries - 1:
-                time.sleep(5)
-                continue
-            return {"error": "timeout", "message": "Request timed out. Please try again."}
-        except Exception as e:
-            return {"error": "exception", "message": str(e)}
-    
-    return {"error": "failed", "message": "Failed after multiple attempts"}
-
-# Model info
 with st.expander("ℹ️ Model Information"):
     st.markdown("""
-    **Model**: ResNet-50 (Fine-tuned)  
-    **Dataset**: Food-101 (101 food categories)  
-    **Training**: 5 epochs, 75K images  
-    **Accuracy**: 41% (vs 1% random baseline)  
-    **HuggingFace**: [View Model](https://huggingface.co/suchithnj12/food101-resnet50)
+    **Your Model**: suchithnj12/food101-resnet50  
+    **Base**: ResNet-50  
+    **Training**: 5 epochs on Food-101  
+    **Accuracy**: 40.8%  
+    **Categories**: 101 food types
     """)
 
 # Main interface
@@ -64,8 +45,7 @@ with col1:
     
     uploaded_file = st.file_uploader(
         "Choose a food image",
-        type=['jpg', 'jpeg', 'png'],
-        help="Upload an image of food"
+        type=['jpg', 'jpeg', 'png']
     )
     
     if uploaded_file:
@@ -76,86 +56,43 @@ with col2:
     st.markdown("### Prediction Results")
     
     if uploaded_file:
-        with st.spinner("Analyzing image..."):
-            # Get image bytes
-            image_bytes = uploaded_file.getvalue()
-            
-            # Query API
-            results = query(image_bytes)
-            
-            # Handle errors
-            if isinstance(results, dict) and "error" in results:
-                if results["error"] == "loading":
-                    st.warning("⏳ Model is loading on HuggingFace servers...")
-                    st.info("First request takes 20-30 seconds. Please wait and try again.")
-                    if st.button("Retry"):
-                        st.rerun()
-                else:
-                    st.error(f"Error: {results.get('message', 'Unknown error')}")
-                    st.info("Please try again in a few seconds.")
-            
-            # Handle successful results
-            elif isinstance(results, list) and len(results) > 0:
-                # Display top prediction
-                top_result = results[0]
-                top_class = top_result['label'].replace('_', ' ').title()
-                top_confidence = top_result['score']
+        with st.spinner("Analyzing..."):
+            try:
+                # Preprocess
+                inputs = processor(images=image, return_tensors="pt")
                 
-                st.markdown(f"### 🎯 {top_class}")
-                st.progress(top_confidence)
-                st.metric("Confidence", f"{top_confidence*100:.1f}%")
+                # Predict
+                with torch.no_grad():
+                    outputs = model(**inputs)
+                    probs = torch.nn.functional.softmax(outputs.logits, dim=1)
+                    top5_probs, top5_indices = torch.topk(probs, 5)
                 
-                st.markdown("---")
+                # Display results
+                st.success("✅ Analysis Complete")
                 
-                # Display top 5
-                st.markdown("#### Top 5 Predictions")
-                for i, result in enumerate(results[:5]):
-                    class_name = result['label'].replace('_', ' ').title()
-                    confidence = result['score']
+                for i in range(5):
+                    label = model.config.id2label[top5_indices[0][i].item()]
+                    score = top5_probs[0][i].item()
+                    
+                    # Format label
+                    label = label.replace('_', ' ').title()
                     
-                    st.markdown(f"**{i+1}. {class_name}**")
-                    st.progress(confidence)
-                    st.caption(f"{confidence*100:.1f}%")
-            else:
-                st.error("Unexpected response from model.")
-                st.info("The model might be initializing. Please wait 30 seconds and try again.")
+                    st.markdown(f"**{i+1}. {label}**")
+                    st.progress(score)
+                    st.caption(f"{score*100:.1f}%")
+                    
+            except Exception as e:
+                st.error(f"Prediction failed: {str(e)}")
     else:
         st.info("👈 Upload an image to get started")
 
-# Sample images section - REMOVED broken image URLs
-st.markdown("---")
-st.markdown("### 💡 Tips for Best Results")
-st.info("""
-- Use clear, well-lit food images
-- Ensure food is the main subject
-- Works best with common dishes
-- First prediction may take 20-30 seconds (model loading)
-""")
-
-# Technical details
 st.markdown("---")
 with st.expander("🔧 Technical Details"):
     st.markdown("""
-    **Architecture**: ResNet-50 with modified classification head
-    
-    **Inference**: HuggingFace Inference API (serverless)
-    
-    **Training Details**:
-    - Optimizer: AdamW
-    - Learning Rate: 2e-5
-    - Batch Size: 64
-    - Mixed Precision: FP16
-    
-    **Performance**:
-    - Test Accuracy: 40.8%
-    - F1 Score: 38.0%
-    - Inference Time: 1-3 seconds (after initial load)
-    
-    **Sample Categories**: 
-    pizza, sushi, hamburger, pasta, steak, salad, ice cream, cake, and 93 more...
-    
-    **Note**: First request takes 20-30 seconds as the model loads on HuggingFace servers.
+    **Model Loading**: Direct from HuggingFace Hub  
+    **Inference**: On HuggingFace Spaces hardware  
+    **Test Accuracy**: 40.8%  
+    **Categories**: apple pie, sushi, pizza, pasta, and 97 more
     """)
 
 st.markdown("---")
-st.caption("Week 1 Complete - Using HuggingFace Inference API")