densnet_spider.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from templates import *\n",
    "from templates_cls import *\n",
    "import monai.transforms as T\n",
    "from monai.utils.misc import first\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from monai.networks.nets import DenseNet121\n",
    "from sklearn.metrics import mean_absolute_error, f1_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading dataset: 100%|██████████| 1185/1185 [15:29<00:00,  1.28it/s]\n"
     ]
    }
   ],
   "source": [
    "batch_size = 32\n",
    "\n",
    "train_spider_files = create_spider_files()\n",
    "train_transforms = T.Compose([\n",
    "        # each image has a single ivd label\n",
    "        T.LoadImaged(keys=['image', 'mask'], ensure_channel_first=True),\n",
    "        T.Orientationd(keys=['image', 'mask'], axcodes='RAS'),\n",
    "        # median dataset spacing\n",
    "        T.Spacingd(keys=['image', 'mask'],pixdim=(3.32, 0.625, 0.625), mode=(\"bilinear\", \"nearest\")),\n",
    "        T.ScaleIntensityRangePercentilesd(keys='image', lower=0, upper=99.5, b_min=0, b_max=1),\n",
    "        # remove other labels from mask\n",
    "        CropMaskByLabel(mask_key='mask', label_key='ivd_label', label_lambda_func=lambda x: x + 200),\n",
    "        # some augmentations\n",
    "        #T.RandGaussianNoised(keys=['image'], mean=0.0, std=0.015, prob=prob),\n",
    "        #T.RandRotated(keys=['image', 'mask'], range_x=30 * (np.pi / 180), mode=[\"bilinear\", \"nearest\"], prob=prob),\n",
    "        # center and crop image around ivd\n",
    "        T.CropForegroundd(keys=['image', 'mask'], source_key='mask', margin=(0, 80, 80), allow_smaller=False),\n",
    "        T.CenterSpatialCropd(keys=['image', 'mask'], roi_size=(-1, 80, 80)),\n",
    "        # get a single slice\n",
    "        T.CenterSpatialCropd(keys=['image', 'mask'], roi_size=(1, -1, -1)),\n",
    "        #T.RandSpatialCropd(keys=['image', 'mask'], roi_size=(1, -1, -1)),\n",
    "        AssertEmptyImaged(),\n",
    "        # resize\n",
    "        T.Resized(keys=['image', 'mask'], spatial_size=(1, 64, 64), anti_aliasing=True),\n",
    "        T.ToTensord(keys=['image', 'mask']),\n",
    "        T.SqueezeDimd(keys=['image', 'mask'], dim=1), \n",
    "])\n",
    "\n",
    "dataset_train = data.CacheDataset(train_spider_files, transform=train_transforms)\n",
    "train_loader = DataLoader(dataset_train, batch_size, shuffle=False, num_workers=8, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Missing grade, skipping IVD\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading dataset: 100%|██████████| 261/261 [02:00<00:00,  2.17it/s]\n"
     ]
    }
   ],
   "source": [
    "val_spider_files = create_spider_files(split=\"validation\")\n",
    "batch_size = 32\n",
    "val_transforms = T.Compose([\n",
    "        # each image has a single ivd label\n",
    "        T.LoadImaged(keys=['image', 'mask'], ensure_channel_first=True),\n",
    "        T.Orientationd(keys=['image', 'mask'], axcodes='RAS'),\n",
    "        # median dataset spacing\n",
    "        T.Spacingd(keys=['image', 'mask'],pixdim=(3.32, 0.625, 0.625), mode=(\"bilinear\", \"nearest\")),\n",
    "        T.ScaleIntensityRangePercentilesd(keys='image', lower=0, upper=99.5, b_min=0, b_max=1),\n",
    "        # remove other labels from mask\n",
    "        CropMaskByLabel(mask_key='mask', label_key='ivd_label', label_lambda_func=lambda x: x + 200),\n",
    "        # center and crop image around ivd\n",
    "        T.CropForegroundd(keys=['image', 'mask'], source_key='mask', margin=(0, 80, 80), allow_smaller=False),\n",
    "        T.CenterSpatialCropd(keys=['image', 'mask'], roi_size=(-1, 80, 80)),\n",
    "        # get a single slice\n",
    "        T.CenterSpatialCropd(keys=['image', 'mask'], roi_size=(1, -1, -1)),\n",
    "        AssertEmptyImaged(),\n",
    "        # resize\n",
    "        T.Resized(keys=['image', 'mask'], spatial_size=(1, 64, 64), anti_aliasing=True),\n",
    "        T.ToTensord(keys=['image', 'mask']),\n",
    "        T.SqueezeDimd(keys=['image', 'mask'], dim=1), \n",
    "])\n",
    "\n",
    "dataset_val = data.CacheDataset(val_spider_files, transform=val_transforms)\n",
    "\n",
    "val_loader = DataLoader(dataset_val, batch_size, shuffle=False, num_workers=8, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_spider_files = create_spider_files(split=\"test\")\n",
    "batch_size = 32\n",
    "val_transforms = T.Compose([\n",
    "        # each image has a single ivd label\n",
    "        T.LoadImaged(keys=['image', 'mask'], ensure_channel_first=True),\n",
    "        T.Orientationd(keys=['image', 'mask'], axcodes='RAS'),\n",
    "        # median dataset spacing\n",
    "        T.Spacingd(keys=['image', 'mask'],pixdim=(3.32, 0.625, 0.625), mode=(\"bilinear\", \"nearest\")),\n",
    "        T.ScaleIntensityRangePercentilesd(keys='image', lower=0, upper=99.5, b_min=0, b_max=1),\n",
    "        # remove other labels from mask\n",
    "        CropMaskByLabel(mask_key='mask', label_key='ivd_label', label_lambda_func=lambda x: x + 200),\n",
    "        # center and crop image around ivd\n",
    "        T.CropForegroundd(keys=['image', 'mask'], source_key='mask', margin=(0, 80, 80), allow_smaller=False),\n",
    "        T.CenterSpatialCropd(keys=['image', 'mask'], roi_size=(-1, 80, 80)),\n",
    "        # get a single slice\n",
    "        T.CenterSpatialCropd(keys=['image', 'mask'], roi_size=(1, -1, -1)),\n",
    "        AssertEmptyImaged(),\n",
    "        # resize\n",
    "        T.Resized(keys=['image', 'mask'], spatial_size=(1, 64, 64), anti_aliasing=True),\n",
    "        T.ToTensord(keys=['image', 'mask']),\n",
    "        T.SqueezeDimd(keys=['image', 'mask'], dim=1), \n",
    "])\n",
    "\n",
    "dataset_val = data.CacheDataset(val_spider_files, transform=val_transforms)\n",
    "\n",
    "val_loader = DataLoader(dataset_val, batch_size, shuffle=False, num_workers=8, pin_memory=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/100, Train Loss: 1.7215, Val Loss: 1.9843, Val MAE: 2.0728\n",
      "Epoch 2/100, Train Loss: 1.5297, Val Loss: 1.8285, Val MAE: 1.7625\n",
      "Epoch 3/100, Train Loss: 1.3695, Val Loss: 1.7144, Val MAE: 1.5977\n",
      "Epoch 4/100, Train Loss: 1.2271, Val Loss: 1.5794, Val MAE: 1.4866\n",
      "Epoch 5/100, Train Loss: 1.0990, Val Loss: 1.4542, Val MAE: 1.4176\n",
      "Epoch 6/100, Train Loss: 0.9809, Val Loss: 1.3546, Val MAE: 1.3678\n",
      "Epoch 7/100, Train Loss: 0.8794, Val Loss: 1.2732, Val MAE: 1.2644\n",
      "Epoch 8/100, Train Loss: 0.7857, Val Loss: 1.1813, Val MAE: 1.1648\n",
      "Epoch 9/100, Train Loss: 0.7033, Val Loss: 1.1009, Val MAE: 1.0920\n",
      "Epoch 10/100, Train Loss: 0.6310, Val Loss: 1.0574, Val MAE: 1.0575\n",
      "Epoch 11/100, Train Loss: 0.5675, Val Loss: 1.0387, Val MAE: 1.0192\n",
      "Epoch 12/100, Train Loss: 0.5126, Val Loss: 0.9766, Val MAE: 0.9579\n",
      "Epoch 13/100, Train Loss: 0.4665, Val Loss: 0.9199, Val MAE: 0.8966\n",
      "Epoch 14/100, Train Loss: 0.4256, Val Loss: 0.9255, Val MAE: 0.9042\n",
      "Epoch 15/100, Train Loss: 0.3949, Val Loss: 0.9323, Val MAE: 0.9234\n",
      "Epoch 16/100, Train Loss: 0.3682, Val Loss: 0.9166, Val MAE: 0.9004\n",
      "Epoch 17/100, Train Loss: 0.3469, Val Loss: 0.8669, Val MAE: 0.8467\n",
      "Epoch 18/100, Train Loss: 0.3297, Val Loss: 0.8357, Val MAE: 0.8314\n",
      "Epoch 19/100, Train Loss: 0.3093, Val Loss: 0.8442, Val MAE: 0.8199\n",
      "Epoch 20/100, Train Loss: 0.2910, Val Loss: 0.8706, Val MAE: 0.8621\n",
      "Epoch 21/100, Train Loss: 0.2816, Val Loss: 0.8903, Val MAE: 0.8774\n",
      "Epoch 22/100, Train Loss: 0.2669, Val Loss: 0.8443, Val MAE: 0.8429\n",
      "Epoch 23/100, Train Loss: 0.2611, Val Loss: 0.8155, Val MAE: 0.8352\n",
      "Epoch 24/100, Train Loss: 0.2634, Val Loss: 0.8046, Val MAE: 0.7969\n",
      "Epoch 25/100, Train Loss: 0.2504, Val Loss: 0.8381, Val MAE: 0.8276\n",
      "Epoch 26/100, Train Loss: 0.2459, Val Loss: 0.8206, Val MAE: 0.8084\n",
      "Epoch 27/100, Train Loss: 0.2281, Val Loss: 0.8216, Val MAE: 0.8276\n",
      "Epoch 28/100, Train Loss: 0.2177, Val Loss: 0.8235, Val MAE: 0.8238\n",
      "Epoch 29/100, Train Loss: 0.2109, Val Loss: 0.8210, Val MAE: 0.8238\n",
      "Epoch 30/100, Train Loss: 0.2058, Val Loss: 0.8202, Val MAE: 0.8161\n",
      "Epoch 31/100, Train Loss: 0.2013, Val Loss: 0.8186, Val MAE: 0.8123\n",
      "Epoch 32/100, Train Loss: 0.1981, Val Loss: 0.8192, Val MAE: 0.8123\n",
      "Epoch 33/100, Train Loss: 0.1953, Val Loss: 0.8180, Val MAE: 0.8084\n",
      "Epoch 34/100, Train Loss: 0.1929, Val Loss: 0.8164, Val MAE: 0.8046\n",
      "Early stopping!\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from monai.networks.nets import DenseNet121\n",
    "from sklearn.metrics import mean_absolute_error\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# Model setup\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=1).to(device)\n",
    "loss_function = nn.L1Loss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-5)\n",
    "\n",
    "# Learning Rate Scheduler\n",
    "scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=25, gamma=0.1)\n",
    "\n",
    "# Early Stopping Setup\n",
    "early_stopping_patience = 10\n",
    "best_val_loss = float('inf')\n",
    "epochs_no_improve = 0\n",
    "\n",
    "max_epochs = 100\n",
    "for epoch in range(max_epochs):\n",
    "    model.train()\n",
    "    train_loss = 0.0\n",
    "    for batch_data in train_loader:\n",
    "        images, labels = batch_data['image'].to(device), batch_data[\"pfirrman_grade\"].float().to(device)\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(images)\n",
    "        loss = loss_function(outputs, labels.unsqueeze(1))\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_loss += loss.item()\n",
    "    train_loss /= len(train_loader)\n",
    "\n",
    "    # Learning rate scheduler step\n",
    "    scheduler.step()\n",
    "\n",
    "    # Validation loop\n",
    "    model.eval()\n",
    "    val_loss = 0.0\n",
    "    total_preds = []\n",
    "    total_labels = []\n",
    "    with torch.no_grad():\n",
    "        for batch_data in val_loader:\n",
    "            images, labels = batch_data['image'].to(device), batch_data['pfirrman_grade'].float().to(device)\n",
    "            outputs = model(images)\n",
    "            loss = loss_function(outputs, labels.unsqueeze(1))\n",
    "            val_loss += loss.item()\n",
    "\n",
    "            preds = outputs.round()\n",
    "            total_preds.extend(preds.cpu().numpy())\n",
    "            total_labels.extend(labels.cpu().numpy())\n",
    "    val_loss /= len(val_loader)\n",
    "    val_mae = mean_absolute_error(total_labels, total_preds)\n",
    "\n",
    "    print(f'Epoch {epoch+1}/{max_epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Val MAE: {val_mae:.4f}')\n",
    "\n",
    "    # Early Stopping Check\n",
    "    if val_loss < best_val_loss:\n",
    "        best_val_loss = val_loss\n",
    "        epochs_no_improve = 0\n",
    "    else:\n",
    "        epochs_no_improve += 1\n",
    "        if epochs_no_improve == early_stopping_patience:\n",
    "            print('Early stopping!')\n",
    "            break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the model state dictionary\n",
    "torch.save(model.state_dict(), \"spider_resnet_reg.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MAE:  0.8309463\n",
      "F1:  0.3000367846158562\n"
     ]
    }
   ],
   "source": [
    "model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=1).to(device)\n",
    "model.load_state_dict(torch.load(\"resent/spider_resnet_reg.pt\"))\n",
    "model.eval()\n",
    "\n",
    "total_preds = []\n",
    "total_labels = []\n",
    "with torch.no_grad():\n",
    "    for batch_data in val_loader:\n",
    "        images, labels = batch_data['image'].to(device), batch_data['pfirrman_grade'].float().to(device)\n",
    "        preds = model(images)\n",
    "        total_preds.extend(preds.cpu().numpy())\n",
    "        total_labels.extend(labels.cpu().numpy())\n",
    "\n",
    "val_mae = mean_absolute_error(total_labels, total_preds)\n",
    "val_f1 = f1_score(total_labels, np.round(total_preds), average=\"macro\")\n",
    "\n",
    "print(\"MAE: \", val_mae)\n",
    "print(\"F1: \", val_f1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/100, Train Loss: 0.5686, Val Loss: 0.5228, Val ACC: 0.8659\n",
      "Epoch 2/100, Train Loss: 0.4911, Val Loss: 0.4779, Val ACC: 0.8659\n",
      "Epoch 3/100, Train Loss: 0.4415, Val Loss: 0.4424, Val ACC: 0.8659\n",
      "Epoch 4/100, Train Loss: 0.4036, Val Loss: 0.4169, Val ACC: 0.8659\n",
      "Epoch 5/100, Train Loss: 0.3697, Val Loss: 0.3993, Val ACC: 0.8659\n",
      "Epoch 6/100, Train Loss: 0.3400, Val Loss: 0.3826, Val ACC: 0.8659\n",
      "Epoch 7/100, Train Loss: 0.3139, Val Loss: 0.3688, Val ACC: 0.8659\n",
      "Epoch 8/100, Train Loss: 0.2918, Val Loss: 0.3580, Val ACC: 0.8659\n",
      "Epoch 9/100, Train Loss: 0.2718, Val Loss: 0.3492, Val ACC: 0.8659\n",
      "Epoch 10/100, Train Loss: 0.2501, Val Loss: 0.3395, Val ACC: 0.8659\n",
      "Epoch 11/100, Train Loss: 0.2307, Val Loss: 0.3338, Val ACC: 0.8621\n",
      "Epoch 12/100, Train Loss: 0.2129, Val Loss: 0.3285, Val ACC: 0.8621\n",
      "Epoch 13/100, Train Loss: 0.1972, Val Loss: 0.3230, Val ACC: 0.8659\n",
      "Epoch 14/100, Train Loss: 0.1817, Val Loss: 0.3165, Val ACC: 0.8621\n",
      "Epoch 15/100, Train Loss: 0.1670, Val Loss: 0.3127, Val ACC: 0.8659\n",
      "Epoch 16/100, Train Loss: 0.1567, Val Loss: 0.3075, Val ACC: 0.8659\n",
      "Epoch 17/100, Train Loss: 0.1431, Val Loss: 0.3052, Val ACC: 0.8659\n",
      "Epoch 18/100, Train Loss: 0.1314, Val Loss: 0.3015, Val ACC: 0.8659\n",
      "Epoch 19/100, Train Loss: 0.1212, Val Loss: 0.2987, Val ACC: 0.8659\n",
      "Epoch 20/100, Train Loss: 0.1120, Val Loss: 0.2967, Val ACC: 0.8659\n",
      "Epoch 21/100, Train Loss: 0.1040, Val Loss: 0.2952, Val ACC: 0.8659\n",
      "Epoch 22/100, Train Loss: 0.0967, Val Loss: 0.2937, Val ACC: 0.8697\n",
      "Epoch 23/100, Train Loss: 0.0902, Val Loss: 0.2929, Val ACC: 0.8659\n",
      "Epoch 24/100, Train Loss: 0.0841, Val Loss: 0.2930, Val ACC: 0.8659\n",
      "Epoch 25/100, Train Loss: 0.0784, Val Loss: 0.2923, Val ACC: 0.8659\n",
      "Epoch 26/100, Train Loss: 0.0744, Val Loss: 0.2918, Val ACC: 0.8697\n",
      "Epoch 27/100, Train Loss: 0.0731, Val Loss: 0.2918, Val ACC: 0.8697\n",
      "Epoch 28/100, Train Loss: 0.0724, Val Loss: 0.2917, Val ACC: 0.8697\n",
      "Epoch 29/100, Train Loss: 0.0719, Val Loss: 0.2916, Val ACC: 0.8697\n",
      "Epoch 30/100, Train Loss: 0.0714, Val Loss: 0.2915, Val ACC: 0.8697\n",
      "Epoch 31/100, Train Loss: 0.0708, Val Loss: 0.2914, Val ACC: 0.8659\n",
      "Epoch 32/100, Train Loss: 0.0703, Val Loss: 0.2914, Val ACC: 0.8659\n",
      "Epoch 33/100, Train Loss: 0.0698, Val Loss: 0.2913, Val ACC: 0.8659\n",
      "Epoch 34/100, Train Loss: 0.0693, Val Loss: 0.2914, Val ACC: 0.8659\n",
      "Epoch 35/100, Train Loss: 0.0687, Val Loss: 0.2913, Val ACC: 0.8659\n",
      "Epoch 36/100, Train Loss: 0.0682, Val Loss: 0.2912, Val ACC: 0.8659\n",
      "Epoch 37/100, Train Loss: 0.0677, Val Loss: 0.2912, Val ACC: 0.8659\n",
      "Epoch 38/100, Train Loss: 0.0672, Val Loss: 0.2911, Val ACC: 0.8659\n",
      "Epoch 39/100, Train Loss: 0.0667, Val Loss: 0.2910, Val ACC: 0.8659\n",
      "Epoch 40/100, Train Loss: 0.0661, Val Loss: 0.2910, Val ACC: 0.8621\n",
      "Epoch 41/100, Train Loss: 0.0656, Val Loss: 0.2910, Val ACC: 0.8621\n",
      "Epoch 42/100, Train Loss: 0.0651, Val Loss: 0.2910, Val ACC: 0.8621\n",
      "Epoch 43/100, Train Loss: 0.0646, Val Loss: 0.2908, Val ACC: 0.8621\n",
      "Epoch 44/100, Train Loss: 0.0641, Val Loss: 0.2909, Val ACC: 0.8621\n",
      "Epoch 45/100, Train Loss: 0.0636, Val Loss: 0.2909, Val ACC: 0.8621\n",
      "Epoch 46/100, Train Loss: 0.0631, Val Loss: 0.2907, Val ACC: 0.8659\n",
      "Epoch 47/100, Train Loss: 0.0626, Val Loss: 0.2908, Val ACC: 0.8659\n",
      "Epoch 48/100, Train Loss: 0.0620, Val Loss: 0.2908, Val ACC: 0.8659\n",
      "Epoch 49/100, Train Loss: 0.0615, Val Loss: 0.2907, Val ACC: 0.8659\n",
      "Epoch 50/100, Train Loss: 0.0610, Val Loss: 0.2908, Val ACC: 0.8659\n",
      "Epoch 51/100, Train Loss: 0.0606, Val Loss: 0.2908, Val ACC: 0.8697\n",
      "Epoch 52/100, Train Loss: 0.0606, Val Loss: 0.2908, Val ACC: 0.8697\n",
      "Epoch 53/100, Train Loss: 0.0605, Val Loss: 0.2908, Val ACC: 0.8697\n",
      "Epoch 54/100, Train Loss: 0.0604, Val Loss: 0.2907, Val ACC: 0.8697\n",
      "Epoch 55/100, Train Loss: 0.0604, Val Loss: 0.2907, Val ACC: 0.8697\n",
      "Epoch 56/100, Train Loss: 0.0603, Val Loss: 0.2907, Val ACC: 0.8697\n",
      "Epoch 57/100, Train Loss: 0.0603, Val Loss: 0.2907, Val ACC: 0.8697\n",
      "Epoch 58/100, Train Loss: 0.0602, Val Loss: 0.2908, Val ACC: 0.8697\n",
      "Epoch 59/100, Train Loss: 0.0602, Val Loss: 0.2907, Val ACC: 0.8697\n",
      "Early stopping!\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from monai.networks.nets import DenseNet121\n",
    "from sklearn.metrics import accuracy_score, f1_score\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# Model setup\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=1).to(device)\n",
    "loss_function = nn.BCEWithLogitsLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-5)\n",
    "\n",
    "# Learning Rate Scheduler\n",
    "scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=25, gamma=0.1)\n",
    "\n",
    "# Early Stopping Setup\n",
    "early_stopping_patience = 10\n",
    "best_val_loss = float('inf')\n",
    "epochs_no_improve = 0\n",
    "\n",
    "max_epochs = 100\n",
    "for epoch in range(max_epochs):\n",
    "    model.train()\n",
    "    train_loss = 0.0\n",
    "    for batch_data in train_loader:\n",
    "        images, labels = batch_data['image'].to(device), batch_data[\"pfirrman_grade\"].to(device)\n",
    "        labels = (labels > 0)\n",
    "        labels = labels.float()\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(images)\n",
    "        loss = loss_function(outputs, labels.unsqueeze(1))\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_loss += loss.item()\n",
    "    train_loss /= len(train_loader)\n",
    "\n",
    "    # Learning rate scheduler step\n",
    "    scheduler.step()\n",
    "\n",
    "    # Validation loop\n",
    "    model.eval()\n",
    "    val_loss = 0.0\n",
    "    total_preds = []\n",
    "    total_labels = []\n",
    "    with torch.no_grad():\n",
    "        for batch_data in val_loader:\n",
    "            images, labels = batch_data['image'].to(device), batch_data['pfirrman_grade'].to(device)\n",
    "            labels = (labels > 0)\n",
    "            labels = labels.float()\n",
    "            outputs = model(images)\n",
    "            loss = loss_function(outputs, labels.unsqueeze(1))\n",
    "            val_loss += loss.item()\n",
    "\n",
    "            preds = torch.sigmoid(outputs)\n",
    "            total_preds.extend(preds.cpu().numpy())\n",
    "            total_labels.extend(labels.cpu().numpy())\n",
    "    val_loss /= len(val_loader)\n",
    "    val_acc = accuracy_score(total_labels, np.round(total_preds))\n",
    "\n",
    "    print(f'Epoch {epoch+1}/{max_epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Val ACC: {val_acc:.4f}')\n",
    "\n",
    "    # Early Stopping Check\n",
    "    if val_loss < best_val_loss:\n",
    "        best_val_loss = val_loss\n",
    "        epochs_no_improve = 0\n",
    "    else:\n",
    "        epochs_no_improve += 1\n",
    "        if epochs_no_improve == early_stopping_patience:\n",
    "            print('Early stopping!')\n",
    "            break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the model state dictionary\n",
    "torch.save(model.state_dict(), \"resent/spider_resnet_cls.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AUC:  0.6108723135271807\n",
      "F1:  0.9276595744680851\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import roc_auc_score, f1_score\n",
    "\n",
    "model = DenseNet121(spatial_dims=2, in_channels=1, out_channels=1).to(device)\n",
    "model.load_state_dict(torch.load(\"resent/spider_resnet_cls.pt\"))\n",
    "model.eval()\n",
    "\n",
    "total_preds = []\n",
    "total_labels = []\n",
    "with torch.no_grad():\n",
    "    for batch_data in val_loader:\n",
    "        images, labels = batch_data['image'].to(device), batch_data['pfirrman_grade'].to(device)\n",
    "        labels = (labels > 0)\n",
    "        labels = labels.float()\n",
    "        outputs = model(images)\n",
    "        preds = torch.sigmoid(outputs).round()\n",
    "        total_preds.extend(preds.cpu().numpy())\n",
    "        total_labels.extend(labels.cpu().numpy())\n",
    "\n",
    "auc = roc_auc_score(total_labels, total_preds)\n",
    "f1 = f1_score(total_labels, total_preds)\n",
    "\n",
    "print(\"AUC: \", auc)\n",
    "print(\"F1: \", f1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "911026d04c504ad5e0c4c11f5deabc0fa44aa59b5237d77193e02811cb4f84ca"
  },
  "kernelspec": {
   "display_name": "Python 3.6.9 64-bit ('myenv': virtualenv)",
   "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.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}