2D_predictions_with_vis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43d13ad5-f188-4b49-be18-e590ad332901",
   "metadata": {},
   "outputs": [],
   "source": [
    "#from segment_anything import SamPredictor, sam_model_registry\n",
    "from models.sam import SamPredictor, sam_model_registry\n",
    "from models.sam.utils.transforms import ResizeLongestSide\n",
    "from skimage.measure import label\n",
    "from models.sam_LoRa import LoRA_Sam\n",
    "#Scientific computing \n",
    "import numpy as np\n",
    "import os\n",
    "#Pytorch packages\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "from torchvision import datasets\n",
    "#Visulization\n",
    "import matplotlib.pyplot as plt\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "#Others\n",
    "from torch.utils.data import DataLoader, Subset\n",
    "from torch.autograd import Variable\n",
    "import matplotlib.pyplot as plt\n",
    "import copy\n",
    "from utils.dataset import Public_dataset\n",
    "import torch.nn.functional as F\n",
    "from torch.nn.functional import one_hot\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "from utils.losses import DiceLoss\n",
    "from utils.dsc import dice_coeff\n",
    "import cv2\n",
    "import monai\n",
    "from utils.utils import vis_image,inverse_normalize,torch_percentile\n",
    "from argparse import Namespace\n",
    "import cfg\n",
    "import PIL\n",
    "import torchio as tio\n",
    "import json\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "arch=\"vit_t\"  # Change this value as needed\n",
    "finetune_type=\"adapter\"\n",
    "dataset_name=\"JHU_prj\"  # Assuming you set this if it's dynamic\n",
    "\n",
    "# Construct the checkpoint directory argument\n",
    "checkpoint_dir= f\"2D-SAM_{arch}_encoderdecoder_{finetune_type}_{dataset_name}_noprompt\"\n",
    "\n",
    "args_path = f\"{checkpoint_dir}/args.json\"\n",
    "\n",
    "# Reading the args from the json file\n",
    "with open(args_path, 'r') as f:\n",
    "    args_dict = json.load(f)\n",
    "\n",
    "\n",
    "# Converting dictionary to Namespace\n",
    "args = Namespace(**args_dict)\n",
    "\n",
    "sam_fine_tune = sam_model_registry[args.arch](args,checkpoint=os.path.join(args.dir_checkpoint,'checkpoint_best.pth'),num_classes=args.num_cls)\n",
    "sam_fine_tune = sam_fine_tune.to('cuda').eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22f35640-a2e5-4f0d-a59d-0fa891062cf7",
   "metadata": {},
   "source": [
    "## evaluate a image volume and save predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1654874-2247-40c0-81f1-1e19dc400c7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to evaluate a single image slice\n",
    "def evaluate_1_slice(image_path, model):\n",
    "    \"\"\"\n",
    "    Evaluates a single image slice using the provided model.\n",
    "\n",
    "    Parameters:\n",
    "    - image_path: Path to the image slice file.\n",
    "    - model: The model used for evaluation.\n",
    "\n",
    "    Returns:\n",
    "    - ori_img: The original image after normalization.\n",
    "    - pred: The prediction from the model.\n",
    "    - Pil_img: The PIL image of the original slice.\n",
    "    \"\"\"\n",
    "    # Load the image\n",
    "    img = Image.open(image_path).convert('RGB')\n",
    "    Pil_img = img.copy()\n",
    "    \n",
    "    # Resize the image to 1024x1024\n",
    "    img = transforms.Resize((1024, 1024))(img)\n",
    "    \n",
    "    # Transform the image to a tensor and normalize\n",
    "    transform_img = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "    ])\n",
    "    img = transform_img(img)\n",
    "    imgs = torch.unsqueeze(transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img), 0).cuda()\n",
    "    \n",
    "    # Perform model inference without gradient calculation\n",
    "    with torch.no_grad():\n",
    "        # Get image embeddings from the image encoder\n",
    "        img_emb = model.image_encoder(imgs)\n",
    "        \n",
    "        # Get sparse and dense embeddings from the prompt encoder\n",
    "        sparse_emb, dense_emb = model.prompt_encoder(\n",
    "            points=None,\n",
    "            boxes=None,\n",
    "            masks=None,\n",
    "        )\n",
    "        \n",
    "        # Get the prediction from the mask decoder\n",
    "        pred, _ = model.mask_decoder(\n",
    "            image_embeddings=img_emb,\n",
    "            image_pe=model.prompt_encoder.get_dense_pe(),\n",
    "            sparse_prompt_embeddings=sparse_emb,\n",
    "            dense_prompt_embeddings=dense_emb,\n",
    "            multimask_output=True,\n",
    "        )\n",
    "        \n",
    "        # Get the most likely prediction\n",
    "        pred = pred.argmax(dim=1)\n",
    "    \n",
    "    # Get the original image after normalization\n",
    "    ori_img = inverse_normalize(imgs.cpu()[0])\n",
    "    \n",
    "    return ori_img, pred, Pil_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c2bf140-fd2c-48aa-8a84-0a5a1dd278af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define paths\n",
    "image_path = 'path_to_your_image_slice.png'  # Replace with the path to your image slice\n",
    "output_dir = 'output_predictions'  # Directory to save predictions\n",
    "\n",
    "# Create the output directory if it doesn't exist\n",
    "Path(output_dir).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# Evaluate the image slice\n",
    "ori_img, pred_1, Pil_img1 = evaluate_1_slice(image_path, model)\n",
    "\n",
    "# Convert the predicted mask to a PIL image\n",
    "mask_pred_1 = ((pred_1).cpu()).float()\n",
    "pil_mask1 = Image.fromarray(np.array(mask_pred_1[0], dtype=np.uint8), 'L').resize(Pil_img1.size, resample=Image.NEAREST)\n",
    "\n",
    "# Save the original image and the predicted mask as PNG\n",
    "ori_img_filename = os.path.join(output_dir, 'original_image.png')\n",
    "mask_img_filename = os.path.join(output_dir, 'predicted_mask.png')\n",
    "\n",
    "Pil_img1.save(ori_img_filename)\n",
    "pil_mask1.save(mask_img_filename)\n",
    "\n",
    "print(f\"Original image saved to {ori_img_filename}\")\n",
    "print(f\"Predicted mask saved to {mask_img_filename}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d415e43d-012b-45bf-91c2-55f671830b6d",
   "metadata": {},
   "source": [
    "## visualize results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9105aa1-da83-4b26-81f0-a445240cea3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and display the original image and the predicted mask\n",
    "fig, axes = plt.subplots(1, 2, figsize=(12, 6))\n",
    "\n",
    "# Display the original slice\n",
    "ori_img_display = np.array(Pil_img1)\n",
    "axes[0].imshow(ori_img_display)\n",
    "axes[0].set_title(\"Original Slice\")\n",
    "axes[0].axis('off')\n",
    "\n",
    "# Display the predicted mask\n",
    "mask_display = np.array(pil_mask1)\n",
    "axes[1].imshow(mask_display, cmap='gray')\n",
    "axes[1].set_title(\"Predicted Mask\")\n",
    "axes[1].axis('off')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60cb1056-fe48-4569-a3dc-7dce67265a59",
   "metadata": {},
   "source": [
    "## visualize results as well as ground truth if ground truth path is provided as well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31c4313a-6955-4233-a004-963dc74c931e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ground_truth_path = 'path_to_your_ground_truth.png'  # Replace with the path to your ground truth mask if available\n",
    "# Visualize the results\n",
    "fig, axes = plt.subplots(1, 3 if ground_truth_path else 2, figsize=(18, 6) if ground_truth_path else (12, 6))\n",
    "\n",
    "# Display the original slice\n",
    "ori_img_display = np.array(Pil_img1)\n",
    "axes[0].imshow(ori_img_display)\n",
    "axes[0].set_title(\"Original Slice\")\n",
    "axes[0].axis('off')\n",
    "\n",
    "# Display the predicted mask\n",
    "mask_display = np.array(pil_mask1)\n",
    "axes[1].imshow(mask_display, cmap='gray')\n",
    "axes[1].set_title(\"Predicted Mask\")\n",
    "axes[1].axis('off')\n",
    "\n",
    "# Display the ground truth mask if provided\n",
    "if ground_truth_path:\n",
    "    ground_truth_img = Image.open(ground_truth_path).convert('L').resize(Pil_img1.size, resample=Image.NEAREST)\n",
    "    ground_truth_display = np.array(ground_truth_img)\n",
    "    axes[2].imshow(ground_truth_display, cmap='gray')\n",
    "    axes[2].set_title(\"Ground Truth Mask\")\n",
    "    axes[2].axis('off')\n",
    "\n",
    "    # Compute and print the Dice Similarity Coefficient for each class\n",
    "    num_classes = 2  # Replace with the actual number of classes\n",
    "    cls_dsc = [0] * num_classes\n",
    "    for cls in range(num_classes):\n",
    "        mask_pred_cls = (mask_pred_1 == cls).float()\n",
    "        mask_gt_cls = (ground_truth_display == cls).float()\n",
    "        cls_dsc[cls] = dice_coeff(mask_pred_cls, mask_gt_cls).item()\n",
    "        print(f\"Dice Similarity Coefficient for class {cls}: {cls_dsc[cls]}\")\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}