Update app.py

app.py

@@ -21,17 +21,27 @@ transform = transforms.Compose([
 # Define the class labels
 class_labels = ['Normal', 'Pneumonia']
 
+# Create a function to preprocess the image
+def preprocess_image(image):
+    # Resize the image to match the model's input shape
+    image = image.resize((224, 224))
+
+    # Convert the image to a tensor
+    image_tensor = transform(image)
+
+    # Add a batch dimension
+    image_tensor = image_tensor.unsqueeze(0)
+
+    return image_tensor
+
 # Create a function to make predictions
 def predict(image):
-    # Convert the image to RGB
-    image = image.convert('RGB')
-    
     # Preprocess the image
-    image = transform(image).unsqueeze(0)
-    
+    preprocessed_image = preprocess_image(image)
+
     # Make the prediction
     with torch.no_grad():
-        output = model(image)
+        output = model(preprocessed_image)
         _, predicted_idx = torch.max(output, 1)
         predicted_label = class_labels[predicted_idx.item()]