Update README.md

README.md

@@ -46,23 +46,39 @@ from transformers import AutoModel
 # Load model
 model = AutoModel.from_pretrained("your-huggingface-username/model-name")
 model.eval()
+```
 
-# Define image preprocessing
+### Transformer Application
+```python
 transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    transforms.Resize((224, 224)),  # Resize image to match input size
+    transforms.ToTensor(),  # Convert image to tensor
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize using ImageNet stats
 ])
+```
+
+### Function to preprocess image and get predictions
+```python
+def predict(image_path):
+    # Load and preprocess the input image
+    image = Image.open(image_path).convert('RGB')  # Ensure RGB format
+    input_tensor = transform(image).unsqueeze(0).to(device)  # Add batch dimension
+
+    # Perform inference
+    with torch.no_grad():
+        outputs = model(input_tensor)  # Forward pass
+        probabilities = torch.nn.functional.softmax(outputs, dim=1)  # Get class probabilities
+    
+    return probabilities.cpu().numpy()[0]  # Return probabilities as a NumPy array
 
-# Load and preprocess image
-image = Image.open("path/to/image.jpg")
-image = transform(image).unsqueeze(0)
+# Test with an example image
+image_path = "your_image_path"  # Replace with your test image path
+class_probs = predict(image_path)
 
-# Perform inference
-with torch.no_grad():
-    output = model(image)
-    prediction = torch.argmax(output, dim=1).item()
-    print(f"Predicted DR severity: {prediction}")
+# Print results
+print(f"Class probabilities: {class_probs}")
+predicted_class = np.argmax(class_probs)  # Get the class with highest probability
+print(f"Predicted class: {predicted_class}")
 ```
 
 ## License