Delete file predict.ipynb

predict.ipynb

@@ -1,104 +0,0 @@
-import os
-import torch
-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')
-    model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224', num_labels=13)
-    model.load_state_dict(torch.load(model_path))
-    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."""
-    try:
-        with open(path, 'rb') as f:
-            img = Image.open(io.BytesIO(f.read()))
-            img = img.convert('RGB')
-        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)
-
-    # Make the 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
-
-    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):
-    """Save the prediction results to an Excel file."""
-    # Flatten the probabilities array and create a DataFrame
-    rows = []
-    for result in results:
-        for idx, prob in enumerate(result['Probabilities']):
-            rows.append({
-                'Image Name': result['Image Name'],
-                'Predicted Class': result['Predicted Class'],
-                'Class Index': 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)
-
-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', 
-                   'meeting', 'speech', 'refugee', 'victory']
-
-    model, feature_extractor, device = load_model(model_path)
-
-    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}')
-    elif os.path.isfile(input_path):
-        # If the input path is a single image, make a 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}')
-    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
-main(input_path, model_path, output_file)