Update app.py

app.py

@@ -47,21 +47,23 @@ model = CustomModel(input_shape=(3,128,128), num_classes=2)
 model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
 
 
-def greet(image):
+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])
+        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']
 
     x = preprocess(image).unsqueeze(0)
     x = model(x)
-    output = torch.nn.functional.softmax(x, dim=1)
-    
+    probabilities = torch.nn.functional.softmax(x, dim=1)
     
     return classes[probabilities.argmax(dim=1).item()]
 
-demo = gr.Interface(fn=greet, inputs="image", outputs="text")
+demo = gr.Interface(fn=predict, inputs="image", outputs="text")
 demo.launch()