Update app.py

app.py

@@ -1,24 +1,24 @@
 import gradio as gr
-import cv2
 import torch
-from model import model
+from PIL import Image
+from model import model 
 from torchvision import transforms
 
-#Load the Model
+# Load your own model
 model.load_state_dict(torch.load('mnist_model.pth'))
-# Set model to evaluation mode
 model.eval()
 
-
 def preprocess_image(image):
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    threshold = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)
-    resized = cv2.resize(threshold, (28, 28), interpolation=cv2.INTER_AREA)
-    tensor = transforms.ToTensor()(resized).unsqueeze(0)
-    tensor = transforms.Normalize((0.5,), (0.5,))(tensor)
+    transform = transforms.Compose([
+        transforms.Resize((28, 28)),
+        transforms.Grayscale(num_output_channels=1),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5,), (0.5,))
+    ])
+    image = Image.fromarray(image)
+    tensor = transform(image).unsqueeze(0)
     return tensor
 
-
 def classify(image):
     tensor = preprocess_image(image)
     with torch.no_grad():
@@ -26,7 +26,6 @@ def classify(image):
         prediction = output.argmax(dim=1, keepdim=True).item()
     return str(prediction)  # Convert prediction to string
 
-
 iface = gr.Interface(
     fn=classify,
     inputs="sketchpad",