ct images also added

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+sample_images/* filter=lfs diff=lfs merge=lfs -text

Brain_Tumor_MRI_Diagnose_plus_openai.ipynb → Brain_Tumor_MRI_CT_Images_Diagnose_plus_openai.ipynb

@@ -1,418 +1,82 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "provenance": [],
-      "gpuType": "T4"
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    },
-    "language_info": {
-      "name": "python"
-    },
-    "accelerator": "GPU",
-    "widgets": {
-      "application/vnd.jupyter.widget-state+json": {
-        "1b0dc97d07a94fca9a24670e3813ff29": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_name": "HBoxModel",
-          "model_module_version": "1.5.0",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "HBoxModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "HBoxView",
-            "box_style": "",
-            "children": [
-              "IPY_MODEL_d19c38e4492c4b9bb4c03fb285df62a2",
-              "IPY_MODEL_815e13d941ee4a26929ca4ca01c8c615",
-              "IPY_MODEL_1fa8e074d5954a6daed1c025aa8191ab"
-            ],
-            "layout": "IPY_MODEL_489118d586fb4c0bb2a148796d39d0e3"
-          }
-        },
-        "d19c38e4492c4b9bb4c03fb285df62a2": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_name": "HTMLModel",
-          "model_module_version": "1.5.0",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "HTMLModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "HTMLView",
-            "description": "",
-            "description_tooltip": null,
-            "layout": "IPY_MODEL_1764933ac18c4c228ad8a852719c4f1a",
-            "placeholder": "​",
-            "style": "IPY_MODEL_07db2fdc39c849bf9fb4be93aca7cf4e",
-            "value": "model.safetensors: 100%"
-          }
-        },
-        "815e13d941ee4a26929ca4ca01c8c615": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_name": "FloatProgressModel",
-          "model_module_version": "1.5.0",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "FloatProgressModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "ProgressView",
-            "bar_style": "success",
-            "description": "",
-            "description_tooltip": null,
-            "layout": "IPY_MODEL_6b82e7b5673a4ad4bb112905754d0cda",
-            "max": 21355344,
-            "min": 0,
-            "orientation": "horizontal",
-            "style": "IPY_MODEL_693ce4424a8240a38ccebaabf030c5ff",
-            "value": 21355344
-          }
-        },
-        "1fa8e074d5954a6daed1c025aa8191ab": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_name": "HTMLModel",
-          "model_module_version": "1.5.0",
-          "state": {
-            "_dom_classes": [],
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "HTMLModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/controls",
-            "_view_module_version": "1.5.0",
-            "_view_name": "HTMLView",
-            "description": "",
-            "description_tooltip": null,
-            "layout": "IPY_MODEL_31420da8c458472c85381cf0e7a7a242",
-            "placeholder": "​",
-            "style": "IPY_MODEL_bd6b1fcf472a4860bbd8adb8a7f8c4ab",
-            "value": " 21.4M/21.4M [00:00<00:00, 53.6MB/s]"
-          }
-        },
-        "489118d586fb4c0bb2a148796d39d0e3": {
-          "model_module": "@jupyter-widgets/base",
-          "model_name": "LayoutModel",
-          "model_module_version": "1.2.0",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": null,
-            "flex": null,
-            "flex_flow": null,
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": null
-          }
-        },
-        "1764933ac18c4c228ad8a852719c4f1a": {
-          "model_module": "@jupyter-widgets/base",
-          "model_name": "LayoutModel",
-          "model_module_version": "1.2.0",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": null,
-            "flex": null,
-            "flex_flow": null,
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": null
-          }
-        },
-        "07db2fdc39c849bf9fb4be93aca7cf4e": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_name": "DescriptionStyleModel",
-          "model_module_version": "1.5.0",
-          "state": {
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "DescriptionStyleModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "StyleView",
-            "description_width": ""
-          }
-        },
-        "6b82e7b5673a4ad4bb112905754d0cda": {
-          "model_module": "@jupyter-widgets/base",
-          "model_name": "LayoutModel",
-          "model_module_version": "1.2.0",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": null,
-            "flex": null,
-            "flex_flow": null,
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": null
-          }
-        },
-        "693ce4424a8240a38ccebaabf030c5ff": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_name": "ProgressStyleModel",
-          "model_module_version": "1.5.0",
-          "state": {
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "ProgressStyleModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "StyleView",
-            "bar_color": null,
-            "description_width": ""
-          }
-        },
-        "31420da8c458472c85381cf0e7a7a242": {
-          "model_module": "@jupyter-widgets/base",
-          "model_name": "LayoutModel",
-          "model_module_version": "1.2.0",
-          "state": {
-            "_model_module": "@jupyter-widgets/base",
-            "_model_module_version": "1.2.0",
-            "_model_name": "LayoutModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "LayoutView",
-            "align_content": null,
-            "align_items": null,
-            "align_self": null,
-            "border": null,
-            "bottom": null,
-            "display": null,
-            "flex": null,
-            "flex_flow": null,
-            "grid_area": null,
-            "grid_auto_columns": null,
-            "grid_auto_flow": null,
-            "grid_auto_rows": null,
-            "grid_column": null,
-            "grid_gap": null,
-            "grid_row": null,
-            "grid_template_areas": null,
-            "grid_template_columns": null,
-            "grid_template_rows": null,
-            "height": null,
-            "justify_content": null,
-            "justify_items": null,
-            "left": null,
-            "margin": null,
-            "max_height": null,
-            "max_width": null,
-            "min_height": null,
-            "min_width": null,
-            "object_fit": null,
-            "object_position": null,
-            "order": null,
-            "overflow": null,
-            "overflow_x": null,
-            "overflow_y": null,
-            "padding": null,
-            "right": null,
-            "top": null,
-            "visibility": null,
-            "width": null
-          }
-        },
-        "bd6b1fcf472a4860bbd8adb8a7f8c4ab": {
-          "model_module": "@jupyter-widgets/controls",
-          "model_name": "DescriptionStyleModel",
-          "model_module_version": "1.5.0",
-          "state": {
-            "_model_module": "@jupyter-widgets/controls",
-            "_model_module_version": "1.5.0",
-            "_model_name": "DescriptionStyleModel",
-            "_view_count": null,
-            "_view_module": "@jupyter-widgets/base",
-            "_view_module_version": "1.2.0",
-            "_view_name": "StyleView",
-            "description_width": ""
-          }
-        }
-      }
-    }
-  },
   "cells": [
     {
       "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "yViCfhqXTD-l",
+        "outputId": "e599cb0e-fb7f-44d0-c84d-82ce2d17c0fe"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "\u001b[?25l   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/832.4 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K   \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m \u001b[32m829.4/832.4 kB\u001b[0m \u001b[31m36.7 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m832.4/832.4 kB\u001b[0m \u001b[31m22.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+            "\u001b[?25h\u001b[?25l   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/983.2 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m983.2/983.2 kB\u001b[0m \u001b[31m52.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+            "\u001b[?25h"
+          ]
+        }
+      ],
       "source": [
         "# Colab: install required packages\n",
         "!pip install -q kagglehub timm torch torchvision scikit-learn pandas pillow matplotlib gradio openai pytorch-lightning"
-      ],
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
-        "id": "yViCfhqXTD-l",
-        "outputId": "901b3218-85e1-4cc5-d380-839d8a3fca86"
+        "id": "B7iabh3oOyNe",
+        "outputId": "fb2920c5-a19b-49c8-814e-640ce263c27b"
       },
-      "execution_count": 10,
       "outputs": [
         {
           "output_type": "stream",
           "name": "stdout",
           "text": [
-            "\u001b[?25l   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/832.4 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m832.4/832.4 kB\u001b[0m \u001b[31m25.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-            "\u001b[?25h\u001b[?25l   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/983.2 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m983.2/983.2 kB\u001b[0m \u001b[31m60.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
-            "\u001b[?25h"
+            "Downloading from https://www.kaggle.com/api/v1/datasets/download/murtozalikhon/brain-tumor-multimodal-image-ct-and-mri?dataset_version_number=1...\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "100%|██████████| 361M/361M [00:02<00:00, 183MB/s]"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Extracting files...\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "\n"
           ]
-        }
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 11,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
         },
-        "id": "B7iabh3oOyNe",
-        "outputId": "b487c9fb-aa5c-4457-e23b-067c427c15ab"
-      },
-      "outputs": [
         {
           "output_type": "stream",
           "name": "stdout",
           "text": [
-            "Using Colab cache for faster access to the 'brain-tumor-classification-mri' dataset.\n",
-            "Dataset path: /kaggle/input/brain-tumor-classification-mri\n",
-            "Training\n",
-            "Testing\n"
+            "Dataset path: /root/.cache/kagglehub/datasets/murtozalikhon/brain-tumor-multimodal-image-ct-and-mri/versions/1\n",
+            "Dataset\n"
           ]
         }
       ],
       "source": [
         "import kagglehub\n",
-        "dataset_ref = \"sartajbhuvaji/brain-tumor-classification-mri\"\n",
+        "dataset_ref = \"murtozalikhon/brain-tumor-multimodal-image-ct-and-mri\"\n",
         "path = kagglehub.dataset_download(dataset_ref)\n",
         "print(\"Dataset path:\", path)\n",
         "\n",
@@ -424,18 +88,14 @@
     },
     {
       "cell_type": "code",
-      "source": [
-        "import torch\n",
-        "print(torch.cuda.is_available())"
-      ],
+      "execution_count": 3,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "gq9zpPkmZVK7",
-        "outputId": "1afaff9b-7d25-4697-eb7a-8e7c77255fb7"
+        "outputId": "174e6148-05a4-4b4b-8ebe-7b62ea806b1b"
       },
-      "execution_count": 20,
       "outputs": [
         {
           "output_type": "stream",
@@ -444,431 +104,274 @@
             "True\n"
           ]
         }
+      ],
+      "source": [
+        "import torch\n",
+        "print(torch.cuda.is_available())"
       ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YvSgspBGU3iG"
+      },
+      "outputs": [],
       "source": [
         "# Create CSV (multi-class) from folders (in-memory)\n",
         "\n",
         "# Change DATA_ROOT to the folder that directly contains glioma_tumor, meningioma_tumor, pituitary_tumor, no_tumor."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# CONFIG - edit these paths before running the pipeline\n",
+        "DATA_ROOT = \"/root/.cache/kagglehub/datasets/murtozalikhon/brain-tumor-multimodal-image-ct-and-mri/versions/1/Dataset\"\n",
+        "# Where to write outputs / checkpoints\n",
+        "OUT_DIR = \"/content/output\"\n",
+        "os.makedirs(OUT_DIR, exist_ok=True)\n",
+        "\n",
+        "# Training hyperparams (change if needed)\n",
+        "IMG_SIZE = 224\n",
+        "BATCH_SIZE = 8\n",
+        "EPOCHS = 10\n",
+        "LR = 1e-4\n",
+        "NUM_WORKERS = 2  # adjust for Colab\n",
+        "TOP_K = 3\n",
+        "\n",
+        "print(\"DATA_ROOT:\", DATA_ROOT)\n"
       ],
       "metadata": {
-        "id": "YvSgspBGU3iG"
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "H5gEKtfMJxVQ",
+        "outputId": "fd555a63-b825-4860-86df-e4bc3f0caa76"
       },
-      "execution_count": 12,
-      "outputs": []
+      "execution_count": 23,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "DATA_ROOT: /root/.cache/kagglehub/datasets/murtozalikhon/brain-tumor-multimodal-image-ct-and-mri/versions/1/Dataset\n"
+          ]
+        }
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": 19,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "1grjHP2ZU5st",
+        "outputId": "050866ba-120a-40d9-a7b3-d33cdcd21394"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Total images: 9618\n",
+            "label\n",
+            "no_tumor      4300\n",
+            "ct_tumor      2318\n",
+            "meningioma    1112\n",
+            "glioma         672\n",
+            "pituitary      629\n",
+            "mri_tumor      587\n",
+            "Name: count, dtype: int64\n",
+            "Wrote CSV: /content/output/combined_image_labels.csv\n"
+          ]
+        }
+      ],
       "source": [
+        "# Build combined CSV for your described dataset layout\n",
+        "# Paste & run in Colab; set PARENT_DIR to your dataset root\n",
+        "\n",
         "import os, pandas as pd\n",
+        "from pathlib import Path\n",
         "\n",
-        "# <-- EDIT this path to match your dataset root that has the 4 class folders\n",
-        "DATA_ROOT = \"/kaggle/input/brain-tumor-classification-mri/Training\"\n",
+        "# EDIT: point this to the parent folder that contains the two modality folders\n",
+        "PARENT_DIR = DATA_ROOT  # ← change this\n",
         "\n",
-        "LABEL_MAP = {\n",
-        "    \"glioma_tumor\": 0,\n",
-        "    \"meningioma_tumor\": 1,\n",
-        "    \"pituitary_tumor\": 2,\n",
-        "    \"no_tumor\": 3\n",
-        "}\n",
+        "# expected modality folder names (we'll match case-insensitively)\n",
+        "# adjust if your folder names differ\n",
+        "CT_FOLDER_KEY = \"ct\"\n",
+        "MRI_FOLDER_KEY = \"mri\"\n",
         "\n",
         "rows = []\n",
-        "for folder in sorted(os.listdir(DATA_ROOT)):\n",
-        "    folder_path = os.path.join(DATA_ROOT, folder)\n",
-        "    if not os.path.isdir(folder_path):\n",
-        "        continue\n",
-        "    if folder not in LABEL_MAP:\n",
-        "        print(\"Skipping unknown folder:\", folder)\n",
+        "for mod_name in sorted(os.listdir(PARENT_DIR)):\n",
+        "    mod_path = os.path.join(PARENT_DIR, mod_name)\n",
+        "    if not os.path.isdir(mod_path):\n",
         "        continue\n",
-        "    label = LABEL_MAP[folder]\n",
-        "    for fname in os.listdir(folder_path):\n",
-        "        if fname.lower().endswith((\".png\",\".jpg\",\".jpeg\")):\n",
-        "            rows.append({\"image\": os.path.join(folder_path, fname), \"label\": label})\n",
+        "    mod_lower = mod_name.lower()\n",
+        "    modality = \"CT\" if CT_FOLDER_KEY in mod_lower else (\"MRI\" if MRI_FOLDER_KEY in mod_lower else mod_name)\n",
+        "    # inside each modality we expect subfolders: 'healthy' and 'tumor'\n",
+        "    for sub in sorted(os.listdir(mod_path)):\n",
+        "        sub_path = os.path.join(mod_path, sub)\n",
+        "        if not os.path.isdir(sub_path):\n",
+        "            continue\n",
+        "        sub_lower = sub.lower()\n",
+        "        # CT modality handling\n",
+        "        if modality == \"CT\":\n",
+        "            if \"healthy\" in sub_lower:\n",
+        "                lab = \"no_tumor\"\n",
+        "            else:\n",
+        "                lab = \"ct_tumor\"\n",
+        "            for fn in os.listdir(sub_path):\n",
+        "                if fn.lower().endswith((\".png\",\".jpg\",\".jpeg\")):\n",
+        "                    rows.append({\"image\": os.path.join(sub_path, fn), \"label\": lab, \"modality\": \"CT\"})\n",
+        "        # MRI modality handling\n",
+        "        elif modality == \"MRI\":\n",
+        "            if \"healthy\" in sub_lower:\n",
+        "                for fn in os.listdir(sub_path):\n",
+        "                    if fn.lower().endswith((\".png\",\".jpg\",\".jpeg\")):\n",
+        "                        rows.append({\"image\": os.path.join(sub_path, fn), \"label\": \"no_tumor\", \"modality\": \"MRI\"})\n",
+        "            else:\n",
+        "                # tumor folder: try to parse subtype from filename, else fallback\n",
+        "                for fn in os.listdir(sub_path):\n",
+        "                    if not fn.lower().endswith((\".png\",\".jpg\",\".jpeg\")):\n",
+        "                        continue\n",
+        "                    low = fn.lower()\n",
+        "                    if \"glioma\" in low:\n",
+        "                        lab = \"glioma\"\n",
+        "                    elif \"meningioma\" in low:\n",
+        "                        lab = \"meningioma\"\n",
+        "                    elif \"pituitary\" in low:\n",
+        "                        lab = \"pituitary\"\n",
+        "                    else:\n",
+        "                        lab = \"mri_tumor\"\n",
+        "                    rows.append({\"image\": os.path.join(sub_path, fn), \"label\": lab, \"modality\": \"MRI\"})\n",
+        "        else:\n",
+        "            # unexpected modality name: treat similarly to MRI fallback\n",
+        "            for fn in os.listdir(sub_path):\n",
+        "                if fn.lower().endswith((\".png\",\".jpg\",\".jpeg\")):\n",
+        "                    rows.append({\"image\": os.path.join(sub_path, fn), \"label\": \"unknown_modality\", \"modality\": mod_name})\n",
         "\n",
         "df = pd.DataFrame(rows).sample(frac=1, random_state=42).reset_index(drop=True)\n",
         "print(\"Total images:\", len(df))\n",
         "print(df.label.value_counts())\n",
-        "df.head()\n"
-      ],
+        "out_csv = os.path.join(OUT_DIR, \"combined_image_labels.csv\")\n",
+        "df.to_csv(out_csv, index=False)\n",
+        "print(\"Wrote CSV:\", out_csv)\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 20,
       "metadata": {
         "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 367
+          "base_uri": "https://localhost:8080/"
         },
-        "id": "1grjHP2ZU5st",
-        "outputId": "a9c67cab-ac1d-46a1-ff2a-d99ed94d6c7a"
+        "id": "OvqCFuB1U5wR",
+        "outputId": "be745f82-a3b2-4b69-ea06-f0615f8cc321"
       },
-      "execution_count": 13,
       "outputs": [
         {
           "output_type": "stream",
           "name": "stdout",
           "text": [
-            "Total images: 2870\n",
-            "label\n",
-            "2    827\n",
-            "0    826\n",
-            "1    822\n",
-            "3    395\n",
-            "Name: count, dtype: int64\n"
+            "Train: 7213 Val: 1202 Test: 1203\n"
           ]
-        },
-        {
-          "output_type": "execute_result",
-          "data": {
-            "text/plain": [
-              "                                               image  label\n",
-              "0  /kaggle/input/brain-tumor-classification-mri/T...      0\n",
-              "1  /kaggle/input/brain-tumor-classification-mri/T...      2\n",
-              "2  /kaggle/input/brain-tumor-classification-mri/T...      0\n",
-              "3  /kaggle/input/brain-tumor-classification-mri/T...      0\n",
-              "4  /kaggle/input/brain-tumor-classification-mri/T...      2"
-            ],
-            "text/html": [
-              "\n",
-              "  <div id=\"df-870f654a-e6d9-4f6c-938c-91d0b4e98446\" class=\"colab-df-container\">\n",
-              "    <div>\n",
-              "<style scoped>\n",
-              "    .dataframe tbody tr th:only-of-type {\n",
-              "        vertical-align: middle;\n",
-              "    }\n",
-              "\n",
-              "    .dataframe tbody tr th {\n",
-              "        vertical-align: top;\n",
-              "    }\n",
-              "\n",
-              "    .dataframe thead th {\n",
-              "        text-align: right;\n",
-              "    }\n",
-              "</style>\n",
-              "<table border=\"1\" class=\"dataframe\">\n",
-              "  <thead>\n",
-              "    <tr style=\"text-align: right;\">\n",
-              "      <th></th>\n",
-              "      <th>image</th>\n",
-              "      <th>label</th>\n",
-              "    </tr>\n",
-              "  </thead>\n",
-              "  <tbody>\n",
-              "    <tr>\n",
-              "      <th>0</th>\n",
-              "      <td>/kaggle/input/brain-tumor-classification-mri/T...</td>\n",
-              "      <td>0</td>\n",
-              "    </tr>\n",
-              "    <tr>\n",
-              "      <th>1</th>\n",
-              "      <td>/kaggle/input/brain-tumor-classification-mri/T...</td>\n",
-              "      <td>2</td>\n",
-              "    </tr>\n",
-              "    <tr>\n",
-              "      <th>2</th>\n",
-              "      <td>/kaggle/input/brain-tumor-classification-mri/T...</td>\n",
-              "      <td>0</td>\n",
-              "    </tr>\n",
-              "    <tr>\n",
-              "      <th>3</th>\n",
-              "      <td>/kaggle/input/brain-tumor-classification-mri/T...</td>\n",
-              "      <td>0</td>\n",
-              "    </tr>\n",
-              "    <tr>\n",
-              "      <th>4</th>\n",
-              "      <td>/kaggle/input/brain-tumor-classification-mri/T...</td>\n",
-              "      <td>2</td>\n",
-              "    </tr>\n",
-              "  </tbody>\n",
-              "</table>\n",
-              "</div>\n",
-              "    <div class=\"colab-df-buttons\">\n",
-              "\n",
-              "  <div class=\"colab-df-container\">\n",
-              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-870f654a-e6d9-4f6c-938c-91d0b4e98446')\"\n",
-              "            title=\"Convert this dataframe to an interactive table.\"\n",
-              "            style=\"display:none;\">\n",
-              "\n",
-              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
-              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
-              "  </svg>\n",
-              "    </button>\n",
-              "\n",
-              "  <style>\n",
-              "    .colab-df-container {\n",
-              "      display:flex;\n",
-              "      gap: 12px;\n",
-              "    }\n",
-              "\n",
-              "    .colab-df-convert {\n",
-              "      background-color: #E8F0FE;\n",
-              "      border: none;\n",
-              "      border-radius: 50%;\n",
-              "      cursor: pointer;\n",
-              "      display: none;\n",
-              "      fill: #1967D2;\n",
-              "      height: 32px;\n",
-              "      padding: 0 0 0 0;\n",
-              "      width: 32px;\n",
-              "    }\n",
-              "\n",
-              "    .colab-df-convert:hover {\n",
-              "      background-color: #E2EBFA;\n",
-              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
-              "      fill: #174EA6;\n",
-              "    }\n",
-              "\n",
-              "    .colab-df-buttons div {\n",
-              "      margin-bottom: 4px;\n",
-              "    }\n",
-              "\n",
-              "    [theme=dark] .colab-df-convert {\n",
-              "      background-color: #3B4455;\n",
-              "      fill: #D2E3FC;\n",
-              "    }\n",
-              "\n",
-              "    [theme=dark] .colab-df-convert:hover {\n",
-              "      background-color: #434B5C;\n",
-              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
-              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
-              "      fill: #FFFFFF;\n",
-              "    }\n",
-              "  </style>\n",
-              "\n",
-              "    <script>\n",
-              "      const buttonEl =\n",
-              "        document.querySelector('#df-870f654a-e6d9-4f6c-938c-91d0b4e98446 button.colab-df-convert');\n",
-              "      buttonEl.style.display =\n",
-              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
-              "\n",
-              "      async function convertToInteractive(key) {\n",
-              "        const element = document.querySelector('#df-870f654a-e6d9-4f6c-938c-91d0b4e98446');\n",
-              "        const dataTable =\n",
-              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
-              "                                                    [key], {});\n",
-              "        if (!dataTable) return;\n",
-              "\n",
-              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
-              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
-              "          + ' to learn more about interactive tables.';\n",
-              "        element.innerHTML = '';\n",
-              "        dataTable['output_type'] = 'display_data';\n",
-              "        await google.colab.output.renderOutput(dataTable, element);\n",
-              "        const docLink = document.createElement('div');\n",
-              "        docLink.innerHTML = docLinkHtml;\n",
-              "        element.appendChild(docLink);\n",
-              "      }\n",
-              "    </script>\n",
-              "  </div>\n",
-              "\n",
-              "\n",
-              "    <div id=\"df-39e3087a-f1c1-4313-b033-254945753780\">\n",
-              "      <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-39e3087a-f1c1-4313-b033-254945753780')\"\n",
-              "                title=\"Suggest charts\"\n",
-              "                style=\"display:none;\">\n",
-              "\n",
-              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
-              "     width=\"24px\">\n",
-              "    <g>\n",
-              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
-              "    </g>\n",
-              "</svg>\n",
-              "      </button>\n",
-              "\n",
-              "<style>\n",
-              "  .colab-df-quickchart {\n",
-              "      --bg-color: #E8F0FE;\n",
-              "      --fill-color: #1967D2;\n",
-              "      --hover-bg-color: #E2EBFA;\n",
-              "      --hover-fill-color: #174EA6;\n",
-              "      --disabled-fill-color: #AAA;\n",
-              "      --disabled-bg-color: #DDD;\n",
-              "  }\n",
-              "\n",
-              "  [theme=dark] .colab-df-quickchart {\n",
-              "      --bg-color: #3B4455;\n",
-              "      --fill-color: #D2E3FC;\n",
-              "      --hover-bg-color: #434B5C;\n",
-              "      --hover-fill-color: #FFFFFF;\n",
-              "      --disabled-bg-color: #3B4455;\n",
-              "      --disabled-fill-color: #666;\n",
-              "  }\n",
-              "\n",
-              "  .colab-df-quickchart {\n",
-              "    background-color: var(--bg-color);\n",
-              "    border: none;\n",
-              "    border-radius: 50%;\n",
-              "    cursor: pointer;\n",
-              "    display: none;\n",
-              "    fill: var(--fill-color);\n",
-              "    height: 32px;\n",
-              "    padding: 0;\n",
-              "    width: 32px;\n",
-              "  }\n",
-              "\n",
-              "  .colab-df-quickchart:hover {\n",
-              "    background-color: var(--hover-bg-color);\n",
-              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
-              "    fill: var(--button-hover-fill-color);\n",
-              "  }\n",
-              "\n",
-              "  .colab-df-quickchart-complete:disabled,\n",
-              "  .colab-df-quickchart-complete:disabled:hover {\n",
-              "    background-color: var(--disabled-bg-color);\n",
-              "    fill: var(--disabled-fill-color);\n",
-              "    box-shadow: none;\n",
-              "  }\n",
-              "\n",
-              "  .colab-df-spinner {\n",
-              "    border: 2px solid var(--fill-color);\n",
-              "    border-color: transparent;\n",
-              "    border-bottom-color: var(--fill-color);\n",
-              "    animation:\n",
-              "      spin 1s steps(1) infinite;\n",
-              "  }\n",
-              "\n",
-              "  @keyframes spin {\n",
-              "    0% {\n",
-              "      border-color: transparent;\n",
-              "      border-bottom-color: var(--fill-color);\n",
-              "      border-left-color: var(--fill-color);\n",
-              "    }\n",
-              "    20% {\n",
-              "      border-color: transparent;\n",
-              "      border-left-color: var(--fill-color);\n",
-              "      border-top-color: var(--fill-color);\n",
-              "    }\n",
-              "    30% {\n",
-              "      border-color: transparent;\n",
-              "      border-left-color: var(--fill-color);\n",
-              "      border-top-color: var(--fill-color);\n",
-              "      border-right-color: var(--fill-color);\n",
-              "    }\n",
-              "    40% {\n",
-              "      border-color: transparent;\n",
-              "      border-right-color: var(--fill-color);\n",
-              "      border-top-color: var(--fill-color);\n",
-              "    }\n",
-              "    60% {\n",
-              "      border-color: transparent;\n",
-              "      border-right-color: var(--fill-color);\n",
-              "    }\n",
-              "    80% {\n",
-              "      border-color: transparent;\n",
-              "      border-right-color: var(--fill-color);\n",
-              "      border-bottom-color: var(--fill-color);\n",
-              "    }\n",
-              "    90% {\n",
-              "      border-color: transparent;\n",
-              "      border-bottom-color: var(--fill-color);\n",
-              "    }\n",
-              "  }\n",
-              "</style>\n",
-              "\n",
-              "      <script>\n",
-              "        async function quickchart(key) {\n",
-              "          const quickchartButtonEl =\n",
-              "            document.querySelector('#' + key + ' button');\n",
-              "          quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
-              "          quickchartButtonEl.classList.add('colab-df-spinner');\n",
-              "          try {\n",
-              "            const charts = await google.colab.kernel.invokeFunction(\n",
-              "                'suggestCharts', [key], {});\n",
-              "          } catch (error) {\n",
-              "            console.error('Error during call to suggestCharts:', error);\n",
-              "          }\n",
-              "          quickchartButtonEl.classList.remove('colab-df-spinner');\n",
-              "          quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
-              "        }\n",
-              "        (() => {\n",
-              "          let quickchartButtonEl =\n",
-              "            document.querySelector('#df-39e3087a-f1c1-4313-b033-254945753780 button');\n",
-              "          quickchartButtonEl.style.display =\n",
-              "            google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
-              "        })();\n",
-              "      </script>\n",
-              "    </div>\n",
-              "\n",
-              "    </div>\n",
-              "  </div>\n"
-            ],
-            "application/vnd.google.colaboratory.intrinsic+json": {
-              "type": "dataframe",
-              "variable_name": "df",
-              "summary": "{\n  \"name\": \"df\",\n  \"rows\": 2870,\n  \"fields\": [\n    {\n      \"column\": \"image\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 2870,\n        \"samples\": [\n          \"/kaggle/input/brain-tumor-classification-mri/Training/glioma_tumor/gg (87).jpg\",\n          \"/kaggle/input/brain-tumor-classification-mri/Training/meningioma_tumor/m2 (51).jpg\",\n          \"/kaggle/input/brain-tumor-classification-mri/Training/meningioma_tumor/m1(109).jpg\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"label\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 1,\n        \"min\": 0,\n        \"max\": 3,\n        \"num_unique_values\": 4,\n        \"samples\": [\n          2,\n          3,\n          0\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}"
-            }
-          },
-          "metadata": {},
-          "execution_count": 13
         }
+      ],
+      "source": [
+        "# Stratified split train/val/test\n",
+        "from sklearn.model_selection import train_test_split\n",
+        "csv_path = os.path.join(OUT_DIR, \"combined_image_labels.csv\")\n",
+        "df = pd.read_csv(csv_path)\n",
+        "\n",
+        "train_df, temp_df = train_test_split(df, test_size=0.25, stratify=df['label'], random_state=42)\n",
+        "val_df, test_df = train_test_split(temp_df, test_size=0.5, stratify=temp_df['label'], random_state=42)\n",
+        "\n",
+        "print(\"Train:\", len(train_df), \"Val:\", len(val_df), \"Test:\", len(test_df))\n",
+        "train_df.to_csv(os.path.join(OUT_DIR, \"train.csv\"), index=False)\n",
+        "val_df.to_csv(os.path.join(OUT_DIR, \"val.csv\"), index=False)\n",
+        "test_df.to_csv(os.path.join(OUT_DIR, \"test.csv\"), index=False)\n"
       ]
     },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "MTmULrSFKdxe"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
     {
       "cell_type": "code",
       "source": [
-        "# Create train / val split (80/20) and display counts\n",
-        "from sklearn.model_selection import train_test_split\n",
-        "\n",
-        "train_df, val_df = train_test_split(df, test_size=0.2, stratify=df['label'], random_state=42)\n",
-        "train_df = train_df.reset_index(drop=True)\n",
-        "val_df = val_df.reset_index(drop=True)\n",
-        "\n",
-        "print(\"Train:\", len(train_df), \"Val:\", len(val_df))\n",
-        "print(train_df.label.value_counts(), val_df.label.value_counts())\n"
+        "# Build label -> index mapping from training labels\n",
+        "labels = sorted(train_df['label'].unique().tolist())\n",
+        "LABEL2IDX = {lab: idx for idx, lab in enumerate(labels)}\n",
+        "IDX2LABEL = {v:k for k,v in LABEL2IDX.items()}\n",
+        "print(\"Labels:\", LABEL2IDX)\n",
+        "NUM_CLASSES = len(labels)\n"
       ],
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
-        "id": "OvqCFuB1U5wR",
-        "outputId": "2964fa89-7d65-4998-dd8f-d21b5d3456f4"
+        "id": "LKnSLphiKd6i",
+        "outputId": "48c16684-e96c-4146-9ba6-1fe0a28e138d"
       },
-      "execution_count": 14,
+      "execution_count": 21,
       "outputs": [
         {
           "output_type": "stream",
           "name": "stdout",
           "text": [
-            "Train: 2296 Val: 574\n",
-            "label\n",
-            "2    662\n",
-            "0    661\n",
-            "1    657\n",
-            "3    316\n",
-            "Name: count, dtype: int64 label\n",
-            "1    165\n",
-            "2    165\n",
-            "0    165\n",
-            "3     79\n",
-            "Name: count, dtype: int64\n"
+            "Labels: {'ct_tumor': 0, 'glioma': 1, 'meningioma': 2, 'mri_tumor': 3, 'no_tumor': 4, 'pituitary': 5}\n"
           ]
         }
       ]
     },
     {
       "cell_type": "code",
-      "source": [
-        "# Define Dataset and Transforms (inline)"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "MYxdFZEiY4_9"
       },
-      "execution_count": 15,
-      "outputs": []
+      "outputs": [],
+      "source": []
     },
     {
       "cell_type": "code",
+      "execution_count": 24,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Nd0TbtAlY7to",
+        "outputId": "63032d78-ebe7-4370-c780-746cee650c8f"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "batch x shape: torch.Size([8, 3, 224, 224]) y shape: torch.Size([8])\n"
+          ]
+        }
+      ],
       "source": [
+        "# Define Dataset and Transforms (inline)\n",
         "import torch\n",
         "from torch.utils.data import Dataset, DataLoader\n",
         "from PIL import Image\n",
         "import torchvision.transforms as T\n",
         "\n",
-        "IMG_SIZE = 224\n",
-        "\n",
+        "# transforms (same for MRI/CT for MVP)\n",
         "train_tf = T.Compose([\n",
         "    T.Resize((IMG_SIZE, IMG_SIZE)),\n",
         "    T.RandomHorizontalFlip(),\n",
         "    T.RandomRotation(10),\n",
-        "    T.ColorJitter(brightness=0.1, contrast=0.1),\n",
+        "    T.ColorJitter(0.08, 0.08, 0.08, 0.02),\n",
         "    T.ToTensor(),\n",
         "    T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])\n",
         "])\n",
@@ -878,117 +381,85 @@
         "    T.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])\n",
         "])\n",
         "\n",
-        "class MedicalMultiClassDataset(Dataset):\n",
-        "    def __init__(self, df, transform=None):\n",
+        "class MultiModalityDataset(Dataset):\n",
+        "    def __init__(self, df, transform=None, label_map=None):\n",
         "        self.df = df.reset_index(drop=True)\n",
         "        self.transform = transform\n",
-        "\n",
+        "        self.label_map = label_map\n",
         "    def __len__(self):\n",
         "        return len(self.df)\n",
-        "\n",
         "    def __getitem__(self, idx):\n",
         "        row = self.df.iloc[idx]\n",
-        "        img = Image.open(row['image']).convert('RGB')\n",
+        "        img_path = row['image']\n",
+        "        img = Image.open(img_path).convert('RGB')   # ensures 3 channels\n",
         "        if self.transform:\n",
         "            img = self.transform(img)\n",
-        "        label = int(row['label'])\n",
+        "        label = self.label_map[row['label']]\n",
         "        return img, label\n",
         "\n",
-        "# quick smoke test\n",
-        "ds = MedicalMultiClassDataset(train_df.iloc[:4], transform=train_tf)\n",
-        "for x,y in ds:\n",
-        "    print(x.shape, y)\n",
-        "    break\n"
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "Nd0TbtAlY7to",
-        "outputId": "78520285-f5f9-4987-ef9d-65be96356ca1"
-      },
-      "execution_count": 16,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "torch.Size([3, 224, 224]) 0\n"
-          ]
-        }
+        "train_ds = MultiModalityDataset(train_df, transform=train_tf, label_map=LABEL2IDX)\n",
+        "val_ds   = MultiModalityDataset(val_df, transform=val_tf, label_map=LABEL2IDX)\n",
+        "test_ds  = MultiModalityDataset(test_df, transform=val_tf, label_map=LABEL2IDX)\n",
+        "\n",
+        "train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=True)\n",
+        "val_loader   = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)\n",
+        "test_loader  = DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)\n",
+        "\n",
+        "# sanity check\n",
+        "x,y = next(iter(train_loader))\n",
+        "print(\"batch x shape:\", x.shape, \"y shape:\", y.shape)"
       ]
     },
     {
       "cell_type": "code",
-      "source": [],
+      "execution_count": null,
       "metadata": {
         "id": "SrJLOm33ZMpC"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": []
     },
     {
       "cell_type": "code",
-      "source": [
-        "# Build dataloaders"
-      ],
+      "execution_count": null,
       "metadata": {
         "id": "bG2JclRDY-Di"
       },
-      "execution_count": 17,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "BATCH_SIZE = 8  # adjust to your GPU memory\n",
-        "\n",
-        "train_loader = DataLoader(MedicalMultiClassDataset(train_df, transform=train_tf), batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True)\n",
-        "val_loader   = DataLoader(MedicalMultiClassDataset(val_df, transform=val_tf), batch_size=BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)\n"
-      ],
-      "metadata": {
-        "id": "Jc4pdqreZGDS"
-      },
-      "execution_count": 18,
-      "outputs": []
+      "outputs": [],
+      "source": []
     },
     {
       "cell_type": "code",
-      "source": [
-        "# Define model, optimizer, loss, device\n",
-        "import timm, torch.nn as nn, torch\n",
-        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
-        "NUM_CLASSES = 4\n",
-        "\n",
-        "model = timm.create_model('efficientnet_b0', pretrained=True, num_classes=NUM_CLASSES)\n",
-        "model = model.to(device)\n",
-        "\n",
-        "loss_fn = nn.CrossEntropyLoss()\n",
-        "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)\n"
-      ],
+      "execution_count": 25,
       "metadata": {
+        "id": "Jc4pdqreZGDS",
         "colab": {
           "base_uri": "https://localhost:8080/",
-          "height": 153,
+          "height": 188,
           "referenced_widgets": [
-            "1b0dc97d07a94fca9a24670e3813ff29",
-            "d19c38e4492c4b9bb4c03fb285df62a2",
-            "815e13d941ee4a26929ca4ca01c8c615",
-            "1fa8e074d5954a6daed1c025aa8191ab",
-            "489118d586fb4c0bb2a148796d39d0e3",
-            "1764933ac18c4c228ad8a852719c4f1a",
-            "07db2fdc39c849bf9fb4be93aca7cf4e",
-            "6b82e7b5673a4ad4bb112905754d0cda",
-            "693ce4424a8240a38ccebaabf030c5ff",
-            "31420da8c458472c85381cf0e7a7a242",
-            "bd6b1fcf472a4860bbd8adb8a7f8c4ab"
+            "fd1f533d7d814e10b252bf66f8c0697e",
+            "243ec1d23b8e469e833db01f89189639",
+            "9bcf41c7fe1e425aae4e37dcfca6a67c",
+            "f09e93d2fdc64071b3b9860fd7782c58",
+            "0a03432f7ce14c8c98090f846a512beb",
+            "1da74b90beb44d85a22ca57ee52b1e9d",
+            "a7c35ed82c20419b9b559b217d924a6a",
+            "0a302207f9884533b49fb150411169d2",
+            "bc007e650c0f4f619946055a6c72d132",
+            "20b5ba3ee19446dc89ddb5222be6f5f2",
+            "1e35273de451460aa8c82df17b3ca23c"
           ]
         },
-        "id": "TRBF6u9FZLnT",
-        "outputId": "2e9229ac-0260-41ce-ea49-9c3982ca7ab6"
+        "outputId": "8f6811b3-e720-4ffd-b49b-45c177d2972e"
       },
-      "execution_count": 19,
       "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Using device: cuda\n"
+          ]
+        },
         {
           "output_type": "stream",
           "name": "stderr",
@@ -1010,31 +481,90 @@
             "application/vnd.jupyter.widget-view+json": {
               "version_major": 2,
               "version_minor": 0,
-              "model_id": "1b0dc97d07a94fca9a24670e3813ff29"
+              "model_id": "fd1f533d7d814e10b252bf66f8c0697e"
             }
           },
           "metadata": {}
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Class weights: tensor([0.6917, 2.3853, 1.4414, 2.7322, 0.3728, 2.5470], device='cuda:0')\n"
+          ]
         }
+      ],
+      "source": [
+        "# Model, loss, optimizer, optional class weights\n",
+        "import timm\n",
+        "import torch.nn as nn\n",
+        "\n",
+        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "print(\"Using device:\", device)\n",
+        "\n",
+        "model = timm.create_model('efficientnet_b0', pretrained=True, num_classes=NUM_CLASSES)\n",
+        "model = model.to(device)\n",
+        "\n",
+        "# Optionally compute class weights if highly imbalanced\n",
+        "from sklearn.utils.class_weight import compute_class_weight\n",
+        "import numpy as np\n",
+        "cls_w = compute_class_weight(class_weight='balanced', classes=np.arange(NUM_CLASSES), y=train_df['label'].map(LABEL2IDX))\n",
+        "cls_w = torch.tensor(cls_w, dtype=torch.float).to(device)\n",
+        "print(\"Class weights:\", cls_w)\n",
+        "\n",
+        "loss_fn = nn.CrossEntropyLoss(weight=cls_w)\n",
+        "optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-4)\n",
+        "scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.5, patience=2)\n"
       ]
     },
     {
       "cell_type": "code",
+      "execution_count": 26,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "nUDYRwowZRGb",
+        "outputId": "bdac381c-00bd-4330-f8d9-f2ba47208e67"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Epoch 1/12 | train_loss 0.7560 train_acc 0.7829 | val_loss 0.4183 val_acc 0.8968 | 1.48 min\n",
+            "Saved best model: 0.8968386023294509\n",
+            "Epoch 2/12 | train_loss 0.3758 train_acc 0.8959 | val_loss 0.3714 val_acc 0.9160 | 1.32 min\n",
+            "Saved best model: 0.9159733777038269\n",
+            "Epoch 3/12 | train_loss 0.2867 train_acc 0.9186 | val_loss 0.3187 val_acc 0.9268 | 1.27 min\n",
+            "Saved best model: 0.9267886855241264\n",
+            "Epoch 4/12 | train_loss 0.2459 train_acc 0.9326 | val_loss 0.2907 val_acc 0.9368 | 1.28 min\n",
+            "Saved best model: 0.9367720465890182\n",
+            "Epoch 5/12 | train_loss 0.2312 train_acc 0.9368 | val_loss 0.3151 val_acc 0.9285 | 1.26 min\n",
+            "Epoch 6/12 | train_loss 0.2106 train_acc 0.9427 | val_loss 0.3328 val_acc 0.9276 | 1.26 min\n",
+            "Epoch 7/12 | train_loss 0.1896 train_acc 0.9473 | val_loss 0.3394 val_acc 0.9301 | 1.28 min\n",
+            "Early stopping triggered. Stopping training.\n"
+          ]
+        }
+      ],
       "source": [
+        "# Training loop with val_loss, val_acc, checkpointing & early stopping\n",
         "from sklearn.metrics import accuracy_score\n",
         "import time, os\n",
+        "import pandas as pd\n",
         "\n",
-        "EPOCHS = 10\n",
-        "CKPT_DIR = \"/content/checkpoints\"\n",
-        "os.makedirs(CKPT_DIR, exist_ok=True)\n",
+        "os.makedirs(OUT_DIR, exist_ok=True)\n",
         "best_val_acc = 0.0\n",
+        "patience = 3\n",
+        "wait = 0\n",
+        "logs = []\n",
         "\n",
         "for epoch in range(1, EPOCHS+1):\n",
+        "    t0 = time.time()\n",
+        "    # train\n",
         "    model.train()\n",
-        "    running_loss = 0.0\n",
+        "    train_loss = 0.0\n",
         "    all_preds, all_labels = [], []\n",
-        "    t0 = time.time()\n",
-        "\n",
-        "    # ---- Training ----\n",
         "    for imgs, labels in train_loader:\n",
         "        imgs = imgs.to(device)\n",
         "        labels = labels.to(device)\n",
@@ -1043,75 +573,125 @@
         "        loss = loss_fn(logits, labels)\n",
         "        loss.backward()\n",
         "        optimizer.step()\n",
-        "        running_loss += loss.item() * imgs.size(0)\n",
+        "        train_loss += loss.item() * imgs.size(0)\n",
         "        preds = logits.argmax(dim=1).cpu().numpy()\n",
         "        all_preds.extend(preds.tolist())\n",
         "        all_labels.extend(labels.cpu().numpy().tolist())\n",
-        "\n",
-        "    train_loss = running_loss / len(train_loader.dataset)\n",
+        "    train_loss = train_loss / len(train_loader.dataset)\n",
         "    train_acc = accuracy_score(all_labels, all_preds)\n",
         "\n",
-        "    # ---- Validation ----\n",
+        "    # val\n",
         "    model.eval()\n",
-        "    val_running_loss = 0.0\n",
+        "    val_loss = 0.0\n",
         "    v_preds, v_labels = [], []\n",
         "    with torch.no_grad():\n",
         "        for imgs, labels in val_loader:\n",
         "            imgs = imgs.to(device)\n",
         "            labels = labels.to(device)\n",
         "            logits = model(imgs)\n",
-        "            loss = loss_fn(logits, labels)   # <-- compute val loss\n",
-        "            val_running_loss += loss.item() * imgs.size(0)\n",
+        "            loss = loss_fn(logits, labels)\n",
+        "            val_loss += loss.item() * imgs.size(0)\n",
         "            preds = logits.argmax(dim=1).cpu().numpy()\n",
         "            v_preds.extend(preds.tolist())\n",
         "            v_labels.extend(labels.cpu().numpy().tolist())\n",
-        "\n",
-        "    val_loss = val_running_loss / len(val_loader.dataset)\n",
+        "    val_loss = val_loss / len(val_loader.dataset)\n",
         "    val_acc = accuracy_score(v_labels, v_preds)\n",
         "\n",
-        "    elapsed = time.time() - t0\n",
-        "    print(f\"Epoch {epoch}/{EPOCHS} \"\n",
-        "          f\"— train_loss: {train_loss:.4f} train_acc: {train_acc:.4f} \"\n",
-        "          f\"val_loss: {val_loss:.4f} val_acc: {val_acc:.4f} \"\n",
-        "          f\"time: {elapsed/60:.2f}m\")\n",
+        "    # scheduler step\n",
+        "    scheduler.step(val_loss)\n",
+        "\n",
+        "    elapsed = (time.time() - t0) / 60\n",
+        "    print(f\"Epoch {epoch}/{EPOCHS} | train_loss {train_loss:.4f} train_acc {train_acc:.4f} | val_loss {val_loss:.4f} val_acc {val_acc:.4f} | {elapsed:.2f} min\")\n",
+        "    logs.append({\"epoch\":epoch,\"train_loss\":train_loss,\"train_acc\":train_acc,\"val_loss\":val_loss,\"val_acc\":val_acc})\n",
         "\n",
-        "    # ---- Save checkpoint ----\n",
-        "    ckpt_path = os.path.join(CKPT_DIR, f\"epoch{epoch}.pth\")\n",
-        "    torch.save(model.state_dict(), ckpt_path)\n",
+        "    # checkpoint\n",
+        "    ckpt = os.path.join(OUT_DIR, f\"epoch{epoch}.pth\")\n",
+        "    torch.save(model.state_dict(), ckpt)\n",
         "    if val_acc > best_val_acc:\n",
         "        best_val_acc = val_acc\n",
-        "        torch.save(model.state_dict(), os.path.join(CKPT_DIR, \"best_model_multiclass.pth\"))\n",
-        "        print(\"Saved best model:\", best_val_acc)\n"
+        "        torch.save(model.state_dict(), os.path.join(OUT_DIR, \"best_model.pth\"))\n",
+        "        print(\"Saved best model:\", best_val_acc)\n",
+        "        wait = 0\n",
+        "    else:\n",
+        "        wait += 1\n",
+        "        if wait >= patience:\n",
+        "            print(\"Early stopping triggered. Stopping training.\")\n",
+        "            break\n",
+        "\n",
+        "# save logs to csv\n",
+        "pd.DataFrame(logs).to_csv(os.path.join(OUT_DIR, \"training_log.csv\"), index=False)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "I2WUIxNnZtVI"
+      },
+      "outputs": [],
+      "source": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Evaluate on test set (classification report + confusion matrix)\n",
+        "from sklearn.metrics import classification_report, confusion_matrix\n",
+        "import numpy as np\n",
+        "\n",
+        "# load best model\n",
+        "best_path = os.path.join(OUT_DIR, \"best_model.pth\")\n",
+        "model.load_state_dict(torch.load(best_path, map_location=device))\n",
+        "model.eval()\n",
+        "\n",
+        "y_true, y_pred = [], []\n",
+        "with torch.no_grad():\n",
+        "    for imgs, labels in test_loader:\n",
+        "        imgs = imgs.to(device)\n",
+        "        labels = labels.to(device)\n",
+        "        logits = model(imgs)\n",
+        "        preds = logits.argmax(dim=1).cpu().numpy()\n",
+        "        y_pred.extend(preds.tolist())\n",
+        "        y_true.extend(labels.cpu().numpy().tolist())\n",
+        "\n",
+        "print(\"Classification Report:\")\n",
+        "print(classification_report(y_true, y_pred, target_names=[IDX2LABEL[i] for i in range(NUM_CLASSES)]))\n",
+        "print(\"Confusion matrix:\")\n",
+        "print(confusion_matrix(y_true, y_pred))\n"
       ],
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
-        "id": "nUDYRwowZRGb",
-        "outputId": "65d7584c-b90e-4a70-bf44-f6747fbe271e"
+        "id": "nbkHkAPELGYq",
+        "outputId": "19358314-2ce1-4bf4-8844-45a0d723100c"
       },
-      "execution_count": 22,
+      "execution_count": 28,
       "outputs": [
         {
           "output_type": "stream",
           "name": "stdout",
           "text": [
-            "Epoch 1/10 — train_loss: 0.0583 train_acc: 0.9800 val_loss: 0.1981 val_acc: 0.9512 time: 0.38m\n",
-            "Saved best model: 0.9512195121951219\n",
-            "Epoch 2/10 — train_loss: 0.0519 train_acc: 0.9817 val_loss: 0.2222 val_acc: 0.9425 time: 0.45m\n",
-            "Epoch 3/10 — train_loss: 0.0665 train_acc: 0.9756 val_loss: 0.1690 val_acc: 0.9582 time: 0.39m\n",
-            "Saved best model: 0.9581881533101045\n",
-            "Epoch 4/10 — train_loss: 0.0410 train_acc: 0.9861 val_loss: 0.1936 val_acc: 0.9564 time: 0.39m\n",
-            "Epoch 5/10 — train_loss: 0.0629 train_acc: 0.9821 val_loss: 0.1469 val_acc: 0.9634 time: 0.39m\n",
-            "Saved best model: 0.9634146341463414\n",
-            "Epoch 6/10 — train_loss: 0.0530 train_acc: 0.9843 val_loss: 0.1630 val_acc: 0.9652 time: 0.41m\n",
-            "Saved best model: 0.9651567944250871\n",
-            "Epoch 7/10 — train_loss: 0.0346 train_acc: 0.9869 val_loss: 0.1726 val_acc: 0.9634 time: 0.41m\n",
-            "Epoch 8/10 — train_loss: 0.0268 train_acc: 0.9917 val_loss: 0.1771 val_acc: 0.9582 time: 0.39m\n",
-            "Epoch 9/10 — train_loss: 0.0349 train_acc: 0.9874 val_loss: 0.1282 val_acc: 0.9756 time: 0.40m\n",
-            "Saved best model: 0.975609756097561\n",
-            "Epoch 10/10 — train_loss: 0.0338 train_acc: 0.9891 val_loss: 0.1267 val_acc: 0.9791 time: 0.39m\n",
-            "Saved best model: 0.9790940766550522\n"
+            "Classification Report:\n",
+            "              precision    recall  f1-score   support\n",
+            "\n",
+            "    ct_tumor       0.98      0.99      0.99       290\n",
+            "      glioma       0.95      0.93      0.94        84\n",
+            "  meningioma       0.77      0.95      0.85       139\n",
+            "   mri_tumor       0.80      0.53      0.64        73\n",
+            "    no_tumor       1.00      0.99      0.99       538\n",
+            "   pituitary       0.99      0.94      0.96        79\n",
+            "\n",
+            "    accuracy                           0.95      1203\n",
+            "   macro avg       0.91      0.89      0.90      1203\n",
+            "weighted avg       0.95      0.95      0.95      1203\n",
+            "\n",
+            "Confusion matrix:\n",
+            "[[288   0   0   0   2   0]\n",
+            " [  0  78   5   1   0   0]\n",
+            " [  0   0 132   7   0   0]\n",
+            " [  0   4  29  39   0   1]\n",
+            " [  5   0   2   0 531   0]\n",
+            " [  0   0   3   2   0  74]]\n"
           ]
         }
       ]
@@ -1120,117 +700,175 @@
       "cell_type": "code",
       "source": [],
       "metadata": {
-        "id": "I2WUIxNnZtVI"
+        "id": "QZCaUaWzLGbS"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "QMaVFK9xLGgu"
       },
       "execution_count": null,
       "outputs": []
     },
     {
       "cell_type": "code",
+      "execution_count": 33,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "H2BqijX2Zupx",
+        "outputId": "3ed96f69-87a9-4b2b-839e-2a8680c672e9"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Sample: /root/.cache/kagglehub/datasets/murtozalikhon/brain-tumor-multimodal-image-ct-and-mri/versions/1/Dataset/Brain Tumor CT scan Images/Tumor/ct_tumor (137).jpg\n",
+            "[{'label': 'ct_tumor', 'score': 0.9998412132263184}, {'label': 'meningioma', 'score': 0.00015001322026364505}, {'label': 'mri_tumor', 'score': 5.044169483880978e-06}]\n"
+          ]
+        }
+      ],
       "source": [
         "# Load best model & test a single image (inference helper)\n",
         "import torch.nn.functional as F\n",
         "from PIL import Image\n",
         "from torchvision import transforms\n",
         "\n",
-        "# load\n",
-        "best_ckpt = \"/content/checkpoints/best_model_multiclass.pth\"\n",
-        "model.load_state_dict(torch.load(best_ckpt, map_location=device))\n",
-        "model.eval()\n",
-        "\n",
-        "IDX2LABEL = {0:\"glioma_tumor\", 1:\"meningioma_tumor\", 2:\"pituitary_tumor\", 3:\"no_tumor\"}\n",
+        "def load_model_for_infer(ckpt_path=os.path.join(OUT_DIR,\"best_model.pth\")):\n",
+        "    m = timm.create_model('efficientnet_b0', pretrained=False, num_classes=NUM_CLASSES)\n",
+        "    m.load_state_dict(torch.load(ckpt_path, map_location=device))\n",
+        "    m.to(device).eval()\n",
+        "    return m\n",
         "\n",
-        "def preprocess(pil_img, img_size=IMG_SIZE):\n",
-        "    tf = transforms.Compose([\n",
-        "        transforms.Resize((img_size, img_size)),\n",
-        "        transforms.ToTensor(),\n",
-        "        transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225])\n",
-        "    ])\n",
-        "    return tf(pil_img).unsqueeze(0).to(device)\n",
-        "\n",
-        "def predict(pil_img, topk=3):\n",
-        "    x = preprocess(pil_img)\n",
+        "# reuse tf from earlier\n",
+        "def predict_topk_pil(model, pil_img, topk=TOP_K):\n",
+        "    x = val_tf(pil_img).unsqueeze(0).to(device)\n",
         "    with torch.no_grad():\n",
         "        logits = model(x)\n",
         "        probs = F.softmax(logits, dim=1).cpu().numpy().ravel()\n",
         "    idxs = probs.argsort()[::-1][:topk]\n",
         "    return [{\"label\": IDX2LABEL[int(i)], \"score\": float(probs[int(i)])} for i in idxs]\n",
         "\n",
-        "# quick test with a sample image from val set\n",
-        "sample_path = val_df.iloc[0].image\n",
-        "print(\"Sample:\", sample_path)\n",
+        "# quick example\n",
+        "infer_model = load_model_for_infer()\n",
+        "sample_path = val_df.iloc[110].image\n",
         "img = Image.open(sample_path).convert(\"RGB\")\n",
-        "preds = predict(img, topk=3)\n",
-        "preds\n"
-      ],
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "H2BqijX2Zupx",
-        "outputId": "8ad2aef5-e047-4054-acca-f50cbc3429be"
-      },
-      "execution_count": 29,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "Sample: /kaggle/input/brain-tumor-classification-mri/Training/meningioma_tumor/m3 (155).jpg\n"
-          ]
-        },
-        {
-          "output_type": "execute_result",
-          "data": {
-            "text/plain": [
-              "[{'label': 'meningioma_tumor', 'score': 0.9999998807907104},\n",
-              " {'label': 'pituitary_tumor', 'score': 5.87046180555717e-08},\n",
-              " {'label': 'glioma_tumor', 'score': 4.05661282343317e-08}]"
-            ]
-          },
-          "metadata": {},
-          "execution_count": 29
-        }
+        "print(\"Sample:\", sample_path)\n",
+        "print(predict_topk_pil(infer_model, img, topk=3))"
       ]
     },
     {
       "cell_type": "code",
-      "source": [],
+      "execution_count": null,
       "metadata": {
         "id": "RAZRNlq2ZutV"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": []
     },
     {
       "cell_type": "code",
-      "source": [],
+      "execution_count": null,
       "metadata": {
         "id": "chDN6CJmc8h0"
       },
+      "outputs": [],
+      "source": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# gradio App plus Openai"
+      ],
+      "metadata": {
+        "id": "LsQ5FIkOOHch"
+      },
       "execution_count": null,
       "outputs": []
     },
     {
       "cell_type": "code",
+      "execution_count": 39,
+      "metadata": {
+        "id": "WSygonZYdvpg"
+      },
+      "outputs": [],
       "source": [
         "sample_paths = [\n",
-        "    \"/kaggle/input/brain-tumor-classification-mri/Testing/glioma_tumor/image(10).jpg\",\n",
-        "    \"/kaggle/input/brain-tumor-classification-mri/Testing/meningioma_tumor/image(10).jpg\",\n",
-        "    \"/kaggle/input/brain-tumor-classification-mri/Testing/no_tumor/image(10).jpg\",\n",
-        "    \"/kaggle/input/brain-tumor-classification-mri/Testing/pituitary_tumor/image(10).jpg\",\n",
+        "    \"/root/.cache/kagglehub/datasets/murtozalikhon/brain-tumor-multimodal-image-ct-and-mri/versions/1/Dataset/Brain Tumor CT scan Images/Tumor/ct_tumor (10).jpg\",\n",
+        "    \"/root/.cache/kagglehub/datasets/murtozalikhon/brain-tumor-multimodal-image-ct-and-mri/versions/1/Dataset/Brain Tumor CT scan Images/Healthy/ct_healthy (1).jpg\",\n",
+        "    \"/root/.cache/kagglehub/datasets/murtozalikhon/brain-tumor-multimodal-image-ct-and-mri/versions/1/Dataset/Brain Tumor MRI images/Tumor/glioma (10).jpg\",\n",
+        "    \"/root/.cache/kagglehub/datasets/murtozalikhon/brain-tumor-multimodal-image-ct-and-mri/versions/1/Dataset/Brain Tumor MRI images/Tumor/meningioma (1056).jpg\",\n",
         "]\n",
         "\n",
         "OPENAI_API_KEY = \"\""
-      ],
-      "metadata": {
-        "id": "WSygonZYdvpg"
-      },
-      "execution_count": 36,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": 43,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 680
+        },
+        "id": "S_rWA1Xcdvx8",
+        "outputId": "55c18697-f04d-45bd-bddc-f59e4a729c3d"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Loading model...\n",
+            "Model loaded on cuda\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "/tmp/ipython-input-2087682989.py:119: UserWarning: You have not specified a value for the `type` parameter. Defaulting to the 'tuples' format for chatbot messages, but this is deprecated and will be removed in a future version of Gradio. Please set type='messages' instead, which uses openai-style dictionaries with 'role' and 'content' keys.\n",
+            "  chatbot = gr.Chatbot(label=\"Assistant\")\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Colab notebook detected. To show errors in colab notebook, set debug=True in launch()\n",
+            "* Running on public URL: https://a1ab9ef721b239799a.gradio.live\n",
+            "\n",
+            "This share link expires in 1 week. For free permanent hosting and GPU upgrades, run `gradio deploy` from the terminal in the working directory to deploy to Hugging Face Spaces (https://huggingface.co/spaces)\n"
+          ]
+        },
+        {
+          "output_type": "display_data",
+          "data": {
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ],
+            "text/html": [
+              "<div><iframe src=\"https://a1ab9ef721b239799a.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+            ]
+          },
+          "metadata": {}
+        },
+        {
+          "output_type": "execute_result",
+          "data": {
+            "text/plain": []
+          },
+          "metadata": {},
+          "execution_count": 43
+        }
+      ],
       "source": [
         "import os\n",
         "from pathlib import Path\n",
@@ -1243,23 +881,33 @@
         "import openai\n",
         "from openai import OpenAI\n",
         "\n",
-        "\n",
-        "# ------------------ USER CONFIG ------------------\n",
-        "CKPT_PATH = \"/content/checkpoints/best_model_multiclass.pth\"   # path to trained checkpoint\n",
+        "# ------------------ USER CONFIG (edit as needed) ------------------\n",
+        "CKPT_PATH = \"/content/output/best_model.pth\"   # path to trained checkpoint\n",
         "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
         "BACKBONE = \"efficientnet_b0\"\n",
-        "NUM_CLASSES = 4\n",
         "IMG_SIZE = 224\n",
         "\n",
-        "# Paste your 4 sample image absolute paths here\n",
-        "# e.g. sample_paths = [\"/content/a.png\", \"/content/b.png\", ...]\n",
-        "sample_paths = sample_paths\n",
+        "# Paste your 4 sample image absolute paths here (or ensure sample_paths variable already exists)\n",
+        "# Example:\n",
+        "# sample_paths = [\"/content/samples/ct1.jpg\", \"/content/samples/mri_glioma1.jpg\", ...]\n",
+        "sample_paths = sample_paths  # keep your existing variable/assignment\n",
         "\n",
-        "OPENAI_API_KEY = OPENAI_API_KEY\n",
+        "OPENAI_API_KEY = OPENAI_API_KEY  # keep as before (e.g., from env or variable)\n",
         "if OPENAI_API_KEY:\n",
         "    openai.api_key = OPENAI_API_KEY\n",
         "\n",
-        "IDX2LABEL = {0: \"glioma_tumor\", 1: \"meningioma_tumor\", 2: \"pituitary_tumor\", 3: \"no_tumor\"}\n",
+        "# ------------------ UPDATED LABEL MAP for combined CT+MRI ------------------\n",
+        "# Minimal change: include ct_tumor and MRI subtypes + no_tumor\n",
+        "IDX2LABEL = {\n",
+        "    0: \"ct_tumor\",\n",
+        "    1: \"glioma\",\n",
+        "    2: \"meningioma\",\n",
+        "    3: \"mri_tumor\",\n",
+        "    4: \"no_tumor\",\n",
+        "    5: \"pituitary\"\n",
+        "}\n",
+        "\n",
+        "NUM_CLASSES = len(IDX2LABEL)\n",
         "# -------------------------------------------------\n",
         "\n",
         "# Validate sample paths\n",
@@ -1295,7 +943,7 @@
         "    idxs = probs.argsort()[::-1][:topk]\n",
         "    return [{\"label\": IDX2LABEL[int(i)], \"score\": float(probs[int(i)])} for i in idxs]\n",
         "\n",
-        "# OpenAI system prompt\n",
+        "# OpenAI system prompt (unchanged)\n",
         "SYSTEM_TEMPLATE = \"\"\"You are an educational radiology assistant. The model predicted:\n",
         "{preds}\n",
         "\n",
@@ -1325,13 +973,12 @@
         "    )\n",
         "    return resp.choices[0].message.content.strip()\n",
         "\n",
-        "\n",
         "# Prepare thumbnails (just images)\n",
         "thumbs = [Image.open(p).convert(\"RGB\") for p in sample_paths]\n",
         "\n",
-        "# Gradio UI\n",
+        "# Gradio UI (unchanged structure)\n",
         "with gr.Blocks() as demo:\n",
-        "    gr.Markdown(\"## Brain MRI multi-class demo — Educational only\")\n",
+        "    gr.Markdown(\"## Brain MRI+CT multi-class demo — Educational only\")\n",
         "    with gr.Row():\n",
         "        with gr.Column():\n",
         "            # Use a simple list of PIL thumbnails for the gallery value\n",
@@ -1349,16 +996,16 @@
         "    selected_path_state = gr.State(value=None)\n",
         "    preds_state = gr.State(value=None)\n",
         "\n",
-        "    # select image from gallery\n",
+        "    # select image from gallery -> store selected path string\n",
         "    def on_select(evt: gr.SelectData):\n",
-        "        idx = evt.index\n",
-        "        path = sample_paths[idx]           # <-- return only the path (string)\n",
+        "        idx = int(evt.index)\n",
+        "        path = sample_paths[idx]\n",
         "        return path\n",
         "\n",
         "    # ensure the gallery select writes only the path to the state\n",
         "    gallery.select(fn=on_select, inputs=None, outputs=[selected_path_state])\n",
         "\n",
-        "    # analyze\n",
+        "    # analyze (unchanged except label names)\n",
         "    def analyze(path):\n",
         "      try:\n",
         "        if path is None:\n",
@@ -1375,7 +1022,7 @@
         "\n",
         "    analyze_btn.click(fn=analyze, inputs=selected_path_state, outputs=[preds_output, chatbot, preds_state])\n",
         "\n",
-        "    # chat\n",
+        "    # chat (keeps same logic)\n",
         "    def chat(chat_history, msg, preds_text):\n",
         "      if preds_text is None:\n",
         "          return chat_history+[(\"AI\",\"Please analyze an image first.\")], \"\"\n",
@@ -1406,74 +1053,391 @@
         "\n",
         "    gr.Markdown(\"⚠️ Educational use only — not a medical diagnosis.\")\n",
         "\n",
-        "demo.launch(share=True)\n"
-      ],
+        "demo.launch(share=True)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "TrS_HlY7d8s3"
+      },
+      "outputs": [],
+      "source": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
       "metadata": {
         "colab": {
-          "base_uri": "https://localhost:8080/",
-          "height": 680
+          "background_save": true
+        },
+        "id": "6Yj2jOOY0swp"
+      },
+      "outputs": [],
+      "source": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "widgets": {
+      "application/vnd.jupyter.widget-state+json": {
+        "fd1f533d7d814e10b252bf66f8c0697e": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "HBoxModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_243ec1d23b8e469e833db01f89189639",
+              "IPY_MODEL_9bcf41c7fe1e425aae4e37dcfca6a67c",
+              "IPY_MODEL_f09e93d2fdc64071b3b9860fd7782c58"
+            ],
+            "layout": "IPY_MODEL_0a03432f7ce14c8c98090f846a512beb"
+          }
+        },
+        "243ec1d23b8e469e833db01f89189639": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "HTMLModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_1da74b90beb44d85a22ca57ee52b1e9d",
+            "placeholder": "​",
+            "style": "IPY_MODEL_a7c35ed82c20419b9b559b217d924a6a",
+            "value": "model.safetensors: 100%"
+          }
+        },
+        "9bcf41c7fe1e425aae4e37dcfca6a67c": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "FloatProgressModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_0a302207f9884533b49fb150411169d2",
+            "max": 21355344,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_bc007e650c0f4f619946055a6c72d132",
+            "value": 21355344
+          }
+        },
+        "f09e93d2fdc64071b3b9860fd7782c58": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "HTMLModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_20b5ba3ee19446dc89ddb5222be6f5f2",
+            "placeholder": "​",
+            "style": "IPY_MODEL_1e35273de451460aa8c82df17b3ca23c",
+            "value": " 21.4M/21.4M [00:00&lt;00:00, 43.5MB/s]"
+          }
+        },
+        "0a03432f7ce14c8c98090f846a512beb": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
         },
-        "id": "S_rWA1Xcdvx8",
-        "outputId": "5735af82-a1ea-4963-d8da-6cc480836bb6"
-      },
-      "execution_count": 50,
-      "outputs": [
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "Loading model...\n",
-            "Model loaded on cuda\n"
-          ]
+        "1da74b90beb44d85a22ca57ee52b1e9d": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
         },
-        {
-          "output_type": "stream",
-          "name": "stderr",
-          "text": [
-            "/tmp/ipython-input-4085922474.py:110: UserWarning: You have not specified a value for the `type` parameter. Defaulting to the 'tuples' format for chatbot messages, but this is deprecated and will be removed in a future version of Gradio. Please set type='messages' instead, which uses openai-style dictionaries with 'role' and 'content' keys.\n",
-            "  chatbot = gr.Chatbot(label=\"Assistant\")\n"
-          ]
+        "a7c35ed82c20419b9b559b217d924a6a": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "DescriptionStyleModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
         },
-        {
-          "output_type": "stream",
-          "name": "stdout",
-          "text": [
-            "Colab notebook detected. To show errors in colab notebook, set debug=True in launch()\n",
-            "* Running on public URL: https://53b5220cc19d12e1dc.gradio.live\n",
-            "\n",
-            "This share link expires in 1 week. For free permanent hosting and GPU upgrades, run `gradio deploy` from the terminal in the working directory to deploy to Hugging Face Spaces (https://huggingface.co/spaces)\n"
-          ]
+        "0a302207f9884533b49fb150411169d2": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
         },
-        {
-          "output_type": "display_data",
-          "data": {
-            "text/plain": [
-              "<IPython.core.display.HTML object>"
-            ],
-            "text/html": [
-              "<div><iframe src=\"https://53b5220cc19d12e1dc.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
-            ]
-          },
-          "metadata": {}
+        "bc007e650c0f4f619946055a6c72d132": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "ProgressStyleModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
         },
-        {
-          "output_type": "execute_result",
-          "data": {
-            "text/plain": []
-          },
-          "metadata": {},
-          "execution_count": 50
+        "20b5ba3ee19446dc89ddb5222be6f5f2": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "1e35273de451460aa8c82df17b3ca23c": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "DescriptionStyleModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
         }
-      ]
-    },
-    {
-      "cell_type": "code",
-      "source": [],
-      "metadata": {
-        "id": "TrS_HlY7d8s3"
-      },
-      "execution_count": null,
-      "outputs": []
+      }
     }
-  ]
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
 }

app.py

@@ -9,28 +9,45 @@ import gradio as gr
 import openai
 from openai import OpenAI
 
+sample_paths = [
+    "sample_images/ct_healthy (1).jpg",
+    "sample_images/ct_tumor (1).png",
+    "sample_images/glioma (1).jpg",
+    "sample_images/pituitary (244).jpg",
+]
+
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+if OPENAI_API_KEY:
+    openai.api_key = OPENAI_API_KEY
 
-# ------------------ USER CONFIG ------------------
-CKPT_PATH = "model/best_model_multiclass.pth"   # path to trained checkpoint
+
+# ------------------ USER CONFIG (edit as needed) ------------------
+CKPT_PATH = "model/best_model.pth"   # path to trained checkpoint
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 BACKBONE = "efficientnet_b0"
-NUM_CLASSES = 4
 IMG_SIZE = 224
 
-# Paste your 4 sample image absolute paths here
-# e.g. sample_paths = ["/content/a.png", "/content/b.png", ...]
-sample_paths = [
-    "sample_images/image.jpg",
-    "sample_images/image(2).jpg",
-    "sample_images/image(3).jpg",
-    "sample_images/image(23).jpg",
-]
+# Paste your 4 sample image absolute paths here (or ensure sample_paths variable already exists)
+# Example:
+# sample_paths = ["/content/samples/ct1.jpg", "/content/samples/mri_glioma1.jpg", ...]
+sample_paths = sample_paths  # keep your existing variable/assignment
 
-OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+OPENAI_API_KEY = OPENAI_API_KEY  # keep as before (e.g., from env or variable)
 if OPENAI_API_KEY:
     openai.api_key = OPENAI_API_KEY
 
-IDX2LABEL = {0: "glioma_tumor", 1: "meningioma_tumor", 2: "pituitary_tumor", 3: "no_tumor"}
+# ------------------ UPDATED LABEL MAP for combined CT+MRI ------------------
+# Minimal change: include ct_tumor and MRI subtypes + no_tumor
+IDX2LABEL = {
+    0: "ct_tumor",
+    1: "glioma",
+    2: "meningioma",
+    3: "mri_tumor",
+    4: "no_tumor",
+    5: "pituitary"
+}
+
+NUM_CLASSES = len(IDX2LABEL)
 # -------------------------------------------------
 
 # Validate sample paths
@@ -66,7 +83,7 @@ def predict_topk_from_pil(pil_img, topk=3):
     idxs = probs.argsort()[::-1][:topk]
     return [{"label": IDX2LABEL[int(i)], "score": float(probs[int(i)])} for i in idxs]
 
-# OpenAI system prompt
+# OpenAI system prompt (unchanged)
 SYSTEM_TEMPLATE = """You are an educational radiology assistant. The model predicted:
 {preds}
 
@@ -96,13 +113,12 @@ def call_openai_seed(preds_text):
     )
     return resp.choices[0].message.content.strip()
 
-
 # Prepare thumbnails (just images)
 thumbs = [Image.open(p).convert("RGB") for p in sample_paths]
 
-# Gradio UI
+# Gradio UI (unchanged structure)
 with gr.Blocks() as demo:
-    gr.Markdown("## Brain MRI multi-class demo — Educational only")
+    gr.Markdown("## Brain MRI+CT multi-class demo — Educational only")
     with gr.Row():
         with gr.Column():
             # Use a simple list of PIL thumbnails for the gallery value
@@ -120,16 +136,16 @@ with gr.Blocks() as demo:
     selected_path_state = gr.State(value=None)
     preds_state = gr.State(value=None)
 
-    # select image from gallery
+    # select image from gallery -> store selected path string
     def on_select(evt: gr.SelectData):
-        idx = evt.index
-        path = sample_paths[idx]           # <-- return only the path (string)
+        idx = int(evt.index)
+        path = sample_paths[idx]
         return path
 
     # ensure the gallery select writes only the path to the state
     gallery.select(fn=on_select, inputs=None, outputs=[selected_path_state])
 
-    # analyze
+    # analyze (unchanged except label names)
     def analyze(path):
       try:
         if path is None:
@@ -146,7 +162,7 @@ with gr.Blocks() as demo:
 
     analyze_btn.click(fn=analyze, inputs=selected_path_state, outputs=[preds_output, chatbot, preds_state])
 
-    # chat
+    # chat (keeps same logic)
     def chat(chat_history, msg, preds_text):
       if preds_text is None:
           return chat_history+[("AI","Please analyze an image first.")], ""
@@ -177,4 +193,4 @@ with gr.Blocks() as demo:
 
     gr.Markdown("⚠️ Educational use only — not a medical diagnosis.")
 
-demo.launch(share=True)
+demo.launch(share=True)

model/{best_model_multiclass.pth → best_model.pth}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:90f3a7461ec275012fc223ca25422cfdcfab75379e2b406841a32f34b3f6619c
-size 16350387
+oid sha256:66488a18d397630b3f6eda091504415329f432fbf90b0a14d1fae5be90bc7864
+size 16356601