Upload 2 files

app.py

@@ -0,0 +1,617 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import streamlit as st
+import numpy as np
+import torchvision.transforms as transforms
+from PIL import Image, ImageDraw
+import os
+import base64
+from io import BytesIO
+
+# Define the neural network model - matching your trained model with 3 input channels
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        # 3 input image channels (RGB), 6 output channels, 5x5 square convolution kernel
+        self.conv1 = nn.Conv2d(3, 6, 5)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        # an affine operation: y = Wx + b
+        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5*5 from image dimension 
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x):
+        # Convolution layer C1: 3 input image channels, 6 output channels,
+        # 5x5 square convolution, it uses RELU activation function, and
+        # outputs a Tensor with size (N, 6, 28, 28), where N is the size of the batch
+        c1 = F.relu(self.conv1(x))
+        # Subsampling layer S2: 2x2 grid, purely functional,
+        # this layer does not have any parameter, and outputs a (N, 6, 14, 14) Tensor
+        s2 = F.max_pool2d(c1, (2, 2))
+        # Convolution layer C3: 6 input channels, 16 output channels,
+        # 5x5 square convolution, it uses RELU activation function, and
+        # outputs a (N, 16, 10, 10) Tensor
+        c3 = F.relu(self.conv2(s2))
+        # Subsampling layer S4: 2x2 grid, purely functional,
+        # this layer does not have any parameter, and outputs a (N, 16, 5, 5) Tensor
+        s4 = F.max_pool2d(c3, 2)
+        # Flatten operation: purely functional, outputs a (N, 400) Tensor
+        s4 = torch.flatten(s4, 1)
+        # Fully connected layer F5: (N, 400) Tensor input,
+        # and outputs a (N, 120) Tensor, it uses RELU activation function
+        f5 = F.relu(self.fc1(s4))
+        # Fully connected layer F6: (N, 120) Tensor input,
+        # and outputs a (N, 84) Tensor, it uses RELU activation function
+        f6 = F.relu(self.fc2(f5))
+        # Gaussian layer OUTPUT: (N, 84) Tensor input, and
+        # outputs a (N, 10) Tensor
+        output = self.fc3(f6)
+        return output
+
+# Initialize the model
+model = Net()
+
+# Load the trained model weights
+def load_model():
+    model_path = "model.pth"  # Update this path to where your model is stored
+    if os.path.exists(model_path):
+        try:
+            # Load the trained model weights
+            # Handle different PyTorch versions
+            try:
+                # For PyTorch 2.6+, we need to set weights_only=False for compatibility
+                model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'), weights_only=False))
+            except TypeError:
+                # For older PyTorch versions that don't support weights_only parameter
+                model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+            print("Loaded trained model weights")
+            return True
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            return False
+    else:
+        print("No trained model found at", model_path)
+        # Initialize with random weights for demonstration
+        for m in model.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+        return False
+
+# Preprocessing function for input images - now handles RGB images
+def preprocess_image(image):
+    # Resize to 32x32 (expected input size for the network)
+    transform = transforms.Compose([
+        transforms.Resize((32, 32)),
+        transforms.ToTensor(),
+    ])
+    
+    image_tensor = transform(image)
+    # Add batch dimension (1, 3, 32, 32)
+    image_tensor = image_tensor.unsqueeze(0)
+    return image_tensor
+
+# Prediction function - matches the PyTorch tutorial exactly
+def predict(image):
+    if image is None:
+        return {f"Class {i}": 0 for i in range(10)}
+    
+    # Preprocess the image
+    input_tensor = preprocess_image(image)
+    
+    # Make prediction - exactly as shown in the PyTorch tutorial
+    model.eval()
+    with torch.no_grad():
+        output = model(input_tensor)
+        # Apply softmax to get probabilities
+        probabilities = F.softmax(output, dim=1)
+        probabilities = probabilities.numpy()[0]
+    
+    # Create labels for CIFAR-10 classes
+    cifar10_classes = ["Airplane", "Automobile", "Bird", "Cat", "Deer", "Dog", "Frog", "Horse", "Ship", "Truck"]
+    
+    # Return as a dictionary
+    return {label: float(prob) for label, prob in zip(cifar10_classes, probabilities)}
+
+# Create example images representing CIFAR-10 classes
+def create_example_images():
+    examples = []
+    example_names = []
+    
+    # CIFAR-10 class names
+    cifar10_classes = ["Airplane", "Automobile", "Bird", "Cat", "Deer", "Dog", "Frog", "Horse", "Ship", "Truck"]
+    
+    # Create simple representations of CIFAR-10 classes
+    for i, class_name in enumerate(cifar10_classes):
+        # Create a 64x64 RGB image for better quality
+        img = Image.new('RGB', (64, 64), color=(255, 255, 255))  # White background
+        draw = ImageDraw.Draw(img)
+        
+        # Draw simple representations of each class
+        if i == 0:  # Airplane
+            # Draw a simple airplane shape
+            draw.polygon([(32, 10), (20, 30), (44, 30)], fill=(169, 169, 169))  # Main body
+            draw.rectangle([25, 30, 39, 35], fill=(105, 105, 105))  # Wings
+            draw.rectangle([30, 35, 34, 45], fill=(128, 128, 128))  # Tail
+        elif i == 1:  # Automobile
+            # Draw a simple car shape
+            draw.rectangle([15, 30, 49, 45], fill=(0, 0, 255))  # Body
+            draw.ellipse([20, 40, 30, 50], fill=(0, 0, 0))  # Wheels
+            draw.ellipse([34, 40, 44, 50], fill=(0, 0, 0))
+            draw.rectangle([25, 20, 39, 30], fill=(0, 0, 255))  # Top
+        elif i == 2:  # Bird
+            # Draw a simple bird shape
+            draw.ellipse([25, 25, 39, 39], fill=(255, 165, 0))  # Body
+            draw.polygon([(32, 15), (25, 25), (39, 25)], fill=(255, 140, 0))  # Head
+            draw.line([20, 30, 10, 20], fill=(255, 165, 0), width=3)  # Wing
+            draw.line([44, 30, 54, 20], fill=(255, 165, 0), width=3)  # Wing
+        elif i == 3:  # Cat
+            # Draw a simple cat shape
+            draw.ellipse([25, 25, 39, 39], fill=(128, 128, 128))  # Body
+            draw.ellipse([30, 20, 40, 30], fill=(169, 169, 169))  # Head
+            draw.polygon([(35, 22), (33, 27), (37, 27)], fill=(0, 0, 0))  # Ear
+            draw.ellipse([32, 28, 34, 30], fill=(0, 0, 0))  # Eye
+        elif i == 4:  # Deer
+            # Draw a simple deer shape
+            draw.ellipse([25, 30, 39, 44], fill=(139, 69, 19))  # Body
+            draw.ellipse([30, 25, 40, 35], fill=(160, 82, 45))  # Head
+            draw.line([35, 15, 40, 25], fill=(139, 69, 19), width=3)  # Antler
+            draw.line([20, 35, 10, 30], fill=(139, 69, 19), width=2)  # Leg
+        elif i == 5:  # Dog
+            # Draw a simple dog shape
+            draw.ellipse([25, 30, 39, 44], fill=(139, 69, 19))  # Body
+            draw.ellipse([30, 25, 40, 35], fill=(160, 82, 45))  # Head
+            draw.ellipse([32, 28, 34, 30], fill=(0, 0, 0))  # Eye
+            draw.ellipse([36, 32, 38, 34], fill=(0, 0, 0))  # Nose
+        elif i == 6:  # Frog
+            # Draw a simple frog shape
+            draw.ellipse([25, 30, 39, 44], fill=(34, 139, 34))  # Body
+            draw.ellipse([30, 25, 40, 35], fill=(0, 100, 0))  # Head
+            draw.ellipse([27, 32, 29, 34], fill=(0, 0, 0))  # Eye
+            draw.ellipse([35, 32, 37, 34], fill=(0, 0, 0))  # Eye
+        elif i == 7:  # Horse
+            # Draw a simple horse shape
+            draw.ellipse([25, 30, 39, 44], fill=(169, 169, 169))  # Body
+            draw.ellipse([35, 20, 45, 30], fill=(128, 128, 128))  # Head
+            draw.line([40, 25, 50, 15], fill=(105, 105, 105), width=3)  # Mane
+        elif i == 8:  # Ship
+            # Draw a simple ship shape
+            draw.polygon([(20, 35), (44, 35), (38, 45), (26, 45)], fill=(139, 69, 19))  # Hull
+            draw.rectangle([30, 20, 34, 35], fill=(169, 169, 169))  # Mast
+            draw.polygon([(30, 20), (32, 15), (34, 20)], fill=(255, 255, 255))  # Sail
+        elif i == 9:  # Truck
+            # Draw a simple truck shape
+            draw.rectangle([15, 25, 49, 45], fill=(255, 0, 0))  # Cab
+            draw.rectangle([25, 15, 45, 25], fill=(255, 0, 0))  # Load area
+            draw.ellipse([20, 40, 30, 50], fill=(0, 0, 0))  # Wheels
+            draw.ellipse([34, 40, 44, 50], fill=(0, 0, 0))
+    
+        examples.append(img)
+        example_names.append(class_name)
+    
+    return examples, example_names
+
+# Function to convert PIL Image to base64 for display
+def image_to_base64(image):
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    return img_str
+
+# Initialize the model
+model_loaded = load_model()
+
+# Create example images
+examples, example_names = create_example_images()
+
+# Streamlit app
+st.set_page_config(
+    page_title="CIFAR-10 Image Classifier",
+    page_icon="🚀",
+    layout="wide"
+)
+
+# Custom CSS with cleaner design
+st.markdown("""
+<style>
+/* Import Google Fonts */
+@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;500;600;700&display=swap');
+
+/* Base styles */
+* {
+    font-family: 'Poppins', sans-serif;
+}
+
+/* Clean background */
+body {
+    background: linear-gradient(135deg, #1a2a6c, #2c3e50);
+    color: white;
+}
+
+/* Main container with clean glassmorphism effect */
+.main-container {
+    background: rgba(255, 255, 255, 0.05);
+    backdrop-filter: blur(10px);
+    border-radius: 20px;
+    border: 1px solid rgba(255, 255, 255, 0.1);
+    box-shadow: 0 8px 32px 0 rgba(0, 0, 0, 0.3);
+    padding: 2rem;
+    margin: 2rem auto;
+    max-width: 1200px;
+}
+
+/* Title with clean gradient */
+.title {
+    background: linear-gradient(90deg, #4facfe 0%, #00f2fe 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+    font-weight: 800;
+    font-size: 2.5rem;
+    text-align: center;
+    margin-bottom: 0.5rem;
+}
+
+/* Subtitle styling */
+.subtitle {
+    text-align: center;
+    color: #a0d2ff;
+    font-size: 1.1rem;
+    margin-bottom: 2rem;
+    opacity: 0.9;
+}
+
+/* Card styling */
+.card {
+    background: rgba(255, 255, 255, 0.05);
+    border-radius: 15px;
+    padding: 1.5rem;
+    margin-bottom: 1.5rem;
+    border: 1px solid rgba(255, 255, 255, 0.1);
+    transition: all 0.3s ease;
+    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.15);
+}
+
+.card:hover {
+    background: rgba(255, 255, 255, 0.08);
+    box-shadow: 0 6px 25px rgba(0, 0, 0, 0.25);
+    transform: translateY(-3px);
+}
+
+/* Section headers */
+.section-header {
+    color: #4facfe;
+    border-bottom: 2px solid #00f2fe;
+    padding-bottom: 0.5rem;
+    margin-bottom: 1rem;
+    font-weight: 600;
+    font-size: 1.3rem;
+}
+
+/* Button styling */
+.stButton > button {
+    background: linear-gradient(90deg, #4facfe 0%, #00f2fe 100%);
+    color: white;
+    border: none;
+    border-radius: 10px;
+    padding: 0.7rem 1.2rem;
+    font-weight: 600;
+    transition: all 0.3s ease;
+    box-shadow: 0 4px 15px rgba(79, 172, 254, 0.3);
+    width: 100%;
+}
+
+.stButton > button:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 6px 20px rgba(79, 172, 254, 0.5);
+}
+
+.stButton > button:active {
+    transform: translateY(1px);
+}
+
+/* File uploader styling */
+.stFileUploader > div {
+    background: rgba(255, 255, 255, 0.05);
+    border-radius: 15px;
+    border: 1px dashed rgba(255, 255, 255, 0.3);
+    padding: 1.5rem;
+    text-align: center;
+}
+
+/* Progress bar styling */
+.stProgress > div > div {
+    background: linear-gradient(90deg, #4facfe 0%, #00f2fe 100%);
+}
+
+/* Result display */
+.result-container {
+    display: flex;
+    flex-wrap: wrap;
+    gap: 0.8rem;
+    justify-content: center;
+}
+
+.result-item {
+    background: rgba(255, 255, 255, 0.08);
+    border-radius: 12px;
+    padding: 1rem;
+    text-align: center;
+    min-width: 110px;
+    transition: all 0.3s ease;
+    border: 1px solid rgba(255, 255, 255, 0.1);
+}
+
+.result-item:hover {
+    background: rgba(79, 172, 254, 0.2);
+    transform: translateY(-3px);
+    box-shadow: 0 5px 15px rgba(0, 0, 0, 0.2);
+}
+
+.result-label {
+    font-weight: 600;
+    margin-bottom: 0.4rem;
+    color: #4facfe;
+    font-size: 0.9rem;
+}
+
+.result-value {
+    font-size: 1.1rem;
+    font-weight: 700;
+    color: white;
+}
+
+/* Example images grid */
+.examples-grid {
+    display: grid;
+    grid-template-columns: repeat(auto-fill, minmax(90px, 1fr));
+    gap: 0.8rem;
+    margin-top: 1rem;
+}
+
+.example-item {
+    cursor: pointer;
+    border-radius: 10px;
+    overflow: hidden;
+    transition: all 0.3s ease;
+    border: 2px solid transparent;
+    background: rgba(255, 255, 255, 0.05);
+}
+
+.example-item:hover {
+    transform: scale(1.05);
+    border-color: #4facfe;
+    box-shadow: 0 5px 15px rgba(0, 0, 0, 0.3);
+    background: rgba(79, 172, 254, 0.1);
+}
+
+.example-item img {
+    border-radius: 8px;
+}
+
+.example-name {
+    text-align: center;
+    margin-top: 5px;
+    font-size: 0.75rem;
+    color: #a0d2ff;
+}
+
+/* Footer */
+.footer {
+    text-align: center;
+    padding: 1.5rem;
+    color: rgba(255, 255, 255, 0.6);
+    font-size: 0.9rem;
+}
+
+/* Responsive design */
+@media (max-width: 768px) {
+    .main-container {
+        padding: 1rem;
+        margin: 1rem;
+    }
+    
+    .title {
+        font-size: 2rem;
+    }
+    
+    .card {
+        padding: 1rem;
+    }
+    
+    .result-item {
+        min-width: 90px;
+        padding: 0.7rem;
+    }
+    
+    .examples-grid {
+        grid-template-columns: repeat(auto-fill, minmax(70px, 1fr));
+    }
+}
+</style>
+""", unsafe_allow_html=True)
+
+# Main app content
+st.markdown('<div class="main-container">', unsafe_allow_html=True)
+
+st.markdown('<h1 class="title">🚀 CIFAR-10 Image Classifier</h1>', unsafe_allow_html=True)
+st.markdown('<p class="subtitle">Convolutional Neural Network for Object Recognition</p>', unsafe_allow_html=True)
+
+# Show model loading status
+if model_loaded:
+    st.success("✅ Model successfully loaded")
+else:
+    st.warning("⚠️ Model not found or error loading. Using random weights for demonstration.")
+
+# Create tabs for better organization
+tab1, tab2, tab3 = st.tabs(["🔍 Classify", "🖼️ Examples", "📚 Information"])
+
+with tab1:
+    # Create two columns for input and output
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        st.markdown('<div class="card">', unsafe_allow_html=True)
+        st.markdown('<h2 class="section-header">📤 Input</h2>', unsafe_allow_html=True)
+        
+        # File uploader
+        uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+        
+        # Display image
+        image = None
+        if uploaded_file is not None:
+            image = Image.open(uploaded_file).convert('RGB')
+            st.image(image, caption="Uploaded Image", use_container_width=True)
+        
+        # Classify button
+        if st.button("Classify Image"):
+            if image is not None:
+                st.session_state.predictions = predict(image)
+            else:
+                st.warning("Please upload an image first")
+        
+        # Clear button
+        if st.button("Clear"):
+            st.session_state.predictions = None
+            st.experimental_rerun()
+        
+        st.markdown('</div>', unsafe_allow_html=True)
+        
+        # Model architecture section
+        st.markdown('<div class="card">', unsafe_allow_html=True)
+        st.markdown('<h2 class="section-header">🎯 Model Architecture</h2>', unsafe_allow_html=True)
+        st.code("""
+Input → Conv2D(3×32×32) → ReLU → MaxPool2D
+     → Conv2D → ReLU → MaxPool2D
+     → Flatten → Linear → ReLU
+     → Linear → ReLU → Linear(10)
+     → Output
+        """, language="text")
+        st.markdown('</div>', unsafe_allow_html=True)
+    
+    with col2:
+        st.markdown('<div class="card">', unsafe_allow_html=True)
+        st.markdown('<h2 class="section-header">📊 Classification Results</h2>', unsafe_allow_html=True)
+        
+        # Display results
+        if "predictions" in st.session_state and st.session_state.predictions:
+            predictions = st.session_state.predictions
+            # Sort predictions by probability
+            sorted_predictions = sorted(predictions.items(), key=lambda x: x[1], reverse=True)
+            
+            # Display top 5 predictions with animated bars
+            st.markdown('<div class="result-container">', unsafe_allow_html=True)
+            for label, prob in sorted_predictions[:5]:
+                st.markdown(f'''
+                    <div class="result-item">
+                        <div class="result-label">{label}</div>
+                        <div class="result-value">{prob:.2f}</div>
+                    </div>
+                ''', unsafe_allow_html=True)
+            st.markdown('</div>', unsafe_allow_html=True)
+            
+            # Display all probabilities in a more detailed way
+            st.subheader("All Class Probabilities")
+            for label, prob in sorted_predictions:
+                st.progress(prob)
+                st.write(f"{label}: {prob:.4f}")
+        else:
+            st.info("Upload an image and click 'Classify Image' to see results")
+        
+        st.markdown('</div>', unsafe_allow_html=True)
+        
+        # Instructions section
+        st.markdown('<div class="card">', unsafe_allow_html=True)
+        st.markdown('<h2 class="section-header">ℹ️ Instructions</h2>', unsafe_allow_html=True)
+        st.markdown("""
+        1. Upload an image using the file uploader
+        2. The image will be automatically resized to 32×32 pixels
+        3. Click "Classify Image" to get predictions
+        4. Results show probabilities for 10 CIFAR-10 classes
+        """)
+        st.markdown('</div>', unsafe_allow_html=True)
+
+with tab2:
+    # Example images section
+    st.markdown('<div class="card">', unsafe_allow_html=True)
+    st.markdown('<h2 class="section-header">🖼️ Example Images</h2>', unsafe_allow_html=True)
+    st.markdown("Click on any example image to classify it:")
+    
+    # Create example grid
+    st.markdown('<div class="examples-grid">', unsafe_allow_html=True)
+    for i, (example_img, example_name) in enumerate(zip(examples, example_names)):
+        # Convert PIL image to base64
+        img_base64 = image_to_base64(example_img)
+        
+        # Create clickable image
+        if st.button(f"example_{i}", key=f"btn_{i}"):
+            st.session_state.predictions = predict(example_img)
+            st.experimental_rerun()
+        
+        st.markdown(f'''
+            <div class="example-item">
+                <img src="data:image/png;base64,{img_base64}" width="100" height="100" alt="{example_name}">
+                <div class="example-name">{example_name}</div>
+            </div>
+        ''', unsafe_allow_html=True)
+    st.markdown('</div>', unsafe_allow_html=True)
+    st.markdown('</div>', unsafe_allow_html=True)
+
+with tab3:
+    # Information sections
+    st.markdown('<div class="card">', unsafe_allow_html=True)
+    st.markdown('<h2 class="section-header">🧪 Testing Different Image Qualities</h2>', unsafe_allow_html=True)
+    st.markdown("""
+    This model is robust to various image conditions:
+    - **Resolution**: Works with images of any resolution (automatically resized to 32×32)
+    - **Contrast**: Handles both high and low contrast images
+    - **Noise**: Can tolerate some image noise
+    - **Rotation**: Some tolerance to slight rotations
+    - **Scale**: Works with objects of different sizes within the image
+    
+    For best results:
+    1. Center the object in the image
+    2. Use clear contrast between the object and background
+    3. Avoid excessive noise or artifacts
+    4. Fill most of the image area with the object
+    """)
+    st.markdown('</div>', unsafe_allow_html=True)
+    
+    st.markdown('<div class="card">', unsafe_allow_html=True)
+    st.markdown('<h2 class="section-header">🎯 CIFAR-10 Classes</h2>', unsafe_allow_html=True)
+    classes_info = """
+    1. **Airplane** - Aircraft flying in the sky
+    2. **Automobile** - Cars and vehicles on the road
+    3. **Bird** - Flying or perched birds
+    4. **Cat** - Domestic cats and felines
+    5. **Deer** - Wild deer and similar animals
+    6. **Dog** - Domestic dogs and canines
+    7. **Frog** - Amphibians like frogs
+    8. **Horse** - Horses and similar animals
+    9. **Ship** - Boats and ships on water
+    10. **Truck** - Trucks and heavy vehicles
+    """
+    st.markdown(classes_info)
+    st.markdown('</div>', unsafe_allow_html=True)
+    
+    # Model architecture section
+    st.markdown('<div class="card">', unsafe_allow_html=True)
+    st.markdown('<h2 class="section-header">🧠 Model Details</h2>', unsafe_allow_html=True)
+    st.markdown("""
+    This convolutional neural network follows the PyTorch CIFAR-10 tutorial architecture:
+    - **Input Layer**: 3×32×32 RGB images
+    - **Convolutional Layers**: 2 layers with ReLU activation
+    - **Pooling Layers**: 2 max-pooling layers
+    - **Fully Connected Layers**: 3 linear layers
+    - **Output Layer**: 10 classes with softmax activation
+    """)
+    st.markdown('</div>', unsafe_allow_html=True)
+
+# Footer
+st.markdown('<div class="footer">', unsafe_allow_html=True)
+st.markdown("Built with ❤️ using Streamlit and PyTorch | Deployable to Hugging Face Spaces")
+st.markdown('</div>', unsafe_allow_html=True)
+
+st.markdown('</div>', unsafe_allow_html=True)

requirements.txt

@@ -0,0 +1,5 @@
+torch>=1.7.0
+torchvision>=0.8.0
+streamlit>=1.25.0
+pillow>=8.0.0
+numpy>=1.19.0