add: sample visualization in tables

colabs/README.md

@@ -29,6 +29,7 @@
 | 🦄 Fine-tune a Torchvision Model with KerasCore | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/wandb/examples/blob/master/colabs/keras/keras_core/torchvision_keras.ipynb) |
 | 🦄 Fine-tune a Timm Model with KerasCore | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/wandb/examples/blob/master/colabs/keras/keras_core/timm_keras.ipynb) |
 | 🦄 Medical Image Classification Tutorial using MonAI and KerasCore | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/wandb/examples/blob/master/colabs/keras/keras_core/monai_medmnist_keras.ipynb) |
+| 🩻 Brain tumor 3D segmentation with MONAI and Weights & Biases | [![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/wandb/examples/blob/main/colabs/monai/3d_brain_tumor_segmentation.ipynb) |
 
 # 🏋🏽‍♂️ W&B Features
 

colabs/monai/3d_brain_tumor_segmentation.ipynb

@@ -6,6 +6,8 @@
    "source": [
     "# Brain tumor 3D segmentation with MONAI and Weights & Biases\n",
     "\n",
+    "[![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/wandb/examples/blob/main/colabs/monai/3d_brain_tumor_segmentation.ipynb)\n",
+    "\n",
     "This tutorial shows how to construct a training workflow of multi-labels 3D brain tumor segmentation task using [MONAI](https://github.com/Project-MONAI/MONAI) and use experiment tracking and data visualization features of [Weights & Biases](https://wandb.ai/site). The tutorial contains the following features:\n",
     "\n",
     "1. Initialize a Weights & Biases run and synchrozize all configs associated with the run for reproducibility.\n",
@@ -117,7 +119,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "wandb.init(project=\"monai-brain-tumor-segmentation\", job_type=\"test\")"
+    "wandb.init(project=\"monai-brain-tumor-segmentation\")"
    ]
   },
   {
@@ -676,7 +678,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Define a W&B Artifact object\n",
+    "artifact = wandb.Artifact(\n",
+    "    name=f\"{wandb.run.id}-checkpoint\", type=\"model\"\n",
+    ")\n",
+    "\n",
     "epoch_progress_bar = tqdm(range(config.max_train_epochs), desc=\"Training:\")\n",
+    "\n",
     "for epoch in epoch_progress_bar:\n",
     "    model.train()\n",
     "    epoch_loss = 0\n",
@@ -741,9 +749,6 @@
     "            torch.save(model.state_dict(), checkpoint_path)\n",
     "            \n",
     "            # Log and versison model checkpoints using W&B artifacts.\n",
-    "            artifact = wandb.Artifact(\n",
-    "                name=f\"{wandb.run.id}-checkpoint\", type=\"model\"\n",
-    "            )\n",
     "            artifact.add_file(local_path=checkpoint_path)\n",
     "            wandb.log_artifact(artifact, aliases=[f\"epoch_{epoch}\"])\n",
     "\n",
@@ -757,7 +762,11 @@
     "                    \"validation/mean_dice_enhanced_tumor\": metric_values_enhanced_tumor[-1],\n",
     "                }\n",
     "            )\n",
-    "            validation_step += 1"
+    "            validation_step += 1\n",
+    "\n",
+    "\n",
+    "# Wait for this artifact to finish logging\n",
+    "artifact.wait()"
    ]
   },
   {
@@ -773,7 +782,56 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 🔱 Inferece"
+    "If we navigate to the artifacts tab in the W&B run dashboard, we will be able to access the different versions of model checkpoint artifacts that we logged during training.\n",
+    "\n",
+    "![](./assets/viz-4.gif)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 🔱 Inference"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Using the artifacts interface, we can select which version of the artifact is the best model checkpoint, in this case, the mean epoch-wise training loss. We can also explore the entire lineage of the artifact and also use the version that we need.\n",
+    "\n",
+    "![](./assets/viz-5.gif)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let us fetch the version of the model artifact with the best epoch-wise mean training loss and load the checkpoint state dictionary to the model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_artifact = wandb.use_artifact(\n",
+    "    \"geekyrakshit/monai-brain-tumor-segmentation/d5ex6n4a-checkpoint:v49\",\n",
+    "    type=\"model\",\n",
+    ")\n",
+    "model_artifact_dir = model_artifact.download()\n",
+    "model.load_state_dict(torch.load(os.path.join(model_artifact_dir, \"model.pth\")))\n",
+    "model.eval()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 📸 Visualizing Predictions and Comparing with the Ground Truth Labels\n",
+    "\n",
+    "In order to visualize the predictions of the pre-trained model and compare them with the corresponding ground-truth segmentation mask using the interactive segmentation mask overlay, let us create another ultility function."
    ]
   },
   {
@@ -847,14 +905,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model_artifact = wandb.use_artifact(\n",
-    "    \"geekyrakshit/monai-brain-tumor-segmentation/d5ex6n4a-checkpoint:v49\",\n",
-    "    type=\"model\",\n",
-    ")\n",
-    "model_artifact_dir = model_artifact.download()\n",
-    "model.load_state_dict(torch.load(os.path.join(model_artifact_dir, \"model.pth\")))\n",
-    "model.eval()\n",
-    "\n",
+    "# create the prediction table\n",
     "prediction_table = wandb.Table(\n",
     "    columns=[\n",
     "        \"Split\",\n",
@@ -875,6 +926,7 @@
     "    ]\n",
     ")\n",
     "\n",
+    "# Perform inference and visualization\n",
     "with torch.no_grad():\n",
     "    config.max_prediction_images_visualized\n",
     "    max_samples = (\n",
@@ -900,7 +952,27 @@
     "            table=prediction_table,\n",
     "        )\n",
     "\n",
-    "    wandb.log({\"Predictions/Tumor-Segmentation-Data\": prediction_table})"
+    "    wandb.log({\"Predictions/Tumor-Segmentation-Data\": prediction_table})\n",
+    "\n",
+    "\n",
+    "# End the experiment\n",
+    "wandb.finish()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let us see how we can analyze and compare the predicted segmentation masks and the ground-truth labels for each class using the interactive segmentation mask overlay.\n",
+    "\n",
+    "![](./assets/viz-6.gif)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can also check out the report [Brain Tumor Segmentation using MONAI and WandB](https://wandb.ai/geekyrakshit/brain-tumor-segmentation/reports/Brain-Tumor-Segmentation-using-MONAI-and-WandB---Vmlldzo0MjUzODIw) for more details regarding training a brain-tumor segmentation model using MONAI and W&B."
    ]
   }
  ],