Upload app.py

app.py

@@ -0,0 +1,138 @@
+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("---")
+st.markdown("<p style='text-align: center; color: #666;'>Created with Streamlit and PyTorch</p>", unsafe_allow_html=True)