Update predict.py

predict.py

@@ -1,118 +1,61 @@
 import os
 import torch
+import pandas as pd
+from torchvision import transforms
 from transformers import ViTForImageClassification, ViTFeatureExtractor
 from PIL import Image
-import io
-import pandas as pd
 
 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')
-    
-    # Load the model with weights mapped to the appropriate device
+    """Load the pre-trained model."""
     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
+    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)
+    model.eval()
+    return model
 
-def safe_load_image(path):
-    """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}")
-        return None
-
-def predict_image_class(image_path, model, feature_extractor, device, class_names):
-    """Predict the class of a given image."""
-    img = safe_load_image(image_path)
-    if img is None:
-        return None, None
-
-    # Preprocess the image
-    inputs = feature_extractor(images=img, return_tensors="pt").to(device)
+def preprocess_image(image_path):
+    """Preprocess the image for prediction."""
+    feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
+    image = Image.open(image_path).convert("RGB")
+    image = feature_extractor(images=image, return_tensors="pt")["pixel_values"]
+    return image
 
-    # Perform prediction
+def predict(model, image_path):
+    """Predict the class probabilities for an image."""
+    image = preprocess_image(image_path)
     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 the class name based on the index
-
-    return predicted_class, probabilities
-
-def predict_images_in_folder(folder_path, model, feature_extractor, device, class_names):
-    """Predict the class of each image in a folder."""
-    results = []
-    for filename in os.listdir(folder_path):
-        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
-            image_path = os.path.join(folder_path, filename)
-            predicted_class, probabilities = predict_image_class(image_path, model, feature_extractor, device, class_names)
-            if predicted_class is not None:
-                results.append({'Image Name': filename, 'Predicted Class': predicted_class, 'Probabilities': probabilities})
-
-    return results
-
-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': 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
+        outputs = model(image).logits
+    probabilities = torch.softmax(outputs, dim=1)
+    return probabilities
 
 def main(input_path, model_path, output_file):
-    """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)
+    """Main function to predict and save results to Excel."""
+    model = load_model(model_path)
+    results = []
 
     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)
-        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, perform direct prediction
-        predicted_class, probabilities = predict_image_class(input_path, model, feature_extractor, device, class_names)
-        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.")
+        for img_name in os.listdir(input_path):
+            img_path = os.path.join(input_path, img_name)
+            if img_path.endswith(('.png', '.jpg', '.jpeg')):  # Check for image file types
+                probs = predict(model, img_path).cpu().numpy()[0]
+                result = {"Image Name": img_name}
+                for i, prob in enumerate(probs):
+                    result[f"Class {i} Probability"] = prob  # Store probabilities
+                results.append(result)
     else:
-        print('Invalid input path. Please provide a valid file or folder path.')
+        # If a single image file is provided
+        probs = predict(model, input_path).cpu().numpy()[0]
+        result = {"Image Name": os.path.basename(input_path)}
+        for i, prob in enumerate(probs):
+            result[f"Class {i} Probability"] = prob  # Store probabilities
+        results.append(result)
+
+    # Create DataFrame and save to Excel
+    df = pd.DataFrame(results)
+    df.to_excel(output_file, index=False)
+    print(f"Results saved to {output_file}")
 
 # Example call
-input_path = '/content/ddd.jpg'  # Replace with your image or folder path
+input_path = '/content/ddd.jpg'  # Replace with your image folder or single image path
 model_path = '/content/model.pth'  # Replace with your model path
-output_file = '/content/predictions.xlsx'  # Name of the output Excel file
+output_file = 'predictions.xlsx'  # Name of the output Excel file
 main(input_path, model_path, output_file)