Create 15_CNN.py

pages/15_CNN.py

@@ -0,0 +1,110 @@
+import streamlit as st
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision.transforms as transforms
+import torchvision
+from torch.utils.data import DataLoader
+from PIL import Image
+import numpy as np
+
+# Define the CNN model
+class Net(nn.Module):
+    def __init__(self):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
+        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
+        self.fc1 = nn.Linear(128 * 4 * 4, 256)
+        self.fc2 = nn.Linear(256, 128)
+        self.fc3 = nn.Linear(128, 10)
+        
+    def forward(self, x):
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = self.pool(F.relu(self.conv3(x)))
+        x = x.view(-1, 128 * 4 * 4)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+# Load pre-trained model (if available)
+def load_model():
+    net = Net()
+    try:
+        net.load_state_dict(torch.load('cnn_model.pth'))
+        net.eval()
+        st.write("Model loaded successfully")
+    except FileNotFoundError:
+        st.write("No pre-trained model found. Please train the model first.")
+    return net
+
+# Function to predict the class of an image
+def predict(image, model):
+    transform = transforms.Compose([
+        transforms.Resize((32, 32)),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+    ])
+    image = transform(image).unsqueeze(0)
+    outputs = model(image)
+    _, predicted = torch.max(outputs.data, 1)
+    return predicted.item()
+
+# Training function
+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 = DataLoader(trainset, batch_size=100, shuffle=True, num_workers=2)
+
+    net = Net()
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
+
+    st.write("Training the model...")
+    for epoch in range(10):
+        running_loss = 0.0
+        for i, data in enumerate(trainloader, 0):
+            inputs, labels = data
+            optimizer.zero_grad()
+            outputs = net(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            running_loss += loss.item()
+            if i % 100 == 99:
+                st.write(f'[Epoch {epoch + 1}, Batch {i + 1}] loss: {running_loss / 100:.3f}')
+                running_loss = 0.0
+
+    st.write('Finished Training')
+    torch.save(net.state_dict(), 'cnn_model.pth')
+    st.write('Model saved as cnn_model.pth')
+    return net
+
+# Streamlit interface
+st.title("CNN Image Classification with CIFAR-10")
+
+mode = st.sidebar.selectbox("Mode", ["Train", "Predict"])
+
+if mode == "Train":
+    if st.button("Train Model"):
+        model = train_model()
+
+if mode == "Predict":
+    model = load_model()
+    uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+        st.image(image, caption='Uploaded Image.', use_column_width=True)
+        st.write("")
+        st.write("Classifying...")
+        class_idx = predict(image, model)
+        classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
+        st.write(f'Prediction: {classes[class_idx]}')