segmentation_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CellViT-Hibou Model Usage Example\n",
    "\n",
    "This notebook showcases the basic usage of the CellViT-Hibou model for segmentation of cell nuclei."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import cv2\n",
    "from torchvision import transforms\n",
    "from hibou import CellViTHibou\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Download weights from [Huggingface](https://huggingface.co/histai/cellvit-hibou-l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt_path = \"INSERT_PATH_TO_CHECKPOINT_HERE\"\n",
    "model = CellViTHibou(\n",
    "    hibou_path=None,  # we don't need to load hibou model separately as it is already included in the checkpoint\n",
    "    num_nuclei_classes=6,\n",
    "    num_tissue_classes=19,\n",
    "    embed_dim=1024,\n",
    "    input_channels=3,\n",
    "    depth=24,\n",
    "    num_heads=16,\n",
    "    extract_layers=[6,12,18,24],\n",
    ")\n",
    "model.load_state_dict(torch.load(ckpt_path))\n",
    "model.eval()\n",
    "model = model.to(device)\n",
    "print(\"Model loaded successfully.\")\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.7068, 0.5755, 0.7220], std=[0.1950, 0.2316, 0.1816]),\n",
    "    transforms.Resize((256, 256))\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load the test image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "orig_image = cv2.imread(\"images/sample.png\", cv2.IMREAD_COLOR)\n",
    "orig_image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)\n",
    "plt.imshow(orig_image)\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Segment the image using the CellViT-Hibou model. Draw the segmentation mask on the image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "global_contour_visualization = orig_image.copy()\n",
    "\n",
    "for i in range(0, orig_image.shape[0], 256):\n",
    "    for j in range(0, orig_image.shape[1], 256):\n",
    "        image = orig_image[i:i+256, j:j+256]\n",
    "        image = transform(image).unsqueeze(0)\n",
    "        with torch.no_grad():\n",
    "            output = model(image.to(device))\n",
    "\n",
    "        output[\"nuclei_binary_map\"] = output[\"nuclei_binary_map\"].softmax(dim=1)\n",
    "        output[\"nuclei_type_map\"] = output[\"nuclei_type_map\"].softmax(dim=1)\n",
    "\n",
    "        for key in output.keys():\n",
    "            if isinstance(output[key], torch.Tensor):\n",
    "                output[key] = output[key].cpu()\n",
    "\n",
    "        (_, instance_types) = model.calculate_instance_map(output, magnification=20)\n",
    "        cells = instance_types[0]\n",
    "        for cell in cells.values():\n",
    "            contour = cell[\"contour\"]\n",
    "            contour = contour + np.array([j, i])\n",
    "            cv2.drawContours(global_contour_visualization, [contour], -1, (255, 0, 0), 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "orig_im_vis = cv2.addWeighted(orig_image, 0.5, global_contour_visualization, 0.5, 0)\n",
    "\n",
    "# Displaying the images\n",
    "plt.figure(figsize=(21, 11))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.title('Original Image with Contours')\n",
    "plt.imshow(orig_im_vis)\n",
    "plt.axis('off')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### In this simple example we used the model only to find the cell nuclei in the image. The model is also trained to predict nuclei types and tissue types. \n",
    "\n",
    "We can check all the outputs for the single cell and tissue type for the patch. For mapping between cell/tissue types and their numerical labels check `configs/dataset_config.yaml` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cell = next(iter(cells.values()))\n",
    "\n",
    "print(\"Cell information:\")\n",
    "for key, value in cell.items():\n",
    "    print(f\"{key}: {value}\")\n",
    "\n",
    "# there is also a tissue type map in the output dictionary\n",
    "print(\"\\nTissue information:\")\n",
    "print(f\"Tissue type: {output['tissue_types'].argmax(dim=1).item()}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "CELL",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}