Update app.py

app.py

@@ -5,6 +5,7 @@ from torchvision import transforms
 import torchvision.models as models
 import torch.nn as nn
 
+
 class DummyModel(nn.Module):
     def __init__(self):
         super(DummyModel, self).__init__()
@@ -80,21 +81,32 @@ class DummyModel(nn.Module):
 
         return dec4
 
+
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = DummyModel()
-model.load_state_dict(torch.load('model.pth'))
+model.load_state_dict(torch.load('model.pth', weights_only=True))
 model = model.to(device)
 model.eval()
 
+
 # Define preprocessing transforms
-preprocess = transforms.Compose([
-    transforms.Resize(150),
+transform = transforms.Compose([
+    transforms.Resize((150, 150)),
     transforms.ToTensor(),
-    transforms.Normalize([0.5), [0.5])
+    transforms.Normalize([0.5], [0.5])
 ])
 
+
+def preprocess(image):
+    image = image.convert('L')
+    image = transform(image)
+    image = image.unsqueeze(0)
+
+    return image
+
+
 def predict(image):
-    image = preprocess(image).to(model.device)
+    image = preprocess(image).to(device)
     with torch.no_grad():
         output = model(image)
 
@@ -102,7 +114,7 @@ def predict(image):
 
     return image
 
-# Create Gradio interface
+
 iface = gr.Interface(fn=predict,
                      inputs=gr.Image(type="pil"),
                      outputs=gr.Image(type="pil"))