Update app.py

app.py

@@ -58,25 +58,25 @@ torch.manual_seed(0)
 #     lr=0.002,
 #     optimizer=torch.optim.Adam,
 #     device=device,
-# )
-model=Cnn()
+import streamlit as st
+from st_canvas import st_canvas
+import torch
+from PIL import Image
+import cv2
+import numpy as np
+from your_model_module import Cnn  # Import your model architecture
+
+# Create an instance of your model (Cnn model)
+model = Cnn()
 
 # Specify the path to the saved model weights
 model_weights_path = 'model_weights.pth'
 
-# Load the model weights
-model.load_state_dict(torch.load(model_weights_path,map_location=torch.device('cpu')))
+# Load the model weights onto a CPU device (if you want to use the CPU)
+model.load_state_dict(torch.load(model_weights_path, map_location=torch.device('cpu')))
 
 # Set the model to evaluation mode for inference
 model.eval()
-# Create a NeuralNetClassifier using the loaded model
-cnn = NeuralNetClassifier(
-    module=model,
-    max_epochs=0,  # Set max_epochs to 0 to avoid additional training
-    lr=0.002,  # You can set this to the learning rate used during training
-    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.)
-)
 
 stroke_width = st.sidebar.slider("Stroke width: ", 1, 35, 32)
 stroke_color = st.sidebar.color_picker("Stroke color hex: ")
@@ -103,16 +103,23 @@ canvas_result = st_canvas(
 
 # Do something interesting with the image data and paths
 if canvas_result.image_data is not None:
-    #st.image(canvas_result.image_data)
     image = canvas_result.image_data
     image1 = image.copy()
     image1 = image1.astype('uint8')
-    image1 = cv2.cvtColor(image1,cv2.COLOR_BGR2GRAY)
-    image1 = cv2.resize(image1,(28,28))
+    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"]))
+    # Convert the image for prediction (assuming image1 is in the right format)
+    image1 = image1[np.newaxis, np.newaxis, ...]  # Add batch and channel dimensions
+
+    # Perform prediction using the pre-trained model
+    with torch.no_grad():
+        tensor_image = torch.tensor(image1, dtype=torch.float32)
+        prediction = model(tensor_image)
+
+    # Display the predicted class
+    predicted_class = prediction.argmax().item()
+    st.title(f"Predicted Class: {predicted_class}")
+