app.py

import streamlit as st
import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
from PIL import Image
from torchvision import transforms
import numpy as np
import matplotlib.pyplot as plt
import io
import requests
import tempfile
import os

# Set page config
st.set_page_config(
    page_title="Dog Breed Classifier",
    page_icon="🐕",
    layout="wide"
)

# Default model URL
DEFAULT_MODEL_URL = "https://huggingface.co/Alamgirapi/dog-breed-convnext-classifier/resolve/main/model.pth"

# Device setup
@st.cache_resource
def setup_device_and_model():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Initialize model
    model = timm.create_model('convnext_base', pretrained=True)
    
    # Define label names
    label_names = ['beagle', 'bulldog', 'dalmatian', 'german-shepherd', 'husky', 'poodle', 'rottweiler']
    
    # Replace head with proper flattening
    model.head = nn.Sequential(
        nn.AdaptiveAvgPool2d(1),
        nn.Flatten(),
        nn.Linear(model.head.in_features, len(label_names))
    )
    
    model = model.to(device)
    
    # Define transform
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    return device, model, label_names, transform

@st.cache_resource
def download_and_load_model(_model, device):
    """Download and load model weights from Hugging Face"""
    try:
        with st.spinner("Downloading model from Hugging Face..."):
            # Download the model file
            response = requests.get(DEFAULT_MODEL_URL)
            response.raise_for_status()
            
            # Save to temporary file
            with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as tmp_file:
                tmp_file.write(response.content)
                tmp_model_path = tmp_file.name
            
            # Load the model weights
            _model.load_state_dict(torch.load(tmp_model_path, map_location=device))
            _model.eval()
            
            # Clean up temporary file
            os.unlink(tmp_model_path)
            
            return True
    except Exception as e:
        st.error(f"Error downloading/loading model: {str(e)}")
        return False

def predict_image(image, model, transform, label_names, device, topk=3):
    """Make predictions on uploaded image"""
    # Transform image
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    img_tensor = transform(image).unsqueeze(0).to(device)
    
    # Predict
    model.eval()
    with torch.no_grad():
        outputs = model(img_tensor)
        probs = F.softmax(outputs, dim=1)
        top_probs, top_idxs = torch.topk(probs, k=topk)
    
    # Convert to CPU for display
    top_probs = top_probs[0].cpu().numpy()
    top_idxs = top_idxs[0].cpu().numpy()
    
    # Build prediction results
    predictions = []
    for idx, prob in zip(top_idxs, top_probs):
        predictions.append({
            'breed': label_names[idx],
            'confidence': prob * 100
        })
    
    return predictions

def create_prediction_chart(predictions):
    """Create a horizontal bar chart for predictions"""
    breeds = [pred['breed'].replace('-', ' ').title() for pred in predictions]
    confidences = [float(pred['confidence']) for pred in predictions]  # Convert to Python float
    
    fig, ax = plt.subplots(figsize=(10, 6))
    bars = ax.barh(breeds, confidences, color=['#1f77b4', '#ff7f0e', '#2ca02c'])
    
    ax.set_xlabel('Confidence (%)')
    ax.set_title('Top 3 Breed Predictions')
    ax.set_xlim(0, 100)
    
    # Add percentage labels on bars
    for i, (bar, conf) in enumerate(zip(bars, confidences)):
        ax.text(bar.get_width() + 1, bar.get_y() + bar.get_height()/2, 
                f'{conf:.1f}%', va='center')
    
    plt.tight_layout()
    return fig

# Main app
def main():
    st.title("🐕 Dog Breed Classifier")
    st.write("Upload an image of a dog to identify its breed!")
    
    # Initialize model and device
    device, model, label_names, transform = setup_device_and_model()
    
    # Download and load the model automatically
    model_loaded = download_and_load_model(model, device)
    
    if model_loaded:
        st.success("✅ Model loaded successfully from Hugging Face!")
    else:
        st.error("❌ Failed to load model. Please refresh the page and try again.")
        return
    
    # Main content
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.header("Upload Image")
        uploaded_file = st.file_uploader(
            "Choose an image file", 
            type=['jpg', 'jpeg', 'png'],
            help="Upload a clear image of a dog for best results"
        )
        
        if uploaded_file is not None:
            # Display uploaded image
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Image", use_container_width=True)
            
            # Show image details
            st.write(f"**Image Size:** {image.size}")
            st.write(f"**Image Mode:** {image.mode}")
    
    with col2:
        st.header("Predictions")
        
        if uploaded_file is not None:
            try:
                with st.spinner("Analyzing image..."):
                    # Make predictions
                    predictions = predict_image(image, model, transform, label_names, device)
                
                # Display results
                st.success("🎉 Analysis Complete!")
                
                # Show top prediction prominently
                top_breed = predictions[0]['breed'].replace('-', ' ').title()
                top_confidence = float(predictions[0]['confidence'])  # Convert to Python float
                
                st.markdown(f"""

                <div style="background-color: #f0f8ff; padding: 20px; border-radius: 10px; border-left: 5px solid #1f77b4;">

                    <h3 style="color: #1f77b4; margin: 0;">🏆 Most Likely Breed</h3>

                    <h2 style="margin: 5px 0;">{top_breed}</h2>

                    <h4 style="color: #666; margin: 0;">Confidence: {top_confidence:.1f}%</h4>

                </div>

                """, unsafe_allow_html=True)
                
                st.write("") # Add some space
                
                # Show all predictions
                st.subheader("All Predictions:")
                for i, pred in enumerate(predictions):
                    breed = pred['breed'].replace('-', ' ').title()
                    confidence = float(pred['confidence'])  # Convert numpy float32 to Python float
                    
                    # Create progress bar
                    st.write(f"**{i+1}. {breed}**")
                    st.progress(confidence/100)
                    st.write(f"Confidence: {confidence:.2f}%")
                    st.write("")
                
                # Show chart
                st.subheader("Prediction Chart:")
                fig = create_prediction_chart(predictions)
                st.pyplot(fig)
                
            except Exception as e:
                st.error(f"Error during prediction: {str(e)}")
                
        else:
            st.info("📤 Please upload an image to start classification.")
    
    # Information section
    with st.expander("ℹ️ About this App"):
        st.write("""

        This app uses a ConvNeXt-Base model trained to classify dog breeds among:

        - Beagle

        - Bulldog  

        - Dalmatian

        - German Shepherd

        - Husky

        - Poodle

        - Rottweiler

        

        **How to use:**

        1. The model is automatically loaded from Hugging Face

        2. Upload a clear image of a dog

        3. View the top 3 breed predictions with confidence scores

        

        **Tips for better results:**

        - Use high-quality, well-lit images

        - Ensure the dog is clearly visible in the image

        - Avoid images with multiple dogs

        """)
    
    # Technical details
    with st.expander("🔧 Technical Details"):
        st.write(f"""

        - **Device:** {device}

        - **Model:** ConvNeXt-Base

        - **Model Source:** Hugging Face (Alamgirapi/dog-breed-convnext-classifier)

        - **Input Size:** 224x224 pixels

        - **Classes:** {len(label_names)}

        - **Framework:** PyTorch + Streamlit

        """)

if __name__ == "__main__":
    main()