Upload predict.ipynb

predict.ipynb

@@ -0,0 +1,179 @@
+{
+  "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",
+        "\n",
+        "    # Load the model with weights mapped to the appropriate device\n",
+        "    model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224', num_labels=13, ignore_mismatched_sizes=True)\n",
+        "\n",
+        "    # Load the state dict with map_location to handle devices properly\n",
+        "    state_dict = torch.load(model_path, map_location=device)\n",
+        "    model.load_state_dict(state_dict, strict=False)  # Use strict=False to ignore size mismatches\n",
+        "\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 potential errors.\"\"\"\n",
+        "    try:\n",
+        "        with open(path, 'rb') as f:\n",
+        "            img = Image.open(io.BytesIO(f.read()))\n",
+        "            img = img.convert('RGB')\n",
+        "            img = img.resize((224, 224))  # Resize the image to (224, 224)\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",
+        "    # Perform 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 the class name based on the 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, class_names):\n",
+        "    \"\"\"Save prediction results to an Excel file.\"\"\"\n",
+        "    # Flatten probability array and create DataFrame\n",
+        "    rows = []\n",
+        "    for result in results:\n",
+        "        # Add each probability with corresponding class name\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': class_names[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",
+        "    print(f'Results saved to {output_file}')  # Confirm saving\n",
+        "\n",
+        "def main(input_path, model_path, output_file):\n",
+        "    \"\"\"Main function to execute image classification predictions, processing 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",
+        "        if results:\n",
+        "            save_results_to_excel(results, output_file, class_names)\n",
+        "        else:\n",
+        "            print(\"No valid images found in the specified folder.\")\n",
+        "    elif os.path.isfile(input_path):\n",
+        "        # If the input path is a single image, perform direct prediction\n",
+        "        predicted_class, probabilities = predict_image_class(input_path, model, feature_extractor, device, class_names)\n",
+        "        if predicted_class is not None:\n",
+        "            print(f'Predicted class for image {os.path.basename(input_path)}: {predicted_class}')\n",
+        "        else:\n",
+        "            print(\"Image could not be processed.\")\n",
+        "    else:\n",
+        "        print('Invalid input path. Please provide a valid file or folder path.')\n"
+      ],
+      "metadata": {
+        "id": "340nVjm4AcDO"
+      },
+      "execution_count": 15,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Example call\n",
+        "input_path = '/content/ddd.jpg'  # Replace with your image or folder path\n",
+        "model_path = '/content/model.pth'  # Replace with your model path\n",
+        "output_file = '/content/predictions.xlsx'  # Name of the output Excel file\n",
+        "main(input_path, model_path, output_file)"
+      ],
+      "metadata": {
+        "id": "CY-fkhjdAeMM",
+        "outputId": "1de8f113-97a5-419d-995b-815ec5391a80",
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        }
+      },
+      "execution_count": 18,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "Some weights of ViTForImageClassification were not initialized from the model checkpoint at google/vit-base-patch16-224 and are newly initialized because the shapes did not match:\n",
+            "- classifier.bias: found shape torch.Size([1000]) in the checkpoint and torch.Size([13]) in the model instantiated\n",
+            "- classifier.weight: found shape torch.Size([1000, 768]) in the checkpoint and torch.Size([13, 768]) in the model instantiated\n",
+            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
+            "<ipython-input-15-1637f10c4a55>:17: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+            "  state_dict = torch.load(model_path, map_location=device)\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Predicted class for image ddd.jpg: infrastructure\n"
+          ]
+        }
+      ]
+    }
+  ]
+}