added image split

.gitignore

@@ -0,0 +1,3 @@
+dataset/
+splits/
+.venv/

training/training.ipynb

@@ -0,0 +1,704 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "46ebdaab",
+   "metadata": {},
+   "source": [
+    "## This is Classification Model Training for DDR dataset. This model uses Densenet-121 a pretrained model.\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ad55f9e5",
+   "metadata": {},
+   "source": [
+    "### STEPS for model training\n",
+    "##### Step 1. Define Preprocessing Function \n",
+    "###### Median Filtering, CLAHE, Gamma Correction, ESRGAN(optional)\n",
+    "##### Step 2. Create Custom Dataset for preprocessing (Pytorch doesn't supports Custom Preprocessing)\n",
+    "##### Step 3. Define Transform (with data augmentation)\n",
+    "##### Step 4. Create datasets and Dataloaders\n",
+    "##### Step 5. Load DenseNet-121 Model\n",
+    "##### Step 6. Define Loss, Optimizer, Scheduler\n",
+    "##### Step 7. Model Training\n",
+    "##### Step 8. Validation and Evaluation\n",
+    "##### Step 9. Save the Model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "58a3a0fd",
+   "metadata": {},
+   "source": [
+    "## Dataset Descrption\n",
+    "\n",
+    "### DDR dataset contains 13,673 fundus images from 147 hospitals, covering 23 provinces in China. The images are classified into 5 classes according to DR severity: none, mild, moderate, severe, and proliferative DR. There is a sixth category which indicates the images with poor quality. The dataset presented here does not include the images with poor quality (sixth category) and all images have been preprocessed to delete the black background. https://www.kaggle.com/datasets/mariaherrerot/ddrdataset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "255d5477",
+   "metadata": {},
+   "source": [
+    "### Import Libraries"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "80aeae51",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import cv2\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "from PIL import Image\n",
+    "from matplotlib import pyplot as plt\n",
+    "from sklearn.model_selection import train_test_split\n",
+    "import shutil \n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from tqdm import tqdm\n",
+    "import torch.optim as optim\n",
+    "from torch.optim.lr_scheduler import StepLR\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "from torchvision import models, transforms"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2a8e8a65",
+   "metadata": {},
+   "source": [
+    "## Load and split to train_test_val DDR Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "99b3890e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "train_labels.csv already exists. Skipping this split.\n",
+      "test_labels.csv already exists. Skipping this split.\n",
+      "Splits and CSVs checked and saved successfully!\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Load CSV\n",
+    "df = pd.read_csv(\"D:\\\\DR_Classification\\\\dataset\\\\DR_grading.csv\")\n",
+    "df['image_path'] = df['id_code'].apply(lambda x: os.path.join(\"D:\\\\DR_Classification\\\\dataset\\\\images\", x))\n",
+    "df['label'] = df['diagnosis']\n",
+    "\n",
+    "# Create output directories\n",
+    "output_root = \"D:\\\\DR_Classification\\\\splits\"\n",
+    "os.makedirs(os.path.join(output_root, \"train\"), exist_ok=True)\n",
+    "os.makedirs(os.path.join(output_root, \"test\"), exist_ok=True)\n",
+    "\n",
+    "# Split: train vs test\n",
+    "train_df, test_df = train_test_split(df, test_size=0.3, stratify=df['label'], random_state=42)\n",
+    "\n",
+    "# Function to copy images and save CSV\n",
+    "def save_split(df_split, split_name):\n",
+    "    split_folder = os.path.join(output_root, split_name)\n",
+    "    new_paths = []\n",
+    "\n",
+    "    # Check if the CSV for the current split already exists\n",
+    "    csv_path = os.path.join(output_root, f\"{split_name}_labels.csv\")\n",
+    "    if os.path.exists(csv_path):\n",
+    "        print(f\"{split_name}_labels.csv already exists. Skipping this split.\")\n",
+    "        return  # Skip processing if the CSV file exists\n",
+    "\n",
+    "    # Proceed with copying images and saving the new CSV if the CSV does not exist\n",
+    "    for _, row in df_split.iterrows():\n",
+    "        src = row['image_path']\n",
+    "        dst = os.path.join(split_folder, os.path.basename(src))\n",
+    "\n",
+    "        if os.path.exists(src):\n",
+    "            shutil.copy(src, dst)\n",
+    "            new_paths.append(dst)\n",
+    "        else:\n",
+    "            print(f\"Warning: Missing image file {src}\")\n",
+    "\n",
+    "    df_split = df_split.copy()\n",
+    "    df_split['new_path'] = new_paths\n",
+    "    df_split[['id_code', 'label', 'new_path']].to_csv(csv_path, index=False)\n",
+    "    print(f\"{split_name}_labels.csv saved successfully!\")\n",
+    "\n",
+    "# Save each split\n",
+    "save_split(train_df, \"train\")\n",
+    "save_split(test_df, \"test\")\n",
+    "\n",
+    "print(\"Splits and CSVs checked and saved successfully!\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "d6bda349",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "0    D:\\DR_Classification\\dataset\\images\\2017041310...\n",
+      "1    D:\\DR_Classification\\dataset\\images\\2017041311...\n",
+      "2    D:\\DR_Classification\\dataset\\images\\2017041311...\n",
+      "3    D:\\DR_Classification\\dataset\\images\\2017041311...\n",
+      "4    D:\\DR_Classification\\dataset\\images\\2017041311...\n",
+      "Name: image_path, dtype: object\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(df['image_path'].head())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "734943e9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Train class distribution:\n",
+      "label\n",
+      "0    4386\n",
+      "2    3134\n",
+      "4     639\n",
+      "1     441\n",
+      "3     165\n",
+      "Name: count, dtype: int64\n",
+      "\n",
+      "Test class distribution:\n",
+      "label\n",
+      "0    1880\n",
+      "2    1343\n",
+      "4     274\n",
+      "1     189\n",
+      "3      71\n",
+      "Name: count, dtype: int64\n",
+      "Total: 12522 images\n",
+      "Train: 8765\n",
+      "Test: 3757\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Count samples per class in each split\n",
+    "print(\"Train class distribution:\")\n",
+    "print(train_df['label'].value_counts())\n",
+    "\n",
+    "print(\"\\nTest class distribution:\")\n",
+    "print(test_df['label'].value_counts())\n",
+    "\n",
+    "# Count the total number of images in each split\n",
+    "print(f\"Total: {len(df)} images\")\n",
+    "print(f\"Train: {len(train_df)}\")\n",
+    "print(f\"Test: {len(test_df)}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "3264f03c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load the split CSVs\n",
+    "train_df = pd.read_csv(\"D:/DR_Classification/splits/train_labels.csv\")\n",
+    "test_df = pd.read_csv(\"D:/DR_Classification/splits/test_labels.csv\")\n",
+    "\n",
+    "# Extract paths and labels\n",
+    "train_paths = train_df['new_path'].tolist()\n",
+    "train_labels = train_df['label'].tolist()\n",
+    "\n",
+    "test_paths = test_df['new_path'].tolist()\n",
+    "test_labels = test_df['label'].tolist()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "85dc2767",
+   "metadata": {},
+   "source": [
+    "### Step 1. Define Preprocessing Functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "0b0e9212",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def apply_median_filter(image):\n",
+    "    return cv2.medianBlur(image, 3)\n",
+    "\n",
+    "def apply_clahe(image):\n",
+    "    lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)\n",
+    "    l, a, b = cv2.split(lab)\n",
+    "    clahe = cv2.createCLAHE(clipLimit=2.0)\n",
+    "    cl = clahe.apply(l)\n",
+    "    merged = cv2.merge((cl, a, b))\n",
+    "    return cv2.cvtColor(merged, cv2.COLOR_LAB2RGB)\n",
+    "\n",
+    "def apply_gamma_correction(image, gamma=1.2):\n",
+    "    invGamma = 1.0 / gamma\n",
+    "    table = np.array([(i / 255.0) ** invGamma * 255 for i in np.arange(0, 256)]).astype(\"uint8\")\n",
+    "    return cv2.LUT(image, table)\n",
+    "\n",
+    "def apply_gaussian_filter(image, kernel_size=(5, 5), sigma=1.0):\n",
+    "    return cv2.GaussianBlur(image, kernel_size, sigma)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1e34f571",
+   "metadata": {},
+   "source": [
+    "#### For the ESRGAN if applicable"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "dd662d1d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# def apply_esrgan(image):\n",
+    "#     # Placeholder: You should replace with actual ESRGAN model output\n",
+    "#     return enhance_with_esrgan(image)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9470983d",
+   "metadata": {},
+   "source": [
+    "### Step 2. Create Custom Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "165f2fc2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Apply customary dataset\n",
+    "class DDRDataset(Dataset):\n",
+    "    def __init__(self, image_paths, labels, transform=None):\n",
+    "        self.image_paths = image_paths\n",
+    "        self.labels = labels\n",
+    "        self.transform = transform\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.image_paths)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        img_path = self.image_paths[idx]\n",
+    "        image = cv2.imread(img_path)\n",
+    "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
+    "\n",
+    "        # Preprocessing steps\n",
+    "        image = apply_median_filter(image)\n",
+    "        image = apply_clahe(image)\n",
+    "        image = apply_gamma_correction(image)\n",
+    "        \n",
+    "        # Apply Gaussian filter\n",
+    "        image = apply_gaussian_filter(image, kernel_size=(5, 5), sigma=1.0)\n",
+    "\n",
+    "        image = Image.fromarray(image)\n",
+    "        if self.transform:\n",
+    "            image = self.transform(image)\n",
+    "\n",
+    "        label = torch.tensor(self.labels[idx], dtype=torch.long)\n",
+    "        return image, label\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7b2fb514",
+   "metadata": {},
+   "source": [
+    "### Step 3. Define Transforms with data augmentations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "0a59f18e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_transform = transforms.Compose([\n",
+    "    transforms.Resize((224, 224)),\n",
+    "    transforms.RandomHorizontalFlip(),\n",
+    "    transforms.RandomVerticalFlip(),  # Add vertical flip\n",
+    "    transforms.RandomRotation(degrees=15),  # Random rotation\n",
+    "    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),  # Random resized crop\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # ImageNet pretraining normalization\n",
+    "])\n",
+    "\n",
+    "val_transform = transforms.Compose([\n",
+    "    transforms.Resize((224, 224)),\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # this is the same as training\n",
+    "])\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e6a2cc7e",
+   "metadata": {},
+   "source": [
+    "### Step 4. Create Datasets and Dataloaders"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "419c07cb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sklearn.model_selection import train_test_split"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "cd96cd1b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Dataset class should accept image paths and labels\n",
+    "train_dataset = DDRDataset(train_paths, train_labels, transform=train_transform)\n",
+    "test_dataset = DDRDataset(test_paths, test_labels, transform=val_transform)\n",
+    "\n",
+    "# Create dataloaders\n",
+    "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
+    "test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "626d78b3",
+   "metadata": {},
+   "source": [
+    "### Step 5. Load DenseNet-121 model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "808006f9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "d:\\DR_Classification\\.venv\\Lib\\site-packages\\torchvision\\models\\_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
+      "  warnings.warn(\n",
+      "d:\\DR_Classification\\.venv\\Lib\\site-packages\\torchvision\\models\\_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=DenseNet121_Weights.IMAGENET1K_V1`. You can also use `weights=DenseNet121_Weights.DEFAULT` to get the most up-to-date weights.\n",
+      "  warnings.warn(msg)\n"
+     ]
+    }
+   ],
+   "source": [
+    "model = models.densenet121(pretrained=True)\n",
+    "model.classifier = nn.Linear(model.classifier.in_features, 5)  # DDR has 5 classes\n",
+    "\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "model = model.to(device)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c9c2bf69",
+   "metadata": {},
+   "source": [
+    "### Step 6. Define Loss, Optimizer, Scheduler, checkpoint"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "0bf18c57",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "criterion = nn.CrossEntropyLoss()\n",
+    "optimizer = optim.AdamW(model.parameters(), lr=1e-3)\n",
+    "scheduler = StepLR(optimizer, step_size=5, gamma=0.1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "576c9583",
+   "metadata": {},
+   "source": [
+    "### Step 7. Train the Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9b76b385",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/15:   0%|          | 0/274 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/15: 100%|██████████| 274/274 [41:09<00:00,  9.01s/it, accuracy=68.1, loss=0.833]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/15\n",
+      "Train Loss: 0.8326 | Train Accuracy: 68.09%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 2/15:   4%|▍         | 12/274 [01:44<37:59,  8.70s/it, accuracy=71.6, loss=0.743]"
+     ]
+    }
+   ],
+   "source": [
+    "def train_model(model, train_loader, criterion, optimizer, scheduler, num_epochs=15, device='cuda'):\n",
+    "    model.train()  # Set model to training mode\n",
+    "    for epoch in range(num_epochs):\n",
+    "        running_train_loss = 0.0\n",
+    "        correct_train = 0\n",
+    "        total_train = 0\n",
+    "\n",
+    "        # Add tqdm progress bar\n",
+    "        with tqdm(total=len(train_loader), desc=f\"Epoch {epoch+1}/{num_epochs}\") as pbar:\n",
+    "            for batch_idx, (inputs, labels) in enumerate(train_loader):\n",
+    "                inputs, labels = inputs.to(device), labels.to(device)\n",
+    "\n",
+    "                optimizer.zero_grad()\n",
+    "                outputs = model(inputs)  # Forward pass\n",
+    "                loss = criterion(outputs, labels)  # Compute loss\n",
+    "                loss.backward()  # Backpropagate\n",
+    "                optimizer.step()  # Update weights\n",
+    "\n",
+    "                running_train_loss += loss.item()\n",
+    "                _, predicted = torch.max(outputs, 1)\n",
+    "                correct_train += (predicted == labels).sum().item()\n",
+    "                total_train += labels.size(0)\n",
+    "\n",
+    "                pbar.set_postfix(loss=running_train_loss / (batch_idx + 1), accuracy=correct_train / total_train * 100)\n",
+    "                pbar.update(1)  # Update progress bar\n",
+    "\n",
+    "        train_loss = running_train_loss / len(train_loader)\n",
+    "        train_acc = correct_train / total_train * 100\n",
+    "\n",
+    "        print(f\"Epoch {epoch+1}/{num_epochs}\")\n",
+    "        print(f\"Train Loss: {train_loss:.4f} | Train Accuracy: {train_acc:.2f}%\")\n",
+    "        scheduler.step()  # Step the learning rate scheduler\n",
+    "\n",
+    "# Call the function\n",
+    "train_model(model, train_loader, criterion, optimizer, scheduler, num_epochs=15, device=device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d439485d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # Training Loop\n",
+    "# num_epochs = 15\n",
+    "# for epoch in range(start_epoch, num_epochs):\n",
+    "#     model.train()\n",
+    "#     running_loss = 0.0\n",
+    "#     correct = 0\n",
+    "#     total = 0\n",
+    "\n",
+    "#     for images, labels in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs}\"):\n",
+    "#         images, labels = images.to(device), labels.to(device)\n",
+    "\n",
+    "#         optimizer.zero_grad()\n",
+    "#         outputs = model(images)\n",
+    "#         loss = criterion(outputs, labels)\n",
+    "#         loss.backward()\n",
+    "\n",
+    "#         optimizer.step()\n",
+    "\n",
+    "#         running_loss += loss.item()\n",
+    "#         _, preds = torch.max(outputs, 1)\n",
+    "#         correct += (preds == labels).sum().item()\n",
+    "#         total += labels.size(0)\n",
+    "\n",
+    "#     # Step the learning rate scheduler\n",
+    "#     scheduler.step()\n",
+    "\n",
+    "#     # Calculate accuracy for this epoch\n",
+    "#     epoch_accuracy = 100 * correct / total\n",
+    "#     print(f\"Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss:.4f}, Accuracy: {epoch_accuracy:.2f}%\")\n",
+    "\n",
+    "#     # Save checkpoint at the end of each epoch\n",
+    "#     save_checkpoint(epoch, model, optimizer, scheduler, running_loss)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fef39141",
+   "metadata": {},
+   "source": [
+    "### Step 8. Validate and Evaluate the Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2090199a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def evaluate_on_test_as_val(model, test_loader, criterion, device='cuda'):\n",
+    "    model.eval()\n",
+    "    running_loss = 0.0\n",
+    "    correct = 0\n",
+    "    total = 0\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for inputs, labels in test_loader:\n",
+    "            inputs, labels = inputs.to(device), labels.to(device)\n",
+    "            outputs = model(inputs)\n",
+    "            loss = criterion(outputs, labels)\n",
+    "\n",
+    "            running_loss += loss.item()\n",
+    "            _, predicted = torch.max(outputs, 1)\n",
+    "            correct += (predicted == labels).sum().item()\n",
+    "            total += labels.size(0)\n",
+    "\n",
+    "    val_loss = running_loss / len(test_loader)\n",
+    "    val_acc = correct / total * 100\n",
+    "    return val_loss, val_acc\n",
+    "\n",
+    "val_loss, val_acc = evaluate_on_test_as_val(model, test_loader, criterion, device)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c5709926",
+   "metadata": {},
+   "source": [
+    "### Step 9. Visualize predictions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "45a03f67",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "\n",
+    "def visualize_predictions(model, dataloader, class_names, device='cuda', num_images=6):\n",
+    "    model.eval()\n",
+    "    images_shown = 0\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for inputs, labels in dataloader:\n",
+    "            inputs = inputs.to(device)\n",
+    "            outputs = model(inputs)\n",
+    "            _, preds = torch.max(outputs, 1)\n",
+    "\n",
+    "            inputs = inputs.cpu()\n",
+    "            labels = labels.cpu()\n",
+    "            preds = preds.cpu()\n",
+    "\n",
+    "            for i in range(inputs.size(0)):\n",
+    "                if images_shown >= num_images:\n",
+    "                    return\n",
+    "                img = inputs[i].permute(1, 2, 0).numpy()\n",
+    "                img = (img - img.min()) / (img.max() - img.min())  # normalize for display\n",
+    "\n",
+    "                plt.imshow(img)\n",
+    "                plt.title(f\"True: {class_names[labels[i]]}, Pred: {class_names[preds[i]]}\")\n",
+    "                plt.axis(\"off\")\n",
+    "                plt.show()\n",
+    "                images_shown += 1\n",
+    "\n",
+    "class_names = ['No DR', 'Mild', 'Moderate', 'Severe', 'Proliferative DR']\n",
+    "visualize_predictions(model, test_loader, class_names, device=device, num_images=6)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c2faf756",
+   "metadata": {},
+   "source": [
+    "### Step 10. Save The Model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "af93f096",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "torch.save(model.state_dict(), \"Densenet121.pth\")\n",
+    "print(\"✅ Final model saved as 'Densenet121.pth'\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}