Update README.md

README.md

@@ -21,7 +21,7 @@ This repository hosts a fine-tuned **ResNet-18-based** model optimized for **pne
 ### **Installation**  
 
 ```bash
-pip install torch torchvision numpy pillow
+pip install torch torchvision pillow
 ```
 
 ### **Loading the Model**  
@@ -30,14 +30,19 @@ pip install torch torchvision numpy pillow
 import torch
 import torchvision.models as models
 
-# Load model architecture
+# Step 1: Define the model architecture (Must match the trained model)
 model = models.resnet18(pretrained=False)
+model.fc = torch.nn.Linear(in_features=512, out_features=2)  # Ensure output matches 2 classes
 
-# Load weights
-model.load_state_dict(torch.load("resnet18_model.pth", map_location="cpu"))
+# Step 2: Load the fine-tuned model weights
+model_path = "/content/chest_xray_pneumonia_model.pth"  # Ensure the file is in the same directory
+model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
 
-# Set to evaluation mode
+# Step 3: Set model to evaluation mode
 model.eval()
+
+print("✅ Model loaded successfully and ready for inference!")
+
 ```
 
 ---
@@ -45,30 +50,34 @@ model.eval()
 ### **🔍 Perform Pneumonia Detection**  
 
 ```python
-from torchvision import transforms
 from PIL import Image
+import torchvision.transforms as transforms
 
-def predict_pneumonia(model, image_path):
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.5], std=[0.5])
-    ])
-    
-    image = Image.open(image_path).convert("L")  # Convert to grayscale
-    image = transform(image).unsqueeze(0)  # Add batch dimension
-    
-    with torch.no_grad():
-        output = model(image)
-    
-    prediction = torch.argmax(output, dim=1).item()
-    return "Pneumonia" if prediction == 1 else "Normal"
-
-# 🔹 **Test Prediction**
-image_path = "chest_xray_sample.jpg"
-result = predict_pneumonia(model, image_path)
-print(f"Prediction: {result}")
-```
+# Load the image
+image_path = "/content/Screenshot 2025-03-04 104637.png"  # Replace with your test image
+image = Image.open(image_path).convert("RGB")  # Ensure 3-channel format
+
+# Define preprocessing (same as used during training)
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),  # Resize to match model input
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# Apply transformations
+image = transform(image).unsqueeze(0)  # Add batch dimension
+
+# Perform inference
+with torch.no_grad():
+    output = model(image)
+
+# Convert output to class prediction
+predicted_class = torch.argmax(output, dim=1).item()
+
+# Map the predicted class to labels (Modify if needed)
+class_labels = {0: "Normal (Healthy)", 1: "Pneumonia"}
+
+print(f"✅ Predicted Class: {class_labels.get(predicted_class, 'Unknown')}")
 
 ---