Update app.py

app.py

@@ -1,138 +1,156 @@
-import streamlit as st
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import torchvision
-import torchvision.transforms as transforms
-from PIL import Image
-import io
-
-# Set page config
-st.set_page_config(page_title="CIFAR-10 Classifier", layout="centered", initial_sidebar_state="collapsed")
-
-# Custom CSS for dark theme
-st.markdown("""
-<style>
-    .stApp {
-        background-color: #0E1117;
-        color: #FAFAFA;
-    }
-    .stButton>button {
-        background-color: #4CAF50;
-        color: white;
-    }
-    .stHeader {
-        background-color: #262730;
-        color: white;
-        padding: 1rem;
-        border-radius: 5px;
-    }
-    .stImage {
-        background-color: #262730;
-        padding: 10px;
-        border-radius: 5px;
-    }
-    .stSuccess {
-        background-color: #262730;
-        color: #4CAF50;
-        padding: 10px;
-        border-radius: 5px;
-    }
-</style>
-""", unsafe_allow_html=True)
-
-# Model definition
-class SimpleCNN(nn.Module):
-    def __init__(self):
-        super(SimpleCNN, self).__init__()
-        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
-        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
-        self.pool = nn.MaxPool2d(2, 2)
-        self.fc1 = nn.Linear(64 * 8 * 8, 512)
-        self.fc2 = nn.Linear(512, 10)
-
-    def forward(self, x):
-        x = self.pool(torch.relu(self.conv1(x)))
-        x = self.pool(torch.relu(self.conv2(x)))
-        x = x.view(-1, 64 * 8 * 8)
-        x = torch.relu(self.fc1(x))
-        x = self.fc2(x)
-        return x
-
-# Function to train the model
-@st.cache_resource
-def train_model():
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
-    ])
-
-    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
-    trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
-
-    model = SimpleCNN()
-    criterion = nn.CrossEntropyLoss()
-    optimizer = optim.Adam(model.parameters(), lr=0.001)
-
-    for epoch in range(5):  # Train for 5 epochs
-        for i, data in enumerate(trainloader, 0):
-            inputs, labels = data
-            optimizer.zero_grad()
-            outputs = model(inputs)
-            loss = criterion(outputs, labels)
-            loss.backward()
-            optimizer.step()
-
-    return model
-
-# Function to load or train the model
-@st.cache_resource
-def get_model():
-    try:
-        model = SimpleCNN()
-        model.load_state_dict(torch.load('cifar10_model.pth'))
-        model.eval()
-    except:
-        model = train_model()
-        torch.save(model.state_dict(), 'cifar10_model.pth')
-    return model
-
-# Streamlit app
-st.markdown("<h1 class='stHeader'>CIFAR-10 Image Classification</h1>", unsafe_allow_html=True)
-st.write("Upload an image to classify it into one of the CIFAR-10 categories.")
-
-# File uploader
-uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-if uploaded_file is not None:
-    # Display uploaded image
-    image = Image.open(uploaded_file)
-    st.markdown("<div class='stImage'>", unsafe_allow_html=True)
-    st.image(image, caption='Uploaded Image', use_column_width=True)
-    st.markdown("</div>", unsafe_allow_html=True)
-
-    # Predict button
-    if st.button('Classify Image'):
-        # Load model
-        model = get_model()
-
-        # Preprocess image
-        transform = transforms.Compose([
-            transforms.Resize((32, 32)),
-            transforms.ToTensor(),
-            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
-        ])
-        input_tensor = transform(image).unsqueeze(0)
-
-        # Make prediction
-        with torch.no_grad():
-            output = model(input_tensor)
-            _, predicted = torch.max(output, 1)
-
-        # Display result
-        classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
-        st.markdown(f"<div class='stSuccess'>Prediction: {classes[predicted.item()]}</div>", unsafe_allow_html=True)
-
-# Footer
-st.markdown("---")
+import streamlit as st
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision
+import torchvision.transforms as transforms
+from PIL import Image
+import io
+
+# Set page config
+st.set_page_config(page_title="CIFAR-10 Classifier", layout="wide", initial_sidebar_state="expanded")
+
+# Custom CSS for dark theme
+st.markdown("""
+<style>
+    .stApp {
+        background-color: #0E1117;
+        color: #FAFAFA;
+    }
+    .stButton>button {
+        background-color: #4CAF50;
+        color: white;
+    }
+    .stHeader {
+        background-color: #262730;
+        color: white;
+        padding: 1rem;
+        border-radius: 5px;
+        margin-bottom: 1rem;
+    }
+    .stImage {
+        background-color: #262730;
+        padding: 10px;
+        border-radius: 5px;
+    }
+    .stSuccess {
+        background-color: #262730;
+        color: #4CAF50;
+        padding: 10px;
+        border-radius: 5px;
+        margin-top: 1rem;
+    }
+    .upload-box {
+        border: 2px dashed #4CAF50;
+        border-radius: 5px;
+        padding: 20px;
+        text-align: center;
+        cursor: pointer;
+    }
+</style>
+""", unsafe_allow_html=True)
+
+# Model definition
+class SimpleCNN(nn.Module):
+    def __init__(self):
+        super(SimpleCNN, self).__init__()
+        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
+        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.fc1 = nn.Linear(64 * 8 * 8, 512)
+        self.fc2 = nn.Linear(512, 10)
+
+    def forward(self, x):
+        x = self.pool(torch.relu(self.conv1(x)))
+        x = self.pool(torch.relu(self.conv2(x)))
+        x = x.view(-1, 64 * 8 * 8)
+        x = torch.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+# Function to train the model
+@st.cache_resource
+def train_model():
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+    ])
+
+    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
+    trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
+
+    model = SimpleCNN()
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+
+    for epoch in range(5):  # Train for 5 epochs
+        for i, data in enumerate(trainloader, 0):
+            inputs, labels = data
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+    return model
+
+# Function to load or train the model
+@st.cache_resource
+def get_model():
+    try:
+        model = SimpleCNN()
+        model.load_state_dict(torch.load('cifar10_model.pth'))
+        model.eval()
+    except:
+        model = train_model()
+        torch.save(model.state_dict(), 'cifar10_model.pth')
+    return model
+
+# Sidebar
+st.sidebar.title("Upload Image")
+uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+
+# Main content
+st.markdown("<h1 class='stHeader'>CIFAR-10 Image Classification</h1>", unsafe_allow_html=True)
+
+# Drag and drop section
+col1, col2, col3 = st.columns([1,2,1])
+# with col2:
+#     st.markdown("<div class='upload-box'>Drag and drop image here</div>", unsafe_allow_html=True)
+
+# Display uploaded image and make prediction
+if uploaded_file is not None:
+    image = Image.open(uploaded_file)
+    col1, col2, col3 = st.columns([1,2,1])
+    with col2:
+        st.markdown("<div class='stImage'>", unsafe_allow_html=True)
+        st.image(image, caption='Uploaded Image', use_column_width=True)
+        st.markdown("</div>", unsafe_allow_html=True)
+
+    # Load model and make prediction
+    model = get_model()
+    transform = transforms.Compose([
+        transforms.Resize((32, 32)),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+    ])
+    input_tensor = transform(image).unsqueeze(0)
+
+    with torch.no_grad():
+        output = model(input_tensor)
+        probabilities = torch.nn.functional.softmax(output[0], dim=0)
+
+    # Display results in sidebar
+    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
+    _, predicted = torch.max(output, 1)
+    st.sidebar.markdown("<div class='stSuccess'>", unsafe_allow_html=True)
+    st.sidebar.write(f"Best Prediction: {classes[predicted.item()]}")
+    st.sidebar.markdown("</div>", unsafe_allow_html=True)
+    
+    st.sidebar.write("Prediction Probabilities:")
+    for i, prob in enumerate(probabilities):
+        st.sidebar.write(f"{classes[i]}: {prob.item():.2%}")
+
+# Footer
+st.markdown("---")
 st.markdown("<p style='text-align: center; color: #666;'>Created with Streamlit and PyTorch</p>", unsafe_allow_html=True)