Update README.md

README.md

@@ -11,14 +11,6 @@ This repository hosts a **fine-tuned ResNet-50 model** for **Vehicle Type Classi
 - **Fine-tuning Framework:** PyTorch (`torchvision.models.resnet50`)
 - **Optimization:** Trained with Adam optimizer and data augmentation for robust performance
 
-## Downloading the Model
-
-You can download the fine-tuned model from the provided link:
-
-```sh
-wget <download_link>/fine_tuned_model.zip
-unzip fine_tuned_model.zip
-```
 
 ## Usage
 
@@ -38,37 +30,54 @@ import torchvision.models as models
 import torchvision.transforms as transforms
 from PIL import Image
 
-# Load the model architecture
-model = models.resnet50(pretrained=False)
-num_ftrs = model.fc.in_features
-model.fc = torch.nn.Linear(num_ftrs, 11)  # 11 vehicle classes
+# Define the model architecture
+resnet50 = models.resnet50(pretrained=False)
+
+# Modify the last layer to match the number of classes (11)
+num_ftrs = resnet50.fc.in_features
+resnet50.fc = torch.nn.Linear(num_ftrs, 11)
+
+# Load trained model weights
+resnet50.load_state_dict(torch.load("fine_tuned_model/pytorch_model.bin"))
+resnet50.eval()  # Set model to evaluation mode
+
+print("Model loaded successfully!")
 
-# Load fine-tuned weights
-model.load_state_dict(torch.load("fine_tuned_model/pytorch_model.bin", map_location=torch.device('cpu')))
-model.eval()  # Set to evaluation mode
 
-# Load class labels
+# Load class names
 with open("fine_tuned_model/classes.txt", "r") as f:
     class_names = f.read().splitlines()
 
-# Define preprocessing transformations
+print("Classes:", class_names)
+
+
+# Define image transformations (same as training)
 transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    transforms.Resize((224, 224)),  # Resize to match ResNet-50 input size
+    transforms.ToTensor(),  
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # Normalization
 ])
 
-# Load and preprocess a test image
-image_path = "path_to_your_image.jpg"  # Change this to your test image path
-image = Image.open(image_path).convert("RGB")
-input_tensor = transform(image).unsqueeze(0)
+# Load the custom image
+image_path = "/kaggle/input/sample-image-1/pickup_truck_sample_image.jpg"  # Change this to your test image path
+image = Image.open(image_path).convert("RGB")  # Open image and convert to RGB
+input_tensor = transform(image).unsqueeze(0)  # Add batch dimension
+
 
-# Make prediction
+
+# Move to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+resnet50 = resnet50.to(device)
+input_tensor = input_tensor.to(device)
+
+# Get predictions
 with torch.no_grad():
-    outputs = model(input_tensor)
-    _, predicted_class = torch.max(outputs, 1)
+    outputs = resnet50(input_tensor)
+    _, predicted_class = torch.max(outputs, 1)  # Get the class with highest score
 
+# Print the result
 print(f"Predicted Vehicle Type: {class_names[predicted_class.item()]}")
+
 ```
 
 ## Performance Metrics