Upload 2 files

predict.ipynb

@@ -0,0 +1,141 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": []
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "source": [
+        "import os\n",
+        "import torch\n",
+        "from transformers import ViTForImageClassification, ViTFeatureExtractor\n",
+        "from PIL import Image\n",
+        "import io\n",
+        "import pandas as pd\n",
+        "\n",
+        "def load_model(model_path):\n",
+        "    \"\"\"Load the pre-trained model and feature extractor.\"\"\"\n",
+        "    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
+        "    feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')\n",
+        "    model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224', num_labels=13)\n",
+        "    model.load_state_dict(torch.load(model_path))\n",
+        "    model = model.to(device)\n",
+        "    model.eval()  # Set the model to evaluation mode\n",
+        "    return model, feature_extractor, device\n",
+        "\n",
+        "def safe_load_image(path):\n",
+        "    \"\"\"Safely load an image, handling possible errors.\"\"\"\n",
+        "    try:\n",
+        "        with open(path, 'rb') as f:\n",
+        "            img = Image.open(io.BytesIO(f.read()))\n",
+        "            img = img.convert('RGB')\n",
+        "        return img\n",
+        "    except Exception as e:\n",
+        "        print(f\"Error loading image {path}: {e}\")\n",
+        "        return None\n",
+        "\n",
+        "def predict_image_class(image_path, model, feature_extractor, device, class_names):\n",
+        "    \"\"\"Predict the class of a given image.\"\"\"\n",
+        "    img = safe_load_image(image_path)\n",
+        "    if img is None:\n",
+        "        return None, None\n",
+        "\n",
+        "    # Preprocess the image\n",
+        "    inputs = feature_extractor(images=img, return_tensors=\"pt\").to(device)\n",
+        "\n",
+        "    # Make the prediction\n",
+        "    with torch.no_grad():\n",
+        "        outputs = model(**inputs).logits\n",
+        "        probabilities = torch.softmax(outputs, dim=1).cpu().numpy()[0]  # Calculate probabilities\n",
+        "        predicted_class_idx = outputs.argmax(dim=1).item()  # Get the predicted class index\n",
+        "        predicted_class = class_names[predicted_class_idx]  # Get class name from index\n",
+        "\n",
+        "    return predicted_class, probabilities\n",
+        "\n",
+        "def predict_images_in_folder(folder_path, model, feature_extractor, device, class_names):\n",
+        "    \"\"\"Predict the class of each image in a folder.\"\"\"\n",
+        "    results = []\n",
+        "    for filename in os.listdir(folder_path):\n",
+        "        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):\n",
+        "            image_path = os.path.join(folder_path, filename)\n",
+        "            predicted_class, probabilities = predict_image_class(image_path, model, feature_extractor, device, class_names)\n",
+        "            if predicted_class is not None:\n",
+        "                results.append({'Image Name': filename, 'Predicted Class': predicted_class, 'Probabilities': probabilities})\n",
+        "\n",
+        "    return results\n",
+        "\n",
+        "def save_results_to_excel(results, output_file):\n",
+        "    \"\"\"Save the prediction results to an Excel file.\"\"\"\n",
+        "    # Flatten the probabilities array and create a DataFrame\n",
+        "    rows = []\n",
+        "    for result in results:\n",
+        "        for idx, prob in enumerate(result['Probabilities']):\n",
+        "            rows.append({\n",
+        "                'Image Name': result['Image Name'],\n",
+        "                'Predicted Class': result['Predicted Class'],\n",
+        "                'Class Index': idx,\n",
+        "                'Probability': prob\n",
+        "            })\n",
+        "\n",
+        "    df = pd.DataFrame(rows)\n",
+        "\n",
+        "    # Sort by probability in descending order\n",
+        "    df = df.sort_values(by='Probability', ascending=False)\n",
+        "\n",
+        "    # Save to Excel\n",
+        "    df.to_excel(output_file, index=False)\n",
+        "\n",
+        "def main(input_path, model_path, output_file):\n",
+        "    \"\"\"Main function to perform image classification prediction, handling single images or folders, and saving results to Excel.\"\"\"\n",
+        "    class_names = ['anti_war_protest', 'combat', 'construction', 'fire', 'human_damage',\n",
+        "                   'humanitarian_aid', 'infrastructure', 'military_parade', 'military_vehicle',\n",
+        "                   'meeting', 'speech', 'refugee', 'victory']\n",
+        "\n",
+        "    model, feature_extractor, device = load_model(model_path)\n",
+        "\n",
+        "    if os.path.isdir(input_path):\n",
+        "        # If the input path is a folder, predict all images in that folder\n",
+        "        results = predict_images_in_folder(input_path, model, feature_extractor, device, class_names)\n",
+        "        save_results_to_excel(results, output_file)\n",
+        "        print(f'Prediction results saved to: {output_file}')\n",
+        "    elif os.path.isfile(input_path):\n",
+        "        # If the input path is a single image, make a direct prediction\n",
+        "        predicted_class, probabilities = predict_image_class(input_path, model, feature_extractor, device, class_names)\n",
+        "        print(f'The predicted class for image {os.path.basename(input_path)}: {predicted_class}')\n",
+        "    else:\n",
+        "        print('Invalid input path. Please provide a valid file or folder path.')\n"
+      ],
+      "metadata": {
+        "id": "340nVjm4AcDO"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Example call\n",
+        "input_path = '/path/to/your/image_or_folder'  # Replace with your image or folder path\n",
+        "model_path = '/kaggle/working/best_modelq.pth'  # Replace with your model path\n",
+        "output_file = 'predictions.xlsx'  # Name of the output Excel file\n",
+        "main(input_path, model_path, output_file)"
+      ],
+      "metadata": {
+        "id": "CY-fkhjdAeMM"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

predict.py

@@ -0,0 +1,113 @@
+# -*- 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
+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)