Update app.py

app.py

@@ -1,3 +1,4 @@
+# importing all the packages
 from sklearn.datasets import fetch_openml
 from sklearn.model_selection import train_test_split
 import numpy as np
@@ -11,59 +12,34 @@ from PIL import Image
 import streamlit as st
 from streamlit_drawable_canvas import st_canvas
 import cv2
+from PIL import Image
+import torchvision.transforms as transforms
 import torch
 from skorch import NeuralNetClassifier
 from torch import nn
 import torch.nn.functional as F
 import matplotlib.pyplot as plt
+from model import Cnn
+
+
+
+
+
+# Reading the data
 mnist = fetch_openml('mnist_784', as_frame=False, cache=False)
 X = mnist.data.astype('float32')
 y = mnist.target.astype('int64')
 X /= 255.0
 
-
 #device = 'cuda' if torch.cuda.is_available() else 'cpu'
-XCnn = X.reshape(-1, 1, 28, 28)
-XCnn_train, XCnn_test, y_train, y_test = train_test_split(XCnn, y, test_size=0.25, random_state=42)
+XCnn = X.reshape(-1, 1, 28, 28) #reshape input 
+XCnn_train, XCnn_test, y_train, y_test = train_test_split(XCnn, y, test_size=0.25, random_state=42) #train test split
 
-from PIL import Image
-import torchvision.transforms as transforms
-class Cnn(nn.Module):
-    def __init__(self, dropout=0.5):
-        super(Cnn, self).__init__()
-        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
-        self.conv2_drop = nn.Dropout2d(p=dropout)
-        self.fc1 = nn.Linear(1600, 100) # 1600 = number channels * width * height
-        self.fc2 = nn.Linear(100, 10)
-        self.fc1_drop = nn.Dropout(p=dropout)
-
-    def forward(self, x):
-        x = torch.relu(F.max_pool2d(self.conv1(x), 2))
-        x = torch.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
-
-        # flatten over channel, height and width = 1600
-        x = x.view(-1, x.size(1) * x.size(2) * x.size(3))
-
-        x = torch.relu(self.fc1_drop(self.fc1(x)))
-        x = torch.softmax(self.fc2(x), dim=-1)
-        return x
 torch.manual_seed(0)
 
 
-# # Create an instance of your model
-# model = NeuralNetClassifier(
-#     Cnn,
-#     max_epochs=10,
-#     lr=0.002,
-#     optimizer=torch.optim.Adam,
-#     device=device,
-import streamlit as st
-import torch
-from PIL import Image
-import cv2
-import numpy as np
 
+# reshape and train test split
 
 #device = 'cuda' if torch.cuda.is_available() else 'cpu'
 XCnn = X.reshape(-1, 1, 28, 28)
@@ -71,37 +47,10 @@ XCnn_train, XCnn_test, y_train, y_test = train_test_split(XCnn, y, test_size=0.2
 
 from PIL import Image
 import torchvision.transforms as transforms
-class Cnn(nn.Module):
-    def __init__(self, dropout=0.5):
-        super(Cnn, self).__init__()
-        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
-        self.conv2_drop = nn.Dropout2d(p=dropout)
-        self.fc1 = nn.Linear(1600, 100) # 1600 = number channels * width * height
-        self.fc2 = nn.Linear(100, 10)
-        self.fc1_drop = nn.Dropout(p=dropout)
-
-    def forward(self, x):
-        x = torch.relu(F.max_pool2d(self.conv1(x), 2))
-        x = torch.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
-
-        # flatten over channel, height and width = 1600
-        x = x.view(-1, x.size(1) * x.size(2) * x.size(3))
-
-        x = torch.relu(self.fc1_drop(self.fc1(x)))
-        x = torch.softmax(self.fc2(x), dim=-1)
-        return x
+
 torch.manual_seed(0)
 
 
-# # Create an instance of your model
-# model = NeuralNetClassifier(
-#     Cnn,
-#     max_epochs=10,
-#     lr=0.002,
-#     optimizer=torch.optim.Adam,
-#     device=device,
-# )
 model=Cnn()
 
 # Specify the path to the saved model weights
@@ -112,6 +61,7 @@ model.load_state_dict(torch.load(model_weights_path,map_location=torch.device('c
 
 # Set the model to evaluation mode for inference
 model.eval()
+
 # Create a NeuralNetClassifier using the loaded model
 cnn = NeuralNetClassifier(
     module=model,
@@ -120,7 +70,9 @@ cnn = NeuralNetClassifier(
     optimizer=torch.optim.Adam,  # You can set the optimizer used during training
     device='cpu'  # You can specify the device ('cpu' for CPU, 'cuda' for GPU, etc.)
 )
+
 cnn.fit(XCnn_train, y_train)
+
 stroke_width = st.sidebar.slider("Stroke width: ", 1, 35, 32)
 stroke_color = st.sidebar.color_picker("Stroke color hex: ")
 bg_color = st.sidebar.color_picker("Background color hex: ", "#eee")
@@ -154,9 +106,9 @@ if canvas_result.image_data is not None:
     image1 = cv2.cvtColor(image1,cv2.COLOR_BGR2GRAY)
     image1 = cv2.resize(image1,(28,28))
     st.image(image1)
-
-
     image1.resize(1,1,28,28)
     st.title(np.argmax(cnn.predict(image1)))
+
+    
 if canvas_result.json_data is not None:
     st.dataframe(pd.json_normalize(canvas_result.json_data["objects"]))