Upload 4 files

app.py

@@ -0,0 +1,18 @@
+from flask import Flask, jsonify, request
+from torch_utils import transform_image
+from torch_utils import get_prediction
+
+
+app = Flask(__name__)
+
+@app.route('/classify', methods=['POST'])
+def classify():
+    if request.method == 'POST':
+        file = request.files['file']
+        img_bytes = file.read()
+        img = transform_image(img_bytes)
+        pred = get_prediction(img)
+        return jsonify({'classification': int(pred[0])})
+    
+if __name__ == '__main__':
+    app.run(port=5000, debug=True)

st.py

@@ -0,0 +1,37 @@
+import streamlit as st
+import requests
+from PIL import Image
+from io import BytesIO
+
+CLASS_LABELS = {
+        0: "airplane",
+        1: "bird", 
+        2: "car", 
+        3: "cat", 
+        4: "deer", 
+        5: "dog", 
+        6: "horse", 
+        7: "monkey", 
+        8: "ship", 
+        9: "truck", 
+}
+
+def get_classification(image_bytes):
+    response = requests.post("http://localhost:5000/classify", files={"file": image_bytes})
+    class_id = response.json()["classification"]
+    return CLASS_LABELS[class_id]
+
+st.title("Image Classification")
+st.write("Upload an image to classify")
+
+uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"])
+
+if uploaded_file is not None:
+    image = Image.open(uploaded_file)
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    if st.button("Classify"):
+        img_bytes = uploaded_file.read()
+        label = get_classification(img_bytes)
+        st.write("Prediction:", label)
+

torch_utils.py

@@ -0,0 +1,53 @@
+import io
+import torch 
+import torch.nn as nn 
+import torch.nn.functional as F
+import torchvision.transforms as transforms 
+from PIL import Image
+
+
+class CNN(nn.Module):
+    def __init__(self):
+        super(CNN, self).__init__()
+        self.conv1 = nn.Conv2d(3, 32, 5)
+        self.conv2 = nn.Conv2d(32, 64, 5)
+        #full layer
+        self.fc1 = nn.Linear(64 * 13 * 13, 128)
+        self.fc2 = nn.Linear(128, 64)
+        self.fc3 = nn.Linear(64, 10)
+
+    def forward(self, x):
+        x = F.max_pool2d(F.relu(self.conv1(x)), (2,2))
+        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
+        x = x.view(-1, self.num_flat_features(x))
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+    
+    def num_flat_features(self, x):
+        size = x.size()[1:]  # all dimensions except the batch dimension
+        num_features = 1
+        for s in size:
+            num_features *= s
+        return num_features
+
+loaded_model = CNN()
+loaded_model.load_state_dict(torch.load("cnn_model.pth")) # it takes the loaded dictionary, not the path file itself
+loaded_model.eval()
+
+
+#transform images
+
+def transform_image(image_bytes):
+    transform = transforms.Compose(
+        [transforms.Resize(64), transforms.CenterCrop(64), transforms.ToTensor()]
+    )
+    image = Image.open(io.BytesIO(image_bytes))
+    return transform(image).unsqueeze(0)
+
+def get_prediction(image_tensor):
+    outputs = loaded_model(image_tensor)
+        # max returns (value ,index)
+    _, predicted = torch.max(outputs.data, 1)
+    return predicted

unv_model.py

@@ -0,0 +1,98 @@
+import torch  
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchvision
+import torchvision.transforms as transforms
+import scipy
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")# assign the device to the model
+
+#hyper-parameters
+learning_rate = 0.005
+batch_size = 128
+hidden_size = 300
+num_classes = 10
+num_epochs = 550
+
+#load data 
+
+transform = transforms.Compose(
+    [transforms.Resize(64), transforms.CenterCrop(64), transforms.ToTensor()]
+)
+
+train_dataset = torchvision.datasets.STL10(root='./dataSTL10', split="train", transform=transform, download=True)
+test_dataset = torchvision.datasets.STL10(root='./dataSTL10', split="test", transform=transform, download=True)
+train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
+test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
+
+# CNN 
+
+class CNN(nn.Module):
+    def __init__(self):
+        super(CNN, self).__init__()
+        self.conv1 = nn.Conv2d(3, 32, 5)
+        self.conv2 = nn.Conv2d(32, 64, 5)
+        #full layer
+        self.fc1 = nn.Linear(64 * 13 * 13, 128)
+        self.fc2 = nn.Linear(128, 64)
+        self.fc3 = nn.Linear(64, num_classes)
+
+    def forward(self, x):
+        x = F.max_pool2d(F.relu(self.conv1(x)), (2,2))
+        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
+        x = x.view(-1, self.num_flat_features(x))
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+    
+    def num_flat_features(self, x):
+        size = x.size()[1:]  # all dimensions except the batch dimension
+        num_features = 1
+        for s in size:
+            num_features *= s
+        return num_features
+    
+cnn = CNN().to(device)
+
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(cnn.parameters(), lr=learning_rate)
+
+# training loop
+
+for epoch in range(num_epochs):
+    for i, (images, labels) in enumerate(train_loader):
+        images = images.to(device)
+        labels = labels.to(device)
+
+        out = cnn(images)
+        loss = criterion(out, labels)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+    if(i+1) % 1 == 0:
+        print(f'epoch: {epoch+1}/{num_epochs} step: {i+1}, loss: loss: {loss.item():.4f}')
+
+
+with torch.no_grad():
+    n_correct = 0
+    n_samples = 0
+    for images, labels in test_loader:
+        images = images.to(device)
+        labels = labels.to(device)
+        outputs = cnn(images)
+        # max returns (value ,index)
+        _, predicted = torch.max(outputs.data, 1)
+        n_samples += labels.size(0)
+        n_correct += (predicted == labels).sum().item()
+
+    acc = 100.0 * n_correct / n_samples
+    print(f'Accuracy of the network on the {n_samples} test images: {acc} %')
+
+
+# Save the model
+torch.save(cnn.state_dict(), "cnn_model.pth")
+