Update app.py

app.py

@@ -49,22 +49,28 @@ model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
 
 
 def predict(image):
-    # Convert NumPy array to PIL Image
-    image = Image.fromarray(np.uint8(image)).convert('RGB')
-
     preprocess = transforms.Compose([
         transforms.Resize((128, 128)),
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
     ])
 
-    classes = ['cat', 'dog']
+    # Ensure the image is a PIL Image
+    image = Image.fromarray(image.astype('uint8'), 'RGB')
 
     x = preprocess(image).unsqueeze(0)
-    x = model(x)
-    probabilities = torch.nn.functional.softmax(x, dim=1)
-    
-    return classes[probabilities.argmax(dim=1).item()]
 
+    # Set model to evaluation mode
+    model.eval()
+
+    with torch.no_grad():  # Use no_grad context for inference to save memory and computations
+        x = model(x)
+        probabilities = torch.nn.functional.softmax(x, dim=1)
+        class_id = probabilities.argmax(dim=1).item()
+
+    classes = ['cat', 'dog']
+    return classes[class_id]
+
+# Update Gradio interface
 demo = gr.Interface(fn=predict, inputs="image", outputs="text")
 demo.launch()