Update predict.py

predict.py

@@ -1,12 +1,3 @@
-# -*- coding: utf-8 -*-
-"""Untitled5.ipynb
-
-Automatically generated by Colab.
-
-Original file is located at
-    https://colab.research.google.com/drive/1kfJMKD93CT0RxoHkh_T3hdcVMovTMHCe
-"""
-
 import os
 import torch
 from transformers import ViTForImageClassification, ViTFeatureExtractor
@@ -18,18 +9,25 @@ def load_model(model_path):
     """Load the pre-trained model and feature extractor."""
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
-    model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224', num_labels=13)
-    model.load_state_dict(torch.load(model_path))
+    
+    # Load the model with weights mapped to the appropriate device
+    model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224', num_labels=13, ignore_mismatched_sizes=True)
+    
+    # Load the state dict with map_location to handle devices properly
+    state_dict = torch.load(model_path, map_location=device)
+    model.load_state_dict(state_dict, strict=False)  # Use strict=False to ignore size mismatches
+
     model = model.to(device)
     model.eval()  # Set the model to evaluation mode
     return model, feature_extractor, device
 
 def safe_load_image(path):
-    """Safely load an image, handling possible errors."""
+    """Safely load an image, handling potential errors."""
     try:
         with open(path, 'rb') as f:
             img = Image.open(io.BytesIO(f.read()))
             img = img.convert('RGB')
+            img = img.resize((224, 224))  # Resize the image to (224, 224)
         return img
     except Exception as e:
         print(f"Error loading image {path}: {e}")
@@ -44,12 +42,12 @@ def predict_image_class(image_path, model, feature_extractor, device, class_name
     # Preprocess the image
     inputs = feature_extractor(images=img, return_tensors="pt").to(device)
 
-    # Make the prediction
+    # Perform prediction
     with torch.no_grad():
         outputs = model(**inputs).logits
         probabilities = torch.softmax(outputs, dim=1).cpu().numpy()[0]  # Calculate probabilities
         predicted_class_idx = outputs.argmax(dim=1).item()  # Get the predicted class index
-        predicted_class = class_names[predicted_class_idx]  # Get class name from index
+        predicted_class = class_names[predicted_class_idx]  # Get the class name based on the index
 
     return predicted_class, probabilities
 
@@ -65,31 +63,33 @@ def predict_images_in_folder(folder_path, model, feature_extractor, device, clas
 
     return results
 
-def save_results_to_excel(results, output_file):
-    """Save the prediction results to an Excel file."""
-    # Flatten the probabilities array and create a DataFrame
+def save_results_to_excel(results, output_file, class_names):
+    """Save prediction results to an Excel file."""
+    # Flatten probability array and create DataFrame
     rows = []
     for result in results:
+        # Add each probability with corresponding class name
         for idx, prob in enumerate(result['Probabilities']):
             rows.append({
                 'Image Name': result['Image Name'],
                 'Predicted Class': result['Predicted Class'],
-                'Class Index': idx,
+                'Class': class_names[idx],
                 'Probability': prob
             })
-
+    
     df = pd.DataFrame(rows)
-
+    
     # Sort by probability in descending order
     df = df.sort_values(by='Probability', ascending=False)
-
+    
     # Save to Excel
     df.to_excel(output_file, index=False)
+    print(f'Results saved to {output_file}')  # Confirm saving
 
 def main(input_path, model_path, output_file):
-    """Main function to perform image classification prediction, handling single images or folders, and saving results to Excel."""
-    class_names = ['anti_war_protest', 'combat', 'construction', 'fire', 'human_damage',
-                   'humanitarian_aid', 'infrastructure', 'military_parade', 'military_vehicle',
+    """Main function to execute image classification predictions, processing single images or folders, and saving results to Excel."""
+    class_names = ['anti_war_protest', 'combat', 'construction', 'fire', 'human_damage', 
+                   'humanitarian_aid', 'infrastructure', 'military_parade', 'military_vehicle', 
                    'meeting', 'speech', 'refugee', 'victory']
 
     model, feature_extractor, device = load_model(model_path)
@@ -97,17 +97,22 @@ def main(input_path, model_path, output_file):
     if os.path.isdir(input_path):
         # If the input path is a folder, predict all images in that folder
         results = predict_images_in_folder(input_path, model, feature_extractor, device, class_names)
-        save_results_to_excel(results, output_file)
-        print(f'Prediction results saved to: {output_file}')
+        if results:
+            save_results_to_excel(results, output_file, class_names)
+        else:
+            print("No valid images found in the specified folder.")
     elif os.path.isfile(input_path):
-        # If the input path is a single image, make a direct prediction
+        # If the input path is a single image, perform direct prediction
         predicted_class, probabilities = predict_image_class(input_path, model, feature_extractor, device, class_names)
-        print(f'The predicted class for image {os.path.basename(input_path)}: {predicted_class}')
+        if predicted_class is not None:
+            print(f'Predicted class for image {os.path.basename(input_path)}: {predicted_class}')
+        else:
+            print("Image could not be processed.")
     else:
         print('Invalid input path. Please provide a valid file or folder path.')
 
 # Example call
-input_path = '/path/to/your/image_or_folder'  # Replace with your image or folder path
-model_path = '/kaggle/working/best_modelq.pth'  # Replace with your model path
-output_file = 'predictions.xlsx'  # Name of the output Excel file
+input_path = '/content/ddd.jpg'  # Replace with your image or folder path
+model_path = '/content/model.pth'  # Replace with your model path
+output_file = '/content/predictions.xlsx'  # Name of the output Excel file
 main(input_path, model_path, output_file)